diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index 314b92643072..6a64c469411d 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -87,6 +87,10 @@ list(FILTER SRC_CU EXCLUDE REGEX "fusedSampling/.*") # Marlin is built as its own architecture-scoped OBJECT library below. list(FILTER SRC_CPP EXCLUDE REGEX "marlin/.*") list(FILTER SRC_CU EXCLUDE REGEX "marlin/.*") +# The MiniMax-M3 decode attention kernels are built as their own +# architecture-scoped OBJECT library below. +list(FILTER SRC_CU EXCLUDE REGEX + "minimaxM3/(fusedDecodeAttn|attnMega|attnC1)\\.cu$") if(NOT ENABLE_MULTI_DEVICE) list(FILTER SRC_CU EXCLUDE REGEX "customAllReduceKernels*.*cu$") @@ -126,10 +130,22 @@ if(MARLIN_SRC) set_cuda_architectures(marlin_src 89 90 120f) endif() +# MiniMax-M3 decode attention: block-scaled tcgen05 MMAs (.scale_vec::1X) exist +# on the arch-specific sm_100a / sm_103a targets only, not on the sm_100f family +# target the rest of this library builds. Other targets compile kernels that +# trap. +add_library( + minimax_m3_attn_src OBJECT minimaxM3/fusedDecodeAttn.cu minimaxM3/attnMega.cu + minimaxM3/attnC1.cu) +set_property(TARGET minimax_m3_attn_src PROPERTY POSITION_INDEPENDENT_CODE ON) +set_property(TARGET minimax_m3_attn_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) +set_cuda_architectures(minimax_m3_attn_src 100 103) + add_library( kernels_src STATIC ${SRC_CPP} ${SRC_CU} ${MOE_KERNELS_SRC} - $<$:$>) + $<$:$> + $) set_property(TARGET kernels_src PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET kernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) target_include_directories( diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/attnC1.cu b/cpp/tensorrt_llm/kernels/minimaxM3/attnC1.cu new file mode 100644 index 000000000000..e375b6c994fd --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/attnC1.cu @@ -0,0 +1,3846 @@ +/* + * 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. + */ + +// MiniMax-M3 attention + o_proj kernel (c1): fused decode attention and o_proj GEMV in one 96-CTA launch. + +#include "tensorrt_llm/kernels/minimaxM3/attnC1.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// The kernel uses tcgen05/TMEM (including block-scaled MMAs, which family targets lack), cluster DSMEM and E2M1 +// conversions: its body is compiled for the arch-specific sm_100a / sm_103a targets only (see the minimax_m3_attn_src +// library in kernels/CMakeLists.txt). Any other target compiles a kernel that traps, so a launch on a build without +// those targets fails loudly instead of returning untouched outputs. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030) \ + && (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FEAT_SM100_ALL) || defined(__CUDA_ARCH_FEAT_SM103_ALL)) +#define MINIMAX_M3_ATTN_DEVICE_ENABLED 1 +#else +#define MINIMAX_M3_ATTN_DEVICE_ENABLED 0 +#endif + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kHeadDim = 128; +constexpr int kRotDim = 64; +constexpr int kPage = 128; +constexpr int kHq = 16; +constexpr int kColK = kHq; // packed head slots: [Q x16 | K | V | idxQ | idxK] +constexpr int kColV = kHq + 1; +constexpr int kColIdxK = kHq + 3; +constexpr int kBlocks = 16; // KV blocks per request == topk == cluster size +constexpr int kComputeWarps = 16; // QK / softmax / PV and the dequant (one 16-byte chunk per thread) +constexpr int kWarps = kComputeWarps + 4; // plus producer-only warps, so new K/V/index-K never queue behind a Q head +constexpr int kThreads = kWarps * 32; +constexpr int kComputeThreads = kComputeWarps * 32; +constexpr int kOPitch = 132; // floats per partial-O row +constexpr float kNegInf = -__builtin_huge_valf(); + +using Params = MinimaxM3AttnC1Params; + +// Optional per-CTA clock64 probes for development (-DM3FA_PROFILE and a prof tensor): [CTA][64] timestamps. +// -DM3FA_PROFILE_GT records %globaltimer (ns, comparable across SMs and kernels) instead of clock64. +#ifdef M3FA_PROFILE_GT +#define M3FA_PROFILE +#define M3FA_NOW(v) asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(v)) +#else +#define M3FA_NOW(v) v = clock64() +#endif +#ifdef M3FA_PROFILE +#define M3FA_PROF_T(i, thread) \ + do \ + { \ + if (p.prof != nullptr && threadIdx.x == (thread)) \ + { \ + unsigned long long now_; \ + M3FA_NOW(now_); \ + p.prof[blockIdx.x * 64 + (i)] = now_; \ + } \ + } while (0) +#else +#define M3FA_PROF_T(i, thread) \ + do \ + { \ + } while (0) +#endif +#define M3FA_PROF(i) M3FA_PROF_T(i, 0) +// clock64 probes (development): slot i of this CTA, written by every calling thread. +#ifdef M3FA_PROFILE +#define M3FA_CLK(i) \ + do \ + { \ + if (p.prof != nullptr) \ + { \ + p.prof[blockIdx.x * 64 + (i)] = clock64(); \ + } \ + } while (0) +#else +#define M3FA_CLK(i) \ + do \ + { \ + } while (0) +#endif + +#ifndef M3FA_MEGA_EARLYCLW +#define M3FA_MEGA_EARLYCLW 1 // pass the cluster barrier (peers' merge mbarriers initialized) early, not before the push +#endif +#ifndef M3FA_V15_EARLYV +#define M3FA_V15_EARLYV 0 // packed path: dequantize V before the PDL wait (switch; off by default) +#endif +#ifndef M3FA_MEGA_TRIGGER +// Megakernel PDL release: 0 o_proj CTAs after their activation (attention CTAs at exit); 1 as 0 with the attention +// CTAs after their KV reads; 2 every CTA right after its griddepcontrol.wait; 3 o_proj CTAs after their epilogue. +#define M3FA_MEGA_TRIGGER 3 +#endif +#ifndef M3FA_TRIGGER_AT +#define M3FA_TRIGGER_AT 2 // where the PDL dependents are released: 0 entry, 1 after the wait, 2 after the KV reads +#endif + +// Generic-proxy shared-memory writes (fp16 tiles, Q, P) become visible to the tensor core (async proxy). +#define M3FA_FENCE_PROXY_ASYNC() asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory") + +__device__ __forceinline__ void cp_async16(void* dst, void const* src) +{ + uint32_t const d = static_cast(__cvta_generic_to_shared(dst)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(d), "l"(src) : "memory"); +} + +__device__ __forceinline__ float ex2(float x) +{ + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ float lg2(float x) +{ + float y; + asm("lg2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +// One packed E2M1 byte -> half2 (low nibble in the low half). +__device__ __forceinline__ uint32_t e2m1x2_to_h2(uint32_t byte) +{ + uint32_t out; + asm("{\n .reg .b8 src;\n mov.b32 {src, _, _, _}, %1;\n cvt.rn.f16x2.e2m1x2 %0, src;\n}\n" : "=r"(out) : "r"(byte)); + return out; +} + +__device__ __forceinline__ uint32_t hmul2_u(uint32_t a, uint32_t b) +{ + __half2 r = __hmul2(*reinterpret_cast<__half2*>(&a), *reinterpret_cast<__half2*>(&b)); + return *reinterpret_cast(&r); +} + +__device__ __forceinline__ uint32_t e4m3_to_h2_splat(uint32_t byte) +{ + __half_raw const h = __nv_cvt_fp8_to_halfraw(static_cast<__nv_fp8_storage_t>(byte), __NV_E4M3); + return static_cast(h.x) | (static_cast(h.x) << 16); +} + +__device__ __forceinline__ uint32_t e4m3x2_to_h2(uint32_t two_bytes) +{ + __half2_raw const h = __nv_cvt_fp8x2_to_halfraw2(static_cast<__nv_fp8x2_storage_t>(two_bytes), __NV_E4M3); + return static_cast(h.x) | (static_cast(h.y) << 16); +} + +// ---- Exact copies of the C++ producer's helpers (fusedQKNormRopeKernel.cu / quantization.cuh) ---- +template +__device__ __forceinline__ uint32_t cvt_sf_group_shfl_mask() +{ + constexpr uint32_t groupSize = static_cast(NUM_THREADS_PER_SF); + constexpr uint32_t groupMask = 0xFFFFFFFFU >> (32U - groupSize); + uint32_t laneId = 0; + asm("mov.u32 %0, %%laneid;" : "=r"(laneId)); + return groupMask << (laneId & ~(groupSize - 1U)); +} + +__device__ __forceinline__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) +{ + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(val) + : "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]), "f"(array[4]), "f"(array[5]), "f"(array[6]), + "f"(array[7])); + return val; +} + +__device__ __forceinline__ uint16_t quantizeMinimaxM3Fp4x4( + float const (&elements)[4], float secondLevelScale, uint8_t* scaleOut) +{ + __nv_bfloat16 rounded[4]; + float localMax = 0.0F; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + rounded[i] = __float2bfloat16_rn(elements[i]); + localMax = fmaxf(localMax, fabsf(__bfloat162float(rounded[i]))); + } + uint32_t const groupMask = cvt_sf_group_shfl_mask<4>(); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 1)); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 2)); + float const scaleValue = secondLevelScale * (localMax * rcp_approx_ftz(6.0F)); + __nv_fp8_e4m3 const encodedScale(scaleValue); + if (scaleOut != nullptr) + { + *scaleOut = encodedScale.__x; + } + float const narrowedScale = static_cast(encodedScale); + float const outputScale + = localMax != 0.0F ? rcp_approx_ftz(narrowedScale * rcp_approx_ftz(secondLevelScale)) : 0.0F; + float quantValues[8] = {}; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + quantValues[i] = __bfloat162float(rounded[i]) * outputScale; + } + return static_cast(fp32_vec_to_e2m1(quantValues)); +} + +// RMSNorm (gemma 1+w) + NeoX RoPE over the first 64 dims, exactly as the producer: lane owns 4 dims. +__device__ __forceinline__ void norm_rope( + float (&e)[4], float const (&w)[4], float eps, float const (&cs)[4], float const (&sn)[4], int lane) +{ + float ss = 0.0F; + ss += e[0] * e[0] + e[1] * e[1]; + ss += e[2] * e[2] + e[3] * e[3]; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + ss = ss + __shfl_xor_sync(0xffffffffu, ss, mask, 32); + } + float const rmsReciprocal = rsqrtf(ss / static_cast(kHeadDim) + eps); +#pragma unroll + for (int i = 0; i < 4; ++i) + { + e[i] *= rmsReciprocal * (1.0F + w[i]); + } + __syncwarp(); + constexpr int kPairOffset = (kRotDim / 2) / 4; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const dim = lane * 4 + i; + float paired = __shfl_xor_sync(0xffffffffu, e[i], kPairOffset); + if (dim < kRotDim) + { + bool const firstHalf = dim < kRotDim / 2; + if (firstHalf) + { + paired = -paired; + } + e[i] = e[i] * cs[i] + paired * sn[i]; + } + } + __syncwarp(); +} + +__device__ __forceinline__ void load_w4(float (&w)[4], __nv_bfloat16 const* src, int lane) +{ + uint2 const raw = *reinterpret_cast(src + lane * 4); + float2 const a = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw.x)); + float2 const b = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw.y)); + w[0] = a.x; + w[1] = a.y; + w[2] = b.x; + w[3] = b.y; +} + +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ uint32_t mapa_shared(uint32_t addr, uint32_t rank) +{ + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n" : "=r"(r) : "r"(addr), "r"(rank)); + return r; +} + +// QLEN query tokens per request; QG of them per cluster (QLEN / QG clusters per request); BPC KV blocks per CTA +// (cluster of kBlocks / BPC CTAs). +// Byte offset of element (r, k) in a UMMA no-swizzle K-major operand tile of 128 k-elements per row: 8x8 core +// matrices of 16-byte rows, 128 bytes apart along k (LBO) and 2048 bytes apart per 8 rows (SBO). The same formula +// with r = token, k = dim stores V so that it is the MN-major (dims contiguous) A operand of O^T = V^T P^T. +__device__ __forceinline__ int kmaj_off(int r, int k) +{ + return (r >> 3) * 2048 + (k >> 3) * 128 + (r & 7) * 16 + (k & 7) * 2; +} + +__device__ __forceinline__ uint64_t umma_desc(uint32_t addr, uint32_t lbo, uint32_t sbo) +{ + return static_cast((addr >> 4) & 0x3FFFu) | (static_cast((lbo >> 4) & 0x3FFFu) << 16) + | (static_cast((sbo >> 4) & 0x3FFFu) << 32) | (1ull << 46); +} + +// kind::f16 instruction descriptor: F32 accumulate, F16 operands, B K-major, M = 128. +__host__ __device__ constexpr uint32_t umma_idesc(int n, int a_mn_major, int b_mn_major) +{ + return (1u << 4) | (static_cast(a_mn_major) << 15) | (static_cast(b_mn_major) << 16) + | (static_cast(n >> 3) << 17) | (static_cast(128 >> 4) << 24); +} + +// 8-row version for the split softmax: the same lane-bit order (16, 8, 4, then 2 and 1 in full) as +// warp_reduce_scatter16, so every row is combined in the same association (fp add is commutative, not associative). +// After it lanes with bits 4..2 = (b4, b3, b2) hold row 4*b4 + 2*b3 + b2. +template +__device__ __forceinline__ float warp_reduce_scatter8(float (&a)[8], int lane) +{ +#pragma unroll + for (int step = 0; step < 3; ++step) + { + int const mask = 16 >> step; + int const half = 4 >> step; + bool const upper = (lane & mask) != 0; +#pragma unroll + for (int i = 0; i < half; ++i) + { + float const send = upper ? a[i] : a[i + half]; + float const keep = upper ? a[i + half] : a[i]; + float const recv = __shfl_xor_sync(0xffffffffu, send, mask); + a[i] = kMax ? fmaxf(keep, recv) : keep + recv; + } + } + float x = a[0]; + float const o2 = __shfl_xor_sync(0xffffffffu, x, 2); + x = kMax ? fmaxf(x, o2) : x + o2; + float const o1 = __shfl_xor_sync(0xffffffffu, x, 1); + return kMax ? fmaxf(x, o1) : x + o1; +} + +// Reduce 16 per-lane values over the warp, one value per row: after the reduce-scatter lane l holds row +// 8*b4 + 4*b3 + 2*b2 + b1 (b_k = bit k of l) reduced over all 32 lanes. 16 shuffles instead of 80. +template +__device__ __forceinline__ float warp_reduce_scatter16(float (&a)[16], int lane) +{ +#pragma unroll + for (int step = 0; step < 4; ++step) + { + int const mask = 16 >> step; + int const half = 8 >> step; + bool const upper = (lane & mask) != 0; +#pragma unroll + for (int i = 0; i < half; ++i) + { + float const send = upper ? a[i] : a[i + half]; + float const keep = upper ? a[i + half] : a[i]; + float const recv = __shfl_xor_sync(0xffffffffu, send, mask); + a[i] = kMax ? fmaxf(keep, recv) : keep + recv; + } + } + float const other = __shfl_xor_sync(0xffffffffu, a[0], 1); + return kMax ? fmaxf(a[0], other) : a[0] + other; +} + +__device__ __forceinline__ void umma_f16( + uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %4, 0;\n tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(accumulate) + : "memory"); +} + +__device__ __forceinline__ void umma_commit(uint32_t bar) +{ + asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];\n" ::"r"(bar) : "memory"); +} + +__device__ __forceinline__ void mbar_wait_cta(uint32_t bar, uint32_t parity) +{ + uint32_t ready = 0; +#ifdef M3FA_DBG_HANG + long long spins = 0; +#endif + while (!ready) + { +#ifdef M3FA_DBG_HANG + if (++spins == (1LL << 24)) + { + printf("M3FA hang: block %d thread %d mbar 0x%x parity %u\n", blockIdx.x, threadIdx.x, bar, parity); + } +#endif + 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"(ready) + : "r"(bar), "r"(parity) + : "memory"); + } +} + +// Spin without suspending (v52's acc_done wait): a suspended try_wait wakes late on a tcgen05.commit arrival. +#ifndef M3FA_EPI_SPIN +#define M3FA_EPI_SPIN 1 +#endif +__device__ __forceinline__ void mbar_wait_epi(uint32_t bar, uint32_t parity) +{ + if (!M3FA_EPI_SPIN) + { + mbar_wait_cta(bar, parity); + return; + } + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(bar), "r"(parity) + : "memory"); + } +} + +__device__ __forceinline__ void tmem_ld8(uint32_t addr, float (&v)[8]) +{ + uint32_t r[8]; + asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n" + : "=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"(addr)); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < 8; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +__device__ __forceinline__ void tmem_ld16(uint32_t addr, 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];\n" + : "=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"(addr)); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < 16; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +// Warp-level fp16 MMA (m16n8k16, fp32 accumulate) and its shared-memory fragment loads. +__device__ __forceinline__ void ldsm_x4(uint32_t addr, uint32_t (&r)[4]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x2(uint32_t addr, uint32_t (&r)[2]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" : "=r"(r[0]), "=r"(r[1]) : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x2_trans(uint32_t addr, uint32_t (&r)[2]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0, %1}, [%2];\n" + : "=r"(r[0]), "=r"(r[1]) + : "r"(addr)); +} + +__device__ __forceinline__ void mma_f16_16816(float (&d)[4], uint32_t const (&a)[4], uint32_t const (&b)[2]) +{ + 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"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +// Register attention core (warp MMAs, no TMEM) for up to two tokens per cluster; the tcgen05 core above that. +// Off by default: on GB300 it measured ~0.25 us slower than the tcgen05 core at M=1 (GEMV + attention chain). +#ifndef M3FA_REG_CORE +#define M3FA_REG_CORE 0 +#endif +template +struct RegCore +{ + static constexpr bool kOn = M3FA_REG_CORE != 0 && QG <= 2; +}; + +template +struct Smem +{ + static constexpr int kRows = 16 * QG; // query rows (token, head) of this cluster = UMMA N + static constexpr int kTok = BPC * kPage; // block tokens per CTA + static constexpr int kTileBytes = kPage * kHeadDim * 2; // one fp16 K or V block tile in UMMA layout + static constexpr int kK16 = 0; + static constexpr int kV16 = kK16 + BPC * kTileBytes; + static constexpr int kQ = kV16 + BPC * kTileBytes; // Q rows, UMMA K-major [row][dim] + // P^T per block, UMMA MN-major [token][row], over the block's K tile: P is written after the QK MMAs (which + // read every K tile) have completed. + static constexpr int kP = kK16; + // Packed V tiles as copied from the cache: [block][128 rows x 64 data bytes | 1 KB of 4x4-quad scales]. + static constexpr int kRawV = kQ + kRows * kHeadDim * 2; + static constexpr int kRawVBlock = kPage * 64 + kPage * 8; + static constexpr int kRowScale = kRawV + BPC * kRawVBlock; + static constexpr int kRowLse = kRowScale + BPC * kRows * 4; + // Per-warp softmax partials [block][row][4 warps]: max, then sum. + static constexpr int kRedM = (kRowLse + BPC * kRows * 4 + 15) / 16 * 16; + static constexpr int kRedL = kRedM + BPC * kRows * 16; + // Partials pushed by every CTA of the cluster: [chunk][token][local head][kOPitch] floats, lse at column 128. + // When they do not fit beside the tiles they overlay the K/V tiles (and P, Q): peers then push only once every + // CTA of the cluster has completed its PV MMAs (a cluster barrier right before the push). + // Register core: per-warp row maxima and sums [block][row][warp of the block], two arrays. + static constexpr int kRedX = kRedL + BPC * kRows * 16; + static constexpr int kRedXBytes = RegCore::kOn ? 2 * BPC * kRows * (16 / BPC) * 4 : 0; + static constexpr int kStatsEnd = kRedX + kRedXBytes; + // Register core: P (fp16, [row][token] core matrices) at the start of its block's K tile, the scaled-free O + // partial ([row][kOPitch] fp32) behind it. + static constexpr int kObufOff = 8192; + static_assert(!RegCore::kOn || (kRows * kPage * 2 <= kObufOff && kObufOff + kRows * kOPitch * 4 <= kTileBytes), + "register-core P and O staging fit in a K tile"); + static constexpr int kRecvBytes = kBlocks * QG * BPC * kOPitch * 4; + static constexpr bool kRecvAlias = kStatsEnd + kRecvBytes + 32 > 227 * 1024; + static constexpr int kRecv = kRecvAlias ? 0 : kStatsEnd; + static constexpr int kMbar = kRecvAlias ? kStatsEnd : kStatsEnd + kRecvBytes; // merge, MMA mbarriers, TMEM addr + static constexpr int kBytes = kMbar + 32; + static constexpr int kTmemCols = BPC * 2 * kRows <= 32 ? 32 + : BPC * 2 * kRows <= 64 ? 64 + : BPC * 2 * kRows <= 128 ? 128 + : BPC * 2 * kRows <= 256 ? 256 + : 512; + static_assert(BPC * 2 * kRows <= 512, "TMEM columns"); + static_assert(!kRecvAlias || kRecvBytes <= kRawV, "aliased partials must fit in the tiles, Q and P"); + static_assert(kRows * kPage * 2 <= kTileBytes, "P^T fits in the K tiles"); + static_assert(kBytes <= 227 * 1024, "shared memory budget"); + static_assert(QLEN % QG == 0 && (BPC == 1 || BPC == 2), "unsupported geometry"); +}; + +// The packed NVFP4 bytes of one 16-byte chunk (32 dims of one block row) and its two block scales. +struct RawChunk +{ + uint4 data; + uint32_t scales; // byte 0: dims [0, 16), byte 1: dims [16, 32) of the chunk +}; + +// Chunk q of a block tile: lanes 8j..8j+7 of a warp take 8 consecutive rows of one 32-dim chunk, so every 16-byte +// fp16 store phase writes 8 distinct core-matrix rows (conflict-free in the UMMA layout). +__device__ __forceinline__ int chunk_row(int q) +{ + return (q >> 5) * 8 + (q & 7); +} + +__device__ __forceinline__ int chunk_col(int q) +{ + return (q >> 3) & 3; +} + +__device__ __forceinline__ RawChunk load_chunk( + uint8_t const* data_page, uint8_t const* scale_page, int row, int cc, bool is_v) +{ + RawChunk c; + c.data = *reinterpret_cast(data_page + row * 64 + cc * 16); + if (is_v) + { + // V scales: 4x4 token-quad order [row / 4][scale column][row % 4]. + uint8_t const* sc = scale_page + (row >> 2) * 32 + (row & 3); + c.scales = static_cast(sc[(cc * 2) * 4]) | (static_cast(sc[(cc * 2 + 1) * 4]) << 8); + } + else + { + c.scales = *reinterpret_cast(scale_page + row * 8 + cc * 2); + } + return c; +} + +__device__ __forceinline__ void store_chunk_fp16(uint8_t* tile, int row, int cc, RawChunk const& c, bool valid) +{ + uint32_t const s0 = e4m3_to_h2_splat(c.scales & 0xFFu); + uint32_t const s1 = e4m3_to_h2_splat((c.scales >> 8) & 0xFFu); + uint32_t const wd[4] = {c.data.x, c.data.y, c.data.z, c.data.w}; + uint32_t out[16]; +#pragma unroll + for (int i = 0; i < 16; ++i) + { + uint32_t const b = (wd[i >> 2] >> (8 * (i & 3))) & 0xFFu; + out[i] = valid ? hmul2_u(e2m1x2_to_h2(b), i < 8 ? s0 : s1) : 0u; + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + *reinterpret_cast(tile + kmaj_off(row, cc * 32 + j * 8)) + = make_uint4(out[4 * j], out[4 * j + 1], out[4 * j + 2], out[4 * j + 3]); + } +} + +// ---------------- Fused QKV projection: block-scaled MXFP8 on tcgen05 (N=2560, K=6144) ---------------- +// out[M, 2560] = mxfp8(x)[M, 6144] @ W^T, the computation of the CUTLASS MXFP8 GEMM: W is the A operand (128 output +// rows per MMA), the quantized activation the B operand (8 token columns), FP32 accumulation in TMEM. A unit is one +// (128-row tile, K split) pair; the 20 x kQkvSplits units run on the last CTAs of the grid. Before the PDL wait a +// CTA streams its unit's weight tiles and scale atoms through a shared-memory ring into TMEM; after it, it quantizes +// its x slice (mxfp8_quantize math), runs the MMAs and stores its FP32 partial rows. +// +// Hand-off without a flag: the partial buffers are double buffered by call parity and hold a sentinel (an FP32 NaN +// the MMA never produces) at rest, so a consumer polls the partials themselves. The producers of call n reset the +// other parity's buffers after their PDL wait; call n - 1, the last reader of those, is complete by then. +constexpr int kQkvN = 2560; +constexpr int kQkvK = 6144; +constexpr int kQkvKTiles = kQkvK / 128; +constexpr int kQkvRowTiles = kQkvN / 128; +#ifndef M3FA_QKV_SPLITS +#define M3FA_QKV_SPLITS 4 +#endif +#ifndef M3FA_QKV_RING +#define M3FA_QKV_RING kQkvMaxTiles // every tile of the unit in flight at once +#endif +#ifndef M3FA_QKV_NACC +#define M3FA_QKV_NACC 1 +#endif +#ifndef M3FA_QKV_SCALED_MMA +#define M3FA_QKV_SCALED_MMA 1 // 0: unscaled per-block MMAs with the scales applied in the epilogue (see below) +#endif +#ifndef M3FA_QKV_PIPE +// 1: the weight TMA is issued first at entry; the activation has its own region, copied right after the PDL wait; +// each weight tile's MMAs issue as soon as it lands (v52's order). 0: tiles into TMEM first, activation over the ring. +#define M3FA_QKV_PIPE 1 +#endif +#ifndef M3FA_QKV_A_SMEM +#define M3FA_QKV_A_SMEM 1 // scaled MMAs only; 1: they read the weight tiles from shared memory (no copy into TMEM) +#endif +// Default: block-scaled kind::mxf8f6f4 MMAs reading A (the weight tiles) in place from the ring, issued per tile as +// it lands (M3FA_QKV_PIPE), one accumulator; M3FA_QKV_A_SMEM=0 copies each tile into TMEM first. Alternative +// (M3FA_QKV_SCALED_MMA=0): one unscaled kind::f8f6f4 MMA per 32-wide K block into its own accumulator, A read in place +// from shared memory, the epilogue applying both UE8M0 scales. On GB300 an M128 x N8 x K32 MMA occupies the tensor pipe +// ~105 cycles block-scaled and ~60 unscaled, independent of N up to 64; the unscaled form's epilogue (320 accumulator +// columns) costs more than it saves. +constexpr bool kQkvScaled = M3FA_QKV_SCALED_MMA != 0; +constexpr bool kQkvASmem = !kQkvScaled || M3FA_QKV_A_SMEM != 0; +constexpr int kQkvSplits = M3FA_QKV_SPLITS; +constexpr int kQkvUnits = kQkvRowTiles * kQkvSplits; +constexpr int kQkvMaxTiles = (kQkvKTiles + kQkvSplits - 1) / kQkvSplits; +constexpr int kQkvRing = M3FA_QKV_RING; +constexpr int kQkvNacc = M3FA_QKV_NACC; +constexpr int kQkvGrid = 112; // 7 clusters of 16 (or 14 of 8): one wave with the attention +constexpr int kMaxTok = 8; +constexpr int kQkvWTile = 128 * 128; // E4M3 bytes of one 128 x 128 weight tile +constexpr int kQkvSfAtom = 512; // one 128-row x 4-block UE8M0 scale atom +constexpr uint32_t kQkvSentinel = 0xFFFFFFFFu; +constexpr int64_t kQkvPartStride = static_cast(kQkvSplits) * kMaxTok * kQkvN; // one parity +// TMEM columns of the one 512-column allocation of a CTA: attention [0, 64), then the GEMM's accumulators, scale +// atoms (A per tile, B per tile) and the A tiles (32 columns of 128 x 128 E4M3 each). +constexpr int kTmemAll = 512; +// Megakernel: the attention CTAs run no projection unit, so the GEMM starts at column 0 of a projection CTA. +constexpr int kTAcc = 0; +constexpr int kTSfa = kTAcc + 16; +constexpr int kTSfb = kTSfa + 4 * kQkvMaxTiles; +constexpr int kTA = (kTSfb + 4 * kQkvMaxTiles + 31) / 32 * 32; +static_assert(!kQkvScaled || kTA + 32 * kQkvMaxTiles <= kTmemAll, "TMEM columns"); +static_assert(kQkvScaled || kTAcc + 32 * kQkvMaxTiles <= kTmemAll, "TMEM columns (one accumulator per K block)"); +static_assert(!kQkvASmem || kQkvRing == kQkvMaxTiles, "shared-memory A needs every tile resident"); +static_assert(kQkvUnits <= kQkvGrid, "at most one unit per CTA"); +static_assert(kQkvNacc == 1 || kQkvNacc == 2 || kQkvNacc == 4, "accumulator count"); +static_assert(!M3FA_QKV_PIPE || kQkvScaled, "the pipelined projection uses the block-scaled MMAs"); +constexpr int kQkvOwnRows = 128 / kQkvSplits; // rows of a tile whose split-K sum a CTA owns +// Megakernel o_proj (v52 m3_mx_gemv for [6144, 2048]: split 2, contiguous k-tiles, one accumulator): unit +// o = CTA - 16 (the last 96 CTAs), row tile o / 2, split o % 2. +constexpr int kOpN = 6144; +constexpr int kOpK = 2048; +constexpr int kOpKTiles = kOpK / 128; +constexpr int kOpSplits = 2; +constexpr int kOpTiles = kOpKTiles / kOpSplits; +constexpr int kOpOwnRows = 128 / kOpSplits; +constexpr int kOpFirstCta = kQkvGrid - kOpN / 128 * kOpSplits; // 16 +constexpr int kOpSfBytes = kOpKTiles * 512; // one parity of the MXFP8 hand-off's scales +#ifndef M3FA_OP_L2PF +#define M3FA_OP_L2PF \ + 0 // 1: QKV CTAs prefetch their o_proj weight slice into L2 before the PDL wait (competes + // with the QKV weight stream for HBM behind the MoE) +#endif +// o_proj TMEM columns (after the QKV unit's MMAs completed): accumulator, A / B scale atoms, A (8 tiles x 32) +// c1 (attention + o_proj kernel): the attention holds columns [0, 32) while the o_proj weights stream in, so the +// o_proj accumulator and scale atoms sit above them. +constexpr int kOAcc = 32; +constexpr int kOSfa = 48; +constexpr int kOSfb = kOSfa + 4 * kOpTiles; +constexpr int kOA = 128; +static_assert(kOA + 32 * kOpTiles <= kTmemAll && kOSfb + 4 * kOpTiles <= kOA, "o_proj TMEM columns"); +constexpr int kQkvMbStride = kMaxTok + 4; // mailbox row pitch (fp32 words) + +__host__ __device__ constexpr int qkv_tile0(int s) +{ + return s * kQkvKTiles / kQkvSplits; +} + +__device__ __forceinline__ int qkv_unit() +{ + int const u = kQkvGrid - 1 - static_cast(blockIdx.x); + return u < kQkvUnits ? u : -1; +} + +// Attention layout plus, for the fused projection, the weight ring in front (the attention tiles reuse that space +// once the ring has drained into TMEM); the activation tiles, the pre-wait V copy and the barriers sit behind. +template +struct SmemX : Smem +{ + using B = Smem; + static constexpr int kRingW = 0; // [ring][128 rows][128 B], 128-byte swizzle + static constexpr int kRingSf = kRingW + kQkvRing * kQkvWTile; // [ring][512 B] + static constexpr int kRingEnd = kRingSf + kQkvRing * kQkvSfAtom; + // The activation tiles overlay the ring: written after the PDL wait, once every weight tile went on into TMEM. + static constexpr int kXt = kRingW; // [tile][8][128 B] + static constexpr int kXsf = kXt + kQkvMaxTiles * 1024; // [tile][512 B] + static constexpr int kXsc = (kRingEnd + 127) / 128 * 128; + // Peers push into recv at any time (an empty CTA right away), so it stays clear of the ring. + static constexpr int kRecv = FQKV ? kXsc : B::kRecv; + static constexpr int kRawV = FQKV ? (kRecv + kBlocks * QG * BPC * kOPitch * 4 + 127) / 128 * 128 : B::kRawV; + // Pipelined projection: the activation tiles and scale atoms in their own region (a projection CTA runs no + // attention, so it overlays the attention's partial-O and V regions), 1024-byte aligned for the swizzle. + static constexpr int kXA = FQKV ? (kRecv + 1023) / 1024 * 1024 : 0; + static constexpr int kXASf = kXA + kQkvMaxTiles * 1024; + static constexpr int kXAEnd = kXASf + kQkvMaxTiles * kQkvSfAtom; + static constexpr int kMbarA = FQKV ? kRawV + BPC * B::kRawVBlock : B::kMbar; + static constexpr int kMbar = FQKV && M3FA_QKV_PIPE && kXAEnd > kMbarA ? (kXAEnd + 127) / 128 * 128 : kMbarA; + static constexpr int kGBar = kMbar + 32; // full[ring], empty[ring], a_ready, acc_done, tile_done[tiles] + static constexpr int kRedBar = kGBar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); // split-K partials in + static constexpr int kOBar = kRedBar + 8; // o_proj: full[8], a_ready, acc_done, reduce + static constexpr int kMailbox = (kOBar + 8 * (kOpTiles + 3) + 15) / 16 * 16; // [split][own row][pitch] fp32 + // o_proj weight ring = ring tiles [0, 8); its activation tiles, scale atoms and mailbox behind it. + static constexpr int kOXt = kRingW + kOpTiles * kQkvWTile; + static constexpr int kOXsf = kOXt + kOpTiles * 1024; + static constexpr int kOMailbox = kOXsf + kOpTiles * kQkvSfAtom; + static_assert(!FQKV || kOMailbox + kOpSplits * kOpOwnRows * kQkvMbStride * 4 <= kRingW + kQkvRing * kQkvWTile, + "o_proj activation and mailbox fit behind its ring"); + static constexpr int kBytes = FQKV ? kMailbox + kQkvSplits * kQkvOwnRows * kQkvMbStride * 4 : B::kBytes; + static constexpr int kTmemCols = FQKV ? kTmemAll : B::kTmemCols; + static constexpr bool kRecvAlias = !FQKV && B::kRecvAlias; + // The epilogue reads the weight scale atoms while the attention already writes its tiles. + static_assert(!FQKV || kQkvScaled || B::kRecv <= kRingSf, "attention regions overlap the scale atoms"); + // The other attention regions overlay the ring: a unit's CTA writes them only once the ring has drained. + static_assert(kBytes <= 227 * 1024, "shared memory budget"); +}; + +// c1: the packed-input attention plus the o_proj unit blockIdx.x (tile bx / 2, split bx % 2). The o_proj weight ring +// and its barriers sit behind the attention layout; the activation tiles, scale atoms and the pair's split-K mailbox +// overlay the K tiles, free once this CTA has pushed its partials (a peer's mailbox push follows that push: its +// activation needs the merges of this cluster, which need this CTA's partials). +template +struct SmemC1 : SmemX +{ + using B = Smem; + static constexpr int kRingW = (B::kBytes + 1023) / 1024 * 1024; + static constexpr int kRingSf = kRingW + kOpTiles * kQkvWTile; + static constexpr int kOBar = kRingSf + kOpTiles * kQkvSfAtom; // full[8], a_ready, acc_done, reduce + static constexpr int kOXt = B::kK16; + static constexpr int kOXsf = kOXt + kOpTiles * 1024; + static constexpr int kOMailbox = kOXsf + kOpTiles * kQkvSfAtom; + static constexpr int kBytes = kOBar + 8 * (kOpTiles + 3); + static constexpr int kTmemCols = kTmemAll; + static_assert(!B::kRecvAlias, "c1: the partials have their own region"); + static_assert( + kOMailbox + kOpSplits * kOpOwnRows * kQkvMbStride * 4 <= B::kV16, "c1: o_proj overlay in the K tiles"); + static_assert(BPC * 2 * B::kRows <= kOAcc, "c1: attention TMEM columns below the o_proj's"); + static_assert(kBytes <= 227 * 1024, "c1: shared memory budget"); +}; + +__device__ __forceinline__ void mbar_init(uint32_t bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(bar), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbar_expect_tx(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 bulk_g2s(void* dst, void const* src, uint32_t bytes, uint32_t bar) +{ + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], " + "%4;\n" ::"r"(smem_u32(dst)), + "l"(src), "r"(bytes), "r"(bar), "l"(0x12F0000000000000ull) // EVICT_FIRST: weights are read once per layer + : "memory"); +} + +__device__ __forceinline__ void tma_load_2d(uint32_t dst, void const* map, int c0, int c1, uint32_t bar) +{ + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;\n" ::"r"(dst), + "l"(map), "r"(c0), "r"(c1), "r"(bar), "l"(0x12F0000000000000ull) + : "memory"); +} + +// K-major operand in the 128-byte swizzle: rows of 128 bytes, 8-row groups 1024 bytes apart. +__device__ __forceinline__ uint64_t umma_desc_sw128(uint32_t addr) +{ + return umma_desc(addr, 16, 1024) | (2ull << 61); +} + +__device__ __forceinline__ void tmem_cp_128x256b(uint32_t taddr, uint64_t desc) +{ + asm volatile("tcgen05.cp.cta_group::1.128x256b [%0], %1;\n" ::"r"(taddr), "l"(desc)); +} + +__device__ __forceinline__ void tmem_cp_sf(uint32_t taddr, uint32_t saddr) +{ + asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;\n" ::"r"(taddr), "l"(umma_desc(saddr, 16, 128))); +} + +// Block-scaled MXFP8 instruction descriptor: E4M3 x E4M3, UE8M0 scales, M = 128, N = 8, scale byte kb. +__host__ __device__ constexpr uint32_t qkv_idesc(int kb) +{ + return (static_cast(kb) << 4) | ((8u >> 3) << 17) | (1u << 23) | ((128u >> 7) << 27) + | (static_cast(kb) << 29); +} + +__device__ __forceinline__ void umma_mx( + uint32_t tmem_d, uint32_t tmem_a, uint64_t bdesc, uint32_t idesc, uint32_t sfa, uint32_t sfb, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], [%1], %2, %3, [%4], [%5], p;\n}\n" :: + "r"(tmem_d), + "r"(tmem_a), "l"(bdesc), "r"(idesc), "r"(sfa), "r"(sfb), "r"(accumulate)); +} + +__device__ __forceinline__ float pow2_from_ue8m0(uint32_t e) +{ + return __uint_as_float(e == 0 ? 0x00400000u : (e << 23)); +} + +// Unscaled E4M3 x E4M3 MMA, FP32 accumulator: M = 128, N = 8, K = 32, a fresh accumulator. +constexpr uint32_t kQkvIdescF8 = (1u << 4) | ((8u >> 3) << 17) | ((128u >> 4) << 24); + +__device__ __forceinline__ void umma_f8_ss(uint32_t tmem_d, uint64_t adesc, uint64_t bdesc) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, 0, 0;\n tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(kQkvIdescF8)); +} + +__device__ __forceinline__ void tmem_ld16_nowait(uint32_t addr, float* v) +{ + 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];\n" + : "=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"(addr)); +#pragma unroll + for (int i = 0; i < 16; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +__device__ __forceinline__ void umma_mx_ss( + uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, uint32_t sfa, uint32_t sfb, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], %1, %2, %3, [%4], [%5], p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(sfa), "r"(sfb), "r"(accumulate)); +} + +__device__ __forceinline__ bool elect_one() +{ + uint32_t pred; + asm volatile("{\n .reg .pred P;\n elect.sync _|P, 0xffffffff;\n selp.u32 %0, 1, 0, P;\n}\n" : "=r"(pred)); + return pred != 0; +} + +__device__ __forceinline__ int qkv_parity(Params const& p) +{ + int v; + asm volatile("ld.relaxed.gpu.global.s32 %0, [%1];\n" : "=r"(v) : "l"(p.counters) : "memory"); + return v & 1; +} + +// Barriers and TMEM (warp 1), then the first ring stages (warp 0). Warps 0 and 1 of every CTA call it. +template +__device__ __forceinline__ void qkv_setup( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int lane, int u) +{ + constexpr bool kPipe = M3FA_QKV_PIPE != 0; + if (kPipe && warp == 0 && lane == 0 && u >= 0) + { + // First thing at entry: this unit's whole weight slice. The issuing thread initializes the ring's full + // barriers itself; every waiter on them passes bar.sync 8 below, which this thread reaches after the init. + asm volatile("prefetch.tensormap [%0];\n" ::"l"(&p.wmap) : "memory"); + int const rt = u / kQkvSplits; + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + for (int i = 0; i < kQkvRing; ++i) + { + mbar_init(gbar + 8 * i, 1); + } + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + for (int i = 0; i < kQkvRing && i < nt; ++i) + { + uint32_t const full = gbar + 8 * i; + mbar_expect_tx(full, kQkvWTile + kQkvSfAtom); + if (p.wpre != nullptr) + { + // Pre-tiled weight: the tile's 16 KB shared-memory image is contiguous (one bulk copy). + bulk_g2s(smem + SX::kRingW + i * kQkvWTile, + p.wpre + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvWTile, kQkvWTile, full); + } + else + { + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap, (kt0 + i) * 128, rt * 128, full); + } + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, full); + } + M3FA_PROF_T(34, 0); + } + if (warp == 1) + { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n" ::"r"(smem_u32(tmem_slot)), + "r"(kTmemAll) + : "memory"); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n" ::: "memory"); + M3FA_PROF_T(35, 32); + if (lane == 0) + { + for (int i = (kPipe && u >= 0) ? kQkvRing : 0; i < 2 * kQkvRing + 2 + kQkvMaxTiles; ++i) + { + // PIPE 2: tile_done[0] (unused by the scaled MMAs) is act_ready, one arrival per thread outside warp 1 + mbar_init(gbar + 8 * i, M3FA_QKV_PIPE == 2 && i == 2 * kQkvRing + 2 ? kThreads - 32 : 1); + } + // The split-K reduce barrier: this CTA's own arrival carries the peers' partial bytes. + uint32_t const red = gbar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); + mbar_init(red, 1); + mbar_expect_tx(red, (kQkvSplits - 1) * kQkvOwnRows * kMaxTok * 4); + uint32_t const ob = smem_u32(smem + SX::kOBar); + for (int i = 0; i < kOpTiles + 3; ++i) + { + mbar_init(ob + 8 * i, 1); + } + mbar_expect_tx(ob + 8 * (kOpTiles + 2), (kOpSplits - 1) * kOpOwnRows * kMaxTok * 4); + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + asm volatile("bar.sync 8, 64;\n" ::: "memory"); + if (!kPipe && warp == 0 && lane == 0 && u >= 0) + { + int const rt = u / kQkvSplits; + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + for (int i = 0; i < kQkvRing && i < nt; ++i) + { + uint32_t const full = gbar + 8 * i; + mbar_expect_tx(full, kQkvWTile + kQkvSfAtom); + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap, (kt0 + i) * 128, rt * 128, full); + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, full); + } + M3FA_PROF_T(34, 0); + } +} + +// The rest of the ring (warp 0, lane 0) and the copies into TMEM (warp 1, lane 0); before the PDL wait. +template +__device__ __forceinline__ void qkv_stream_lane( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int u) +{ + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + uint32_t const full = gbar; + uint32_t const a_ready = gbar + 16 * kQkvRing; + if (warp == 0 || !kQkvScaled) + { + return; // every tile was issued in qkv_setup; the unscaled MMAs read them in place + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + // Fully unrolled with compile-time offsets: tcgen05 operands stay in uniform registers (a per-instruction + // conversion of computed operands costs ~150 cycles per MMA or copy). + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint32_t const sf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + mbar_wait_cta(full + 8 * i, 0); + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if constexpr (!kQkvASmem) + { +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kTA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + } + tmem_cp_sf(tm + kTSfa + i * 4, sf0 + i * kQkvSfAtom); + } + } + M3FA_PROF_T(31, 32); + umma_commit(a_ready); +} + +__device__ __forceinline__ bool mbar_test(uint32_t bar, uint32_t parity) +{ + uint32_t ready; + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(bar), "r"(parity) + : "memory"); + return ready != 0; +} + +// PIPE 2, warp 1's elected thread from before the PDL wait to the last MMA. Tiles that land before the activation +// is ready go on into TMEM (A and scale atom), off the critical path; once the activation is in (act_ready), the +// MMAs issue in tile order: the copied tiles with A in TMEM, every later tile as it lands with A read in place from +// the ring. The MMA sequence (and so the accumulation order) is the same either way. +template +__device__ __forceinline__ void qkv_hybrid(Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int u) +{ + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + uint32_t const act_ready = gbar + 8 * (2 * kQkvRing + 2); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + SX::kXA)); + uint32_t const wsf0 = smem_u32(smem + SX::kRingSf); + uint32_t const xsf0 = smem_u32(smem + SX::kXASf); + int copied = 0; // tiles [0, copied) are in TMEM + bool act = false; +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt && !act) + { + while (!mbar_test(gbar + 8 * i, 0)) + { + if (mbar_test(act_ready, 0)) + { + act = true; + break; + } + } + if (!act) + { + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kTA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + copied = i + 1; + } + } + } + while (!act) + { + act = mbar_test(act_ready, 0); + } + M3FA_PROF_T(21, 32); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + bool const in_tmem = i < copied; + if (!in_tmem) + { + mbar_wait_epi(gbar + 8 * i, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + } + uint32_t const sfb = tm + kTSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + uint64_t const bd = b0 + ((i * 1024 + kb * 32) >> 4); + if (in_tmem) + { + umma_mx(tm + kTAcc, tm + kTA + j * 8, bd, qkv_idesc(kb), tm + kTSfa + i * 4, sfb, j > 0); + } + else + { + umma_mx_ss(tm + kTAcc, w0 + ((i * kQkvWTile + kb * 32) >> 4), bd, qkv_idesc(kb), tm + kTSfa + i * 4, + sfb, j > 0); + } + } + } + } + M3FA_PROF_T(22, 32); + M3FA_PROF_T(31, 32); + umma_commit(gbar + 16 * kQkvRing); // a_ready (for its waiters) + umma_commit(gbar + 16 * kQkvRing + 8); // acc_done +} + +// The other lanes of warps 0 and 1 hold at __syncwarp rather than at griddepcontrol.wait, which would stall the +// whole warp (lane 0's loop included) until the producer grid completes. +template +__device__ __forceinline__ void qkv_stream( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int lane, int u) +{ + if (M3FA_QKV_PIPE == 2 && u >= 0 && warp == 1) + { + if (elect_one()) + { + qkv_hybrid(p, smem, gbar, tmem_slot, u); + } + __syncwarp(); + return; + } + if (u < 0 || warp > 1 || M3FA_QKV_PIPE) + { + return; // pipelined: the MMA thread takes each tile as it lands + } + if (lane == 0) + { + qkv_stream_lane(p, smem, gbar, tmem_slot, warp, u); + } + __syncwarp(); +} + +// After the PDL wait. Every thread quantizes a piece of the x slice into the B tiles (named barrier 9); warp 1 then +// issues the scale copies and the MMAs. +template +__device__ __forceinline__ void qkv_mma( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int tid, int u) +{ + int const warp = tid >> 5; + int const s = u % kQkvSplits; + int const kt0 = qkv_tile0(s); + int const nt = qkv_tile0(s + 1) - kt0; + { + // The MXFP8 activation slice (E4M3 rows of the unit's k-tiles in 128-byte-swizzled B tiles, and their scale + // atoms), copied with plain 16-byte loads over the overlaid ring once its tiles went on into TMEM: every load + // first, then the stores (one L2 round trip). + constexpr int kXDst = M3FA_QKV_PIPE ? SX::kXA : SX::kXt; + constexpr int kXSfDst = M3FA_QKV_PIPE ? SX::kXASf : SX::kXsf; + if (!M3FA_QKV_PIPE) + { + mbar_wait_cta(gbar + 16 * kQkvRing, 0); // a_ready: the ring's weight tiles are in TMEM + } + constexpr int kXItems = kQkvMaxTiles * 8 * kMaxTok; // (tile, 16-byte chunk, token) + constexpr int kSfItems = kQkvMaxTiles * (kQkvSfAtom / 16); + constexpr bool kHyb = M3FA_QKV_PIPE == 2; + constexpr int kCopiers = kHyb ? kThreads - 32 : kThreads; // PIPE 2: warp 1 is issuing the tiles + if (kHyb && warp == 1) + { + return; + } + int const ctid = kHyb && tid >= 64 ? tid - 32 : tid; + constexpr int kPer = (kXItems + kSfItems + kCopiers - 1) / kCopiers; + uint4 v[kPer]; + int dst[kPer]; +#pragma unroll + for (int k = 0; k < kPer; ++k) + { + int const q = ctid + kCopiers * k; + dst[k] = -1; + if (q < kXItems) + { + int const t = q % kMaxTok; + int const c = (q / kMaxTok) % 8; + int const i = q / (kMaxTok * 8); + if (i < nt && t < p.num_tokens) + { + v[k] = *reinterpret_cast( + p.xq + static_cast(t) * p.ldxq + (kt0 + i) * 128 + c * 16); + dst[k] = kXDst + i * 1024 + t * 128 + ((c ^ t) * 16); + } + } + else if (q < kXItems + kSfItems) + { + int const j = q - kXItems; + int const i = j / (kQkvSfAtom / 16); + int const c = j % (kQkvSfAtom / 16); + if (i < nt) + { + v[k] = *reinterpret_cast( + p.xsf_act + static_cast(kt0 + i) * kQkvSfAtom + c * 16); + dst[k] = kXSfDst + i * kQkvSfAtom + c * 16; + } + } + } +#pragma unroll + for (int k = 0; k < kPer; ++k) + { + if (dst[k] >= 0) + { + *reinterpret_cast(smem + dst[k]) = v[k]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if (kHyb) + { + asm volatile("mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];\n" ::"r"(gbar + 8 * (2 * kQkvRing + 2)) + : "memory"); + return; + } + if (warp != 1) + { + asm volatile("bar.arrive 9, %0;\n" ::"n"(kThreads) : "memory"); + } + } + if (warp == 1) + { + M3FA_PROF_T(32, 32); + asm volatile("bar.sync 9, %0;\n" ::"n"(kThreads) : "memory"); + M3FA_PROF_T(21, 32); + // Operands computed by the whole warp (warp-uniform, kept in uniform registers); one elected lane issues. + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + (M3FA_QKV_PIPE ? SX::kXA : SX::kXt))); + uint64_t const a0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + if (elect_one()) + { + M3FA_CLK(41); + if constexpr (!kQkvScaled) + { +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + mbar_wait_cta(gbar + 8 * i, 0); // weight tile i landed + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + umma_f8_ss(tm + kTAcc + (i * 4 + kb) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4), + b0 + ((i * 1024 + kb * 32) >> 4)); + } + umma_commit(gbar + 8 * (2 * kQkvRing + 2 + i)); // tile i's blocks are in TMEM + } + } + } + else + { + uint32_t const xsf0 = smem_u32(smem + (M3FA_QKV_PIPE ? SX::kXASf : SX::kXsf)); + uint32_t const wsf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + if constexpr (M3FA_QKV_PIPE != 0) + { + // Tile i as it lands: its A copy and scale atom into TMEM, then its MMAs (issue order). + mbar_wait_cta(gbar + 8 * i, 0); + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + if (i == nt - 1) + { + M3FA_CLK(47); // the unit's last weight tile landed (clock64 builds) + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if constexpr (!kQkvASmem) + { +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b( + tm + kTA + (i * 4 + kb) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + } + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + } + uint32_t const sfb = tm + kTSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + uint64_t const bd = b0 + ((i * 1024 + kb * 32) >> 4); + if constexpr (kQkvASmem) + { + umma_mx_ss(tm + kTAcc + (j % kQkvNacc) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4), bd, + qkv_idesc(kb), tm + kTSfa + i * 4, sfb, j >= kQkvNacc); + } + else + { + umma_mx(tm + kTAcc + (j % kQkvNacc) * 8, tm + kTA + j * 8, bd, qkv_idesc(kb), + tm + kTSfa + i * 4, sfb, j >= kQkvNacc); + } + } + } + } + } + M3FA_CLK(42); + if (M3FA_QKV_PIPE) + { + M3FA_PROF_T(31, 32); + umma_commit(gbar + 16 * kQkvRing); // a_ready (kept for its waiters: every tile is consumed) + } + umma_commit(gbar + 16 * kQkvRing + 8); // acc_done + } + M3FA_PROF_T(22, 32); + } +} + +// Warps 16..19 (TMEM lanes 32 (warp % 4) + lane = the unit's rows): partial rows out, the other parity reset. +template +__device__ __forceinline__ void qkv_epilogue( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int tid, int u, int par, unsigned int calls) +{ + int const warp = tid >> 5; + int const lane = tid & 31; + if constexpr (kQkvScaled) + { + mbar_wait_epi(gbar + 16 * kQkvRing + 8, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + M3FA_PROF_T(24, 512); + if (tid == 512) + { + M3FA_CLK(43); + } + int const r = (warp & 3) * 32 + lane; // row within the unit's tile + uint32_t const taddr = *tmem_slot + (((warp & 3) * 32) << 16) + kTAcc; + float val[kMaxTok]; + if constexpr (kQkvScaled) + { + float v[32]; + tmem_ld16(taddr, *reinterpret_cast(v)); + if constexpr (kQkvNacc == 4) + { + tmem_ld16(taddr + 16, *reinterpret_cast(v + 16)); + } +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = kQkvNacc == 4 ? (v[t] + v[8 + t]) + (v[16 + t] + v[24 + t]) + : kQkvNacc == 2 ? v[t] + v[8 + t] + : v[t]; + } + } + else + { + // Sum over K blocks in order: block result x (weight scale x activation scale). + int const nt = qkv_tile0(u % kQkvSplits + 1) - qkv_tile0(u % kQkvSplits); + float const* xsc = reinterpret_cast(smem + SX::kXsc); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = 0.0f; + } +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + // Tile i's four blocks as soon as its MMAs complete: the epilogue trails the MMAs by one tile. + mbar_wait_cta(gbar + 8 * (2 * kQkvRing + 2 + i), 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + float d[32]; + tmem_ld16_nowait(taddr + i * 32, d); + tmem_ld16_nowait(taddr + i * 32 + 16, d + 16); + uint32_t const wsw = *reinterpret_cast( + smem + SX::kRingSf + i * kQkvSfAtom + (r % 32) * 16 + (r / 32) * 4); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + float const ws = pow2_from_ue8m0((wsw >> (8 * kb)) & 0xFFu); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = fmaf(d[kb * 8 + t], ws * xsc[(i * 4 + kb) * kMaxTok + t], val[t]); + } + } + } + } + } + // v52's split-K reduce: split s owns rows [32 s, 32 s + 32) of the tile; the other splits push their fp32 partials + // of those rows into its mailbox (st.async, completing its reduce barrier's bytes); the owner adds the splits in + // rank order from 0.0 and stores bf16. + int const s = u % kQkvSplits; + int const owner = r / kQkvOwnRows; + int const local = r % kQkvOwnRows; + uint32_t const red = gbar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); + uint32_t cta_rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + // Units are numbered down from the last CTA: split q of this tile is CTA rank cta_rank - (q - s). + if (owner != s) + { + uint32_t const peer = cta_rank + s - owner; + uint32_t const dst + = mapa_shared(smem_u32(smem + SX::kMailbox) + ((s * kQkvOwnRows + local) * kQkvMbStride) * 4, peer); + uint32_t const pbar = mapa_shared(red, peer); +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( + dst + 16 * g), + "r"(__float_as_uint(val[4 * g])), "r"(__float_as_uint(val[4 * g + 1])), + "r"(__float_as_uint(val[4 * g + 2])), "r"(__float_as_uint(val[4 * g + 3])), "r"(pbar) + : "memory"); + } + } + else + { + M3FA_PROF_T(26, tid - lane); + // test_wait (no suspend): completions by peers' st.async bytes wake a suspended try_wait late. + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(red) + : "memory"); + } + M3FA_PROF_T(27, tid - lane); + float const* mb = reinterpret_cast(smem + SX::kMailbox); + int const row = (u / kQkvSplits) * 128 + r; + // Lamport hand-off: this call's rows go to buffer (calls & 1) with relaxed stores, the consumers poll the + // values themselves (0xFFFF, a NaN the rounding of a finite sum never produces, marks "not yet"); the other + // buffer, last read by the previous call (complete: this grid passed its PDL wait), is reset to the marker. + uint16_t* cur = reinterpret_cast(p.qkv_out) + (calls & 1u) * kMaxTok * kQkvN; + uint16_t* nxt = reinterpret_cast(p.qkv_out) + ((calls & 1u) ^ 1u) * kMaxTok * kQkvN; +#pragma unroll + float qsum[kMaxTok]; +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + float4 part[kQkvSplits]; +#pragma unroll + for (int q = 0; q < kQkvSplits; ++q) + { + part[q] = *reinterpret_cast(mb + (q * kQkvOwnRows + local) * kQkvMbStride + 4 * g); + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float sum = 0.0F; +#pragma unroll + for (int q = 0; q < kQkvSplits; ++q) + { + float const pj = j == 0 ? part[q].x : j == 1 ? part[q].y : j == 2 ? part[q].z : part[q].w; + sum = sum + (q == s ? val[4 * g + j] : pj); + } + qsum[4 * g + j] = sum; + } + } +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + float const sum = qsum[t]; + if (t < p.num_tokens) + { + __nv_bfloat16 const b = __float2bfloat16_rn(sum); + uint16_t bits = *reinterpret_cast(&b); + // Never the "not yet" marker: a 0xFFFF NaN sum is stored as the NaN 0x7FFF. Only a NaN projection row + // changes (its K/V cache bytes then differ from v52 + p6 in the NaN's encoding); finite rows are exact. + bits = bits == 0xFFFFu ? static_cast(0x7FFFu) : bits; +#ifndef M3FA_QKV_REL +#define M3FA_QKV_REL 0 // 1: the row store is a release (tried: whether it becomes visible to the pollers sooner) +#endif + if (M3FA_QKV_REL) + { + asm volatile( + "st.release.gpu.global.u16 [%0], %1;\n" ::"l"(cur + static_cast(t) * kQkvN + row), + "h"(bits) + : "memory"); + } + else + { + asm volatile( + "st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(cur + static_cast(t) * kQkvN + row), + "h"(bits) + : "memory"); + } + } + asm volatile("st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(nxt + static_cast(t) * kQkvN + row), + "h"(static_cast(0xFFFFu)) + : "memory"); + } + M3FA_PROF_T(28, tid - lane); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + M3FA_PROF_T(25, 512); + if (tid == 512) + { + M3FA_CLK(46); + } +} + +// ---------------- Megakernel o_proj (v52 m3_mx_gemv bit for bit: split 2, contiguous k-tiles, nacc 1) ---------------- +// The unit's weight tiles and scale atoms into the ring (warp 0, lane 0). From L2 when the CTA prefetched them there. +template +__device__ __forceinline__ void op_issue(Params const& p, uint8_t* smem, int ou) +{ + int const rt = ou / kOpSplits; + int const kt0 = (ou % kOpSplits) * kOpTiles; + uint32_t const ob = smem_u32(smem + SX::kOBar); + for (int i = 0; i < kOpTiles; ++i) + { + mbar_expect_tx(ob + 8 * i, kQkvWTile + kQkvSfAtom); + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap_o, (kt0 + i) * 128, rt * 128, ob + 8 * i); + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf_o + (static_cast(rt) * kOpKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, ob + 8 * i); + } +} + +// L2 prefetch of the unit's weight tiles and scale atoms (one lane; nothing waits for it). +// The QKV units' share of the window-1 L2 prefetch (the next consumer's weights), issued once their own QKV weight +// stream has landed (acc_done): the attention core reads little from HBM. +__device__ __forceinline__ void window_prefetch_l2(Params const& p, int u) +{ + constexpr long long kChunk = 32768; + long long const n0 = (p.pf_b0 + kChunk - 1) / kChunk, n1 = (p.pf_b1 + kChunk - 1) / kChunk; + uint64_t pol = 0; + if (p.pf_hint) + { + asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, 1.0;\n" : "=l"(pol)); + } + for (long long c = u; c < n0 + n1; c += kQkvUnits) + { + long long const a = c < n0 ? p.pf_a0 + c * kChunk : p.pf_a1 + (c - n0) * kChunk; + long long const rem = c < n0 ? p.pf_b0 - c * kChunk : p.pf_b1 - (c - n0) * kChunk; + unsigned int const n = static_cast(rem < kChunk ? rem : kChunk) & ~15u; + if (n == 0) + { + continue; + } + if (p.pf_hint) + { + asm volatile("cp.async.bulk.prefetch.L2.global.L2::cache_hint [%0], %1, %2;\n" ::"l"(a), "r"(n), "l"(pol) + : "memory"); + } + else + { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" ::"l"(a), "r"(n) : "memory"); + } + } +} + +__device__ __forceinline__ void op_prefetch_l2(Params const& p, int ou) +{ + int const rt = ou / kOpSplits; + int const kt0 = (ou % kOpSplits) * kOpTiles; + for (int i = 0; i < kOpTiles; ++i) + { + asm volatile("cp.async.bulk.prefetch.tensor.2d.L2.global [%0, {%1, %2}];\n" ::"l"(&p.wmap_o), + "r"((kt0 + i) * 128), "r"(rt * 128) + : "memory"); + } + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" ::"l"( + p.wsf_o + (static_cast(rt) * kOpKTiles + kt0) * kQkvSfAtom), + "r"(kOpTiles * kQkvSfAtom) + : "memory"); +} + +// The weight tiles into TMEM as they land (warp 1, elected lane, operands warp-uniform). +template +__device__ __forceinline__ void op_stream(uint8_t* smem, uint32_t tm) +{ + uint32_t const ob = smem_u32(smem + SX::kOBar); + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint32_t const sf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kOpTiles; ++i) + { + mbar_wait_cta(ob + 8 * i, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kOA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + tmem_cp_sf(tm + kOSfa + i * 4, sf0 + i * kQkvSfAtom); + } + umma_commit(ob + 8 * kOpTiles); // a_ready +} + +// The attention's MXFP8 rows of this unit's k-tiles (the heads [kt0, kt0 + 8)), Lamport-polled value by value: 0xFF +// is NaN in both E4M3 and UE8M0, which the merge never writes (see the merge). Every thread of the CTA, one item each. +template +__device__ __forceinline__ void op_act(Params const& p, uint8_t* smem, int ou, int tid, unsigned int calls) +{ + int const kt0 = (ou % kOpSplits) * kOpTiles; + uint8_t const* q8 = p.q8 + (calls & 1u) * kMaxTok * kOpK; + uint8_t const* sf = p.sf + (calls & 1u) * kOpSfBytes; + constexpr int kXItems = kOpTiles * 8 * kMaxTok; + constexpr int kSfItems = kOpTiles * (kQkvSfAtom / 16); + for (int q = tid; q < kXItems + kSfItems; q += kThreads) + { + if (q < kXItems) + { + int const t = q % kMaxTok; + int const c = (q / kMaxTok) % 8; + int const i = q / (kMaxTok * 8); + if (t < p.num_tokens) + { + uint4 v; + bool pending; + do + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(q8 + static_cast(t) * kOpK + (kt0 + i) * 128 + c * 16) + : "memory"); + uint32_t const w[4] = {v.x, v.y, v.z, v.w}; + pending = false; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + // a byte equal to 0xFF: (x & 0x7F7F7F7F) + 0x01010101 sets bit 7 of exactly the 0x7F/0xFF bytes + uint32_t const hit = (((w[j] & 0x7F7F7F7Fu) + 0x01010101u) & w[j] & 0x80808080u); + pending |= hit != 0u; + } + } while (pending); + *reinterpret_cast(smem + SX::kOXt + i * 1024 + t * 128 + ((c ^ t) * 16)) = v; + } + } + else if (q < kXItems + kSfItems) + { + int const j = q - kXItems; + int const i = j / (kQkvSfAtom / 16); + int const c = j % (kQkvSfAtom / 16); + uint8_t const* src = sf + (kt0 + i) * kQkvSfAtom + c * 16; + uint4 v; + if (c < p.num_tokens) + { + // Row c of the atom: its 4 block scales in bytes [16 c, 16 c + 4). + do + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(src) + : "memory"); + } while ((((v.x & 0x7F7F7F7Fu) + 0x01010101u) & v.x & 0x80808080u) != 0u); + } + else + { + v = *reinterpret_cast(src); + } + *reinterpret_cast(smem + SX::kOXsf + i * kQkvSfAtom + c * 16) = v; + } + } +} + +// The MMAs (warp 1, elected lane): v52's order (per k-tile: B scale copy, 4 MMAs; one accumulator). +template +__device__ __forceinline__ void op_mma(uint8_t* smem, uint32_t tm) +{ + uint32_t const ob = smem_u32(smem + SX::kOBar); + mbar_wait_cta(ob + 8 * kOpTiles, 0); // a_ready + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + SX::kOXt)); + uint32_t const xsf0 = smem_u32(smem + SX::kOXsf); +#pragma unroll + for (int i = 0; i < kOpTiles; ++i) + { + uint32_t const sfb = tm + kOSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + umma_mx(tm + kOAcc, tm + kOA + j * 8, b0 + ((i * 1024 + kb * 32) >> 4), qkv_idesc(kb), tm + kOSfa + i * 4, + sfb, j > 0); + } + } + umma_commit(ob + 8 * (kOpTiles + 1)); // acc_done +} + +// Warps 16..19: the pair's split-K reduce (owner rows [64 s, 64 s + 64)) and the bf16 rows of y. +template +__device__ __forceinline__ void op_epilogue(Params const& p, uint8_t* smem, uint32_t tm, int tid, int ou) +{ + int const warp = tid >> 5; + int const lane = tid & 31; + uint32_t const ob = smem_u32(smem + SX::kOBar); + // The store's kernel parameters, fetched before the wait: a first touch of their constant-bank line costs a miss. + __nv_bfloat16* y_out = p.y_out; + int num_tokens = p.num_tokens; + int ar_mode = p.ar_mode; + asm volatile("" : "+l"(y_out), "+r"(num_tokens), "+r"(ar_mode)); +#ifndef M3FA_OP_PARPRE +#define M3FA_OP_PARPRE 1 // read the all-reduce slot parities before the accumulator wait (a global load off the tail) +#endif + // The parity flags are stable here: the all-reduce that flips them last completed before this grid's PDL wait, + // and the next one reads them only after this kernel's pushes. + int par_pre[kMaxTok]; + { + int const row_pre = (ou / kOpSplits) * 128 + ((tid >> 5) & 3) * 32 + (tid & 31); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + par_pre[t] = 0; + if (M3FA_OP_PARPRE && ar_mode != 0 && t < num_tokens) + { + asm volatile("ld.volatile.global.s32 %0, [%1];\n" + : "=r"(par_pre[t]) + : "l"(p.ar_flags + t * 8 + row_pre / (kOpN / 8)) + : "memory"); + } + } + } + mbar_wait_epi(ob + 8 * (kOpTiles + 1), 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF_T(51, kComputeThreads); + int const r = (warp & 3) * 32 + lane; + float v[16]; + tmem_ld16(tm + (((warp & 3) * 32) << 16) + kOAcc, v); + int const s = ou % kOpSplits; + int const owner = r / kOpOwnRows; + int const local = r % kOpOwnRows; + uint32_t const red = ob + 8 * (kOpTiles + 2); + uint32_t cta_rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + if (owner != s) + { + uint32_t const peer = cta_rank + owner - s; // o_proj units count up with the CTA index + uint32_t const dst + = mapa_shared(smem_u32(smem + SX::kOMailbox) + ((s * kOpOwnRows + local) * kQkvMbStride) * 4, peer); + uint32_t const pbar = mapa_shared(red, peer); +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( + dst + 16 * g), + "r"(__float_as_uint(v[4 * g])), "r"(__float_as_uint(v[4 * g + 1])), "r"(__float_as_uint(v[4 * g + 2])), + "r"(__float_as_uint(v[4 * g + 3])), "r"(pbar) + : "memory"); + } + } + else + { + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(red) + : "memory"); + } + M3FA_PROF_T(52, tid - lane); + float const* mb = reinterpret_cast(smem + SX::kOMailbox); + int const row = (ou / kOpSplits) * 128 + r; + float sums[kMaxTok]; +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + float4 part[kOpSplits]; +#pragma unroll + for (int q = 0; q < kOpSplits; ++q) + { + part[q] = *reinterpret_cast(mb + (q * kOpOwnRows + local) * kQkvMbStride + 4 * g); + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float sum = 0.0F; +#pragma unroll + for (int q = 0; q < kOpSplits; ++q) + { + float const pj = j == 0 ? part[q].x : j == 1 ? part[q].y : j == 2 ? part[q].z : part[q].w; + sum = sum + (q == s ? v[4 * g + j] : pj); + } + sums[4 * g + j] = sum; + } + } + M3FA_PROF_T(54, tid - lane); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + float const sum = sums[t]; +#ifndef M3FA_EXP +#define M3FA_EXP 0 +#endif + if (t < num_tokens) + { + if (M3FA_EXP == 1) + { + asm volatile( + "st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(y_out + static_cast(t) * kOpN + row), + "h"(__bfloat16_as_ushort(__float2bfloat16_rn(sum))) + : "memory"); + } + else + { + y_out[static_cast(t) * kOpN + row] = __float2bfloat16_rn(sum); + } + } + if (ar_mode != 0) + { + // m3_mx_gemv's push into the all-reduce exchange: lane pairs hold row pairs; the even lane writes the + // bf16x2 word into this rank's slot of the (token, column slice) parity the AR's flags select. + float const hi = __shfl_xor_sync(0xffffffffu, sum, 1); + if ((lane & 1) == 0 && t < num_tokens) + { + int par; + if (M3FA_OP_PARPRE) + { + par = par_pre[t]; + } + else + { + asm volatile("ld.volatile.global.s32 %0, [%1];\n" + : "=r"(par) + : "l"(p.ar_flags + t * 8 + row / (kOpN / 8)) + : "memory"); + } + int64_t const off = static_cast(par & 1) * p.ar_maxt * p.ar_world * (kOpN / 2) + + (static_cast(t) * p.ar_world + p.ar_rank) * (kOpN / 2) + row / 2; + __nv_bfloat162 w2 = __floats2bfloat162_rn(sum, hi); + uint32_t word = *reinterpret_cast(&w2); + // -0.0 halves to +0.0: 0x80000000 marks an empty exchange word (as m3_mx_gemv's _bf16x2_word) + word = (word & 0xFFFFu) == 0x8000u ? (word & 0xFFFF0000u) : word; + word = (word >> 16) == 0x8000u ? (word & 0x0000FFFFu) : word; + asm volatile("st.global.b32 [%0], %1;\n" ::"l"(p.ar_ws + off), "r"(word) : "memory"); + } + } + } + M3FA_PROF_T(53, tid - lane); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); +} + +// Packed-QKV element group [token][col * 128 + 4 * lane, +4): the split partials, polled until every one is stored, +// summed in split order and rounded to BF16. +__device__ __forceinline__ void qkv_load(Params const& p, int par, int tok, int col, int lane, float (&e)[4]) +{ + float const* base = p.part + par * kQkvPartStride + static_cast(tok) * kQkvN + col * kHeadDim + lane * 4; + uint4 pr[kQkvSplits]; + long long spins = 0; + bool ready; + do + { + ready = true; +#pragma unroll + for (int s = 0; s < kQkvSplits; ++s) + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(pr[s].x), "=r"(pr[s].y), "=r"(pr[s].z), "=r"(pr[s].w) + : "l"(base + static_cast(s) * kMaxTok * kQkvN) + : "memory"); + ready &= pr[s].x != kQkvSentinel && pr[s].y != kQkvSentinel && pr[s].z != kQkvSentinel + && pr[s].w != kQkvSentinel; + } + if (++spins > (1LL << 26)) + { + __trap(); + } + } while (!ready); + float s0 = __uint_as_float(pr[0].x), s1 = __uint_as_float(pr[0].y), s2 = __uint_as_float(pr[0].z), + s3 = __uint_as_float(pr[0].w); +#pragma unroll + for (int s = 1; s < kQkvSplits; ++s) + { + s0 += __uint_as_float(pr[s].x); + s1 += __uint_as_float(pr[s].y); + s2 += __uint_as_float(pr[s].z); + s3 += __uint_as_float(pr[s].w); + } + e[0] = __bfloat162float(__float2bfloat16_rn(s0)); + e[1] = __bfloat162float(__float2bfloat16_rn(s1)); + e[2] = __bfloat162float(__float2bfloat16_rn(s2)); + e[3] = __bfloat162float(__float2bfloat16_rn(s3)); +} + +// Every CTA once it has read the call parity and the partials it needs: the last one flips the parity. +__device__ __forceinline__ void qkv_done(Params const& p) +{ + if (atomicAdd(p.counters + 1, 1) == kQkvGrid - 1) + { + atomicExch(p.counters + 1, 0); + atomicXor(p.counters, 1); + } +} + +} // namespace + +namespace +{ + +// FQKV: the kernel also computes the QKV(+index) projection from the layer input (see qkv_mma) and runs on a fixed +// kQkvGrid-CTA grid; the first attn_clusters clusters then run the attention on the projection's output. +#ifndef M3FA_C1_OPISSUE +#define M3FA_C1_OPISSUE \ + 4 // c1: o_proj weight TMA issued 0: at entry (before the PDL wait), 1: right after the PDL + // wait, 2: once the attention's K/V reads landed, 3: after the merge, 4: before the push, 5: after it +#endif +#ifndef M3FA_C1_TRIGGER +#define M3FA_C1_TRIGGER 3 // c1: dependents released 2: at the attention's usual point, 3: after the o_proj pushes +#endif +template +__global__ void __launch_bounds__(kThreads, 1) m3_fused_decode_attn_kernel(__grid_constant__ Params const p) +{ +#if MINIMAX_M3_ATTN_DEVICE_ENABLED + static_assert(!OPJ || (!FQKV && MXFP8), "c1: packed input, MXFP8 hand-off"); + using S = std::conditional_t, SmemX>; + constexpr bool kAttnTrigger = !FQKV && !(OPJ && M3FA_C1_TRIGGER == 3); + constexpr int kRows = S::kRows; + constexpr int kTok = S::kTok; + constexpr int kCs = kBlocks / BPC; // cluster size + constexpr int kGroups = QLEN / QG; // clusters per request + constexpr uint32_t kPushBytes = (kHeadDim + 1) * 4; // one partial row: 128 outputs and the LSE + constexpr int kTasks = kHq * QG + 3 * QLEN; // producer tasks: 16 Q heads x QG, [K, V, index-K] x QLEN + constexpr int kMaxT = (kTasks + kWarps - 1) / kWarps; + constexpr int kChunks = BPC * kPage * 4; // 16-byte packed chunks per K (or V) tile set + constexpr int kKPer = (kChunks + kComputeThreads - 1) / kComputeThreads; // K chunks per thread [0, 512) + constexpr bool kRegCore = RegCore::kOn; // warp-MMA attention core (see RegCore) + // V tiles (raw copy before the PDL wait, dequantized during QK and the softmax): on the packed path every warp + // outside the softmax groups takes part (at M=1, warps 4..19), on the fused-projection path the producer warps. +#ifndef M3FA_RS2 +#define M3FA_RS2 1 // megakernel at one 16-row chunk: the softmax and push rows split over two warp groups (8 each) +#endif + // Split softmax: the megakernel and multi-row (spec verify) steps; at one-row decode on the packed path it + // measured slower (it takes four warps from the V dequantization there). + constexpr bool kRs2 = M3FA_RS2 && BPC * QG == 1 && !RegCore::kOn && (FQKV || QLEN > 1); +#ifndef M3FA_MEGA_WARPMMA +#define M3FA_MEGA_WARPMMA 1 +#endif + // Warp-issued MMAs and the early cluster barrier: measured gains in the megakernel and on multi-row / paired-block + // steps, but not at one-row decode on 16-CTA clusters (the packed path's (1, 1, 1)), where they stay off. + constexpr bool kTuned = FQKV || QLEN > 1 || BPC > 1; + constexpr bool kWarpMma = M3FA_MEGA_WARPMMA && kTuned; + constexpr bool kEarlyClw = M3FA_MEGA_EARLYCLW && kTuned; + constexpr int kSoftmaxGroups = kRs2 ? 2 : (BPC * QG) < 4 ? (BPC * QG) : 4; +#ifndef M3FA_MEGA_VWIDE +#define M3FA_MEGA_VWIDE \ + 1 // megakernel attention CTAs (no projection unit): every warp outside the softmax copies and + // dequantizes V, as on the packed path (not only the four producer warps) +#endif + constexpr int kVStart = (kRegCore || (FQKV && !M3FA_MEGA_VWIDE)) ? kComputeThreads : 128 * kSoftmaxGroups; + constexpr int kVThreads = kThreads - kVStart; // V chunks: the warps outside the softmax groups + constexpr int kVPer = (kChunks + kVThreads - 1) / kVThreads; + constexpr int kPairs = BPC * QG; // (block, 16-row chunk) pairs of the softmax and push + constexpr int kPairsPerGroup = (kPairs + 3) / 4; // per group of four compute warps + // Named barrier 6: the softmax warps (P written) and the producer warps (V written) before the PV issue. + constexpr int kPvBarThreads = 128 * kSoftmaxGroups + kVThreads; + extern __shared__ __align__(128) uint8_t smem[]; + float* row_scale = reinterpret_cast(smem + S::kRowScale); + float* row_lse = reinterpret_cast(smem + S::kRowLse); + float* red_m = reinterpret_cast(smem + S::kRedM); + float* red_l = reinterpret_cast(smem + S::kRedL); + float* recv = reinterpret_cast(smem + S::kRecv); + uint32_t const mbar = smem_u32(smem + S::kMbar); + uint32_t const mma_bar = mbar + 8; + uint32_t* tmem_slot = reinterpret_cast(smem + S::kMbar + 16); + + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const cr = blockIdx.x % kCs; // cluster rank; owns blocks [cr * BPC, cr * BPC + BPC) + int const ci = blockIdx.x / kCs; + int const req = ci / kGroups; + int const tfirst = (ci % kGroups) * QG; // this cluster's query tokens [tfirst, tfirst + QG) of the request + int const tok0 = req * QLEN; + uint32_t const gbar = smem_u32(smem + S::kGBar); // FQKV: weight ring, A ready, accumulators done + int const qkv_u = FQKV ? qkv_unit() : -1; // FQKV: this CTA's projection unit, or -1 + + // ---------------- Prologue: nothing here depends on the producer kernel ---------------- +#if M3FA_TRIGGER_AT == 0 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (!FQKV) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + M3FA_PROF(0); +#ifdef M3FA_PROFILE + if (p.prof != nullptr && tid == 0) + { + uint32_t smid; + asm volatile("mov.u32 %0, %%smid;" : "=r"(smid)); + p.prof[blockIdx.x * 64 + 40] = smid; + } +#endif + if constexpr (FQKV) + { + // The projection weights stream into TMEM while the producer (the all-reduce) still runs. + if (warp < 2) + { + qkv_setup(p, smem, gbar, tmem_slot, warp, lane, qkv_u); + } + if (ci >= p.attn_clusters) + { + // A projection-only CTA. Every reduce barrier of the cluster is initialized before any peer pushes. + asm volatile("barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;\n" :: + : "memory"); + int const op_u = static_cast(blockIdx.x) - kOpFirstCta; // this CTA's o_proj unit (>= 0 here) +#ifndef M3FA_OP_EARLY +#define M3FA_OP_EARLY 0 // 1: o_proj-only CTAs stream their weights from launch (0: after the PDL wait) +#endif + if (M3FA_OP_EARLY && qkv_u < 0 && warp == 0 && lane == 0) + { + op_issue(p, smem, op_u); // an o_proj-only CTA: its weights stream from launch on + } + if (M3FA_OP_L2PF && qkv_u >= 0 && warp == 0 && lane == 0) + { + op_prefetch_l2(p, op_u); // into L2 during the all-reduce, behind this CTA's QKV ring + } + qkv_stream(p, smem, gbar, tmem_slot, warp, lane, qkv_u); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (M3FA_MEGA_TRIGGER == 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + // This CTA's call counter: read after the PDL wait (the previous call's CTA increments it at its end), and + // early, so its latency stays off the epilogue's critical path. + unsigned int const calls = p.flags[32 + blockIdx.x]; + M3FA_PROF(1); + if (!M3FA_OP_EARLY && qkv_u < 0 && warp == 0 && lane == 0) + { + op_issue(p, smem, op_u); + } + if (qkv_u >= 0) + { + int const par = 0; + qkv_mma(p, smem, gbar, tmem_slot, tid, qkv_u); + if (warp >= kComputeWarps) + { + qkv_epilogue(p, smem, gbar, tmem_slot, tid, qkv_u, par, calls); + } + if (warp == 0 && lane == 0) + { + // The ring (and the activation over it) is free once the QKV MMAs completed. + mbar_wait_cta(gbar + 16 * kQkvRing + 8, 0); + op_issue(p, smem, op_u); + if (p.pf_b0 > 0 || p.pf_b1 > 0) + { + window_prefetch_l2(p, qkv_u); + } + } + } + __syncthreads(); // the QKV accumulator reads are done: TMEM columns [0, 512) are the o_proj's + M3FA_PROF(2); + uint32_t const tm = *tmem_slot; + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_stream(smem, tm); + } + __syncwarp(); + } + op_act(p, smem, op_u, tid, calls); + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + M3FA_PROF_T(36, 0); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Late release: the dependents launch once every unit has its attention input (the rel_v4 rule). + if (M3FA_MEGA_TRIGGER < 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + __syncthreads(); + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_mma(smem, tm); + M3FA_PROF_T(50, threadIdx.x); + } + __syncwarp(); + } + if (warp >= kComputeWarps) + { + op_epilogue(p, smem, tm, tid, op_u); + } + M3FA_PROF_T(37, kComputeThreads); + __syncthreads(); + M3FA_PROF_T(55, 0); + if (M3FA_MEGA_TRIGGER == 3) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // after this CTA's y rows / pushes + } + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" ::"r"(*tmem_slot), "r"(kTmemAll) + : "memory"); + M3FA_PROF_T(56, 32); + } + if (tid == 0) + { + p.flags[32 + blockIdx.x] = calls + 1u; // a plain store: no load on the exit path + } + M3FA_PROF(39); +#ifndef M3FA_FINAL_CLBAR +#define M3FA_FINAL_CLBAR 0 +#endif + if (M3FA_FINAL_CLBAR) + { + // No peer may exit while its partials can still be pushed into it (DSMEM lifetime). + asm volatile("barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;\n" :: + : "memory"); + } + M3FA_PROF(38); + return; + } + } +#ifndef M3FA_MEGA_ATTN_WAIT_FIRST +#define M3FA_MEGA_ATTN_WAIT_FIRST 0 // 1: the attention CTAs read nothing (metadata, caches) before their PDL wait +#endif +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (FQKV && M3FA_MEGA_ATTN_WAIT_FIRST) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } +#endif + if (!FQKV && warp == kWarps - 1) // a producer warp: it has nothing else to do before the PDL wait + { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n" ::"r"(smem_u32(tmem_slot)), + "r"(S::kTmemCols) + : "memory"); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n" ::: "memory"); + } + int const seq_len = p.seq_lens[req]; + int page[BPC]; +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + int const b = cr * BPC + h; + page[h] = b < p.bt_cols ? p.block_table[static_cast(req) * p.bt_stride + b] : 0; + } + // New-token cache slots, for the tokens up to this cluster's last one. + int slot[QLEN]; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + slot[t] = t < tfirst + QG ? p.out_cache_loc[tok0 + t] : -1; + } + // This warp's producer tasks and, for every lane that applies RoPE, the task token's cos/sin. + float cs[kMaxT][4], sn[kMaxT][4]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + int const t = i < kHq * QG ? tfirst + i / kHq : (i - kHq * QG) / 3; + if (i < kTasks && lane < 16 && t < QLEN) + { + int64_t const row = static_cast(p.position_ids[tok0 + t]) * kRotDim; +#pragma unroll + for (int e = 0; e < 4; ++e) + { + int const coefficient = (lane * 4 + e) & (kRotDim / 2 - 1); + cs[k][e] = p.rope[row + coefficient]; + sn[k][e] = p.rope[row + kRotDim / 2 + coefficient]; + } + } + else + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + cs[k][e] = 1.0F; + sn[k][e] = 0.0F; + } + } + } + float qw[4], kw[4], ikw[4]; + load_w4(qw, p.q_w, lane); + load_w4(kw, p.k_w, lane); + load_w4(ikw, p.ik_w, lane); + float const kqs = p.kv_quant_scale[1]; + float const vqs = p.kv_quant_scale[2]; + float const sl2e = (p.sm_scale * 1.4426950409f) * p.k_gs[0]; + float const vgs = p.v_gs[0]; + // Longest attended length among this cluster's tokens, and each block's rows within it. + int const kv_max = max(seq_len - QLEN + tfirst + QG, 0); + int nvh[BPC]; +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + nvh[h] = min(max(kv_max - (cr * BPC + h) * kPage, 0), kPage); + } + bool const any_valid = nvh[0] > 0; + if (tid == 32) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(mbar), "r"(1)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(mma_bar), "r"(1)); + asm volatile( + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(mbar), "r"(kBlocks * QG * BPC * kPushBytes) + : "memory"); + if constexpr (OPJ) + { + uint32_t const ob = smem_u32(smem + S::kOBar); + for (int i = 0; i < kOpTiles + 3; ++i) + { + mbar_init(ob + 8 * i, 1); + } + mbar_expect_tx(ob + 8 * (kOpTiles + 2), (kOpSplits - 1) * kOpOwnRows * kMaxTok * 4); + } + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + if constexpr (OPJ) + { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if (M3FA_C1_OPISSUE == 0) + { + op_issue(p, smem, static_cast(blockIdx.x)); + } + } + } + // This thread's packed K / V chunks, held in registers until they are dequantized after the PDL wait. + RawChunk kraw[kKPer]; +#pragma unroll + for (int k = 0; k < kKPer; ++k) + { + int const q = tid + kComputeThreads * k; + kraw[k] = RawChunk{make_uint4(0u, 0u, 0u, 0u), 0u}; + int const h = q / (kPage * 4); + if (q < kChunks && tid < kComputeThreads && nvh[h] > 0) + { + kraw[k] = load_chunk(p.kv_data + static_cast(page[h]) * p.d_page, + p.kv_scale + static_cast(page[h]) * p.s_page, chunk_row(q % (kPage * 4)), + chunk_col(q % (kPage * 4)), false); + } + } + // V goes through cp.async into shared memory: an asynchronous copy does not hold the producer warps up the way + // an outstanding register load does, and they need V only after QK has been issued. + if (tid >= kVStart) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + if (nvh[h] > 0) + { + uint8_t* dst = smem + S::kRawV + h * S::kRawVBlock; + uint8_t const* data = p.kv_data + static_cast(page[h]) * p.d_page + p.d_plane; + uint8_t const* scales = p.kv_scale + static_cast(page[h]) * p.s_page + p.s_plane; + for (int c = tid - kVStart; c < S::kRawVBlock / 16; c += kVThreads) + { + cp_async16(dst + c * 16, c < kPage * 4 ? data + c * 16 : scales + (c - kPage * 4) * 16); + } + } + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + if (warp == kWarps - 1) + { + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + // Every CTA's mbarrier is initialized before any peer pushes into it; the matching wait sits right before the + // push, so this barrier's latency hides behind the attention (with aliased partials it moves there). + if constexpr (!S::kRecvAlias) + { + asm volatile("barrier.cluster.arrive.relaxed.aligned;\n" ::: "memory"); + } + int kv_len[QG]; +#pragma unroll + for (int tl = 0; tl < QG; ++tl) + { + kv_len[tl] = max(seq_len - QLEN + tfirst + tl + 1, 0); + } + int const cta_tok0 = cr * kTok; // request position of this CTA's first block row + // New token t's row in this CTA's tiles (block * 128 + row), or -1: the CTA owns the block holding the token's + // position, and the token is attended by this cluster (t <= its last token). Every thread derives the same. + int new_row[QLEN]; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + int const pos = seq_len - QLEN + t; + int const blk = min(pos >> 7, kBlocks - 1); + int const spage = slot[t] / kPage; + bool const mine = t < tfirst + QG && pos >= 0 && blk / BPC == cr && slot[t] >= 0 && spage < p.num_pages; + new_row[t] = mine && spage == page[blk % BPC] ? (blk % BPC) * kPage + (slot[t] & (kPage - 1)) : -1; + } + auto is_new_row = [&](int r) + { + bool hit = false; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + hit |= new_row[t] == r; + } + return hit; + }; + // Cache writes of the new K/V/index-K are issued after this CTA has pushed its partials: only the fp16 tile rows + // are on the attention's path. kind 0: none, 1: NVFP4 data + scale, 2: FP8 index-K. + int wr_kind[kMaxT]; + uint8_t* wr_ptr[kMaxT]; + uint8_t* wr_sptr[kMaxT]; + uint32_t wr_val[kMaxT]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + wr_kind[k] = 0; + wr_ptr[k] = nullptr; + wr_sptr[k] = nullptr; + wr_val[k] = 0u; + } + uint32_t tmem = 0; + auto dequant_k = [&]() + { + if (tid < kComputeThreads) + { +#pragma unroll + for (int k = 0; k < kKPer; ++k) + { + int const q = tid + kComputeThreads * k; + if (q < kChunks) + { + int const h = q / (kPage * 4); + int const row = chunk_row(q % (kPage * 4)); + if (!is_new_row(h * kPage + row)) + { + store_chunk_fp16( + smem + S::kK16 + h * S::kTileBytes, row, chunk_col(q % (kPage * 4)), kraw[k], row < nvh[h]); + } + } + } + } + }; + // The producer warps' V tiles (raw tiles from cp.async, dequantized to fp16), then their PV-barrier arrival. + auto dequant_v = [&]() + { +#ifdef M3FA_DBG_NO_VDQ + if constexpr (!kRegCore) + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + return; +#endif + asm volatile("cp.async.wait_all;\n" ::: "memory"); + asm volatile("bar.sync 4, %0;\n" ::"n"(kVThreads) : "memory"); +#pragma unroll + for (int k = 0; k < kVPer; ++k) + { + int const q = tid - kVStart + kVThreads * k; + if (q < kChunks) + { + int const h = q / (kPage * 4); + int const row = chunk_row(q % (kPage * 4)); + int const cc = chunk_col(q % (kPage * 4)); + if (!is_new_row(h * kPage + row)) + { + uint8_t const* rawv = smem + S::kRawV + h * S::kRawVBlock; + uint8_t const* sc = rawv + kPage * 64 + (row >> 2) * 32 + (row & 3); + RawChunk c; + c.data = *reinterpret_cast(rawv + row * 64 + cc * 16); + c.scales + = static_cast(sc[(cc * 2) * 4]) | (static_cast(sc[(cc * 2 + 1) * 4]) << 8); + store_chunk_fp16(smem + S::kV16 + h * S::kTileBytes, row, cc, c, row < nvh[h]); + } + } + } + M3FA_FENCE_PROXY_ASYNC(); // V tiles -> tensor core reads (ordered by the barrier before the PV issue) + if constexpr (!kRegCore) + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + M3FA_PROF_T(16, kComputeThreads); + M3FA_PROF_T(17, kThreads - 1); + }; + if constexpr (FQKV) + { + qkv_stream(p, smem, gbar, tmem_slot, warp, lane, qkv_u); + } + if constexpr (!FQKV) + { + if (M3FA_V15_EARLYV && any_valid && tid >= kVStart) + { + dequant_v(); + } + if constexpr (kEarlyClw && !S::kRecvAlias) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); + } + } + M3FA_PROF(11); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (FQKV && M3FA_MEGA_TRIGGER == 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#if M3FA_TRIGGER_AT == 1 + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +#endif +#endif + // FQKV (megakernel): this call's QKV rows sit in buffer (calls & 1), polled value by value below. + unsigned int const mega_calls = (FQKV || OPJ) ? p.flags[32 + blockIdx.x] : 0u; + if (OPJ && M3FA_C1_OPISSUE == 1 && tid == 32) + { + op_issue(p, smem, static_cast(blockIdx.x)); + } + __nv_bfloat16 const* qkv_src = FQKV ? p.packed + (mega_calls & 1u) * kMaxTok * kQkvN : p.packed; + M3FA_PROF(1); + int const par = 0; + if constexpr (FQKV) + { + if (qkv_u >= 0) + { + qkv_mma(p, smem, gbar, tmem_slot, tid, qkv_u); + } + } + + if constexpr (FQKV) + { + if (qkv_u >= 0) + { + if (warp >= kComputeWarps) + { + qkv_epilogue(p, smem, gbar, tmem_slot, tid, qkv_u, par, mega_calls); + } + else + { + // The attention regions overlay the weight ring: every weight byte has gone on into TMEM (or, with + // the MMAs reading the ring, every MMA has completed) first. + mbar_wait_cta(gbar + 16 * kQkvRing + (kQkvASmem ? 8 : 0), 0); + } + } + if (any_valid) + { + dequant_k(); + } +#ifndef M3FA_MEGA_EARLYV +#define M3FA_MEGA_EARLYV 1 // attention CTAs dequantize V here, while the QKV rows are still on their way +#endif +#ifndef M3FA_MEGA_EARLYCLW +#define M3FA_MEGA_EARLYCLW 1 // and pass the cluster barrier (peers' merge mbarriers initialized) here too +#endif + if (M3FA_MEGA_EARLYV && any_valid && tid >= kVStart) + { + dequant_v(); + } + if constexpr (kEarlyClw && !S::kRecvAlias) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); + } + M3FA_PROF(14); + } + + if (any_valid) // CTA-uniform + { + // ---------------- Producer ---------------- + // Q rows of this cluster's tokens in every CTA. The new K/V of every token up to this cluster's last one + // whose position falls in this CTA's blocks go straight into the fp16 tiles; the CTA writes the caches, and + // the index-K, only for its own tokens. + uint2 raw[kMaxT]; + float fq[FQKV ? kMaxT : 1][4]; + bool act[kMaxT]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + raw[k] = make_uint2(0u, 0u); + act[k] = false; + if (i < kTasks) + { + int t, col; + bool active = true; + if (i >= kHq * QG) + { + int const kv = i - kHq * QG; + t = kv / 3; + int const kind = kv % 3; + col = kind == 0 ? kColK : (kind == 1 ? kColV : kColIdxK); + int const pos = seq_len - QLEN + t; + active = t < tfirst + QG && pos >= 0 && min(pos >> 7, kBlocks - 1) / BPC == cr && slot[t] >= 0 + && slot[t] / kPage < p.num_pages && (kind != 2 || t >= tfirst); + } + else + { + t = tfirst + i / kHq; + col = i % kHq; + } + act[k] = active; + if (active) + { + uint2 const* src = reinterpret_cast( + qkv_src + static_cast(tok0 + t) * p.packed_stride + col * kHeadDim + lane * 4); + if constexpr (FQKV) + { + uint2 v; + bool pending; + do + { + asm volatile("ld.relaxed.gpu.global.v2.b32 {%0, %1}, [%2];\n" + : "=r"(v.x), "=r"(v.y) + : "l"(src) + : "memory"); + pending = (v.x & 0xFFFFu) == 0xFFFFu || (v.x >> 16) == 0xFFFFu || (v.y & 0xFFFFu) == 0xFFFFu + || (v.y >> 16) == 0xFFFFu; + } while (pending); + raw[k] = v; + } + else + { + raw[k] = *src; + } + } + } + } + M3FA_PROF(15); + // While the packed-QKV loads are in flight: dequantize this thread's K chunks to fp16 (exact: e2m1 x e4m3 + // fits fp16). Rows past the longest attended length are zeroed, as the Triton producer masks its loads; the + // new tokens' rows are written by their producer tasks instead. (Building the tiles before the PDL wait was + // measured slower: the cache loads issued at entry are still in flight there, so the wait moves later.) + if constexpr (!FQKV) + { + dequant_k(); + } +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + if (i >= kTasks || !act[k]) + { + continue; + } + float e[4]; + { + float2 const a = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw[k].x)); + float2 const b = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw[k].y)); + e[0] = a.x; + e[1] = a.y; + e[2] = b.x; + e[3] = b.y; + } + if (i < kHq * QG) + { + norm_rope(e, qw, p.eps, cs[k], sn[k], lane); + __nv_fp8x2_e4m3 const lo(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const hi(make_float2(e[2], e[3])); + uint2 hv; + hv.x = e4m3x2_to_h2(lo.__x); + hv.y = e4m3x2_to_h2(hi.__x); + *reinterpret_cast(smem + S::kQ + kmaj_off(i, lane * 4)) = hv; // row i = token * 16 + head + continue; + } + int const t = (i - kHq * QG) / 3; + int const kind = (i - kHq * QG) % 3; + int const spage = slot[t] / kPage; + int const within = slot[t] & (kPage - 1); + if (kind == 2) + { + norm_rope(e, ikw, p.eps, cs[k], sn[k], lane); + __nv_fp8x2_e4m3 const lo(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const hi(make_float2(e[2], e[3])); + wr_kind[k] = 2; + wr_val[k] = static_cast(lo.__x) | (static_cast(hi.__x) << 16); + wr_ptr[k] = p.idx_k + static_cast(spage) * p.i_page + static_cast(within) * p.i_tok + + lane * 4; + continue; + } + bool const isV = kind == 1; + if (!isV) + { + norm_rope(e, kw, p.eps, cs[k], sn[k], lane); + } + uint8_t sbyte = 0; + uint16_t const fp4 = quantizeMinimaxM3Fp4x4(e, isV ? vqs : kqs, &sbyte); + if (t >= tfirst) + { + int const scale_col = lane / 4; + int const scale_within = isV ? (within / 4) * 32 + scale_col * 4 + within % 4 : within * 8 + scale_col; + wr_kind[k] = 1; + wr_val[k] = static_cast(fp4) | (static_cast(sbyte) << 16); + wr_ptr[k] = p.kv_data + static_cast(spage) * p.d_page + (isV ? p.d_plane : 0) + within * 64 + + lane * 2; + wr_sptr[k] = p.kv_scale + static_cast(spage) * p.s_page + (isV ? p.s_plane : 0) + scale_within; + } + if (new_row[t] >= 0) + { + // The attention reads this page: its fp16 tile row is the freshly quantized row, dequantized like + // the rest of the tile. + int const h = new_row[t] / kPage; + int const row = new_row[t] % kPage; + uint32_t const sc2 = e4m3_to_h2_splat(sbyte); + uint2 hv; + hv.x = hmul2_u(e2m1x2_to_h2(fp4 & 0xFFu), sc2); + hv.y = hmul2_u(e2m1x2_to_h2(fp4 >> 8), sc2); + *reinterpret_cast( + smem + (isV ? S::kV16 : S::kK16) + h * S::kTileBytes + kmaj_off(row, lane * 4)) + = hv; + } + } + // K tiles, Q rows and patched rows -> tensor core reads. Only threads that wrote fence: the fence also + // waits for the thread's outstanding loads, and the producer warps still have their V chunks in flight. + bool wrote = tid < kComputeThreads; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + wrote |= act[k]; + } + if (wrote) + { + M3FA_FENCE_PROXY_ASYNC(); + } + M3FA_PROF(13); + __syncthreads(); +#if M3FA_TRIGGER_AT == 2 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Every K/V byte this CTA reads has arrived (the V chunks sit in registers): release the dependents now so + // their weight streams do not queue ahead of the KV reads (the megakernel releases late, from its o_proj). + if constexpr (kAttnTrigger || (FQKV && M3FA_MEGA_TRIGGER == 1)) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + if (OPJ && M3FA_C1_OPISSUE == 2 && tid == 32) + { + op_issue(p, smem, static_cast(blockIdx.x)); + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + tmem = *tmem_slot; + M3FA_PROF(2); + // ---------------- S^T = K Q^T per block on the tensor core: TMEM lane = token, column = query row ---------- +#ifndef M3FA_MEGA_WARPMMA +#define M3FA_MEGA_WARPMMA 1 // attention MMAs issued by warp 0's elected lane with warp-uniform operands (not tid 0) +#endif + if (!kRegCore && (kWarpMma ? warp == 0 : tid == 0)) + { + // Operands computed by the whole warp stay in uniform registers; one elected lane issues. + uint32_t const tm_u = tmem; + bool const issue = kWarpMma ? elect_one() : true; + if (issue) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + uint32_t const a0 = smem_u32(smem + S::kK16 + h * S::kTileBytes); + uint32_t const b0 = smem_u32(smem + S::kQ); +#pragma unroll + for (int ks = 0; ks < kHeadDim / 16; ++ks) + { + umma_f16(tm_u + h * kRows, umma_desc(a0 + ks * 256, 128, 2048), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 0, 0), ks > 0); + } + } + umma_commit(mma_bar); + } + if (kWarpMma) + { + __syncwarp(); + } + } + if (tid >= kVStart && !(FQKV && M3FA_MEGA_EARLYV) && !(!FQKV && M3FA_V15_EARLYV)) + { + // The producer warps dequantize V while QK and the softmax run. + dequant_v(); + } + if constexpr (kRegCore) + { + // ---------------- Register core: warp w of block h = w / kWpb owns 8 * kNt tokens for S and 8 * kNt + // output dims for O; m16n8k16 fp16 MMAs with fp32 accumulators, the softmax in the S fragments ---------- + constexpr int kWpb = kComputeWarps / BPC; // warps per block + constexpr int kNt = 16 / kWpb; // 8-token (S) or 8-dim (O) tiles per warp + float* redx_m = reinterpret_cast(smem + S::kRedX); + float* redx_l = redx_m + BPC * kRows * kWpb; + if (warp < kComputeWarps) + { + int const h = warp / kWpb; + int const wib = warp % kWpb; + int const g = lane >> 2; + int const q4 = lane & 3; + uint32_t const ktile = smem_u32(smem + S::kK16 + h * S::kTileBytes); + uint32_t const vtile = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const qrows = smem_u32(smem + S::kQ); + uint8_t* ptile = smem + S::kK16 + h * S::kTileBytes; // P over this block's K tile, after QK + auto block_bar = [&]() + { + if (h == 0) + { + asm volatile("bar.sync 14, %0;\n" ::"n"(kWpb * 32) : "memory"); + } + else + { + asm volatile("bar.sync 15, %0;\n" ::"n"(kWpb * 32) : "memory"); + } + }; + // S = Q K^T: rows (c, g / g + 8), tokens 8 tb + 2 q4 (+1). + float sv[QG][kNt][4]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + sv[c][nt][e] = 0.0F; + } + } + } +#pragma unroll + for (int ks = 0; ks < kHeadDim / 16; ++ks) + { + uint32_t bk[kNt][2]; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const tb = wib * kNt + nt; + ldsm_x2(ktile + tb * 2048 + (ks * 2 + ((lane >> 3) & 1)) * 128 + (lane & 7) * 16, bk[nt]); + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { + uint32_t aq[4]; + ldsm_x4( + qrows + (c * 2 + ((lane >> 3) & 1)) * 2048 + (ks * 2 + (lane >> 4)) * 128 + (lane & 7) * 16, + aq); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + mma_f16_16816(sv[c][nt], aq, bk[nt]); + } + } + } + M3FA_PROF(7); + // Row maxima (log2 domain, masked as the producer masks), quad then across the block's warps. + float mrow[QG][2]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { + float mx0 = kNegInf, mx1 = kNegInf; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + int const tok = cta_tok0 + h * kPage + (wib * kNt + nt) * 8 + 2 * q4 + (e & 1); + bool const valid = tok < kv_max && tok < kv_len[c]; + float const v = valid ? sv[c][nt][e] * sl2e : kNegInf; + sv[c][nt][e] = v; + if (e < 2) + { + mx0 = fmaxf(mx0, v); + } + else + { + mx1 = fmaxf(mx1, v); + } + } + } + mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 1)); + mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 2)); + mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 1)); + mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 2)); + if (q4 == 0) + { + redx_m[(h * kRows + c * 16 + g) * kWpb + wib] = mx0; + redx_m[(h * kRows + c * 16 + g + 8) * kWpb + wib] = mx1; + } + } + block_bar(); // every warp of the block is past QK (P may overwrite the K tile) and wrote its maxima + M3FA_PROF(8); +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int half = 0; half < 2; ++half) + { + float const* rm = redx_m + (h * kRows + c * 16 + g + 8 * half) * kWpb; + float m = kNegInf; +#pragma unroll + for (int w4 = 0; w4 < kWpb; w4 += 4) + { + float4 const v = *reinterpret_cast(rm + w4); + m = fmaxf(m, fmaxf(fmaxf(v.x, v.y), fmaxf(v.z, v.w))); + } + mrow[c][half] = m; + } + // P = exp2(s - m) as fp16 into [row][token] core matrices; the fp32 values feed the row sums. + float l0 = 0.0F, l1 = 0.0F; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + float a[4]; +#pragma unroll + for (int e = 0; e < 4; ++e) + { + float const m = mrow[c][e >> 1]; + a[e] = m == kNegInf ? 0.0F : ex2(sv[c][nt][e] - m); + } + l0 += a[0] + a[1]; + l1 += a[2] + a[3]; + int const t0 = (wib * kNt + nt) * 8 + 2 * q4; + int const r0 = c * 16 + g; + __half2 const p0 = __floats2half2_rn(a[0], a[1]); + __half2 const p1 = __floats2half2_rn(a[2], a[3]); + *reinterpret_cast<__half2*>( + ptile + (r0 >> 3) * 2048 + (t0 >> 3) * 128 + (r0 & 7) * 16 + (t0 & 7) * 2) + = p0; + *reinterpret_cast<__half2*>( + ptile + ((r0 + 8) >> 3) * 2048 + (t0 >> 3) * 128 + (r0 & 7) * 16 + (t0 & 7) * 2) + = p1; + } + l0 += __shfl_xor_sync(0xffffffffu, l0, 1); + l0 += __shfl_xor_sync(0xffffffffu, l0, 2); + l1 += __shfl_xor_sync(0xffffffffu, l1, 1); + l1 += __shfl_xor_sync(0xffffffffu, l1, 2); + if (q4 == 0) + { + redx_l[(h * kRows + c * 16 + g) * kWpb + wib] = l0; + redx_l[(h * kRows + c * 16 + g + 8) * kWpb + wib] = l1; + } + } + M3FA_PROF_T(19, 0); + } + __syncthreads(); // P, the V tiles (producer warps) and the row sums visible + M3FA_PROF(9); + if (warp < kComputeWarps) + { + int const h = warp / kWpb; + int const wib = warp % kWpb; + int const g = lane >> 2; + int const q4 = lane & 3; + uint32_t const vtile = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const ptile = smem_u32(smem + S::kK16 + h * S::kTileBytes); + float* obuf = reinterpret_cast(smem + S::kK16 + h * S::kTileBytes + S::kObufOff); + // O = P V: rows (c, g / g + 8), dims 8 dt + 2 q4 (+1). + float ov[QG][kNt][4]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + ov[c][nt][e] = 0.0F; + } + } + } +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + uint32_t bv[kNt][2]; + int const tok = ks * 16 + ((lane >> 3) & 1) * 8 + (lane & 7); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const dt = wib * kNt + nt; + ldsm_x2_trans(vtile + (tok >> 3) * 2048 + dt * 128 + (tok & 7) * 16, bv[nt]); + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { + uint32_t ap[4]; + ldsm_x4( + ptile + (c * 2 + ((lane >> 3) & 1)) * 2048 + (ks * 2 + (lane >> 4)) * 128 + (lane & 7) * 16, + ap); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + mma_f16_16816(ov[c][nt], ap, bv[nt]); + } + } + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const d = (wib * kNt + nt) * 8 + 2 * q4; + int const r0 = c * 16 + g; + *reinterpret_cast(obuf + r0 * kOPitch + d) = make_float2(ov[c][nt][0], ov[c][nt][1]); + *reinterpret_cast(obuf + (r0 + 8) * kOPitch + d) + = make_float2(ov[c][nt][2], ov[c][nt][3]); + } + } + } + else if (tid - kComputeThreads < BPC * kRows) + { + // Row statistics of every (block, row): LSE and the PV scale, as the producer derives them. + int const r = tid - kComputeThreads; + float const* rm = redx_m + r * kWpb; + float const* rl = redx_l + r * kWpb; + float m = kNegInf; + float l = 0.0F; +#pragma unroll + for (int w4 = 0; w4 < kWpb; w4 += 4) + { + float4 const vm = *reinterpret_cast(rm + w4); + float4 const vl = *reinterpret_cast(rl + w4); + m = fmaxf(m, fmaxf(fmaxf(vm.x, vm.y), fmaxf(vm.z, vm.w))); + l += (vl.x + vl.y) + (vl.z + vl.w); + } + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[r] = scale * vgs; + row_lse[r] = lse; + } + __syncthreads(); // O partials and row statistics visible to the push + } + else + { + if constexpr (kRs2) + { + // ---------------- Split softmax: warps 0..3 rows 0..7, warps 4..7 rows 8..15 (thread = block token) + // -------- + if (warp < 8) + { + int const rg = warp / 4; + int const wq = warp % 4; + int const tk = wq * 32 + lane; + auto rg_bar = [&]() + { + if (rg == 0) + { + asm volatile("bar.sync 10, 128;\n" ::: "memory"); + } + else + { + asm volatile("bar.sync 11, 128;\n" ::: "memory"); + } + }; + int const rs8 = ((lane >> 4) & 1) * 4 + ((lane >> 3) & 1) * 2 + ((lane >> 2) & 1); + mbar_wait_cta(mma_bar, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF(7); + float v[8]; + tmem_ld8(tmem + ((wq * 32) << 16) + 8 * rg, v); + bool const valid = cta_tok0 + tk < kv_max && cta_tok0 + tk < kv_len[0]; + float a[8]; +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] = valid ? v[j2] * sl2e : kNegInf; + a[j2] = v[j2]; + } + float const mw = warp_reduce_scatter8(a, lane); + if ((lane & 3) == 0) + { + red_m[(8 * rg + rs8) * 4 + wq] = mw; + } + rg_bar(); + M3FA_PROF(8); + uint8_t* pdst = smem + S::kP + (tk >> 3) * 128 + (tk & 7) * 16; + float mrow = kNegInf; + if (lane < 8) + { + float4 const mq = *reinterpret_cast(red_m + (8 * rg + lane) * 4); + mrow = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + } + float pa[8]; +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + float const m = __shfl_sync(0xffffffffu, mrow, j2); + pa[j2] = m == kNegInf ? 0.0F : ex2(v[j2] - m); + } + if (tid == 0) + { + M3FA_CLK(41); + } + uint32_t pw[4]; +#pragma unroll + for (int j2 = 0; j2 < 4; ++j2) + { + __half2 const hp = __floats2half2_rn(pa[2 * j2], pa[2 * j2 + 1]); + pw[j2] = *reinterpret_cast(&hp); + } + *reinterpret_cast(pdst + rg * 2048) = make_uint4(pw[0], pw[1], pw[2], pw[3]); + M3FA_PROF_T(18, 0); + if (tid == 0) + { + M3FA_CLK(42); + } + M3FA_FENCE_PROXY_ASYNC(); // P -> tensor core reads + if (tid == 0) + { + M3FA_CLK(43); + } + if (warp == 0) + { + asm volatile("bar.sync 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + if (tid == 0) + { + M3FA_CLK(44); + } + } + else + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + if (kWarpMma ? warp == 0 : tid == 0) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm_u = tmem; + uint32_t const a0 = smem_u32(smem + S::kV16); + uint32_t const b0 = smem_u32(smem + S::kP); + if (kWarpMma ? elect_one() : true) + { +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + umma_f16(tm_u + BPC * kRows, umma_desc(a0 + ks * 4096, 2048, 128), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 1, 1), ks > 0); + } + umma_commit(mma_bar); + } + if (kWarpMma) + { + __syncwarp(); + } + } + M3FA_PROF_T(19, 0); + float const lw = warp_reduce_scatter8(pa, lane); + if ((lane & 3) == 0) + { + red_l[(8 * rg + rs8) * 4 + wq] = lw; + } + rg_bar(); + M3FA_PROF_T(20, 0); + if (wq == 0 && lane < 8) + { + int const r = 8 * rg + lane; + float4 const mq = *reinterpret_cast(red_m + r * 4); + float4 const lq = *reinterpret_cast(red_l + r * 4); + float const m = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + float const l = (lq.x + lq.y) + (lq.z + lq.w); + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[r] = scale * vgs; + row_lse[r] = lse; + } + } + __syncthreads(); // row scales and LSEs visible (PV was issued above) + M3FA_PROF(9); + if (warp < 8) + { + mbar_wait_cta(mma_bar, 1); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + } + else + { + // ---------------- Softmax: thread = block token; the (block, 16-row chunk) pairs are spread over the + // four groups of four compute warps (each group covers the 128 TMEM lanes = block tokens) + // ---------------- + if (warp < kComputeWarps && warp / 4 < kPairs) + { + int const g = warp / 4; + int const wq = warp % 4; + int const tk = wq * 32 + lane; + auto group_bar = [&]() + { + if (g == 0) + { + asm volatile("bar.sync 10, 128;\n" ::: "memory"); + } + else if (g == 1) + { + asm volatile("bar.sync 11, 128;\n" ::: "memory"); + } + else if (g == 2) + { + asm volatile("bar.sync 12, 128;\n" ::: "memory"); + } + else + { + asm volatile("bar.sync 13, 128;\n" ::: "memory"); + } + }; + // Row r(lane) this lane holds after the reduce-scatter: bits 4..1 of the lane. + int const rs_row + = ((lane >> 4) & 1) * 8 + ((lane >> 3) & 1) * 4 + ((lane >> 2) & 1) * 2 + ((lane >> 1) & 1); + mbar_wait_cta(mma_bar, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF(7); + float v[kPairsPerGroup][16]; +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + tmem_ld16(tmem + ((wq * 32) << 16) + h * kRows + c * 16, v[k]); + bool const valid + = cta_tok0 + h * kPage + tk < kv_max && cta_tok0 + h * kPage + tk < kv_len[c]; + float a[16]; +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[k][j] = valid ? v[k][j] * sl2e : kNegInf; + a[j] = v[k][j]; + } + float const mw = warp_reduce_scatter16(a, lane); + if ((lane & 1) == 0) + { + red_m[((h * QG + c) * 16 + rs_row) * 4 + wq] = mw; + } + } + } + group_bar(); + M3FA_PROF(8); + float pa[kPairsPerGroup][16]; +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + uint8_t* pdst = smem + S::kP + h * kRows * kPage * 2 + (tk >> 3) * 128 + (tk & 7) * 16; + // Row maxima: lane j (< 16) reduces row j's four warp partials, then broadcasts it. + float mrow = kNegInf; + if (lane < 16) + { + float4 const mq + = *reinterpret_cast(red_m + ((h * QG + c) * 16 + lane) * 4); + mrow = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + } +#pragma unroll + for (int j = 0; j < 16; ++j) + { + float const m = __shfl_sync(0xffffffffu, mrow, j); + pa[k][j] = m == kNegInf ? 0.0F : ex2(v[k][j] - m); + } + if (tid == 0) + { + M3FA_CLK(41); + } + uint32_t pw[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) + { + __half2 const hp = __floats2half2_rn(pa[k][2 * j], pa[k][2 * j + 1]); + pw[j] = *reinterpret_cast(&hp); + } + // P^T as the MN-major B operand: 8 query rows of one token are one 16-byte core-matrix row. + *reinterpret_cast(pdst + (2 * c) * 2048) = make_uint4(pw[0], pw[1], pw[2], pw[3]); + *reinterpret_cast(pdst + (2 * c + 1) * 2048) + = make_uint4(pw[4], pw[5], pw[6], pw[7]); + M3FA_PROF_T(18, 0); + } + } + if (tid == 0) + { + M3FA_CLK(42); + } + M3FA_FENCE_PROXY_ASYNC(); // P -> tensor core reads + if (tid == 0) + { + M3FA_CLK(43); + } + // PV starts as soon as P and the V tiles are in: thread 0 waits for the softmax warps and the + // producer warps (V); the row sums and statistics below overlap the MMAs. + if (warp == 0) + { + asm volatile("bar.sync 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + else + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + if (kWarpMma ? warp == 0 : tid == 0) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm_u = tmem; + if (kWarpMma ? elect_one() : true) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + uint32_t const a0 = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const b0 = smem_u32(smem + S::kP + h * kRows * kPage * 2); +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + umma_f16(tm_u + (BPC + h) * kRows, umma_desc(a0 + ks * 4096, 2048, 128), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 1, 1), ks > 0); + } + } + umma_commit(mma_bar); + } + if (kWarpMma) + { + __syncwarp(); + } + } + M3FA_PROF_T(19, 0); +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + float const lw = warp_reduce_scatter16(pa[k], lane); + if ((lane & 1) == 0) + { + red_l[((h * QG + c) * 16 + rs_row) * 4 + wq] = lw; + } + } + } + group_bar(); + M3FA_PROF_T(20, 0); + if (wq == 0 && lane < 16) + { +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + int const r = (h * QG + c) * 16 + lane; + float4 const mq = *reinterpret_cast(red_m + r * 4); + float4 const lq = *reinterpret_cast(red_l + r * 4); + float const m = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + float const l = (lq.x + lq.y) + (lq.z + lq.w); + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[h * kRows + c * 16 + lane] = scale * vgs; + row_lse[h * kRows + c * 16 + lane] = lse; + } + } + } + } + __syncthreads(); // row scales and LSEs visible (PV was issued above) + M3FA_PROF(9); + if (warp < kComputeWarps && warp / 4 < kPairs) + { + mbar_wait_cta(mma_bar, 1); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + } + } + M3FA_PROF(3); + } + else + { + asm volatile("cp.async.wait_all;\n" ::: "memory"); + if (OPJ && M3FA_C1_OPISSUE == 2 && tid == 32) + { + op_issue(p, smem, static_cast(blockIdx.x)); + } +#if M3FA_TRIGGER_AT == 2 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (kAttnTrigger || (FQKV && M3FA_MEGA_TRIGGER == 1)) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + for (int r = tid; r < BPC * kRows; r += kThreads) + { + row_lse[r] = kNegInf; + row_scale[r] = 0.0F; + } + __syncthreads(); + } + + // ---------------- Push every partial row to the CTA that merges its head (CTA r merges heads r*BPC..) -------- + if constexpr (S::kRecvAlias) + { + // Peers' partials overlay this CTA's tiles: every CTA's PV MMAs have completed. + asm volatile("barrier.cluster.arrive.aligned;\n" ::: "memory"); + } + if (!(kEarlyClw && !S::kRecvAlias)) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); // peers' mbarriers are initialized + } + M3FA_PROF(4); + if (OPJ && M3FA_C1_OPISSUE == 4 && tid == 32) + { + op_issue(p, smem, static_cast(blockIdx.x)); + } + uint32_t const recv_base = smem_u32(recv); + if (kRs2 && warp < 8) + { + // Split push: warp group rg pushes rows 8 rg .. 8 rg + 7 (row j goes to CTA j, which merges head j). + int const rg = warp / 4; + int const d = (warp % 4) * 32 + lane; + float v[8]; + if (any_valid) + { + tmem_ld8(tmem + ((warp % 4) * 32 << 16) + BPC * kRows + 8 * rg, v); + } + else + { +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] = 0.0F; + } + } +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] *= row_scale[8 * rg + j2]; + } +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + uint32_t const dst_cta = 8 * rg + j2; + uint32_t const row_addr = recv_base + static_cast((cr * kOPitch + d) * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr, dst_cta)), + "r"(__float_as_uint(v[j2])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + else if (!kRs2 && warp < kComputeWarps && warp / 4 < kPairs) + { + int const g = warp / 4; + int const d = (warp % 4) * 32 + lane; // output dim +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + int const chunk = cr * BPC + h; + float v[16]; + if (any_valid && kRegCore) + { + float const* obuf + = reinterpret_cast(smem + S::kK16 + h * S::kTileBytes + S::kObufOff); +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] = obuf[(c * 16 + j) * kOPitch + d]; + } + } + else if (any_valid) + { + tmem_ld16(tmem + ((warp % 4) * 32 << 16) + (BPC + h) * kRows + c * 16, v); + } + else + { +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] = 0.0F; + } + } +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] *= row_scale[h * kRows + c * 16 + j]; + } + // One scalar store per row: a vector store's lanes would address several destination CTAs at once, + // which measured several times slower. +#pragma unroll + for (int j = 0; j < 16; ++j) + { + uint32_t const dst_cta = j / BPC; + uint32_t const row_addr + = recv_base + static_cast((((chunk * QG + c) * BPC + j % BPC) * kOPitch + d) * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr, dst_cta)), + "r"(__float_as_uint(v[j])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + } + } + else if (tid >= kComputeThreads) + { + // Every partial row's LSE (BPC * kRows of them; more than the producer threads for the widest geometries). + for (int r = tid - kComputeThreads; r < BPC * kRows; r += kThreads - kComputeThreads) + { + int const h = r / kRows; + int const c = (r % kRows) / 16; + int const j = r % 16; + uint32_t const dst_cta = j / BPC; + uint32_t const row_addr + = recv_base + static_cast((((cr * BPC + h) * QG + c) * BPC + j % BPC) * kOPitch * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr + kHeadDim * 4, dst_cta)), + "r"(__float_as_uint(row_lse[r])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + M3FA_PROF(10); + if (OPJ && M3FA_C1_OPISSUE == 5 && tid == 32) + { + op_issue(p, smem, static_cast(blockIdx.x)); + } +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + if (wr_kind[k] == 1) + { + *reinterpret_cast(wr_ptr[k]) = static_cast(wr_val[k] & 0xFFFFu); + if (lane % 4 == 0) + { + *wr_sptr[k] = static_cast(wr_val[k] >> 16); + } + } + else if (wr_kind[k] == 2) + { + *reinterpret_cast(wr_ptr[k]) = wr_val[k]; + } + } + { + uint32_t ready = 0; +#ifdef M3FA_DBG_HANG + long long spins = 0; +#endif + while (!ready) + { +#ifdef M3FA_DBG_HANG + if (++spins == (1LL << 24)) + { + printf("M3FA hang: block %d thread %d merge recv\n", blockIdx.x, threadIdx.x); + } +#endif + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(mbar) + : "memory"); + } + } + M3FA_PROF(5); + + // ---------------- Merge heads [cr * BPC, cr * BPC + BPC) (split-K merge kernel's LSE weighting) ---------------- + constexpr int kMergeRows = QG * BPC; // (token, local head) pairs, 16 chunk weights each + // Wide geometries: chunk weights once per merged row (into the softmax scratch, free by now), then the weighted + // sums. Narrow ones recompute the weights per output element: no extra barrier on the critical path. + constexpr bool kRowWeights = kMergeRows >= 4; + float* wmerge = red_m; // [kMergeRows][kBlocks]; red_m holds BPC * kRows * 4 = kMergeRows * 64 floats + [[maybe_unused]] auto chunk_weights = [&](int mr, float(&lc)[kBlocks]) + { + float lmax = kNegInf; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + lc[c] = recv[(c * kMergeRows + mr) * kOPitch + kHeadDim]; + lmax = fmaxf(lmax, lc[c]); + } + lmax = lmax == kNegInf ? 0.0F : lmax; + float denom = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + lc[c] = ex2(lc[c] - lmax); // now the chunk weight; 0 for an empty chunk + denom += lc[c]; + } + denom = denom > 0.0F ? denom : 1.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + // -0.0 marks an empty chunk (skipped below, as the merge kernel skips zero weights). + lc[c] = lc[c] != 0.0F ? __fdividef(lc[c], denom) : -0.0F; + } + }; + if constexpr (kRowWeights) + { + for (int mr = tid; mr < kMergeRows; mr += kThreads) + { + float lc[kBlocks]; + chunk_weights(mr, lc); +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wmerge[mr * kBlocks + c] = lc[c]; + } + } + __syncthreads(); + } +#pragma unroll + for (int idx = tid; idx < kMergeRows * kHeadDim; idx += kThreads) + { + int const mr = idx >> 7; + int const d = idx & (kHeadDim - 1); + float wn[kBlocks]; + if constexpr (kRowWeights) + { +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wn[c] = wmerge[mr * kBlocks + c]; + } + } + else + { +#ifndef M3FA_MEGA_MERGE_LOCAL +#define M3FA_MEGA_MERGE_LOCAL \ + 0 // 1 (measured +0.6 us slower): every merging thread derives the 16 chunk weights itself (broadcast shared loads, + // independent exps): no shuffle chain; the same arithmetic as chunk_weights +#endif + if constexpr (M3FA_MEGA_MERGE_LOCAL != 0) + { + chunk_weights(mr, wn); + } + else + { + // Narrow geometries: the warp (32 dims of one merged row) derives the 16 chunk weights together, lane c + // holding chunk c's, and broadcasts them. + float lc = (lane < kBlocks) ? recv[(lane * kMergeRows + mr) * kOPitch + kHeadDim] : kNegInf; + float lmax = lc; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + lmax = fmaxf(lmax, __shfl_xor_sync(0xffffffffu, lmax, o)); + } + lmax = lmax == kNegInf ? 0.0F : lmax; + float const wc = lane < kBlocks ? ex2(lc - lmax) : 0.0F; + // The denominator in chunk order from 0.0, as chunk_weights (and the merge kernel) add it: a butterfly + // reduction rounds differently in the last bit. + float denom = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + denom += __shfl_sync(0xffffffffu, wc, c); + } + denom = denom > 0.0F ? denom : 1.0F; + float const wl = wc != 0.0F ? __fdividef(wc, denom) : -0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wn[c] = __shfl_sync(0xffffffffu, wl, c); + } + } + } + if (tid == 0) + { + M3FA_CLK(45); + } + float acc = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + float const prod = recv[(c * kMergeRows + mr) * kOPitch + d] * wn[c]; + acc = __float_as_uint(wn[c]) != 0x80000000u ? acc + prod : acc; + } + if (tid == 0) + { + M3FA_CLK(46); + } + int const tl = mr / BPC; + int const head = cr * BPC + mr % BPC; + __nv_bfloat16 const ob = __float2bfloat16_rn(acc); + int const out_row = tok0 + tfirst + tl; + p.out[static_cast(out_row) * p.out_stride + head * kHeadDim + d] = ob; + if constexpr (MXFP8) + { + float const x = __bfloat162float(ob); + float amax = fabsf(x); +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, mask)); + } + float const sfv = amax * rcp_approx_ftz(448.0F); + int const bits = __float_as_int(sfv); + int ex = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0 ? 1 : 0); + ex = min(ex, 254); + float scale = __int_as_float((254 - ex) << 23); + scale = amax != 0.0F ? scale : 0.0F; + __nv_fp8_e4m3 const q8v(x * scale); + int64_t const q8off = static_cast(out_row) * (p.num_sf_k_tiles * 128) + head * kHeadDim + d; + if constexpr (FQKV || OPJ) + { + // Megakernel hand-off to the o_proj units (Lamport): this call's parity buffer, the other one reset to + // 0xFF. 0xFF is never written: a NaN code is canonicalized to 0x7F; the UE8M0 code is at most 254. + uint8_t const code = q8v.__x == 0xFFu ? static_cast(0x7Fu) : q8v.__x; + int64_t const cur = (mega_calls & 1u) * kMaxTok * kOpK; + int64_t const nxt = ((mega_calls & 1u) ^ 1u) * kMaxTok * kOpK; + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.q8 + cur + q8off), "h"(static_cast(code)) + : "memory"); + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.q8 + nxt + q8off), "h"(static_cast(0xFFu)) + : "memory"); + } + else + { + p.q8[q8off] = q8v.__x; + } + if (lane == 0) + { + int const k_sf = head * (kHeadDim / 32) + (d >> 5); + int64_t const off = static_cast(out_row / 128) * p.num_sf_k_tiles * 512 + (k_sf / 4) * 512 + + (out_row % 32) * 16 + ((out_row % 128) / 32) * 4 + k_sf % 4; + if constexpr (FQKV || OPJ) + { + int64_t const cur = (mega_calls & 1u) * kOpSfBytes; + int64_t const nxt = ((mega_calls & 1u) ^ 1u) * kOpSfBytes; + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.sf + cur + off), "h"(static_cast(ex)) + : "memory"); + asm volatile("st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.sf + nxt + off), + "h"(static_cast(0xFFu)) + : "memory"); + } + else + { + p.sf[off] = static_cast(ex); + } + } + } + } + M3FA_PROF(60); + if constexpr (OPJ) + { + // ---------------- c1: o_proj unit blockIdx.x (v52 m3_mx_gemv bit for bit, mega_v7's phase) ---------------- + if (any_valid) + { + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + __syncthreads(); // the merge's reads done; the K tiles are free for the activation overlay + int const op_u = static_cast(blockIdx.x); + if (M3FA_C1_OPISSUE == 3 && tid == 32) + { + op_issue(p, smem, op_u); + } + uint32_t const tm = *tmem_slot; + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_stream(smem, tm); + } + __syncwarp(); + } + op_act(p, smem, op_u, tid, mega_calls); + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + M3FA_PROF_T(36, 0); + __syncthreads(); + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_mma(smem, tm); + M3FA_PROF_T(50, threadIdx.x); + } + __syncwarp(); + } + if (warp >= kComputeWarps) + { + op_epilogue(p, smem, tm, tid, op_u); + } + M3FA_PROF_T(37, kComputeThreads); + __syncthreads(); +#if M3FA_C1_TRIGGER == 3 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // after this CTA's y rows / pushes +#endif + if (warp == kWarps - 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" ::"r"(tm), "r"(S::kTmemCols) + : "memory"); + } + if (tid == 0) + { + p.flags[32 + blockIdx.x] = mega_calls + 1u; + } + M3FA_PROF(6); + return; + } + if (any_valid) + { + // Every TMEM read completed above (tcgen05.wait::ld); release the columns. + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + __syncthreads(); + if (FQKV && tid == 0) + { + p.flags[32 + blockIdx.x] = mega_calls + 1u; + } + M3FA_PROF(61); + if (warp == kWarps - 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" ::"r"(*tmem_slot), "r"(S::kTmemCols) + : "memory"); + } + M3FA_PROF(6); +#else + __trap(); +#endif // MINIMAX_M3_ATTN_DEVICE_ENABLED +} + +template +cudaError_t launch_impl(Params const& prm, int num_reqs, bool pdl, cudaStream_t stream, int* max_clusters) +{ + auto kernel = m3_fused_decode_attn_kernel; + constexpr int smem = std::conditional_t, SmemX>::kBytes; + constexpr int kCs = kBlocks / BPC; + static bool const configured = [&]() + { + bool ok = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem) == cudaSuccess; + ok = ok + && (kCs <= 8 + || cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) == cudaSuccess); + if (!ok) + { + cudaGetLastError(); + } + return ok; + }(); + if (!configured) + { + return cudaErrorInvalidConfiguration; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(FQKV ? kQkvGrid : kCs * (QLEN / QG) * num_reqs); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeClusterDimension; + attr[0].val.clusterDim.x = kCs; + attr[0].val.clusterDim.y = 1; + attr[0].val.clusterDim.z = 1; + attr[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[1].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + if (max_clusters != nullptr) + { + cfg.numAttrs = 1; + cudaError_t const err = cudaOccupancyMaxActiveClusters(max_clusters, kernel, &cfg); + if (err != cudaSuccess) + { + cudaGetLastError(); + } + return err; + } + cfg.numAttrs = 2; + return cudaLaunchKernelEx(&cfg, kernel, prm); +} + +// The one geometry the target launches: one request verifying 6 query tokens, 6 clusters of 16 CTAs (96 o_proj +// units). +cudaError_t dispatch(Params const& prm, int qlen, int qg, int bpc, int num_reqs, bool pdl, cudaStream_t stream) +{ + if (qlen == 6 && qg == 1 && bpc == 1) + { + return launch_impl<6, 1, 1, true, false, true>(prm, num_reqs, pdl, stream, nullptr); + } + return cudaErrorInvalidValue; +} + +static_assert(kMaxTok == minimax_m3_attn_c1::kMaxTokens && kOpN == minimax_m3_attn_c1::kOpN + && kOpK == minimax_m3_attn_c1::kOpK && kOpSfBytes == minimax_m3_attn_c1::kOpSfBytes + && kOpN / 128 * kOpSplits == minimax_m3_attn_c1::kCtas, + "header dimensions"); + +} // namespace + +namespace +{ +// TMA descriptor of a [n, k] E4M3 weight (128 x 128 boxes, 128-byte swizzle), cached per weight address. +bool weight_map(CUtensorMap& map, void const* w, int k, int n) +{ + static std::mutex mu; + static std::unordered_map cache; + std::lock_guard lock(mu); + auto it = cache.find(w); + if (it != cache.end()) + { + map = it->second; + return true; + } + static PFN_cuTensorMapEncodeTiled_v12000 encode = []() + { + void* fn = nullptr; + cudaDriverEntryPointQueryResult q; + cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled", &fn, 12000, cudaEnableDefault, &q); + return reinterpret_cast(fn); + }(); + if (encode == nullptr) + { + return false; + } + cuuint64_t dims[2] = {static_cast(k), static_cast(n)}; + cuuint64_t strides[1] = {static_cast(k)}; + cuuint32_t box[2] = {128, 128}; + cuuint32_t elem[2] = {1, 1}; + CUresult const r = encode(&map, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(w), dims, strides, box, elem, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_L2_256B, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + if (r != CUDA_SUCCESS) + { + return false; + } + cache.emplace(w, map); + return true; +} +} // namespace + +cudaError_t invokeMinimaxM3AttnC1(MinimaxM3AttnC1Params& params, void const* oWeight, int qlen, int qg, int bpc, + int numReqs, bool pdl, cudaStream_t stream) +{ + if (!weight_map(params.wmap_o, oWeight, kOpK, kOpN)) + { + return cudaErrorInvalidValue; + } + return dispatch(params, qlen, qg, bpc, numReqs, pdl, stream); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/attnC1.h b/cpp/tensorrt_llm/kernels/minimaxM3/attnC1.h new file mode 100644 index 000000000000..f13de72e14fb --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/attnC1.h @@ -0,0 +1,110 @@ +/* + * 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 +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace minimax_m3_attn_c1 +{ +constexpr int kCtas = 96; // attention CTAs == o_proj units (48 row tiles x 2 K splits) +constexpr int kMaxTokens = 8; // rows of each call-parity buffer of the MXFP8 hand-off +constexpr int kOpN = 6144; // o_proj output width +constexpr int kOpK = 2048; // o_proj input width (16 heads x 128) +constexpr int kOpSfBytes = 16 * 512; // one parity of the MXFP8 hand-off's swizzled scales +} // namespace minimax_m3_attn_c1 + +//! Kernel parameters of the attention + o_proj (c1) kernel (the device code's Params). +struct MinimaxM3AttnC1Params +{ + __nv_bfloat16 const* packed; + int64_t packed_stride; + uint8_t* kv_data; + int64_t d_page, d_plane; + uint8_t* kv_scale; + int64_t s_page, s_plane; + uint8_t* idx_k; + int64_t i_page, i_tok; + int64_t num_pages; + int const* out_cache_loc; + int const* position_ids; + int const* block_table; + int64_t bt_stride; + int bt_cols; + int const* seq_lens; + float const* kv_quant_scale; + float const* k_gs; + float const* v_gs; + __nv_bfloat16 const* q_w; + __nv_bfloat16 const* k_w; + __nv_bfloat16 const* ik_w; + float const* rope; + float eps; + float sm_scale; + __nv_bfloat16* out; + int64_t out_stride; + uint8_t* q8; + uint8_t* sf; + int num_sf_k_tiles; + unsigned long long* prof; // optional per-CTA clock64 probes (-DM3FA_PROFILE builds only) + // Fused QKV projection (FQKV kernels only): out[M, 2560] = mxfp8(x)[M, 6144] @ W^T. + __nv_bfloat16 const* x; // [M, 6144] layer input (the all-reduce output) + int ldx; + int num_tokens; // M + uint8_t const* wsf; // UE8M0 weight scales, 128x4-swizzled + float* part; // [2 parities][splits][8 tokens][2560] partial sums, sentinel at rest + int* counters; // [0]: call parity, [1]: CTAs done with this call (zero at rest) + int attn_clusters; // leading clusters that run the attention + uint8_t const* xq; // mega: the layer input's MXFP8 values [M, 6144] (row stride ldxq) ... + uint8_t const* xsf_act; // ... and its 128x4-swizzled UE8M0 scale atoms (atom k at byte 512 k) + int ldxq; + __nv_bfloat16* qkv_out; // mega: the packed QKV rows [2 parities][8][2560] the attention reads + uint8_t const* wsf_o; // mega: the o_proj weight's scale atoms + __nv_bfloat16* y_out; // mega: the o_proj output [M, 6144] + int* ar_ws; // mega, ar_mode: the next m3_mc_allreduce's multicast exchange (int32 words)... + int const* ar_flags; // ... its call flags (parity of a (token, column slice) slot = flags[t * 8 + slice] & 1) + int ar_rank; + int ar_world; + int ar_maxt; + int ar_mode; // 0: none; 2: push the o_proj rows into the exchange (as m3_mx_gemv's ar mode 2) + unsigned int* flags; // mega: [0, 20) tile counters (4 per call), [32 + CTA] the attention CTAs' targets + long long pf_a0, pf_b0, pf_a1, pf_b1; // mega: L2 prefetch regions issued in the attention window (bytes 0: none) + int pf_hint; // 1: evict_last policy on the prefetch + uint8_t const* wpre; // mega: optional pre-tiled QKV weight [N/128][K/128][16 KB swizzled smem image], or null + CUtensorMap wmap; // the [2560, 6144] E4M3 weight, 128 x 128 boxes, 128-byte swizzle + CUtensorMap wmap_o; // mega: the [6144, 2048] E4M3 o_proj weight, same boxes +}; + +//! Launch the c1 kernel on `stream`. The caller fills `params` except `wmap_o`, the o_proj weight's TMA +//! descriptor (cached per address). Only (qlen, qg, bpc) = (6, 1, 1) with one request is compiled; anything else, or +//! a failed TMA descriptor encode, returns cudaErrorInvalidValue. +cudaError_t invokeMinimaxM3AttnC1(MinimaxM3AttnC1Params& params, void const* oWeight, int qlen, int qg, int bpc, + int numReqs, bool pdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.cu b/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.cu new file mode 100644 index 000000000000..a4c8dd923f6f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.cu @@ -0,0 +1,3673 @@ +/* + * 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. + */ + +// MiniMax-M3 attention megakernel: QKV GEMV + fused decode attention + o_proj GEMV in one launch. + +#include "tensorrt_llm/kernels/minimaxM3/attnMega.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +// The kernel uses tcgen05/TMEM (including block-scaled MMAs, which family targets lack), cluster DSMEM and E2M1 +// conversions: its body is compiled for the arch-specific sm_100a / sm_103a targets only (see the minimax_m3_attn_src +// library in kernels/CMakeLists.txt). Any other target compiles a kernel that traps, so a launch on a build without +// those targets fails loudly instead of returning untouched outputs. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030) \ + && (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FEAT_SM100_ALL) || defined(__CUDA_ARCH_FEAT_SM103_ALL)) +#define MINIMAX_M3_ATTN_DEVICE_ENABLED 1 +#else +#define MINIMAX_M3_ATTN_DEVICE_ENABLED 0 +#endif + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kHeadDim = 128; +constexpr int kRotDim = 64; +constexpr int kPage = 128; +constexpr int kHq = 16; +constexpr int kColK = kHq; // packed head slots: [Q x16 | K | V | idxQ | idxK] +constexpr int kColV = kHq + 1; +constexpr int kColIdxK = kHq + 3; +constexpr int kBlocks = 16; // KV blocks per request == topk == cluster size +constexpr int kComputeWarps = 16; // QK / softmax / PV and the dequant (one 16-byte chunk per thread) +constexpr int kWarps = kComputeWarps + 4; // plus producer-only warps, so new K/V/index-K never queue behind a Q head +constexpr int kThreads = kWarps * 32; +constexpr int kComputeThreads = kComputeWarps * 32; +constexpr int kOPitch = 132; // floats per partial-O row +constexpr float kNegInf = -__builtin_huge_valf(); + +using Params = MinimaxM3AttnMegaParams; + +// Optional per-CTA clock64 probes for development (-DM3FA_PROFILE and a prof tensor): [CTA][64] timestamps. +// -DM3FA_PROFILE_GT records %globaltimer (ns, comparable across SMs and kernels) instead of clock64. +#ifdef M3FA_PROFILE_GT +#define M3FA_PROFILE +#define M3FA_NOW(v) asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(v)) +#else +#define M3FA_NOW(v) v = clock64() +#endif +#ifdef M3FA_PROFILE +#define M3FA_PROF_T(i, thread) \ + do \ + { \ + if (p.prof != nullptr && threadIdx.x == (thread)) \ + { \ + unsigned long long now_; \ + M3FA_NOW(now_); \ + p.prof[blockIdx.x * 64 + (i)] = now_; \ + } \ + } while (0) +#else +#define M3FA_PROF_T(i, thread) \ + do \ + { \ + } while (0) +#endif +#define M3FA_PROF(i) M3FA_PROF_T(i, 0) +// clock64 probes (development): slot i of this CTA, written by every calling thread. +#ifdef M3FA_PROFILE +#define M3FA_CLK(i) \ + do \ + { \ + if (p.prof != nullptr) \ + { \ + p.prof[blockIdx.x * 64 + (i)] = clock64(); \ + } \ + } while (0) +#else +#define M3FA_CLK(i) \ + do \ + { \ + } while (0) +#endif + +#ifndef M3FA_MEGA_TRIGGER +// Megakernel PDL release: 0 o_proj CTAs after their activation (attention CTAs at exit); 1 as 0 with the attention +// CTAs after their KV reads; 2 every CTA right after its griddepcontrol.wait; 3 o_proj CTAs after their epilogue. +#define M3FA_MEGA_TRIGGER 3 +#endif +#ifndef M3FA_TRIGGER_AT +#define M3FA_TRIGGER_AT 2 // where the PDL dependents are released: 0 entry, 1 after the wait, 2 after the KV reads +#endif + +// Generic-proxy shared-memory writes (fp16 tiles, Q, P) become visible to the tensor core (async proxy). +#define M3FA_FENCE_PROXY_ASYNC() asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory") + +__device__ __forceinline__ void cp_async16(void* dst, void const* src) +{ + uint32_t const d = static_cast(__cvta_generic_to_shared(dst)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(d), "l"(src) : "memory"); +} + +__device__ __forceinline__ float ex2(float x) +{ + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ float lg2(float x) +{ + float y; + asm("lg2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +// One packed E2M1 byte -> half2 (low nibble in the low half). +__device__ __forceinline__ uint32_t e2m1x2_to_h2(uint32_t byte) +{ + uint32_t out; + asm("{\n .reg .b8 src;\n mov.b32 {src, _, _, _}, %1;\n cvt.rn.f16x2.e2m1x2 %0, src;\n}\n" : "=r"(out) : "r"(byte)); + return out; +} + +__device__ __forceinline__ uint32_t hmul2_u(uint32_t a, uint32_t b) +{ + __half2 r = __hmul2(*reinterpret_cast<__half2*>(&a), *reinterpret_cast<__half2*>(&b)); + return *reinterpret_cast(&r); +} + +__device__ __forceinline__ uint32_t e4m3_to_h2_splat(uint32_t byte) +{ + __half_raw const h = __nv_cvt_fp8_to_halfraw(static_cast<__nv_fp8_storage_t>(byte), __NV_E4M3); + return static_cast(h.x) | (static_cast(h.x) << 16); +} + +__device__ __forceinline__ uint32_t e4m3x2_to_h2(uint32_t two_bytes) +{ + __half2_raw const h = __nv_cvt_fp8x2_to_halfraw2(static_cast<__nv_fp8x2_storage_t>(two_bytes), __NV_E4M3); + return static_cast(h.x) | (static_cast(h.y) << 16); +} + +// ---- Exact copies of the C++ producer's helpers (fusedQKNormRopeKernel.cu / quantization.cuh) ---- +template +__device__ __forceinline__ uint32_t cvt_sf_group_shfl_mask() +{ + constexpr uint32_t groupSize = static_cast(NUM_THREADS_PER_SF); + constexpr uint32_t groupMask = 0xFFFFFFFFU >> (32U - groupSize); + uint32_t laneId = 0; + asm("mov.u32 %0, %%laneid;" : "=r"(laneId)); + return groupMask << (laneId & ~(groupSize - 1U)); +} + +__device__ __forceinline__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) +{ + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(val) + : "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]), "f"(array[4]), "f"(array[5]), "f"(array[6]), + "f"(array[7])); + return val; +} + +__device__ __forceinline__ uint16_t quantizeMinimaxM3Fp4x4( + float const (&elements)[4], float secondLevelScale, uint8_t* scaleOut) +{ + __nv_bfloat16 rounded[4]; + float localMax = 0.0F; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + rounded[i] = __float2bfloat16_rn(elements[i]); + localMax = fmaxf(localMax, fabsf(__bfloat162float(rounded[i]))); + } + uint32_t const groupMask = cvt_sf_group_shfl_mask<4>(); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 1)); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 2)); + float const scaleValue = secondLevelScale * (localMax * rcp_approx_ftz(6.0F)); + __nv_fp8_e4m3 const encodedScale(scaleValue); + if (scaleOut != nullptr) + { + *scaleOut = encodedScale.__x; + } + float const narrowedScale = static_cast(encodedScale); + float const outputScale + = localMax != 0.0F ? rcp_approx_ftz(narrowedScale * rcp_approx_ftz(secondLevelScale)) : 0.0F; + float quantValues[8] = {}; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + quantValues[i] = __bfloat162float(rounded[i]) * outputScale; + } + return static_cast(fp32_vec_to_e2m1(quantValues)); +} + +// RMSNorm (gemma 1+w) + NeoX RoPE over the first 64 dims, exactly as the producer: lane owns 4 dims. +__device__ __forceinline__ void norm_rope( + float (&e)[4], float const (&w)[4], float eps, float const (&cs)[4], float const (&sn)[4], int lane) +{ + float ss = 0.0F; + ss += e[0] * e[0] + e[1] * e[1]; + ss += e[2] * e[2] + e[3] * e[3]; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + ss = ss + __shfl_xor_sync(0xffffffffu, ss, mask, 32); + } + float const rmsReciprocal = rsqrtf(ss / static_cast(kHeadDim) + eps); +#pragma unroll + for (int i = 0; i < 4; ++i) + { + e[i] *= rmsReciprocal * (1.0F + w[i]); + } + __syncwarp(); + constexpr int kPairOffset = (kRotDim / 2) / 4; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const dim = lane * 4 + i; + float paired = __shfl_xor_sync(0xffffffffu, e[i], kPairOffset); + if (dim < kRotDim) + { + bool const firstHalf = dim < kRotDim / 2; + if (firstHalf) + { + paired = -paired; + } + e[i] = e[i] * cs[i] + paired * sn[i]; + } + } + __syncwarp(); +} + +__device__ __forceinline__ void load_w4(float (&w)[4], __nv_bfloat16 const* src, int lane) +{ + uint2 const raw = *reinterpret_cast(src + lane * 4); + float2 const a = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw.x)); + float2 const b = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw.y)); + w[0] = a.x; + w[1] = a.y; + w[2] = b.x; + w[3] = b.y; +} + +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ uint32_t mapa_shared(uint32_t addr, uint32_t rank) +{ + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n" : "=r"(r) : "r"(addr), "r"(rank)); + return r; +} + +// QLEN query tokens per request; QG of them per cluster (QLEN / QG clusters per request); BPC KV blocks per CTA +// (cluster of kBlocks / BPC CTAs). +// Byte offset of element (r, k) in a UMMA no-swizzle K-major operand tile of 128 k-elements per row: 8x8 core +// matrices of 16-byte rows, 128 bytes apart along k (LBO) and 2048 bytes apart per 8 rows (SBO). The same formula +// with r = token, k = dim stores V so that it is the MN-major (dims contiguous) A operand of O^T = V^T P^T. +__device__ __forceinline__ int kmaj_off(int r, int k) +{ + return (r >> 3) * 2048 + (k >> 3) * 128 + (r & 7) * 16 + (k & 7) * 2; +} + +__device__ __forceinline__ uint64_t umma_desc(uint32_t addr, uint32_t lbo, uint32_t sbo) +{ + return static_cast((addr >> 4) & 0x3FFFu) | (static_cast((lbo >> 4) & 0x3FFFu) << 16) + | (static_cast((sbo >> 4) & 0x3FFFu) << 32) | (1ull << 46); +} + +// kind::f16 instruction descriptor: F32 accumulate, F16 operands, B K-major, M = 128. +__host__ __device__ constexpr uint32_t umma_idesc(int n, int a_mn_major, int b_mn_major) +{ + return (1u << 4) | (static_cast(a_mn_major) << 15) | (static_cast(b_mn_major) << 16) + | (static_cast(n >> 3) << 17) | (static_cast(128 >> 4) << 24); +} + +// 8-row version for the split softmax: the same lane-bit order (16, 8, 4, then 2 and 1 in full) as +// warp_reduce_scatter16, so every row is combined in the same association (fp add is commutative, not associative). +// After it lanes with bits 4..2 = (b4, b3, b2) hold row 4*b4 + 2*b3 + b2. +template +__device__ __forceinline__ float warp_reduce_scatter8(float (&a)[8], int lane) +{ +#pragma unroll + for (int step = 0; step < 3; ++step) + { + int const mask = 16 >> step; + int const half = 4 >> step; + bool const upper = (lane & mask) != 0; +#pragma unroll + for (int i = 0; i < half; ++i) + { + float const send = upper ? a[i] : a[i + half]; + float const keep = upper ? a[i + half] : a[i]; + float const recv = __shfl_xor_sync(0xffffffffu, send, mask); + a[i] = kMax ? fmaxf(keep, recv) : keep + recv; + } + } + float x = a[0]; + float const o2 = __shfl_xor_sync(0xffffffffu, x, 2); + x = kMax ? fmaxf(x, o2) : x + o2; + float const o1 = __shfl_xor_sync(0xffffffffu, x, 1); + return kMax ? fmaxf(x, o1) : x + o1; +} + +// Reduce 16 per-lane values over the warp, one value per row: after the reduce-scatter lane l holds row +// 8*b4 + 4*b3 + 2*b2 + b1 (b_k = bit k of l) reduced over all 32 lanes. 16 shuffles instead of 80. +template +__device__ __forceinline__ float warp_reduce_scatter16(float (&a)[16], int lane) +{ +#pragma unroll + for (int step = 0; step < 4; ++step) + { + int const mask = 16 >> step; + int const half = 8 >> step; + bool const upper = (lane & mask) != 0; +#pragma unroll + for (int i = 0; i < half; ++i) + { + float const send = upper ? a[i] : a[i + half]; + float const keep = upper ? a[i + half] : a[i]; + float const recv = __shfl_xor_sync(0xffffffffu, send, mask); + a[i] = kMax ? fmaxf(keep, recv) : keep + recv; + } + } + float const other = __shfl_xor_sync(0xffffffffu, a[0], 1); + return kMax ? fmaxf(a[0], other) : a[0] + other; +} + +__device__ __forceinline__ void umma_f16( + uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %4, 0;\n tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(accumulate) + : "memory"); +} + +__device__ __forceinline__ void umma_commit(uint32_t bar) +{ + asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];\n" ::"r"(bar) : "memory"); +} + +__device__ __forceinline__ void mbar_wait_cta(uint32_t bar, uint32_t parity) +{ + uint32_t ready = 0; +#ifdef M3FA_DBG_HANG + long long spins = 0; +#endif + while (!ready) + { +#ifdef M3FA_DBG_HANG + if (++spins == (1LL << 24)) + { + printf("M3FA hang: block %d thread %d mbar 0x%x parity %u\n", blockIdx.x, threadIdx.x, bar, parity); + } +#endif + 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"(ready) + : "r"(bar), "r"(parity) + : "memory"); + } +} + +// Spin without suspending (v52's acc_done wait): a suspended try_wait wakes late on a tcgen05.commit arrival. +#ifndef M3FA_EPI_SPIN +#define M3FA_EPI_SPIN 1 +#endif +__device__ __forceinline__ void mbar_wait_epi(uint32_t bar, uint32_t parity) +{ + if (!M3FA_EPI_SPIN) + { + mbar_wait_cta(bar, parity); + return; + } + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(bar), "r"(parity) + : "memory"); + } +} + +__device__ __forceinline__ void tmem_ld8(uint32_t addr, float (&v)[8]) +{ + uint32_t r[8]; + asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n" + : "=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"(addr)); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < 8; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +__device__ __forceinline__ void tmem_ld16(uint32_t addr, 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];\n" + : "=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"(addr)); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < 16; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +// Warp-level fp16 MMA (m16n8k16, fp32 accumulate) and its shared-memory fragment loads. +__device__ __forceinline__ void ldsm_x4(uint32_t addr, uint32_t (&r)[4]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x2(uint32_t addr, uint32_t (&r)[2]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" : "=r"(r[0]), "=r"(r[1]) : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x2_trans(uint32_t addr, uint32_t (&r)[2]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0, %1}, [%2];\n" + : "=r"(r[0]), "=r"(r[1]) + : "r"(addr)); +} + +__device__ __forceinline__ void mma_f16_16816(float (&d)[4], uint32_t const (&a)[4], uint32_t const (&b)[2]) +{ + 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"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +// Register attention core (warp MMAs, no TMEM) for up to two tokens per cluster; the tcgen05 core above that. +// Off by default: on GB300 it measured ~0.25 us slower than the tcgen05 core at M=1 (GEMV + attention chain). +#ifndef M3FA_REG_CORE +#define M3FA_REG_CORE 0 +#endif +template +struct RegCore +{ + static constexpr bool kOn = M3FA_REG_CORE != 0 && QG <= 2; +}; + +template +struct Smem +{ + static constexpr int kRows = 16 * QG; // query rows (token, head) of this cluster = UMMA N + static constexpr int kTok = BPC * kPage; // block tokens per CTA + static constexpr int kTileBytes = kPage * kHeadDim * 2; // one fp16 K or V block tile in UMMA layout + static constexpr int kK16 = 0; + static constexpr int kV16 = kK16 + BPC * kTileBytes; + static constexpr int kQ = kV16 + BPC * kTileBytes; // Q rows, UMMA K-major [row][dim] + // P^T per block, UMMA MN-major [token][row], over the block's K tile: P is written after the QK MMAs (which + // read every K tile) have completed. + static constexpr int kP = kK16; + // Packed V tiles as copied from the cache: [block][128 rows x 64 data bytes | 1 KB of 4x4-quad scales]. + static constexpr int kRawV = kQ + kRows * kHeadDim * 2; + static constexpr int kRawVBlock = kPage * 64 + kPage * 8; + static constexpr int kRowScale = kRawV + BPC * kRawVBlock; + static constexpr int kRowLse = kRowScale + BPC * kRows * 4; + // Per-warp softmax partials [block][row][4 warps]: max, then sum. + static constexpr int kRedM = (kRowLse + BPC * kRows * 4 + 15) / 16 * 16; + static constexpr int kRedL = kRedM + BPC * kRows * 16; + // Partials pushed by every CTA of the cluster: [chunk][token][local head][kOPitch] floats, lse at column 128. + // When they do not fit beside the tiles they overlay the K/V tiles (and P, Q): peers then push only once every + // CTA of the cluster has completed its PV MMAs (a cluster barrier right before the push). + // Register core: per-warp row maxima and sums [block][row][warp of the block], two arrays. + static constexpr int kRedX = kRedL + BPC * kRows * 16; + static constexpr int kRedXBytes = RegCore::kOn ? 2 * BPC * kRows * (16 / BPC) * 4 : 0; + static constexpr int kStatsEnd = kRedX + kRedXBytes; + // Register core: P (fp16, [row][token] core matrices) at the start of its block's K tile, the scaled-free O + // partial ([row][kOPitch] fp32) behind it. + static constexpr int kObufOff = 8192; + static_assert(!RegCore::kOn || (kRows * kPage * 2 <= kObufOff && kObufOff + kRows * kOPitch * 4 <= kTileBytes), + "register-core P and O staging fit in a K tile"); + static constexpr int kRecvBytes = kBlocks * QG * BPC * kOPitch * 4; + static constexpr bool kRecvAlias = kStatsEnd + kRecvBytes + 32 > 227 * 1024; + static constexpr int kRecv = kRecvAlias ? 0 : kStatsEnd; + static constexpr int kMbar = kRecvAlias ? kStatsEnd : kStatsEnd + kRecvBytes; // merge, MMA mbarriers, TMEM addr + static constexpr int kBytes = kMbar + 32; + static constexpr int kTmemCols = BPC * 2 * kRows <= 32 ? 32 + : BPC * 2 * kRows <= 64 ? 64 + : BPC * 2 * kRows <= 128 ? 128 + : BPC * 2 * kRows <= 256 ? 256 + : 512; + static_assert(BPC * 2 * kRows <= 512, "TMEM columns"); + static_assert(!kRecvAlias || kRecvBytes <= kRawV, "aliased partials must fit in the tiles, Q and P"); + static_assert(kRows * kPage * 2 <= kTileBytes, "P^T fits in the K tiles"); + static_assert(kBytes <= 227 * 1024, "shared memory budget"); + static_assert(QLEN % QG == 0 && (BPC == 1 || BPC == 2), "unsupported geometry"); +}; + +// The packed NVFP4 bytes of one 16-byte chunk (32 dims of one block row) and its two block scales. +struct RawChunk +{ + uint4 data; + uint32_t scales; // byte 0: dims [0, 16), byte 1: dims [16, 32) of the chunk +}; + +// Chunk q of a block tile: lanes 8j..8j+7 of a warp take 8 consecutive rows of one 32-dim chunk, so every 16-byte +// fp16 store phase writes 8 distinct core-matrix rows (conflict-free in the UMMA layout). +__device__ __forceinline__ int chunk_row(int q) +{ + return (q >> 5) * 8 + (q & 7); +} + +__device__ __forceinline__ int chunk_col(int q) +{ + return (q >> 3) & 3; +} + +__device__ __forceinline__ RawChunk load_chunk( + uint8_t const* data_page, uint8_t const* scale_page, int row, int cc, bool is_v) +{ + RawChunk c; + c.data = *reinterpret_cast(data_page + row * 64 + cc * 16); + if (is_v) + { + // V scales: 4x4 token-quad order [row / 4][scale column][row % 4]. + uint8_t const* sc = scale_page + (row >> 2) * 32 + (row & 3); + c.scales = static_cast(sc[(cc * 2) * 4]) | (static_cast(sc[(cc * 2 + 1) * 4]) << 8); + } + else + { + c.scales = *reinterpret_cast(scale_page + row * 8 + cc * 2); + } + return c; +} + +__device__ __forceinline__ void store_chunk_fp16(uint8_t* tile, int row, int cc, RawChunk const& c, bool valid) +{ + uint32_t const s0 = e4m3_to_h2_splat(c.scales & 0xFFu); + uint32_t const s1 = e4m3_to_h2_splat((c.scales >> 8) & 0xFFu); + uint32_t const wd[4] = {c.data.x, c.data.y, c.data.z, c.data.w}; + uint32_t out[16]; +#pragma unroll + for (int i = 0; i < 16; ++i) + { + uint32_t const b = (wd[i >> 2] >> (8 * (i & 3))) & 0xFFu; + out[i] = valid ? hmul2_u(e2m1x2_to_h2(b), i < 8 ? s0 : s1) : 0u; + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + *reinterpret_cast(tile + kmaj_off(row, cc * 32 + j * 8)) + = make_uint4(out[4 * j], out[4 * j + 1], out[4 * j + 2], out[4 * j + 3]); + } +} + +// ---------------- Fused QKV projection: block-scaled MXFP8 on tcgen05 (N=2560, K=6144) ---------------- +// out[M, 2560] = mxfp8(x)[M, 6144] @ W^T, the computation of the CUTLASS MXFP8 GEMM: W is the A operand (128 output +// rows per MMA), the quantized activation the B operand (8 token columns), FP32 accumulation in TMEM. A unit is one +// (128-row tile, K split) pair; the 20 x kQkvSplits units run on the last CTAs of the grid. Before the PDL wait a +// CTA streams its unit's weight tiles and scale atoms through a shared-memory ring into TMEM; after it, it quantizes +// its x slice (mxfp8_quantize math), runs the MMAs and stores its FP32 partial rows. +// +// Hand-off without a flag: the partial buffers are double buffered by call parity and hold a sentinel (an FP32 NaN +// the MMA never produces) at rest, so a consumer polls the partials themselves. The producers of call n reset the +// other parity's buffers after their PDL wait; call n - 1, the last reader of those, is complete by then. +constexpr int kQkvN = 2560; +constexpr int kQkvK = 6144; +constexpr int kQkvKTiles = kQkvK / 128; +constexpr int kQkvRowTiles = kQkvN / 128; +#ifndef M3FA_QKV_SPLITS +#define M3FA_QKV_SPLITS 4 +#endif +#ifndef M3FA_QKV_RING +#define M3FA_QKV_RING kQkvMaxTiles // every tile of the unit in flight at once +#endif +#ifndef M3FA_QKV_NACC +#define M3FA_QKV_NACC 1 +#endif +#ifndef M3FA_QKV_SCALED_MMA +#define M3FA_QKV_SCALED_MMA 1 // 0: unscaled per-block MMAs with the scales applied in the epilogue (see below) +#endif +#ifndef M3FA_QKV_PIPE +// 1: the weight TMA is issued first at entry; the activation has its own region, copied right after the PDL wait; +// each weight tile's MMAs issue as soon as it lands (v52's order). 0: tiles into TMEM first, activation over the ring. +#define M3FA_QKV_PIPE 1 +#endif +#ifndef M3FA_QKV_A_SMEM +#define M3FA_QKV_A_SMEM 1 // scaled MMAs only; 1: they read the weight tiles from shared memory (no copy into TMEM) +#endif +// Default: block-scaled kind::mxf8f6f4 MMAs reading A (the weight tiles) in place from the ring, issued per tile as +// it lands (M3FA_QKV_PIPE), one accumulator; M3FA_QKV_A_SMEM=0 copies each tile into TMEM first. Alternative +// (M3FA_QKV_SCALED_MMA=0): one unscaled kind::f8f6f4 MMA per 32-wide K block into its own accumulator, A read in place +// from shared memory, the epilogue applying both UE8M0 scales. On GB300 an M128 x N8 x K32 MMA occupies the tensor pipe +// ~105 cycles block-scaled and ~60 unscaled, independent of N up to 64; the unscaled form's epilogue (320 accumulator +// columns) costs more than it saves. +constexpr bool kQkvScaled = M3FA_QKV_SCALED_MMA != 0; +constexpr bool kQkvASmem = !kQkvScaled || M3FA_QKV_A_SMEM != 0; +constexpr int kQkvSplits = M3FA_QKV_SPLITS; +constexpr int kQkvUnits = kQkvRowTiles * kQkvSplits; +constexpr int kQkvMaxTiles = (kQkvKTiles + kQkvSplits - 1) / kQkvSplits; +constexpr int kQkvRing = M3FA_QKV_RING; +constexpr int kQkvNacc = M3FA_QKV_NACC; +constexpr int kQkvGrid = 112; // 7 clusters of 16 (or 14 of 8): one wave with the attention +constexpr int kMaxTok = 8; +constexpr int kQkvWTile = 128 * 128; // E4M3 bytes of one 128 x 128 weight tile +constexpr int kQkvSfAtom = 512; // one 128-row x 4-block UE8M0 scale atom +constexpr uint32_t kQkvSentinel = 0xFFFFFFFFu; +constexpr int64_t kQkvPartStride = static_cast(kQkvSplits) * kMaxTok * kQkvN; // one parity +// TMEM columns of the one 512-column allocation of a CTA: attention [0, 64), then the GEMM's accumulators, scale +// atoms (A per tile, B per tile) and the A tiles (32 columns of 128 x 128 E4M3 each). +constexpr int kTmemAll = 512; +// Megakernel: the attention CTAs run no projection unit, so the GEMM starts at column 0 of a projection CTA. +constexpr int kTAcc = 0; +constexpr int kTSfa = kTAcc + 16; +constexpr int kTSfb = kTSfa + 4 * kQkvMaxTiles; +constexpr int kTA = (kTSfb + 4 * kQkvMaxTiles + 31) / 32 * 32; +static_assert(!kQkvScaled || kTA + 32 * kQkvMaxTiles <= kTmemAll, "TMEM columns"); +static_assert(kQkvScaled || kTAcc + 32 * kQkvMaxTiles <= kTmemAll, "TMEM columns (one accumulator per K block)"); +static_assert(!kQkvASmem || kQkvRing == kQkvMaxTiles, "shared-memory A needs every tile resident"); +static_assert(kQkvUnits <= kQkvGrid, "at most one unit per CTA"); +static_assert(kQkvNacc == 1 || kQkvNacc == 2 || kQkvNacc == 4, "accumulator count"); +static_assert(!M3FA_QKV_PIPE || kQkvScaled, "the pipelined projection uses the block-scaled MMAs"); +constexpr int kQkvOwnRows = 128 / kQkvSplits; // rows of a tile whose split-K sum a CTA owns +// Megakernel o_proj (v52 m3_mx_gemv for [6144, 2048]: split 2, contiguous k-tiles, one accumulator): unit +// o = CTA - 16 (the last 96 CTAs), row tile o / 2, split o % 2. +constexpr int kOpN = 6144; +constexpr int kOpK = 2048; +constexpr int kOpKTiles = kOpK / 128; +constexpr int kOpSplits = 2; +constexpr int kOpTiles = kOpKTiles / kOpSplits; +constexpr int kOpOwnRows = 128 / kOpSplits; +constexpr int kOpFirstCta = kQkvGrid - kOpN / 128 * kOpSplits; // 16 +constexpr int kOpSfBytes = kOpKTiles * 512; // one parity of the MXFP8 hand-off's scales +#ifndef M3FA_OP_L2PF +#define M3FA_OP_L2PF \ + 0 // 1: QKV CTAs prefetch their o_proj weight slice into L2 before the PDL wait (competes + // with the QKV weight stream for HBM behind the MoE) +#endif +// o_proj TMEM columns (after the QKV unit's MMAs completed): accumulator, A / B scale atoms, A (8 tiles x 32) +constexpr int kOAcc = 0; +constexpr int kOSfa = 16; +constexpr int kOSfb = kOSfa + 4 * kOpTiles; +constexpr int kOA = 128; +static_assert(kOA + 32 * kOpTiles <= kTmemAll && kOSfb + 4 * kOpTiles <= kOA, "o_proj TMEM columns"); +constexpr int kQkvMbStride = kMaxTok + 4; // mailbox row pitch (fp32 words) + +__host__ __device__ constexpr int qkv_tile0(int s) +{ + return s * kQkvKTiles / kQkvSplits; +} + +__device__ __forceinline__ int qkv_unit() +{ + int const u = kQkvGrid - 1 - static_cast(blockIdx.x); + return u < kQkvUnits ? u : -1; +} + +// Attention layout plus, for the fused projection, the weight ring in front (the attention tiles reuse that space +// once the ring has drained into TMEM); the activation tiles, the pre-wait V copy and the barriers sit behind. +template +struct SmemX : Smem +{ + using B = Smem; + static constexpr int kRingW = 0; // [ring][128 rows][128 B], 128-byte swizzle + static constexpr int kRingSf = kRingW + kQkvRing * kQkvWTile; // [ring][512 B] + static constexpr int kRingEnd = kRingSf + kQkvRing * kQkvSfAtom; + // The activation tiles overlay the ring: written after the PDL wait, once every weight tile went on into TMEM. + static constexpr int kXt = kRingW; // [tile][8][128 B] + static constexpr int kXsf = kXt + kQkvMaxTiles * 1024; // [tile][512 B] + static constexpr int kXsc = (kRingEnd + 127) / 128 * 128; + // Peers push into recv at any time (an empty CTA right away), so it stays clear of the ring. + static constexpr int kRecv = FQKV ? kXsc : B::kRecv; + static constexpr int kRawV = FQKV ? (kRecv + kBlocks * QG * BPC * kOPitch * 4 + 127) / 128 * 128 : B::kRawV; + // Pipelined projection: the activation tiles and scale atoms in their own region (a projection CTA runs no + // attention, so it overlays the attention's partial-O and V regions), 1024-byte aligned for the swizzle. + static constexpr int kXA = FQKV ? (kRecv + 1023) / 1024 * 1024 : 0; + static constexpr int kXASf = kXA + kQkvMaxTiles * 1024; + static constexpr int kXAEnd = kXASf + kQkvMaxTiles * kQkvSfAtom; + static constexpr int kMbarA = FQKV ? kRawV + BPC * B::kRawVBlock : B::kMbar; + static constexpr int kMbar = FQKV && M3FA_QKV_PIPE && kXAEnd > kMbarA ? (kXAEnd + 127) / 128 * 128 : kMbarA; + static constexpr int kGBar = kMbar + 32; // full[ring], empty[ring], a_ready, acc_done, tile_done[tiles] + static constexpr int kRedBar = kGBar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); // split-K partials in + static constexpr int kOBar = kRedBar + 8; // o_proj: full[8], a_ready, acc_done, reduce + static constexpr int kMailbox = (kOBar + 8 * (kOpTiles + 3) + 15) / 16 * 16; // [split][own row][pitch] fp32 + // o_proj weight ring = ring tiles [0, 8); its activation tiles, scale atoms and mailbox behind it. + static constexpr int kOXt = kRingW + kOpTiles * kQkvWTile; + static constexpr int kOXsf = kOXt + kOpTiles * 1024; + static constexpr int kOMailbox = kOXsf + kOpTiles * kQkvSfAtom; + static_assert(!FQKV || kOMailbox + kOpSplits * kOpOwnRows * kQkvMbStride * 4 <= kRingW + kQkvRing * kQkvWTile, + "o_proj activation and mailbox fit behind its ring"); + static constexpr int kBytes = FQKV ? kMailbox + kQkvSplits * kQkvOwnRows * kQkvMbStride * 4 : B::kBytes; + static constexpr int kTmemCols = FQKV ? kTmemAll : B::kTmemCols; + static constexpr bool kRecvAlias = !FQKV && B::kRecvAlias; + // The epilogue reads the weight scale atoms while the attention already writes its tiles. + static_assert(!FQKV || kQkvScaled || B::kRecv <= kRingSf, "attention regions overlap the scale atoms"); + // The other attention regions overlay the ring: a unit's CTA writes them only once the ring has drained. + static_assert(kBytes <= 227 * 1024, "shared memory budget"); +}; + +__device__ __forceinline__ void mbar_init(uint32_t bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(bar), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbar_expect_tx(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 bulk_g2s(void* dst, void const* src, uint32_t bytes, uint32_t bar) +{ + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], " + "%4;\n" ::"r"(smem_u32(dst)), + "l"(src), "r"(bytes), "r"(bar), "l"(0x12F0000000000000ull) // EVICT_FIRST: weights are read once per layer + : "memory"); +} + +__device__ __forceinline__ void tma_load_2d(uint32_t dst, void const* map, int c0, int c1, uint32_t bar) +{ + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;\n" ::"r"(dst), + "l"(map), "r"(c0), "r"(c1), "r"(bar), "l"(0x12F0000000000000ull) + : "memory"); +} + +// K-major operand in the 128-byte swizzle: rows of 128 bytes, 8-row groups 1024 bytes apart. +__device__ __forceinline__ uint64_t umma_desc_sw128(uint32_t addr) +{ + return umma_desc(addr, 16, 1024) | (2ull << 61); +} + +__device__ __forceinline__ void tmem_cp_128x256b(uint32_t taddr, uint64_t desc) +{ + asm volatile("tcgen05.cp.cta_group::1.128x256b [%0], %1;\n" ::"r"(taddr), "l"(desc)); +} + +__device__ __forceinline__ void tmem_cp_sf(uint32_t taddr, uint32_t saddr) +{ + asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;\n" ::"r"(taddr), "l"(umma_desc(saddr, 16, 128))); +} + +// Block-scaled MXFP8 instruction descriptor: E4M3 x E4M3, UE8M0 scales, M = 128, N = 8, scale byte kb. +__host__ __device__ constexpr uint32_t qkv_idesc(int kb) +{ + return (static_cast(kb) << 4) | ((8u >> 3) << 17) | (1u << 23) | ((128u >> 7) << 27) + | (static_cast(kb) << 29); +} + +__device__ __forceinline__ void umma_mx( + uint32_t tmem_d, uint32_t tmem_a, uint64_t bdesc, uint32_t idesc, uint32_t sfa, uint32_t sfb, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], [%1], %2, %3, [%4], [%5], p;\n}\n" :: + "r"(tmem_d), + "r"(tmem_a), "l"(bdesc), "r"(idesc), "r"(sfa), "r"(sfb), "r"(accumulate)); +} + +__device__ __forceinline__ float pow2_from_ue8m0(uint32_t e) +{ + return __uint_as_float(e == 0 ? 0x00400000u : (e << 23)); +} + +// Unscaled E4M3 x E4M3 MMA, FP32 accumulator: M = 128, N = 8, K = 32, a fresh accumulator. +constexpr uint32_t kQkvIdescF8 = (1u << 4) | ((8u >> 3) << 17) | ((128u >> 4) << 24); + +__device__ __forceinline__ void umma_f8_ss(uint32_t tmem_d, uint64_t adesc, uint64_t bdesc) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, 0, 0;\n tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(kQkvIdescF8)); +} + +__device__ __forceinline__ void tmem_ld16_nowait(uint32_t addr, float* v) +{ + 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];\n" + : "=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"(addr)); +#pragma unroll + for (int i = 0; i < 16; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +__device__ __forceinline__ void umma_mx_ss( + uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, uint32_t sfa, uint32_t sfb, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], %1, %2, %3, [%4], [%5], p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(sfa), "r"(sfb), "r"(accumulate)); +} + +__device__ __forceinline__ bool elect_one() +{ + uint32_t pred; + asm volatile("{\n .reg .pred P;\n elect.sync _|P, 0xffffffff;\n selp.u32 %0, 1, 0, P;\n}\n" : "=r"(pred)); + return pred != 0; +} + +__device__ __forceinline__ int qkv_parity(Params const& p) +{ + int v; + asm volatile("ld.relaxed.gpu.global.s32 %0, [%1];\n" : "=r"(v) : "l"(p.counters) : "memory"); + return v & 1; +} + +// Barriers and TMEM (warp 1), then the first ring stages (warp 0). Warps 0 and 1 of every CTA call it. +template +__device__ __forceinline__ void qkv_setup( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int lane, int u) +{ + constexpr bool kPipe = M3FA_QKV_PIPE != 0; + if (kPipe && warp == 0 && lane == 0 && u >= 0) + { + // First thing at entry: this unit's whole weight slice. The issuing thread initializes the ring's full + // barriers itself; every waiter on them passes bar.sync 8 below, which this thread reaches after the init. + asm volatile("prefetch.tensormap [%0];\n" ::"l"(&p.wmap) : "memory"); + int const rt = u / kQkvSplits; + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + for (int i = 0; i < kQkvRing; ++i) + { + mbar_init(gbar + 8 * i, 1); + } + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + for (int i = 0; i < kQkvRing && i < nt; ++i) + { + uint32_t const full = gbar + 8 * i; + mbar_expect_tx(full, kQkvWTile + kQkvSfAtom); + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap, (kt0 + i) * 128, rt * 128, full); + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, full); + } + M3FA_PROF_T(34, 0); + } + if (warp == 1) + { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n" ::"r"(smem_u32(tmem_slot)), + "r"(kTmemAll) + : "memory"); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n" ::: "memory"); + M3FA_PROF_T(35, 32); + if (lane == 0) + { + for (int i = (kPipe && u >= 0) ? kQkvRing : 0; i < 2 * kQkvRing + 2 + kQkvMaxTiles; ++i) + { + // PIPE 2: tile_done[0] (unused by the scaled MMAs) is act_ready, one arrival per thread outside warp 1 + mbar_init(gbar + 8 * i, M3FA_QKV_PIPE == 2 && i == 2 * kQkvRing + 2 ? kThreads - 32 : 1); + } + // The split-K reduce barrier: this CTA's own arrival carries the peers' partial bytes. + uint32_t const red = gbar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); + mbar_init(red, 1); + mbar_expect_tx(red, (kQkvSplits - 1) * kQkvOwnRows * kMaxTok * 4); + uint32_t const ob = smem_u32(smem + SX::kOBar); + for (int i = 0; i < kOpTiles + 3; ++i) + { + mbar_init(ob + 8 * i, 1); + } + mbar_expect_tx(ob + 8 * (kOpTiles + 2), (kOpSplits - 1) * kOpOwnRows * kMaxTok * 4); + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + asm volatile("bar.sync 8, 64;\n" ::: "memory"); + if (!kPipe && warp == 0 && lane == 0 && u >= 0) + { + int const rt = u / kQkvSplits; + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + for (int i = 0; i < kQkvRing && i < nt; ++i) + { + uint32_t const full = gbar + 8 * i; + mbar_expect_tx(full, kQkvWTile + kQkvSfAtom); + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap, (kt0 + i) * 128, rt * 128, full); + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, full); + } + M3FA_PROF_T(34, 0); + } +} + +// The rest of the ring (warp 0, lane 0) and the copies into TMEM (warp 1, lane 0); before the PDL wait. +template +__device__ __forceinline__ void qkv_stream_lane( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int u) +{ + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + uint32_t const full = gbar; + uint32_t const a_ready = gbar + 16 * kQkvRing; + if (warp == 0 || !kQkvScaled) + { + return; // every tile was issued in qkv_setup; the unscaled MMAs read them in place + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + // Fully unrolled with compile-time offsets: tcgen05 operands stay in uniform registers (a per-instruction + // conversion of computed operands costs ~150 cycles per MMA or copy). + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint32_t const sf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + mbar_wait_cta(full + 8 * i, 0); + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if constexpr (!kQkvASmem) + { +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kTA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + } + tmem_cp_sf(tm + kTSfa + i * 4, sf0 + i * kQkvSfAtom); + } + } + M3FA_PROF_T(31, 32); + umma_commit(a_ready); +} + +__device__ __forceinline__ bool mbar_test(uint32_t bar, uint32_t parity) +{ + uint32_t ready; + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(bar), "r"(parity) + : "memory"); + return ready != 0; +} + +// PIPE 2, warp 1's elected thread from before the PDL wait to the last MMA. Tiles that land before the activation +// is ready go on into TMEM (A and scale atom), off the critical path; once the activation is in (act_ready), the +// MMAs issue in tile order: the copied tiles with A in TMEM, every later tile as it lands with A read in place from +// the ring. The MMA sequence (and so the accumulation order) is the same either way. +template +__device__ __forceinline__ void qkv_hybrid(Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int u) +{ + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + uint32_t const act_ready = gbar + 8 * (2 * kQkvRing + 2); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + SX::kXA)); + uint32_t const wsf0 = smem_u32(smem + SX::kRingSf); + uint32_t const xsf0 = smem_u32(smem + SX::kXASf); + int copied = 0; // tiles [0, copied) are in TMEM + bool act = false; +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt && !act) + { + while (!mbar_test(gbar + 8 * i, 0)) + { + if (mbar_test(act_ready, 0)) + { + act = true; + break; + } + } + if (!act) + { + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kTA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + copied = i + 1; + } + } + } + while (!act) + { + act = mbar_test(act_ready, 0); + } + M3FA_PROF_T(21, 32); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + bool const in_tmem = i < copied; + if (!in_tmem) + { + mbar_wait_epi(gbar + 8 * i, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + } + uint32_t const sfb = tm + kTSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + uint64_t const bd = b0 + ((i * 1024 + kb * 32) >> 4); + if (in_tmem) + { + umma_mx(tm + kTAcc, tm + kTA + j * 8, bd, qkv_idesc(kb), tm + kTSfa + i * 4, sfb, j > 0); + } + else + { + umma_mx_ss(tm + kTAcc, w0 + ((i * kQkvWTile + kb * 32) >> 4), bd, qkv_idesc(kb), tm + kTSfa + i * 4, + sfb, j > 0); + } + } + } + } + M3FA_PROF_T(22, 32); + M3FA_PROF_T(31, 32); + umma_commit(gbar + 16 * kQkvRing); // a_ready (for its waiters) + umma_commit(gbar + 16 * kQkvRing + 8); // acc_done +} + +// The other lanes of warps 0 and 1 hold at __syncwarp rather than at griddepcontrol.wait, which would stall the +// whole warp (lane 0's loop included) until the producer grid completes. +template +__device__ __forceinline__ void qkv_stream( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int lane, int u) +{ + if (M3FA_QKV_PIPE == 2 && u >= 0 && warp == 1) + { + if (elect_one()) + { + qkv_hybrid(p, smem, gbar, tmem_slot, u); + } + __syncwarp(); + return; + } + if (u < 0 || warp > 1 || M3FA_QKV_PIPE) + { + return; // pipelined: the MMA thread takes each tile as it lands + } + if (lane == 0) + { + qkv_stream_lane(p, smem, gbar, tmem_slot, warp, u); + } + __syncwarp(); +} + +// After the PDL wait. Every thread quantizes a piece of the x slice into the B tiles (named barrier 9); warp 1 then +// issues the scale copies and the MMAs. +template +__device__ __forceinline__ void qkv_mma( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int tid, int u) +{ + int const warp = tid >> 5; + int const s = u % kQkvSplits; + int const kt0 = qkv_tile0(s); + int const nt = qkv_tile0(s + 1) - kt0; + { + // The MXFP8 activation slice (E4M3 rows of the unit's k-tiles in 128-byte-swizzled B tiles, and their scale + // atoms), copied with plain 16-byte loads over the overlaid ring once its tiles went on into TMEM: every load + // first, then the stores (one L2 round trip). + constexpr int kXDst = M3FA_QKV_PIPE ? SX::kXA : SX::kXt; + constexpr int kXSfDst = M3FA_QKV_PIPE ? SX::kXASf : SX::kXsf; + if (!M3FA_QKV_PIPE) + { + mbar_wait_cta(gbar + 16 * kQkvRing, 0); // a_ready: the ring's weight tiles are in TMEM + } + constexpr int kXItems = kQkvMaxTiles * 8 * kMaxTok; // (tile, 16-byte chunk, token) + constexpr int kSfItems = kQkvMaxTiles * (kQkvSfAtom / 16); + constexpr bool kHyb = M3FA_QKV_PIPE == 2; + constexpr int kCopiers = kHyb ? kThreads - 32 : kThreads; // PIPE 2: warp 1 is issuing the tiles + if (kHyb && warp == 1) + { + return; + } + int const ctid = kHyb && tid >= 64 ? tid - 32 : tid; + constexpr int kPer = (kXItems + kSfItems + kCopiers - 1) / kCopiers; + uint4 v[kPer]; + int dst[kPer]; +#pragma unroll + for (int k = 0; k < kPer; ++k) + { + int const q = ctid + kCopiers * k; + dst[k] = -1; + if (q < kXItems) + { + int const t = q % kMaxTok; + int const c = (q / kMaxTok) % 8; + int const i = q / (kMaxTok * 8); + if (i < nt && t < p.num_tokens) + { + v[k] = *reinterpret_cast( + p.xq + static_cast(t) * p.ldxq + (kt0 + i) * 128 + c * 16); + dst[k] = kXDst + i * 1024 + t * 128 + ((c ^ t) * 16); + } + } + else if (q < kXItems + kSfItems) + { + int const j = q - kXItems; + int const i = j / (kQkvSfAtom / 16); + int const c = j % (kQkvSfAtom / 16); + if (i < nt) + { + v[k] = *reinterpret_cast( + p.xsf_act + static_cast(kt0 + i) * kQkvSfAtom + c * 16); + dst[k] = kXSfDst + i * kQkvSfAtom + c * 16; + } + } + } +#pragma unroll + for (int k = 0; k < kPer; ++k) + { + if (dst[k] >= 0) + { + *reinterpret_cast(smem + dst[k]) = v[k]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if (kHyb) + { + asm volatile("mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];\n" ::"r"(gbar + 8 * (2 * kQkvRing + 2)) + : "memory"); + return; + } + if (warp != 1) + { + asm volatile("bar.arrive 9, %0;\n" ::"n"(kThreads) : "memory"); + } + } + if (warp == 1) + { + M3FA_PROF_T(32, 32); + asm volatile("bar.sync 9, %0;\n" ::"n"(kThreads) : "memory"); + M3FA_PROF_T(21, 32); + // Operands computed by the whole warp (warp-uniform, kept in uniform registers); one elected lane issues. + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + (M3FA_QKV_PIPE ? SX::kXA : SX::kXt))); + uint64_t const a0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + if (elect_one()) + { + M3FA_CLK(41); + if constexpr (!kQkvScaled) + { +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + mbar_wait_cta(gbar + 8 * i, 0); // weight tile i landed + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + umma_f8_ss(tm + kTAcc + (i * 4 + kb) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4), + b0 + ((i * 1024 + kb * 32) >> 4)); + } + umma_commit(gbar + 8 * (2 * kQkvRing + 2 + i)); // tile i's blocks are in TMEM + } + } + } + else + { + uint32_t const xsf0 = smem_u32(smem + (M3FA_QKV_PIPE ? SX::kXASf : SX::kXsf)); + uint32_t const wsf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + if constexpr (M3FA_QKV_PIPE != 0) + { + // Tile i as it lands: its A copy and scale atom into TMEM, then its MMAs (issue order). + mbar_wait_cta(gbar + 8 * i, 0); + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + if (i == nt - 1) + { + M3FA_CLK(47); // the unit's last weight tile landed (clock64 builds) + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if constexpr (!kQkvASmem) + { +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b( + tm + kTA + (i * 4 + kb) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + } + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + } + uint32_t const sfb = tm + kTSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + uint64_t const bd = b0 + ((i * 1024 + kb * 32) >> 4); + if constexpr (kQkvASmem) + { + umma_mx_ss(tm + kTAcc + (j % kQkvNacc) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4), bd, + qkv_idesc(kb), tm + kTSfa + i * 4, sfb, j >= kQkvNacc); + } + else + { + umma_mx(tm + kTAcc + (j % kQkvNacc) * 8, tm + kTA + j * 8, bd, qkv_idesc(kb), + tm + kTSfa + i * 4, sfb, j >= kQkvNacc); + } + } + } + } + } + M3FA_CLK(42); + if (M3FA_QKV_PIPE) + { + M3FA_PROF_T(31, 32); + umma_commit(gbar + 16 * kQkvRing); // a_ready (kept for its waiters: every tile is consumed) + } + umma_commit(gbar + 16 * kQkvRing + 8); // acc_done + } + M3FA_PROF_T(22, 32); + } +} + +// Warps 16..19 (TMEM lanes 32 (warp % 4) + lane = the unit's rows): partial rows out, the other parity reset. +template +__device__ __forceinline__ void qkv_epilogue( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int tid, int u, int par, unsigned int calls) +{ + int const warp = tid >> 5; + int const lane = tid & 31; + if constexpr (kQkvScaled) + { + mbar_wait_epi(gbar + 16 * kQkvRing + 8, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + M3FA_PROF_T(24, 512); + if (tid == 512) + { + M3FA_CLK(43); + } + int const r = (warp & 3) * 32 + lane; // row within the unit's tile + uint32_t const taddr = *tmem_slot + (((warp & 3) * 32) << 16) + kTAcc; + float val[kMaxTok]; + if constexpr (kQkvScaled) + { + float v[32]; + tmem_ld16(taddr, *reinterpret_cast(v)); + if constexpr (kQkvNacc == 4) + { + tmem_ld16(taddr + 16, *reinterpret_cast(v + 16)); + } +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = kQkvNacc == 4 ? (v[t] + v[8 + t]) + (v[16 + t] + v[24 + t]) + : kQkvNacc == 2 ? v[t] + v[8 + t] + : v[t]; + } + } + else + { + // Sum over K blocks in order: block result x (weight scale x activation scale). + int const nt = qkv_tile0(u % kQkvSplits + 1) - qkv_tile0(u % kQkvSplits); + float const* xsc = reinterpret_cast(smem + SX::kXsc); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = 0.0f; + } +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + // Tile i's four blocks as soon as its MMAs complete: the epilogue trails the MMAs by one tile. + mbar_wait_cta(gbar + 8 * (2 * kQkvRing + 2 + i), 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + float d[32]; + tmem_ld16_nowait(taddr + i * 32, d); + tmem_ld16_nowait(taddr + i * 32 + 16, d + 16); + uint32_t const wsw = *reinterpret_cast( + smem + SX::kRingSf + i * kQkvSfAtom + (r % 32) * 16 + (r / 32) * 4); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + float const ws = pow2_from_ue8m0((wsw >> (8 * kb)) & 0xFFu); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = fmaf(d[kb * 8 + t], ws * xsc[(i * 4 + kb) * kMaxTok + t], val[t]); + } + } + } + } + } + // v52's split-K reduce: split s owns rows [32 s, 32 s + 32) of the tile; the other splits push their fp32 partials + // of those rows into its mailbox (st.async, completing its reduce barrier's bytes); the owner adds the splits in + // rank order from 0.0 and stores bf16. + int const s = u % kQkvSplits; + int const owner = r / kQkvOwnRows; + int const local = r % kQkvOwnRows; + uint32_t const red = gbar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); + uint32_t cta_rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + // Units are numbered down from the last CTA: split q of this tile is CTA rank cta_rank - (q - s). + if (owner != s) + { + uint32_t const peer = cta_rank + s - owner; + uint32_t const dst + = mapa_shared(smem_u32(smem + SX::kMailbox) + ((s * kQkvOwnRows + local) * kQkvMbStride) * 4, peer); + uint32_t const pbar = mapa_shared(red, peer); +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( + dst + 16 * g), + "r"(__float_as_uint(val[4 * g])), "r"(__float_as_uint(val[4 * g + 1])), + "r"(__float_as_uint(val[4 * g + 2])), "r"(__float_as_uint(val[4 * g + 3])), "r"(pbar) + : "memory"); + } + } + else + { + M3FA_PROF_T(26, tid - lane); + // test_wait (no suspend): completions by peers' st.async bytes wake a suspended try_wait late. + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(red) + : "memory"); + } + M3FA_PROF_T(27, tid - lane); + float const* mb = reinterpret_cast(smem + SX::kMailbox); + int const row = (u / kQkvSplits) * 128 + r; + // Lamport hand-off: this call's rows go to buffer (calls & 1) with relaxed stores, the consumers poll the + // values themselves (0xFFFF, a NaN the rounding of a finite sum never produces, marks "not yet"); the other + // buffer, last read by the previous call (complete: this grid passed its PDL wait), is reset to the marker. + uint16_t* cur = reinterpret_cast(p.qkv_out) + (calls & 1u) * kMaxTok * kQkvN; + uint16_t* nxt = reinterpret_cast(p.qkv_out) + ((calls & 1u) ^ 1u) * kMaxTok * kQkvN; +#pragma unroll + float qsum[kMaxTok]; +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + float4 part[kQkvSplits]; +#pragma unroll + for (int q = 0; q < kQkvSplits; ++q) + { + part[q] = *reinterpret_cast(mb + (q * kQkvOwnRows + local) * kQkvMbStride + 4 * g); + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float sum = 0.0F; +#pragma unroll + for (int q = 0; q < kQkvSplits; ++q) + { + float const pj = j == 0 ? part[q].x : j == 1 ? part[q].y : j == 2 ? part[q].z : part[q].w; + sum = sum + (q == s ? val[4 * g + j] : pj); + } + qsum[4 * g + j] = sum; + } + } +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + float const sum = qsum[t]; + if (t < p.num_tokens) + { + __nv_bfloat16 const b = __float2bfloat16_rn(sum); + uint16_t bits = *reinterpret_cast(&b); + // Never the "not yet" marker: a 0xFFFF NaN sum is stored as the NaN 0x7FFF. Only a NaN projection row + // changes (its K/V cache bytes then differ from v52 + p6 in the NaN's encoding); finite rows are exact. + bits = bits == 0xFFFFu ? static_cast(0x7FFFu) : bits; +#ifndef M3FA_QKV_REL +#define M3FA_QKV_REL 0 // 1: the row store is a release (tried: whether it becomes visible to the pollers sooner) +#endif + if (M3FA_QKV_REL) + { + asm volatile( + "st.release.gpu.global.u16 [%0], %1;\n" ::"l"(cur + static_cast(t) * kQkvN + row), + "h"(bits) + : "memory"); + } + else + { + asm volatile( + "st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(cur + static_cast(t) * kQkvN + row), + "h"(bits) + : "memory"); + } + } + asm volatile("st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(nxt + static_cast(t) * kQkvN + row), + "h"(static_cast(0xFFFFu)) + : "memory"); + } + M3FA_PROF_T(28, tid - lane); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + M3FA_PROF_T(25, 512); + if (tid == 512) + { + M3FA_CLK(46); + } +} + +// ---------------- Megakernel o_proj (v52 m3_mx_gemv bit for bit: split 2, contiguous k-tiles, nacc 1) ---------------- +// The unit's weight tiles and scale atoms into the ring (warp 0, lane 0). From L2 when the CTA prefetched them there. +template +__device__ __forceinline__ void op_issue(Params const& p, uint8_t* smem, int ou) +{ + int const rt = ou / kOpSplits; + int const kt0 = (ou % kOpSplits) * kOpTiles; + uint32_t const ob = smem_u32(smem + SX::kOBar); + for (int i = 0; i < kOpTiles; ++i) + { + mbar_expect_tx(ob + 8 * i, kQkvWTile + kQkvSfAtom); + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap_o, (kt0 + i) * 128, rt * 128, ob + 8 * i); + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf_o + (static_cast(rt) * kOpKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, ob + 8 * i); + } +} + +// L2 prefetch of the unit's weight tiles and scale atoms (one lane; nothing waits for it). +// The QKV units' share of the window-1 L2 prefetch (the next consumer's weights), issued once their own QKV weight +// stream has landed (acc_done): the attention core reads little from HBM. +__device__ __forceinline__ void window_prefetch_l2(Params const& p, int u) +{ + constexpr long long kChunk = 32768; + long long const n0 = (p.pf_b0 + kChunk - 1) / kChunk, n1 = (p.pf_b1 + kChunk - 1) / kChunk; + uint64_t pol = 0; + if (p.pf_hint) + { + asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, 1.0;\n" : "=l"(pol)); + } + for (long long c = u; c < n0 + n1; c += kQkvUnits) + { + long long const a = c < n0 ? p.pf_a0 + c * kChunk : p.pf_a1 + (c - n0) * kChunk; + long long const rem = c < n0 ? p.pf_b0 - c * kChunk : p.pf_b1 - (c - n0) * kChunk; + unsigned int const n = static_cast(rem < kChunk ? rem : kChunk) & ~15u; + if (n == 0) + { + continue; + } + if (p.pf_hint) + { + asm volatile("cp.async.bulk.prefetch.L2.global.L2::cache_hint [%0], %1, %2;\n" ::"l"(a), "r"(n), "l"(pol) + : "memory"); + } + else + { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" ::"l"(a), "r"(n) : "memory"); + } + } +} + +__device__ __forceinline__ void op_prefetch_l2(Params const& p, int ou) +{ + int const rt = ou / kOpSplits; + int const kt0 = (ou % kOpSplits) * kOpTiles; + for (int i = 0; i < kOpTiles; ++i) + { + asm volatile("cp.async.bulk.prefetch.tensor.2d.L2.global [%0, {%1, %2}];\n" ::"l"(&p.wmap_o), + "r"((kt0 + i) * 128), "r"(rt * 128) + : "memory"); + } + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" ::"l"( + p.wsf_o + (static_cast(rt) * kOpKTiles + kt0) * kQkvSfAtom), + "r"(kOpTiles * kQkvSfAtom) + : "memory"); +} + +// The weight tiles into TMEM as they land (warp 1, elected lane, operands warp-uniform). +template +__device__ __forceinline__ void op_stream(uint8_t* smem, uint32_t tm) +{ + uint32_t const ob = smem_u32(smem + SX::kOBar); + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint32_t const sf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kOpTiles; ++i) + { + mbar_wait_cta(ob + 8 * i, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kOA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + tmem_cp_sf(tm + kOSfa + i * 4, sf0 + i * kQkvSfAtom); + } + umma_commit(ob + 8 * kOpTiles); // a_ready +} + +// The attention's MXFP8 rows of this unit's k-tiles (the heads [kt0, kt0 + 8)), Lamport-polled value by value: 0xFF +// is NaN in both E4M3 and UE8M0, which the merge never writes (see the merge). Every thread of the CTA, one item each. +template +__device__ __forceinline__ void op_act(Params const& p, uint8_t* smem, int ou, int tid, unsigned int calls) +{ + int const kt0 = (ou % kOpSplits) * kOpTiles; + uint8_t const* q8 = p.q8 + (calls & 1u) * kMaxTok * kOpK; + uint8_t const* sf = p.sf + (calls & 1u) * kOpSfBytes; + constexpr int kXItems = kOpTiles * 8 * kMaxTok; + constexpr int kSfItems = kOpTiles * (kQkvSfAtom / 16); + for (int q = tid; q < kXItems + kSfItems; q += kThreads) + { + if (q < kXItems) + { + int const t = q % kMaxTok; + int const c = (q / kMaxTok) % 8; + int const i = q / (kMaxTok * 8); + if (t < p.num_tokens) + { + uint4 v; + bool pending; + do + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(q8 + static_cast(t) * kOpK + (kt0 + i) * 128 + c * 16) + : "memory"); + uint32_t const w[4] = {v.x, v.y, v.z, v.w}; + pending = false; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + // a byte equal to 0xFF: (x & 0x7F7F7F7F) + 0x01010101 sets bit 7 of exactly the 0x7F/0xFF bytes + uint32_t const hit = (((w[j] & 0x7F7F7F7Fu) + 0x01010101u) & w[j] & 0x80808080u); + pending |= hit != 0u; + } + } while (pending); + *reinterpret_cast(smem + SX::kOXt + i * 1024 + t * 128 + ((c ^ t) * 16)) = v; + } + } + else if (q < kXItems + kSfItems) + { + int const j = q - kXItems; + int const i = j / (kQkvSfAtom / 16); + int const c = j % (kQkvSfAtom / 16); + uint8_t const* src = sf + (kt0 + i) * kQkvSfAtom + c * 16; + uint4 v; + if (c < p.num_tokens) + { + // Row c of the atom: its 4 block scales in bytes [16 c, 16 c + 4). + do + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(src) + : "memory"); + } while ((((v.x & 0x7F7F7F7Fu) + 0x01010101u) & v.x & 0x80808080u) != 0u); + } + else + { + v = *reinterpret_cast(src); + } + *reinterpret_cast(smem + SX::kOXsf + i * kQkvSfAtom + c * 16) = v; + } + } +} + +// The MMAs (warp 1, elected lane): v52's order (per k-tile: B scale copy, 4 MMAs; one accumulator). +template +__device__ __forceinline__ void op_mma(uint8_t* smem, uint32_t tm) +{ + uint32_t const ob = smem_u32(smem + SX::kOBar); + mbar_wait_cta(ob + 8 * kOpTiles, 0); // a_ready + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + SX::kOXt)); + uint32_t const xsf0 = smem_u32(smem + SX::kOXsf); +#pragma unroll + for (int i = 0; i < kOpTiles; ++i) + { + uint32_t const sfb = tm + kOSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + umma_mx(tm + kOAcc, tm + kOA + j * 8, b0 + ((i * 1024 + kb * 32) >> 4), qkv_idesc(kb), tm + kOSfa + i * 4, + sfb, j > 0); + } + } + umma_commit(ob + 8 * (kOpTiles + 1)); // acc_done +} + +// Warps 16..19: the pair's split-K reduce (owner rows [64 s, 64 s + 64)) and the bf16 rows of y. +template +__device__ __forceinline__ void op_epilogue(Params const& p, uint8_t* smem, uint32_t tm, int tid, int ou) +{ + int const warp = tid >> 5; + int const lane = tid & 31; + uint32_t const ob = smem_u32(smem + SX::kOBar); + // The store's kernel parameters, fetched before the wait: a first touch of their constant-bank line costs a miss. + __nv_bfloat16* y_out = p.y_out; + int num_tokens = p.num_tokens; + int ar_mode = p.ar_mode; + asm volatile("" : "+l"(y_out), "+r"(num_tokens), "+r"(ar_mode)); +#ifndef M3FA_OP_PARPRE +#define M3FA_OP_PARPRE 1 // read the all-reduce slot parities before the accumulator wait (a global load off the tail) +#endif + // The parity flags are stable here: the all-reduce that flips them last completed before this grid's PDL wait, + // and the next one reads them only after this kernel's pushes. + int par_pre[kMaxTok]; + { + int const row_pre = (ou / kOpSplits) * 128 + ((tid >> 5) & 3) * 32 + (tid & 31); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + par_pre[t] = 0; + if (M3FA_OP_PARPRE && ar_mode != 0 && t < num_tokens) + { + asm volatile("ld.volatile.global.s32 %0, [%1];\n" + : "=r"(par_pre[t]) + : "l"(p.ar_flags + t * 8 + row_pre / (kOpN / 8)) + : "memory"); + } + } + } + mbar_wait_epi(ob + 8 * (kOpTiles + 1), 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF_T(51, kComputeThreads); + int const r = (warp & 3) * 32 + lane; + float v[16]; + tmem_ld16(tm + (((warp & 3) * 32) << 16) + kOAcc, v); + int const s = ou % kOpSplits; + int const owner = r / kOpOwnRows; + int const local = r % kOpOwnRows; + uint32_t const red = ob + 8 * (kOpTiles + 2); + uint32_t cta_rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + if (owner != s) + { + uint32_t const peer = cta_rank + owner - s; // o_proj units count up with the CTA index + uint32_t const dst + = mapa_shared(smem_u32(smem + SX::kOMailbox) + ((s * kOpOwnRows + local) * kQkvMbStride) * 4, peer); + uint32_t const pbar = mapa_shared(red, peer); +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( + dst + 16 * g), + "r"(__float_as_uint(v[4 * g])), "r"(__float_as_uint(v[4 * g + 1])), "r"(__float_as_uint(v[4 * g + 2])), + "r"(__float_as_uint(v[4 * g + 3])), "r"(pbar) + : "memory"); + } + } + else + { + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(red) + : "memory"); + } + M3FA_PROF_T(52, tid - lane); + float const* mb = reinterpret_cast(smem + SX::kOMailbox); + int const row = (ou / kOpSplits) * 128 + r; + float sums[kMaxTok]; +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + float4 part[kOpSplits]; +#pragma unroll + for (int q = 0; q < kOpSplits; ++q) + { + part[q] = *reinterpret_cast(mb + (q * kOpOwnRows + local) * kQkvMbStride + 4 * g); + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float sum = 0.0F; +#pragma unroll + for (int q = 0; q < kOpSplits; ++q) + { + float const pj = j == 0 ? part[q].x : j == 1 ? part[q].y : j == 2 ? part[q].z : part[q].w; + sum = sum + (q == s ? v[4 * g + j] : pj); + } + sums[4 * g + j] = sum; + } + } + M3FA_PROF_T(54, tid - lane); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + float const sum = sums[t]; +#ifndef M3FA_EXP +#define M3FA_EXP 0 +#endif + if (t < num_tokens) + { + if (M3FA_EXP == 1) + { + asm volatile( + "st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(y_out + static_cast(t) * kOpN + row), + "h"(__bfloat16_as_ushort(__float2bfloat16_rn(sum))) + : "memory"); + } + else + { + y_out[static_cast(t) * kOpN + row] = __float2bfloat16_rn(sum); + } + } + if (ar_mode != 0) + { + // m3_mx_gemv's push into the all-reduce exchange: lane pairs hold row pairs; the even lane writes the + // bf16x2 word into this rank's slot of the (token, column slice) parity the AR's flags select. + float const hi = __shfl_xor_sync(0xffffffffu, sum, 1); + if ((lane & 1) == 0 && t < num_tokens) + { + int par; + if (M3FA_OP_PARPRE) + { + par = par_pre[t]; + } + else + { + asm volatile("ld.volatile.global.s32 %0, [%1];\n" + : "=r"(par) + : "l"(p.ar_flags + t * 8 + row / (kOpN / 8)) + : "memory"); + } + int64_t const off = static_cast(par & 1) * p.ar_maxt * p.ar_world * (kOpN / 2) + + (static_cast(t) * p.ar_world + p.ar_rank) * (kOpN / 2) + row / 2; + __nv_bfloat162 w2 = __floats2bfloat162_rn(sum, hi); + uint32_t word = *reinterpret_cast(&w2); + // -0.0 halves to +0.0: 0x80000000 marks an empty exchange word (as m3_mx_gemv's _bf16x2_word) + word = (word & 0xFFFFu) == 0x8000u ? (word & 0xFFFF0000u) : word; + word = (word >> 16) == 0x8000u ? (word & 0x0000FFFFu) : word; + asm volatile("st.global.b32 [%0], %1;\n" ::"l"(p.ar_ws + off), "r"(word) : "memory"); + } + } + } + M3FA_PROF_T(53, tid - lane); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); +} + +// Packed-QKV element group [token][col * 128 + 4 * lane, +4): the split partials, polled until every one is stored, +// summed in split order and rounded to BF16. +__device__ __forceinline__ void qkv_load(Params const& p, int par, int tok, int col, int lane, float (&e)[4]) +{ + float const* base = p.part + par * kQkvPartStride + static_cast(tok) * kQkvN + col * kHeadDim + lane * 4; + uint4 pr[kQkvSplits]; + long long spins = 0; + bool ready; + do + { + ready = true; +#pragma unroll + for (int s = 0; s < kQkvSplits; ++s) + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(pr[s].x), "=r"(pr[s].y), "=r"(pr[s].z), "=r"(pr[s].w) + : "l"(base + static_cast(s) * kMaxTok * kQkvN) + : "memory"); + ready &= pr[s].x != kQkvSentinel && pr[s].y != kQkvSentinel && pr[s].z != kQkvSentinel + && pr[s].w != kQkvSentinel; + } + if (++spins > (1LL << 26)) + { + __trap(); + } + } while (!ready); + float s0 = __uint_as_float(pr[0].x), s1 = __uint_as_float(pr[0].y), s2 = __uint_as_float(pr[0].z), + s3 = __uint_as_float(pr[0].w); +#pragma unroll + for (int s = 1; s < kQkvSplits; ++s) + { + s0 += __uint_as_float(pr[s].x); + s1 += __uint_as_float(pr[s].y); + s2 += __uint_as_float(pr[s].z); + s3 += __uint_as_float(pr[s].w); + } + e[0] = __bfloat162float(__float2bfloat16_rn(s0)); + e[1] = __bfloat162float(__float2bfloat16_rn(s1)); + e[2] = __bfloat162float(__float2bfloat16_rn(s2)); + e[3] = __bfloat162float(__float2bfloat16_rn(s3)); +} + +// Every CTA once it has read the call parity and the partials it needs: the last one flips the parity. +__device__ __forceinline__ void qkv_done(Params const& p) +{ + if (atomicAdd(p.counters + 1, 1) == kQkvGrid - 1) + { + atomicExch(p.counters + 1, 0); + atomicXor(p.counters, 1); + } +} + +} // namespace + +namespace +{ + +// FQKV: the kernel also computes the QKV(+index) projection from the layer input (see qkv_mma) and runs on a fixed +// kQkvGrid-CTA grid; the first attn_clusters clusters then run the attention on the projection's output. +template +__global__ void __launch_bounds__(kThreads, 1) m3_fused_decode_attn_kernel(__grid_constant__ Params const p) +{ +#if MINIMAX_M3_ATTN_DEVICE_ENABLED + using S = SmemX; + constexpr int kRows = S::kRows; + constexpr int kTok = S::kTok; + constexpr int kCs = kBlocks / BPC; // cluster size + constexpr int kGroups = QLEN / QG; // clusters per request + constexpr uint32_t kPushBytes = (kHeadDim + 1) * 4; // one partial row: 128 outputs and the LSE + constexpr int kTasks = kHq * QG + 3 * QLEN; // producer tasks: 16 Q heads x QG, [K, V, index-K] x QLEN + constexpr int kMaxT = (kTasks + kWarps - 1) / kWarps; + constexpr int kChunks = BPC * kPage * 4; // 16-byte packed chunks per K (or V) tile set + constexpr int kKPer = (kChunks + kComputeThreads - 1) / kComputeThreads; // K chunks per thread [0, 512) + constexpr bool kRegCore = RegCore::kOn; // warp-MMA attention core (see RegCore) + // V tiles (raw copy before the PDL wait, dequantized during QK and the softmax): on the packed path every warp + // outside the softmax groups takes part (at M=1, warps 4..19), on the fused-projection path the producer warps. +#ifndef M3FA_RS2 +#define M3FA_RS2 1 // megakernel at one 16-row chunk: the softmax and push rows split over two warp groups (8 each) +#endif + constexpr bool kRs2 = FQKV && M3FA_RS2 && BPC * QG == 1 && !RegCore::kOn; + constexpr int kSoftmaxGroups = kRs2 ? 2 : (BPC * QG) < 4 ? (BPC * QG) : 4; +#ifndef M3FA_MEGA_VWIDE +#define M3FA_MEGA_VWIDE \ + 1 // megakernel attention CTAs (no projection unit): every warp outside the softmax copies and + // dequantizes V, as on the packed path (not only the four producer warps) +#endif + constexpr int kVStart = (kRegCore || (FQKV && !M3FA_MEGA_VWIDE)) ? kComputeThreads : 128 * kSoftmaxGroups; + constexpr int kVThreads = kThreads - kVStart; // V chunks: the warps outside the softmax groups + constexpr int kVPer = (kChunks + kVThreads - 1) / kVThreads; + constexpr int kPairs = BPC * QG; // (block, 16-row chunk) pairs of the softmax and push + constexpr int kPairsPerGroup = (kPairs + 3) / 4; // per group of four compute warps + // Named barrier 6: the softmax warps (P written) and the producer warps (V written) before the PV issue. + constexpr int kPvBarThreads = 128 * kSoftmaxGroups + kVThreads; + extern __shared__ __align__(128) uint8_t smem[]; + float* row_scale = reinterpret_cast(smem + S::kRowScale); + float* row_lse = reinterpret_cast(smem + S::kRowLse); + float* red_m = reinterpret_cast(smem + S::kRedM); + float* red_l = reinterpret_cast(smem + S::kRedL); + float* recv = reinterpret_cast(smem + S::kRecv); + uint32_t const mbar = smem_u32(smem + S::kMbar); + uint32_t const mma_bar = mbar + 8; + uint32_t* tmem_slot = reinterpret_cast(smem + S::kMbar + 16); + + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const cr = blockIdx.x % kCs; // cluster rank; owns blocks [cr * BPC, cr * BPC + BPC) + int const ci = blockIdx.x / kCs; + int const req = ci / kGroups; + int const tfirst = (ci % kGroups) * QG; // this cluster's query tokens [tfirst, tfirst + QG) of the request + int const tok0 = req * QLEN; + uint32_t const gbar = smem_u32(smem + S::kGBar); // FQKV: weight ring, A ready, accumulators done + int const qkv_u = FQKV ? qkv_unit() : -1; // FQKV: this CTA's projection unit, or -1 + + // ---------------- Prologue: nothing here depends on the producer kernel ---------------- +#if M3FA_TRIGGER_AT == 0 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (!FQKV) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + M3FA_PROF(0); +#ifdef M3FA_PROFILE + if (p.prof != nullptr && tid == 0) + { + uint32_t smid; + asm volatile("mov.u32 %0, %%smid;" : "=r"(smid)); + p.prof[blockIdx.x * 64 + 40] = smid; + } +#endif + if constexpr (FQKV) + { + // The projection weights stream into TMEM while the producer (the all-reduce) still runs. + if (warp < 2) + { + qkv_setup(p, smem, gbar, tmem_slot, warp, lane, qkv_u); + } + if (ci >= p.attn_clusters) + { + // A projection-only CTA. Every reduce barrier of the cluster is initialized before any peer pushes. + asm volatile("barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;\n" :: + : "memory"); + int const op_u = static_cast(blockIdx.x) - kOpFirstCta; // this CTA's o_proj unit (>= 0 here) +#ifndef M3FA_OP_EARLY +#define M3FA_OP_EARLY 0 // 1: o_proj-only CTAs stream their weights from launch (0: after the PDL wait) +#endif + if (M3FA_OP_EARLY && qkv_u < 0 && warp == 0 && lane == 0) + { + op_issue(p, smem, op_u); // an o_proj-only CTA: its weights stream from launch on + } + if (M3FA_OP_L2PF && qkv_u >= 0 && warp == 0 && lane == 0) + { + op_prefetch_l2(p, op_u); // into L2 during the all-reduce, behind this CTA's QKV ring + } + qkv_stream(p, smem, gbar, tmem_slot, warp, lane, qkv_u); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (M3FA_MEGA_TRIGGER == 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + // This CTA's call counter: read after the PDL wait (the previous call's CTA increments it at its end), and + // early, so its latency stays off the epilogue's critical path. + unsigned int const calls = p.flags[32 + blockIdx.x]; + M3FA_PROF(1); + if (!M3FA_OP_EARLY && qkv_u < 0 && warp == 0 && lane == 0) + { + op_issue(p, smem, op_u); + } + if (qkv_u >= 0) + { + int const par = 0; + qkv_mma(p, smem, gbar, tmem_slot, tid, qkv_u); + if (warp >= kComputeWarps) + { + qkv_epilogue(p, smem, gbar, tmem_slot, tid, qkv_u, par, calls); + } + if (warp == 0 && lane == 0) + { + // The ring (and the activation over it) is free once the QKV MMAs completed. + mbar_wait_cta(gbar + 16 * kQkvRing + 8, 0); + op_issue(p, smem, op_u); + if (p.pf_b0 > 0 || p.pf_b1 > 0) + { + window_prefetch_l2(p, qkv_u); + } + } + } + __syncthreads(); // the QKV accumulator reads are done: TMEM columns [0, 512) are the o_proj's + M3FA_PROF(2); + uint32_t const tm = *tmem_slot; + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_stream(smem, tm); + } + __syncwarp(); + } + op_act(p, smem, op_u, tid, calls); + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + M3FA_PROF_T(36, 0); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Late release: the dependents launch once every unit has its attention input (the rel_v4 rule). + if (M3FA_MEGA_TRIGGER < 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + __syncthreads(); + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_mma(smem, tm); + M3FA_PROF_T(50, threadIdx.x); + } + __syncwarp(); + } + if (warp >= kComputeWarps) + { + op_epilogue(p, smem, tm, tid, op_u); + } + M3FA_PROF_T(37, kComputeThreads); + __syncthreads(); + M3FA_PROF_T(55, 0); + if (M3FA_MEGA_TRIGGER == 3) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // after this CTA's y rows / pushes + } + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" ::"r"(*tmem_slot), "r"(kTmemAll) + : "memory"); + M3FA_PROF_T(56, 32); + } + if (tid == 0) + { + p.flags[32 + blockIdx.x] = calls + 1u; // a plain store: no load on the exit path + } + M3FA_PROF(39); +#ifndef M3FA_FINAL_CLBAR +#define M3FA_FINAL_CLBAR 0 +#endif + if (M3FA_FINAL_CLBAR) + { + // No peer may exit while its partials can still be pushed into it (DSMEM lifetime). + asm volatile("barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;\n" :: + : "memory"); + } + M3FA_PROF(38); + return; + } + } +#ifndef M3FA_MEGA_ATTN_WAIT_FIRST +#define M3FA_MEGA_ATTN_WAIT_FIRST 0 // 1: the attention CTAs read nothing (metadata, caches) before their PDL wait +#endif +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (FQKV && M3FA_MEGA_ATTN_WAIT_FIRST) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } +#endif + if (!FQKV && warp == kWarps - 1) // a producer warp: it has nothing else to do before the PDL wait + { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n" ::"r"(smem_u32(tmem_slot)), + "r"(S::kTmemCols) + : "memory"); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n" ::: "memory"); + } + int const seq_len = p.seq_lens[req]; + int page[BPC]; +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + int const b = cr * BPC + h; + page[h] = b < p.bt_cols ? p.block_table[static_cast(req) * p.bt_stride + b] : 0; + } + // New-token cache slots, for the tokens up to this cluster's last one. + int slot[QLEN]; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + slot[t] = t < tfirst + QG ? p.out_cache_loc[tok0 + t] : -1; + } + // This warp's producer tasks and, for every lane that applies RoPE, the task token's cos/sin. + float cs[kMaxT][4], sn[kMaxT][4]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + int const t = i < kHq * QG ? tfirst + i / kHq : (i - kHq * QG) / 3; + if (i < kTasks && lane < 16 && t < QLEN) + { + int64_t const row = static_cast(p.position_ids[tok0 + t]) * kRotDim; +#pragma unroll + for (int e = 0; e < 4; ++e) + { + int const coefficient = (lane * 4 + e) & (kRotDim / 2 - 1); + cs[k][e] = p.rope[row + coefficient]; + sn[k][e] = p.rope[row + kRotDim / 2 + coefficient]; + } + } + else + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + cs[k][e] = 1.0F; + sn[k][e] = 0.0F; + } + } + } + float qw[4], kw[4], ikw[4]; + load_w4(qw, p.q_w, lane); + load_w4(kw, p.k_w, lane); + load_w4(ikw, p.ik_w, lane); + float const kqs = p.kv_quant_scale[1]; + float const vqs = p.kv_quant_scale[2]; + float const sl2e = (p.sm_scale * 1.4426950409f) * p.k_gs[0]; + float const vgs = p.v_gs[0]; + // Longest attended length among this cluster's tokens, and each block's rows within it. + int const kv_max = max(seq_len - QLEN + tfirst + QG, 0); + int nvh[BPC]; +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + nvh[h] = min(max(kv_max - (cr * BPC + h) * kPage, 0), kPage); + } + bool const any_valid = nvh[0] > 0; + if (tid == 32) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(mbar), "r"(1)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(mma_bar), "r"(1)); + asm volatile( + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(mbar), "r"(kBlocks * QG * BPC * kPushBytes) + : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + } + // This thread's packed K / V chunks, held in registers until they are dequantized after the PDL wait. + RawChunk kraw[kKPer]; +#pragma unroll + for (int k = 0; k < kKPer; ++k) + { + int const q = tid + kComputeThreads * k; + kraw[k] = RawChunk{make_uint4(0u, 0u, 0u, 0u), 0u}; + int const h = q / (kPage * 4); + if (q < kChunks && tid < kComputeThreads && nvh[h] > 0) + { + kraw[k] = load_chunk(p.kv_data + static_cast(page[h]) * p.d_page, + p.kv_scale + static_cast(page[h]) * p.s_page, chunk_row(q % (kPage * 4)), + chunk_col(q % (kPage * 4)), false); + } + } + // V goes through cp.async into shared memory: an asynchronous copy does not hold the producer warps up the way + // an outstanding register load does, and they need V only after QK has been issued. + if (tid >= kVStart) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + if (nvh[h] > 0) + { + uint8_t* dst = smem + S::kRawV + h * S::kRawVBlock; + uint8_t const* data = p.kv_data + static_cast(page[h]) * p.d_page + p.d_plane; + uint8_t const* scales = p.kv_scale + static_cast(page[h]) * p.s_page + p.s_plane; + for (int c = tid - kVStart; c < S::kRawVBlock / 16; c += kVThreads) + { + cp_async16(dst + c * 16, c < kPage * 4 ? data + c * 16 : scales + (c - kPage * 4) * 16); + } + } + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + if (warp == kWarps - 1) + { + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + // Every CTA's mbarrier is initialized before any peer pushes into it; the matching wait sits right before the + // push, so this barrier's latency hides behind the attention (with aliased partials it moves there). + if constexpr (!S::kRecvAlias) + { + asm volatile("barrier.cluster.arrive.relaxed.aligned;\n" ::: "memory"); + } + int kv_len[QG]; +#pragma unroll + for (int tl = 0; tl < QG; ++tl) + { + kv_len[tl] = max(seq_len - QLEN + tfirst + tl + 1, 0); + } + int const cta_tok0 = cr * kTok; // request position of this CTA's first block row + // New token t's row in this CTA's tiles (block * 128 + row), or -1: the CTA owns the block holding the token's + // position, and the token is attended by this cluster (t <= its last token). Every thread derives the same. + int new_row[QLEN]; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + int const pos = seq_len - QLEN + t; + int const blk = min(pos >> 7, kBlocks - 1); + int const spage = slot[t] / kPage; + bool const mine = t < tfirst + QG && pos >= 0 && blk / BPC == cr && slot[t] >= 0 && spage < p.num_pages; + new_row[t] = mine && spage == page[blk % BPC] ? (blk % BPC) * kPage + (slot[t] & (kPage - 1)) : -1; + } + auto is_new_row = [&](int r) + { + bool hit = false; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + hit |= new_row[t] == r; + } + return hit; + }; + // Cache writes of the new K/V/index-K are issued after this CTA has pushed its partials: only the fp16 tile rows + // are on the attention's path. kind 0: none, 1: NVFP4 data + scale, 2: FP8 index-K. + int wr_kind[kMaxT]; + uint8_t* wr_ptr[kMaxT]; + uint8_t* wr_sptr[kMaxT]; + uint32_t wr_val[kMaxT]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + wr_kind[k] = 0; + wr_ptr[k] = nullptr; + wr_sptr[k] = nullptr; + wr_val[k] = 0u; + } + uint32_t tmem = 0; + auto dequant_k = [&]() + { + if (tid < kComputeThreads) + { +#pragma unroll + for (int k = 0; k < kKPer; ++k) + { + int const q = tid + kComputeThreads * k; + if (q < kChunks) + { + int const h = q / (kPage * 4); + int const row = chunk_row(q % (kPage * 4)); + if (!is_new_row(h * kPage + row)) + { + store_chunk_fp16( + smem + S::kK16 + h * S::kTileBytes, row, chunk_col(q % (kPage * 4)), kraw[k], row < nvh[h]); + } + } + } + } + }; + // The producer warps' V tiles (raw tiles from cp.async, dequantized to fp16), then their PV-barrier arrival. + auto dequant_v = [&]() + { +#ifdef M3FA_DBG_NO_VDQ + if constexpr (!kRegCore) + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + return; +#endif + asm volatile("cp.async.wait_all;\n" ::: "memory"); + asm volatile("bar.sync 4, %0;\n" ::"n"(kVThreads) : "memory"); +#pragma unroll + for (int k = 0; k < kVPer; ++k) + { + int const q = tid - kVStart + kVThreads * k; + if (q < kChunks) + { + int const h = q / (kPage * 4); + int const row = chunk_row(q % (kPage * 4)); + int const cc = chunk_col(q % (kPage * 4)); + if (!is_new_row(h * kPage + row)) + { + uint8_t const* rawv = smem + S::kRawV + h * S::kRawVBlock; + uint8_t const* sc = rawv + kPage * 64 + (row >> 2) * 32 + (row & 3); + RawChunk c; + c.data = *reinterpret_cast(rawv + row * 64 + cc * 16); + c.scales + = static_cast(sc[(cc * 2) * 4]) | (static_cast(sc[(cc * 2 + 1) * 4]) << 8); + store_chunk_fp16(smem + S::kV16 + h * S::kTileBytes, row, cc, c, row < nvh[h]); + } + } + } + M3FA_FENCE_PROXY_ASYNC(); // V tiles -> tensor core reads (ordered by the barrier before the PV issue) + if constexpr (!kRegCore) + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + M3FA_PROF_T(16, kComputeThreads); + M3FA_PROF_T(17, kThreads - 1); + }; + if constexpr (FQKV) + { + qkv_stream(p, smem, gbar, tmem_slot, warp, lane, qkv_u); + } + M3FA_PROF(11); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (FQKV && M3FA_MEGA_TRIGGER == 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#if M3FA_TRIGGER_AT == 1 + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +#endif +#endif + // FQKV (megakernel): this call's QKV rows sit in buffer (calls & 1), polled value by value below. + unsigned int const mega_calls = FQKV ? p.flags[32 + blockIdx.x] : 0u; + __nv_bfloat16 const* qkv_src = FQKV ? p.packed + (mega_calls & 1u) * kMaxTok * kQkvN : p.packed; + M3FA_PROF(1); + int const par = 0; + if constexpr (FQKV) + { + if (qkv_u >= 0) + { + qkv_mma(p, smem, gbar, tmem_slot, tid, qkv_u); + } + } + + if constexpr (FQKV) + { + if (qkv_u >= 0) + { + if (warp >= kComputeWarps) + { + qkv_epilogue(p, smem, gbar, tmem_slot, tid, qkv_u, par, mega_calls); + } + else + { + // The attention regions overlay the weight ring: every weight byte has gone on into TMEM (or, with + // the MMAs reading the ring, every MMA has completed) first. + mbar_wait_cta(gbar + 16 * kQkvRing + (kQkvASmem ? 8 : 0), 0); + } + } + if (any_valid) + { + dequant_k(); + } +#ifndef M3FA_MEGA_EARLYV +#define M3FA_MEGA_EARLYV 1 // attention CTAs dequantize V here, while the QKV rows are still on their way +#endif +#ifndef M3FA_MEGA_EARLYCLW +#define M3FA_MEGA_EARLYCLW 1 // and pass the cluster barrier (peers' merge mbarriers initialized) here too +#endif + if (M3FA_MEGA_EARLYV && any_valid && tid >= kVStart) + { + dequant_v(); + } + if constexpr (M3FA_MEGA_EARLYCLW && !S::kRecvAlias) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); + } + M3FA_PROF(14); + } + + if (any_valid) // CTA-uniform + { + // ---------------- Producer ---------------- + // Q rows of this cluster's tokens in every CTA. The new K/V of every token up to this cluster's last one + // whose position falls in this CTA's blocks go straight into the fp16 tiles; the CTA writes the caches, and + // the index-K, only for its own tokens. + uint2 raw[kMaxT]; + float fq[FQKV ? kMaxT : 1][4]; + bool act[kMaxT]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + raw[k] = make_uint2(0u, 0u); + act[k] = false; + if (i < kTasks) + { + int t, col; + bool active = true; + if (i >= kHq * QG) + { + int const kv = i - kHq * QG; + t = kv / 3; + int const kind = kv % 3; + col = kind == 0 ? kColK : (kind == 1 ? kColV : kColIdxK); + int const pos = seq_len - QLEN + t; + active = t < tfirst + QG && pos >= 0 && min(pos >> 7, kBlocks - 1) / BPC == cr && slot[t] >= 0 + && slot[t] / kPage < p.num_pages && (kind != 2 || t >= tfirst); + } + else + { + t = tfirst + i / kHq; + col = i % kHq; + } + act[k] = active; + if (active) + { + uint2 const* src = reinterpret_cast( + qkv_src + static_cast(tok0 + t) * p.packed_stride + col * kHeadDim + lane * 4); + if constexpr (FQKV) + { + uint2 v; + bool pending; + do + { + asm volatile("ld.relaxed.gpu.global.v2.b32 {%0, %1}, [%2];\n" + : "=r"(v.x), "=r"(v.y) + : "l"(src) + : "memory"); + pending = (v.x & 0xFFFFu) == 0xFFFFu || (v.x >> 16) == 0xFFFFu || (v.y & 0xFFFFu) == 0xFFFFu + || (v.y >> 16) == 0xFFFFu; + } while (pending); + raw[k] = v; + } + else + { + raw[k] = *src; + } + } + } + } + M3FA_PROF(15); + // While the packed-QKV loads are in flight: dequantize this thread's K chunks to fp16 (exact: e2m1 x e4m3 + // fits fp16). Rows past the longest attended length are zeroed, as the Triton producer masks its loads; the + // new tokens' rows are written by their producer tasks instead. (Building the tiles before the PDL wait was + // measured slower: the cache loads issued at entry are still in flight there, so the wait moves later.) + if constexpr (!FQKV) + { + dequant_k(); + } +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + if (i >= kTasks || !act[k]) + { + continue; + } + float e[4]; + { + float2 const a = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw[k].x)); + float2 const b = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw[k].y)); + e[0] = a.x; + e[1] = a.y; + e[2] = b.x; + e[3] = b.y; + } + if (i < kHq * QG) + { + norm_rope(e, qw, p.eps, cs[k], sn[k], lane); + __nv_fp8x2_e4m3 const lo(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const hi(make_float2(e[2], e[3])); + uint2 hv; + hv.x = e4m3x2_to_h2(lo.__x); + hv.y = e4m3x2_to_h2(hi.__x); + *reinterpret_cast(smem + S::kQ + kmaj_off(i, lane * 4)) = hv; // row i = token * 16 + head + continue; + } + int const t = (i - kHq * QG) / 3; + int const kind = (i - kHq * QG) % 3; + int const spage = slot[t] / kPage; + int const within = slot[t] & (kPage - 1); + if (kind == 2) + { + norm_rope(e, ikw, p.eps, cs[k], sn[k], lane); + __nv_fp8x2_e4m3 const lo(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const hi(make_float2(e[2], e[3])); + wr_kind[k] = 2; + wr_val[k] = static_cast(lo.__x) | (static_cast(hi.__x) << 16); + wr_ptr[k] = p.idx_k + static_cast(spage) * p.i_page + static_cast(within) * p.i_tok + + lane * 4; + continue; + } + bool const isV = kind == 1; + if (!isV) + { + norm_rope(e, kw, p.eps, cs[k], sn[k], lane); + } + uint8_t sbyte = 0; + uint16_t const fp4 = quantizeMinimaxM3Fp4x4(e, isV ? vqs : kqs, &sbyte); + if (t >= tfirst) + { + int const scale_col = lane / 4; + int const scale_within = isV ? (within / 4) * 32 + scale_col * 4 + within % 4 : within * 8 + scale_col; + wr_kind[k] = 1; + wr_val[k] = static_cast(fp4) | (static_cast(sbyte) << 16); + wr_ptr[k] = p.kv_data + static_cast(spage) * p.d_page + (isV ? p.d_plane : 0) + within * 64 + + lane * 2; + wr_sptr[k] = p.kv_scale + static_cast(spage) * p.s_page + (isV ? p.s_plane : 0) + scale_within; + } + if (new_row[t] >= 0) + { + // The attention reads this page: its fp16 tile row is the freshly quantized row, dequantized like + // the rest of the tile. + int const h = new_row[t] / kPage; + int const row = new_row[t] % kPage; + uint32_t const sc2 = e4m3_to_h2_splat(sbyte); + uint2 hv; + hv.x = hmul2_u(e2m1x2_to_h2(fp4 & 0xFFu), sc2); + hv.y = hmul2_u(e2m1x2_to_h2(fp4 >> 8), sc2); + *reinterpret_cast( + smem + (isV ? S::kV16 : S::kK16) + h * S::kTileBytes + kmaj_off(row, lane * 4)) + = hv; + } + } + // K tiles, Q rows and patched rows -> tensor core reads. Only threads that wrote fence: the fence also + // waits for the thread's outstanding loads, and the producer warps still have their V chunks in flight. + bool wrote = tid < kComputeThreads; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + wrote |= act[k]; + } + if (wrote) + { + M3FA_FENCE_PROXY_ASYNC(); + } + M3FA_PROF(13); + __syncthreads(); +#if M3FA_TRIGGER_AT == 2 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Every K/V byte this CTA reads has arrived (the V chunks sit in registers): release the dependents now so + // their weight streams do not queue ahead of the KV reads (the megakernel releases late, from its o_proj). + if constexpr (!FQKV || M3FA_MEGA_TRIGGER == 1) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + tmem = *tmem_slot; + M3FA_PROF(2); + // ---------------- S^T = K Q^T per block on the tensor core: TMEM lane = token, column = query row ---------- +#ifndef M3FA_MEGA_WARPMMA +#define M3FA_MEGA_WARPMMA 1 // attention MMAs issued by warp 0's elected lane with warp-uniform operands (not tid 0) +#endif + if (!kRegCore && (FQKV && M3FA_MEGA_WARPMMA ? warp == 0 : tid == 0)) + { + // Operands computed by the whole warp stay in uniform registers; one elected lane issues. + uint32_t const tm_u = tmem; + bool const issue = FQKV && M3FA_MEGA_WARPMMA ? elect_one() : true; + if (issue) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + uint32_t const a0 = smem_u32(smem + S::kK16 + h * S::kTileBytes); + uint32_t const b0 = smem_u32(smem + S::kQ); +#pragma unroll + for (int ks = 0; ks < kHeadDim / 16; ++ks) + { + umma_f16(tm_u + h * kRows, umma_desc(a0 + ks * 256, 128, 2048), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 0, 0), ks > 0); + } + } + umma_commit(mma_bar); + } + if (FQKV && M3FA_MEGA_WARPMMA) + { + __syncwarp(); + } + } + if (tid >= kVStart && !(FQKV && M3FA_MEGA_EARLYV)) + { + // The producer warps dequantize V while QK and the softmax run. + dequant_v(); + } + if constexpr (kRegCore) + { + // ---------------- Register core: warp w of block h = w / kWpb owns 8 * kNt tokens for S and 8 * kNt + // output dims for O; m16n8k16 fp16 MMAs with fp32 accumulators, the softmax in the S fragments ---------- + constexpr int kWpb = kComputeWarps / BPC; // warps per block + constexpr int kNt = 16 / kWpb; // 8-token (S) or 8-dim (O) tiles per warp + float* redx_m = reinterpret_cast(smem + S::kRedX); + float* redx_l = redx_m + BPC * kRows * kWpb; + if (warp < kComputeWarps) + { + int const h = warp / kWpb; + int const wib = warp % kWpb; + int const g = lane >> 2; + int const q4 = lane & 3; + uint32_t const ktile = smem_u32(smem + S::kK16 + h * S::kTileBytes); + uint32_t const vtile = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const qrows = smem_u32(smem + S::kQ); + uint8_t* ptile = smem + S::kK16 + h * S::kTileBytes; // P over this block's K tile, after QK + auto block_bar = [&]() + { + if (h == 0) + { + asm volatile("bar.sync 14, %0;\n" ::"n"(kWpb * 32) : "memory"); + } + else + { + asm volatile("bar.sync 15, %0;\n" ::"n"(kWpb * 32) : "memory"); + } + }; + // S = Q K^T: rows (c, g / g + 8), tokens 8 tb + 2 q4 (+1). + float sv[QG][kNt][4]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + sv[c][nt][e] = 0.0F; + } + } + } +#pragma unroll + for (int ks = 0; ks < kHeadDim / 16; ++ks) + { + uint32_t bk[kNt][2]; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const tb = wib * kNt + nt; + ldsm_x2(ktile + tb * 2048 + (ks * 2 + ((lane >> 3) & 1)) * 128 + (lane & 7) * 16, bk[nt]); + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { + uint32_t aq[4]; + ldsm_x4( + qrows + (c * 2 + ((lane >> 3) & 1)) * 2048 + (ks * 2 + (lane >> 4)) * 128 + (lane & 7) * 16, + aq); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + mma_f16_16816(sv[c][nt], aq, bk[nt]); + } + } + } + M3FA_PROF(7); + // Row maxima (log2 domain, masked as the producer masks), quad then across the block's warps. + float mrow[QG][2]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { + float mx0 = kNegInf, mx1 = kNegInf; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + int const tok = cta_tok0 + h * kPage + (wib * kNt + nt) * 8 + 2 * q4 + (e & 1); + bool const valid = tok < kv_max && tok < kv_len[c]; + float const v = valid ? sv[c][nt][e] * sl2e : kNegInf; + sv[c][nt][e] = v; + if (e < 2) + { + mx0 = fmaxf(mx0, v); + } + else + { + mx1 = fmaxf(mx1, v); + } + } + } + mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 1)); + mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 2)); + mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 1)); + mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 2)); + if (q4 == 0) + { + redx_m[(h * kRows + c * 16 + g) * kWpb + wib] = mx0; + redx_m[(h * kRows + c * 16 + g + 8) * kWpb + wib] = mx1; + } + } + block_bar(); // every warp of the block is past QK (P may overwrite the K tile) and wrote its maxima + M3FA_PROF(8); +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int half = 0; half < 2; ++half) + { + float const* rm = redx_m + (h * kRows + c * 16 + g + 8 * half) * kWpb; + float m = kNegInf; +#pragma unroll + for (int w4 = 0; w4 < kWpb; w4 += 4) + { + float4 const v = *reinterpret_cast(rm + w4); + m = fmaxf(m, fmaxf(fmaxf(v.x, v.y), fmaxf(v.z, v.w))); + } + mrow[c][half] = m; + } + // P = exp2(s - m) as fp16 into [row][token] core matrices; the fp32 values feed the row sums. + float l0 = 0.0F, l1 = 0.0F; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + float a[4]; +#pragma unroll + for (int e = 0; e < 4; ++e) + { + float const m = mrow[c][e >> 1]; + a[e] = m == kNegInf ? 0.0F : ex2(sv[c][nt][e] - m); + } + l0 += a[0] + a[1]; + l1 += a[2] + a[3]; + int const t0 = (wib * kNt + nt) * 8 + 2 * q4; + int const r0 = c * 16 + g; + __half2 const p0 = __floats2half2_rn(a[0], a[1]); + __half2 const p1 = __floats2half2_rn(a[2], a[3]); + *reinterpret_cast<__half2*>( + ptile + (r0 >> 3) * 2048 + (t0 >> 3) * 128 + (r0 & 7) * 16 + (t0 & 7) * 2) + = p0; + *reinterpret_cast<__half2*>( + ptile + ((r0 + 8) >> 3) * 2048 + (t0 >> 3) * 128 + (r0 & 7) * 16 + (t0 & 7) * 2) + = p1; + } + l0 += __shfl_xor_sync(0xffffffffu, l0, 1); + l0 += __shfl_xor_sync(0xffffffffu, l0, 2); + l1 += __shfl_xor_sync(0xffffffffu, l1, 1); + l1 += __shfl_xor_sync(0xffffffffu, l1, 2); + if (q4 == 0) + { + redx_l[(h * kRows + c * 16 + g) * kWpb + wib] = l0; + redx_l[(h * kRows + c * 16 + g + 8) * kWpb + wib] = l1; + } + } + M3FA_PROF_T(19, 0); + } + __syncthreads(); // P, the V tiles (producer warps) and the row sums visible + M3FA_PROF(9); + if (warp < kComputeWarps) + { + int const h = warp / kWpb; + int const wib = warp % kWpb; + int const g = lane >> 2; + int const q4 = lane & 3; + uint32_t const vtile = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const ptile = smem_u32(smem + S::kK16 + h * S::kTileBytes); + float* obuf = reinterpret_cast(smem + S::kK16 + h * S::kTileBytes + S::kObufOff); + // O = P V: rows (c, g / g + 8), dims 8 dt + 2 q4 (+1). + float ov[QG][kNt][4]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + ov[c][nt][e] = 0.0F; + } + } + } +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + uint32_t bv[kNt][2]; + int const tok = ks * 16 + ((lane >> 3) & 1) * 8 + (lane & 7); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const dt = wib * kNt + nt; + ldsm_x2_trans(vtile + (tok >> 3) * 2048 + dt * 128 + (tok & 7) * 16, bv[nt]); + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { + uint32_t ap[4]; + ldsm_x4( + ptile + (c * 2 + ((lane >> 3) & 1)) * 2048 + (ks * 2 + (lane >> 4)) * 128 + (lane & 7) * 16, + ap); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + mma_f16_16816(ov[c][nt], ap, bv[nt]); + } + } + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const d = (wib * kNt + nt) * 8 + 2 * q4; + int const r0 = c * 16 + g; + *reinterpret_cast(obuf + r0 * kOPitch + d) = make_float2(ov[c][nt][0], ov[c][nt][1]); + *reinterpret_cast(obuf + (r0 + 8) * kOPitch + d) + = make_float2(ov[c][nt][2], ov[c][nt][3]); + } + } + } + else if (tid - kComputeThreads < BPC * kRows) + { + // Row statistics of every (block, row): LSE and the PV scale, as the producer derives them. + int const r = tid - kComputeThreads; + float const* rm = redx_m + r * kWpb; + float const* rl = redx_l + r * kWpb; + float m = kNegInf; + float l = 0.0F; +#pragma unroll + for (int w4 = 0; w4 < kWpb; w4 += 4) + { + float4 const vm = *reinterpret_cast(rm + w4); + float4 const vl = *reinterpret_cast(rl + w4); + m = fmaxf(m, fmaxf(fmaxf(vm.x, vm.y), fmaxf(vm.z, vm.w))); + l += (vl.x + vl.y) + (vl.z + vl.w); + } + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[r] = scale * vgs; + row_lse[r] = lse; + } + __syncthreads(); // O partials and row statistics visible to the push + } + else + { + if constexpr (kRs2) + { + // ---------------- Split softmax: warps 0..3 rows 0..7, warps 4..7 rows 8..15 (thread = block token) + // -------- + if (warp < 8) + { + int const rg = warp / 4; + int const wq = warp % 4; + int const tk = wq * 32 + lane; + auto rg_bar = [&]() + { + if (rg == 0) + { + asm volatile("bar.sync 10, 128;\n" ::: "memory"); + } + else + { + asm volatile("bar.sync 11, 128;\n" ::: "memory"); + } + }; + int const rs8 = ((lane >> 4) & 1) * 4 + ((lane >> 3) & 1) * 2 + ((lane >> 2) & 1); + mbar_wait_cta(mma_bar, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF(7); + float v[8]; + tmem_ld8(tmem + ((wq * 32) << 16) + 8 * rg, v); + bool const valid = cta_tok0 + tk < kv_max && cta_tok0 + tk < kv_len[0]; + float a[8]; +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] = valid ? v[j2] * sl2e : kNegInf; + a[j2] = v[j2]; + } + float const mw = warp_reduce_scatter8(a, lane); + if ((lane & 3) == 0) + { + red_m[(8 * rg + rs8) * 4 + wq] = mw; + } + rg_bar(); + M3FA_PROF(8); + uint8_t* pdst = smem + S::kP + (tk >> 3) * 128 + (tk & 7) * 16; + float mrow = kNegInf; + if (lane < 8) + { + float4 const mq = *reinterpret_cast(red_m + (8 * rg + lane) * 4); + mrow = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + } + float pa[8]; +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + float const m = __shfl_sync(0xffffffffu, mrow, j2); + pa[j2] = m == kNegInf ? 0.0F : ex2(v[j2] - m); + } + if (tid == 0) + { + M3FA_CLK(41); + } + uint32_t pw[4]; +#pragma unroll + for (int j2 = 0; j2 < 4; ++j2) + { + __half2 const hp = __floats2half2_rn(pa[2 * j2], pa[2 * j2 + 1]); + pw[j2] = *reinterpret_cast(&hp); + } + *reinterpret_cast(pdst + rg * 2048) = make_uint4(pw[0], pw[1], pw[2], pw[3]); + M3FA_PROF_T(18, 0); + if (tid == 0) + { + M3FA_CLK(42); + } + M3FA_FENCE_PROXY_ASYNC(); // P -> tensor core reads + if (tid == 0) + { + M3FA_CLK(43); + } + if (warp == 0) + { + asm volatile("bar.sync 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + if (tid == 0) + { + M3FA_CLK(44); + } + } + else + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + if (M3FA_MEGA_WARPMMA ? warp == 0 : tid == 0) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm_u = tmem; + uint32_t const a0 = smem_u32(smem + S::kV16); + uint32_t const b0 = smem_u32(smem + S::kP); + if (M3FA_MEGA_WARPMMA ? elect_one() : true) + { +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + umma_f16(tm_u + BPC * kRows, umma_desc(a0 + ks * 4096, 2048, 128), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 1, 1), ks > 0); + } + umma_commit(mma_bar); + } + if (M3FA_MEGA_WARPMMA) + { + __syncwarp(); + } + } + M3FA_PROF_T(19, 0); + float const lw = warp_reduce_scatter8(pa, lane); + if ((lane & 3) == 0) + { + red_l[(8 * rg + rs8) * 4 + wq] = lw; + } + rg_bar(); + M3FA_PROF_T(20, 0); + if (wq == 0 && lane < 8) + { + int const r = 8 * rg + lane; + float4 const mq = *reinterpret_cast(red_m + r * 4); + float4 const lq = *reinterpret_cast(red_l + r * 4); + float const m = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + float const l = (lq.x + lq.y) + (lq.z + lq.w); + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[r] = scale * vgs; + row_lse[r] = lse; + } + } + __syncthreads(); // row scales and LSEs visible (PV was issued above) + M3FA_PROF(9); + if (warp < 8) + { + mbar_wait_cta(mma_bar, 1); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + } + else + { + // ---------------- Softmax: thread = block token; the (block, 16-row chunk) pairs are spread over the + // four groups of four compute warps (each group covers the 128 TMEM lanes = block tokens) + // ---------------- + if (warp < kComputeWarps && warp / 4 < kPairs) + { + int const g = warp / 4; + int const wq = warp % 4; + int const tk = wq * 32 + lane; + auto group_bar = [&]() + { + if (g == 0) + { + asm volatile("bar.sync 10, 128;\n" ::: "memory"); + } + else if (g == 1) + { + asm volatile("bar.sync 11, 128;\n" ::: "memory"); + } + else if (g == 2) + { + asm volatile("bar.sync 12, 128;\n" ::: "memory"); + } + else + { + asm volatile("bar.sync 13, 128;\n" ::: "memory"); + } + }; + // Row r(lane) this lane holds after the reduce-scatter: bits 4..1 of the lane. + int const rs_row + = ((lane >> 4) & 1) * 8 + ((lane >> 3) & 1) * 4 + ((lane >> 2) & 1) * 2 + ((lane >> 1) & 1); + mbar_wait_cta(mma_bar, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF(7); + float v[kPairsPerGroup][16]; +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + tmem_ld16(tmem + ((wq * 32) << 16) + h * kRows + c * 16, v[k]); + bool const valid + = cta_tok0 + h * kPage + tk < kv_max && cta_tok0 + h * kPage + tk < kv_len[c]; + float a[16]; +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[k][j] = valid ? v[k][j] * sl2e : kNegInf; + a[j] = v[k][j]; + } + float const mw = warp_reduce_scatter16(a, lane); + if ((lane & 1) == 0) + { + red_m[((h * QG + c) * 16 + rs_row) * 4 + wq] = mw; + } + } + } + group_bar(); + M3FA_PROF(8); + float pa[kPairsPerGroup][16]; +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + uint8_t* pdst = smem + S::kP + h * kRows * kPage * 2 + (tk >> 3) * 128 + (tk & 7) * 16; + // Row maxima: lane j (< 16) reduces row j's four warp partials, then broadcasts it. + float mrow = kNegInf; + if (lane < 16) + { + float4 const mq + = *reinterpret_cast(red_m + ((h * QG + c) * 16 + lane) * 4); + mrow = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + } +#pragma unroll + for (int j = 0; j < 16; ++j) + { + float const m = __shfl_sync(0xffffffffu, mrow, j); + pa[k][j] = m == kNegInf ? 0.0F : ex2(v[k][j] - m); + } + if (tid == 0) + { + M3FA_CLK(41); + } + uint32_t pw[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) + { + __half2 const hp = __floats2half2_rn(pa[k][2 * j], pa[k][2 * j + 1]); + pw[j] = *reinterpret_cast(&hp); + } + // P^T as the MN-major B operand: 8 query rows of one token are one 16-byte core-matrix row. + *reinterpret_cast(pdst + (2 * c) * 2048) = make_uint4(pw[0], pw[1], pw[2], pw[3]); + *reinterpret_cast(pdst + (2 * c + 1) * 2048) + = make_uint4(pw[4], pw[5], pw[6], pw[7]); + M3FA_PROF_T(18, 0); + } + } + if (tid == 0) + { + M3FA_CLK(42); + } + M3FA_FENCE_PROXY_ASYNC(); // P -> tensor core reads + if (tid == 0) + { + M3FA_CLK(43); + } + // PV starts as soon as P and the V tiles are in: thread 0 waits for the softmax warps and the + // producer warps (V); the row sums and statistics below overlap the MMAs. + if (warp == 0) + { + asm volatile("bar.sync 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + else + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + if (tid == 0) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + uint32_t const a0 = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const b0 = smem_u32(smem + S::kP + h * kRows * kPage * 2); +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + umma_f16(tmem + (BPC + h) * kRows, umma_desc(a0 + ks * 4096, 2048, 128), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 1, 1), ks > 0); + } + } + umma_commit(mma_bar); + } + M3FA_PROF_T(19, 0); +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + float const lw = warp_reduce_scatter16(pa[k], lane); + if ((lane & 1) == 0) + { + red_l[((h * QG + c) * 16 + rs_row) * 4 + wq] = lw; + } + } + } + group_bar(); + M3FA_PROF_T(20, 0); + if (wq == 0 && lane < 16) + { +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + int const r = (h * QG + c) * 16 + lane; + float4 const mq = *reinterpret_cast(red_m + r * 4); + float4 const lq = *reinterpret_cast(red_l + r * 4); + float const m = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + float const l = (lq.x + lq.y) + (lq.z + lq.w); + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[h * kRows + c * 16 + lane] = scale * vgs; + row_lse[h * kRows + c * 16 + lane] = lse; + } + } + } + } + __syncthreads(); // row scales and LSEs visible (PV was issued above) + M3FA_PROF(9); + if (warp < kComputeWarps && warp / 4 < kPairs) + { + mbar_wait_cta(mma_bar, 1); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + } + } + M3FA_PROF(3); + } + else + { + asm volatile("cp.async.wait_all;\n" ::: "memory"); +#if M3FA_TRIGGER_AT == 2 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (!FQKV || M3FA_MEGA_TRIGGER == 1) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + for (int r = tid; r < BPC * kRows; r += kThreads) + { + row_lse[r] = kNegInf; + row_scale[r] = 0.0F; + } + __syncthreads(); + } + + // ---------------- Push every partial row to the CTA that merges its head (CTA r merges heads r*BPC..) -------- + if constexpr (S::kRecvAlias) + { + // Peers' partials overlay this CTA's tiles: every CTA's PV MMAs have completed. + asm volatile("barrier.cluster.arrive.aligned;\n" ::: "memory"); + } + if (!(FQKV && M3FA_MEGA_EARLYCLW && !S::kRecvAlias)) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); // peers' mbarriers are initialized + } + M3FA_PROF(4); + uint32_t const recv_base = smem_u32(recv); + if (kRs2 && warp < 8) + { + // Split push: warp group rg pushes rows 8 rg .. 8 rg + 7 (row j goes to CTA j, which merges head j). + int const rg = warp / 4; + int const d = (warp % 4) * 32 + lane; + float v[8]; + if (any_valid) + { + tmem_ld8(tmem + ((warp % 4) * 32 << 16) + BPC * kRows + 8 * rg, v); + } + else + { +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] = 0.0F; + } + } +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] *= row_scale[8 * rg + j2]; + } +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + uint32_t const dst_cta = 8 * rg + j2; + uint32_t const row_addr = recv_base + static_cast((cr * kOPitch + d) * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr, dst_cta)), + "r"(__float_as_uint(v[j2])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + else if (!kRs2 && warp < kComputeWarps && warp / 4 < kPairs) + { + int const g = warp / 4; + int const d = (warp % 4) * 32 + lane; // output dim +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + int const chunk = cr * BPC + h; + float v[16]; + if (any_valid && kRegCore) + { + float const* obuf + = reinterpret_cast(smem + S::kK16 + h * S::kTileBytes + S::kObufOff); +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] = obuf[(c * 16 + j) * kOPitch + d]; + } + } + else if (any_valid) + { + tmem_ld16(tmem + ((warp % 4) * 32 << 16) + (BPC + h) * kRows + c * 16, v); + } + else + { +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] = 0.0F; + } + } +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] *= row_scale[h * kRows + c * 16 + j]; + } + // One scalar store per row: a vector store's lanes would address several destination CTAs at once, + // which measured several times slower. +#pragma unroll + for (int j = 0; j < 16; ++j) + { + uint32_t const dst_cta = j / BPC; + uint32_t const row_addr + = recv_base + static_cast((((chunk * QG + c) * BPC + j % BPC) * kOPitch + d) * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr, dst_cta)), + "r"(__float_as_uint(v[j])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + } + } + else if (tid >= kComputeThreads) + { + // Every partial row's LSE (BPC * kRows of them; more than the producer threads for the widest geometries). + for (int r = tid - kComputeThreads; r < BPC * kRows; r += kThreads - kComputeThreads) + { + int const h = r / kRows; + int const c = (r % kRows) / 16; + int const j = r % 16; + uint32_t const dst_cta = j / BPC; + uint32_t const row_addr + = recv_base + static_cast((((cr * BPC + h) * QG + c) * BPC + j % BPC) * kOPitch * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr + kHeadDim * 4, dst_cta)), + "r"(__float_as_uint(row_lse[r])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + M3FA_PROF(10); +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + if (wr_kind[k] == 1) + { + *reinterpret_cast(wr_ptr[k]) = static_cast(wr_val[k] & 0xFFFFu); + if (lane % 4 == 0) + { + *wr_sptr[k] = static_cast(wr_val[k] >> 16); + } + } + else if (wr_kind[k] == 2) + { + *reinterpret_cast(wr_ptr[k]) = wr_val[k]; + } + } + { + uint32_t ready = 0; +#ifdef M3FA_DBG_HANG + long long spins = 0; +#endif + while (!ready) + { +#ifdef M3FA_DBG_HANG + if (++spins == (1LL << 24)) + { + printf("M3FA hang: block %d thread %d merge recv\n", blockIdx.x, threadIdx.x); + } +#endif + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(mbar) + : "memory"); + } + } + M3FA_PROF(5); + + // ---------------- Merge heads [cr * BPC, cr * BPC + BPC) (split-K merge kernel's LSE weighting) ---------------- + constexpr int kMergeRows = QG * BPC; // (token, local head) pairs, 16 chunk weights each + // Wide geometries: chunk weights once per merged row (into the softmax scratch, free by now), then the weighted + // sums. Narrow ones recompute the weights per output element: no extra barrier on the critical path. + constexpr bool kRowWeights = kMergeRows >= 4; + float* wmerge = red_m; // [kMergeRows][kBlocks]; red_m holds BPC * kRows * 4 = kMergeRows * 64 floats + [[maybe_unused]] auto chunk_weights = [&](int mr, float(&lc)[kBlocks]) + { + float lmax = kNegInf; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + lc[c] = recv[(c * kMergeRows + mr) * kOPitch + kHeadDim]; + lmax = fmaxf(lmax, lc[c]); + } + lmax = lmax == kNegInf ? 0.0F : lmax; + float denom = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + lc[c] = ex2(lc[c] - lmax); // now the chunk weight; 0 for an empty chunk + denom += lc[c]; + } + denom = denom > 0.0F ? denom : 1.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + // -0.0 marks an empty chunk (skipped below, as the merge kernel skips zero weights). + lc[c] = lc[c] != 0.0F ? __fdividef(lc[c], denom) : -0.0F; + } + }; + if constexpr (kRowWeights) + { + for (int mr = tid; mr < kMergeRows; mr += kThreads) + { + float lc[kBlocks]; + chunk_weights(mr, lc); +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wmerge[mr * kBlocks + c] = lc[c]; + } + } + __syncthreads(); + } +#pragma unroll + for (int idx = tid; idx < kMergeRows * kHeadDim; idx += kThreads) + { + int const mr = idx >> 7; + int const d = idx & (kHeadDim - 1); + float wn[kBlocks]; + if constexpr (kRowWeights) + { +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wn[c] = wmerge[mr * kBlocks + c]; + } + } + else + { +#ifndef M3FA_MEGA_MERGE_LOCAL +#define M3FA_MEGA_MERGE_LOCAL \ + 0 // 1 (measured +0.6 us slower): every merging thread derives the 16 chunk weights itself (broadcast shared loads, + // independent exps): no shuffle chain; the same arithmetic as chunk_weights +#endif + if constexpr (M3FA_MEGA_MERGE_LOCAL != 0) + { + chunk_weights(mr, wn); + } + else + { + // Narrow geometries: the warp (32 dims of one merged row) derives the 16 chunk weights together, lane c + // holding chunk c's, and broadcasts them. + float lc = (lane < kBlocks) ? recv[(lane * kMergeRows + mr) * kOPitch + kHeadDim] : kNegInf; + float lmax = lc; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + lmax = fmaxf(lmax, __shfl_xor_sync(0xffffffffu, lmax, o)); + } + lmax = lmax == kNegInf ? 0.0F : lmax; + float const wc = lane < kBlocks ? ex2(lc - lmax) : 0.0F; + // The denominator in chunk order from 0.0, as chunk_weights (and the merge kernel) add it: a butterfly + // reduction rounds differently in the last bit. + float denom = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + denom += __shfl_sync(0xffffffffu, wc, c); + } + denom = denom > 0.0F ? denom : 1.0F; + float const wl = wc != 0.0F ? __fdividef(wc, denom) : -0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wn[c] = __shfl_sync(0xffffffffu, wl, c); + } + } + } + if (tid == 0) + { + M3FA_CLK(45); + } + float acc = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + float const prod = recv[(c * kMergeRows + mr) * kOPitch + d] * wn[c]; + acc = __float_as_uint(wn[c]) != 0x80000000u ? acc + prod : acc; + } + if (tid == 0) + { + M3FA_CLK(46); + } + int const tl = mr / BPC; + int const head = cr * BPC + mr % BPC; + __nv_bfloat16 const ob = __float2bfloat16_rn(acc); + int const out_row = tok0 + tfirst + tl; + p.out[static_cast(out_row) * p.out_stride + head * kHeadDim + d] = ob; + if constexpr (MXFP8) + { + float const x = __bfloat162float(ob); + float amax = fabsf(x); +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, mask)); + } + float const sfv = amax * rcp_approx_ftz(448.0F); + int const bits = __float_as_int(sfv); + int ex = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0 ? 1 : 0); + ex = min(ex, 254); + float scale = __int_as_float((254 - ex) << 23); + scale = amax != 0.0F ? scale : 0.0F; + __nv_fp8_e4m3 const q8v(x * scale); + int64_t const q8off = static_cast(out_row) * (p.num_sf_k_tiles * 128) + head * kHeadDim + d; + if constexpr (FQKV) + { + // Megakernel hand-off to the o_proj units (Lamport): this call's parity buffer, the other one reset to + // 0xFF. 0xFF is never written: a NaN code is canonicalized to 0x7F; the UE8M0 code is at most 254. + uint8_t const code = q8v.__x == 0xFFu ? static_cast(0x7Fu) : q8v.__x; + int64_t const cur = (mega_calls & 1u) * kMaxTok * kOpK; + int64_t const nxt = ((mega_calls & 1u) ^ 1u) * kMaxTok * kOpK; + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.q8 + cur + q8off), "h"(static_cast(code)) + : "memory"); + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.q8 + nxt + q8off), "h"(static_cast(0xFFu)) + : "memory"); + } + else + { + p.q8[q8off] = q8v.__x; + } + if (lane == 0) + { + int const k_sf = head * (kHeadDim / 32) + (d >> 5); + int64_t const off = static_cast(out_row / 128) * p.num_sf_k_tiles * 512 + (k_sf / 4) * 512 + + (out_row % 32) * 16 + ((out_row % 128) / 32) * 4 + k_sf % 4; + if constexpr (FQKV) + { + int64_t const cur = (mega_calls & 1u) * kOpSfBytes; + int64_t const nxt = ((mega_calls & 1u) ^ 1u) * kOpSfBytes; + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.sf + cur + off), "h"(static_cast(ex)) + : "memory"); + asm volatile("st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.sf + nxt + off), + "h"(static_cast(0xFFu)) + : "memory"); + } + else + { + p.sf[off] = static_cast(ex); + } + } + } + } + M3FA_PROF(60); + if (any_valid) + { + // Every TMEM read completed above (tcgen05.wait::ld); release the columns. + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + __syncthreads(); + if (FQKV && tid == 0) + { + p.flags[32 + blockIdx.x] = mega_calls + 1u; + } + M3FA_PROF(61); + if (warp == kWarps - 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" ::"r"(*tmem_slot), "r"(S::kTmemCols) + : "memory"); + } + M3FA_PROF(6); +#else + __trap(); +#endif // MINIMAX_M3_ATTN_DEVICE_ENABLED +} + +template +cudaError_t launch_impl(Params const& prm, int num_reqs, bool pdl, cudaStream_t stream, int* max_clusters) +{ + auto kernel = m3_fused_decode_attn_kernel; + constexpr int smem = SmemX::kBytes; + constexpr int kCs = kBlocks / BPC; + static bool const configured = [&]() + { + bool ok = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem) == cudaSuccess; + ok = ok + && (kCs <= 8 + || cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) == cudaSuccess); + if (!ok) + { + cudaGetLastError(); + } + return ok; + }(); + if (!configured) + { + return cudaErrorInvalidConfiguration; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(FQKV ? kQkvGrid : kCs * (QLEN / QG) * num_reqs); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeClusterDimension; + attr[0].val.clusterDim.x = kCs; + attr[0].val.clusterDim.y = 1; + attr[0].val.clusterDim.z = 1; + attr[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[1].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + if (max_clusters != nullptr) + { + cfg.numAttrs = 1; + cudaError_t const err = cudaOccupancyMaxActiveClusters(max_clusters, kernel, &cfg); + if (err != cudaSuccess) + { + cudaGetLastError(); + } + return err; + } + cfg.numAttrs = 2; + return cudaLaunchKernelEx(&cfg, kernel, prm); +} + +// The one geometry the target launches: one request, one query token, one 16-CTA attention cluster. +cudaError_t dispatch(Params const& prm, int qlen, int qg, int bpc, int num_reqs, bool pdl, cudaStream_t stream) +{ + if (qlen == 1 && qg == 1 && bpc == 1) + { + return launch_impl<1, 1, 1, true, true>(prm, num_reqs, pdl, stream, nullptr); + } + return cudaErrorInvalidValue; +} + +static_assert(kQkvGrid == minimax_m3_attn_mega::kGrid && kMaxTok == minimax_m3_attn_mega::kMaxTokens + && kQkvN == minimax_m3_attn_mega::kQkvN && kQkvK == minimax_m3_attn_mega::kQkvK + && kOpN == minimax_m3_attn_mega::kOpN && kOpK == minimax_m3_attn_mega::kOpK + && kOpSfBytes == minimax_m3_attn_mega::kOpSfBytes, + "header dimensions"); + +} // namespace + +namespace +{ +// TMA descriptor of a [n, k] E4M3 weight (128 x 128 boxes, 128-byte swizzle), cached per weight address. +bool weight_map(CUtensorMap& map, void const* w, int k, int n) +{ + static std::mutex mu; + static std::unordered_map cache; + std::lock_guard lock(mu); + auto it = cache.find(w); + if (it != cache.end()) + { + map = it->second; + return true; + } + static PFN_cuTensorMapEncodeTiled_v12000 encode = []() + { + void* fn = nullptr; + cudaDriverEntryPointQueryResult q; + cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled", &fn, 12000, cudaEnableDefault, &q); + return reinterpret_cast(fn); + }(); + if (encode == nullptr) + { + return false; + } + cuuint64_t dims[2] = {static_cast(k), static_cast(n)}; + cuuint64_t strides[1] = {static_cast(k)}; + cuuint32_t box[2] = {128, 128}; + cuuint32_t elem[2] = {1, 1}; + CUresult const r = encode(&map, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(w), dims, strides, box, elem, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_L2_256B, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + if (r != CUDA_SUCCESS) + { + return false; + } + cache.emplace(w, map); + return true; +} +} // namespace + +cudaError_t invokeMinimaxM3AttnMega(MinimaxM3AttnMegaParams& params, void const* qkvWeight, void const* oWeight, + int qlen, int qg, int bpc, int numReqs, bool pdl, cudaStream_t stream) +{ + if (!weight_map(params.wmap, qkvWeight, kQkvK, kQkvN) || !weight_map(params.wmap_o, oWeight, kOpK, kOpN)) + { + return cudaErrorInvalidValue; + } + params.packed = params.qkv_out; + params.packed_stride = kQkvN; + params.attn_clusters = numReqs * (qlen / qg); + return dispatch(params, qlen, qg, bpc, numReqs, pdl, stream); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.h b/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.h new file mode 100644 index 000000000000..cc751911b3c0 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.h @@ -0,0 +1,112 @@ +/* + * 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 +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace minimax_m3_attn_mega +{ +constexpr int kGrid = 112; // CTAs of the persistent grid (7 clusters of 16) +constexpr int kMaxTokens = 8; // rows of each call-parity buffer +constexpr int kQkvN = 2560; // packed QKV(+index) width +constexpr int kQkvK = 6144; // hidden size +constexpr int kOpN = 6144; // o_proj output width +constexpr int kOpK = 2048; // o_proj input width (16 heads x 128) +constexpr int kOpSfBytes = 16 * 512; // one parity of the MXFP8 hand-off's swizzled scales +} // namespace minimax_m3_attn_mega + +//! Kernel parameters of the attention megakernel (the device code's Params). +struct MinimaxM3AttnMegaParams +{ + __nv_bfloat16 const* packed; + int64_t packed_stride; + uint8_t* kv_data; + int64_t d_page, d_plane; + uint8_t* kv_scale; + int64_t s_page, s_plane; + uint8_t* idx_k; + int64_t i_page, i_tok; + int64_t num_pages; + int const* out_cache_loc; + int const* position_ids; + int const* block_table; + int64_t bt_stride; + int bt_cols; + int const* seq_lens; + float const* kv_quant_scale; + float const* k_gs; + float const* v_gs; + __nv_bfloat16 const* q_w; + __nv_bfloat16 const* k_w; + __nv_bfloat16 const* ik_w; + float const* rope; + float eps; + float sm_scale; + __nv_bfloat16* out; + int64_t out_stride; + uint8_t* q8; + uint8_t* sf; + int num_sf_k_tiles; + unsigned long long* prof; // optional per-CTA clock64 probes (-DM3FA_PROFILE builds only) + // Fused QKV projection (FQKV kernels only): out[M, 2560] = mxfp8(x)[M, 6144] @ W^T. + __nv_bfloat16 const* x; // [M, 6144] layer input (the all-reduce output) + int ldx; + int num_tokens; // M + uint8_t const* wsf; // UE8M0 weight scales, 128x4-swizzled + float* part; // [2 parities][splits][8 tokens][2560] partial sums, sentinel at rest + int* counters; // [0]: call parity, [1]: CTAs done with this call (zero at rest) + int attn_clusters; // leading clusters that run the attention + uint8_t const* xq; // mega: the layer input's MXFP8 values [M, 6144] (row stride ldxq) ... + uint8_t const* xsf_act; // ... and its 128x4-swizzled UE8M0 scale atoms (atom k at byte 512 k) + int ldxq; + __nv_bfloat16* qkv_out; // mega: the packed QKV rows [2 parities][8][2560] the attention reads + uint8_t const* wsf_o; // mega: the o_proj weight's scale atoms + __nv_bfloat16* y_out; // mega: the o_proj output [M, 6144] + int* ar_ws; // mega, ar_mode: the next m3_mc_allreduce's multicast exchange (int32 words)... + int const* ar_flags; // ... its call flags (parity of a (token, column slice) slot = flags[t * 8 + slice] & 1) + int ar_rank; + int ar_world; + int ar_maxt; + int ar_mode; // 0: none; 2: push the o_proj rows into the exchange (as m3_mx_gemv's ar mode 2) + unsigned int* flags; // mega: [0, 20) tile counters (4 per call), [32 + CTA] the attention CTAs' targets + long long pf_a0, pf_b0, pf_a1, pf_b1; // mega: L2 prefetch regions issued in the attention window (bytes 0: none) + int pf_hint; // 1: evict_last policy on the prefetch + CUtensorMap wmap; // the [2560, 6144] E4M3 weight, 128 x 128 boxes, 128-byte swizzle + CUtensorMap wmap_o; // mega: the [6144, 2048] E4M3 o_proj weight, same boxes +}; + +//! Launch the megakernel on `stream`. The caller fills `params` except `wmap`, `wmap_o`, `packed`, +//! `packed_stride` and `attn_clusters`, which this derives (the TMA descriptors of the two weights, cached per +//! address). Only (qlen, qg, bpc) = (1, 1, 1) with the MXFP8 hand-off is compiled; anything else, or a failed TMA +//! descriptor encode, returns cudaErrorInvalidValue. +cudaError_t invokeMinimaxM3AttnMega(MinimaxM3AttnMegaParams& params, void const* qkvWeight, void const* oWeight, + int qlen, int qg, int bpc, int numReqs, bool pdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/draftStep.cu b/cpp/tensorrt_llm/kernels/minimaxM3/draftStep.cu new file mode 100644 index 000000000000..4df471b7a78f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/draftStep.cu @@ -0,0 +1,576 @@ +/* + * 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/minimaxM3/draftStep.h" + +#include +#include +#include + +// One pass of the fused greedy Eagle3 draft loop of MiniMax-M3 (GB300, TP4, <= 8 requests per pass). +// +// One kernel replaces the draft lm_head GEMM, the logits all-gather / float cast / argmax of the in-tree greedy draft +// sampling and the draft loop's per-pass bookkeeping, at the in-tree precision: +// +// * The vocab-sharded BF16 lm_head weight is stored in mma.m16n8k16 A-fragment order (the same values): per 16-row +// tile and k16 step, 512 bytes; lane (g, j) holds 16 bytes = {row g k 2j..2j+1, row g+8 k 2j.., row g k 2j+8.., +// row g+8 k 2j+8..}. +// * Or, by default, in a lossless 12-bit coding of those fragments (75% of the bytes; the target head's coding, see +// targetHead.cu): per 16-row tile and unit of 4 k16 steps a lane reads three 16-byte words, the sign/mantissa bytes +// of steps 0-1, of steps 2-3, then 4 code words (exponent - base, 4 bits; 15 = an exception). The decoded values +// are the checkpoint's BF16 bits; an exception decodes to +0 and the segment's flush adds the exact product x * w +// from the tile's list (products of BF16 values are exact in FP32). +// * A persistent grid (one 16-warp CTA per SM) streams whole tiles over the full K with BF16 MMAs and FP32 +// accumulation; the BF16 head input rows sit in shared memory in B-fragment order, unchanged. +// * Each logit is rounded to BF16 (the in-tree logits are the BF16 GEMM output) and becomes a 64-bit key +// ordered(value) << 32 | ~global_index, with -0 taken as +0: the larger value wins and ties pick the lowest index, as +// torch.max / torch.argmax over the gathered logits do. Keys are max-reduced per row: warp shuffles, then shared +// memory, then atomicMax. +// * The last CTA to finish pushes the rank's per-row key into every rank's exchange slot through a multicast mapping +// (Lamport: 0 = empty, a key is never 0), polls its own slots for all ranks, takes the global argmax, re-arms the +// slots, and runs the pass's bookkeeping. Call n uses slot half n % 2. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3 +{ +namespace draftStep +{ +namespace +{ +constexpr int kUnroll = 8; // k16 chunks in flight per warp +constexpr int kChunkBytes = 512; // one tile's k16 step: 32 lanes x 16 bytes + +// Coded units (4 k16 steps, 3 x 16 bytes per lane) in flight per warp: 3 up to 4 rows, 2 at 8 rows. +template +constexpr int codedUnroll() +{ + return MP >= 8 ? 2 : 3; +} + +__device__ __forceinline__ void mma_bf16( + 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)); +} + +// BF16 pair p (values 2p, 2p + 1) of a lane's 8 coded values: sign/mantissa bytes sm (4 per word), codes word cw. +__device__ __forceinline__ uint32_t coded_decode_pair(uint32_t sm_word, int byte0, uint32_t cw, int p, uint32_t base) +{ + uint32_t const b = __byte_perm(sm_word, 0u, byte0 == 0 ? 0x4140 : 0x4342); // {b0, 0, b1, 0} + uint32_t const q = (cw >> (4 * p)) & 0x000f000fu; + uint32_t const e = q * 128u + base; // exponent field of both halves + // Code 15 (an exception, sign/mantissa byte 0) decodes to +0. + return ((b & 0x007f007fu) | ((b << 8) & 0x80008000u) | e) & __vcmpne2(q, 0x000f000fu); +} + +// One coded unit (4 k16 steps) against the BF16 B fragments: steps alternate between the two accumulators, as the +// k16 steps of the plain BF16 head do. +template +__device__ __forceinline__ void chunk_mma_coded(uint4 const& sm01, uint4 const& sm23, uint4 const& cw, uint32_t base, + uint2 const* xs, int unit, int j, int g, bool b_active, float (&da)[4], float (&db)[4]) +{ + uint32_t const smw[8] = {sm01.x, sm01.y, sm01.z, sm01.w, sm23.x, sm23.y, sm23.z, sm23.w}; + uint32_t const cws[4] = {cw.x, cw.y, cw.z, cw.w}; +#pragma unroll + for (int ci = 0; ci < 4; ++ci) + { + uint2 b = make_uint2(0u, 0u); + if (b_active) + b = xs[((unit * 4 + ci) * 4 + j) * MP + g]; + uint32_t const w0 = coded_decode_pair(smw[2 * ci], 0, cws[ci], 0, base); + uint32_t const w1 = coded_decode_pair(smw[2 * ci], 2, cws[ci], 1, base); + uint32_t const w2 = coded_decode_pair(smw[2 * ci + 1], 0, cws[ci], 2, base); + uint32_t const w3 = coded_decode_pair(smw[2 * ci + 1], 2, cws[ci], 3, base); + if (ci & 1) + mma_bf16(db, w0, w1, w2, w3, b.x, b.y); + else + mma_bf16(da, w0, w1, w2, w3, b.x, b.y); + } +} + +// Exceptions of `tile` with k in [kb, ke): v += x * w (exact) for the owner lanes (row g or g + 8; tokens 2j, +// 2j + 1 < M), x from the staged B fragments. Warp-uniform (shuffles). +template +__device__ __forceinline__ void coded_exceptions( + CodedArgs const& cd, int tile, int kb, int ke, uint2 const* xs, int M, int lane, float (&v)[4]) +{ + int const g = lane >> 2, j = lane & 3; + int const e0 = __ldg(cd.exc_off + tile), e1 = __ldg(cd.exc_off + tile + 1); + for (int b0 = e0; b0 < e1; b0 += 32) + { + uint32_t pos = 0xffffffffu; + float2 wv = make_float2(0.f, 0.f); + if (b0 + lane < e1) + { + pos = __ldg(cd.exc_pos + b0 + lane); + wv = __ldg(cd.exc_w + b0 + lane); + } + int const cnt = min(32, e1 - b0); + for (int q = 0; q < cnt; ++q) + { + uint32_t const pq = __shfl_sync(0xffffffffu, pos, q); + float const w0 = __shfl_sync(0xffffffffu, wv.x, q); + int const r = static_cast(pq >> 16), k = static_cast(pq & 0xffffu); + if (k >= kb && k < ke && (r & 7) == g) + { + int const st = k >> 4, rr = k & 15; + bool const hi = r >= 8; +#pragma unroll + for (int q2 = 0; q2 < 2; ++q2) + { + int const n = 2 * j + q2; + if (n < M) + { + uint2 const bw = xs[(st * 4 + ((rr & 7) >> 1)) * MP + n]; + uint32_t const word = rr >= 8 ? bw.y : bw.x; + float const xv = __uint_as_float((rr & 1) ? (word & 0xffff0000u) : (word << 16)); + float const d = __fmul_rn(xv, w0); + if (hi) + v[2 + q2] += d; + else + v[q2] += d; + } + } + } + __syncwarp(); + } + } +} + +__device__ __forceinline__ uint4 ldg_stream(uint4 const* p) +{ + uint4 r; + asm volatile("ld.global.nc.L1::no_allocate.L2::256B.v4.u32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(p)); + return r; +} + +// The argmax value of a logit: its BF16 rounding (the in-tree logits' dtype), -0 as +0. +__device__ __forceinline__ float key_value(float v) +{ + float const r = __bfloat162float(__float2bfloat16(v)); + return r == 0.f ? 0.f : r; +} + +__device__ __forceinline__ unsigned long long make_key(float v, uint32_t index) +{ + uint32_t u = __float_as_uint(v); + u = (u & 0x80000000u) ? ~u : (u | 0x80000000u); + return (static_cast(u) << 32) | static_cast(0xFFFFFFFFu - index); +} + +__device__ __forceinline__ unsigned long long shfl_max(unsigned long long k, int o) +{ + unsigned long long const ok = __shfl_xor_sync(0xffffffffu, k, o); + return ok > k ? ok : k; +} + +// One k16 step of a tile: the lane's A fragment w against the activation B fragment of step `chunk` (rows g < MP). +// Even and odd steps accumulate separately (two independent MMA chains); the tile's sum is da + db. +template +__device__ __forceinline__ void chunk_mma( + uint4 const& w, uint2 const* xs, int chunk, int j, int g, bool b_active, float (&da)[4], float (&db)[4]) +{ + uint2 b = make_uint2(0u, 0u); + if (b_active) + b = xs[(chunk * 4 + j) * MP + g]; + if (chunk & 1) + mma_bf16(db, w.x, w.y, w.z, w.w, b.x, b.y); + else + mma_bf16(da, w.x, w.y, w.z, w.w, b.x, b.y); +} + +// Add a warp's segment of a tile into the CTA's accumulator (with a coded head, the segment's exceptions first). +template +__device__ __forceinline__ void flush_tile( + float4* acc, float (&da)[4], float (&db)[4], Params const& p, int tile, int kb, int ke, uint2 const* xs, int lane) +{ + float v[4] = {da[0] + db[0], da[1] + db[1], da[2] + db[2], da[3] + db[3]}; + if constexpr (CODED) + coded_exceptions(p.cd, tile, kb, ke, xs, p.M, lane, v); + float* a = reinterpret_cast(acc); + atomicAdd(a + 0, v[0]); + atomicAdd(a + 1, v[1]); + atomicAdd(a + 2, v[2]); + atomicAdd(a + 3, v[3]); +#pragma unroll + for (int e = 0; e < 4; ++e) + { + da[e] = 0.f; + db[e] = 0.f; + } +} + +// Balanced streaming GEMV over the fragment-ordered weight, persistent grid of 16-warp CTAs. CTA c owns tiles +// [t_begin, t_end); its (tile, chunk) space is split into 16 equal contiguous ranges, one per warp, so every warp +// streams the same number of bytes. A tile straddles at most two warps; their partial sums are added into shared +// memory (two addends: the order does not change the sum). MP = rows padded to 1 / 2 / 4 / 8. A chunk is one k16 +// step of the BF16 head, or one coded unit (4 k16 steps) of the coded head (CODED: no shared-memory prologue). +template +__global__ void __launch_bounds__(kWarps * 32, 1) + draft_step_kernel(Params const p, int kchunks, int pf_chunks, int pf_off) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ __align__(16) uint2 xs[]; // [K/16 steps][4 (j)][MP (g)], then the weight prologue + __shared__ float4 acc_s[kMaxTilesPerCta][32]; + __shared__ unsigned long long red[kWarps][kMaxTokens]; + __shared__ __align__(8) unsigned long long pf_bar[kWarps]; + __shared__ bool s_last; + __shared__ int s_tok[kMaxTokens]; + + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + int const g = lane >> 2; + int const j = lane & 3; + constexpr int kChunkK = CODED ? 64 : 16; // k per chunk + constexpr int kWords = CODED ? 3 : 1; // 16-byte words per lane and chunk + constexpr int KU = CODED ? codedUnroll() : kUnroll; + int const K = kchunks * kChunkK; + int const M = p.M; + int const t_begin = static_cast((static_cast(blockIdx.x) * p.n_tiles) / gridDim.x); + int const t_end = static_cast((static_cast(blockIdx.x + 1) * p.n_tiles) / gridDim.x); + int const n_local = t_end - t_begin; + int const span = n_local * kchunks; + int const q_begin = warp * span / kWarps; + int const q_end = (warp + 1) * span / kWarps; + long long const base = static_cast(t_begin) * kchunks; // first chunk of this CTA + uint4 const* const wf = p.wf; + + // Weights are static: start this warp's stream before waiting on the producer of x. The warp's first pf_chunks + // chunks come in by one bulk copy into shared memory (dynamic offset pf_off), which fills the grid-dependency + // wait and the staging; the register ring starts after them. + int const n_pf = CODED ? 0 : min(pf_chunks, q_end - q_begin); + uint4* const wpf = reinterpret_cast(reinterpret_cast(xs) + pf_off) + warp * pf_chunks * 32; + if (n_pf > 0 && lane == 0) + { + uint32_t const bar = static_cast(__cvta_generic_to_shared(&pf_bar[warp])); + uint32_t const bytes = static_cast(n_pf) * kChunkBytes; + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(bar) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(bar), "r"(bytes) : "memory"); + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + static_cast(__cvta_generic_to_shared(wpf))), + "l"(wf + (base + q_begin) * 32), "r"(bytes), "r"(bar) + : "memory"); + } + __syncwarp(); + int const q_reg = q_begin + n_pf; + uint4 buf[KU][kWords]; +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q_reg + u < q_end) + { +#pragma unroll + for (int l = 0; l < kWords; ++l) + buf[u][l] = ldg_stream(wf + ((base + q_reg + u) * kWords + l) * 32 + lane); + } + } + for (int i = threadIdx.x; i < n_local * 32; i += blockDim.x) + acc_s[i / 32][i % 32] = make_float4(0.f, 0.f, 0.f, 0.f); + + asm volatile("griddepcontrol.wait;" ::: "memory"); + + // Stage the head input rows in B-fragment order: item (k16 step S, row gg) = two 16-byte loads of x, four + // fragment words {x[2jj], x[2jj + 1]}, {x[2jj + 8], x[2jj + 9]} (BF16 pairs, moved unchanged). Rows >= M are 0. + constexpr int kSU = MP >= 4 ? 4 : 2; + int const n_items = (K / 16) * MP; + for (int i0 = threadIdx.x; i0 < n_items; i0 += kSU * blockDim.x) + { + uint4 xv[kSU][2]; +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + int const gg = i % MP; + int const S = i / MP; + xv[u][0] = make_uint4(0u, 0u, 0u, 0u); + xv[u][1] = make_uint4(0u, 0u, 0u, 0u); + if (i < n_items && gg < M) + { + long long const xr = p.x_rows != nullptr ? p.x_rows[gg] : gg; + uint4 const* xp = reinterpret_cast(p.x + xr * K + 16 * S); + xv[u][0] = xp[0]; + xv[u][1] = xp[1]; + } + } +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + if (i >= n_items) + continue; + int const gg = i % MP; + int const S = i / MP; + uint32_t const lo[4] = {xv[u][0].x, xv[u][0].y, xv[u][0].z, xv[u][0].w}; + uint32_t const hi[4] = {xv[u][1].x, xv[u][1].y, xv[u][1].z, xv[u][1].w}; +#pragma unroll + for (int jj = 0; jj < 4; ++jj) + xs[(S * 4 + jj) * MP + gg] = make_uint2(lo[jj], hi[jj]); + } + } + __syncthreads(); + + bool const b_active = g < MP; + float da[4] = {0.f, 0.f, 0.f, 0.f}; + float db[4] = {0.f, 0.f, 0.f, 0.f}; + int tl = q_begin / kchunks; // local tile of the current accumulator + int c = q_begin - tl * kchunks; // chunk within that tile + int cs = c; // first chunk of the warp's segment of that tile + if (n_pf > 0) + { + uint32_t const bar = static_cast(__cvta_generic_to_shared(&pf_bar[warp])); + asm volatile( + "{\n .reg .pred p;\n PF_WAIT%=:\n mbarrier.try_wait.parity.shared::cta.b64 p, [%0], 0;\n" + " @!p bra PF_WAIT%=;\n}\n" ::"r"(bar) + : "memory"); + for (int q = q_begin; q < q_reg; ++q) + { + chunk_mma(wpf[(q - q_begin) * 32 + lane], xs, c, j, g, b_active, da, db); + if (++c == kchunks) + { + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, K, xs, lane); + c = 0; + cs = 0; + ++tl; + } + } + } + for (int q = q_reg; q < q_end; q += KU) + { + uint4 nxt[KU][kWords]; +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q + KU + u < q_end) + { +#pragma unroll + for (int l = 0; l < kWords; ++l) + nxt[u][l] = ldg_stream(wf + ((base + q + KU + u) * kWords + l) * 32 + lane); + } + } +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q + u < q_end) + { + if constexpr (CODED) + chunk_mma_coded(buf[u][0], buf[u][kWords > 1 ? 1 : 0], buf[u][kWords > 2 ? 2 : 0], p.cd.base, + xs, c, j, g, b_active, da, db); + else + chunk_mma(buf[u][0], xs, c, j, g, b_active, da, db); + if (++c == kchunks) + { + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, K, xs, lane); + c = 0; + cs = 0; + ++tl; + } + } + } +#pragma unroll + for (int u = 0; u < KU; ++u) +#pragma unroll + for (int l = 0; l < kWords; ++l) + buf[u][l] = nxt[u][l]; + } + if (c != 0) + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, c * kChunkK, xs, lane); + __syncthreads(); + + // Accumulator lane (g, j) of a tile: rows g / g + 8 (x, y / z, w) of tokens 2j / 2j + 1 (x, z / y, w). + unsigned long long best0 = 0ull, best1 = 0ull; // tokens 2j, 2j + 1 + for (int t = warp; t < n_local; t += kWarps) + { + float4 const v = acc_s[t][lane]; + int const r0 = (t_begin + t) * 16 + g; + uint32_t const i0 = static_cast(p.vocab_offset + r0); + unsigned long long k0 = make_key(key_value(v.x), i0); + unsigned long long const k0b = make_key(key_value(v.z), i0 + 8); + unsigned long long k1 = make_key(key_value(v.y), i0); + unsigned long long const k1b = make_key(key_value(v.w), i0 + 8); + k0 = k0b > k0 ? k0b : k0; + k1 = k1b > k1 ? k1b : k1; + best0 = k0 > best0 ? k0 : best0; + best1 = k1 > best1 ? k1 : best1; + } +#pragma unroll + for (int o = 4; o < 32; o <<= 1) + { + best0 = shfl_max(best0, o); + best1 = shfl_max(best1, o); + } + asm volatile("griddepcontrol.launch_dependents;"); + if (lane < 4) + { + red[warp][2 * lane] = best0; + red[warp][2 * lane + 1] = best1; + } + __syncthreads(); + if (threadIdx.x < M) + { + unsigned long long b = 0ull; +#pragma unroll + for (int w = 0; w < kWarps; ++w) + b = red[w][threadIdx.x] > b ? red[w][threadIdx.x] : b; + if (b != 0ull) + atomicMax(p.keys + threadIdx.x, b); + } + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + unsigned int const prev = atomicAdd(p.counters, 1u); + s_last = prev == gridDim.x - 1; + } + __syncthreads(); + if (!s_last) + return; + + __threadfence(); + unsigned int const call = *reinterpret_cast(p.counters + 1); + int const half = static_cast(call & 1u); + if (threadIdx.x < M) + { + int const t = threadIdx.x; + unsigned long long const mine = atomicExch(p.keys + t, 0ull); + unsigned long long* slot_mc = p.ws_mc + (half * kMaxTokens + t) * p.world; + unsigned long long* slot_uc = p.ws_uc + (half * kMaxTokens + t) * p.world; + slot_mc[p.rank] = mine; + unsigned long long best = 0ull; + for (int r = 0; r < p.world; ++r) + { + unsigned long long v; + do + { + v = *reinterpret_cast(slot_uc + r); + } while (v == 0ull); + best = v > best ? v : best; + slot_uc[r] = 0ull; + } + s_tok[t] = static_cast(0xFFFFFFFFu - static_cast(best & 0xFFFFFFFFull)); + if (p.out_keys != nullptr) + p.out_keys[t] = best; + } + __syncthreads(); + if (threadIdx.x == 0) + { + p.counters[0] = 0u; + p.counters[1] = call + 1u; + } + if (threadIdx.x < M) + { + StepArgs const& st = p.st; + int const b = threadIdx.x; + int const tok = s_tok[b]; + st.draft_out[b * st.draft_stride + st.step] = tok; + st.ids_next[b] = tok; + long long const r = p.x_rows != nullptr ? p.x_rows[b] : b; + st.pos_next[b] = st.pos_src[r] + 1; + if (st.step == 0 && st.num_acc != nullptr) + { + st.kv_lens[b] -= st.draft_len - st.num_acc[b]; + st.seq_lens[b] = 1; + } + else if (st.step > 0) + { + st.kv_lens[b] += 1; + } + } +#endif +} + +template +cudaError_t launch(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + auto kern = draft_step_kernel; + size_t const xs_bytes = static_cast(K / 16) * 4 * MP * sizeof(uint2); + // Shared-memory weight prologue: chunks per warp that fit next to the kernel's static + dynamic shared memory. + static int const sStaticSmem = [] + { + cudaFuncAttributes fa{}; + return cudaFuncGetAttributes(&fa, draft_step_kernel) == cudaSuccess + ? static_cast(fa.sharedSizeBytes) + : -1; + }(); + static int const sMaxSmemPerBlock = [] + { + int dev = 0, maxSmem = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&maxSmem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + return maxSmem; + }(); + int pf = 0; + if (!CODED && sStaticSmem >= 0) + { + long long const room + = static_cast(sMaxSmemPerBlock) - sStaticSmem - static_cast(xs_bytes) - 1024; + long long const want = static_cast(kPrefetchKb) * 1024; + pf = static_cast(std::max(0ll, std::min(room, want)) / (kWarps * kChunkBytes)); + } + size_t const smem = xs_bytes + static_cast(pf) * kWarps * kChunkBytes; + cudaError_t err = cudaFuncSetAttribute(kern, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem)); + if (err != cudaSuccess) + return err; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(static_cast(grid)); + cfg.blockDim = dim3(kWarps * 32); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, kern, p, K / (CODED ? 64 : 16), pf, static_cast(xs_bytes)); +} + +template +cudaError_t dispatchRows(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + if (p.M <= 1) + return launch<1, CODED>(p, K, grid, pdl, stream); + if (p.M <= 2) + return launch<2, CODED>(p, K, grid, pdl, stream); + if (p.M <= 4) + return launch<4, CODED>(p, K, grid, pdl, stream); + return launch<8, CODED>(p, K, grid, pdl, stream); +} + +} // namespace + +cudaError_t invokeDraftStep(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + bool const coded = p.cd.exc_off != nullptr; + if (p.M < 1 || p.M > kMaxTokens || K % (coded ? 64 : 16) != 0 || grid < 1) + return cudaErrorInvalidValue; + return coded ? dispatchRows(p, K, grid, pdl, stream) : dispatchRows(p, K, grid, pdl, stream); +} + +} // namespace draftStep +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/draftStep.h b/cpp/tensorrt_llm/kernels/minimaxM3/draftStep.h new file mode 100644 index 000000000000..d6f4c1dd67d9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/draftStep.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 +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3 +{ +//! One pass of the fused greedy Eagle3 draft loop of the MiniMax-M3 NVFP4 TP4 decode path (sm_103a): the BF16 +//! vocab-sharded draft head, the cross-rank argmax and the draft loop's per-pass bookkeeping in one kernel. See +//! draftStep.cu. +namespace draftStep +{ + +constexpr int kMaxTokens = 8; // head rows (requests) of one launch +constexpr int kWarps = 16; // warps per CTA +constexpr int kMaxTilesPerCta = 32; // 16-row weight tiles one CTA can own +//! Shared-memory weight prologue per CTA, in KB (capped by the free shared memory). +constexpr int kPrefetchKb = 112; + +// Draft-pass epilogue (last CTA, one thread per request row b): stores the token as draft token `step`, stages the +// next pass's input id and position, and advances the draft KV lengths the way the in-tree draft loop does between +// passes (pass 0 with num_acc: drop the rejected verify rows and switch to one token per request; later passes: +1). +struct StepArgs +{ + int* draft_out; // [nreq, draft_stride]: column `step` receives the token + int draft_stride; + int step; + int* ids_next; // [nreq] + int const* pos_src; // positions of the rows the head read (indexed like x) + int* pos_next; // [nreq] = pos_src[x_rows[b] or b] + 1 + int* kv_lens; // [nreq] draft KV lengths, updated in place + int const* num_acc; // pass 0: accepted counts (nullptr: no rollback) + int* seq_lens; // pass 0: set to 1 + int draft_len; +}; + +// The 12-bit coded head (see draftStep.cu): the per-tile exception lists and the exponent base. exc_off == nullptr: +// the head is the plain BF16 fragment-ordered copy. +struct CodedArgs +{ + int const* exc_off; // [n_tiles + 1] exception list offsets + uint32_t const* exc_pos; // row in tile << 16 | k, per exception + float2 const* exc_w; // {w, 0} per exception: the exact weight + uint32_t base; // exponent base << 7 in both BF16 halves (base << 7 | base << 23) +}; + +// Everything a draft-step launch takes (see invokeDraftStep). +struct Params +{ + __nv_bfloat16 const* x; // [*, K] head input rows + long long const* x_rows; // optional [M]: head input row of request b (pass 0: the last accepted verify row) + int M; // requests, <= kMaxTokens + uint4 const* wf; // [n_tiles][K / 16][32 lanes] 16-byte BF16 A fragments, or the 12-bit coded stream + // [n_tiles][K / 64][3][32 lanes] when cd.exc_off is set (see draftStep.cu) + int n_tiles; + int vocab_offset; // global id of this rank's first head row + unsigned long long* keys; + unsigned int* counters; + unsigned long long* ws_uc; // this rank's exchange slots [2][kMaxTokens][world] + unsigned long long* ws_mc; // their multicast address + int rank; + int world; + unsigned long long* out_keys; // optional [M]: the winning global key of each row + StepArgs st; + CodedArgs cd; +}; + +//! Launch the draft step over `grid` 512-thread CTAs (K = the hidden size, a multiple of 16; of 64 for a coded head). +cudaError_t invokeDraftStep(Params const& p, int K, int grid, bool pdl, cudaStream_t stream); + +} // namespace draftStep +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.cu b/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.cu new file mode 100644 index 000000000000..4383bb0fb2e6 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.cu @@ -0,0 +1,3679 @@ +/* + * 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. + */ + +// Fused short-sequence decode attention of the MiniMax-M3 MSA sparse layers (packed QKV input). + +#include "tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +// The kernel uses tcgen05/TMEM (including block-scaled MMAs, which family targets lack), cluster DSMEM and E2M1 +// conversions: its body is compiled for the arch-specific sm_100a / sm_103a targets only (see the minimax_m3_attn_src +// library in kernels/CMakeLists.txt). Any other target compiles a kernel that traps, so a launch on a build without +// those targets fails loudly instead of returning untouched outputs. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030) \ + && (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FEAT_SM100_ALL) || defined(__CUDA_ARCH_FEAT_SM103_ALL)) +#define MINIMAX_M3_ATTN_DEVICE_ENABLED 1 +#else +#define MINIMAX_M3_ATTN_DEVICE_ENABLED 0 +#endif + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kHeadDim = 128; +constexpr int kRotDim = 64; +constexpr int kPage = 128; +constexpr int kHq = 16; +constexpr int kColK = kHq; // packed head slots: [Q x16 | K | V | idxQ | idxK] +constexpr int kColV = kHq + 1; +constexpr int kColIdxK = kHq + 3; +constexpr int kBlocks = 16; // KV blocks per request == topk == cluster size +constexpr int kComputeWarps = 16; // QK / softmax / PV and the dequant (one 16-byte chunk per thread) +constexpr int kWarps = kComputeWarps + 4; // plus producer-only warps, so new K/V/index-K never queue behind a Q head +constexpr int kThreads = kWarps * 32; +constexpr int kComputeThreads = kComputeWarps * 32; +constexpr int kOPitch = 132; // floats per partial-O row +constexpr float kNegInf = -__builtin_huge_valf(); + +using Params = MinimaxM3FusedDecodeAttnParams; + +// Optional per-CTA clock64 probes for development (-DM3FA_PROFILE and a prof tensor): [CTA][64] timestamps. +// -DM3FA_PROFILE_GT records %globaltimer (ns, comparable across SMs and kernels) instead of clock64. +#ifdef M3FA_PROFILE_GT +#define M3FA_PROFILE +#define M3FA_NOW(v) asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(v)) +#else +#define M3FA_NOW(v) v = clock64() +#endif +#ifdef M3FA_PROFILE +#define M3FA_PROF_T(i, thread) \ + do \ + { \ + if (p.prof != nullptr && threadIdx.x == (thread)) \ + { \ + unsigned long long now_; \ + M3FA_NOW(now_); \ + p.prof[blockIdx.x * 64 + (i)] = now_; \ + } \ + } while (0) +#else +#define M3FA_PROF_T(i, thread) \ + do \ + { \ + } while (0) +#endif +#define M3FA_PROF(i) M3FA_PROF_T(i, 0) +// clock64 probes (development): slot i of this CTA, written by every calling thread. +#ifdef M3FA_PROFILE +#define M3FA_CLK(i) \ + do \ + { \ + if (p.prof != nullptr) \ + { \ + p.prof[blockIdx.x * 64 + (i)] = clock64(); \ + } \ + } while (0) +#else +#define M3FA_CLK(i) \ + do \ + { \ + } while (0) +#endif + +#ifndef M3FA_MEGA_EARLYCLW +#define M3FA_MEGA_EARLYCLW 1 // pass the cluster barrier (peers' merge mbarriers initialized) early, not before the push +#endif +#ifndef M3FA_V15_EARLYV +#define M3FA_V15_EARLYV 0 // packed path: dequantize V before the PDL wait (switch; off by default) +#endif +#ifndef M3FA_MEGA_TRIGGER +// Megakernel PDL release: 0 o_proj CTAs after their activation (attention CTAs at exit); 1 as 0 with the attention +// CTAs after their KV reads; 2 every CTA right after its griddepcontrol.wait; 3 o_proj CTAs after their epilogue. +#define M3FA_MEGA_TRIGGER 3 +#endif +#ifndef M3FA_TRIGGER_AT +#define M3FA_TRIGGER_AT 2 // where the PDL dependents are released: 0 entry, 1 after the wait, 2 after the KV reads +#endif + +// Generic-proxy shared-memory writes (fp16 tiles, Q, P) become visible to the tensor core (async proxy). +#define M3FA_FENCE_PROXY_ASYNC() asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory") + +__device__ __forceinline__ void cp_async16(void* dst, void const* src) +{ + uint32_t const d = static_cast(__cvta_generic_to_shared(dst)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(d), "l"(src) : "memory"); +} + +__device__ __forceinline__ float ex2(float x) +{ + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ float lg2(float x) +{ + float y; + asm("lg2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +// One packed E2M1 byte -> half2 (low nibble in the low half). +__device__ __forceinline__ uint32_t e2m1x2_to_h2(uint32_t byte) +{ + uint32_t out; + asm("{\n .reg .b8 src;\n mov.b32 {src, _, _, _}, %1;\n cvt.rn.f16x2.e2m1x2 %0, src;\n}\n" : "=r"(out) : "r"(byte)); + return out; +} + +__device__ __forceinline__ uint32_t hmul2_u(uint32_t a, uint32_t b) +{ + __half2 r = __hmul2(*reinterpret_cast<__half2*>(&a), *reinterpret_cast<__half2*>(&b)); + return *reinterpret_cast(&r); +} + +__device__ __forceinline__ uint32_t e4m3_to_h2_splat(uint32_t byte) +{ + __half_raw const h = __nv_cvt_fp8_to_halfraw(static_cast<__nv_fp8_storage_t>(byte), __NV_E4M3); + return static_cast(h.x) | (static_cast(h.x) << 16); +} + +__device__ __forceinline__ uint32_t e4m3x2_to_h2(uint32_t two_bytes) +{ + __half2_raw const h = __nv_cvt_fp8x2_to_halfraw2(static_cast<__nv_fp8x2_storage_t>(two_bytes), __NV_E4M3); + return static_cast(h.x) | (static_cast(h.y) << 16); +} + +// ---- Exact copies of the C++ producer's helpers (fusedQKNormRopeKernel.cu / quantization.cuh) ---- +template +__device__ __forceinline__ uint32_t cvt_sf_group_shfl_mask() +{ + constexpr uint32_t groupSize = static_cast(NUM_THREADS_PER_SF); + constexpr uint32_t groupMask = 0xFFFFFFFFU >> (32U - groupSize); + uint32_t laneId = 0; + asm("mov.u32 %0, %%laneid;" : "=r"(laneId)); + return groupMask << (laneId & ~(groupSize - 1U)); +} + +__device__ __forceinline__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) +{ + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(val) + : "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]), "f"(array[4]), "f"(array[5]), "f"(array[6]), + "f"(array[7])); + return val; +} + +__device__ __forceinline__ uint16_t quantizeMinimaxM3Fp4x4( + float const (&elements)[4], float secondLevelScale, uint8_t* scaleOut) +{ + __nv_bfloat16 rounded[4]; + float localMax = 0.0F; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + rounded[i] = __float2bfloat16_rn(elements[i]); + localMax = fmaxf(localMax, fabsf(__bfloat162float(rounded[i]))); + } + uint32_t const groupMask = cvt_sf_group_shfl_mask<4>(); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 1)); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 2)); + float const scaleValue = secondLevelScale * (localMax * rcp_approx_ftz(6.0F)); + __nv_fp8_e4m3 const encodedScale(scaleValue); + if (scaleOut != nullptr) + { + *scaleOut = encodedScale.__x; + } + float const narrowedScale = static_cast(encodedScale); + float const outputScale + = localMax != 0.0F ? rcp_approx_ftz(narrowedScale * rcp_approx_ftz(secondLevelScale)) : 0.0F; + float quantValues[8] = {}; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + quantValues[i] = __bfloat162float(rounded[i]) * outputScale; + } + return static_cast(fp32_vec_to_e2m1(quantValues)); +} + +// RMSNorm (gemma 1+w) + NeoX RoPE over the first 64 dims, exactly as the producer: lane owns 4 dims. +__device__ __forceinline__ void norm_rope( + float (&e)[4], float const (&w)[4], float eps, float const (&cs)[4], float const (&sn)[4], int lane) +{ + float ss = 0.0F; + ss += e[0] * e[0] + e[1] * e[1]; + ss += e[2] * e[2] + e[3] * e[3]; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + ss = ss + __shfl_xor_sync(0xffffffffu, ss, mask, 32); + } + float const rmsReciprocal = rsqrtf(ss / static_cast(kHeadDim) + eps); +#pragma unroll + for (int i = 0; i < 4; ++i) + { + e[i] *= rmsReciprocal * (1.0F + w[i]); + } + __syncwarp(); + constexpr int kPairOffset = (kRotDim / 2) / 4; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const dim = lane * 4 + i; + float paired = __shfl_xor_sync(0xffffffffu, e[i], kPairOffset); + if (dim < kRotDim) + { + bool const firstHalf = dim < kRotDim / 2; + if (firstHalf) + { + paired = -paired; + } + e[i] = e[i] * cs[i] + paired * sn[i]; + } + } + __syncwarp(); +} + +__device__ __forceinline__ void load_w4(float (&w)[4], __nv_bfloat16 const* src, int lane) +{ + uint2 const raw = *reinterpret_cast(src + lane * 4); + float2 const a = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw.x)); + float2 const b = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw.y)); + w[0] = a.x; + w[1] = a.y; + w[2] = b.x; + w[3] = b.y; +} + +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ uint32_t mapa_shared(uint32_t addr, uint32_t rank) +{ + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n" : "=r"(r) : "r"(addr), "r"(rank)); + return r; +} + +// QLEN query tokens per request; QG of them per cluster (QLEN / QG clusters per request); BPC KV blocks per CTA +// (cluster of kBlocks / BPC CTAs). +// Byte offset of element (r, k) in a UMMA no-swizzle K-major operand tile of 128 k-elements per row: 8x8 core +// matrices of 16-byte rows, 128 bytes apart along k (LBO) and 2048 bytes apart per 8 rows (SBO). The same formula +// with r = token, k = dim stores V so that it is the MN-major (dims contiguous) A operand of O^T = V^T P^T. +__device__ __forceinline__ int kmaj_off(int r, int k) +{ + return (r >> 3) * 2048 + (k >> 3) * 128 + (r & 7) * 16 + (k & 7) * 2; +} + +__device__ __forceinline__ uint64_t umma_desc(uint32_t addr, uint32_t lbo, uint32_t sbo) +{ + return static_cast((addr >> 4) & 0x3FFFu) | (static_cast((lbo >> 4) & 0x3FFFu) << 16) + | (static_cast((sbo >> 4) & 0x3FFFu) << 32) | (1ull << 46); +} + +// kind::f16 instruction descriptor: F32 accumulate, F16 operands, B K-major, M = 128. +__host__ __device__ constexpr uint32_t umma_idesc(int n, int a_mn_major, int b_mn_major) +{ + return (1u << 4) | (static_cast(a_mn_major) << 15) | (static_cast(b_mn_major) << 16) + | (static_cast(n >> 3) << 17) | (static_cast(128 >> 4) << 24); +} + +// 8-row version for the split softmax: the same lane-bit order (16, 8, 4, then 2 and 1 in full) as +// warp_reduce_scatter16, so every row is combined in the same association (fp add is commutative, not associative). +// After it lanes with bits 4..2 = (b4, b3, b2) hold row 4*b4 + 2*b3 + b2. +template +__device__ __forceinline__ float warp_reduce_scatter8(float (&a)[8], int lane) +{ +#pragma unroll + for (int step = 0; step < 3; ++step) + { + int const mask = 16 >> step; + int const half = 4 >> step; + bool const upper = (lane & mask) != 0; +#pragma unroll + for (int i = 0; i < half; ++i) + { + float const send = upper ? a[i] : a[i + half]; + float const keep = upper ? a[i + half] : a[i]; + float const recv = __shfl_xor_sync(0xffffffffu, send, mask); + a[i] = kMax ? fmaxf(keep, recv) : keep + recv; + } + } + float x = a[0]; + float const o2 = __shfl_xor_sync(0xffffffffu, x, 2); + x = kMax ? fmaxf(x, o2) : x + o2; + float const o1 = __shfl_xor_sync(0xffffffffu, x, 1); + return kMax ? fmaxf(x, o1) : x + o1; +} + +// Reduce 16 per-lane values over the warp, one value per row: after the reduce-scatter lane l holds row +// 8*b4 + 4*b3 + 2*b2 + b1 (b_k = bit k of l) reduced over all 32 lanes. 16 shuffles instead of 80. +template +__device__ __forceinline__ float warp_reduce_scatter16(float (&a)[16], int lane) +{ +#pragma unroll + for (int step = 0; step < 4; ++step) + { + int const mask = 16 >> step; + int const half = 8 >> step; + bool const upper = (lane & mask) != 0; +#pragma unroll + for (int i = 0; i < half; ++i) + { + float const send = upper ? a[i] : a[i + half]; + float const keep = upper ? a[i + half] : a[i]; + float const recv = __shfl_xor_sync(0xffffffffu, send, mask); + a[i] = kMax ? fmaxf(keep, recv) : keep + recv; + } + } + float const other = __shfl_xor_sync(0xffffffffu, a[0], 1); + return kMax ? fmaxf(a[0], other) : a[0] + other; +} + +__device__ __forceinline__ void umma_f16( + uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %4, 0;\n tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(accumulate) + : "memory"); +} + +__device__ __forceinline__ void umma_commit(uint32_t bar) +{ + asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];\n" ::"r"(bar) : "memory"); +} + +__device__ __forceinline__ void mbar_wait_cta(uint32_t bar, uint32_t parity) +{ + uint32_t ready = 0; +#ifdef M3FA_DBG_HANG + long long spins = 0; +#endif + while (!ready) + { +#ifdef M3FA_DBG_HANG + if (++spins == (1LL << 24)) + { + printf("M3FA hang: block %d thread %d mbar 0x%x parity %u\n", blockIdx.x, threadIdx.x, bar, parity); + } +#endif + 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"(ready) + : "r"(bar), "r"(parity) + : "memory"); + } +} + +// Spin without suspending (v52's acc_done wait): a suspended try_wait wakes late on a tcgen05.commit arrival. +#ifndef M3FA_EPI_SPIN +#define M3FA_EPI_SPIN 1 +#endif +__device__ __forceinline__ void mbar_wait_epi(uint32_t bar, uint32_t parity) +{ + if (!M3FA_EPI_SPIN) + { + mbar_wait_cta(bar, parity); + return; + } + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(bar), "r"(parity) + : "memory"); + } +} + +__device__ __forceinline__ void tmem_ld8(uint32_t addr, float (&v)[8]) +{ + uint32_t r[8]; + asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n" + : "=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"(addr)); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < 8; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +__device__ __forceinline__ void tmem_ld16(uint32_t addr, 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];\n" + : "=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"(addr)); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < 16; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +// Warp-level fp16 MMA (m16n8k16, fp32 accumulate) and its shared-memory fragment loads. +__device__ __forceinline__ void ldsm_x4(uint32_t addr, uint32_t (&r)[4]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x2(uint32_t addr, uint32_t (&r)[2]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" : "=r"(r[0]), "=r"(r[1]) : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x2_trans(uint32_t addr, uint32_t (&r)[2]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0, %1}, [%2];\n" + : "=r"(r[0]), "=r"(r[1]) + : "r"(addr)); +} + +__device__ __forceinline__ void mma_f16_16816(float (&d)[4], uint32_t const (&a)[4], uint32_t const (&b)[2]) +{ + 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"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +// Register attention core (warp MMAs, no TMEM) for up to two tokens per cluster; the tcgen05 core above that. +// Off by default: on GB300 it measured ~0.25 us slower than the tcgen05 core at M=1 (GEMV + attention chain). +#ifndef M3FA_REG_CORE +#define M3FA_REG_CORE 0 +#endif +template +struct RegCore +{ + static constexpr bool kOn = M3FA_REG_CORE != 0 && QG <= 2; +}; + +template +struct Smem +{ + static constexpr int kRows = 16 * QG; // query rows (token, head) of this cluster = UMMA N + static constexpr int kTok = BPC * kPage; // block tokens per CTA + static constexpr int kTileBytes = kPage * kHeadDim * 2; // one fp16 K or V block tile in UMMA layout + static constexpr int kK16 = 0; + static constexpr int kV16 = kK16 + BPC * kTileBytes; + static constexpr int kQ = kV16 + BPC * kTileBytes; // Q rows, UMMA K-major [row][dim] + // P^T per block, UMMA MN-major [token][row], over the block's K tile: P is written after the QK MMAs (which + // read every K tile) have completed. + static constexpr int kP = kK16; + // Packed V tiles as copied from the cache: [block][128 rows x 64 data bytes | 1 KB of 4x4-quad scales]. + static constexpr int kRawV = kQ + kRows * kHeadDim * 2; + static constexpr int kRawVBlock = kPage * 64 + kPage * 8; + static constexpr int kRowScale = kRawV + BPC * kRawVBlock; + static constexpr int kRowLse = kRowScale + BPC * kRows * 4; + // Per-warp softmax partials [block][row][4 warps]: max, then sum. + static constexpr int kRedM = (kRowLse + BPC * kRows * 4 + 15) / 16 * 16; + static constexpr int kRedL = kRedM + BPC * kRows * 16; + // Partials pushed by every CTA of the cluster: [chunk][token][local head][kOPitch] floats, lse at column 128. + // When they do not fit beside the tiles they overlay the K/V tiles (and P, Q): peers then push only once every + // CTA of the cluster has completed its PV MMAs (a cluster barrier right before the push). + // Register core: per-warp row maxima and sums [block][row][warp of the block], two arrays. + static constexpr int kRedX = kRedL + BPC * kRows * 16; + static constexpr int kRedXBytes = RegCore::kOn ? 2 * BPC * kRows * (16 / BPC) * 4 : 0; + static constexpr int kStatsEnd = kRedX + kRedXBytes; + // Register core: P (fp16, [row][token] core matrices) at the start of its block's K tile, the scaled-free O + // partial ([row][kOPitch] fp32) behind it. + static constexpr int kObufOff = 8192; + static_assert(!RegCore::kOn || (kRows * kPage * 2 <= kObufOff && kObufOff + kRows * kOPitch * 4 <= kTileBytes), + "register-core P and O staging fit in a K tile"); + static constexpr int kRecvBytes = kBlocks * QG * BPC * kOPitch * 4; + static constexpr bool kRecvAlias = kStatsEnd + kRecvBytes + 32 > 227 * 1024; + static constexpr int kRecv = kRecvAlias ? 0 : kStatsEnd; + static constexpr int kMbar = kRecvAlias ? kStatsEnd : kStatsEnd + kRecvBytes; // merge, MMA mbarriers, TMEM addr + static constexpr int kBytes = kMbar + 32; + static constexpr int kTmemCols = BPC * 2 * kRows <= 32 ? 32 + : BPC * 2 * kRows <= 64 ? 64 + : BPC * 2 * kRows <= 128 ? 128 + : BPC * 2 * kRows <= 256 ? 256 + : 512; + static_assert(BPC * 2 * kRows <= 512, "TMEM columns"); + static_assert(!kRecvAlias || kRecvBytes <= kRawV, "aliased partials must fit in the tiles, Q and P"); + static_assert(kRows * kPage * 2 <= kTileBytes, "P^T fits in the K tiles"); + static_assert(kBytes <= 227 * 1024, "shared memory budget"); + static_assert(QLEN % QG == 0 && (BPC == 1 || BPC == 2), "unsupported geometry"); +}; + +// The packed NVFP4 bytes of one 16-byte chunk (32 dims of one block row) and its two block scales. +struct RawChunk +{ + uint4 data; + uint32_t scales; // byte 0: dims [0, 16), byte 1: dims [16, 32) of the chunk +}; + +// Chunk q of a block tile: lanes 8j..8j+7 of a warp take 8 consecutive rows of one 32-dim chunk, so every 16-byte +// fp16 store phase writes 8 distinct core-matrix rows (conflict-free in the UMMA layout). +__device__ __forceinline__ int chunk_row(int q) +{ + return (q >> 5) * 8 + (q & 7); +} + +__device__ __forceinline__ int chunk_col(int q) +{ + return (q >> 3) & 3; +} + +__device__ __forceinline__ RawChunk load_chunk( + uint8_t const* data_page, uint8_t const* scale_page, int row, int cc, bool is_v) +{ + RawChunk c; + c.data = *reinterpret_cast(data_page + row * 64 + cc * 16); + if (is_v) + { + // V scales: 4x4 token-quad order [row / 4][scale column][row % 4]. + uint8_t const* sc = scale_page + (row >> 2) * 32 + (row & 3); + c.scales = static_cast(sc[(cc * 2) * 4]) | (static_cast(sc[(cc * 2 + 1) * 4]) << 8); + } + else + { + c.scales = *reinterpret_cast(scale_page + row * 8 + cc * 2); + } + return c; +} + +__device__ __forceinline__ void store_chunk_fp16(uint8_t* tile, int row, int cc, RawChunk const& c, bool valid) +{ + uint32_t const s0 = e4m3_to_h2_splat(c.scales & 0xFFu); + uint32_t const s1 = e4m3_to_h2_splat((c.scales >> 8) & 0xFFu); + uint32_t const wd[4] = {c.data.x, c.data.y, c.data.z, c.data.w}; + uint32_t out[16]; +#pragma unroll + for (int i = 0; i < 16; ++i) + { + uint32_t const b = (wd[i >> 2] >> (8 * (i & 3))) & 0xFFu; + out[i] = valid ? hmul2_u(e2m1x2_to_h2(b), i < 8 ? s0 : s1) : 0u; + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + *reinterpret_cast(tile + kmaj_off(row, cc * 32 + j * 8)) + = make_uint4(out[4 * j], out[4 * j + 1], out[4 * j + 2], out[4 * j + 3]); + } +} + +// ---------------- Fused QKV projection: block-scaled MXFP8 on tcgen05 (N=2560, K=6144) ---------------- +// out[M, 2560] = mxfp8(x)[M, 6144] @ W^T, the computation of the CUTLASS MXFP8 GEMM: W is the A operand (128 output +// rows per MMA), the quantized activation the B operand (8 token columns), FP32 accumulation in TMEM. A unit is one +// (128-row tile, K split) pair; the 20 x kQkvSplits units run on the last CTAs of the grid. Before the PDL wait a +// CTA streams its unit's weight tiles and scale atoms through a shared-memory ring into TMEM; after it, it quantizes +// its x slice (mxfp8_quantize math), runs the MMAs and stores its FP32 partial rows. +// +// Hand-off without a flag: the partial buffers are double buffered by call parity and hold a sentinel (an FP32 NaN +// the MMA never produces) at rest, so a consumer polls the partials themselves. The producers of call n reset the +// other parity's buffers after their PDL wait; call n - 1, the last reader of those, is complete by then. +constexpr int kQkvN = 2560; +constexpr int kQkvK = 6144; +constexpr int kQkvKTiles = kQkvK / 128; +constexpr int kQkvRowTiles = kQkvN / 128; +#ifndef M3FA_QKV_SPLITS +#define M3FA_QKV_SPLITS 4 +#endif +#ifndef M3FA_QKV_RING +#define M3FA_QKV_RING kQkvMaxTiles // every tile of the unit in flight at once +#endif +#ifndef M3FA_QKV_NACC +#define M3FA_QKV_NACC 1 +#endif +#ifndef M3FA_QKV_SCALED_MMA +#define M3FA_QKV_SCALED_MMA 1 // 0: unscaled per-block MMAs with the scales applied in the epilogue (see below) +#endif +#ifndef M3FA_QKV_PIPE +// 1: the weight TMA is issued first at entry; the activation has its own region, copied right after the PDL wait; +// each weight tile's MMAs issue as soon as it lands (v52's order). 0: tiles into TMEM first, activation over the ring. +#define M3FA_QKV_PIPE 1 +#endif +#ifndef M3FA_QKV_A_SMEM +#define M3FA_QKV_A_SMEM 1 // scaled MMAs only; 1: they read the weight tiles from shared memory (no copy into TMEM) +#endif +// Default: block-scaled kind::mxf8f6f4 MMAs reading A (the weight tiles) in place from the ring, issued per tile as +// it lands (M3FA_QKV_PIPE), one accumulator; M3FA_QKV_A_SMEM=0 copies each tile into TMEM first. Alternative +// (M3FA_QKV_SCALED_MMA=0): one unscaled kind::f8f6f4 MMA per 32-wide K block into its own accumulator, A read in place +// from shared memory, the epilogue applying both UE8M0 scales. On GB300 an M128 x N8 x K32 MMA occupies the tensor pipe +// ~105 cycles block-scaled and ~60 unscaled, independent of N up to 64; the unscaled form's epilogue (320 accumulator +// columns) costs more than it saves. +constexpr bool kQkvScaled = M3FA_QKV_SCALED_MMA != 0; +constexpr bool kQkvASmem = !kQkvScaled || M3FA_QKV_A_SMEM != 0; +constexpr int kQkvSplits = M3FA_QKV_SPLITS; +constexpr int kQkvUnits = kQkvRowTiles * kQkvSplits; +constexpr int kQkvMaxTiles = (kQkvKTiles + kQkvSplits - 1) / kQkvSplits; +constexpr int kQkvRing = M3FA_QKV_RING; +constexpr int kQkvNacc = M3FA_QKV_NACC; +constexpr int kQkvGrid = 112; // 7 clusters of 16 (or 14 of 8): one wave with the attention +constexpr int kMaxTok = 8; +constexpr int kQkvWTile = 128 * 128; // E4M3 bytes of one 128 x 128 weight tile +constexpr int kQkvSfAtom = 512; // one 128-row x 4-block UE8M0 scale atom +constexpr uint32_t kQkvSentinel = 0xFFFFFFFFu; +constexpr int64_t kQkvPartStride = static_cast(kQkvSplits) * kMaxTok * kQkvN; // one parity +// TMEM columns of the one 512-column allocation of a CTA: attention [0, 64), then the GEMM's accumulators, scale +// atoms (A per tile, B per tile) and the A tiles (32 columns of 128 x 128 E4M3 each). +constexpr int kTmemAll = 512; +// Megakernel: the attention CTAs run no projection unit, so the GEMM starts at column 0 of a projection CTA. +constexpr int kTAcc = 0; +constexpr int kTSfa = kTAcc + 16; +constexpr int kTSfb = kTSfa + 4 * kQkvMaxTiles; +constexpr int kTA = (kTSfb + 4 * kQkvMaxTiles + 31) / 32 * 32; +static_assert(!kQkvScaled || kTA + 32 * kQkvMaxTiles <= kTmemAll, "TMEM columns"); +static_assert(kQkvScaled || kTAcc + 32 * kQkvMaxTiles <= kTmemAll, "TMEM columns (one accumulator per K block)"); +static_assert(!kQkvASmem || kQkvRing == kQkvMaxTiles, "shared-memory A needs every tile resident"); +static_assert(kQkvUnits <= kQkvGrid, "at most one unit per CTA"); +static_assert(kQkvNacc == 1 || kQkvNacc == 2 || kQkvNacc == 4, "accumulator count"); +static_assert(!M3FA_QKV_PIPE || kQkvScaled, "the pipelined projection uses the block-scaled MMAs"); +constexpr int kQkvOwnRows = 128 / kQkvSplits; // rows of a tile whose split-K sum a CTA owns +// Megakernel o_proj (v52 m3_mx_gemv for [6144, 2048]: split 2, contiguous k-tiles, one accumulator): unit +// o = CTA - 16 (the last 96 CTAs), row tile o / 2, split o % 2. +constexpr int kOpN = 6144; +constexpr int kOpK = 2048; +constexpr int kOpKTiles = kOpK / 128; +constexpr int kOpSplits = 2; +constexpr int kOpTiles = kOpKTiles / kOpSplits; +constexpr int kOpOwnRows = 128 / kOpSplits; +constexpr int kOpFirstCta = kQkvGrid - kOpN / 128 * kOpSplits; // 16 +constexpr int kOpSfBytes = kOpKTiles * 512; // one parity of the MXFP8 hand-off's scales +#ifndef M3FA_OP_L2PF +#define M3FA_OP_L2PF \ + 0 // 1: QKV CTAs prefetch their o_proj weight slice into L2 before the PDL wait (competes + // with the QKV weight stream for HBM behind the MoE) +#endif +// o_proj TMEM columns (after the QKV unit's MMAs completed): accumulator, A / B scale atoms, A (8 tiles x 32) +constexpr int kOAcc = 0; +constexpr int kOSfa = 16; +constexpr int kOSfb = kOSfa + 4 * kOpTiles; +constexpr int kOA = 128; +static_assert(kOA + 32 * kOpTiles <= kTmemAll && kOSfb + 4 * kOpTiles <= kOA, "o_proj TMEM columns"); +constexpr int kQkvMbStride = kMaxTok + 4; // mailbox row pitch (fp32 words) + +__host__ __device__ constexpr int qkv_tile0(int s) +{ + return s * kQkvKTiles / kQkvSplits; +} + +__device__ __forceinline__ int qkv_unit() +{ + int const u = kQkvGrid - 1 - static_cast(blockIdx.x); + return u < kQkvUnits ? u : -1; +} + +// Attention layout plus, for the fused projection, the weight ring in front (the attention tiles reuse that space +// once the ring has drained into TMEM); the activation tiles, the pre-wait V copy and the barriers sit behind. +template +struct SmemX : Smem +{ + using B = Smem; + static constexpr int kRingW = 0; // [ring][128 rows][128 B], 128-byte swizzle + static constexpr int kRingSf = kRingW + kQkvRing * kQkvWTile; // [ring][512 B] + static constexpr int kRingEnd = kRingSf + kQkvRing * kQkvSfAtom; + // The activation tiles overlay the ring: written after the PDL wait, once every weight tile went on into TMEM. + static constexpr int kXt = kRingW; // [tile][8][128 B] + static constexpr int kXsf = kXt + kQkvMaxTiles * 1024; // [tile][512 B] + static constexpr int kXsc = (kRingEnd + 127) / 128 * 128; + // Peers push into recv at any time (an empty CTA right away), so it stays clear of the ring. + static constexpr int kRecv = FQKV ? kXsc : B::kRecv; + static constexpr int kRawV = FQKV ? (kRecv + kBlocks * QG * BPC * kOPitch * 4 + 127) / 128 * 128 : B::kRawV; + // Pipelined projection: the activation tiles and scale atoms in their own region (a projection CTA runs no + // attention, so it overlays the attention's partial-O and V regions), 1024-byte aligned for the swizzle. + static constexpr int kXA = FQKV ? (kRecv + 1023) / 1024 * 1024 : 0; + static constexpr int kXASf = kXA + kQkvMaxTiles * 1024; + static constexpr int kXAEnd = kXASf + kQkvMaxTiles * kQkvSfAtom; + static constexpr int kMbarA = FQKV ? kRawV + BPC * B::kRawVBlock : B::kMbar; + static constexpr int kMbar = FQKV && M3FA_QKV_PIPE && kXAEnd > kMbarA ? (kXAEnd + 127) / 128 * 128 : kMbarA; + static constexpr int kGBar = kMbar + 32; // full[ring], empty[ring], a_ready, acc_done, tile_done[tiles] + static constexpr int kRedBar = kGBar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); // split-K partials in + static constexpr int kOBar = kRedBar + 8; // o_proj: full[8], a_ready, acc_done, reduce + static constexpr int kMailbox = (kOBar + 8 * (kOpTiles + 3) + 15) / 16 * 16; // [split][own row][pitch] fp32 + // o_proj weight ring = ring tiles [0, 8); its activation tiles, scale atoms and mailbox behind it. + static constexpr int kOXt = kRingW + kOpTiles * kQkvWTile; + static constexpr int kOXsf = kOXt + kOpTiles * 1024; + static constexpr int kOMailbox = kOXsf + kOpTiles * kQkvSfAtom; + static_assert(!FQKV || kOMailbox + kOpSplits * kOpOwnRows * kQkvMbStride * 4 <= kRingW + kQkvRing * kQkvWTile, + "o_proj activation and mailbox fit behind its ring"); + static constexpr int kBytes = FQKV ? kMailbox + kQkvSplits * kQkvOwnRows * kQkvMbStride * 4 : B::kBytes; + static constexpr int kTmemCols = FQKV ? kTmemAll : B::kTmemCols; + static constexpr bool kRecvAlias = !FQKV && B::kRecvAlias; + // The epilogue reads the weight scale atoms while the attention already writes its tiles. + static_assert(!FQKV || kQkvScaled || B::kRecv <= kRingSf, "attention regions overlap the scale atoms"); + // The other attention regions overlay the ring: a unit's CTA writes them only once the ring has drained. + static_assert(kBytes <= 227 * 1024, "shared memory budget"); +}; + +__device__ __forceinline__ void mbar_init(uint32_t bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(bar), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbar_expect_tx(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 bulk_g2s(void* dst, void const* src, uint32_t bytes, uint32_t bar) +{ + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], " + "%4;\n" ::"r"(smem_u32(dst)), + "l"(src), "r"(bytes), "r"(bar), "l"(0x12F0000000000000ull) // EVICT_FIRST: weights are read once per layer + : "memory"); +} + +__device__ __forceinline__ void tma_load_2d(uint32_t dst, void const* map, int c0, int c1, uint32_t bar) +{ + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;\n" ::"r"(dst), + "l"(map), "r"(c0), "r"(c1), "r"(bar), "l"(0x12F0000000000000ull) + : "memory"); +} + +// K-major operand in the 128-byte swizzle: rows of 128 bytes, 8-row groups 1024 bytes apart. +__device__ __forceinline__ uint64_t umma_desc_sw128(uint32_t addr) +{ + return umma_desc(addr, 16, 1024) | (2ull << 61); +} + +__device__ __forceinline__ void tmem_cp_128x256b(uint32_t taddr, uint64_t desc) +{ + asm volatile("tcgen05.cp.cta_group::1.128x256b [%0], %1;\n" ::"r"(taddr), "l"(desc)); +} + +__device__ __forceinline__ void tmem_cp_sf(uint32_t taddr, uint32_t saddr) +{ + asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;\n" ::"r"(taddr), "l"(umma_desc(saddr, 16, 128))); +} + +// Block-scaled MXFP8 instruction descriptor: E4M3 x E4M3, UE8M0 scales, M = 128, N = 8, scale byte kb. +__host__ __device__ constexpr uint32_t qkv_idesc(int kb) +{ + return (static_cast(kb) << 4) | ((8u >> 3) << 17) | (1u << 23) | ((128u >> 7) << 27) + | (static_cast(kb) << 29); +} + +__device__ __forceinline__ void umma_mx( + uint32_t tmem_d, uint32_t tmem_a, uint64_t bdesc, uint32_t idesc, uint32_t sfa, uint32_t sfb, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], [%1], %2, %3, [%4], [%5], p;\n}\n" :: + "r"(tmem_d), + "r"(tmem_a), "l"(bdesc), "r"(idesc), "r"(sfa), "r"(sfb), "r"(accumulate)); +} + +__device__ __forceinline__ float pow2_from_ue8m0(uint32_t e) +{ + return __uint_as_float(e == 0 ? 0x00400000u : (e << 23)); +} + +// Unscaled E4M3 x E4M3 MMA, FP32 accumulator: M = 128, N = 8, K = 32, a fresh accumulator. +constexpr uint32_t kQkvIdescF8 = (1u << 4) | ((8u >> 3) << 17) | ((128u >> 4) << 24); + +__device__ __forceinline__ void umma_f8_ss(uint32_t tmem_d, uint64_t adesc, uint64_t bdesc) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, 0, 0;\n tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(kQkvIdescF8)); +} + +__device__ __forceinline__ void tmem_ld16_nowait(uint32_t addr, float* v) +{ + 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];\n" + : "=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"(addr)); +#pragma unroll + for (int i = 0; i < 16; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +__device__ __forceinline__ void umma_mx_ss( + uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, uint32_t sfa, uint32_t sfb, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], %1, %2, %3, [%4], [%5], p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(sfa), "r"(sfb), "r"(accumulate)); +} + +__device__ __forceinline__ bool elect_one() +{ + uint32_t pred; + asm volatile("{\n .reg .pred P;\n elect.sync _|P, 0xffffffff;\n selp.u32 %0, 1, 0, P;\n}\n" : "=r"(pred)); + return pred != 0; +} + +__device__ __forceinline__ int qkv_parity(Params const& p) +{ + int v; + asm volatile("ld.relaxed.gpu.global.s32 %0, [%1];\n" : "=r"(v) : "l"(p.counters) : "memory"); + return v & 1; +} + +// Barriers and TMEM (warp 1), then the first ring stages (warp 0). Warps 0 and 1 of every CTA call it. +template +__device__ __forceinline__ void qkv_setup( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int lane, int u) +{ + constexpr bool kPipe = M3FA_QKV_PIPE != 0; + if (kPipe && warp == 0 && lane == 0 && u >= 0) + { + // First thing at entry: this unit's whole weight slice. The issuing thread initializes the ring's full + // barriers itself; every waiter on them passes bar.sync 8 below, which this thread reaches after the init. + asm volatile("prefetch.tensormap [%0];\n" ::"l"(&p.wmap) : "memory"); + int const rt = u / kQkvSplits; + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + for (int i = 0; i < kQkvRing; ++i) + { + mbar_init(gbar + 8 * i, 1); + } + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + for (int i = 0; i < kQkvRing && i < nt; ++i) + { + uint32_t const full = gbar + 8 * i; + mbar_expect_tx(full, kQkvWTile + kQkvSfAtom); + if (p.wpre != nullptr) + { + // Pre-tiled weight: the tile's 16 KB shared-memory image is contiguous (one bulk copy). + bulk_g2s(smem + SX::kRingW + i * kQkvWTile, + p.wpre + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvWTile, kQkvWTile, full); + } + else + { + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap, (kt0 + i) * 128, rt * 128, full); + } + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, full); + } + M3FA_PROF_T(34, 0); + } + if (warp == 1) + { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n" ::"r"(smem_u32(tmem_slot)), + "r"(kTmemAll) + : "memory"); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n" ::: "memory"); + M3FA_PROF_T(35, 32); + if (lane == 0) + { + for (int i = (kPipe && u >= 0) ? kQkvRing : 0; i < 2 * kQkvRing + 2 + kQkvMaxTiles; ++i) + { + // PIPE 2: tile_done[0] (unused by the scaled MMAs) is act_ready, one arrival per thread outside warp 1 + mbar_init(gbar + 8 * i, M3FA_QKV_PIPE == 2 && i == 2 * kQkvRing + 2 ? kThreads - 32 : 1); + } + // The split-K reduce barrier: this CTA's own arrival carries the peers' partial bytes. + uint32_t const red = gbar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); + mbar_init(red, 1); + mbar_expect_tx(red, (kQkvSplits - 1) * kQkvOwnRows * kMaxTok * 4); + uint32_t const ob = smem_u32(smem + SX::kOBar); + for (int i = 0; i < kOpTiles + 3; ++i) + { + mbar_init(ob + 8 * i, 1); + } + mbar_expect_tx(ob + 8 * (kOpTiles + 2), (kOpSplits - 1) * kOpOwnRows * kMaxTok * 4); + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + asm volatile("bar.sync 8, 64;\n" ::: "memory"); + if (!kPipe && warp == 0 && lane == 0 && u >= 0) + { + int const rt = u / kQkvSplits; + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + for (int i = 0; i < kQkvRing && i < nt; ++i) + { + uint32_t const full = gbar + 8 * i; + mbar_expect_tx(full, kQkvWTile + kQkvSfAtom); + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap, (kt0 + i) * 128, rt * 128, full); + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, full); + } + M3FA_PROF_T(34, 0); + } +} + +// The rest of the ring (warp 0, lane 0) and the copies into TMEM (warp 1, lane 0); before the PDL wait. +template +__device__ __forceinline__ void qkv_stream_lane( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int u) +{ + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + uint32_t const full = gbar; + uint32_t const a_ready = gbar + 16 * kQkvRing; + if (warp == 0 || !kQkvScaled) + { + return; // every tile was issued in qkv_setup; the unscaled MMAs read them in place + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + // Fully unrolled with compile-time offsets: tcgen05 operands stay in uniform registers (a per-instruction + // conversion of computed operands costs ~150 cycles per MMA or copy). + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint32_t const sf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + mbar_wait_cta(full + 8 * i, 0); + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if constexpr (!kQkvASmem) + { +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kTA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + } + tmem_cp_sf(tm + kTSfa + i * 4, sf0 + i * kQkvSfAtom); + } + } + M3FA_PROF_T(31, 32); + umma_commit(a_ready); +} + +__device__ __forceinline__ bool mbar_test(uint32_t bar, uint32_t parity) +{ + uint32_t ready; + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(bar), "r"(parity) + : "memory"); + return ready != 0; +} + +// PIPE 2, warp 1's elected thread from before the PDL wait to the last MMA. Tiles that land before the activation +// is ready go on into TMEM (A and scale atom), off the critical path; once the activation is in (act_ready), the +// MMAs issue in tile order: the copied tiles with A in TMEM, every later tile as it lands with A read in place from +// the ring. The MMA sequence (and so the accumulation order) is the same either way. +template +__device__ __forceinline__ void qkv_hybrid(Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int u) +{ + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + uint32_t const act_ready = gbar + 8 * (2 * kQkvRing + 2); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + SX::kXA)); + uint32_t const wsf0 = smem_u32(smem + SX::kRingSf); + uint32_t const xsf0 = smem_u32(smem + SX::kXASf); + int copied = 0; // tiles [0, copied) are in TMEM + bool act = false; +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt && !act) + { + while (!mbar_test(gbar + 8 * i, 0)) + { + if (mbar_test(act_ready, 0)) + { + act = true; + break; + } + } + if (!act) + { + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kTA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + copied = i + 1; + } + } + } + while (!act) + { + act = mbar_test(act_ready, 0); + } + M3FA_PROF_T(21, 32); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + bool const in_tmem = i < copied; + if (!in_tmem) + { + mbar_wait_epi(gbar + 8 * i, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + } + uint32_t const sfb = tm + kTSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + uint64_t const bd = b0 + ((i * 1024 + kb * 32) >> 4); + if (in_tmem) + { + umma_mx(tm + kTAcc, tm + kTA + j * 8, bd, qkv_idesc(kb), tm + kTSfa + i * 4, sfb, j > 0); + } + else + { + umma_mx_ss(tm + kTAcc, w0 + ((i * kQkvWTile + kb * 32) >> 4), bd, qkv_idesc(kb), tm + kTSfa + i * 4, + sfb, j > 0); + } + } + } + } + M3FA_PROF_T(22, 32); + M3FA_PROF_T(31, 32); + umma_commit(gbar + 16 * kQkvRing); // a_ready (for its waiters) + umma_commit(gbar + 16 * kQkvRing + 8); // acc_done +} + +// The other lanes of warps 0 and 1 hold at __syncwarp rather than at griddepcontrol.wait, which would stall the +// whole warp (lane 0's loop included) until the producer grid completes. +template +__device__ __forceinline__ void qkv_stream( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int lane, int u) +{ + if (M3FA_QKV_PIPE == 2 && u >= 0 && warp == 1) + { + if (elect_one()) + { + qkv_hybrid(p, smem, gbar, tmem_slot, u); + } + __syncwarp(); + return; + } + if (u < 0 || warp > 1 || M3FA_QKV_PIPE) + { + return; // pipelined: the MMA thread takes each tile as it lands + } + if (lane == 0) + { + qkv_stream_lane(p, smem, gbar, tmem_slot, warp, u); + } + __syncwarp(); +} + +// After the PDL wait. Every thread quantizes a piece of the x slice into the B tiles (named barrier 9); warp 1 then +// issues the scale copies and the MMAs. +template +__device__ __forceinline__ void qkv_mma( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int tid, int u) +{ + int const warp = tid >> 5; + int const s = u % kQkvSplits; + int const kt0 = qkv_tile0(s); + int const nt = qkv_tile0(s + 1) - kt0; + { + // The MXFP8 activation slice (E4M3 rows of the unit's k-tiles in 128-byte-swizzled B tiles, and their scale + // atoms), copied with plain 16-byte loads over the overlaid ring once its tiles went on into TMEM: every load + // first, then the stores (one L2 round trip). + constexpr int kXDst = M3FA_QKV_PIPE ? SX::kXA : SX::kXt; + constexpr int kXSfDst = M3FA_QKV_PIPE ? SX::kXASf : SX::kXsf; + if (!M3FA_QKV_PIPE) + { + mbar_wait_cta(gbar + 16 * kQkvRing, 0); // a_ready: the ring's weight tiles are in TMEM + } + constexpr int kXItems = kQkvMaxTiles * 8 * kMaxTok; // (tile, 16-byte chunk, token) + constexpr int kSfItems = kQkvMaxTiles * (kQkvSfAtom / 16); + constexpr bool kHyb = M3FA_QKV_PIPE == 2; + constexpr int kCopiers = kHyb ? kThreads - 32 : kThreads; // PIPE 2: warp 1 is issuing the tiles + if (kHyb && warp == 1) + { + return; + } + int const ctid = kHyb && tid >= 64 ? tid - 32 : tid; + constexpr int kPer = (kXItems + kSfItems + kCopiers - 1) / kCopiers; + uint4 v[kPer]; + int dst[kPer]; +#pragma unroll + for (int k = 0; k < kPer; ++k) + { + int const q = ctid + kCopiers * k; + dst[k] = -1; + if (q < kXItems) + { + int const t = q % kMaxTok; + int const c = (q / kMaxTok) % 8; + int const i = q / (kMaxTok * 8); + if (i < nt && t < p.num_tokens) + { + v[k] = *reinterpret_cast( + p.xq + static_cast(t) * p.ldxq + (kt0 + i) * 128 + c * 16); + dst[k] = kXDst + i * 1024 + t * 128 + ((c ^ t) * 16); + } + } + else if (q < kXItems + kSfItems) + { + int const j = q - kXItems; + int const i = j / (kQkvSfAtom / 16); + int const c = j % (kQkvSfAtom / 16); + if (i < nt) + { + v[k] = *reinterpret_cast( + p.xsf_act + static_cast(kt0 + i) * kQkvSfAtom + c * 16); + dst[k] = kXSfDst + i * kQkvSfAtom + c * 16; + } + } + } +#pragma unroll + for (int k = 0; k < kPer; ++k) + { + if (dst[k] >= 0) + { + *reinterpret_cast(smem + dst[k]) = v[k]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if (kHyb) + { + asm volatile("mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];\n" ::"r"(gbar + 8 * (2 * kQkvRing + 2)) + : "memory"); + return; + } + if (warp != 1) + { + asm volatile("bar.arrive 9, %0;\n" ::"n"(kThreads) : "memory"); + } + } + if (warp == 1) + { + M3FA_PROF_T(32, 32); + asm volatile("bar.sync 9, %0;\n" ::"n"(kThreads) : "memory"); + M3FA_PROF_T(21, 32); + // Operands computed by the whole warp (warp-uniform, kept in uniform registers); one elected lane issues. + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + (M3FA_QKV_PIPE ? SX::kXA : SX::kXt))); + uint64_t const a0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + if (elect_one()) + { + M3FA_CLK(41); + if constexpr (!kQkvScaled) + { +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + mbar_wait_cta(gbar + 8 * i, 0); // weight tile i landed + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + umma_f8_ss(tm + kTAcc + (i * 4 + kb) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4), + b0 + ((i * 1024 + kb * 32) >> 4)); + } + umma_commit(gbar + 8 * (2 * kQkvRing + 2 + i)); // tile i's blocks are in TMEM + } + } + } + else + { + uint32_t const xsf0 = smem_u32(smem + (M3FA_QKV_PIPE ? SX::kXASf : SX::kXsf)); + uint32_t const wsf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + if constexpr (M3FA_QKV_PIPE != 0) + { + // Tile i as it lands: its A copy and scale atom into TMEM, then its MMAs (issue order). + mbar_wait_cta(gbar + 8 * i, 0); + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + if (i == nt - 1) + { + M3FA_CLK(47); // the unit's last weight tile landed (clock64 builds) + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if constexpr (!kQkvASmem) + { +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b( + tm + kTA + (i * 4 + kb) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + } + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + } + uint32_t const sfb = tm + kTSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + uint64_t const bd = b0 + ((i * 1024 + kb * 32) >> 4); + if constexpr (kQkvASmem) + { + umma_mx_ss(tm + kTAcc + (j % kQkvNacc) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4), bd, + qkv_idesc(kb), tm + kTSfa + i * 4, sfb, j >= kQkvNacc); + } + else + { + umma_mx(tm + kTAcc + (j % kQkvNacc) * 8, tm + kTA + j * 8, bd, qkv_idesc(kb), + tm + kTSfa + i * 4, sfb, j >= kQkvNacc); + } + } + } + } + } + M3FA_CLK(42); + if (M3FA_QKV_PIPE) + { + M3FA_PROF_T(31, 32); + umma_commit(gbar + 16 * kQkvRing); // a_ready (kept for its waiters: every tile is consumed) + } + umma_commit(gbar + 16 * kQkvRing + 8); // acc_done + } + M3FA_PROF_T(22, 32); + } +} + +// Warps 16..19 (TMEM lanes 32 (warp % 4) + lane = the unit's rows): partial rows out, the other parity reset. +template +__device__ __forceinline__ void qkv_epilogue( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int tid, int u, int par, unsigned int calls) +{ + int const warp = tid >> 5; + int const lane = tid & 31; + if constexpr (kQkvScaled) + { + mbar_wait_epi(gbar + 16 * kQkvRing + 8, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + M3FA_PROF_T(24, 512); + if (tid == 512) + { + M3FA_CLK(43); + } + int const r = (warp & 3) * 32 + lane; // row within the unit's tile + uint32_t const taddr = *tmem_slot + (((warp & 3) * 32) << 16) + kTAcc; + float val[kMaxTok]; + if constexpr (kQkvScaled) + { + float v[32]; + tmem_ld16(taddr, *reinterpret_cast(v)); + if constexpr (kQkvNacc == 4) + { + tmem_ld16(taddr + 16, *reinterpret_cast(v + 16)); + } +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = kQkvNacc == 4 ? (v[t] + v[8 + t]) + (v[16 + t] + v[24 + t]) + : kQkvNacc == 2 ? v[t] + v[8 + t] + : v[t]; + } + } + else + { + // Sum over K blocks in order: block result x (weight scale x activation scale). + int const nt = qkv_tile0(u % kQkvSplits + 1) - qkv_tile0(u % kQkvSplits); + float const* xsc = reinterpret_cast(smem + SX::kXsc); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = 0.0f; + } +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + // Tile i's four blocks as soon as its MMAs complete: the epilogue trails the MMAs by one tile. + mbar_wait_cta(gbar + 8 * (2 * kQkvRing + 2 + i), 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + float d[32]; + tmem_ld16_nowait(taddr + i * 32, d); + tmem_ld16_nowait(taddr + i * 32 + 16, d + 16); + uint32_t const wsw = *reinterpret_cast( + smem + SX::kRingSf + i * kQkvSfAtom + (r % 32) * 16 + (r / 32) * 4); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + float const ws = pow2_from_ue8m0((wsw >> (8 * kb)) & 0xFFu); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = fmaf(d[kb * 8 + t], ws * xsc[(i * 4 + kb) * kMaxTok + t], val[t]); + } + } + } + } + } + // v52's split-K reduce: split s owns rows [32 s, 32 s + 32) of the tile; the other splits push their fp32 partials + // of those rows into its mailbox (st.async, completing its reduce barrier's bytes); the owner adds the splits in + // rank order from 0.0 and stores bf16. + int const s = u % kQkvSplits; + int const owner = r / kQkvOwnRows; + int const local = r % kQkvOwnRows; + uint32_t const red = gbar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); + uint32_t cta_rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + // Units are numbered down from the last CTA: split q of this tile is CTA rank cta_rank - (q - s). + if (owner != s) + { + uint32_t const peer = cta_rank + s - owner; + uint32_t const dst + = mapa_shared(smem_u32(smem + SX::kMailbox) + ((s * kQkvOwnRows + local) * kQkvMbStride) * 4, peer); + uint32_t const pbar = mapa_shared(red, peer); +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( + dst + 16 * g), + "r"(__float_as_uint(val[4 * g])), "r"(__float_as_uint(val[4 * g + 1])), + "r"(__float_as_uint(val[4 * g + 2])), "r"(__float_as_uint(val[4 * g + 3])), "r"(pbar) + : "memory"); + } + } + else + { + M3FA_PROF_T(26, tid - lane); + // test_wait (no suspend): completions by peers' st.async bytes wake a suspended try_wait late. + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(red) + : "memory"); + } + M3FA_PROF_T(27, tid - lane); + float const* mb = reinterpret_cast(smem + SX::kMailbox); + int const row = (u / kQkvSplits) * 128 + r; + // Lamport hand-off: this call's rows go to buffer (calls & 1) with relaxed stores, the consumers poll the + // values themselves (0xFFFF, a NaN the rounding of a finite sum never produces, marks "not yet"); the other + // buffer, last read by the previous call (complete: this grid passed its PDL wait), is reset to the marker. + uint16_t* cur = reinterpret_cast(p.qkv_out) + (calls & 1u) * kMaxTok * kQkvN; + uint16_t* nxt = reinterpret_cast(p.qkv_out) + ((calls & 1u) ^ 1u) * kMaxTok * kQkvN; +#pragma unroll + float qsum[kMaxTok]; +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + float4 part[kQkvSplits]; +#pragma unroll + for (int q = 0; q < kQkvSplits; ++q) + { + part[q] = *reinterpret_cast(mb + (q * kQkvOwnRows + local) * kQkvMbStride + 4 * g); + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float sum = 0.0F; +#pragma unroll + for (int q = 0; q < kQkvSplits; ++q) + { + float const pj = j == 0 ? part[q].x : j == 1 ? part[q].y : j == 2 ? part[q].z : part[q].w; + sum = sum + (q == s ? val[4 * g + j] : pj); + } + qsum[4 * g + j] = sum; + } + } +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + float const sum = qsum[t]; + if (t < p.num_tokens) + { + __nv_bfloat16 const b = __float2bfloat16_rn(sum); + uint16_t bits = *reinterpret_cast(&b); + // Never the "not yet" marker: a 0xFFFF NaN sum is stored as the NaN 0x7FFF. Only a NaN projection row + // changes (its K/V cache bytes then differ from v52 + p6 in the NaN's encoding); finite rows are exact. + bits = bits == 0xFFFFu ? static_cast(0x7FFFu) : bits; +#ifndef M3FA_QKV_REL +#define M3FA_QKV_REL 0 // 1: the row store is a release (tried: whether it becomes visible to the pollers sooner) +#endif + if (M3FA_QKV_REL) + { + asm volatile( + "st.release.gpu.global.u16 [%0], %1;\n" ::"l"(cur + static_cast(t) * kQkvN + row), + "h"(bits) + : "memory"); + } + else + { + asm volatile( + "st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(cur + static_cast(t) * kQkvN + row), + "h"(bits) + : "memory"); + } + } + asm volatile("st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(nxt + static_cast(t) * kQkvN + row), + "h"(static_cast(0xFFFFu)) + : "memory"); + } + M3FA_PROF_T(28, tid - lane); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + M3FA_PROF_T(25, 512); + if (tid == 512) + { + M3FA_CLK(46); + } +} + +// ---------------- Megakernel o_proj (v52 m3_mx_gemv bit for bit: split 2, contiguous k-tiles, nacc 1) ---------------- +// The unit's weight tiles and scale atoms into the ring (warp 0, lane 0). From L2 when the CTA prefetched them there. +template +__device__ __forceinline__ void op_issue(Params const& p, uint8_t* smem, int ou) +{ + int const rt = ou / kOpSplits; + int const kt0 = (ou % kOpSplits) * kOpTiles; + uint32_t const ob = smem_u32(smem + SX::kOBar); + for (int i = 0; i < kOpTiles; ++i) + { + mbar_expect_tx(ob + 8 * i, kQkvWTile + kQkvSfAtom); + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap_o, (kt0 + i) * 128, rt * 128, ob + 8 * i); + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf_o + (static_cast(rt) * kOpKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, ob + 8 * i); + } +} + +// L2 prefetch of the unit's weight tiles and scale atoms (one lane; nothing waits for it). +// The QKV units' share of the window-1 L2 prefetch (the next consumer's weights), issued once their own QKV weight +// stream has landed (acc_done): the attention core reads little from HBM. +__device__ __forceinline__ void window_prefetch_l2(Params const& p, int u) +{ + constexpr long long kChunk = 32768; + long long const n0 = (p.pf_b0 + kChunk - 1) / kChunk, n1 = (p.pf_b1 + kChunk - 1) / kChunk; + uint64_t pol = 0; + if (p.pf_hint) + { + asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, 1.0;\n" : "=l"(pol)); + } + for (long long c = u; c < n0 + n1; c += kQkvUnits) + { + long long const a = c < n0 ? p.pf_a0 + c * kChunk : p.pf_a1 + (c - n0) * kChunk; + long long const rem = c < n0 ? p.pf_b0 - c * kChunk : p.pf_b1 - (c - n0) * kChunk; + unsigned int const n = static_cast(rem < kChunk ? rem : kChunk) & ~15u; + if (n == 0) + { + continue; + } + if (p.pf_hint) + { + asm volatile("cp.async.bulk.prefetch.L2.global.L2::cache_hint [%0], %1, %2;\n" ::"l"(a), "r"(n), "l"(pol) + : "memory"); + } + else + { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" ::"l"(a), "r"(n) : "memory"); + } + } +} + +__device__ __forceinline__ void op_prefetch_l2(Params const& p, int ou) +{ + int const rt = ou / kOpSplits; + int const kt0 = (ou % kOpSplits) * kOpTiles; + for (int i = 0; i < kOpTiles; ++i) + { + asm volatile("cp.async.bulk.prefetch.tensor.2d.L2.global [%0, {%1, %2}];\n" ::"l"(&p.wmap_o), + "r"((kt0 + i) * 128), "r"(rt * 128) + : "memory"); + } + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" ::"l"( + p.wsf_o + (static_cast(rt) * kOpKTiles + kt0) * kQkvSfAtom), + "r"(kOpTiles * kQkvSfAtom) + : "memory"); +} + +// The weight tiles into TMEM as they land (warp 1, elected lane, operands warp-uniform). +template +__device__ __forceinline__ void op_stream(uint8_t* smem, uint32_t tm) +{ + uint32_t const ob = smem_u32(smem + SX::kOBar); + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint32_t const sf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kOpTiles; ++i) + { + mbar_wait_cta(ob + 8 * i, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kOA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + tmem_cp_sf(tm + kOSfa + i * 4, sf0 + i * kQkvSfAtom); + } + umma_commit(ob + 8 * kOpTiles); // a_ready +} + +// The attention's MXFP8 rows of this unit's k-tiles (the heads [kt0, kt0 + 8)), Lamport-polled value by value: 0xFF +// is NaN in both E4M3 and UE8M0, which the merge never writes (see the merge). Every thread of the CTA, one item each. +template +__device__ __forceinline__ void op_act(Params const& p, uint8_t* smem, int ou, int tid, unsigned int calls) +{ + int const kt0 = (ou % kOpSplits) * kOpTiles; + uint8_t const* q8 = p.q8 + (calls & 1u) * kMaxTok * kOpK; + uint8_t const* sf = p.sf + (calls & 1u) * kOpSfBytes; + constexpr int kXItems = kOpTiles * 8 * kMaxTok; + constexpr int kSfItems = kOpTiles * (kQkvSfAtom / 16); + for (int q = tid; q < kXItems + kSfItems; q += kThreads) + { + if (q < kXItems) + { + int const t = q % kMaxTok; + int const c = (q / kMaxTok) % 8; + int const i = q / (kMaxTok * 8); + if (t < p.num_tokens) + { + uint4 v; + bool pending; + do + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(q8 + static_cast(t) * kOpK + (kt0 + i) * 128 + c * 16) + : "memory"); + uint32_t const w[4] = {v.x, v.y, v.z, v.w}; + pending = false; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + // a byte equal to 0xFF: (x & 0x7F7F7F7F) + 0x01010101 sets bit 7 of exactly the 0x7F/0xFF bytes + uint32_t const hit = (((w[j] & 0x7F7F7F7Fu) + 0x01010101u) & w[j] & 0x80808080u); + pending |= hit != 0u; + } + } while (pending); + *reinterpret_cast(smem + SX::kOXt + i * 1024 + t * 128 + ((c ^ t) * 16)) = v; + } + } + else if (q < kXItems + kSfItems) + { + int const j = q - kXItems; + int const i = j / (kQkvSfAtom / 16); + int const c = j % (kQkvSfAtom / 16); + uint8_t const* src = sf + (kt0 + i) * kQkvSfAtom + c * 16; + uint4 v; + if (c < p.num_tokens) + { + // Row c of the atom: its 4 block scales in bytes [16 c, 16 c + 4). + do + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(src) + : "memory"); + } while ((((v.x & 0x7F7F7F7Fu) + 0x01010101u) & v.x & 0x80808080u) != 0u); + } + else + { + v = *reinterpret_cast(src); + } + *reinterpret_cast(smem + SX::kOXsf + i * kQkvSfAtom + c * 16) = v; + } + } +} + +// The MMAs (warp 1, elected lane): v52's order (per k-tile: B scale copy, 4 MMAs; one accumulator). +template +__device__ __forceinline__ void op_mma(uint8_t* smem, uint32_t tm) +{ + uint32_t const ob = smem_u32(smem + SX::kOBar); + mbar_wait_cta(ob + 8 * kOpTiles, 0); // a_ready + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + SX::kOXt)); + uint32_t const xsf0 = smem_u32(smem + SX::kOXsf); +#pragma unroll + for (int i = 0; i < kOpTiles; ++i) + { + uint32_t const sfb = tm + kOSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + umma_mx(tm + kOAcc, tm + kOA + j * 8, b0 + ((i * 1024 + kb * 32) >> 4), qkv_idesc(kb), tm + kOSfa + i * 4, + sfb, j > 0); + } + } + umma_commit(ob + 8 * (kOpTiles + 1)); // acc_done +} + +// Warps 16..19: the pair's split-K reduce (owner rows [64 s, 64 s + 64)) and the bf16 rows of y. +template +__device__ __forceinline__ void op_epilogue(Params const& p, uint8_t* smem, uint32_t tm, int tid, int ou) +{ + int const warp = tid >> 5; + int const lane = tid & 31; + uint32_t const ob = smem_u32(smem + SX::kOBar); + // The store's kernel parameters, fetched before the wait: a first touch of their constant-bank line costs a miss. + __nv_bfloat16* y_out = p.y_out; + int num_tokens = p.num_tokens; + int ar_mode = p.ar_mode; + asm volatile("" : "+l"(y_out), "+r"(num_tokens), "+r"(ar_mode)); +#ifndef M3FA_OP_PARPRE +#define M3FA_OP_PARPRE 1 // read the all-reduce slot parities before the accumulator wait (a global load off the tail) +#endif + // The parity flags are stable here: the all-reduce that flips them last completed before this grid's PDL wait, + // and the next one reads them only after this kernel's pushes. + int par_pre[kMaxTok]; + { + int const row_pre = (ou / kOpSplits) * 128 + ((tid >> 5) & 3) * 32 + (tid & 31); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + par_pre[t] = 0; + if (M3FA_OP_PARPRE && ar_mode != 0 && t < num_tokens) + { + asm volatile("ld.volatile.global.s32 %0, [%1];\n" + : "=r"(par_pre[t]) + : "l"(p.ar_flags + t * 8 + row_pre / (kOpN / 8)) + : "memory"); + } + } + } + mbar_wait_epi(ob + 8 * (kOpTiles + 1), 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF_T(51, kComputeThreads); + int const r = (warp & 3) * 32 + lane; + float v[16]; + tmem_ld16(tm + (((warp & 3) * 32) << 16) + kOAcc, v); + int const s = ou % kOpSplits; + int const owner = r / kOpOwnRows; + int const local = r % kOpOwnRows; + uint32_t const red = ob + 8 * (kOpTiles + 2); + uint32_t cta_rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + if (owner != s) + { + uint32_t const peer = cta_rank + owner - s; // o_proj units count up with the CTA index + uint32_t const dst + = mapa_shared(smem_u32(smem + SX::kOMailbox) + ((s * kOpOwnRows + local) * kQkvMbStride) * 4, peer); + uint32_t const pbar = mapa_shared(red, peer); +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( + dst + 16 * g), + "r"(__float_as_uint(v[4 * g])), "r"(__float_as_uint(v[4 * g + 1])), "r"(__float_as_uint(v[4 * g + 2])), + "r"(__float_as_uint(v[4 * g + 3])), "r"(pbar) + : "memory"); + } + } + else + { + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(red) + : "memory"); + } + M3FA_PROF_T(52, tid - lane); + float const* mb = reinterpret_cast(smem + SX::kOMailbox); + int const row = (ou / kOpSplits) * 128 + r; + float sums[kMaxTok]; +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + float4 part[kOpSplits]; +#pragma unroll + for (int q = 0; q < kOpSplits; ++q) + { + part[q] = *reinterpret_cast(mb + (q * kOpOwnRows + local) * kQkvMbStride + 4 * g); + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float sum = 0.0F; +#pragma unroll + for (int q = 0; q < kOpSplits; ++q) + { + float const pj = j == 0 ? part[q].x : j == 1 ? part[q].y : j == 2 ? part[q].z : part[q].w; + sum = sum + (q == s ? v[4 * g + j] : pj); + } + sums[4 * g + j] = sum; + } + } + M3FA_PROF_T(54, tid - lane); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + float const sum = sums[t]; +#ifndef M3FA_EXP +#define M3FA_EXP 0 +#endif + if (t < num_tokens) + { + if (M3FA_EXP == 1) + { + asm volatile( + "st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(y_out + static_cast(t) * kOpN + row), + "h"(__bfloat16_as_ushort(__float2bfloat16_rn(sum))) + : "memory"); + } + else + { + y_out[static_cast(t) * kOpN + row] = __float2bfloat16_rn(sum); + } + } + if (ar_mode != 0) + { + // m3_mx_gemv's push into the all-reduce exchange: lane pairs hold row pairs; the even lane writes the + // bf16x2 word into this rank's slot of the (token, column slice) parity the AR's flags select. + float const hi = __shfl_xor_sync(0xffffffffu, sum, 1); + if ((lane & 1) == 0 && t < num_tokens) + { + int par; + if (M3FA_OP_PARPRE) + { + par = par_pre[t]; + } + else + { + asm volatile("ld.volatile.global.s32 %0, [%1];\n" + : "=r"(par) + : "l"(p.ar_flags + t * 8 + row / (kOpN / 8)) + : "memory"); + } + int64_t const off = static_cast(par & 1) * p.ar_maxt * p.ar_world * (kOpN / 2) + + (static_cast(t) * p.ar_world + p.ar_rank) * (kOpN / 2) + row / 2; + __nv_bfloat162 w2 = __floats2bfloat162_rn(sum, hi); + uint32_t word = *reinterpret_cast(&w2); + // -0.0 halves to +0.0: 0x80000000 marks an empty exchange word (as m3_mx_gemv's _bf16x2_word) + word = (word & 0xFFFFu) == 0x8000u ? (word & 0xFFFF0000u) : word; + word = (word >> 16) == 0x8000u ? (word & 0x0000FFFFu) : word; + asm volatile("st.global.b32 [%0], %1;\n" ::"l"(p.ar_ws + off), "r"(word) : "memory"); + } + } + } + M3FA_PROF_T(53, tid - lane); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); +} + +// Packed-QKV element group [token][col * 128 + 4 * lane, +4): the split partials, polled until every one is stored, +// summed in split order and rounded to BF16. +__device__ __forceinline__ void qkv_load(Params const& p, int par, int tok, int col, int lane, float (&e)[4]) +{ + float const* base = p.part + par * kQkvPartStride + static_cast(tok) * kQkvN + col * kHeadDim + lane * 4; + uint4 pr[kQkvSplits]; + long long spins = 0; + bool ready; + do + { + ready = true; +#pragma unroll + for (int s = 0; s < kQkvSplits; ++s) + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(pr[s].x), "=r"(pr[s].y), "=r"(pr[s].z), "=r"(pr[s].w) + : "l"(base + static_cast(s) * kMaxTok * kQkvN) + : "memory"); + ready &= pr[s].x != kQkvSentinel && pr[s].y != kQkvSentinel && pr[s].z != kQkvSentinel + && pr[s].w != kQkvSentinel; + } + if (++spins > (1LL << 26)) + { + __trap(); + } + } while (!ready); + float s0 = __uint_as_float(pr[0].x), s1 = __uint_as_float(pr[0].y), s2 = __uint_as_float(pr[0].z), + s3 = __uint_as_float(pr[0].w); +#pragma unroll + for (int s = 1; s < kQkvSplits; ++s) + { + s0 += __uint_as_float(pr[s].x); + s1 += __uint_as_float(pr[s].y); + s2 += __uint_as_float(pr[s].z); + s3 += __uint_as_float(pr[s].w); + } + e[0] = __bfloat162float(__float2bfloat16_rn(s0)); + e[1] = __bfloat162float(__float2bfloat16_rn(s1)); + e[2] = __bfloat162float(__float2bfloat16_rn(s2)); + e[3] = __bfloat162float(__float2bfloat16_rn(s3)); +} + +// Every CTA once it has read the call parity and the partials it needs: the last one flips the parity. +__device__ __forceinline__ void qkv_done(Params const& p) +{ + if (atomicAdd(p.counters + 1, 1) == kQkvGrid - 1) + { + atomicExch(p.counters + 1, 0); + atomicXor(p.counters, 1); + } +} + +} // namespace + +namespace +{ + +// FQKV: the kernel also computes the QKV(+index) projection from the layer input (see qkv_mma) and runs on a fixed +// kQkvGrid-CTA grid; the first attn_clusters clusters then run the attention on the projection's output. +template +__global__ void __launch_bounds__(kThreads, 1) m3_fused_decode_attn_kernel(__grid_constant__ Params const p) +{ +#if MINIMAX_M3_ATTN_DEVICE_ENABLED + using S = SmemX; + constexpr int kRows = S::kRows; + constexpr int kTok = S::kTok; + constexpr int kCs = kBlocks / BPC; // cluster size + constexpr int kGroups = QLEN / QG; // clusters per request + constexpr uint32_t kPushBytes = (kHeadDim + 1) * 4; // one partial row: 128 outputs and the LSE + constexpr int kTasks = kHq * QG + 3 * QLEN; // producer tasks: 16 Q heads x QG, [K, V, index-K] x QLEN + constexpr int kMaxT = (kTasks + kWarps - 1) / kWarps; + constexpr int kChunks = BPC * kPage * 4; // 16-byte packed chunks per K (or V) tile set + constexpr int kKPer = (kChunks + kComputeThreads - 1) / kComputeThreads; // K chunks per thread [0, 512) + constexpr bool kRegCore = RegCore::kOn; // warp-MMA attention core (see RegCore) + // V tiles (raw copy before the PDL wait, dequantized during QK and the softmax): on the packed path every warp + // outside the softmax groups takes part (at M=1, warps 4..19), on the fused-projection path the producer warps. +#ifndef M3FA_RS2 +#define M3FA_RS2 1 // megakernel at one 16-row chunk: the softmax and push rows split over two warp groups (8 each) +#endif + // Split softmax: the megakernel and multi-row (spec verify) steps; at one-row decode on the packed path it + // measured slower (it takes four warps from the V dequantization there). + constexpr bool kRs2 = M3FA_RS2 && BPC * QG == 1 && !RegCore::kOn && (FQKV || QLEN > 1); +#ifndef M3FA_MEGA_WARPMMA +#define M3FA_MEGA_WARPMMA 1 +#endif + // Warp-issued MMAs and the early cluster barrier: measured gains in the megakernel and on multi-row / paired-block + // steps, but not at one-row decode on 16-CTA clusters (the packed path's (1, 1, 1)), where they stay off. + constexpr bool kTuned = FQKV || QLEN > 1 || BPC > 1; + constexpr bool kWarpMma = M3FA_MEGA_WARPMMA && kTuned; + constexpr bool kEarlyClw = M3FA_MEGA_EARLYCLW && kTuned; + constexpr int kSoftmaxGroups = kRs2 ? 2 : (BPC * QG) < 4 ? (BPC * QG) : 4; +#ifndef M3FA_MEGA_VWIDE +#define M3FA_MEGA_VWIDE \ + 1 // megakernel attention CTAs (no projection unit): every warp outside the softmax copies and + // dequantizes V, as on the packed path (not only the four producer warps) +#endif + constexpr int kVStart = (kRegCore || (FQKV && !M3FA_MEGA_VWIDE)) ? kComputeThreads : 128 * kSoftmaxGroups; + constexpr int kVThreads = kThreads - kVStart; // V chunks: the warps outside the softmax groups + constexpr int kVPer = (kChunks + kVThreads - 1) / kVThreads; + constexpr int kPairs = BPC * QG; // (block, 16-row chunk) pairs of the softmax and push + constexpr int kPairsPerGroup = (kPairs + 3) / 4; // per group of four compute warps + // Named barrier 6: the softmax warps (P written) and the producer warps (V written) before the PV issue. + constexpr int kPvBarThreads = 128 * kSoftmaxGroups + kVThreads; + extern __shared__ __align__(128) uint8_t smem[]; + float* row_scale = reinterpret_cast(smem + S::kRowScale); + float* row_lse = reinterpret_cast(smem + S::kRowLse); + float* red_m = reinterpret_cast(smem + S::kRedM); + float* red_l = reinterpret_cast(smem + S::kRedL); + float* recv = reinterpret_cast(smem + S::kRecv); + uint32_t const mbar = smem_u32(smem + S::kMbar); + uint32_t const mma_bar = mbar + 8; + uint32_t* tmem_slot = reinterpret_cast(smem + S::kMbar + 16); + + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const cr = blockIdx.x % kCs; // cluster rank; owns blocks [cr * BPC, cr * BPC + BPC) + int const ci = blockIdx.x / kCs; + int const req = ci / kGroups; + int const tfirst = (ci % kGroups) * QG; // this cluster's query tokens [tfirst, tfirst + QG) of the request + int const tok0 = req * QLEN; + uint32_t const gbar = smem_u32(smem + S::kGBar); // FQKV: weight ring, A ready, accumulators done + int const qkv_u = FQKV ? qkv_unit() : -1; // FQKV: this CTA's projection unit, or -1 + + // ---------------- Prologue: nothing here depends on the producer kernel ---------------- +#if M3FA_TRIGGER_AT == 0 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (!FQKV) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + M3FA_PROF(0); +#ifdef M3FA_PROFILE + if (p.prof != nullptr && tid == 0) + { + uint32_t smid; + asm volatile("mov.u32 %0, %%smid;" : "=r"(smid)); + p.prof[blockIdx.x * 64 + 40] = smid; + } +#endif + if constexpr (FQKV) + { + // The projection weights stream into TMEM while the producer (the all-reduce) still runs. + if (warp < 2) + { + qkv_setup(p, smem, gbar, tmem_slot, warp, lane, qkv_u); + } + if (ci >= p.attn_clusters) + { + // A projection-only CTA. Every reduce barrier of the cluster is initialized before any peer pushes. + asm volatile("barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;\n" :: + : "memory"); + int const op_u = static_cast(blockIdx.x) - kOpFirstCta; // this CTA's o_proj unit (>= 0 here) +#ifndef M3FA_OP_EARLY +#define M3FA_OP_EARLY 0 // 1: o_proj-only CTAs stream their weights from launch (0: after the PDL wait) +#endif + if (M3FA_OP_EARLY && qkv_u < 0 && warp == 0 && lane == 0) + { + op_issue(p, smem, op_u); // an o_proj-only CTA: its weights stream from launch on + } + if (M3FA_OP_L2PF && qkv_u >= 0 && warp == 0 && lane == 0) + { + op_prefetch_l2(p, op_u); // into L2 during the all-reduce, behind this CTA's QKV ring + } + qkv_stream(p, smem, gbar, tmem_slot, warp, lane, qkv_u); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (M3FA_MEGA_TRIGGER == 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + // This CTA's call counter: read after the PDL wait (the previous call's CTA increments it at its end), and + // early, so its latency stays off the epilogue's critical path. + unsigned int const calls = p.flags[32 + blockIdx.x]; + M3FA_PROF(1); + if (!M3FA_OP_EARLY && qkv_u < 0 && warp == 0 && lane == 0) + { + op_issue(p, smem, op_u); + } + if (qkv_u >= 0) + { + int const par = 0; + qkv_mma(p, smem, gbar, tmem_slot, tid, qkv_u); + if (warp >= kComputeWarps) + { + qkv_epilogue(p, smem, gbar, tmem_slot, tid, qkv_u, par, calls); + } + if (warp == 0 && lane == 0) + { + // The ring (and the activation over it) is free once the QKV MMAs completed. + mbar_wait_cta(gbar + 16 * kQkvRing + 8, 0); + op_issue(p, smem, op_u); + if (p.pf_b0 > 0 || p.pf_b1 > 0) + { + window_prefetch_l2(p, qkv_u); + } + } + } + __syncthreads(); // the QKV accumulator reads are done: TMEM columns [0, 512) are the o_proj's + M3FA_PROF(2); + uint32_t const tm = *tmem_slot; + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_stream(smem, tm); + } + __syncwarp(); + } + op_act(p, smem, op_u, tid, calls); + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + M3FA_PROF_T(36, 0); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Late release: the dependents launch once every unit has its attention input (the rel_v4 rule). + if (M3FA_MEGA_TRIGGER < 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + __syncthreads(); + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_mma(smem, tm); + M3FA_PROF_T(50, threadIdx.x); + } + __syncwarp(); + } + if (warp >= kComputeWarps) + { + op_epilogue(p, smem, tm, tid, op_u); + } + M3FA_PROF_T(37, kComputeThreads); + __syncthreads(); + M3FA_PROF_T(55, 0); + if (M3FA_MEGA_TRIGGER == 3) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // after this CTA's y rows / pushes + } + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" ::"r"(*tmem_slot), "r"(kTmemAll) + : "memory"); + M3FA_PROF_T(56, 32); + } + if (tid == 0) + { + p.flags[32 + blockIdx.x] = calls + 1u; // a plain store: no load on the exit path + } + M3FA_PROF(39); +#ifndef M3FA_FINAL_CLBAR +#define M3FA_FINAL_CLBAR 0 +#endif + if (M3FA_FINAL_CLBAR) + { + // No peer may exit while its partials can still be pushed into it (DSMEM lifetime). + asm volatile("barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;\n" :: + : "memory"); + } + M3FA_PROF(38); + return; + } + } +#ifndef M3FA_MEGA_ATTN_WAIT_FIRST +#define M3FA_MEGA_ATTN_WAIT_FIRST 0 // 1: the attention CTAs read nothing (metadata, caches) before their PDL wait +#endif +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (FQKV && M3FA_MEGA_ATTN_WAIT_FIRST) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } +#endif + if (!FQKV && warp == kWarps - 1) // a producer warp: it has nothing else to do before the PDL wait + { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n" ::"r"(smem_u32(tmem_slot)), + "r"(S::kTmemCols) + : "memory"); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n" ::: "memory"); + } + int const seq_len = p.seq_lens[req]; + int page[BPC]; +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + int const b = cr * BPC + h; + page[h] = b < p.bt_cols ? p.block_table[static_cast(req) * p.bt_stride + b] : 0; + } + // New-token cache slots, for the tokens up to this cluster's last one. + int slot[QLEN]; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + slot[t] = t < tfirst + QG ? p.out_cache_loc[tok0 + t] : -1; + } + // This warp's producer tasks and, for every lane that applies RoPE, the task token's cos/sin. + float cs[kMaxT][4], sn[kMaxT][4]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + int const t = i < kHq * QG ? tfirst + i / kHq : (i - kHq * QG) / 3; + if (i < kTasks && lane < 16 && t < QLEN) + { + int64_t const row = static_cast(p.position_ids[tok0 + t]) * kRotDim; +#pragma unroll + for (int e = 0; e < 4; ++e) + { + int const coefficient = (lane * 4 + e) & (kRotDim / 2 - 1); + cs[k][e] = p.rope[row + coefficient]; + sn[k][e] = p.rope[row + kRotDim / 2 + coefficient]; + } + } + else + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + cs[k][e] = 1.0F; + sn[k][e] = 0.0F; + } + } + } + float qw[4], kw[4], ikw[4]; + load_w4(qw, p.q_w, lane); + load_w4(kw, p.k_w, lane); + load_w4(ikw, p.ik_w, lane); + float const kqs = p.kv_quant_scale[1]; + float const vqs = p.kv_quant_scale[2]; + float const sl2e = (p.sm_scale * 1.4426950409f) * p.k_gs[0]; + float const vgs = p.v_gs[0]; + // Longest attended length among this cluster's tokens, and each block's rows within it. + int const kv_max = max(seq_len - QLEN + tfirst + QG, 0); + int nvh[BPC]; +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + nvh[h] = min(max(kv_max - (cr * BPC + h) * kPage, 0), kPage); + } + bool const any_valid = nvh[0] > 0; + if (tid == 32) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(mbar), "r"(1)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(mma_bar), "r"(1)); + asm volatile( + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(mbar), "r"(kBlocks * QG * BPC * kPushBytes) + : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + } + // This thread's packed K / V chunks, held in registers until they are dequantized after the PDL wait. + RawChunk kraw[kKPer]; +#pragma unroll + for (int k = 0; k < kKPer; ++k) + { + int const q = tid + kComputeThreads * k; + kraw[k] = RawChunk{make_uint4(0u, 0u, 0u, 0u), 0u}; + int const h = q / (kPage * 4); + if (q < kChunks && tid < kComputeThreads && nvh[h] > 0) + { + kraw[k] = load_chunk(p.kv_data + static_cast(page[h]) * p.d_page, + p.kv_scale + static_cast(page[h]) * p.s_page, chunk_row(q % (kPage * 4)), + chunk_col(q % (kPage * 4)), false); + } + } + // V goes through cp.async into shared memory: an asynchronous copy does not hold the producer warps up the way + // an outstanding register load does, and they need V only after QK has been issued. + if (tid >= kVStart) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + if (nvh[h] > 0) + { + uint8_t* dst = smem + S::kRawV + h * S::kRawVBlock; + uint8_t const* data = p.kv_data + static_cast(page[h]) * p.d_page + p.d_plane; + uint8_t const* scales = p.kv_scale + static_cast(page[h]) * p.s_page + p.s_plane; + for (int c = tid - kVStart; c < S::kRawVBlock / 16; c += kVThreads) + { + cp_async16(dst + c * 16, c < kPage * 4 ? data + c * 16 : scales + (c - kPage * 4) * 16); + } + } + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + if (warp == kWarps - 1) + { + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + // Every CTA's mbarrier is initialized before any peer pushes into it; the matching wait sits right before the + // push, so this barrier's latency hides behind the attention (with aliased partials it moves there). + if constexpr (!S::kRecvAlias) + { + asm volatile("barrier.cluster.arrive.relaxed.aligned;\n" ::: "memory"); + } + int kv_len[QG]; +#pragma unroll + for (int tl = 0; tl < QG; ++tl) + { + kv_len[tl] = max(seq_len - QLEN + tfirst + tl + 1, 0); + } + int const cta_tok0 = cr * kTok; // request position of this CTA's first block row + // New token t's row in this CTA's tiles (block * 128 + row), or -1: the CTA owns the block holding the token's + // position, and the token is attended by this cluster (t <= its last token). Every thread derives the same. + int new_row[QLEN]; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + int const pos = seq_len - QLEN + t; + int const blk = min(pos >> 7, kBlocks - 1); + int const spage = slot[t] / kPage; + bool const mine = t < tfirst + QG && pos >= 0 && blk / BPC == cr && slot[t] >= 0 && spage < p.num_pages; + new_row[t] = mine && spage == page[blk % BPC] ? (blk % BPC) * kPage + (slot[t] & (kPage - 1)) : -1; + } + auto is_new_row = [&](int r) + { + bool hit = false; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + hit |= new_row[t] == r; + } + return hit; + }; + // Cache writes of the new K/V/index-K are issued after this CTA has pushed its partials: only the fp16 tile rows + // are on the attention's path. kind 0: none, 1: NVFP4 data + scale, 2: FP8 index-K. + int wr_kind[kMaxT]; + uint8_t* wr_ptr[kMaxT]; + uint8_t* wr_sptr[kMaxT]; + uint32_t wr_val[kMaxT]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + wr_kind[k] = 0; + wr_ptr[k] = nullptr; + wr_sptr[k] = nullptr; + wr_val[k] = 0u; + } + uint32_t tmem = 0; + auto dequant_k = [&]() + { + if (tid < kComputeThreads) + { +#pragma unroll + for (int k = 0; k < kKPer; ++k) + { + int const q = tid + kComputeThreads * k; + if (q < kChunks) + { + int const h = q / (kPage * 4); + int const row = chunk_row(q % (kPage * 4)); + if (!is_new_row(h * kPage + row)) + { + store_chunk_fp16( + smem + S::kK16 + h * S::kTileBytes, row, chunk_col(q % (kPage * 4)), kraw[k], row < nvh[h]); + } + } + } + } + }; + // The producer warps' V tiles (raw tiles from cp.async, dequantized to fp16), then their PV-barrier arrival. + auto dequant_v = [&]() + { +#ifdef M3FA_DBG_NO_VDQ + if constexpr (!kRegCore) + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + return; +#endif + asm volatile("cp.async.wait_all;\n" ::: "memory"); + asm volatile("bar.sync 4, %0;\n" ::"n"(kVThreads) : "memory"); +#pragma unroll + for (int k = 0; k < kVPer; ++k) + { + int const q = tid - kVStart + kVThreads * k; + if (q < kChunks) + { + int const h = q / (kPage * 4); + int const row = chunk_row(q % (kPage * 4)); + int const cc = chunk_col(q % (kPage * 4)); + if (!is_new_row(h * kPage + row)) + { + uint8_t const* rawv = smem + S::kRawV + h * S::kRawVBlock; + uint8_t const* sc = rawv + kPage * 64 + (row >> 2) * 32 + (row & 3); + RawChunk c; + c.data = *reinterpret_cast(rawv + row * 64 + cc * 16); + c.scales + = static_cast(sc[(cc * 2) * 4]) | (static_cast(sc[(cc * 2 + 1) * 4]) << 8); + store_chunk_fp16(smem + S::kV16 + h * S::kTileBytes, row, cc, c, row < nvh[h]); + } + } + } + M3FA_FENCE_PROXY_ASYNC(); // V tiles -> tensor core reads (ordered by the barrier before the PV issue) + if constexpr (!kRegCore) + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + M3FA_PROF_T(16, kComputeThreads); + M3FA_PROF_T(17, kThreads - 1); + }; + if constexpr (FQKV) + { + qkv_stream(p, smem, gbar, tmem_slot, warp, lane, qkv_u); + } + if constexpr (!FQKV) + { + if (M3FA_V15_EARLYV && any_valid && tid >= kVStart) + { + dequant_v(); + } + if constexpr (kEarlyClw && !S::kRecvAlias) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); + } + } + M3FA_PROF(11); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (FQKV && M3FA_MEGA_TRIGGER == 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#if M3FA_TRIGGER_AT == 1 + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +#endif +#endif + // FQKV (megakernel): this call's QKV rows sit in buffer (calls & 1), polled value by value below. + unsigned int const mega_calls = FQKV ? p.flags[32 + blockIdx.x] : 0u; + __nv_bfloat16 const* qkv_src = FQKV ? p.packed + (mega_calls & 1u) * kMaxTok * kQkvN : p.packed; + M3FA_PROF(1); + int const par = 0; + if constexpr (FQKV) + { + if (qkv_u >= 0) + { + qkv_mma(p, smem, gbar, tmem_slot, tid, qkv_u); + } + } + + if constexpr (FQKV) + { + if (qkv_u >= 0) + { + if (warp >= kComputeWarps) + { + qkv_epilogue(p, smem, gbar, tmem_slot, tid, qkv_u, par, mega_calls); + } + else + { + // The attention regions overlay the weight ring: every weight byte has gone on into TMEM (or, with + // the MMAs reading the ring, every MMA has completed) first. + mbar_wait_cta(gbar + 16 * kQkvRing + (kQkvASmem ? 8 : 0), 0); + } + } + if (any_valid) + { + dequant_k(); + } +#ifndef M3FA_MEGA_EARLYV +#define M3FA_MEGA_EARLYV 1 // attention CTAs dequantize V here, while the QKV rows are still on their way +#endif +#ifndef M3FA_MEGA_EARLYCLW +#define M3FA_MEGA_EARLYCLW 1 // and pass the cluster barrier (peers' merge mbarriers initialized) here too +#endif + if (M3FA_MEGA_EARLYV && any_valid && tid >= kVStart) + { + dequant_v(); + } + if constexpr (kEarlyClw && !S::kRecvAlias) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); + } + M3FA_PROF(14); + } + + if (any_valid) // CTA-uniform + { + // ---------------- Producer ---------------- + // Q rows of this cluster's tokens in every CTA. The new K/V of every token up to this cluster's last one + // whose position falls in this CTA's blocks go straight into the fp16 tiles; the CTA writes the caches, and + // the index-K, only for its own tokens. + uint2 raw[kMaxT]; + float fq[FQKV ? kMaxT : 1][4]; + bool act[kMaxT]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + raw[k] = make_uint2(0u, 0u); + act[k] = false; + if (i < kTasks) + { + int t, col; + bool active = true; + if (i >= kHq * QG) + { + int const kv = i - kHq * QG; + t = kv / 3; + int const kind = kv % 3; + col = kind == 0 ? kColK : (kind == 1 ? kColV : kColIdxK); + int const pos = seq_len - QLEN + t; + active = t < tfirst + QG && pos >= 0 && min(pos >> 7, kBlocks - 1) / BPC == cr && slot[t] >= 0 + && slot[t] / kPage < p.num_pages && (kind != 2 || t >= tfirst); + } + else + { + t = tfirst + i / kHq; + col = i % kHq; + } + act[k] = active; + if (active) + { + uint2 const* src = reinterpret_cast( + qkv_src + static_cast(tok0 + t) * p.packed_stride + col * kHeadDim + lane * 4); + if constexpr (FQKV) + { + uint2 v; + bool pending; + do + { + asm volatile("ld.relaxed.gpu.global.v2.b32 {%0, %1}, [%2];\n" + : "=r"(v.x), "=r"(v.y) + : "l"(src) + : "memory"); + pending = (v.x & 0xFFFFu) == 0xFFFFu || (v.x >> 16) == 0xFFFFu || (v.y & 0xFFFFu) == 0xFFFFu + || (v.y >> 16) == 0xFFFFu; + } while (pending); + raw[k] = v; + } + else + { + raw[k] = *src; + } + } + } + } + M3FA_PROF(15); + // While the packed-QKV loads are in flight: dequantize this thread's K chunks to fp16 (exact: e2m1 x e4m3 + // fits fp16). Rows past the longest attended length are zeroed, as the Triton producer masks its loads; the + // new tokens' rows are written by their producer tasks instead. (Building the tiles before the PDL wait was + // measured slower: the cache loads issued at entry are still in flight there, so the wait moves later.) + if constexpr (!FQKV) + { + dequant_k(); + } +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + if (i >= kTasks || !act[k]) + { + continue; + } + float e[4]; + { + float2 const a = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw[k].x)); + float2 const b = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw[k].y)); + e[0] = a.x; + e[1] = a.y; + e[2] = b.x; + e[3] = b.y; + } + if (i < kHq * QG) + { + norm_rope(e, qw, p.eps, cs[k], sn[k], lane); + __nv_fp8x2_e4m3 const lo(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const hi(make_float2(e[2], e[3])); + uint2 hv; + hv.x = e4m3x2_to_h2(lo.__x); + hv.y = e4m3x2_to_h2(hi.__x); + *reinterpret_cast(smem + S::kQ + kmaj_off(i, lane * 4)) = hv; // row i = token * 16 + head + continue; + } + int const t = (i - kHq * QG) / 3; + int const kind = (i - kHq * QG) % 3; + int const spage = slot[t] / kPage; + int const within = slot[t] & (kPage - 1); + if (kind == 2) + { + norm_rope(e, ikw, p.eps, cs[k], sn[k], lane); + __nv_fp8x2_e4m3 const lo(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const hi(make_float2(e[2], e[3])); + wr_kind[k] = 2; + wr_val[k] = static_cast(lo.__x) | (static_cast(hi.__x) << 16); + wr_ptr[k] = p.idx_k + static_cast(spage) * p.i_page + static_cast(within) * p.i_tok + + lane * 4; + continue; + } + bool const isV = kind == 1; + if (!isV) + { + norm_rope(e, kw, p.eps, cs[k], sn[k], lane); + } + uint8_t sbyte = 0; + uint16_t const fp4 = quantizeMinimaxM3Fp4x4(e, isV ? vqs : kqs, &sbyte); + if (t >= tfirst) + { + int const scale_col = lane / 4; + int const scale_within = isV ? (within / 4) * 32 + scale_col * 4 + within % 4 : within * 8 + scale_col; + wr_kind[k] = 1; + wr_val[k] = static_cast(fp4) | (static_cast(sbyte) << 16); + wr_ptr[k] = p.kv_data + static_cast(spage) * p.d_page + (isV ? p.d_plane : 0) + within * 64 + + lane * 2; + wr_sptr[k] = p.kv_scale + static_cast(spage) * p.s_page + (isV ? p.s_plane : 0) + scale_within; + } + if (new_row[t] >= 0) + { + // The attention reads this page: its fp16 tile row is the freshly quantized row, dequantized like + // the rest of the tile. + int const h = new_row[t] / kPage; + int const row = new_row[t] % kPage; + uint32_t const sc2 = e4m3_to_h2_splat(sbyte); + uint2 hv; + hv.x = hmul2_u(e2m1x2_to_h2(fp4 & 0xFFu), sc2); + hv.y = hmul2_u(e2m1x2_to_h2(fp4 >> 8), sc2); + *reinterpret_cast( + smem + (isV ? S::kV16 : S::kK16) + h * S::kTileBytes + kmaj_off(row, lane * 4)) + = hv; + } + } + // K tiles, Q rows and patched rows -> tensor core reads. Only threads that wrote fence: the fence also + // waits for the thread's outstanding loads, and the producer warps still have their V chunks in flight. + bool wrote = tid < kComputeThreads; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + wrote |= act[k]; + } + if (wrote) + { + M3FA_FENCE_PROXY_ASYNC(); + } + M3FA_PROF(13); + __syncthreads(); +#if M3FA_TRIGGER_AT == 2 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Every K/V byte this CTA reads has arrived (the V chunks sit in registers): release the dependents now so + // their weight streams do not queue ahead of the KV reads (the megakernel releases late, from its o_proj). + if constexpr (!FQKV || M3FA_MEGA_TRIGGER == 1) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + tmem = *tmem_slot; + M3FA_PROF(2); + // ---------------- S^T = K Q^T per block on the tensor core: TMEM lane = token, column = query row ---------- +#ifndef M3FA_MEGA_WARPMMA +#define M3FA_MEGA_WARPMMA 1 // attention MMAs issued by warp 0's elected lane with warp-uniform operands (not tid 0) +#endif + if (!kRegCore && (kWarpMma ? warp == 0 : tid == 0)) + { + // Operands computed by the whole warp stay in uniform registers; one elected lane issues. + uint32_t const tm_u = tmem; + bool const issue = kWarpMma ? elect_one() : true; + if (issue) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + uint32_t const a0 = smem_u32(smem + S::kK16 + h * S::kTileBytes); + uint32_t const b0 = smem_u32(smem + S::kQ); +#pragma unroll + for (int ks = 0; ks < kHeadDim / 16; ++ks) + { + umma_f16(tm_u + h * kRows, umma_desc(a0 + ks * 256, 128, 2048), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 0, 0), ks > 0); + } + } + umma_commit(mma_bar); + } + if (kWarpMma) + { + __syncwarp(); + } + } + if (tid >= kVStart && !(FQKV && M3FA_MEGA_EARLYV) && !(!FQKV && M3FA_V15_EARLYV)) + { + // The producer warps dequantize V while QK and the softmax run. + dequant_v(); + } + if constexpr (kRegCore) + { + // ---------------- Register core: warp w of block h = w / kWpb owns 8 * kNt tokens for S and 8 * kNt + // output dims for O; m16n8k16 fp16 MMAs with fp32 accumulators, the softmax in the S fragments ---------- + constexpr int kWpb = kComputeWarps / BPC; // warps per block + constexpr int kNt = 16 / kWpb; // 8-token (S) or 8-dim (O) tiles per warp + float* redx_m = reinterpret_cast(smem + S::kRedX); + float* redx_l = redx_m + BPC * kRows * kWpb; + if (warp < kComputeWarps) + { + int const h = warp / kWpb; + int const wib = warp % kWpb; + int const g = lane >> 2; + int const q4 = lane & 3; + uint32_t const ktile = smem_u32(smem + S::kK16 + h * S::kTileBytes); + uint32_t const vtile = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const qrows = smem_u32(smem + S::kQ); + uint8_t* ptile = smem + S::kK16 + h * S::kTileBytes; // P over this block's K tile, after QK + auto block_bar = [&]() + { + if (h == 0) + { + asm volatile("bar.sync 14, %0;\n" ::"n"(kWpb * 32) : "memory"); + } + else + { + asm volatile("bar.sync 15, %0;\n" ::"n"(kWpb * 32) : "memory"); + } + }; + // S = Q K^T: rows (c, g / g + 8), tokens 8 tb + 2 q4 (+1). + float sv[QG][kNt][4]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + sv[c][nt][e] = 0.0F; + } + } + } +#pragma unroll + for (int ks = 0; ks < kHeadDim / 16; ++ks) + { + uint32_t bk[kNt][2]; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const tb = wib * kNt + nt; + ldsm_x2(ktile + tb * 2048 + (ks * 2 + ((lane >> 3) & 1)) * 128 + (lane & 7) * 16, bk[nt]); + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { + uint32_t aq[4]; + ldsm_x4( + qrows + (c * 2 + ((lane >> 3) & 1)) * 2048 + (ks * 2 + (lane >> 4)) * 128 + (lane & 7) * 16, + aq); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + mma_f16_16816(sv[c][nt], aq, bk[nt]); + } + } + } + M3FA_PROF(7); + // Row maxima (log2 domain, masked as the producer masks), quad then across the block's warps. + float mrow[QG][2]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { + float mx0 = kNegInf, mx1 = kNegInf; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + int const tok = cta_tok0 + h * kPage + (wib * kNt + nt) * 8 + 2 * q4 + (e & 1); + bool const valid = tok < kv_max && tok < kv_len[c]; + float const v = valid ? sv[c][nt][e] * sl2e : kNegInf; + sv[c][nt][e] = v; + if (e < 2) + { + mx0 = fmaxf(mx0, v); + } + else + { + mx1 = fmaxf(mx1, v); + } + } + } + mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 1)); + mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 2)); + mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 1)); + mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 2)); + if (q4 == 0) + { + redx_m[(h * kRows + c * 16 + g) * kWpb + wib] = mx0; + redx_m[(h * kRows + c * 16 + g + 8) * kWpb + wib] = mx1; + } + } + block_bar(); // every warp of the block is past QK (P may overwrite the K tile) and wrote its maxima + M3FA_PROF(8); +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int half = 0; half < 2; ++half) + { + float const* rm = redx_m + (h * kRows + c * 16 + g + 8 * half) * kWpb; + float m = kNegInf; +#pragma unroll + for (int w4 = 0; w4 < kWpb; w4 += 4) + { + float4 const v = *reinterpret_cast(rm + w4); + m = fmaxf(m, fmaxf(fmaxf(v.x, v.y), fmaxf(v.z, v.w))); + } + mrow[c][half] = m; + } + // P = exp2(s - m) as fp16 into [row][token] core matrices; the fp32 values feed the row sums. + float l0 = 0.0F, l1 = 0.0F; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + float a[4]; +#pragma unroll + for (int e = 0; e < 4; ++e) + { + float const m = mrow[c][e >> 1]; + a[e] = m == kNegInf ? 0.0F : ex2(sv[c][nt][e] - m); + } + l0 += a[0] + a[1]; + l1 += a[2] + a[3]; + int const t0 = (wib * kNt + nt) * 8 + 2 * q4; + int const r0 = c * 16 + g; + __half2 const p0 = __floats2half2_rn(a[0], a[1]); + __half2 const p1 = __floats2half2_rn(a[2], a[3]); + *reinterpret_cast<__half2*>( + ptile + (r0 >> 3) * 2048 + (t0 >> 3) * 128 + (r0 & 7) * 16 + (t0 & 7) * 2) + = p0; + *reinterpret_cast<__half2*>( + ptile + ((r0 + 8) >> 3) * 2048 + (t0 >> 3) * 128 + (r0 & 7) * 16 + (t0 & 7) * 2) + = p1; + } + l0 += __shfl_xor_sync(0xffffffffu, l0, 1); + l0 += __shfl_xor_sync(0xffffffffu, l0, 2); + l1 += __shfl_xor_sync(0xffffffffu, l1, 1); + l1 += __shfl_xor_sync(0xffffffffu, l1, 2); + if (q4 == 0) + { + redx_l[(h * kRows + c * 16 + g) * kWpb + wib] = l0; + redx_l[(h * kRows + c * 16 + g + 8) * kWpb + wib] = l1; + } + } + M3FA_PROF_T(19, 0); + } + __syncthreads(); // P, the V tiles (producer warps) and the row sums visible + M3FA_PROF(9); + if (warp < kComputeWarps) + { + int const h = warp / kWpb; + int const wib = warp % kWpb; + int const g = lane >> 2; + int const q4 = lane & 3; + uint32_t const vtile = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const ptile = smem_u32(smem + S::kK16 + h * S::kTileBytes); + float* obuf = reinterpret_cast(smem + S::kK16 + h * S::kTileBytes + S::kObufOff); + // O = P V: rows (c, g / g + 8), dims 8 dt + 2 q4 (+1). + float ov[QG][kNt][4]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + ov[c][nt][e] = 0.0F; + } + } + } +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + uint32_t bv[kNt][2]; + int const tok = ks * 16 + ((lane >> 3) & 1) * 8 + (lane & 7); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const dt = wib * kNt + nt; + ldsm_x2_trans(vtile + (tok >> 3) * 2048 + dt * 128 + (tok & 7) * 16, bv[nt]); + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { + uint32_t ap[4]; + ldsm_x4( + ptile + (c * 2 + ((lane >> 3) & 1)) * 2048 + (ks * 2 + (lane >> 4)) * 128 + (lane & 7) * 16, + ap); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + mma_f16_16816(ov[c][nt], ap, bv[nt]); + } + } + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const d = (wib * kNt + nt) * 8 + 2 * q4; + int const r0 = c * 16 + g; + *reinterpret_cast(obuf + r0 * kOPitch + d) = make_float2(ov[c][nt][0], ov[c][nt][1]); + *reinterpret_cast(obuf + (r0 + 8) * kOPitch + d) + = make_float2(ov[c][nt][2], ov[c][nt][3]); + } + } + } + else if (tid - kComputeThreads < BPC * kRows) + { + // Row statistics of every (block, row): LSE and the PV scale, as the producer derives them. + int const r = tid - kComputeThreads; + float const* rm = redx_m + r * kWpb; + float const* rl = redx_l + r * kWpb; + float m = kNegInf; + float l = 0.0F; +#pragma unroll + for (int w4 = 0; w4 < kWpb; w4 += 4) + { + float4 const vm = *reinterpret_cast(rm + w4); + float4 const vl = *reinterpret_cast(rl + w4); + m = fmaxf(m, fmaxf(fmaxf(vm.x, vm.y), fmaxf(vm.z, vm.w))); + l += (vl.x + vl.y) + (vl.z + vl.w); + } + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[r] = scale * vgs; + row_lse[r] = lse; + } + __syncthreads(); // O partials and row statistics visible to the push + } + else + { + if constexpr (kRs2) + { + // ---------------- Split softmax: warps 0..3 rows 0..7, warps 4..7 rows 8..15 (thread = block token) + // -------- + if (warp < 8) + { + int const rg = warp / 4; + int const wq = warp % 4; + int const tk = wq * 32 + lane; + auto rg_bar = [&]() + { + if (rg == 0) + { + asm volatile("bar.sync 10, 128;\n" ::: "memory"); + } + else + { + asm volatile("bar.sync 11, 128;\n" ::: "memory"); + } + }; + int const rs8 = ((lane >> 4) & 1) * 4 + ((lane >> 3) & 1) * 2 + ((lane >> 2) & 1); + mbar_wait_cta(mma_bar, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF(7); + float v[8]; + tmem_ld8(tmem + ((wq * 32) << 16) + 8 * rg, v); + bool const valid = cta_tok0 + tk < kv_max && cta_tok0 + tk < kv_len[0]; + float a[8]; +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] = valid ? v[j2] * sl2e : kNegInf; + a[j2] = v[j2]; + } + float const mw = warp_reduce_scatter8(a, lane); + if ((lane & 3) == 0) + { + red_m[(8 * rg + rs8) * 4 + wq] = mw; + } + rg_bar(); + M3FA_PROF(8); + uint8_t* pdst = smem + S::kP + (tk >> 3) * 128 + (tk & 7) * 16; + float mrow = kNegInf; + if (lane < 8) + { + float4 const mq = *reinterpret_cast(red_m + (8 * rg + lane) * 4); + mrow = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + } + float pa[8]; +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + float const m = __shfl_sync(0xffffffffu, mrow, j2); + pa[j2] = m == kNegInf ? 0.0F : ex2(v[j2] - m); + } + if (tid == 0) + { + M3FA_CLK(41); + } + uint32_t pw[4]; +#pragma unroll + for (int j2 = 0; j2 < 4; ++j2) + { + __half2 const hp = __floats2half2_rn(pa[2 * j2], pa[2 * j2 + 1]); + pw[j2] = *reinterpret_cast(&hp); + } + *reinterpret_cast(pdst + rg * 2048) = make_uint4(pw[0], pw[1], pw[2], pw[3]); + M3FA_PROF_T(18, 0); + if (tid == 0) + { + M3FA_CLK(42); + } + M3FA_FENCE_PROXY_ASYNC(); // P -> tensor core reads + if (tid == 0) + { + M3FA_CLK(43); + } + if (warp == 0) + { + asm volatile("bar.sync 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + if (tid == 0) + { + M3FA_CLK(44); + } + } + else + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + if (kWarpMma ? warp == 0 : tid == 0) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm_u = tmem; + uint32_t const a0 = smem_u32(smem + S::kV16); + uint32_t const b0 = smem_u32(smem + S::kP); + if (kWarpMma ? elect_one() : true) + { +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + umma_f16(tm_u + BPC * kRows, umma_desc(a0 + ks * 4096, 2048, 128), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 1, 1), ks > 0); + } + umma_commit(mma_bar); + } + if (kWarpMma) + { + __syncwarp(); + } + } + M3FA_PROF_T(19, 0); + float const lw = warp_reduce_scatter8(pa, lane); + if ((lane & 3) == 0) + { + red_l[(8 * rg + rs8) * 4 + wq] = lw; + } + rg_bar(); + M3FA_PROF_T(20, 0); + if (wq == 0 && lane < 8) + { + int const r = 8 * rg + lane; + float4 const mq = *reinterpret_cast(red_m + r * 4); + float4 const lq = *reinterpret_cast(red_l + r * 4); + float const m = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + float const l = (lq.x + lq.y) + (lq.z + lq.w); + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[r] = scale * vgs; + row_lse[r] = lse; + } + } + __syncthreads(); // row scales and LSEs visible (PV was issued above) + M3FA_PROF(9); + if (warp < 8) + { + mbar_wait_cta(mma_bar, 1); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + } + else + { + // ---------------- Softmax: thread = block token; the (block, 16-row chunk) pairs are spread over the + // four groups of four compute warps (each group covers the 128 TMEM lanes = block tokens) + // ---------------- + if (warp < kComputeWarps && warp / 4 < kPairs) + { + int const g = warp / 4; + int const wq = warp % 4; + int const tk = wq * 32 + lane; + auto group_bar = [&]() + { + if (g == 0) + { + asm volatile("bar.sync 10, 128;\n" ::: "memory"); + } + else if (g == 1) + { + asm volatile("bar.sync 11, 128;\n" ::: "memory"); + } + else if (g == 2) + { + asm volatile("bar.sync 12, 128;\n" ::: "memory"); + } + else + { + asm volatile("bar.sync 13, 128;\n" ::: "memory"); + } + }; + // Row r(lane) this lane holds after the reduce-scatter: bits 4..1 of the lane. + int const rs_row + = ((lane >> 4) & 1) * 8 + ((lane >> 3) & 1) * 4 + ((lane >> 2) & 1) * 2 + ((lane >> 1) & 1); + mbar_wait_cta(mma_bar, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF(7); + float v[kPairsPerGroup][16]; +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + tmem_ld16(tmem + ((wq * 32) << 16) + h * kRows + c * 16, v[k]); + bool const valid + = cta_tok0 + h * kPage + tk < kv_max && cta_tok0 + h * kPage + tk < kv_len[c]; + float a[16]; +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[k][j] = valid ? v[k][j] * sl2e : kNegInf; + a[j] = v[k][j]; + } + float const mw = warp_reduce_scatter16(a, lane); + if ((lane & 1) == 0) + { + red_m[((h * QG + c) * 16 + rs_row) * 4 + wq] = mw; + } + } + } + group_bar(); + M3FA_PROF(8); + float pa[kPairsPerGroup][16]; +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + uint8_t* pdst = smem + S::kP + h * kRows * kPage * 2 + (tk >> 3) * 128 + (tk & 7) * 16; + // Row maxima: lane j (< 16) reduces row j's four warp partials, then broadcasts it. + float mrow = kNegInf; + if (lane < 16) + { + float4 const mq + = *reinterpret_cast(red_m + ((h * QG + c) * 16 + lane) * 4); + mrow = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + } +#pragma unroll + for (int j = 0; j < 16; ++j) + { + float const m = __shfl_sync(0xffffffffu, mrow, j); + pa[k][j] = m == kNegInf ? 0.0F : ex2(v[k][j] - m); + } + if (tid == 0) + { + M3FA_CLK(41); + } + uint32_t pw[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) + { + __half2 const hp = __floats2half2_rn(pa[k][2 * j], pa[k][2 * j + 1]); + pw[j] = *reinterpret_cast(&hp); + } + // P^T as the MN-major B operand: 8 query rows of one token are one 16-byte core-matrix row. + *reinterpret_cast(pdst + (2 * c) * 2048) = make_uint4(pw[0], pw[1], pw[2], pw[3]); + *reinterpret_cast(pdst + (2 * c + 1) * 2048) + = make_uint4(pw[4], pw[5], pw[6], pw[7]); + M3FA_PROF_T(18, 0); + } + } + if (tid == 0) + { + M3FA_CLK(42); + } + M3FA_FENCE_PROXY_ASYNC(); // P -> tensor core reads + if (tid == 0) + { + M3FA_CLK(43); + } + // PV starts as soon as P and the V tiles are in: thread 0 waits for the softmax warps and the + // producer warps (V); the row sums and statistics below overlap the MMAs. + if (warp == 0) + { + asm volatile("bar.sync 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + else + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + if (kWarpMma ? warp == 0 : tid == 0) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm_u = tmem; + if (kWarpMma ? elect_one() : true) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + uint32_t const a0 = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const b0 = smem_u32(smem + S::kP + h * kRows * kPage * 2); +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + umma_f16(tm_u + (BPC + h) * kRows, umma_desc(a0 + ks * 4096, 2048, 128), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 1, 1), ks > 0); + } + } + umma_commit(mma_bar); + } + if (kWarpMma) + { + __syncwarp(); + } + } + M3FA_PROF_T(19, 0); +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + float const lw = warp_reduce_scatter16(pa[k], lane); + if ((lane & 1) == 0) + { + red_l[((h * QG + c) * 16 + rs_row) * 4 + wq] = lw; + } + } + } + group_bar(); + M3FA_PROF_T(20, 0); + if (wq == 0 && lane < 16) + { +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + int const r = (h * QG + c) * 16 + lane; + float4 const mq = *reinterpret_cast(red_m + r * 4); + float4 const lq = *reinterpret_cast(red_l + r * 4); + float const m = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + float const l = (lq.x + lq.y) + (lq.z + lq.w); + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[h * kRows + c * 16 + lane] = scale * vgs; + row_lse[h * kRows + c * 16 + lane] = lse; + } + } + } + } + __syncthreads(); // row scales and LSEs visible (PV was issued above) + M3FA_PROF(9); + if (warp < kComputeWarps && warp / 4 < kPairs) + { + mbar_wait_cta(mma_bar, 1); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + } + } + M3FA_PROF(3); + } + else + { + asm volatile("cp.async.wait_all;\n" ::: "memory"); +#if M3FA_TRIGGER_AT == 2 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (!FQKV || M3FA_MEGA_TRIGGER == 1) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + for (int r = tid; r < BPC * kRows; r += kThreads) + { + row_lse[r] = kNegInf; + row_scale[r] = 0.0F; + } + __syncthreads(); + } + + // ---------------- Push every partial row to the CTA that merges its head (CTA r merges heads r*BPC..) -------- + if constexpr (S::kRecvAlias) + { + // Peers' partials overlay this CTA's tiles: every CTA's PV MMAs have completed. + asm volatile("barrier.cluster.arrive.aligned;\n" ::: "memory"); + } + if (!(kEarlyClw && !S::kRecvAlias)) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); // peers' mbarriers are initialized + } + M3FA_PROF(4); + uint32_t const recv_base = smem_u32(recv); + if (kRs2 && warp < 8) + { + // Split push: warp group rg pushes rows 8 rg .. 8 rg + 7 (row j goes to CTA j, which merges head j). + int const rg = warp / 4; + int const d = (warp % 4) * 32 + lane; + float v[8]; + if (any_valid) + { + tmem_ld8(tmem + ((warp % 4) * 32 << 16) + BPC * kRows + 8 * rg, v); + } + else + { +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] = 0.0F; + } + } +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] *= row_scale[8 * rg + j2]; + } +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + uint32_t const dst_cta = 8 * rg + j2; + uint32_t const row_addr = recv_base + static_cast((cr * kOPitch + d) * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr, dst_cta)), + "r"(__float_as_uint(v[j2])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + else if (!kRs2 && warp < kComputeWarps && warp / 4 < kPairs) + { + int const g = warp / 4; + int const d = (warp % 4) * 32 + lane; // output dim +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + int const chunk = cr * BPC + h; + float v[16]; + if (any_valid && kRegCore) + { + float const* obuf + = reinterpret_cast(smem + S::kK16 + h * S::kTileBytes + S::kObufOff); +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] = obuf[(c * 16 + j) * kOPitch + d]; + } + } + else if (any_valid) + { + tmem_ld16(tmem + ((warp % 4) * 32 << 16) + (BPC + h) * kRows + c * 16, v); + } + else + { +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] = 0.0F; + } + } +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] *= row_scale[h * kRows + c * 16 + j]; + } + // One scalar store per row: a vector store's lanes would address several destination CTAs at once, + // which measured several times slower. +#pragma unroll + for (int j = 0; j < 16; ++j) + { + uint32_t const dst_cta = j / BPC; + uint32_t const row_addr + = recv_base + static_cast((((chunk * QG + c) * BPC + j % BPC) * kOPitch + d) * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr, dst_cta)), + "r"(__float_as_uint(v[j])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + } + } + else if (tid >= kComputeThreads) + { + // Every partial row's LSE (BPC * kRows of them; more than the producer threads for the widest geometries). + for (int r = tid - kComputeThreads; r < BPC * kRows; r += kThreads - kComputeThreads) + { + int const h = r / kRows; + int const c = (r % kRows) / 16; + int const j = r % 16; + uint32_t const dst_cta = j / BPC; + uint32_t const row_addr + = recv_base + static_cast((((cr * BPC + h) * QG + c) * BPC + j % BPC) * kOPitch * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr + kHeadDim * 4, dst_cta)), + "r"(__float_as_uint(row_lse[r])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + M3FA_PROF(10); +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + if (wr_kind[k] == 1) + { + *reinterpret_cast(wr_ptr[k]) = static_cast(wr_val[k] & 0xFFFFu); + if (lane % 4 == 0) + { + *wr_sptr[k] = static_cast(wr_val[k] >> 16); + } + } + else if (wr_kind[k] == 2) + { + *reinterpret_cast(wr_ptr[k]) = wr_val[k]; + } + } + { + uint32_t ready = 0; +#ifdef M3FA_DBG_HANG + long long spins = 0; +#endif + while (!ready) + { +#ifdef M3FA_DBG_HANG + if (++spins == (1LL << 24)) + { + printf("M3FA hang: block %d thread %d merge recv\n", blockIdx.x, threadIdx.x); + } +#endif + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(mbar) + : "memory"); + } + } + M3FA_PROF(5); + + // ---------------- Merge heads [cr * BPC, cr * BPC + BPC) (split-K merge kernel's LSE weighting) ---------------- + constexpr int kMergeRows = QG * BPC; // (token, local head) pairs, 16 chunk weights each + // Wide geometries: chunk weights once per merged row (into the softmax scratch, free by now), then the weighted + // sums. Narrow ones recompute the weights per output element: no extra barrier on the critical path. + constexpr bool kRowWeights = kMergeRows >= 4; + float* wmerge = red_m; // [kMergeRows][kBlocks]; red_m holds BPC * kRows * 4 = kMergeRows * 64 floats + [[maybe_unused]] auto chunk_weights = [&](int mr, float(&lc)[kBlocks]) + { + float lmax = kNegInf; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + lc[c] = recv[(c * kMergeRows + mr) * kOPitch + kHeadDim]; + lmax = fmaxf(lmax, lc[c]); + } + lmax = lmax == kNegInf ? 0.0F : lmax; + float denom = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + lc[c] = ex2(lc[c] - lmax); // now the chunk weight; 0 for an empty chunk + denom += lc[c]; + } + denom = denom > 0.0F ? denom : 1.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + // -0.0 marks an empty chunk (skipped below, as the merge kernel skips zero weights). + lc[c] = lc[c] != 0.0F ? __fdividef(lc[c], denom) : -0.0F; + } + }; + if constexpr (kRowWeights) + { + for (int mr = tid; mr < kMergeRows; mr += kThreads) + { + float lc[kBlocks]; + chunk_weights(mr, lc); +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wmerge[mr * kBlocks + c] = lc[c]; + } + } + __syncthreads(); + } +#pragma unroll + for (int idx = tid; idx < kMergeRows * kHeadDim; idx += kThreads) + { + int const mr = idx >> 7; + int const d = idx & (kHeadDim - 1); + float wn[kBlocks]; + if constexpr (kRowWeights) + { +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wn[c] = wmerge[mr * kBlocks + c]; + } + } + else + { +#ifndef M3FA_MEGA_MERGE_LOCAL +#define M3FA_MEGA_MERGE_LOCAL \ + 0 // 1 (measured +0.6 us slower): every merging thread derives the 16 chunk weights itself (broadcast shared loads, + // independent exps): no shuffle chain; the same arithmetic as chunk_weights +#endif + if constexpr (M3FA_MEGA_MERGE_LOCAL != 0) + { + chunk_weights(mr, wn); + } + else + { + // Narrow geometries: the warp (32 dims of one merged row) derives the 16 chunk weights together, lane c + // holding chunk c's, and broadcasts them. + float lc = (lane < kBlocks) ? recv[(lane * kMergeRows + mr) * kOPitch + kHeadDim] : kNegInf; + float lmax = lc; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + lmax = fmaxf(lmax, __shfl_xor_sync(0xffffffffu, lmax, o)); + } + lmax = lmax == kNegInf ? 0.0F : lmax; + float const wc = lane < kBlocks ? ex2(lc - lmax) : 0.0F; + // The denominator in chunk order from 0.0, as chunk_weights (and the merge kernel) add it: a butterfly + // reduction rounds differently in the last bit. + float denom = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + denom += __shfl_sync(0xffffffffu, wc, c); + } + denom = denom > 0.0F ? denom : 1.0F; + float const wl = wc != 0.0F ? __fdividef(wc, denom) : -0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wn[c] = __shfl_sync(0xffffffffu, wl, c); + } + } + } + if (tid == 0) + { + M3FA_CLK(45); + } + float acc = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + float const prod = recv[(c * kMergeRows + mr) * kOPitch + d] * wn[c]; + acc = __float_as_uint(wn[c]) != 0x80000000u ? acc + prod : acc; + } + if (tid == 0) + { + M3FA_CLK(46); + } + int const tl = mr / BPC; + int const head = cr * BPC + mr % BPC; + __nv_bfloat16 const ob = __float2bfloat16_rn(acc); + int const out_row = tok0 + tfirst + tl; + p.out[static_cast(out_row) * p.out_stride + head * kHeadDim + d] = ob; + if constexpr (MXFP8) + { + float const x = __bfloat162float(ob); + float amax = fabsf(x); +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, mask)); + } + float const sfv = amax * rcp_approx_ftz(448.0F); + int const bits = __float_as_int(sfv); + int ex = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0 ? 1 : 0); + ex = min(ex, 254); + float scale = __int_as_float((254 - ex) << 23); + scale = amax != 0.0F ? scale : 0.0F; + __nv_fp8_e4m3 const q8v(x * scale); + int64_t const q8off = static_cast(out_row) * (p.num_sf_k_tiles * 128) + head * kHeadDim + d; + if constexpr (FQKV) + { + // Megakernel hand-off to the o_proj units (Lamport): this call's parity buffer, the other one reset to + // 0xFF. 0xFF is never written: a NaN code is canonicalized to 0x7F; the UE8M0 code is at most 254. + uint8_t const code = q8v.__x == 0xFFu ? static_cast(0x7Fu) : q8v.__x; + int64_t const cur = (mega_calls & 1u) * kMaxTok * kOpK; + int64_t const nxt = ((mega_calls & 1u) ^ 1u) * kMaxTok * kOpK; + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.q8 + cur + q8off), "h"(static_cast(code)) + : "memory"); + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.q8 + nxt + q8off), "h"(static_cast(0xFFu)) + : "memory"); + } + else + { + p.q8[q8off] = q8v.__x; + } + if (lane == 0) + { + int const k_sf = head * (kHeadDim / 32) + (d >> 5); + int64_t const off = static_cast(out_row / 128) * p.num_sf_k_tiles * 512 + (k_sf / 4) * 512 + + (out_row % 32) * 16 + ((out_row % 128) / 32) * 4 + k_sf % 4; + if constexpr (FQKV) + { + int64_t const cur = (mega_calls & 1u) * kOpSfBytes; + int64_t const nxt = ((mega_calls & 1u) ^ 1u) * kOpSfBytes; + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.sf + cur + off), "h"(static_cast(ex)) + : "memory"); + asm volatile("st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.sf + nxt + off), + "h"(static_cast(0xFFu)) + : "memory"); + } + else + { + p.sf[off] = static_cast(ex); + } + } + } + } + M3FA_PROF(60); + if (any_valid) + { + // Every TMEM read completed above (tcgen05.wait::ld); release the columns. + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + __syncthreads(); + if (FQKV && tid == 0) + { + p.flags[32 + blockIdx.x] = mega_calls + 1u; + } + M3FA_PROF(61); + if (warp == kWarps - 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" ::"r"(*tmem_slot), "r"(S::kTmemCols) + : "memory"); + } + M3FA_PROF(6); +#else + __trap(); +#endif // MINIMAX_M3_ATTN_DEVICE_ENABLED +} + +template +cudaError_t launch_impl(Params const& prm, int num_reqs, bool pdl, cudaStream_t stream, int* max_clusters) +{ + auto kernel = m3_fused_decode_attn_kernel; + constexpr int smem = SmemX::kBytes; + constexpr int kCs = kBlocks / BPC; + static bool const configured = [&]() + { + bool ok = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem) == cudaSuccess; + ok = ok + && (kCs <= 8 + || cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) == cudaSuccess); + if (!ok) + { + cudaGetLastError(); + } + return ok; + }(); + if (!configured) + { + return cudaErrorInvalidConfiguration; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(FQKV ? kQkvGrid : kCs * (QLEN / QG) * num_reqs); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeClusterDimension; + attr[0].val.clusterDim.x = kCs; + attr[0].val.clusterDim.y = 1; + attr[0].val.clusterDim.z = 1; + attr[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[1].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + if (max_clusters != nullptr) + { + cfg.numAttrs = 1; + cudaError_t const err = cudaOccupancyMaxActiveClusters(max_clusters, kernel, &cfg); + if (err != cudaSuccess) + { + cudaGetLastError(); + } + return err; + } + cfg.numAttrs = 2; + return cudaLaunchKernelEx(&cfg, kernel, prm); +} + +// The (QLEN, QG, BPC) geometries the target launches: default_geometry() of every step shape it routes here (one +// query token per 16-CTA cluster, or KV blocks paired into 8-CTA clusters), always with the MXFP8 hand-off. +#define M3_FA_GEOMETRIES(X) \ + X(1, 1, 1) \ + X(1, 1, 2) X(2, 1, 1) X(2, 1, 2) X(3, 1, 1) X(4, 1, 1) X(5, 1, 1) X(6, 1, 1) X(6, 1, 2) X(7, 1, 1) X(8, 1, 1) + +cudaError_t dispatch(Params const& prm, int qlen, int qg, int bpc, int num_reqs, bool pdl, cudaStream_t stream) +{ +#define M3_FA_CASE(Q, G, B) \ + if (qlen == Q && qg == G && bpc == B) \ + return launch_impl(prm, num_reqs, pdl, stream, nullptr); + M3_FA_GEOMETRIES(M3_FA_CASE) +#undef M3_FA_CASE + return cudaErrorInvalidValue; +} + +} // namespace + +bool minimaxM3FusedDecodeAttnSupports(int qlen, int qg, int bpc) +{ +#define M3_FA_CASE(Q, G, B) \ + if (qlen == Q && qg == G && bpc == B) \ + return true; + M3_FA_GEOMETRIES(M3_FA_CASE) +#undef M3_FA_CASE + return false; +} + +cudaError_t invokeMinimaxM3FusedDecodeAttn( + MinimaxM3FusedDecodeAttnParams const& params, int qlen, int qg, int bpc, int numReqs, bool pdl, cudaStream_t stream) +{ + return dispatch(params, qlen, qg, bpc, numReqs, pdl, stream); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.h b/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.h new file mode 100644 index 000000000000..e1967d82f556 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.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 +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +//! Kernel parameters of the fused decode attention (the device code's Params). +struct MinimaxM3FusedDecodeAttnParams +{ + __nv_bfloat16 const* packed; + int64_t packed_stride; + uint8_t* kv_data; + int64_t d_page, d_plane; + uint8_t* kv_scale; + int64_t s_page, s_plane; + uint8_t* idx_k; + int64_t i_page, i_tok; + int64_t num_pages; + int const* out_cache_loc; + int const* position_ids; + int const* block_table; + int64_t bt_stride; + int bt_cols; + int const* seq_lens; + float const* kv_quant_scale; + float const* k_gs; + float const* v_gs; + __nv_bfloat16 const* q_w; + __nv_bfloat16 const* k_w; + __nv_bfloat16 const* ik_w; + float const* rope; + float eps; + float sm_scale; + __nv_bfloat16* out; + int64_t out_stride; + uint8_t* q8; + uint8_t* sf; + int num_sf_k_tiles; + unsigned long long* prof; // optional per-CTA clock64 probes (-DM3FA_PROFILE builds only) + // Fused QKV projection (FQKV kernels only): out[M, 2560] = mxfp8(x)[M, 6144] @ W^T. + __nv_bfloat16 const* x; // [M, 6144] layer input (the all-reduce output) + int ldx; + int num_tokens; // M + uint8_t const* wsf; // UE8M0 weight scales, 128x4-swizzled + float* part; // [2 parities][splits][8 tokens][2560] partial sums, sentinel at rest + int* counters; // [0]: call parity, [1]: CTAs done with this call (zero at rest) + int attn_clusters; // leading clusters that run the attention + uint8_t const* xq; // mega: the layer input's MXFP8 values [M, 6144] (row stride ldxq) ... + uint8_t const* xsf_act; // ... and its 128x4-swizzled UE8M0 scale atoms (atom k at byte 512 k) + int ldxq; + __nv_bfloat16* qkv_out; // mega: the packed QKV rows [2 parities][8][2560] the attention reads + uint8_t const* wsf_o; // mega: the o_proj weight's scale atoms + __nv_bfloat16* y_out; // mega: the o_proj output [M, 6144] + int* ar_ws; // mega, ar_mode: the next m3_mc_allreduce's multicast exchange (int32 words)... + int const* ar_flags; // ... its call flags (parity of a (token, column slice) slot = flags[t * 8 + slice] & 1) + int ar_rank; + int ar_world; + int ar_maxt; + int ar_mode; // 0: none; 2: push the o_proj rows into the exchange (as m3_mx_gemv's ar mode 2) + unsigned int* flags; // mega: [0, 20) tile counters (4 per call), [32 + CTA] the attention CTAs' targets + long long pf_a0, pf_b0, pf_a1, pf_b1; // mega: L2 prefetch regions issued in the attention window (bytes 0: none) + int pf_hint; // 1: evict_last policy on the prefetch + uint8_t const* wpre; // mega: optional pre-tiled QKV weight [N/128][K/128][16 KB swizzled smem image], or null + CUtensorMap wmap; // the [2560, 6144] E4M3 weight, 128 x 128 boxes, 128-byte swizzle + CUtensorMap wmap_o; // mega: the [6144, 2048] E4M3 o_proj weight, same boxes +}; + +//! Whether the (decode_query_len, query tokens per cluster, KV blocks per CTA) geometry is compiled. The MXFP8 +//! hand-off (q8 / sf) is part of every compiled variant. +bool minimaxM3FusedDecodeAttnSupports(int qlen, int qg, int bpc); + +//! Launch the fused norm/RoPE/KV-insert + identity-selection sparse decode + split-K merge kernel on `stream`. +//! `params` holds the packed QKV rows, the paged caches and the outputs (see the field comments). Returns +//! cudaErrorInvalidValue for a geometry minimaxM3FusedDecodeAttnSupports() rejects. +cudaError_t invokeMinimaxM3FusedDecodeAttn(MinimaxM3FusedDecodeAttnParams const& params, int qlen, int qg, int bpc, + int numReqs, bool pdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/mcAllreduceNorm.cu b/cpp/tensorrt_llm/kernels/minimaxM3/mcAllreduceNorm.cu new file mode 100644 index 000000000000..fc60fb3035a9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/mcAllreduceNorm.cu @@ -0,0 +1,563 @@ +/* + * 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/minimaxM3/mcAllreduceNorm.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmException.h" + +#include +#include +#include +#include + +// The kernels use sm_100-family features (thread-block clusters, cp.async.bulk, st.async, mbarrier transaction +// counts, block-scaled conversions). Every other architecture of the build compiles them to empty bodies; the host +// entry points refuse to launch on those devices. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030) \ + && (defined(__CUDA_ARCH_FEAT_SM100_ALL) || defined(__CUDA_ARCH_FEAT_SM103_ALL) \ + || defined(__CUDA_ARCH_FAMILY_SPECIFIC__)) +#define MINIMAX_M3_KERNELS_ENABLED 1 +#else +#define MINIMAX_M3_KERNELS_ENABLED 0 +#endif + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace cg = cooperative_groups; + +namespace +{ + +constexpr int kEmpty = static_cast(0x80000000u); +constexpr int kFinMaxTopk = 8; // MoE finalize: top-k rows loaded together up to this k + +__device__ __forceinline__ uint32_t canon_bf16x2(uint32_t w) +{ + if ((w & 0xFFFFu) == 0x8000u) + w &= 0xFFFF0000u; + if ((w >> 16) == 0x8000u) + w &= 0x0000FFFFu; + return w; +} + +__device__ __forceinline__ float bf16_lo(uint32_t w) +{ + return __uint_as_float(w << 16); +} + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ float bf16_hi(uint32_t w) +{ + return __uint_as_float(w & 0xFFFF0000u); +} + +__device__ __forceinline__ int4 ld_volatile_v4(int const* p) +{ + int4 v; + asm volatile("ld.volatile.global.v4.s32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p) + : "memory"); + return v; +} + +struct FinalizeArgs +{ + __nv_bfloat16 const* fc2; // [permuted rows, H] routed-expert FC2 output (null: plain input x) + int const* perm; // [M, topk] expanded index -> permuted row, -1 = dropped + void const* scales; // [M, topk] routing weights (fp32 or bf16) + int scale_fp32; + int topk; +}; + +// grid = min(M, MAXT) * CLUSTER CTAs, cluster (CLUSTER, 1, 1), THREADS threads. +// OPTS (each bit-identical): 1 residual + norm weight loaded at entry (before the poll); 2 cluster RMS exchange by +// st.async + mbarrier complete_tx instead of cluster.sync; 4 call parity rotated right after the poll; 8 max-smem +// carveout so the next kernel's CTAs co-reside early. +template +__global__ void __launch_bounds__(THREADS) mc_allreduce_norm_kernel(__nv_bfloat16 const* __restrict__ x, + __nv_bfloat16 const* __restrict__ x2, FinalizeArgs fin, int M, int H, __nv_bfloat16 const* __restrict__ residual, + __nv_bfloat16 const* __restrict__ norm_w, float eps, int* __restrict__ ws_uc, int* __restrict__ ws_mc, + int* __restrict__ flags, int rank, int world, __nv_bfloat16* __restrict__ out_normed, + __nv_bfloat16* __restrict__ out_residual, uint8_t* __restrict__ out_q, uint8_t* __restrict__ out_sf, int sf_k_tiles, + int prepushed) +{ +#if MINIMAX_M3_KERNELS_ENABLED + __shared__ float warp_sq[THREADS / 32]; + __shared__ alignas(16) float cluster_sq[CLUSTER]; + __shared__ alignas(8) unsigned long long xbar; + cg::cluster_group cluster = cg::this_cluster(); + if constexpr ((OPTS & 2) != 0) + { + if (threadIdx.x == 0) + { + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(b) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); + } + int const crank = static_cast(cluster.block_rank()); + int const t = blockIdx.x / CLUSTER; + int const cols = H / CLUSTER; + int const c0 = crank * cols + threadIdx.x * 8; + bool const active = t < M && threadIdx.x * 8 < cols; + int64_t const half_words = static_cast(MAXT) * world * (H / 2); + + // prepushed 1: the producer (m3_mx_gemv) already pushed this rank's rows into the exchange slots; 2: and it + // launched this grid only after its pushes (its grid wait passed), so the slots are polled without a grid wait; + // 3: as 2, with x in set A and x2 in set B. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (prepushed < 2) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + int const parity = *reinterpret_cast(flags + blockIdx.x) & 1; + uint4 rraw = make_uint4(0u, 0u, 0u, 0u); + uint4 wraw = make_uint4(0u, 0u, 0u, 0u); + if constexpr ((OPTS & 1) != 0) + { + // the residual row was written grids earlier (prepushed >= 2: the producer launched this grid only after its + // own grid wait; otherwise the grid wait above) + if (active) + { + rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + wraw = *reinterpret_cast(norm_w + c0); + } + } + // Push this rank's slice into every rank's buffer. + if (active && !prepushed) + { + uint4 raw; + if (fin.fc2 != nullptr) + { + // MoE finalize: fp32 sum of the routing-weighted top-k expert rows plus the shared expert, rounded once. + float f[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + if (fin.topk <= kFinMaxTopk) + { + // All top-k indices and scales, then every expert-row load in flight at once, then the same in-order + // FP32 accumulation as the loop below (one memory round trip instead of topk; bit-identical). + int prow[kFinMaxTopk]; + float sc[kFinMaxTopk]; +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + prow[k] = -1; + sc[k] = 0.f; + if (k < fin.topk) + { + int const e = t * fin.topk + k; + prow[k] = fin.perm[e]; + sc[k] = fin.scale_fp32 ? static_cast(fin.scales)[e] + : __bfloat162float(static_cast<__nv_bfloat16 const*>(fin.scales)[e]); + } + } + uint4 rv[kFinMaxTopk]; +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + if (prow[k] >= 0) + rv[k] = *reinterpret_cast(fin.fc2 + static_cast(prow[k]) * H + c0); + } +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + if (prow[k] < 0) + continue; + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv[k]); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] += v.x * sc[k]; + f[2 * q + 1] += v.y * sc[k]; + } + } + } + else + for (int k = 0; k < fin.topk; ++k) + { + int const e = t * fin.topk + k; + int const prow = fin.perm[e]; + if (prow < 0) + continue; + float const sc = fin.scale_fp32 + ? static_cast(fin.scales)[e] + : __bfloat162float(static_cast<__nv_bfloat16 const*>(fin.scales)[e]); + uint4 const rv = *reinterpret_cast(fin.fc2 + static_cast(prow) * H + c0); + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] += v.x * sc; + f[2 * q + 1] += v.y * sc; + } + } + if (x2 != nullptr) + { + // As stock: the routed output is rounded to BF16 first, then the shared expert is added (one more + // BF16 rounding of the FP32 sum). + uint4 const rv = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const routed = __bfloat1622float2(__floats2bfloat162_rn(f[2 * q], f[2 * q + 1])); + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] = routed.x + v.x; + f[2 * q + 1] = routed.y + v.y; + } + } + __nv_bfloat162* o = reinterpret_cast<__nv_bfloat162*>(&raw); +#pragma unroll + for (int q = 0; q < 4; ++q) + o[q] = __floats2bfloat162_rn(f[2 * q], f[2 * q + 1]); + } + else + { + raw = *reinterpret_cast(x + static_cast(t) * H + c0); + } + if (fin.fc2 == nullptr && x2 != nullptr) + { + uint4 const raw2 = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162* a = reinterpret_cast<__nv_bfloat162*>(&raw); + __nv_bfloat162 const* b = reinterpret_cast<__nv_bfloat162 const*>(&raw2); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const fa = __bfloat1622float2(a[q]); + float2 const fb = __bfloat1622float2(b[q]); + a[q] = __floats2bfloat162_rn(fa.x + fb.x, fa.y + fb.y); + } + } + int64_t const slot = parity * half_words + (static_cast(t) * world + rank) * (H / 2) + c0 / 2; + *reinterpret_cast(ws_mc + slot) + = make_int4(static_cast(canon_bf16x2(raw.x)), static_cast(canon_bf16x2(raw.y)), + static_cast(canon_bf16x2(raw.z)), static_cast(canon_bf16x2(raw.w))); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float s[8] = {}; + float sq = 0.f; + if (active) + { + // The residual row (written grids earlier) and every rank's slot are loaded together; slots still empty are + // re-polled. Ranks are then summed in rank order, as before. + if constexpr ((OPTS & 1) == 0) + rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + constexpr int kMaxWorld = 8; + int* const p0 = ws_uc + parity * half_words + static_cast(t) * world * (H / 2) + c0 / 2; + int4 v[kMaxWorld]; + int4 v2[kMaxWorld]; + bool const two = prepushed == 3; + int* const p1 = p0 + 2 * half_words; // set B + uint32_t pending = (1u << world) - 1u; + uint32_t pending2 = two ? pending : 0u; + long long spins = 0; + while ((pending | pending2) != 0u) + { +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if ((pending >> r) & 1u) + v[r] = ld_volatile_v4(p0 + static_cast(r) * (H / 2)); + if ((pending2 >> r) & 1u) + v2[r] = ld_volatile_v4(p1 + static_cast(r) * (H / 2)); + } +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if (((pending >> r) & 1u) && v[r].x != kEmpty && v[r].y != kEmpty && v[r].z != kEmpty + && v[r].w != kEmpty) + pending &= ~(1u << r); + if (((pending2 >> r) & 1u) && v2[r].x != kEmpty && v2[r].y != kEmpty && v2[r].z != kEmpty + && v2[r].w != kEmpty) + pending2 &= ~(1u << r); + } + if (++spins > (1LL << 31)) + __trap(); + } +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if (r < world) + { + *reinterpret_cast(p0 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t ww[4] = {static_cast(v[r].x), static_cast(v[r].y), + static_cast(v[r].z), static_cast(v[r].w)}; + if (two) + { + *reinterpret_cast(p1 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t const w2[4] = {static_cast(v2[r].x), static_cast(v2[r].y), + static_cast(v2[r].z), static_cast(v2[r].w)}; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + __nv_bfloat162 const u + = __floats2bfloat162_rn(bf16_lo(ww[q]) + bf16_lo(w2[q]), bf16_hi(ww[q]) + bf16_hi(w2[q])); + ww[q] = *reinterpret_cast(&u); + } + } +#pragma unroll + for (int q = 0; q < 4; ++q) + { + s[2 * q] += bf16_lo(ww[q]); + s[2 * q + 1] += bf16_hi(ww[q]); + } + } + } + __nv_bfloat162 const* rr = reinterpret_cast<__nv_bfloat162 const*>(&rraw); + uint32_t upd[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const rf = __bfloat1622float2(rr[q]); + __nv_bfloat162 const u = __floats2bfloat162_rn(s[2 * q] + rf.x, s[2 * q + 1] + rf.y); + upd[q] = *reinterpret_cast(&u); + float2 const uf = __bfloat1622float2(u); + s[2 * q] = uf.x; + s[2 * q + 1] = uf.y; + sq += uf.x * uf.x + uf.y * uf.y; + } + *reinterpret_cast(out_residual + static_cast(t) * H + c0) + = make_uint4(upd[0], upd[1], upd[2], upd[3]); + } + if constexpr ((OPTS & 4) != 0) + { + // every CTA (active or not) rotates its call parity once its own poll is done (read only by later grids) + if (threadIdx.x == 0) + flags[blockIdx.x] = parity ^ 1; + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffffu, sq, off); + if ((threadIdx.x & 31) == 0) + warp_sq[threadIdx.x / 32] = sq; + __syncthreads(); + if constexpr ((OPTS & 2) != 0) + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); // peers' xbar initialized (long since) + if (threadIdx.x == 0) + { + float total = 0.f; +#pragma unroll + for (int w = 0; w < THREADS / 32; ++w) + total += warp_sq[w]; + if constexpr ((OPTS & 2) == 0) + for (int c = 0; c < CLUSTER; ++c) + *cluster.map_shared_rank(&cluster_sq[crank], c) = total; + else + { + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + uint32_t const d = static_cast(__cvta_generic_to_shared(&cluster_sq[crank])); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(b), "r"(CLUSTER * 4) + : "memory"); + for (int c = 0; c < CLUSTER; ++c) + { + uint32_t rd, rb; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rd) : "r"(d), "r"(c)); + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rb) : "r"(b), "r"(c)); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(rd), + "r"(__float_as_uint(total)), "r"(rb) + : "memory"); + } + } + } + if constexpr ((OPTS & 2) == 0) + cluster.sync(); + else + { + // every peer's sum landed (st.async complete_tx on this CTA's barrier); no cluster-wide release-arrive + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + uint32_t done = 0; + while (!done) + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}" + : "=r"(done) + : "r"(b) + : "memory"); + } + if (active) + { + float total = 0.f; +#pragma unroll + for (int c = 0; c < CLUSTER; ++c) + total += cluster_sq[c]; + float const inv = rsqrtf(total / static_cast(H) + eps); + if constexpr ((OPTS & 1) == 0) + wraw = *reinterpret_cast(norm_w + c0); + __nv_bfloat162 const* wv = reinterpret_cast<__nv_bfloat162 const*>(&wraw); + uint32_t o[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const wf = __bfloat1622float2(wv[q]); + __nv_bfloat162 const n = __floats2bfloat162_rn(s[2 * q] * inv * wf.x, s[2 * q + 1] * inv * wf.y); + o[q] = *reinterpret_cast(&n); + float2 const nf = __bfloat1622float2(n); + s[2 * q] = nf.x; + s[2 * q + 1] = nf.y; + } + *reinterpret_cast(out_normed + static_cast(t) * H + c0) = make_uint4(o[0], o[1], o[2], o[3]); + } + if (out_q != nullptr) + { + // MXFP8 copy of the normed row for the next GEMM, as trtllm::mxfp8_quantize(normed, True): four threads + // hold one 32-column scale block. + float amax = 0.f; + if (active) + { +#pragma unroll + for (int e = 0; e < 8; ++e) + amax = fmaxf(amax, fabsf(s[e])); + } + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 2)); + if (active) + { + __nv_fp8_e8m0 code; + code.__x = __nv_cvt_float_to_e8m0(amax * rcp_approx_ftz(448.0f), __NV_SATFINITE, cudaRoundPosInf); + float const out_scale = amax != 0.f ? rcp_approx_ftz(static_cast(code)) : 0.0f; + uint32_t packed[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) + { + __nv_fp8x2_e4m3 const lo(make_float2(s[4 * h] * out_scale, s[4 * h + 1] * out_scale)); + __nv_fp8x2_e4m3 const hi(make_float2(s[4 * h + 2] * out_scale, s[4 * h + 3] * out_scale)); + packed[h] = static_cast(lo.__x) | (static_cast(hi.__x) << 16); + } + *reinterpret_cast(out_q + static_cast(t) * H + c0) = make_uint2(packed[0], packed[1]); + if ((threadIdx.x & 3) == 0) + { + int const kb = c0 / 32; + out_sf[static_cast(t / 128) * sf_k_tiles * 512 + (kb / 4) * 512 + (t % 32) * 16 + + ((t % 128) / 32) * 4 + kb % 4] + = code.__x; + } + } + } + if constexpr ((OPTS & 4) == 0) + { + __syncthreads(); + if (threadIdx.x == 0) + flags[blockIdx.x] = parity ^ 1; + } +#endif +} + +template +void launch(MinimaxM3McAllreduceNormParams const& p, bool pdl, cudaStream_t stream) +{ + int const M = p.numTokens, H = p.hidden; + TLLM_CHECK_WITH_INFO(H % (CLUSTER * 8) == 0 && H / CLUSTER / 8 <= THREADS, "unsupported hidden size %d", H); + TLLM_CHECK_WITH_INFO(M >= 1 && M <= MAXT, "%d tokens do not fit the %d-token bucket", M, MAXT); + auto kernel = mc_allreduce_norm_kernel; + static bool const configured = [&]() + { + bool ok = CLUSTER <= 8 + || cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) == cudaSuccess; + if ((OPTS & 8) != 0) + ok = ok + && cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared) + == cudaSuccess; + if (!ok) + cudaGetLastError(); + return ok; + }(); + TLLM_CHECK_WITH_INFO(configured, "cluster size not supported"); + FinalizeArgs const fin{p.fc2, p.perm, p.scales, p.scaleFp32 ? 1 : 0, p.topk}; + cudaLaunchConfig_t cfg = {}; + // The grid is sized to the call's rows: min(M, MAXT) clusters. Each (row, slice) slot keeps its own parity flag, + // read by the producers before they push and rotated only by the slot's CTA after it polled: a slot that a call + // does not use keeps its parity, and every rank runs the same calls, so the parities stay consistent when M + // changes between calls. + cfg.gridDim = dim3((M < MAXT ? M : MAXT) * CLUSTER); + cfg.blockDim = dim3(THREADS); + cfg.stream = stream; + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + attr[1].id = cudaLaunchAttributeClusterDimension; + attr[1].val.clusterDim.x = CLUSTER; + attr[1].val.clusterDim.y = 1; + attr[1].val.clusterDim.z = 1; + cfg.attrs = attr; + cfg.numAttrs = 2; + cudaError_t const err + = cudaLaunchKernelEx(&cfg, kernel, p.x, p.x2, fin, M, H, p.residual, p.normWeight, p.eps, p.wsUc, p.wsMc, + p.flags, p.rank, p.world, p.outNormed, p.outResidual, p.outQuant, p.outSf, (H / 32 + 3) / 4, p.prepushed); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +constexpr int kCluster = 8; +constexpr int kThreads = 128; + +} // namespace + +bool isMinimaxM3McAllreduceNormSupported(int32_t maxTokens, int32_t opts) +{ + return (maxTokens == 8 && (opts == 7 || opts == 15)) || (maxTokens == 16 && opts == 15) + || (maxTokens == 48 && opts == 15); +} + +void invokeMinimaxM3McAllreduceNorm( + MinimaxM3McAllreduceNormParams const& params, int32_t maxTokens, int32_t opts, bool pdl, cudaStream_t stream) +{ + if (maxTokens == 8 && opts == 7) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 15) + { + launch(params, pdl, stream); + } + else if (maxTokens == 16 && opts == 15) + { + launch(params, pdl, stream); + } + else if (maxTokens == 48 && opts == 15) + { + launch(params, pdl, stream); + } + else + { + TLLM_THROW("unsupported (bucket %d, opts %d)", maxTokens, opts); + } +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/mcAllreduceNorm.h b/cpp/tensorrt_llm/kernels/minimaxM3/mcAllreduceNorm.h new file mode 100644 index 000000000000..b80a78c9e173 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/mcAllreduceNorm.h @@ -0,0 +1,76 @@ +/* + * 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 +{ + +//! Operands of the MiniMax-M3 multicast all-reduce + residual + RMSNorm kernel (decode, at most 48 tokens). +//! +//! Per rank the kernel forms the partial `bf16(x + x2)` (or, in finalize mode, +//! `bf16(bf16(sum_k scale * fc2[perm]) + x2)`, the routed sum rounded before the shared expert is added, as stock), +//! pushes it through the multicast mapping `wsMc` into every rank's exchange, polls its own unicast view `wsUc`, +//! sums the ranks in FP32 in rank order, adds the residual (BF16 round) and applies RMSNorm. With `outQuant` it +//! also writes the MXFP8 copy of the normed rows (E4M3 [numTokens, hidden] and swizzled UE8M0 scales). +struct MinimaxM3McAllreduceNormParams +{ + __nv_bfloat16 const* x{nullptr}; //!< [numTokens, hidden] partial; in finalize mode the FC2 rows [rows, hidden]. + __nv_bfloat16 const* x2{nullptr}; //!< Optional second addend [numTokens, hidden]. + __nv_bfloat16 const* fc2{nullptr}; //!< Finalize mode: the permuted FC2 output (same buffer as x); null otherwise. + int32_t const* perm{nullptr}; //!< Finalize mode: [numTokens, topk] expanded index -> permuted row, -1 dropped. + void const* scales{nullptr}; //!< Finalize mode: [numTokens, topk] routing weights (fp32 or bf16). + bool scaleFp32{false}; + int32_t topk{0}; + int32_t numTokens{0}; + int32_t hidden{0}; + __nv_bfloat16 const* residual{nullptr}; //!< [numTokens, hidden] + __nv_bfloat16 const* normWeight{nullptr}; //!< [hidden] + float eps{0.F}; + int32_t* wsUc{nullptr}; //!< This rank's unicast view of the exchange (int32 words, re-armed after reading). + int32_t* wsMc{nullptr}; //!< The multicast mapping of the exchange. + int32_t* flags{nullptr}; //!< Per-CTA call-parity counters, [maxTokens * 8]. + int32_t rank{0}; + int32_t world{0}; + __nv_bfloat16* outNormed{nullptr}; + __nv_bfloat16* outResidual{nullptr}; + uint8_t* outQuant{nullptr}; //!< Optional E4M3 [numTokens, hidden]. + uint8_t* outSf{nullptr}; //!< Optional UE8M0 scales, 128x4-swizzled. + //! 0: push x (+ x2) here; 2: a producer already pushed x and launched this grid after its pushes; 3: as 2, with x + //! in exchange set A and x2 in set B. + int32_t prepushed{0}; +}; + +//! The (bucket, epilogue options) pairs this launcher is built for: (8, 7), (8, 15), (16, 15) and (48, 15). +bool isMinimaxM3McAllreduceNormSupported(int32_t maxTokens, int32_t opts); + +//! Launch with 8-CTA clusters of 128 threads, one cluster per token. `maxTokens` selects the exchange bucket; +//! `opts` the epilogue variant (every variant is bit-identical). +void invokeMinimaxM3McAllreduceNorm( + MinimaxM3McAllreduceNormParams const& params, int32_t maxTokens, int32_t opts, bool pdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/msaKvUpdate.cu b/cpp/tensorrt_llm/kernels/minimaxM3/msaKvUpdate.cu new file mode 100644 index 000000000000..2fb9f768c1bd --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/msaKvUpdate.cu @@ -0,0 +1,70 @@ +/* + * 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/minimaxM3/msaKvUpdate.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ +__global__ void __launch_bounds__(256) + msa_kv_update_kernel(int const* __restrict__ kv_lens, int const* __restrict__ staged, int const* __restrict__ qo, + int const* __restrict__ q_batch_row, int const* __restrict__ q_intra, int const* __restrict__ table, + long long table_ld, int ncols, int page, int batch, int total_q, int* __restrict__ seq_lens_out, + int* __restrict__ cu_kv, int* __restrict__ out_loc, int* __restrict__ n_valid) +{ + for (int b = threadIdx.x; b < batch; b += blockDim.x) + seq_lens_out[b] = min(kv_lens[b], staged[b]); + if (cu_kv != nullptr && threadIdx.x == 0) + { + int acc = 0; + cu_kv[0] = 0; + for (int b = 0; b < batch; ++b) + { + acc += min(kv_lens[b], staged[b]); + cu_kv[b + 1] = acc; + } + } + for (int t = threadIdx.x; t < total_q; t += blockDim.x) + { + int const r = q_batch_row[t]; + int const pos = min(kv_lens[r], staged[r]) - qo[r] + q_intra[t]; + int const ps = max(pos, 0); + int const col = min(ps / page, ncols - 1); + out_loc[t] = table[r * table_ld + col] * page + ps % page; + if (n_valid != nullptr) + n_valid[t] = (max(pos + 1, 0) + page - 1) / page; + } +} +} // namespace + +cudaError_t invokeMinimaxM3MsaKvUpdate(int32_t const* kvLens, int32_t const* staged, int32_t const* qo, + int32_t const* qBatchRow, int32_t const* qIntra, int32_t const* table, int64_t tableLd, int32_t ncols, int32_t page, + int32_t batch, int32_t totalQ, int32_t* seqLensOut, int32_t* cuKv, int32_t* outLoc, int32_t* nValid, + cudaStream_t stream) +{ + msa_kv_update_kernel<<<1, 256, 0, stream>>>(kvLens, staged, qo, qBatchRow, qIntra, table, + static_cast(tableLd), ncols, page, batch, totalQ, seqLensOut, cuKv, outLoc, nValid); + return cudaGetLastError(); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/msaKvUpdate.h b/cpp/tensorrt_llm/kernels/minimaxM3/msaKvUpdate.h new file mode 100644 index 000000000000..d703cef6df73 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/msaKvUpdate.h @@ -0,0 +1,45 @@ +/* + * 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 +{ + +//! MiniMax-M3 MSA KV-length correction in one single-CTA launch, for all int32 operands: +//! +//! kvTrue[b] = min(kvLens[b], staged[b]) -> seqLensOut[b]; cuKv = exclusive prefix sum (optional) +//! pos[t] = kvTrue[r] - qo[r] + qIntra[t] (r = qBatchRow[t]) +//! outLoc[t] = table[r, min(max(pos, 0) / page, ncols - 1)] * page + max(pos, 0) % page +//! nValid[t] = (max(pos + 1, 0) + page - 1) / page (optional) +//! +//! cuKv and nValid may be null. table is row-major with row stride tableLd (elements) and ncols columns. +cudaError_t invokeMinimaxM3MsaKvUpdate(int32_t const* kvLens, int32_t const* staged, int32_t const* qo, + int32_t const* qBatchRow, int32_t const* qIntra, int32_t const* table, int64_t tableLd, int32_t ncols, int32_t page, + int32_t batch, int32_t totalQ, int32_t* seqLensOut, int32_t* cuKv, int32_t* outLoc, int32_t* nValid, + cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/mxfp8LowMGemm.cu b/cpp/tensorrt_llm/kernels/minimaxM3/mxfp8LowMGemm.cu new file mode 100644 index 000000000000..2e942d5b5a57 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/mxfp8LowMGemm.cu @@ -0,0 +1,820 @@ +/* + * 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/minimaxM3/mxfp8LowMGemm.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmException.h" + +#include +#include +#include +#include + +// The kernels use sm_100-family features (thread-block clusters, cp.async.bulk, st.async, mbarrier transaction +// counts, block-scaled conversions). Every other architecture of the build compiles them to empty bodies; the host +// entry points refuse to launch on those devices. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030) \ + && (defined(__CUDA_ARCH_FEAT_SM100_ALL) || defined(__CUDA_ARCH_FEAT_SM103_ALL) \ + || defined(__CUDA_ARCH_FAMILY_SPECIFIC__)) +#define MINIMAX_M3_KERNELS_ENABLED 1 +#else +#define MINIMAX_M3_KERNELS_ENABLED 0 +#endif + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace cg = cooperative_groups; + +namespace +{ + +constexpr int kMaxM = 48; +constexpr int kRowPad = 32; +constexpr int kEmptyWord = static_cast(0x80000000u); + +// Layer-boundary reduction inputs of the AR prologue (see m3_ar_fused.py); unused otherwise. +struct ArArgs +{ + int const* uc; + int const* cur_buf; + int64_t buf_words; + int world; + __nv_bfloat16 const* residual; + __nv_bfloat16 const* norm_w; + float eps; + __nv_bfloat16* residual_out; + uint8_t* xq_g; // [8][K] quantized activation published by tile 0 + uint8_t* xc_g; // [8][K / 32] its UE8M0 scales + int* flags; // [CLUSTER_K] per K-slice publish generation (call + 1) +}; + +// Optional L2 prefetch of the paged KV blocks the step's sparse decode will read (issued before the PDL wait, so it +// overlaps the producer kernel); all-null when unused. +struct PfArgs +{ + int const* block_table; // [rows, bt_stride] page ids + int const* seq_lens; // [rows] KV length per row + int rows; + int max_blocks; + int bt_stride; + char const* data; // page-major KV data pool + long long data_stride; // bytes between pages + int data_bytes; // contiguous bytes per page + char const* scale; // page-major block-scale pool (or null) + long long scale_stride; + int scale_bytes; +}; + +// Already-quantized activation (E4M3 rows [M, K] + UE8M0 scales in the 128x4-swizzled layout of +// trtllm::mxfp8_quantize(x, True)); all-null to quantize the BF16 activation in-kernel. +struct PreqArgs +{ + uint8_t const* xq; + uint8_t const* xsf; + int sf_k_tiles; // ceil(K / 128) +}; + +__device__ __forceinline__ void prefetch_kv_pages(PfArgs const& pf) +{ + if (pf.block_table == nullptr) + return; + int const units = pf.rows * pf.max_blocks; + for (int u = blockIdx.x; u < units; u += gridDim.x) + { + int const r = u / pf.max_blocks; + int const b = u - r * pf.max_blocks; + if (b * 128 >= pf.seq_lens[r]) + continue; + long long const page = pf.block_table[static_cast(r) * pf.bt_stride + b]; + char const* d = pf.data + page * pf.data_stride; + for (int off = threadIdx.x * 128; off < pf.data_bytes; off += blockDim.x * 128) + asm volatile("prefetch.global.L2::evict_last [%0];" ::"l"(d + off)); + if (pf.scale != nullptr) + { + char const* sc = pf.scale + page * pf.scale_stride; + for (int off = threadIdx.x * 128; off < pf.scale_bytes; off += blockDim.x * 128) + asm volatile("prefetch.global.L2::evict_last [%0];" ::"l"(sc + off)); + } + } +} + +__device__ __forceinline__ uint32_t mapa_shared(uint32_t addr, uint32_t target_rank) +{ + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n" : "=r"(r) : "r"(addr), "r"(target_rank)); + return r; +} + +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ float pow2_from_ue8m0(uint32_t e) +{ + return __uint_as_float(e == 0 ? 0x00400000u : (e << 23)); +} + +__device__ __forceinline__ void mma_e4m3( + 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.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\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), "f"(0.f), "f"(0.f), "f"(0.f), "f"(0.f)); +} + +__device__ __forceinline__ void bulk_g2s(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];\n" ::"r"( + smem_u32(dst)), + "l"(src), "r"(bytes), "r"(smem_u32(bar)) + : "memory"); +} + +__device__ __forceinline__ void mbar_wait(uint64_t* bar, uint32_t parity) +{ + asm volatile( + "{\n .reg .pred p;\n WAIT_%=:\n mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n @!p bra " + "WAIT_%=;\n}\n" ::"r"(smem_u32(bar)), + "r"(parity) + : "memory"); +} + +// grid: (N / (16 * RT) * CLUSTER_K); cluster (CLUSTER_K, 1, 1); block WARPS warps. +// RT: 16-row weight tiles per CTA (amortizes the activation fetch + quantization over more rows). +// NT: number of 8-row activation tiles (MMA n=8 each), 1 for M <= 8 and 2 for M <= 16. +template +__global__ void __launch_bounds__(WARPS * 32) mxfp8_low_m_gemm_tma_kernel(__nv_bfloat16 const* __restrict__ x, int ldx, + int M, int N, int K, uint8_t const* __restrict__ w, uint8_t const* __restrict__ wsf, int num_k_tiles, + __nv_bfloat16* __restrict__ out, int ldo, int red_aliased, ArArgs ar, PfArgs pf, PreqArgs preq) +{ +#if MINIMAX_M3_KERNELS_ENABLED + constexpr int kThreads = WARPS * 32; + constexpr int kRows = 16 * RT; + constexpr int kTile = kRows * 8 * NT; + extern __shared__ __align__(128) uint8_t smem[]; + __shared__ __align__(8) uint64_t bar; + __shared__ __align__(8) uint64_t red_bar; + + int const rank = CLUSTER_K > 1 ? static_cast(cg::this_cluster().block_rank()) : 0; + int const tile = blockIdx.x / CLUSTER_K; + int const row0 = tile * kRows; + int const kc = K / CLUSTER_K; // columns owned by this CTA + int const kc_blocks = kc / 32; + int const k_begin = rank * kc; + int const kb_begin = k_begin / 32; + int const row_stride = kc + kRowPad; + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + + uint8_t* w_s = smem; // [kRows][row_stride] + uint8_t* sf_s = w_s + kRows * row_stride; // [RT][kc_blocks / 4][256] + uint8_t* xq_s = sf_s + RT * (kc_blocks / 4) * 256; // [M][kc] + uint8_t* xc_s = xq_s + M * kc; // [M][kc_blocks] + // Peer partial tiles (root only), then the warp partials unless they reuse the dead weight tile. + float* cluster_red = reinterpret_cast(xc_s + ((M * kc_blocks + 15) / 16) * 16); + float* red = red_aliased ? reinterpret_cast(w_s) : cluster_red + (CLUSTER_K - 1) * kTile; + // AR: the reduced + residual-added activation slice, BF16 [M][kc]. + __nv_bfloat16* u_s = reinterpret_cast<__nv_bfloat16*>( + reinterpret_cast(cluster_red + (CLUSTER_K - 1) * kTile) + (red_aliased ? 0 : WARPS * kTile * 4)); + __shared__ float tok_sq[kMaxM]; + __shared__ float peer_sq[CLUSTER_K][kMaxM]; + + uint32_t const w_bytes = kRows * kc; + uint32_t const sf_bytes = RT * (kc_blocks / 4) * 256u; + if (threadIdx.x == 0) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" ::"r"(smem_u32(&bar))); + if (CLUSTER_K > 1 && rank == 0) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" ::"r"(smem_u32(&red_bar))); + } + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + asm volatile( + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smem_u32(&bar)), "r"(w_bytes + sf_bytes) + : "memory"); + if (CLUSTER_K > 1 && rank == 0) + { + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smem_u32(&red_bar)), + "r"(static_cast((CLUSTER_K - 1) * kTile * sizeof(float))) + : "memory"); + } + } + __syncthreads(); + // The root's reduction barrier must be initialized before any peer pushes into it. + if constexpr (CLUSTER_K > 1) + { + asm volatile("barrier.cluster.arrive.release.aligned;\n" ::: "memory"); + } + // Weight rows and scale chunks do not depend on the producer kernel. + if (warp == 0) + { + for (int r = lane; r < kRows; r += 32) + { + bulk_g2s(w_s + r * row_stride, w + static_cast(row0 + r) * K + k_begin, kc, &bar); + } + int const chunks = kc_blocks / 4; + for (int t = lane; t < RT * chunks; t += 32) + { + int const rt = t / chunks; + int const c = t % chunks; + int const rr = row0 + 16 * rt; + int64_t const sf_base = static_cast(rr / 128) * num_k_tiles * 512 + (rr % 32) * 16; + bulk_g2s(sf_s + t * 256, wsf + sf_base + static_cast(kb_begin / 4 + c) * 512, 256, &bar); + } + } + prefetch_kv_pages(pf); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + if constexpr (AR) + { + // Layer-boundary reduction: sum the ranks' pushed partials for this CTA's K slice, add the residual + // (BF16 round, as the unfused add), then normalize with the RMS statistic of the whole row (the + // cluster spans K) and quantize like the unfused path. + if (threadIdx.x < kMaxM) + tok_sq[threadIdx.x] = 0.f; + __syncthreads(); + int const call = *reinterpret_cast(ar.cur_buf); + if (tile == 0) + { + int* base = const_cast(ar.uc) + static_cast(call & 1) * ar.buf_words; + int const half_k = K / 2; + for (int p = threadIdx.x; p < M * kc_blocks; p += kThreads) + { + int const m = p / kc_blocks; + int const b = p % kc_blocks; + int const k0 = k_begin + b * 32; + float v[32]; +#pragma unroll + for (int e = 0; e < 32; ++e) + v[e] = 0.f; + constexpr int kWorld = 4; + int4 wv[kWorld][4]; + int* src0 = base + static_cast(m) * kWorld * half_k + k0 / 2; +#pragma unroll + for (int r = 0; r < kWorld; ++r) + { +#pragma unroll + for (int q = 0; q < 4; ++q) + { + asm volatile("ld.volatile.global.v4.s32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(wv[r][q].x), "=r"(wv[r][q].y), "=r"(wv[r][q].z), "=r"(wv[r][q].w) + : "l"(src0 + r * half_k + 4 * q) + : "memory"); + } + } + long long spins = 0; + while (true) + { + bool pending = false; +#pragma unroll + for (int r = 0; r < kWorld; ++r) + { +#pragma unroll + for (int q = 0; q < 4; ++q) + { + int4 const c = wv[r][q]; + if (c.x == kEmptyWord || c.y == kEmptyWord || c.z == kEmptyWord || c.w == kEmptyWord) + { + pending = true; + asm volatile("ld.volatile.global.v4.s32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(wv[r][q].x), "=r"(wv[r][q].y), "=r"(wv[r][q].z), "=r"(wv[r][q].w) + : "l"(src0 + r * half_k + 4 * q) + : "memory"); + } + } + } + if (!pending) + break; + if (++spins > (1LL << 30)) + __trap(); + } +#pragma unroll + for (int r = 0; r < kWorld; ++r) + { +#pragma unroll + for (int q = 0; q < 4; ++q) + { + *reinterpret_cast(src0 + r * half_k + 4 * q) + = make_int4(kEmptyWord, kEmptyWord, kEmptyWord, kEmptyWord); + uint32_t const ww[4] = {static_cast(wv[r][q].x), static_cast(wv[r][q].y), + static_cast(wv[r][q].z), static_cast(wv[r][q].w)}; +#pragma unroll + for (int h = 0; h < 4; ++h) + { + v[q * 8 + 2 * h] += __uint_as_float(ww[h] << 16); + v[q * 8 + 2 * h + 1] += __uint_as_float(ww[h] & 0xFFFF0000u); + } + } + } + uint4 const* rsrc = reinterpret_cast(ar.residual + static_cast(m) * K + k0); + uint4* rdst = reinterpret_cast(ar.residual_out + static_cast(m) * K + k0); + uint4* udst = reinterpret_cast(u_s + m * kc + b * 32); + float sq = 0.f; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + uint4 const rraw = rsrc[q]; + __nv_bfloat162 const* rr = reinterpret_cast<__nv_bfloat162 const*>(&rraw); + uint32_t uw[4]; +#pragma unroll + for (int h = 0; h < 4; ++h) + { + float2 const rf = __bfloat1622float2(rr[h]); + __nv_bfloat162 const u + = __floats2bfloat162_rn(v[q * 8 + 2 * h] + rf.x, v[q * 8 + 2 * h + 1] + rf.y); + uw[h] = *reinterpret_cast(&u); + float2 const uf = __bfloat1622float2(u); + sq += uf.x * uf.x + uf.y * uf.y; + } + uint4 const uu = make_uint4(uw[0], uw[1], uw[2], uw[3]); + udst[q] = uu; + if (tile == 0) + rdst[q] = uu; + } + atomicAdd(&tok_sq[m], sq); + } + __syncthreads(); + if constexpr (CLUSTER_K > 1) + { + cg::cluster_group cl = cg::this_cluster(); + asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory"); + if (threadIdx.x < M) + { + for (int c = 0; c < CLUSTER_K; ++c) + { + *cl.map_shared_rank(&peer_sq[rank][threadIdx.x], c) = tok_sq[threadIdx.x]; + } + } + asm volatile("barrier.cluster.arrive.release.aligned;\n" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory"); + // Re-arm for the split-K reduction's wait below. + asm volatile("barrier.cluster.arrive.release.aligned;\n" ::: "memory"); + } + else + { + if (threadIdx.x < M) + peer_sq[0][threadIdx.x] = tok_sq[threadIdx.x]; + __syncthreads(); + } + for (int p = threadIdx.x; p < M * kc_blocks; p += kThreads) + { + int const m = p / kc_blocks; + int const b = p % kc_blocks; + float total = 0.f; +#pragma unroll + for (int c = 0; c < CLUSTER_K; ++c) + total += peer_sq[c][m]; + float const inv = rsqrtf(total / static_cast(K) + ar.eps); + uint4 const* usrc = reinterpret_cast(u_s + m * kc + b * 32); + uint4 const* wsrc = reinterpret_cast(ar.norm_w + k_begin + b * 32); + float v[32]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + uint4 const uraw = usrc[q]; + uint4 const wraw = wsrc[q]; + __nv_bfloat162 const* uu = reinterpret_cast<__nv_bfloat162 const*>(&uraw); + __nv_bfloat162 const* wwv = reinterpret_cast<__nv_bfloat162 const*>(&wraw); +#pragma unroll + for (int h = 0; h < 4; ++h) + { + float2 const uf = __bfloat1622float2(uu[h]); + float2 const wf = __bfloat1622float2(wwv[h]); + float2 const nf + = __bfloat1622float2(__floats2bfloat162_rn(uf.x * inv * wf.x, uf.y * inv * wf.y)); + v[q * 8 + 2 * h] = nf.x; + v[q * 8 + 2 * h + 1] = nf.y; + } + } + float amax = 0.f; +#pragma unroll + for (int t = 0; t < 32; ++t) + amax = fmaxf(amax, fabsf(v[t])); + float const sf = amax * rcp_approx_ftz(448.0f); + __nv_fp8_e8m0 code; + code.__x = __nv_cvt_float_to_e8m0(sf, __NV_SATFINITE, cudaRoundPosInf); + float const out_scale = amax != 0.f ? rcp_approx_ftz(static_cast(code)) : 0.0f; + uint32_t packed[8]; +#pragma unroll + for (int t = 0; t < 8; ++t) + { + __nv_fp8x2_e4m3 const lo(make_float2(v[4 * t] * out_scale, v[4 * t + 1] * out_scale)); + __nv_fp8x2_e4m3 const hi(make_float2(v[4 * t + 2] * out_scale, v[4 * t + 3] * out_scale)); + packed[t] = static_cast(lo.__x) | (static_cast(hi.__x) << 16); + } + uint4* dst = reinterpret_cast(xq_s + m * kc + b * 32); + dst[0] = make_uint4(packed[0], packed[1], packed[2], packed[3]); + dst[1] = make_uint4(packed[4], packed[5], packed[6], packed[7]); + xc_s[m * kc_blocks + b] = code.__x; + } + // Publish the quantized slice for the other tiles. + for (int p = threadIdx.x; p < M * kc_blocks; p += kThreads) + { + int const m = p / kc_blocks; + int const bb = p % kc_blocks; + uint4 const* src = reinterpret_cast(xq_s + m * kc + bb * 32); + uint4* dst = reinterpret_cast(ar.xq_g + static_cast(m) * K + k_begin + bb * 32); + dst[0] = src[0]; + dst[1] = src[1]; + ar.xc_g[static_cast(m) * (K / 32) + kb_begin + bb] = xc_s[m * kc_blocks + bb]; + } + __syncthreads(); + if (threadIdx.x == 0) + { + asm volatile("fence.acq_rel.gpu;" ::: "memory"); + asm volatile("st.release.gpu.global.s32 [%0], %1;" ::"l"(ar.flags + rank), "r"(call + 1) : "memory"); + } + } + else + { + if (threadIdx.x == 0) + { + int v; + long long spins = 0; + do + { + asm volatile("ld.acquire.gpu.global.s32 %0, [%1];" : "=r"(v) : "l"(ar.flags + rank) : "memory"); + if (++spins > (1LL << 31)) + __trap(); + } while (v - (call + 1) < 0); + } + __syncthreads(); + for (int p = threadIdx.x; p < M * kc_blocks; p += kThreads) + { + int const m = p / kc_blocks; + int const bb = p % kc_blocks; + uint4 const* src + = reinterpret_cast(ar.xq_g + static_cast(m) * K + k_begin + bb * 32); + uint4* dst = reinterpret_cast(xq_s + m * kc + bb * 32); + dst[0] = __ldcg(src); + dst[1] = __ldcg(src + 1); + xc_s[m * kc_blocks + bb] = __ldcg(ar.xc_g + static_cast(m) * (K / 32) + kb_begin + bb); + } + if constexpr (CLUSTER_K > 1) + { + // Keep the cluster barrier phases aligned with tile 0's prologue. + asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory"); + asm volatile("barrier.cluster.arrive.release.aligned;\n" ::: "memory"); + } + } + } + else if (preq.xq != nullptr) + { + // Stage this CTA's K slice of the pre-quantized activation. + int const vec_per_row = kc / 16; + int const vecs = M * vec_per_row; +#pragma unroll 4 + for (int p = threadIdx.x; p < vecs; p += kThreads) + { + int const m = p / vec_per_row; + int const v = p - m * vec_per_row; + reinterpret_cast(xq_s + m * kc)[v] + = __ldcg(reinterpret_cast(preq.xq + static_cast(m) * K + k_begin) + v); + } + for (int p = threadIdx.x; p < M * kc_blocks; p += kThreads) + { + int const m = p / kc_blocks; + int const b = p - m * kc_blocks; + int const kb = kb_begin + b; + xc_s[m * kc_blocks + b] = __ldcg(preq.xsf + static_cast(m / 128) * preq.sf_k_tiles * 512 + + (kb / 4) * 512 + (m % 32) * 16 + ((m % 128) / 32) * 4 + kb % 4); + } + } + else + { + // Quantize the activation columns owned by this CTA (native mxfp8_quantize math). Each thread + // issues the loads of all its blocks before converting any, so the CTA pays one L2 round trip. + int const items = M * kc_blocks; + for (int p0 = threadIdx.x; p0 < items; p0 += 2 * kThreads) + { + uint4 raw[2][4]; +#pragma unroll + for (int u = 0; u < 2; ++u) + { + int const p = p0 + u * kThreads; + if (p < items) + { + int const m = p / kc_blocks; + int const b = p % kc_blocks; + uint4 const* src + = reinterpret_cast(x + static_cast(m) * ldx + k_begin + b * 32); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + raw[u][q] = src[q]; + } + } + } +#pragma unroll + for (int u = 0; u < 2; ++u) + { + int const p = p0 + u * kThreads; + if (p >= items) + { + continue; + } + int const m = p / kc_blocks; + int const b = p % kc_blocks; + float v[32]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&raw[u][q]); +#pragma unroll + for (int t = 0; t < 4; ++t) + { + float2 const f = __bfloat1622float2(h[t]); + v[q * 8 + t * 2] = f.x; + v[q * 8 + t * 2 + 1] = f.y; + } + } + float amax = 0.f; +#pragma unroll + for (int t = 0; t < 32; ++t) + { + amax = fmaxf(amax, fabsf(v[t])); + } + float const sf = amax * rcp_approx_ftz(448.0f); + __nv_fp8_e8m0 code; + code.__x = __nv_cvt_float_to_e8m0(sf, __NV_SATFINITE, cudaRoundPosInf); + float const out_scale = amax != 0.f ? rcp_approx_ftz(static_cast(code)) : 0.0f; + uint32_t packed[8]; +#pragma unroll + for (int t = 0; t < 8; ++t) + { + __nv_fp8x2_e4m3 const lo(make_float2(v[4 * t] * out_scale, v[4 * t + 1] * out_scale)); + __nv_fp8x2_e4m3 const hi(make_float2(v[4 * t + 2] * out_scale, v[4 * t + 3] * out_scale)); + packed[t] = static_cast(lo.__x) | (static_cast(hi.__x) << 16); + } + uint4* dst = reinterpret_cast(xq_s + m * kc + b * 32); + dst[0] = make_uint4(packed[0], packed[1], packed[2], packed[3]); + dst[1] = make_uint4(packed[4], packed[5], packed[6], packed[7]); + xc_s[m * kc_blocks + b] = code.__x; + } + } + } + __syncthreads(); + mbar_wait(&bar, 0); + + int const g = lane >> 2; + int const j = lane & 3; + int const col0 = 2 * j; + int const bpw = kc_blocks / WARPS; + int const chunks = kc_blocks / 4; + float acc[RT][NT][4] = {}; +#pragma unroll 2 + for (int i = 0; i < bpw; ++i) + { + int const b = warp * bpw + i; + int const kb = kb_begin + b; + uint2 xb[NT]; + float xs0[NT], xs1[NT]; +#pragma unroll + for (int t = 0; t < NT; ++t) + { + int const row = 8 * t + g; + int const c0 = 8 * t + col0; + xb[t] = row < M ? *reinterpret_cast(xq_s + row * kc + b * 32 + 8 * j) : make_uint2(0, 0); + xs0[t] = c0 < M ? pow2_from_ue8m0(xc_s[c0 * kc_blocks + b]) : 0.f; + xs1[t] = c0 + 1 < M ? pow2_from_ue8m0(xc_s[(c0 + 1) * kc_blocks + b]) : 0.f; + } +#pragma unroll + for (int rt = 0; rt < RT; ++rt) + { + uint8_t const* wa_s = w_s + (16 * rt + g) * row_stride + 8 * j + b * 32; + uint2 const wa = *reinterpret_cast(wa_s); + uint2 const wb = *reinterpret_cast(wa_s + 8 * row_stride); + int const inner_m = ((row0 + 16 * rt) % 128) / 32; + uint8_t const* sft = sf_s + (rt * chunks + b / 4) * 256 + inner_m * 4 + kb % 4; + float const wsa = pow2_from_ue8m0(sft[g * 16]); + float const wsb = pow2_from_ue8m0(sft[(g + 8) * 16]); +#pragma unroll + for (int t = 0; t < NT; ++t) + { + float d[4]; + mma_e4m3(d, wa.x, wb.x, wa.y, wb.y, xb[t].x, xb[t].y); + acc[rt][t][0] += d[0] * (wsa * xs0[t]); + acc[rt][t][1] += d[1] * (wsa * xs1[t]); + acc[rt][t][2] += d[2] * (wsb * xs0[t]); + acc[rt][t][3] += d[3] * (wsb * xs1[t]); + } + } + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + if (red_aliased) + { + __syncthreads(); + } +#pragma unroll + for (int rt = 0; rt < RT; ++rt) + { +#pragma unroll + for (int t = 0; t < NT; ++t) + { + int const ra = 16 * rt + g; + red[warp * kTile + ra * 8 * NT + 8 * t + col0] = acc[rt][t][0]; + red[warp * kTile + ra * 8 * NT + 8 * t + col0 + 1] = acc[rt][t][1]; + red[warp * kTile + (ra + 8) * 8 * NT + 8 * t + col0] = acc[rt][t][2]; + red[warp * kTile + (ra + 8) * 8 * NT + 8 * t + col0 + 1] = acc[rt][t][3]; + } + } + __syncthreads(); + + // Warp reduction in fixed order. Peers push their partial tile into the root's shared memory with + // st.async (completing bytes on the root's mbarrier) and exit; the root sums ranks in order. + if constexpr (CLUSTER_K > 1) + { + asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory"); + } + if (CLUSTER_K == 1 || rank == 0) + { + if constexpr (CLUSTER_K > 1) + { + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(smem_u32(&red_bar)) + : "memory"); + } + } + __syncthreads(); + for (int idx = threadIdx.x; idx < kTile; idx += kThreads) + { + int const r = idx / (8 * NT); + int const m = idx % (8 * NT); + if (m < M) + { + float total = 0.f; +#pragma unroll + for (int ww = 0; ww < WARPS; ++ww) + { + total += red[ww * kTile + idx]; + } +#pragma unroll + for (int c = 1; c < CLUSTER_K; ++c) + { + total += cluster_red[(c - 1) * kTile + idx]; + } + out[static_cast(m) * ldo + row0 + r] = __float2bfloat16_rn(total); + } + } + } + else + { + uint32_t const remote_bar = mapa_shared(smem_u32(&red_bar), 0); + uint32_t const remote_dst = mapa_shared(smem_u32(cluster_red + (rank - 1) * kTile), 0); + for (int idx = threadIdx.x; idx < kTile; idx += kThreads) + { + float s = 0.f; +#pragma unroll + for (int ww = 0; ww < WARPS; ++ww) + { + s += red[ww * kTile + idx]; + } + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + remote_dst + idx * 4), + "r"(__float_as_uint(s)), "r"(remote_bar) + : "memory"); + } + } +#endif +} + +template +void launch(__nv_bfloat16 const* x, int ldx, int M, int N, int K, uint8_t const* w, uint8_t const* wsf, + __nv_bfloat16* out, int ldo, bool pdl, cudaStream_t stream) +{ + constexpr bool AR = false; + int const kc = K / CLUSTER_K; + int const kc_blocks = kc / 32; + TLLM_CHECK_WITH_INFO(N % (16 * RT) == 0 && K % (CLUSTER_K * 128) == 0 && kc_blocks % WARPS == 0, + "unsupported shape N=%d K=%d", N, K); + constexpr int kTile = 16 * RT * 8 * NT; + size_t const w_bytes = 16 * RT * (kc + kRowPad); + size_t const red_bytes = static_cast(WARPS) * kTile * sizeof(float); + int const red_aliased = red_bytes <= w_bytes ? 1 : 0; + size_t const smem = w_bytes + RT * (kc_blocks / 4) * 256 + M * kc + ((M * kc_blocks + 15) / 16) * 16 + + (CLUSTER_K - 1) * kTile * sizeof(float) + (red_aliased ? 0 : red_bytes); + auto kernel = mxfp8_low_m_gemm_tma_kernel; + // Configure once, outside any CUDA graph capture, to the largest supported request. + static bool const configured = [&]() + { + int device = 0, max_optin = 0; + cudaGetDevice(&device); + cudaDeviceGetAttribute(&max_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device); + cudaFuncAttributes attrs = {}; + if (cudaFuncGetAttributes(&attrs, kernel) != cudaSuccess) + { + cudaGetLastError(); + return false; + } + bool const ok = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + max_optin - static_cast(attrs.sharedSizeBytes)) + == cudaSuccess; + if (!ok) + { + cudaGetLastError(); + } + return ok; + }(); + TLLM_CHECK_WITH_INFO(configured, "failed to configure dynamic shared memory"); + cudaFuncAttributes fattrs = {}; + cudaFuncGetAttributes(&fattrs, kernel); + TLLM_CHECK_WITH_INFO(static_cast(smem) <= fattrs.maxDynamicSharedSizeBytes, + "shared memory request exceeds the configured limit"); + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(N / (16 * RT) * CLUSTER_K); + cfg.blockDim = dim3(WARPS * 32); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + attr[1].id = cudaLaunchAttributeClusterDimension; + attr[1].val.clusterDim.x = CLUSTER_K; + attr[1].val.clusterDim.y = 1; + attr[1].val.clusterDim.z = 1; + cfg.attrs = attr; + cfg.numAttrs = 2; + cudaError_t const err = cudaLaunchKernelEx( + &cfg, kernel, x, ldx, M, N, K, w, wsf, (K / 32 + 3) / 4, out, ldo, red_aliased, ArArgs{}, PfArgs{}, PreqArgs{}); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +} // namespace + +bool isMinimaxM3Mxfp8LowMGemmSupported(int32_t warps, int32_t clusterK, int32_t rowTiles, int32_t m) +{ + bool const built = (warps == 4 && clusterK == 4 && rowTiles == 2) || (warps == 8 && clusterK == 2 && rowTiles == 4) + || (warps == 4 && clusterK == 4 && rowTiles == 1) || (warps == 4 && clusterK == 2 && rowTiles == 2); + return built && m >= 1 && m <= 8; +} + +void invokeMinimaxM3Mxfp8LowMGemm(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t n, int32_t k, + uint8_t const* w, uint8_t const* wsf, __nv_bfloat16* out, int32_t ldo, int32_t warps, int32_t clusterK, + int32_t rowTiles, bool pdl, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(isMinimaxM3Mxfp8LowMGemmSupported(warps, clusterK, rowTiles, m), + "unsupported configuration (warps %d, cluster_k %d, row_tiles %d) at M=%d", warps, clusterK, rowTiles, m); + // One 8-row activation tile (NT = 1) covers every built row count. + if (warps == 4 && clusterK == 4 && rowTiles == 2) + { + launch<4, 4, 1, 2>(x, ldx, m, n, k, w, wsf, out, ldo, pdl, stream); + } + else if (warps == 8 && clusterK == 2 && rowTiles == 4) + { + launch<8, 2, 1, 4>(x, ldx, m, n, k, w, wsf, out, ldo, pdl, stream); + } + else if (warps == 4 && clusterK == 4 && rowTiles == 1) + { + launch<4, 4, 1, 1>(x, ldx, m, n, k, w, wsf, out, ldo, pdl, stream); + } + else + { + launch<4, 2, 1, 2>(x, ldx, m, n, k, w, wsf, out, ldo, pdl, stream); + } +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/mxfp8LowMGemm.h b/cpp/tensorrt_llm/kernels/minimaxM3/mxfp8LowMGemm.h new file mode 100644 index 000000000000..433cb6a1fe07 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/mxfp8LowMGemm.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 +{ + +//! Whether `invokeMinimaxM3Mxfp8LowMGemm` is built for this (warps, clusterK, rowTiles) configuration at M rows: +//! (4, 4, 2), (8, 2, 4), (4, 4, 1) and (4, 2, 2), each for 1 <= M <= 8. +bool isMinimaxM3Mxfp8LowMGemmSupported(int32_t warps, int32_t clusterK, int32_t rowTiles, int32_t m); + +//! Fused MXFP8 activation quantization + low-M MXFP8 GEMM: `out[M, N] = mxfp8(x)[M, K] @ w[N, K]^T` in BF16. +//! +//! `x` is BF16 [M, K] with row stride `ldx`; `w` is E4M3 [N, K] row-major with UE8M0 scales `wsf` in the 128x4 +//! swizzled layout of `mxfp8_quantize(w, True)`; `out` is BF16 [M, N] with row stride `ldo`. Each CTA owns +//! 16 * rowTiles weight rows and K / clusterK columns; `warps` warps per CTA. +void invokeMinimaxM3Mxfp8LowMGemm(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t n, int32_t k, + uint8_t const* w, uint8_t const* wsf, __nv_bfloat16* out, int32_t ldo, int32_t warps, int32_t clusterK, + int32_t rowTiles, bool pdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/producerPw.cu b/cpp/tensorrt_llm/kernels/minimaxM3/producerPw.cu new file mode 100644 index 000000000000..fb3e4bfb3a89 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/producerPw.cu @@ -0,0 +1,324 @@ +/* + * 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/minimaxM3/producerPw.h" + +#include +#include +#include + +// MiniMax-M3 NVFP4 QKV + indexer norm/RoPE/KV-insert producer with an early PDL trigger. +// +// Same math and outputs as minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel (fusedQKNormRopeKernel.cu, built with +// --use_fast_math like the rest of the library), specialized to 1 KV head, 1 index head, head_dim 128, rotary_dim 64, +// page 128, NVFP4 main KV and FP8 index-K, with 8 (TP8) or 16 (TP4) query heads per rank. One warp per (token, +// packed head). +// +// The decode producer is latency bound: after griddepcontrol.wait the native kernel runs a chain of dependent global +// loads (packed row -> norm weight -> position id -> RoPE row -> cache slot -> quant scale). This kernel reads the +// constant norm weights and KV quant scale before griddepcontrol.wait and issues the packed row, position id and cache +// slot loads together after it, followed by the RoPE row. With `early` it also calls griddepcontrol.launch_dependents +// on entry: consumers read its outputs only after their own griddepcontrol.wait, which waits for this grid to complete, +// so the early trigger only lets the next kernel launch and run its pre-wait prologue while this one runs. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3 +{ +namespace producerPw +{ +namespace +{ + +constexpr int kHeadDim = 128; +constexpr int kRotaryDim = 64; +constexpr int kPageSize = 128; +constexpr int kElems = kHeadDim / 32; +constexpr int kHeadsKV = 1; +constexpr int kHeadsIndex = 1; +constexpr int kBlock = 256; + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ uint32_t sf_group_mask4() +{ + uint32_t laneId; + asm("mov.u32 %0, %%laneid;" : "=r"(laneId)); + return 0xFu << (laneId & ~3u); +} + +__device__ __forceinline__ uint32_t fp32x8_to_e2m1(float (&a)[8]) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(val) + : "f"(a[0]), "f"(a[1]), "f"(a[2]), "f"(a[3]), "f"(a[4]), "f"(a[5]), "f"(a[6]), "f"(a[7])); + return val; +#else + return 0; +#endif +} + +__device__ __forceinline__ void store_fp8x4(__nv_fp8_e4m3* out, int64_t off, float const (&e)[kElems]) +{ + __nv_fp8x2_e4m3 const low(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const high(make_float2(e[2], e[3])); + *reinterpret_cast(out + off) = static_cast(low.__x) | (static_cast(high.__x) << 16); +} + +__device__ __forceinline__ uint16_t quantize_fp4x4(float const (&e)[4], float secondLevelScale, uint8_t* scaleOut) +{ + __nv_bfloat16 rounded[4]; + float localMax = 0.0F; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + rounded[i] = __float2bfloat16_rn(e[i]); + localMax = fmaxf(localMax, fabsf(__bfloat162float(rounded[i]))); + } + uint32_t const groupMask = sf_group_mask4(); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 1)); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 2)); + float const scaleValue = secondLevelScale * (localMax * rcp_approx_ftz(6.0F)); + __nv_fp8_e4m3 const encodedScale(scaleValue); + if (scaleOut != nullptr) + *scaleOut = encodedScale.__x; + float const narrowedScale = static_cast(encodedScale); + float const outputScale + = localMax != 0.0F ? rcp_approx_ftz(narrowedScale * rcp_approx_ftz(secondLevelScale)) : 0.0F; + float q[8] = {}; +#pragma unroll + for (int i = 0; i < 4; ++i) + q[i] = __bfloat162float(rounded[i]) * outputScale; + return static_cast(fp32x8_to_e2m1(q)); +} + +template +__global__ void __launch_bounds__(kBlock) m3_producer_kernel(Params p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (kEarly) + { + asm volatile("griddepcontrol.launch_dependents;"); + } +#endif + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarp = blockIdx.x * (kBlock / 32) + warpId; + int const tokenIdx = globalWarp / kTotalHeads; + int const localHead = globalWarp % kTotalHeads; + if (tokenIdx >= p.numTokens) + return; + + constexpr int kBegin = kHeadsQ; + constexpr int vBegin = kBegin + kHeadsKV; + constexpr int indexQBegin = vBegin + kHeadsKV; + constexpr int indexKHead = indexQBegin + kHeadsIndex; + bool const isQ = localHead < kBegin; + bool const isK = localHead >= kBegin && localHead < vBegin; + bool const isV = localHead >= vBegin && localHead < indexQBegin; + bool const isIndexQ = localHead >= indexQBegin && localHead < indexKHead; + bool const isIndexK = localHead == indexKHead; + + // Only constant data is read before the wait: norm weights and the KV quant scale. + float w[kElems] = {}; + float quantScale = 0.0F; + __nv_bfloat16 const* normWeight + = isQ ? p.qWeight : (isK ? p.kWeight : (isIndexQ ? p.indexQWeight : p.indexKWeight)); + if (!isV) + { + uint2 const wv = *reinterpret_cast(normWeight + laneId * kElems); + float2 const w01 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.x)); + float2 const w23 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.y)); + w[0] = w01.x; + w[1] = w01.y; + w[2] = w23.x; + w[3] = w23.y; + } + if (isK || isV) + { + quantScale = p.kvQuantScale[isV ? 2 : 1]; + } + + int64_t const inputOffset = (static_cast(tokenIdx) * kTotalHeads + localHead) * kHeadDim + laneId * kElems; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + // Position ids and cache slots are per-step metadata, read after the wait. The packed row, the position id and + // the slot are independent loads, and the RoPE row load is issued before the norm reduction, so the + // native kernel's serial chain of dependent loads becomes two round trips. + uint2 const packed = *reinterpret_cast(p.packed + inputOffset); + bool const rotates = !isV && laneId * kElems < kRotaryDim; + int const positionId = rotates ? p.positionIds[tokenIdx] : 0; + int const slot = (isK || isV || isIndexK) ? p.outCacheLoc[tokenIdx] : 0; + float cosine[kElems] = {}; + float sine[kElems] = {}; + if (rotates) + { + int64_t const ropeRow = static_cast(positionId) * kRotaryDim; + int const coefficient = (laneId * kElems) % (kRotaryDim / 2); + float4 const c = *reinterpret_cast(p.rotaryCosSin + ropeRow + coefficient); + float4 const s = *reinterpret_cast(p.rotaryCosSin + ropeRow + kRotaryDim / 2 + coefficient); + cosine[0] = c.x; + cosine[1] = c.y; + cosine[2] = c.z; + cosine[3] = c.w; + sine[0] = s.x; + sine[1] = s.y; + sine[2] = s.z; + sine[3] = s.w; + } + + float e[kElems]; + float sumSquares = 0.0F; +#pragma unroll + for (int pair = 0; pair < 2; ++pair) + { + float2 const v = __bfloat1622float2( + *reinterpret_cast<__nv_bfloat162 const*>(reinterpret_cast(&packed) + pair)); + e[2 * pair] = v.x; + e[2 * pair + 1] = v.y; + if (!isV) + sumSquares += v.x * v.x + v.y * v.y; + } + + if (!isV) + { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + sumSquares += __shfl_xor_sync(0xffffffff, sumSquares, mask, 32); + float const rmsReciprocal = rsqrtf(sumSquares / static_cast(kHeadDim) + p.eps); +#pragma unroll + for (int i = 0; i < kElems; ++i) + { + e[i] *= rmsReciprocal * (1.0F + w[i]); + } + constexpr int kPairOffset = (kRotaryDim / 2) / kElems; +#pragma unroll + for (int i = 0; i < kElems; ++i) + { + int const dim = laneId * kElems + i; + float paired = __shfl_xor_sync(0xffffffff, e[i], kPairOffset); + if (dim < kRotaryDim) + { + if (dim < kRotaryDim / 2) + paired = -paired; + // The native kernel's contraction of e * cos + paired * sin; spelled out so this build matches it. + e[i] = __fmaf_rn(paired, sine[i], __fmul_rn(e[i], cosine[i])); + } + } + } + + if (isQ) + { + store_fp8x4(p.qOut, (static_cast(tokenIdx) * kHeadsQ + localHead) * kHeadDim + laneId * kElems, e); + return; + } + if (isIndexQ) + { + store_fp8x4(p.indexQOut, + (static_cast(tokenIdx) * kHeadsIndex + localHead - indexQBegin) * kHeadDim + laneId * kElems, e); + return; + } + + if (slot < 0) + return; + int const page = slot / kPageSize; + if (page >= p.numPages) + return; + int const withinPage = slot & (kPageSize - 1); + if (isIndexK) + { + store_fp8x4(p.indexK, + static_cast(page) * p.indexPageStride + static_cast(withinPage) * p.indexTokenStride + + laneId * kElems, + e); + return; + } + + int const plane = isV ? 1 : 0; + int const scaleColumn = laneId / 4; + int64_t const scaleWithinPage = isV + ? static_cast(withinPage / 4) * 32 + static_cast(scaleColumn) * 4 + withinPage % 4 + : static_cast(withinPage) * 8 + scaleColumn; + int64_t const scaleOffset = static_cast(page) * p.scalePageStride + + static_cast(plane) * p.scalePlaneStride + scaleWithinPage; + uint8_t* scaleOut = laneId % 4 == 0 ? p.kvScale + scaleOffset : nullptr; + uint16_t const packedFp4 = quantize_fp4x4(e, quantScale, scaleOut); + int64_t const dataOffset = static_cast(page) * p.dataPageStride + + static_cast(plane) * p.dataPlaneStride + static_cast(withinPage) * p.dataTokenStride + + laneId * 2; + *reinterpret_cast(p.kvData + dataOffset) = packedFp4; +#endif +} + +template +cudaError_t launch(Params const& p, int numTokens, bool pdl, cudaStream_t stream) +{ + constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; + int const warps = numTokens * kTotalHeads; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3((warps + kBlock / 32 - 1) / (kBlock / 32)); + cfg.blockDim = dim3(kBlock); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, m3_producer_kernel, p); +} + +} // namespace + +cudaError_t invokeProducerPw(Params const& p, int numTokens, int headsQ, bool pdl, bool early, cudaStream_t stream) +{ + if (numTokens <= 0) + return cudaSuccess; + if (headsQ == 16) + return early ? launch<16, true>(p, numTokens, pdl, stream) : launch<16, false>(p, numTokens, pdl, stream); + if (headsQ == 8) + return early ? launch<8, true>(p, numTokens, pdl, stream) : launch<8, false>(p, numTokens, pdl, stream); + return cudaErrorInvalidValue; +} + +} // namespace producerPw +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/producerPw.h b/cpp/tensorrt_llm/kernels/minimaxM3/producerPw.h new file mode 100644 index 000000000000..6d6848e0e199 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/producerPw.h @@ -0,0 +1,69 @@ +/* + * 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 +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3 +{ +//! The MiniMax-M3 NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer with its constant loads before the PDL +//! grid-dependency wait and an optional early dependents trigger (sm_103a). Same outputs, bit for bit, as +//! minimaxM3Nvfp4QKVIndexerNormRopeKVInsert. See producerPw.cu. +namespace producerPw +{ + +struct Params +{ + __nv_bfloat16 const* packed; + __nv_fp8_e4m3* qOut; + __nv_fp8_e4m3* indexQOut; + uint8_t* kvData; + uint8_t* kvScale; + __nv_fp8_e4m3* indexK; + int const* outCacheLoc; + float const* kvQuantScale; + int64_t dataPageStride, dataPlaneStride, dataHeadStride, dataTokenStride; + int64_t scalePageStride, scalePlaneStride, scaleHeadStride; + int64_t indexPageStride, indexTokenStride, numPages; + int numTokens; + float eps; + __nv_bfloat16 const* qWeight; + __nv_bfloat16 const* kWeight; + __nv_bfloat16 const* indexQWeight; + __nv_bfloat16 const* indexKWeight; + float const* rotaryCosSin; + int const* positionIds; +}; + +//! Launch over `numTokens` tokens, one warp per (token, packed head); headsQ is 8 or 16. +cudaError_t invokeProducerPw(Params const& p, int numTokens, int headsQ, bool pdl, bool early, cudaStream_t stream); + +} // namespace producerPw +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.cu b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.cu new file mode 100644 index 000000000000..d12b765e5ecf --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.cu @@ -0,0 +1,606 @@ +/* + * 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/minimaxM3/routerGemv.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" + +#include +#include +#include + +// The kernels use sm_100-family features (thread-block clusters, cp.async.bulk, st.async, mbarrier transaction +// counts, block-scaled conversions). Every other architecture of the build compiles them to empty bodies; the host +// entry points refuse to launch on those devices. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030) \ + && (defined(__CUDA_ARCH_FEAT_SM100_ALL) || defined(__CUDA_ARCH_FEAT_SM103_ALL) \ + || defined(__CUDA_ARCH_FAMILY_SPECIFIC__)) +#define MINIMAX_M3_KERNELS_ENABLED 1 +#else +#define MINIMAX_M3_KERNELS_ENABLED 0 +#endif + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kMaxM = 16; + +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void bulk_g2s(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];\n" ::"r"( + smem_u32(dst)), + "l"(src), "r"(bytes), "r"(smem_u32(bar)) + : "memory"); +} + +__device__ __forceinline__ void mbar_wait(uint64_t* bar, uint32_t parity) +{ + asm volatile( + "{\n .reg .pred p;\n WAIT_%=:\n mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n @!p bra " + "WAIT_%=;\n}\n" ::"r"(smem_u32(bar)), + "r"(parity) + : "memory"); +} + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ uint32_t e2m1x8(float2 const (&v)[4]) +{ + uint32_t r; + asm("{\n .reg .b8 b0, b1, b2, b3;\n" + " cvt.rn.satfinite.e2m1x2.f32 b0, %2, %1;\n cvt.rn.satfinite.e2m1x2.f32 b1, %4, %3;\n" + " cvt.rn.satfinite.e2m1x2.f32 b2, %6, %5;\n cvt.rn.satfinite.e2m1x2.f32 b3, %8, %7;\n" + " mov.b32 %0, {b0, b1, b2, b3};\n}\n" + : "=r"(r) + : "f"(v[0].x), "f"(v[0].y), "f"(v[1].x), "f"(v[1].y), "f"(v[2].x), "f"(v[2].y), "f"(v[3].x), "f"(v[3].y)); + return r; +} + +// NVFP4 (16-element E4M3 block scale, per-tensor global scale) quantization of one 16-value block, bit-identical +// to trtllm fp4_quantize (cvt_warp_fp16_to_fp4, linear scale layout). +__device__ __forceinline__ void quant_block16(__nv_bfloat16 const* src, float sf_scale, uint2* q, uint8_t* sf) +{ + uint4 const v0 = *reinterpret_cast(src); + uint4 const v1 = *reinterpret_cast(src + 8); + __nv_bfloat162 h[8]; + *reinterpret_cast(&h[0]) = v0; + *reinterpret_cast(&h[4]) = v1; + __nv_bfloat162 mx = __habs2(h[0]); +#pragma unroll + for (int i = 1; i < 8; ++i) + mx = __hmax2(mx, __habs2(h[i])); + float const vec_max = fmaxf(__low2float(mx), __high2float(mx)); + float sf_value = sf_scale * (vec_max * rcp_approx_ftz(6.0f)); + __nv_fp8_e4m3 const tmp = __nv_fp8_e4m3(sf_value); + sf_value = static_cast(tmp); + float const out_scale = vec_max != 0.f ? rcp_approx_ftz(sf_value * rcp_approx_ftz(sf_scale)) : 0.f; + float2 lo[4], hi[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + lo[i] = __bfloat1622float2(h[i]); + lo[i].x *= out_scale; + lo[i].y *= out_scale; + hi[i] = __bfloat1622float2(h[i + 4]); + hi[i].x *= out_scale; + hi[i].y *= out_scale; + } + *q = make_uint2(e2m1x8(lo), e2m1x8(hi)); + *sf = tmp.__x; +} + +// grid: E / ROWS; block: 256 threads. smem: ROWS * K fp32 weights + M * K bf16 activations. +// With q_out, the CTAs also emit the NVFP4 activation ([M, K / 2] packed, [M, K / 16] E4M3 scales) for the MoE. +template +__global__ void __launch_bounds__(256) m3_router_gemv_kernel(__nv_bfloat16 const* __restrict__ x, int ldx, + float const* __restrict__ w, float* __restrict__ out, int M, int K, int E, uint8_t* __restrict__ q_out, + uint8_t* __restrict__ sf_out, float const* __restrict__ sf_scale) +{ +#if MINIMAX_M3_KERNELS_ENABLED + extern __shared__ __align__(128) uint8_t smem[]; + __shared__ __align__(8) uint64_t w_bar; + __shared__ __align__(8) uint64_t x_bar; + __shared__ float red[8][ROWS][kMaxM]; + float* w_s = reinterpret_cast(smem); + __nv_bfloat16* x_s = reinterpret_cast<__nv_bfloat16*>(smem + ROWS * K * sizeof(float)); + int const e0 = blockIdx.x * ROWS; + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + if (threadIdx.x == 0) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" ::"r"(smem_u32(&w_bar))); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" ::"r"(smem_u32(&x_bar))); + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smem_u32(&w_bar)), + "r"(static_cast(ROWS * K * sizeof(float))) + : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smem_u32(&x_bar)), + "r"(static_cast(M * K * sizeof(__nv_bfloat16))) + : "memory"); + } + __syncthreads(); + // Router weights do not depend on the producer kernel. + if (warp == 0 && lane < ROWS) + { + bulk_g2s(w_s + lane * K, w + static_cast(e0 + lane) * K, K * sizeof(float), &w_bar); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + if (warp == 1 && lane < M) + { + bulk_g2s(x_s + lane * K, x + static_cast(lane) * ldx, K * sizeof(__nv_bfloat16), &x_bar); + } + mbar_wait(&x_bar, 0); + if (q_out != nullptr) + { + float const s = *sf_scale; + int const kblocks = K / 16; + for (int b = blockIdx.x * blockDim.x + threadIdx.x; b < M * kblocks; b += gridDim.x * blockDim.x) + { + int const m = b / kblocks; + int const kb = b - m * kblocks; + quant_block16(x_s + m * K + kb * 16, s, + reinterpret_cast(q_out + static_cast(m) * (K / 2) + kb * 8), + sf_out + static_cast(m) * kblocks + kb); + } + } + mbar_wait(&w_bar, 0); + + float acc[ROWS][kMaxM]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) + { +#pragma unroll + for (int m = 0; m < kMaxM; ++m) + { + acc[r][m] = 0.f; + } + } + // Each thread owns 4 consecutive columns per 1024-column stride. + for (int k = threadIdx.x * 4; k < K; k += 256 * 4) + { + float4 wv[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) + { + wv[r] = *reinterpret_cast(w_s + r * K + k); + } +#pragma unroll + for (int m = 0; m < kMaxM; ++m) + { + if (m < M) + { + uint2 const raw = *reinterpret_cast(x_s + m * K + k); + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&raw); + float2 const a = __bfloat1622float2(h[0]); + float2 const b = __bfloat1622float2(h[1]); +#pragma unroll + for (int r = 0; r < ROWS; ++r) + { + acc[r][m] = fmaf(a.x, wv[r].x, acc[r][m]); + acc[r][m] = fmaf(a.y, wv[r].y, acc[r][m]); + acc[r][m] = fmaf(b.x, wv[r].z, acc[r][m]); + acc[r][m] = fmaf(b.y, wv[r].w, acc[r][m]); + } + } + } + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +#pragma unroll + for (int r = 0; r < ROWS; ++r) + { +#pragma unroll + for (int m = 0; m < kMaxM; ++m) + { + if (m < M) + { + float v = acc[r][m]; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + v += __shfl_xor_sync(0xFFFFFFFFu, v, o); + } + if (lane == 0) + { + red[warp][r][m] = v; + } + } + } + } + __syncthreads(); + if (threadIdx.x < ROWS * M) + { + int const r = threadIdx.x / M; + int const m = threadIdx.x % M; + float s = 0.f; +#pragma unroll + for (int ww = 0; ww < 8; ++ww) + { + s += red[ww][r][m]; + } + out[static_cast(m) * E + e0 + r] = s; + } +#endif +} + +// ---------------- 9..48 tokens: TF32 tensor-core router (same precision as the native TF32 GEMM) ---------------- +// grid: 8 expert tiles (16 experts) x CL CTAs, one CL-CTA cluster per tile, 256 threads. CTA r of a cluster computes +// the [16 experts x M] partials of its 16 / CL K slices (16 slices of 384), each over the same 8-warp split. CTA r +// owns 16 / CL experts of the tile: every CTA st.asyncs its partials for those experts into the owner's mailbox, and +// the owner sums the 16 slices in slice order, so the logits do not depend on CL. Optionally each CTA also emits the +// NVFP4 activation for its K range and the token rows m with m % 8 == expert tile. +constexpr int kTcExperts = 16; +constexpr int kTcSlices = 16; +constexpr int kTcMaxM = 48; + +__device__ __forceinline__ uint32_t f2tf32(float f) +{ + uint32_t r; + asm("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(r) : "f"(f)); + return r; +} + +__device__ __forceinline__ uint32_t mapa_u32(void const* p, uint32_t rank) +{ + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n" : "=r"(r) : "r"(smem_u32(p)), "r"(rank)); + return r; +} + +template +__global__ void __launch_bounds__(256) m3_router_tc_kernel(__nv_bfloat16 const* __restrict__ x, int ldx, + float const* __restrict__ w, float* __restrict__ out, int M, int K, uint8_t* __restrict__ q_out, + uint8_t* __restrict__ sf_out, float const* __restrict__ sf_scale) +{ +#if MINIMAX_M3_KERNELS_ENABLED + extern __shared__ __align__(128) uint8_t smem[]; + __shared__ __align__(8) uint64_t w_bar; + __shared__ __align__(8) uint64_t x_bar; + __shared__ __align__(8) uint64_t mb_bar; + constexpr int SPC = kTcSlices / CL; // K slices per CTA + constexpr int EPC = kTcExperts / CL; // experts owned per CTA + __shared__ float red[8][kTcExperts][kTcMaxM]; + __shared__ float part[SPC][kTcExperts][kTcMaxM]; // this CTA's slice partials + __shared__ float mailbox[CL][SPC][EPC][kTcMaxM]; // [source CTA][slice][owned expert][m] + int const tile = blockIdx.x / CL; + int const crank = blockIdx.x % CL; + int const ks = K / kTcSlices; // 384 + int const kc = SPC * ks; + int const k0 = crank * kc; + int const e0 = tile * kTcExperts; + int const wst = kc + 4; // padded row strides (bank-conflict-free fragment loads) + int const xst = kc + 8; + float* w_s = reinterpret_cast(smem); // [16][wst] + __nv_bfloat16* x_s = reinterpret_cast<__nv_bfloat16*>(smem + kTcExperts * wst * 4); // [M][xst] + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + if (threadIdx.x == 0) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" ::"r"(smem_u32(&w_bar))); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" ::"r"(smem_u32(&x_bar))); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" ::"r"(smem_u32(&mb_bar))); + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smem_u32(&w_bar)), + "r"(static_cast(kTcExperts * kc * 4)) + : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smem_u32(&x_bar)), + "r"(static_cast(M * kc * 2)) + : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smem_u32(&mb_bar)), + "r"(static_cast((CL - 1) * SPC * EPC * M * 4)) + : "memory"); + } + __syncthreads(); + // The peers' mailbox barriers must be initialized before this CTA's st.async lands there. + asm volatile("barrier.cluster.arrive.release.aligned;\n" ::: "memory"); + if (warp == 0 && lane < kTcExperts) + bulk_g2s(w_s + lane * wst, w + static_cast(e0 + lane) * K + k0, kc * 4, &w_bar); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + if (warp == 1) + { + for (int m = lane; m < M; m += 32) + bulk_g2s(x_s + m * xst, x + static_cast(m) * ldx + k0, kc * 2, &x_bar); + } + mbar_wait(&x_bar, 0); + if (q_out != nullptr) + { + float const sc = *sf_scale; + int const kb_cta = kc / 16; + for (int b = threadIdx.x; b < M * kb_cta; b += blockDim.x) + { + int const m = b / kb_cta; + if (m % 8 != tile) + continue; // 8 expert tiles share the row quantization + int const kb = b - m * kb_cta; + int const kbg = k0 / 16 + kb; + quant_block16(x_s + m * xst + kb * 16, sc, + reinterpret_cast(q_out + static_cast(m) * (K / 2) + kbg * 8), + sf_out + static_cast(m) * (K / 16) + kbg); + } + } + mbar_wait(&w_bar, 0); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + // Per slice h: warp w owns k8 steps [w * steps, (w + 1) * steps) of the slice; all token tiles (n = 8 each). + int const g = lane >> 2; + int const t = lane & 3; + int const ntiles = (M + 7) / 8; + int const steps = ks / 8 / 8; +#pragma unroll 1 + for (int h = 0; h < SPC; ++h) + { + float acc[6][4]; +#pragma unroll + for (int i = 0; i < 6; ++i) + acc[i][0] = acc[i][1] = acc[i][2] = acc[i][3] = 0.f; + for (int st = warp * steps; st < (warp + 1) * steps; ++st) + { + int const kk = h * ks + st * 8; + uint32_t const a0 = f2tf32(w_s[g * wst + kk + t]); + uint32_t const a1 = f2tf32(w_s[(g + 8) * wst + kk + t]); + uint32_t const a2 = f2tf32(w_s[g * wst + kk + t + 4]); + uint32_t const a3 = f2tf32(w_s[(g + 8) * wst + kk + t + 4]); +#pragma unroll + for (int nt = 0; nt < 6; ++nt) + { + if (nt < ntiles) + { + int const m = nt * 8 + g; + float bx0 = 0.f, bx1 = 0.f; + if (m < M) + { + bx0 = __bfloat162float(x_s[m * xst + kk + t]); + bx1 = __bfloat162float(x_s[m * xst + kk + t + 4]); + } + uint32_t const b0 = f2tf32(bx0), b1 = f2tf32(bx1); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(acc[nt][0]), "+f"(acc[nt][1]), "+f"(acc[nt][2]), "+f"(acc[nt][3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); + } + } + } +#pragma unroll + for (int nt = 0; nt < 6; ++nt) + { + if (nt < ntiles) + { + int const c = nt * 8 + 2 * t; + red[warp][g][c] = acc[nt][0]; + red[warp][g][c + 1] = acc[nt][1]; + red[warp][g + 8][c] = acc[nt][2]; + red[warp][g + 8][c + 1] = acc[nt][3]; + } + } + __syncthreads(); + for (int i = threadIdx.x; i < kTcExperts * M; i += blockDim.x) + { + int const e = i / M, m = i % M; + float v = 0.f; +#pragma unroll + for (int ww = 0; ww < 8; ++ww) + v += red[ww][e][m]; + part[h][e][m] = v; + } + __syncthreads(); + } + asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory"); + // Expert e of the tile belongs to CTA e / EPC: its slice partials go to that CTA's mailbox[crank]. + for (int i = threadIdx.x; i < SPC * kTcExperts * M; i += blockDim.x) + { + int const h = i / (kTcExperts * M); + int const r = i - h * kTcExperts * M; + int const e = r / M, m = r % M; + int const owner = e / EPC; + float const v = part[h][e][m]; + float* dst = &mailbox[crank][h][e % EPC][m]; + if (owner == crank) + { + *dst = v; + } + else + { + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_u32(dst, owner)), + "r"(__float_as_uint(v)), "r"(mapa_u32(&mb_bar, owner)) + : "memory"); + } + } + __syncthreads(); + { + uint32_t done = 0; + while (!done) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(done) + : "r"(smem_u32(&mb_bar)) + : "memory"); + } + } + for (int i = threadIdx.x; i < EPC * M; i += blockDim.x) + { + int const le = i / M, m = i % M; + float v = 0.f; +#pragma unroll + for (int src = 0; src < CL; ++src) + { +#pragma unroll + for (int h = 0; h < SPC; ++h) + v += mailbox[src][h][le][m]; + } + out[static_cast(m) * (kTcExperts * 8) + e0 + EPC * crank + le] = v; + } +#endif +} + +template +void launchGemv(__nv_bfloat16 const* x, int ldx, int M, int K, float const* w, int E, float* out, uint8_t* q_out, + uint8_t* sf_out, float const* sf_scale, bool pdl, cudaStream_t stream) +{ + size_t const smem = static_cast(ROWS) * K * sizeof(float) + static_cast(M) * K * 2; + auto kernel = m3_router_gemv_kernel; + static bool const configured = [&]() + { + int device = 0, max_optin = 0; + cudaGetDevice(&device); + cudaDeviceGetAttribute(&max_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device); + cudaFuncAttributes attrs = {}; + if (cudaFuncGetAttributes(&attrs, kernel) != cudaSuccess) + { + return false; + } + return cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + max_optin - static_cast(attrs.sharedSizeBytes)) + == cudaSuccess; + }(); + TLLM_CHECK_WITH_INFO(configured, "failed to configure dynamic shared memory"); + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(E / ROWS); + cfg.blockDim = dim3(256); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + cudaError_t const err = cudaLaunchKernelEx(&cfg, kernel, x, ldx, w, out, M, K, E, q_out, sf_out, sf_scale); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +} // namespace + +void invokeMinimaxM3RouterGemv(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t k, float const* w, int32_t e, + float* out, uint8_t* qOut, uint8_t* sfOut, float const* sfScale, bool pdl, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(m >= 1 && m <= kMaxM && k % 1024 == 0, "unsupported router GEMV shape M=%d K=%d", m, k); + launchGemv<1>(x, ldx, m, k, w, e, out, qOut, sfOut, sfScale, pdl, stream); +} + +void invokeMinimaxM3RouterTc(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t k, float const* w, float* out, + uint8_t* qOut, uint8_t* sfOut, float const* sfScale, bool pdl, cudaStream_t stream) +{ + int64_t const M = m, K = k; + TLLM_CHECK_WITH_INFO( + M >= 1 && M <= kTcMaxM && K % (kTcSlices * 64) == 0, "unsupported router shape M=%d K=%d", m, k); + int const ks = K / kTcSlices; + auto launch_tc = [&](auto kernel, int cl) -> cudaError_t + { + int const kc = (kTcSlices / cl) * ks; + size_t const smem = static_cast(kTcExperts) * (kc + 4) * 4 + static_cast(M) * (kc + 8) * 2; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(8 * cl); + cfg.blockDim = dim3(256); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + attr[1].id = cudaLaunchAttributeClusterDimension; + attr[1].val.clusterDim.x = cl; + attr[1].val.clusterDim.y = 1; + attr[1].val.clusterDim.z = 1; + cfg.attrs = attr; + cfg.numAttrs = 2; + return cudaLaunchKernelEx( + &cfg, kernel, x, ldx, w, out, static_cast(M), static_cast(K), qOut, sfOut, sfScale); + }; + // 16-CTA clusters (128 CTAs, one K slice each) when the device co-schedules all 8, else 8-CTA clusters. + static int const cluster = [&]() + { + int device = 0, max_optin = 0; + cudaGetDevice(&device); + cudaDeviceGetAttribute(&max_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device); + auto configure = [&](auto kernel, bool non_portable) + { + cudaFuncAttributes attrs = {}; + if (cudaFuncGetAttributes(&attrs, kernel) != cudaSuccess) + return false; + if (non_portable + && cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) != cudaSuccess) + return false; + return cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + max_optin - static_cast(attrs.sharedSizeBytes)) + == cudaSuccess; + }; + TLLM_CHECK_WITH_INFO(configure(m3_router_tc_kernel<8>, false), "failed to configure the router kernel"); + if (!configure(m3_router_tc_kernel<16>, true)) + { + cudaGetLastError(); + return 8; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(8 * 16); + cfg.blockDim = dim3(256); + cfg.dynamicSmemBytes + = static_cast(kTcExperts) * (ks + 4) * 4 + static_cast(kTcMaxM) * (ks + 8) * 2; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeClusterDimension; + attr[0].val.clusterDim.x = 16; + attr[0].val.clusterDim.y = 1; + attr[0].val.clusterDim.z = 1; + cfg.attrs = attr; + cfg.numAttrs = 1; + int clusters = 0; + if (cudaOccupancyMaxActiveClusters(&clusters, m3_router_tc_kernel<16>, &cfg) != cudaSuccess) + { + cudaGetLastError(); + return 8; + } + return clusters >= 8 ? 16 : 8; + }(); + cudaError_t const err + = cluster == 16 ? launch_tc(m3_router_tc_kernel<16>, 16) : launch_tc(m3_router_tc_kernel<8>, 8); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.h b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.h new file mode 100644 index 000000000000..6f4467883b24 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.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 +{ + +//! Low-M fp32 router GEMV: `out[M, E] = float(x[M, K]) @ w[E, K]^T` with exact fp32 FMA, one CTA per expert, +//! 1 <= M <= 16, K % 1024 == 0. With `qOut` the CTAs also emit the NVFP4 activation (packed e2m1 [M, K / 2], linear +//! E4M3 scales [M, K / 16]) for global scale `*sfScale`, bit-identical to `fp4_quantize(x, sfScale, 16, False, +//! False)`. +void invokeMinimaxM3RouterGemv(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t k, float const* w, int32_t e, + float* out, uint8_t* qOut, uint8_t* sfOut, float const* sfScale, bool pdl, cudaStream_t stream); + +//! TF32 tensor-core router for 128 experts and 1 <= M <= 48 (K % 1024 == 0): `out[M, 128]` with the weights rounded +//! to TF32 (round to nearest, ties away) and fp32 accumulation, independent of the cluster size picked. Optional +//! NVFP4 activation as in `invokeMinimaxM3RouterGemv`. +void invokeMinimaxM3RouterTc(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t k, float const* w, float* out, + uint8_t* qOut, uint8_t* sfOut, float const* sfScale, bool pdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.cu b/cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.cu new file mode 100644 index 000000000000..91b4c38d0fda --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.cu @@ -0,0 +1,525 @@ +/* + * 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/minimaxM3/sampledAccept.h" + +#include "tensorrt_llm/common/assert.h" + +#include +#include + +// Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). +// +// Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the +// strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p +// active must stay on the generic path; the caller decides. +// +// Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL +// vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z +// (the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it +// is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens +// heavier than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the +// nucleus with probability >= top_p, so most rows decide on their second exchange. +// +// Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, +// staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local +// candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the +// candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based +// Philox keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one +// 4-word record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and +// reads the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the +// same order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing +// and two call halves keep a rank one call ahead from overwriting a slot still being read. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3 +{ +namespace sampledAccept +{ +namespace +{ +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / 32; +constexpr int kMaxRec = kMaxWorld * kMaxSplits; +constexpr int kMaxAttempts = 16; // exchange slots per row, reused cyclically by the rounds (no round cap) +constexpr int kEx = 4; // words per record: key, logit of key, f2, f3 +constexpr int kTagBits = 64 - 32 - kIdBits; +constexpr float kTempEpsilon = 1e-6f; // the fused sampling kernel's 1/T guard + +// Multicast exchange, per call half: [kMaxAttempts][kMaxRows][world * kMaxSplits][kEx] u64. +__host__ __device__ inline long long half_words(int world) +{ + return static_cast(kMaxAttempts) * kMaxRows * world * kMaxSplits * kEx; +} + +__device__ __forceinline__ uint32_t ordered(float v) +{ + uint32_t const b = __float_as_uint(v); + return (b & 0x80000000u) ? ~b : (b | 0x80000000u); +} + +// Philox4x32-10. +__device__ __forceinline__ uint4 philox(uint4 c, uint2 k) +{ +#pragma unroll + for (int i = 0; i < 10; ++i) + { + uint32_t const hi0 = __umulhi(0xD2511F53u, c.x), lo0 = 0xD2511F53u * c.x; + uint32_t const hi1 = __umulhi(0xCD9E8D57u, c.z), lo1 = 0xCD9E8D57u * c.z; + c = make_uint4(hi1 ^ c.y ^ k.x, lo1, hi0 ^ c.w ^ k.y, lo0); + k.x += 0x9E3779B9u; + k.y += 0xBB67AE85u; + } + return c; +} + +// u in (0, 1): (k + 1/2) / 2^23, exact in fp32. +__device__ __forceinline__ float uniform(uint32_t bits) +{ + return (static_cast(bits >> 9) + 0.5f) * (1.0f / 8388608.0f); +} + +// Key on the wire: perturbed log-mass, then the candidate token id; 0 = no candidate. Tag bits are added on the wire. +__device__ __forceinline__ unsigned long long gkey(float lg, int id) +{ + return (static_cast(ordered(lg)) << kIdBits) + | static_cast((1u << kIdBits) - 1u - static_cast(id)); +} + +__device__ __forceinline__ int key_token(unsigned long long k) +{ + return static_cast((1u << kIdBits) - 1u - static_cast(k & ((1u << kIdBits) - 1u))); +} + +__device__ __forceinline__ void st_sys(unsigned long long* p, unsigned long long v) +{ + asm volatile("st.global.relaxed.sys.u64 [%0], %1;" ::"l"(p), "l"(v) : "memory"); +} + +__device__ __forceinline__ unsigned long long ld_sys(unsigned long long const* p) +{ + unsigned long long v; + asm volatile("ld.global.relaxed.sys.u64 %0, [%1];" : "=l"(v) : "l"(p) : "memory"); + return v; +} + +template +struct Shared +{ + __nv_bfloat16 row[kMaxVs / kSplits]; + float redf[kWarps]; + unsigned long long redk[kWarps]; + unsigned long long rk[kMaxRec]; + float rv[kMaxRec], r2[kMaxRec], r3[kMaxRec]; + float wsum[kWarps]; + int owner; + int pick; + float pickx; + bool last; +}; + +template +__device__ __forceinline__ float block_max(float v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o)); + if ((threadIdx.x & 31) == 0) + s.redf[threadIdx.x >> 5] = v; + __syncthreads(); + float r = s.redf[0]; + for (int i = 1; i < kWarps; ++i) + r = fmaxf(r, s.redf[i]); + __syncthreads(); + return r; +} + +// Fixed-order block sum: identical result for identical inputs. +template +__device__ __forceinline__ float block_sum(float v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + if ((threadIdx.x & 31) == 0) + s.redf[threadIdx.x >> 5] = v; + __syncthreads(); + float r = 0.f; + for (int i = 0; i < kWarps; ++i) + r += s.redf[i]; + __syncthreads(); + return r; +} + +template +__device__ __forceinline__ unsigned long long block_max_u64(unsigned long long v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + unsigned long long const w = __shfl_xor_sync(0xffffffffu, v, o); + v = w > v ? w : v; + } + if ((threadIdx.x & 31) == 0) + s.redk[threadIdx.x >> 5] = v; + __syncthreads(); + unsigned long long r = s.redk[0]; + for (int i = 1; i < kWarps; ++i) + r = s.redk[i] > r ? s.redk[i] : r; + __syncthreads(); + return r; +} + +// Exclusive prefix of v over the threads in thread order, and the total. +template +__device__ __forceinline__ float block_excl_scan(float v, float& total, S& s) +{ + int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + float inc = v; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) + { + float const t = __shfl_up_sync(0xffffffffu, inc, o); + if (lane >= o) + inc += t; + } + if (lane == 31) + s.wsum[warp] = inc; + __syncthreads(); + float base = 0.f; + total = 0.f; + for (int i = 0; i < kWarps; ++i) + { + base += i < warp ? s.wsum[i] : 0.f; + total += s.wsum[i]; + } + __syncthreads(); + return base + inc - v; +} + +// Every CTA of the grid must be co-resident: a row's CTAs wait for each other's records. +template +__global__ void __launch_bounds__(kThreads, 2) sampled_accept_kernel(Args const a) +{ + extern __shared__ __align__(16) unsigned char smem[]; + Shared& s = *reinterpret_cast*>(smem); + float constexpr kNegInf = -__builtin_huge_valf(); + int const tid = threadIdx.x; + int const row = blockIdx.x / kSplits; + int const split = blockIdx.x % kSplits; + int const world = a.world; + int const nrec = world * kSplits; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + unsigned int const call = *reinterpret_cast(a.counters + 1); + // Wire tags: never 0 (fresh memory), different between calls of the same half and between consecutive laps of + // the slot ring (round r uses slot r % kMaxAttempts and lap parity (r / kMaxAttempts) & 1). A rank decides round + // r only after reading every record of round r, so no rank is more than one round ahead and a slot is rewritten + // only after every reader has consumed it; the lap bit tells the new record from the previous lap's. + auto ktag_of = [&](int r) -> unsigned long long + { + unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; + return (((call >> 1) << 1) | lap) % ((1u << kTagBits) - 1u) + 1u; + }; + auto ftag_of = [&](int r) -> unsigned long long + { + unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; + return static_cast(((((call >> 1) << 1) | lap) & 0x7FFFFFFFu) | 0x80000000u) << 32; + }; + long long const rec_stride = static_cast(world) * kMaxSplits * kEx; + unsigned long long* const xu + = a.ws_uc + static_cast(call & 1u) * half_words(world) + static_cast(row) * rec_stride; + unsigned long long* const xm = a.ws_mc + static_cast(call & 1u) * half_words(world) + + static_cast(row) * rec_stride + (a.rank * kMaxSplits + split) * kEx; + long long const round_stride = static_cast(kMaxRows) * rec_stride; + + float const tempInv = 1.0f / (a.temperatures[row] + kTempEpsilon); + float const topP = a.top_ps[row]; + unsigned long long const seed = static_cast(a.seeds[row]); + unsigned long long const offset = static_cast(a.offsets[row]); + + // ---- stage this CTA's slice and its max ------------------------------------------------------------------------- + __nv_bfloat16 const* src = a.logits + row * a.ld; + int const sl_len = ((a.Vs + kSplits - 1) / kSplits + 7) / 8 * 8; + int const s0 = min(split * sl_len, a.Vs); + int const n = min(s0 + sl_len, a.Vs) - s0; + int const id0 = a.vocab_offset + s0; + float tmax = kNegInf; + bool const vec = (a.ld % 8) == 0 && (reinterpret_cast(a.logits) & 15) == 0; + int const nvec = vec ? n / 8 : 0; + for (int v = tid; v < nvec; v += kThreads) + { + uint4 const w = *reinterpret_cast(src + s0 + 8 * v); + *reinterpret_cast(&s.row[8 * v]) = w; + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&w); +#pragma unroll + for (int j = 0; j < 8; ++j) + tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(e[j]), tempInv)); + } + for (int i = nvec * 8 + tid; i < n; i += kThreads) + { + s.row[i] = src[s0 + i]; + tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(src[s0 + i]), tempInv)); + } + float const smax = block_max(tmax, s); // the staged slice is visible to all threads after this barrier + float M = 0.f, wc = 0.f, target = 0.f, cv = 0.f; // global max, candidate weight and logit, top_p * Z + unsigned long long ckey = 0ULL; + int const gid = a.rank * kMaxSplits + split; + for (int r = 0;; ++r) + { + // ---- local candidate by inverse CDF over the slice's tokens heavier than the candidate (w > 0 in round 0) + float const ref = r == 0 ? smax : M; + float const wmin = r == 0 ? 0.f : wc; + float ts = 0.f; + for (int i = tid; i < n; i += kThreads) + { + float const w = __expf(__fmul_rn(__bfloat162float(s.row[i]), tempInv) - ref); + ts += w > wmin ? w : 0.f; + } + float total; + float const excl = block_excl_scan(ts, total, s); + uint4 const rnd = philox(make_uint4(static_cast(gid), static_cast(row + 64 * r), + static_cast(offset), static_cast(offset >> 32)), + make_uint2(static_cast(seed), static_cast(seed >> 32))); + float const tgt = uniform(rnd.x) * total; + if (tid == 0) + s.owner = -1; + __syncthreads(); + if (ts > 0.f && excl <= tgt) + atomicMax(&s.owner, tid); + __syncthreads(); + if (tid == s.owner) + { + float acc = excl; + int pick = -1; + float px = 0.f; + for (int i = tid; i < n; i += kThreads) + { + float const x = __fmul_rn(__bfloat162float(s.row[i]), tempInv); + float const w = __expf(x - ref); + if (w > wmin) + { + acc += w; + pick = i; + px = x; + if (acc > tgt) + break; + } + } + s.pick = pick; + s.pickx = px; + } + __syncthreads(); + unsigned long long sb = 0ULL; + if (total > 0.f && s.owner >= 0) + sb = gkey(__logf(total) + ref - logf(-logf(uniform(rnd.y))), id0 + s.pick); + float const smab = total; + unsigned long long const ktag = ktag_of(r); + unsigned long long const ftag = ftag_of(r); + unsigned long long* const xr = xu + (r % kMaxAttempts) * round_stride; + if (tid == 0) + { + unsigned long long* const xw = xm + (r % kMaxAttempts) * round_stride; + st_sys(xw + 0, (sb << kTagBits) | ktag); + st_sys(xw + 1, ftag | __float_as_uint(sb != 0ULL ? s.pickx : 0.f)); + st_sys(xw + 2, ftag | __float_as_uint(r == 0 ? smax : smab)); + st_sys(xw + 3, ftag | __float_as_uint(r == 0 ? smab : 0.f)); + } + // ---- every record of the row, every rank --------------------------------------------------------------- + if (tid < nrec * kEx) + { + int const q = tid / kEx, f = tid % kEx; + unsigned long long v; + if (f == 0) + { + do + v = ld_sys(xr + (q / kSplits * kMaxSplits + q % kSplits) * kEx); + while ((v & ((1ULL << kTagBits) - 1ULL)) != ktag); + s.rk[q] = v >> kTagBits; + } + else + { + do + v = ld_sys(xr + (q / kSplits * kMaxSplits + q % kSplits) * kEx + f); + while ((v & 0xFFFFFFFF00000000ULL) != ftag); + float const fv = __uint_as_float(static_cast(v)); + (f == 1 ? s.rv : f == 2 ? s.r2 : s.r3)[q] = fv; + } + } + __syncthreads(); + // ---- the decision, identical in every CTA of the row on every rank: fixed record order ----------------- + unsigned long long gk = 0ULL; + float gv = 0.f; + for (int q = 0; q < nrec; ++q) + { + if (s.rk[q] > gk) + { + gk = s.rk[q]; + gv = s.rv[q]; + } + } + bool accept; + if (r == 0) + { + M = kNegInf; + for (int q = 0; q < nrec; ++q) + M = fmaxf(M, s.r2[q]); + float Z = 0.f; + for (int q = 0; q < nrec; ++q) + Z += s.r2[q] == kNegInf ? 0.f : s.r3[q] * __expf(s.r2[q] - M); + target = topP * Z; + accept = !(topP < 1.0f); // no nucleus cut: the first draw stands + ckey = gk; + cv = gv; + } + else + { + float S = 0.f; // mass strictly above the current candidate + for (int q = 0; q < nrec; ++q) + S += s.r2[q]; + accept = S <= target || gk == 0ULL; + if (!accept) + { + ckey = gk; + cv = gv; + } + } + wc = __expf(cv - M); + __syncthreads(); // records consumed before the next round's writes to s + if (accept) + { + if (split != 0) + return; + if (tid == 0) + { + a.tokens[row] = key_token(ckey); + if (a.extra_attempts != nullptr && r > 1) + atomicAdd(a.extra_attempts, r - 1); + } + break; + } + } + + // ---- split 0: the last row to finish runs the strict acceptance ---------------------------------------------- +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + __threadfence(); + __syncthreads(); + if (tid == 0) + s.last = atomicAdd(a.counters, 1u) == static_cast(a.rows) - 1; + __syncthreads(); + if (!s.last) + return; + __threadfence(); + if (tid == 0) + { + a.counters[0] = 0u; + a.counters[1] = call + 1u; + } + if (tid < a.nreq) + { + int const* tok = a.tokens; + int const b = tid; + int const k1 = a.draft_len + 1; + int nacc = 1; + bool run = true; + for (int k = 0; k < a.draft_len; ++k) + { + int const tt = *reinterpret_cast(tok + b * k1 + k); + a.acc_tokens[b * k1 + k] = tt; + run = run && a.draft[b * a.draft_len + k] == tt; + nacc += run ? 1 : 0; + } + a.acc_tokens[b * k1 + a.draft_len] = *reinterpret_cast(tok + b * k1 + a.draft_len); + if (a.force_vec != nullptr) + nacc = a.force_vec[b]; + else if (a.force_total > 0) + nacc = a.force_total; + a.num_acc[b] = nacc; + a.gather_ids[b] = static_cast(a.batch_idx[b]) * k1 + nacc - 1; + if (a.next_new != nullptr) + a.next_new[b * a.next_stride] + = *reinterpret_cast(tok + a.batch_idx[b] * k1 + nacc - 1); + } +} + +template +cudaError_t launch(Args const& a, bool use_pdl, cudaStream_t stream) +{ + static int max_ctas = 0; + if (max_ctas == 0) + { + int dev = 0, sms = 0, per_sm = 0; + cudaError_t err = cudaGetDevice(&dev); + if (err == cudaSuccess) + { + err = cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); + } + if (err == cudaSuccess) + { + err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &per_sm, sampled_accept_kernel, kThreads, sizeof(Shared)); + } + if (err != cudaSuccess) + { + return err; + } + max_ctas = per_sm * sms; + } + TLLM_CHECK_WITH_INFO(a.rows * kSplits <= max_ctas, "m3_sampled_accept: %d CTAs cannot all be co-resident (max %d)", + a.rows * kSplits, max_ctas); + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(a.rows * kSplits); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = sizeof(Shared); + cfg.stream = stream; + cudaLaunchAttribute attr; + attr.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr.val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; + cfg.attrs = &attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, sampled_accept_kernel, a); +} + +} // namespace + +long long sampledAcceptWsWords(int world) +{ + return 2 * half_words(world); +} + +cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream) +{ + // TP4 caps 2 / 4 / 8 (8 / 16 / 32 rows) and TP8 cap 1 (4 rows): 8 CTAs per row, at most 256 CTAs (2 per SM); + // more rows: 4 CTAs per row. + if (a.rows <= 32) + { + return launch<8>(a, usePdl, stream); + } + return launch<4>(a, usePdl, stream); +} + +} // namespace sampledAccept +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.h b/cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.h new file mode 100644 index 000000000000..1dc211a6f8a8 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.h @@ -0,0 +1,79 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3 +{ +//! Fused sampled (temperature + top-p) verification of the MiniMax-M3 Eagle3 target head over the TP-sharded +//! vocabulary, decode rows <= kMaxRows; see sampledAccept.cu. +namespace sampledAccept +{ + +constexpr int kMaxSplits = 8; // CTAs per verify row (at most) +constexpr int kMaxRows = 48; // verify rows per call +constexpr int kMaxWorld = 8; // TP ranks +constexpr int kMaxVs = 51200; // vocab shard width +constexpr int kIdBits = 18; // token ids < 2^18 + +struct Args +{ + __nv_bfloat16 const* logits; // [rows, Vs] this rank's shard + long long ld; + int rows, Vs, vocab_offset; + float const* temperatures; + float const* top_ps; + long long const* seeds; + long long const* offsets; + unsigned long long* ws_uc; + unsigned long long* ws_mc; + unsigned int* counters; // [0] finished rows, [1] call + int* tokens; // [rows] sampled token per row + int* extra_attempts; // optional [1]: rounds beyond the second, summed over rows + int rank, world; + int const* draft; + int const* batch_idx; + int* acc_tokens; + int* num_acc; + long long* gather_ids; + int draft_len, nreq, force_total; + int* next_new; + int next_stride; + int const* force_vec; +}; + +//! Words (uint64) of the multicast exchange workspace for `world` ranks: two call halves. +long long sampledAcceptWsWords(int world); + +//! Launch the sampled verification over a.rows * (8 or 4) co-resident CTAs (8 per row up to 32 rows). +cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream); + +} // namespace sampledAccept +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/selectLong.cu b/cpp/tensorrt_llm/kernels/minimaxM3/selectLong.cu new file mode 100644 index 000000000000..39351f0ad3ce --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/selectLong.cu @@ -0,0 +1,235 @@ +/* + * 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. + */ + +// Top-16 block select for MiniMax-M3 decode rows of 129 to 9216 blocks. +// +// One 1024-thread CTA per (query, KV head) row holds the row's ordered score keys in registers (2, 4 or 9 per +// thread, by the row's valid count). The weakest of the 16 maxima of the interleaved groups `block % 16` is a lower +// bound on the 16th-largest key, so the top 16 is among the keys at or above it; those are compacted to shared +// memory and ranked exactly by (key, block). More than 1024 keys tied at the bound (degenerate rows) rank every +// candidate against the whole row (re-read). + +#include "tensorrt_llm/kernels/minimaxM3/selectLong.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr float kInitScore = 1.0e30f; +constexpr float kLocalScore = 1.0e29f; +constexpr uint32_t kNegInfKey = 0xff800000u; +constexpr int kItems = 9; // keys per thread: rows up to 9216 blocks (1M tokens + draft tokens) + +// Descending order key: a smaller key is a larger score (NaN above +inf), as in the reference op. +__device__ __forceinline__ uint32_t orderedKey(float x) +{ + uint32_t const b = __float_as_uint(x); + return (b & 0x80000000u) ? b : (~b & 0x7fffffffu); +} + +__device__ __forceinline__ float applyForcing(float raw, int block, int initBlocks, long long localStart) +{ + if (block >= localStart) + return kLocalScore; + if (block < initBlocks) + return kInitScore; + return raw; +} + +struct SelectSmem +{ + uint32_t groupMin[16]; + unsigned long long cand[1024]; + int numCand; + unsigned long long picked[16]; +}; + +// One row with up to kItems * 1024 valid blocks; the caller picks kItems from the row's valid count. +template +__device__ __forceinline__ void selectRow( + SelectSmem& sm, float const* rowBase, long long blockStride, int valid, int initBlocks, long long localStart) +{ + int const tid = threadIdx.x, lane = tid & 31; + uint32_t key[kItems]; + uint32_t localMin = 0xffffffffu; +#pragma unroll + for (int s = 0; s < kItems; ++s) + { + int const block = s * 1024 + tid; + key[s] = block < valid ? orderedKey(applyForcing(rowBase[block * blockStride], block, initBlocks, localStart)) + : 0xffffffffu; + localMin = min(localMin, key[s]); + } + if (valid <= 16) + { + if (tid < valid) + sm.picked[tid] = (static_cast(key[0]) << 32) | static_cast(tid); + return; + } + // Group g = block % 16 = tid % 16: lanes l and l ^ 16 of every warp. + localMin = min(localMin, __shfl_xor_sync(0xffffffffu, localMin, 16)); + if (lane < 16) + atomicMin(&sm.groupMin[lane], localMin); + __syncthreads(); + uint32_t bound = 0; +#pragma unroll + for (int g = 0; g < 16; ++g) + bound = max(bound, sm.groupMin[g]); +#pragma unroll + for (int s = 0; s < kItems; ++s) + { + int const block = s * 1024 + tid; + bool const isCand = block < valid && key[s] <= bound; + unsigned const ballot = __ballot_sync(0xffffffffu, isCand); + int base = 0; + if (lane == 0 && ballot) + base = atomicAdd(&sm.numCand, __popc(ballot)); + base = __shfl_sync(0xffffffffu, base, 0); + if (isCand) + { + int const i = base + __popc(ballot & ((1u << lane) - 1)); + if (i < 1024) + sm.cand[i] = (static_cast(key[s]) << 32) | static_cast(block); + } + } + __syncthreads(); + int const n = sm.numCand; + if (n <= 1024) + { + for (int i = tid; i < n; i += 1024) + { + unsigned long long const mine = sm.cand[i]; + int rank = 0; + for (int j = 0; j < n; ++j) + rank += sm.cand[j] < mine; + if (rank < 16) + sm.picked[rank] = mine; + } + return; + } + // More than 1024 keys tied at the bound (degenerate rows): rank every candidate against the whole row. +#pragma unroll + for (int s = 0; s < kItems; ++s) + { + int const block = s * 1024 + tid; + if (block < valid && key[s] <= bound) + { + int rank = 0; + for (int j = 0; j < valid; ++j) + { + uint32_t const other = orderedKey(applyForcing(rowBase[j * blockStride], j, initBlocks, localStart)); + rank += other < key[s] || (other == key[s] && j < block); + } + if (rank < 16) + sm.picked[rank] = (static_cast(key[s]) << 32) | static_cast(block); + } + } +} + +__global__ void __launch_bounds__(1024) + m3_select_long_kernel(float const* __restrict__ scores, long long headStride, long long blockStride, + long long queryStride, int const* __restrict__ nValid, int* __restrict__ out, long long outQueryStride, + long long outHeadStride, int numKvHeads, int numBlocks, int initBlocks, int localBlocks, bool earlyTrigger) +{ + __shared__ SelectSmem sm; + int const row = blockIdx.x, query = row / numKvHeads, head = row % numKvHeads; + int const tid = threadIdx.x, lane = tid & 31; + if (tid < 16) + { + sm.groupMin[tid] = 0xffffffffu; + sm.picked[tid] = ~0ull; + } + if (tid == 0) + sm.numCand = 0; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaGridDependencySynchronize(); + // The sparse decode that follows stages its table-independent inputs before its own grid wait; past this grid's + // wait those inputs are final, so it may launch now rather than at this grid's exit. + if (earlyTrigger) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + int const raw = nValid[query]; + int const valid = max(0, min(raw, numBlocks)); + long long const localStart = max(static_cast(raw) - localBlocks, 0LL); + float const* rowBase = scores + head * headStride + query * queryStride; + __syncthreads(); + if (valid <= 2048) + selectRow<2>(sm, rowBase, blockStride, valid, initBlocks, localStart); + else if (valid <= 4096) + selectRow<4>(sm, rowBase, blockStride, valid, initBlocks, localStart); + else + selectRow(sm, rowBase, blockStride, valid, initBlocks, localStart); + __syncthreads(); + if (tid >= 32) + return; + uint32_t blockKey = ~0u; + if (lane < 16) + { + unsigned long long const c = sm.picked[lane]; + if (c != ~0ull && static_cast(c >> 32) != kNegInfKey) + blockKey = static_cast(c & 0xffffffffu); + } +#pragma unroll + for (int size = 2; size <= 32; size <<= 1) + { +#pragma unroll + for (int stride = size >> 1; stride > 0; stride >>= 1) + { + uint32_t const other = __shfl_xor_sync(0xffffffffu, blockKey, stride); + bool const ascending = (lane & size) == 0, lower = (lane & stride) == 0; + blockKey = (lower == ascending) ? min(blockKey, other) : max(blockKey, other); + } + } + if (lane < 16) + out[query * outQueryStride + head * outHeadStride + lane] = blockKey == ~0u ? -1 : static_cast(blockKey); +} + +} // namespace + +cudaError_t invokeMinimaxM3SelectLong(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride, + int32_t const* nValidBlocks, int32_t* out, int64_t outQueryStride, int64_t outHeadStride, int32_t numKvHeads, + int32_t numBlocks, int32_t totalQ, int32_t initBlocks, int32_t localBlocks, bool usePdl, bool earlyTrigger, + cudaStream_t stream) +{ + if (numBlocks > kMinimaxM3SelectLongMaxBlocks) + { + return cudaErrorInvalidValue; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(totalQ * numKvHeads); + cfg.blockDim = dim3(1024); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = usePdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, m3_select_long_kernel, scores, static_cast(headStride), + static_cast(blockStride), static_cast(queryStride), nValidBlocks, out, + static_cast(outQueryStride), static_cast(outHeadStride), numKvHeads, numBlocks, + initBlocks, localBlocks, earlyTrigger); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/selectLong.h b/cpp/tensorrt_llm/kernels/minimaxM3/selectLong.h new file mode 100644 index 000000000000..258954235317 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/selectLong.h @@ -0,0 +1,45 @@ +/* + * 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 +{ + +//! Largest row (blocks per query) the long-row top-16 select covers: 9 keys per thread of one 1024-thread CTA. +constexpr int32_t kMinimaxM3SelectLongMaxBlocks = 9216; + +//! MiniMax-M3 top-16 block select over per-KV-head scores [numKvHeads, numBlocks, totalQ] (element strides given), +//! one 1024-thread CTA per (query, KV head). Same contract as minimax_m3_select_blocks: init / local forcing, the +//! per-query valid-block count, ties to the smaller block, forced -inf picks emitted as -1, ascending ids with -1 +//! tail padding. Writes out[query * outQueryStride + head * outHeadStride + 0..15]. With earlyTrigger the dependent +//! grid is released right after this grid's dependency wait instead of at its exit. +cudaError_t invokeMinimaxM3SelectLong(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride, + int32_t const* nValidBlocks, int32_t* out, int64_t outQueryStride, int64_t outHeadStride, int32_t numKvHeads, + int32_t numBlocks, int32_t totalQ, int32_t initBlocks, int32_t localBlocks, bool usePdl, bool earlyTrigger, + cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/shardAccept.cu b/cpp/tensorrt_llm/kernels/minimaxM3/shardAccept.cu new file mode 100644 index 000000000000..afc70e0e2c35 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/shardAccept.cu @@ -0,0 +1,204 @@ +/* + * 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/minimaxM3/shardAccept.h" + +#include +#include +#include + +// Fused greedy Eagle3 verification for the MiniMax-M3 target head (GB300, TP-sharded vocabulary, <= 48 verify rows). +// +// One kernel replaces the all-gather of the target logits, the argmax and the strict-acceptance ops: +// +// * Each logit of this rank's BF16 shard becomes a 64-bit key ordered(value) << 32 | ~global_index (ties pick the +// lowest index, as torch.argmax), max-reduced per row: warp shuffles, then atomicMax. +// * The last CTA to finish pushes the rank's per-row key into every rank's exchange slot through a multicast mapping +// (Lamport: 0 = empty, the key is never 0), polls its own slots for all ranks, takes the global argmax, re-arms the +// slots and runs the strict acceptance. Call n uses slot half n % 2. +// +// The keys order the BF16 values exactly, so the tokens are the argmax of the gathered logits. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3 +{ +namespace shardAccept +{ +namespace +{ +__device__ __forceinline__ unsigned long long make_key(float v, uint32_t index) +{ + uint32_t u = __float_as_uint(v); + // -0.0 and +0.0 compare equal (torch.argmax takes the lower index among them): one key for both. + u = (u == 0x80000000u) ? 0u : u; + u = (u & 0x80000000u) ? ~u : (u | 0x80000000u); + return (static_cast(u) << 32) | static_cast(0xFFFFFFFFu - index); +} + +__device__ __forceinline__ unsigned long long shfl_max(unsigned long long k, int o) +{ + unsigned long long const ok = __shfl_xor_sync(0xffffffffu, k, o); + return ok > k ? ok : k; +} + +// Greedy verification from this rank's logits shard (verify-row count <= kMaxRows): per-row argmax keys over the +// shard, the last CTA exchanges them across ranks (multicast slots [2][kMaxRows][world], Lamport 0 = empty) and runs +// the acceptance epilogue. +__global__ void __launch_bounds__(kThreads, 1) shard_accept_kernel(__nv_bfloat16 const* __restrict__ logits, int M, + int Vs, long long ld, int vocab_offset, unsigned long long* __restrict__ keys, unsigned int* __restrict__ counters, + unsigned long long* __restrict__ ws_uc, unsigned long long* __restrict__ ws_mc, int rank, int world, AcceptArgs acc) +{ + __shared__ bool s_last; + __shared__ int s_tok[kMaxRows]; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + int const nwarps = blockDim.x >> 5; + int const groups = (Vs + 7) / 8; + int const g0 = static_cast((static_cast(blockIdx.x) * groups) / gridDim.x); + int const g1 = static_cast((static_cast(blockIdx.x + 1) * groups) / gridDim.x); + bool const vec = (ld % 8) == 0; + for (int r = warp; r < M; r += nwarps) + { + unsigned long long best = 0ull; + __nv_bfloat16 const* row = logits + r * ld; + for (int gi = g0 + lane; gi < g1; gi += 32) + { + int const c = gi * 8; + if (vec && c + 8 <= Vs) + { + uint4 const raw = *reinterpret_cast(row + c); + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); +#pragma unroll + for (int q = 0; q < 8; ++q) + { + unsigned long long const k + = make_key(__bfloat162float(e[q]), static_cast(vocab_offset + c + q)); + best = k > best ? k : best; + } + } + else + { + for (int q = 0; q < 8 && c + q < Vs; ++q) + { + unsigned long long const k + = make_key(__bfloat162float(row[c + q]), static_cast(vocab_offset + c + q)); + best = k > best ? k : best; + } + } + } +#pragma unroll + for (int o = 1; o < 32; o <<= 1) + best = shfl_max(best, o); + if (lane == 0 && best != 0ull) + atomicMax(keys + r, best); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + unsigned int const prev = atomicAdd(counters, 1u); + s_last = prev == gridDim.x - 1; + } + __syncthreads(); + if (!s_last) + return; + __threadfence(); + unsigned int const call = *reinterpret_cast(counters + 1); + int const half = static_cast(call & 1u); + for (int t = threadIdx.x; t < M; t += blockDim.x) + { + unsigned long long const mine = atomicExch(keys + t, 0ull); + unsigned long long* slot_mc = ws_mc + (half * kMaxRows + t) * world; + unsigned long long* slot_uc = ws_uc + (half * kMaxRows + t) * world; + slot_mc[rank] = mine; + unsigned long long best = 0ull; + for (int r = 0; r < world; ++r) + { + unsigned long long v; + do + { + v = *reinterpret_cast(slot_uc + r); + } while (v == 0ull); + best = v > best ? v : best; + slot_uc[r] = 0ull; + } + s_tok[t] = static_cast(0xFFFFFFFFu - static_cast(best & 0xFFFFFFFFull)); + } + __syncthreads(); + if (threadIdx.x == 0) + { + counters[0] = 0u; + counters[1] = call + 1u; + } + if (threadIdx.x < acc.nreq) + { + int const b = threadIdx.x; + int const k1 = acc.draft_len + 1; + int n = 1; + bool run = true; + for (int k = 0; k < acc.draft_len; ++k) + { + int const tt = s_tok[b * k1 + k]; + acc.acc_tokens[b * k1 + k] = tt; + run = run && acc.draft[b * acc.draft_len + k] == tt; + n += run ? 1 : 0; + } + acc.acc_tokens[b * k1 + acc.draft_len] = s_tok[b * k1 + acc.draft_len]; + if (acc.force_vec != nullptr) + n = acc.force_vec[b]; + else if (acc.force_total > 0) + n = acc.force_total; + acc.num_acc[b] = n; + acc.gather_ids[b] = static_cast(acc.batch_idx[b]) * k1 + n - 1; + if (acc.next_new != nullptr) + acc.next_new[b * acc.next_stride] = s_tok[acc.batch_idx[b] * k1 + n - 1]; + } +} + +} // namespace + +cudaError_t invokeShardAccept(__nv_bfloat16 const* logits, int M, int Vs, long long ld, int vocabOffset, + unsigned long long* keys, unsigned int* counters, unsigned long long* wsUc, unsigned long long* wsMc, int rank, + int world, AcceptArgs const& acc, int numCtas, cudaStream_t stream) +{ + cudaLaunchAttribute attr[1]; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(static_cast(numCtas)); + cfg.blockDim = dim3(kThreads); + cfg.stream = stream; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = 1; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx( + &cfg, shard_accept_kernel, logits, M, Vs, ld, vocabOffset, keys, counters, wsUc, wsMc, rank, world, acc); +} + +} // namespace shardAccept +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/shardAccept.h b/cpp/tensorrt_llm/kernels/minimaxM3/shardAccept.h new file mode 100644 index 000000000000..e41982b9d357 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/shardAccept.h @@ -0,0 +1,67 @@ +/* + * 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 +{ +namespace minimaxM3 +{ +//! Fused greedy Eagle3 verification of the MiniMax-M3 NVFP4 TP4 decode path (sm_103a): the cross-rank argmax of the +//! target logits shard and the strict acceptance in one kernel. See shardAccept.cu. +namespace shardAccept +{ + +constexpr int kMaxRows = 48; // verify rows of one launch +constexpr int kThreads = 512; // threads per CTA + +// Greedy-acceptance epilogue run by the last CTA once the global per-row argmax is known (rows = nreq requests x +// (draft_len + 1) verify positions, request-major). Mirrors the strict path of SpecWorkerBase: accepted = the target +// tokens, count = 1 + length of the matching draft prefix, optionally forced; plus the first draft pass's row index. +struct AcceptArgs +{ + int const* draft; // [nreq, draft_len] previous draft tokens + int const* batch_idx; // [nreq] request row in the verify layout + int* acc_tokens; // [nreq, draft_len + 1] + int* num_acc; // [nreq] + long long* gather_ids; // [nreq] last accepted row + int draft_len; + int nreq; + int force_total; // > 0 overrides the accepted count + int* next_new; // optional [nreq, next_stride]: column 0 = last accepted token (next step's input) + int next_stride; + int const* force_vec; // optional [nreq] per-request accepted count; overrides force_total +}; + +//! Greedy verification from this rank's BF16 logits shard [M, Vs] (row stride ld), M <= kMaxRows, over numCtas +//! kThreads-thread CTAs, followed by the strict acceptance epilogue. +cudaError_t invokeShardAccept(__nv_bfloat16 const* logits, int M, int Vs, long long ld, int vocabOffset, + unsigned long long* keys, unsigned int* counters, unsigned long long* wsUc, unsigned long long* wsMc, int rank, + int world, AcceptArgs const& acc, int numCtas, cudaStream_t stream); + +} // namespace shardAccept +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/sparseDecodeLc.cu b/cpp/tensorrt_llm/kernels/minimaxM3/sparseDecodeLc.cu new file mode 100644 index 000000000000..e01fa78927a5 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/sparseDecodeLc.cu @@ -0,0 +1,561 @@ +/* + * 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. + */ + +// MiniMax-M3 long-context sparse decode attention in one launch. +// +// Grid: one cluster of 16 CTAs per query token; CTA `r` attends the token's `r`-th selected block (the Triton +// producer's chunk `r` at 16 chunks), keeps the partial in shared memory, and after a cluster barrier merges head +// `r` from all 16 partials over DSMEM, writing the BF16 row and, optionally, its MXFP8 copy and swizzled UE8M0 +// scales exactly as the Triton merge's `_store_mxfp8_head` lays them out. +// +// Per-block math follows the Triton producer: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16, QK and +// PV on FP16 tensor cores with FP32 accumulation, log2-domain softmax with either the linear denominator or the +// log-domain recurrence. + +#include "tensorrt_llm/kernels/minimaxM3/sparseDecodeLc.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kHeads = 16; +constexpr int kDim = 128; +constexpr int kPage = 128; +constexpr int kTopk = 16; +constexpr int kThreads = 128; +constexpr int kPitch = kDim + 8; // halves per smem row: 272 B, conflict-free ldmatrix +constexpr float kNegInf = -__builtin_huge_valf(); + +using Params = MinimaxM3SparseDecodeLcParams; + +struct Smem +{ + __half qs[kHeads * kPitch]; + __half vs[kPage * kPitch]; + __half ps[kHeads * kPitch]; + float recv[kTopk][kDim]; // head `rank` partials pushed by the 16 CTAs, chunk-major + float recv_lse[kTopk]; + float red[4][kHeads]; + float rowm[kHeads]; + float rowl[kHeads]; + alignas(8) unsigned long long rbar; // the split-K partials of head `rank` landed (complete_tx) +}; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ uint32_t e2m1x2_to_f16x2(uint32_t byte) +{ + uint32_t out; + asm("{\n .reg .b8 s;\n cvt.u8.u32 s, %1;\n cvt.rn.f16x2.e2m1x2 %0, s;\n}" : "=r"(out) : "r"(byte)); + return out; +} + +__device__ __forceinline__ uint32_t e4m3x2_to_f16x2(uint16_t pair) +{ + uint32_t out; + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(out) : "h"(pair)); + return out; +} + +__device__ __forceinline__ float ex2(float x) +{ + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ float lg2(float x) +{ + float y; + asm("lg2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ void ldsm_x4(uint32_t (&r)[4], uint32_t addr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x4_t(uint32_t (&r)[4], uint32_t addr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void mma(float (&c)[4], uint32_t const (&a)[4], 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};" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ uint32_t map_cluster(void const* local, uint32_t rank) +{ + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(smem_u32(local)), "r"(rank)); + return remote; +} + +// Posted DSMEM stores that complete_tx on the receiving CTA's mbarrier `bar`. +__device__ __forceinline__ void st_async_v2(uint32_t addr, float a, float b, uint32_t bar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];" ::"r"(addr), + "f"(a), "f"(b), "r"(bar) + : "memory"); +} + +__device__ __forceinline__ void st_async_f32(uint32_t addr, float a, uint32_t bar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(addr), + "r"(__float_as_uint(a)), "r"(bar) + : "memory"); +} + +// One packed NVFP4 row (64 bytes, 128 values) with its 8 E4M3 scales -> 128 fp16 values in smem. +__device__ __forceinline__ void dequant_row(__half* dst, uint8_t const* packed, uint8_t const (&sc)[8], bool valid) +{ + uint4 const* src = reinterpret_cast(packed); + uint32_t* d = reinterpret_cast(dst); +#pragma unroll + for (int c = 0; c < 4; ++c) + { + uint4 w = valid ? __ldg(src + c) : make_uint4(0, 0, 0, 0); + uint32_t words[4] = {w.x, w.y, w.z, w.w}; + // 16 bytes = 32 values = scale columns 2c and 2c + 1. +#pragma unroll + for (int h = 0; h < 2; ++h) + { + uint8_t const s = valid ? sc[2 * c + h] : 0; + uint32_t const sh = e4m3x2_to_f16x2(static_cast(s | (s << 8))); + __half2 const scale = *reinterpret_cast<__half2 const*>(&sh); +#pragma unroll + for (int j = 0; j < 2; ++j) + { + uint32_t const word = words[2 * h + j]; +#pragma unroll + for (int b = 0; b < 4; ++b) + { + uint32_t const v2 = e2m1x2_to_f16x2((word >> (8 * b)) & 0xFF); + __half2 const x = __hmul2(*reinterpret_cast<__half2 const*>(&v2), scale); + d[c * 16 + h * 8 + j * 4 + b] = *reinterpret_cast(&x); + } + } + } + } +} +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +template +__global__ void __launch_bounds__(kThreads) m3_lc_attn_kernel(__grid_constant__ Params const p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ __align__(128) uint8_t smem_raw[]; + Smem& s = *reinterpret_cast(smem_raw); + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const tok = blockIdx.x / kTopk; + uint32_t rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(rank)); + + if (tid == 0) + { + // One arrival (here) plus every cluster CTA's partial row (kDim floats) and LSE for head `rank`. + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smem_u32(&s.rbar)) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smem_u32(&s.rbar)), + "r"(kTopk * (kDim + 1) * 4) + : "memory"); + } + // Every peer must be running before its shared memory takes a DSMEM store: arrive now, wait before pushing. + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); + + // Everything that does not depend on the top-16 table is read before the grid dependency: the request's length + // (written before the step's graph runs) and the query rows (written by the producer, which completed before the + // select ahead of this grid could release it: the select triggers its dependents only after its own grid wait). + int const req = tok / p.dql; + int const qoff = tok - req * p.dql; + int const seq = p.seq_lens[req]; + int const kv_len = max(seq - p.dql + qoff + 1, 0); + int const nblocks = (kv_len + kPage - 1) / kPage; + int const real_topk = min(kTopk, nblocks); + // q: 16 heads x 128 e4m3 -> fp16; thread = (head, 16-value chunk). + { + int const h = tid >> 3; + int const c = (tid & 7) * 16; + uint4 const w = *reinterpret_cast(p.q + tok * p.q_tok + h * p.q_head + c); + uint32_t const words[4] = {w.x, w.y, w.z, w.w}; + uint32_t* d = reinterpret_cast(s.qs + h * kPitch + c); +#pragma unroll + for (int i = 0; i < 4; ++i) + { + d[2 * i] = e4m3x2_to_f16x2(static_cast(words[i] & 0xFFFF)); + d[2 * i + 1] = e4m3x2_to_f16x2(static_cast(words[i] >> 16)); + } + } + + asm volatile("griddepcontrol.wait;" ::: "memory"); + if constexpr (EARLY) + { + asm volatile("griddepcontrol.launch_dependents;"); + } + int blk = static_cast(rank) < real_topk ? p.topk[tok * p.t_tok + rank * p.t_k] : -1; + bool const active = blk >= 0; + int64_t const page_pw = !active ? 0 : p.block_table[req * p.bt_stride + blk]; + + float m_row[2] = {kNegInf, kNegInf}; + float acc[4][4]; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + acc[j][i] = 0.F; + float l_row[2] = {0.F, 0.F}; + + // K stays packed in registers: the MMA B fragment of token n = lane / 4 at k-step s is the f16x2 pair of + // packed bytes 8s + q and 8s + 4 + q (q = lane % 4), both under scale column s. + uint32_t kw[4][16]; + uint2 ksc[4]; + int64_t page = 0; + if (active) + { + page = page_pw; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + int const n = warp * 32 + j * 8 + (lane >> 2); + bool const valid = blk * kPage + n < kv_len; + uint4 const* src = reinterpret_cast(p.k + page * p.k_blk + n * 64); +#pragma unroll + for (int c = 0; c < 4; ++c) + { + uint4 const w = valid ? __ldg(src + c) : make_uint4(0, 0, 0, 0); + kw[j][4 * c] = w.x; + kw[j][4 * c + 1] = w.y; + kw[j][4 * c + 2] = w.z; + kw[j][4 * c + 3] = w.w; + } + ksc[j] = valid ? __ldg(reinterpret_cast(p.ksc + page * p.ksc_blk + n * 8)) : make_uint2(0, 0); + } + bool const valid = blk * kPage + tid < kv_len; + uint8_t vsc[8]; + if (valid) + { + uint8_t const* vb = p.vsc + page * p.vsc_blk + (tid >> 2) * 32 + (tid & 3); +#pragma unroll + for (int c = 0; c < 8; ++c) + vsc[c] = __ldg(vb + 4 * c); + } + dequant_row(s.vs + tid * kPitch, p.v + page * p.v_blk + tid * 64, vsc, valid); + } + __syncthreads(); + + if (active) + { + float const scale_log2 = p.sm_scale * 1.4426950409F * __ldg(p.kgs); + int const q4 = (lane & 3) * 8; + // S = q K^T for this warp's 32 tokens: rows g / g + 8, columns 2 * (lane % 4) + {0, 1} of each n-tile. + float sacc[4][4]; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + sacc[j][i] = 0.F; +#pragma unroll + for (int ks = 0; ks < kDim / 16; ++ks) + { + uint32_t a[4]; + ldsm_x4(a, smem_u32(s.qs + (lane & 15) * kPitch + ks * 16 + (lane >> 4) * 8)); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + uint32_t const scw = ks < 4 ? ksc[j].x : ksc[j].y; + uint8_t const sc = static_cast(scw >> (8 * (ks & 3))); + uint32_t const sh = e4m3x2_to_f16x2(static_cast(sc | (sc << 8))); + __half2 const scale = *reinterpret_cast<__half2 const*>(&sh); + uint32_t const r0 = e2m1x2_to_f16x2((kw[j][2 * ks] >> q4) & 0xFF); + uint32_t const r1 = e2m1x2_to_f16x2((kw[j][2 * ks + 1] >> q4) & 0xFF); + __half2 const h0 = __hmul2(*reinterpret_cast<__half2 const*>(&r0), scale); + __half2 const h1 = __hmul2(*reinterpret_cast<__half2 const*>(&r1), scale); + mma(sacc[j], a, *reinterpret_cast(&h0), *reinterpret_cast(&h1)); + } + } + int const g = lane >> 2; + float mloc[2] = {kNegInf, kNegInf}; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const t = warp * 32 + j * 8 + 2 * (lane & 3) + (i & 1); + float const v = blk * kPage + t < kv_len ? 0.F + sacc[j][i] * scale_log2 : kNegInf; + sacc[j][i] = v; + mloc[i >> 1] = fmaxf(mloc[i >> 1], v); + } +#pragma unroll + for (int r = 0; r < 2; ++r) + { + mloc[r] = fmaxf(mloc[r], __shfl_xor_sync(0xffffffff, mloc[r], 1)); + mloc[r] = fmaxf(mloc[r], __shfl_xor_sync(0xffffffff, mloc[r], 2)); + } + if ((lane & 3) == 0) + { + s.red[warp][g] = mloc[0]; + s.red[warp][g + 8] = mloc[1]; + } + __syncthreads(); + if (tid < kHeads) + s.rowm[tid] = fmaxf(fmaxf(s.red[0][tid], s.red[1][tid]), fmaxf(s.red[2][tid], s.red[3][tid])); + __syncthreads(); + m_row[0] = s.rowm[g]; + m_row[1] = s.rowm[g + 8]; + float lloc[2] = {0.F, 0.F}; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; i += 2) + { + float const p0 = ex2(sacc[j][i] - m_row[i >> 1]); + float const p1 = ex2(sacc[j][i + 1] - m_row[i >> 1]); + lloc[i >> 1] += p0 + p1; + int const row = g + (i >> 1) * 8; + int const t = warp * 32 + j * 8 + 2 * (lane & 3); + *reinterpret_cast<__half2*>(s.ps + row * kPitch + t) = __floats2half2_rn(p0, p1); + } +#pragma unroll + for (int r = 0; r < 2; ++r) + { + lloc[r] += __shfl_xor_sync(0xffffffff, lloc[r], 1); + lloc[r] += __shfl_xor_sync(0xffffffff, lloc[r], 2); + } + __syncthreads(); // rowm reads done before red is reused; ps complete + if ((lane & 3) == 0) + { + s.red[warp][g] = lloc[0]; + s.red[warp][g + 8] = lloc[1]; + } + // O[:, warp's 32 dims] = P V over all 128 tokens. +#pragma unroll + for (int kk = 0; kk < kPage; kk += 16) + { + uint32_t a[4]; + ldsm_x4(a, smem_u32(s.ps + (lane & 15) * kPitch + kk + (lane >> 4) * 8)); +#pragma unroll + for (int jp = 0; jp < 2; ++jp) + { + uint32_t b[4]; + int const t = kk + (lane & 7) + ((lane >> 3) & 1) * 8; + int const d = warp * 32 + jp * 16 + (lane >> 4) * 8; + ldsm_x4_t(b, smem_u32(s.vs + t * kPitch + d)); + mma(acc[2 * jp], a, b[0], b[1]); + mma(acc[2 * jp + 1], a, b[2], b[3]); + } + } + __syncthreads(); + if (tid < kHeads) + s.rowl[tid] = (s.red[0][tid] + s.red[1][tid]) + (s.red[2][tid] + s.red[3][tid]); + __syncthreads(); + l_row[0] = s.rowl[g]; + l_row[1] = s.rowl[g + 8]; + } + + // Normalize the partial and push head h's row with its LSE (log2) into CTA h's recv[rank] slot, as the Triton + // producer stores chunk `rank`; st.async stores that complete_tx on CTA h's barrier. + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); + { + int const g = lane >> 2; + float const vgs = __ldg(p.vgs); +#pragma unroll + for (int r = 0; r < 2; ++r) + { + float scale, lse; + if constexpr (LINEAR) + { + scale = l_row[r] > 0.F ? 1.0F / l_row[r] : 0.F; + lse = l_row[r] > 0.F ? m_row[r] + lg2(l_row[r]) : kNegInf; + } + else + { + float const lse_i = active ? m_row[r] + lg2(ex2(kNegInf - m_row[r]) + l_row[r]) : kNegInf; + scale = lse_i > kNegInf ? ex2(m_row[r] - lse_i) : 0.F; + lse = lse_i; + } + scale = scale * vgs; + int const row = g + r * 8; + uint32_t const dst = map_cluster(&s.recv[rank][0], row); + uint32_t const rb = map_cluster(&s.rbar, row); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + int const d = warp * 32 + j * 8 + 2 * (lane & 3); + st_async_v2(dst + d * 4, acc[j][2 * r] * scale, acc[j][2 * r + 1] * scale, rb); + } + if (warp == 0 && (lane & 3) == 0) + st_async_f32(map_cluster(&s.recv_lse[rank], row), lse, rb); + } + } + { + // Every CTA's partial row and LSE for head `rank` landed (complete_tx on this CTA's barrier). + uint32_t done = 0; + while (!done) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}" + : "=r"(done) + : "r"(smem_u32(&s.rbar)) + : "memory"); + } + } + + // CTA `rank` merges head `rank` across the 16 partials; thread = output dim. + int const h = static_cast(rank); + float lse_c[kTopk]; + float lse_max = kNegInf; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + { + lse_c[c] = s.recv_lse[c]; + lse_max = fmaxf(lse_max, lse_c[c]); + } + lse_max = lse_max == kNegInf ? 0.F : lse_max; + float denom = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + { + lse_c[c] = ex2(lse_c[c] - lse_max); + denom += lse_c[c]; + } + denom = denom > 0.F ? denom : 1.F; + float o = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + o += s.recv[c][tid] * __fdiv_rn(lse_c[c], denom); + __nv_bfloat16 const ob = __float2bfloat16_rn(o); + p.out[tok * p.out_stride + h * kDim + tid] = ob; + + if constexpr (MXFP8) + { + float const x = __bfloat162float(ob); + float amax = fabsf(x); +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, off)); + float rcp; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(rcp) : "f"(448.0F)); + float const sfv = amax * rcp; + int const bits = __float_as_int(sfv); + int e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0); + e = min(e, 254); + float const qs = amax != 0.F ? __int_as_float((254 - e) << 23) : 0.F; + float const xq = x * qs; + uint16_t q2; + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(q2) : "f"(0.F), "f"(xq)); + p.q8[static_cast(tok) * (p.num_sf_k_tiles * 128) + h * kDim + tid] = static_cast(q2 & 0xFF); + if (lane == 0) + { + int const k_sf = h * 4 + warp; + int64_t const off = static_cast(tok / 128) * p.num_sf_k_tiles * 512 + (k_sf / 4) * 512 + + (tok % 32) * 16 + ((tok % 128) / 32) * 4 + k_sf % 4; + p.sf[off] = static_cast(e); + } + } + if constexpr (!EARLY) + { + asm volatile("griddepcontrol.launch_dependents;"); + } +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +template +cudaError_t configure() +{ + auto kernel = m3_lc_attn_kernel; + cudaError_t err = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(Smem)); + if (err != cudaSuccess) + { + return err; + } + return cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1); +} + +template +cudaError_t launch(Params const& prm, int num_tokens, bool use_pdl, cudaStream_t stream) +{ + // Kernel attributes are set once per process, on the first launch (outside CUDA-graph capture: warmup runs + // every variant eagerly first). + static cudaError_t const configured = configure(); + if (configured != cudaSuccess) + { + return configured; + } + auto kernel = m3_lc_attn_kernel; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(num_tokens * kTopk); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = sizeof(Smem); + cfg.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kTopk; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; + cfg.attrs = attrs; + cfg.numAttrs = 2; + return cudaLaunchKernelEx(&cfg, kernel, prm); +} + +} // namespace + +cudaError_t invokeMinimaxM3SparseDecodeLc(MinimaxM3SparseDecodeLcParams const& params, int32_t numTokens, + bool linearSoftmax, bool mxfp8, bool earlyTrigger, bool usePdl, cudaStream_t stream) +{ + if (!earlyTrigger) + { + return cudaErrorNotSupported; + } + if (linearSoftmax) + { + return mxfp8 ? launch(params, numTokens, usePdl, stream) + : launch(params, numTokens, usePdl, stream); + } + return mxfp8 ? launch(params, numTokens, usePdl, stream) + : launch(params, numTokens, usePdl, stream); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/sparseDecodeLc.h b/cpp/tensorrt_llm/kernels/minimaxM3/sparseDecodeLc.h new file mode 100644 index 000000000000..659531f725bf --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/sparseDecodeLc.h @@ -0,0 +1,70 @@ +/* + * 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 +{ + +//! Operands of the MiniMax-M3 one-launch long-context sparse decode. Strides are in elements of the pointed-to type. +//! Geometry is fixed: 16 query heads on one KV head, head_dim 128, page 128, top-k 16. +struct MinimaxM3SparseDecodeLcParams +{ + uint8_t const* q; // e4m3 [T, 16, 128] + int64_t q_tok, q_head; + uint8_t const* k; // packed e2m1 [pages, 1, 128, 64] + uint8_t const* v; + int64_t k_blk, v_blk; + uint8_t const* ksc; // e4m3 [pages, 1, 128, 8], token-major + uint8_t const* vsc; // e4m3, 4x4 token-quad order + int64_t ksc_blk, vsc_blk; + float const* kgs; + float const* vgs; + int const* topk; // [T, 16] view (kv head 0) + int64_t t_tok, t_k; + int const* block_table; + int64_t bt_stride; + int const* seq_lens; + int dql; + float sm_scale; + __nv_bfloat16* out; // [T, 2048] + int64_t out_stride; + uint8_t* q8; // e4m3 [rows, num_sf_k_tiles * 128] + uint8_t* sf; + int num_sf_k_tiles; +}; + +//! One cluster of 16 CTAs per query token: CTA r attends the token's r-th selected block, then the cluster merges +//! the 16 partials over DSMEM into the BF16 row (and, when mxfp8, its MXFP8 copy with swizzled UE8M0 scales). +//! +//! Instantiated variants: linearSoftmax in {false, true} x mxfp8 in {false, true}, with earlyTrigger (release PDL +//! dependents at entry). Returns cudaErrorNotSupported for any other variant. The device code is SM100-family only; +//! on other architectures the kernel compiles to an empty body, so callers check the device first. +cudaError_t invokeMinimaxM3SparseDecodeLc(MinimaxM3SparseDecodeLcParams const& params, int32_t numTokens, + bool linearSoftmax, bool mxfp8, bool earlyTrigger, bool usePdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu b/cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu new file mode 100644 index 000000000000..4069da14faf0 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu @@ -0,0 +1,425 @@ +/* + * 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/minimaxM3/specTail.h" + +#include +#include +#include + +// Small fused kernels of the MiniMax-M3 Eagle3 step (GB300, decode). +// +// chunk_norm / capture_norm apply the drafter's per-layer fc_norm RMSNorms to the captured target hidden states with +// the in-tree RMSNorm arithmetic (fp32 statistic, (x * rsqrt) * w, one bf16 rounding); step_gather / slot_scatter move +// the step's integer outputs between the spec sampler's store and the engine's input buffers. +// +// draft_input replaces a fused draft pass's embedding lookup, its two RMSNorms (input_layernorm on the embedding, +// hidden_norm on the carried hidden state, the same arithmetic) and their concatenation with one launch. It optionally +// gathers the carried hidden rows (the second pass reads the last accepted row of each request) and writes them out +// as the residual of the decoder layer. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3 +{ +namespace specTail +{ +namespace +{ +constexpr int kVec = 8; // bf16 per 16-byte access + +__device__ __forceinline__ float block_sum(float v, float* red) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + int const warp = threadIdx.x >> 5; + if ((threadIdx.x & 31) == 0) + red[warp] = v; + __syncthreads(); + float s = 0.f; +#pragma unroll + for (int w = 0; w < kThreads / 32; ++w) + s += red[w]; + return s; +} + +// blockIdx.x = row, blockIdx.y = 0 (embedding half) / 1 (hidden half). H % (kThreads * kVec) == 0 is not required; +// H % kVec == 0 and H <= kMaxPerThread * kThreads * kVec are. +template +__global__ void __launch_bounds__(kThreads) + draft_input_kernel(int const* __restrict__ ids, int vocab, __nv_bfloat16 const* __restrict__ table, + __nv_bfloat16 const* __restrict__ hidden, long long const* __restrict__ hidden_rows, + __nv_bfloat16 const* __restrict__ w_embed, __nv_bfloat16 const* __restrict__ w_hidden, float eps_embed, + float eps_hidden, int H, __nv_bfloat16* __restrict__ out, __nv_bfloat16* __restrict__ residual) +{ + __shared__ float red[kThreads / 32]; + int const row = blockIdx.x; + bool const is_hidden = blockIdx.y == 1; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + __nv_bfloat16 const* src; + bool valid = true; + if (is_hidden) + { + long long const r = hidden_rows ? hidden_rows[row] : row; + src = hidden + r * H; + } + else + { + int const id = ids[row]; + valid = id >= 0 && id < vocab; + src = table + static_cast(valid ? id : 0) * H; + } + float v[kIters][kVec]; + float ss = 0.f; +#pragma unroll + for (int it = 0; it < kIters; ++it) + { + int const c = (it * kThreads + threadIdx.x) * kVec; + if (c < H) + { + uint4 const raw = *reinterpret_cast(src + c); + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); +#pragma unroll + for (int j = 0; j < kVec; ++j) + { + v[it][j] = valid ? __bfloat162float(e[j]) : 0.f; + ss += v[it][j] * v[it][j]; + } + if (is_hidden && residual) + *reinterpret_cast(residual + static_cast(row) * H + c) = raw; + } + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float const rs = rsqrtf(block_sum(ss, red) / H + (is_hidden ? eps_hidden : eps_embed)); + __nv_bfloat16 const* w = is_hidden ? w_hidden : w_embed; + __nv_bfloat16* dst = out + static_cast(row) * 2 * H + (is_hidden ? H : 0); +#pragma unroll + for (int it = 0; it < kIters; ++it) + { + int const c = (it * kThreads + threadIdx.x) * kVec; + if (c < H) + { + uint4 const wr = *reinterpret_cast(w + c); + __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); + uint4 o; + __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); +#pragma unroll + for (int j = 0; j < kVec; ++j) + oe[j] = __float2bfloat16(v[it][j] * rs * __bfloat162float(we[j])); + *reinterpret_cast(dst + c) = o; + } + } +} + +// blockIdx.x = row, blockIdx.y = chunk c: dst[row, c*H:(c+1)*H] = rmsnorm(src[row, c*H:(c+1)*H]) * w[c]. +template +__global__ void __launch_bounds__(kThreads) chunk_norm_kernel(__nv_bfloat16 const* __restrict__ src, + long long src_stride, __nv_bfloat16 const* __restrict__ w, float eps, int H, int C, __nv_bfloat16* __restrict__ dst) +{ + __shared__ float red[kThreads / 32]; + int const row = blockIdx.x; + int const c = blockIdx.y; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + __nv_bfloat16 const* in = src + row * src_stride + static_cast(c) * H; + float v[kIters][kVec]; + float ss = 0.f; +#pragma unroll + for (int it = 0; it < kIters; ++it) + { + int const col = (it * kThreads + threadIdx.x) * kVec; + if (col < H) + { + uint4 const raw = *reinterpret_cast(in + col); + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); +#pragma unroll + for (int j = 0; j < kVec; ++j) + { + v[it][j] = __bfloat162float(e[j]); + ss += v[it][j] * v[it][j]; + } + } + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float const rs = rsqrtf(block_sum(ss, red) / H + eps); + __nv_bfloat16 const* wc = w + static_cast(c) * H; + __nv_bfloat16* out = dst + (static_cast(row) * C + c) * H; +#pragma unroll + for (int it = 0; it < kIters; ++it) + { + int const col = (it * kThreads + threadIdx.x) * kVec; + if (col < H) + { + uint4 const wr = *reinterpret_cast(wc + col); + __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); + uint4 o; + __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); +#pragma unroll + for (int j = 0; j < kVec; ++j) + oe[j] = __float2bfloat16(v[it][j] * rs * __bfloat162float(we[j])); + *reinterpret_cast(out + col) = o; + } + } +} + +// Eagle3 hidden-state capture with the drafter's fc_norm folded in (one row per CTA): +// dst[row, :H] = rmsnorm(bf16(h[row] + r[row])) * w (r optional) +// with chunk_norm's arithmetic, so fc reads the same values the capture + chunk_norm pair produced. +template +__global__ void __launch_bounds__(kThreads) capture_norm_kernel(__nv_bfloat16 const* __restrict__ h, long long h_stride, + __nv_bfloat16 const* __restrict__ r, long long r_stride, __nv_bfloat16 const* __restrict__ w, float eps, int H, + __nv_bfloat16* __restrict__ dst, long long dst_stride) +{ + __shared__ float red[kThreads / 32]; + int const row = blockIdx.x; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + float v[kIters][kVec]; + float ss = 0.f; +#pragma unroll + for (int it = 0; it < kIters; ++it) + { + int const col = (it * kThreads + threadIdx.x) * kVec; + if (col < H) + { + uint4 const raw = *reinterpret_cast(h + row * h_stride + col); + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); + uint4 rraw = make_uint4(0u, 0u, 0u, 0u); + if (r != nullptr) + rraw = *reinterpret_cast(r + row * r_stride + col); + __nv_bfloat16 const* re = reinterpret_cast<__nv_bfloat16 const*>(&rraw); +#pragma unroll + for (int j = 0; j < kVec; ++j) + { + float const x = __bfloat162float(e[j]); + v[it][j] = r != nullptr ? __bfloat162float(__float2bfloat16(x + __bfloat162float(re[j]))) : x; + ss += v[it][j] * v[it][j]; + } + } + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float const rs = rsqrtf(block_sum(ss, red) / H + eps); + __nv_bfloat16* out = dst + row * dst_stride; +#pragma unroll + for (int it = 0; it < kIters; ++it) + { + int const col = (it * kThreads + threadIdx.x) * kVec; + if (col < H) + { + uint4 const wr = *reinterpret_cast(w + col); + __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); + uint4 o; + __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); +#pragma unroll + for (int j = 0; j < kVec; ++j) + oe[j] = __float2bfloat16(v[it][j] * rs * __bfloat162float(we[j])); + *reinterpret_cast(out + col) = o; + } + } +} + +// Overlap-scheduler step inputs from the previous step's device outputs (one launch for the engine's gathers): +// ids[ids0 + p*T + t] = new_tok[t * nt_s0 + slot_p * nt_s1] +// draft[d0 + p*D + d] = next_draft[slot_p * nd_s0 + d * nd_s1] +// pos_off[i] = i in [b0*T, (b0+P)*T) ? lens[ppi[i - b0*T]] : 0 (i < n_pos) +// kv_off[i] = i in [b0, b0+P) ? lens[slot_{i-b0}] - T : 0 (i < n_kv) +__global__ void step_gather_kernel(int const* __restrict__ new_tok, long long nt_s0, long long nt_s1, + int const* __restrict__ next_draft, long long nd_s0, long long nd_s1, int const* __restrict__ lens, + int const* __restrict__ slots, int const* __restrict__ ppi, int P, int T, int D, int* __restrict__ ids, int ids0, + int* __restrict__ draft, int d0, int* __restrict__ pos_off, int n_pos, int* __restrict__ kv_off, int n_kv, int b0) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + int const tid = blockIdx.x * blockDim.x + threadIdx.x; + int const stride = gridDim.x * blockDim.x; + for (int i = tid; i < P * T; i += stride) + { + int const p = i / T, t = i - p * T; + ids[ids0 + i] = new_tok[t * nt_s0 + static_cast(slots[p]) * nt_s1]; + } + for (int i = tid; i < P * D; i += stride) + { + int const p = i / D, d = i - p * D; + draft[d0 + i] = next_draft[static_cast(slots[p]) * nd_s0 + d * nd_s1]; + } + for (int i = tid; i < n_pos; i += stride) + { + int const k = i - b0 * T; + pos_off[i] = (k >= 0 && k < P * T) ? lens[ppi[k]] : 0; + } + for (int i = tid; i < n_kv; i += stride) + { + int const k = i - b0; + kv_off[i] = (k >= 0 && k < P) ? lens[slots[k]] - T : 0; + } +} + +// Sampler stores from one step's outputs (one launch for its index_copy_s), zero-padding to the store widths: +// a[slot_i][c] for c < wa (from src_a[i][c] if c < sa_w), b likewise, lens[slot_i] = src_lens[i], d likewise. +__global__ void slot_scatter_kernel(long long const* __restrict__ slots, int n, int const* __restrict__ sa, + long long sa_s0, long long sa_s1, int sa_w, int* __restrict__ da, long long da_s0, long long da_s1, int wa, + int const* __restrict__ sb, long long sb_s0, long long sb_s1, int sb_w, int* __restrict__ db, long long db_s0, + long long db_s1, int wb, int const* __restrict__ sl, int* __restrict__ dl, int const* __restrict__ sd, + long long sd_s0, long long sd_s1, int sd_w, int* __restrict__ dd, long long dd_s0, long long dd_s1, int wd) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + int const i = blockIdx.x; + if (i >= n) + return; + long long const s = slots[i]; + for (int c = threadIdx.x; c < wa; c += blockDim.x) + da[s * da_s0 + c * da_s1] = c < sa_w ? sa[i * sa_s0 + c * sa_s1] : 0; + for (int c = threadIdx.x; c < wb; c += blockDim.x) + db[s * db_s0 + c * db_s1] = c < sb_w ? sb[i * sb_s0 + c * sb_s1] : 0; + for (int c = threadIdx.x; c < wd; c += blockDim.x) + dd[s * dd_s0 + c * dd_s1] = c < sd_w ? sd[i * sd_s0 + c * sd_s1] : 0; + if (threadIdx.x == 0) + dl[s] = sl[i]; +} + +cudaLaunchConfig_t small_cfg(int grid, int block, bool pdl, cudaLaunchAttribute* attr, cudaStream_t stream) +{ + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(grid); + cfg.blockDim = dim3(block); + cfg.stream = stream; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cfg; +} + +int norm_iters(int H) +{ + return (H / kVec + kThreads - 1) / kThreads; +} + +} // namespace + +cudaError_t invokeDraftInput(int const* ids, int vocab, __nv_bfloat16 const* table, __nv_bfloat16 const* hidden, + long long const* rows, __nv_bfloat16 const* wEmbed, __nv_bfloat16 const* wHidden, float epsEmbed, float epsHidden, + int H, int M, __nv_bfloat16* out, __nv_bfloat16* residual, bool pdl, cudaStream_t stream) +{ + int const iters = norm_iters(H); + cudaLaunchAttribute attr[1]; + cudaLaunchConfig_t cfg = small_cfg(1, kThreads, pdl, attr, stream); + cfg.gridDim = dim3(M, 2); +#define LAUNCH(I) \ + cudaLaunchKernelEx(&cfg, draft_input_kernel, ids, vocab, table, hidden, rows, wEmbed, wHidden, epsEmbed, \ + epsHidden, H, out, residual) + switch (iters) + { + case 0: + case 1: return LAUNCH(1); + case 2: return LAUNCH(2); + case 3: return LAUNCH(3); + case 4: return LAUNCH(4); + default: return cudaErrorInvalidValue; + } +#undef LAUNCH +} + +cudaError_t invokeChunkNorm(__nv_bfloat16 const* src, long long srcStride, __nv_bfloat16 const* w, float eps, int H, + int C, int M, __nv_bfloat16* dst, bool pdl, cudaStream_t stream) +{ + int const iters = norm_iters(H); + cudaLaunchAttribute attr[1]; + cudaLaunchConfig_t cfg = small_cfg(1, kThreads, pdl, attr, stream); + cfg.gridDim = dim3(M, C); +#define LAUNCH(I) cudaLaunchKernelEx(&cfg, chunk_norm_kernel, src, srcStride, w, eps, H, C, dst) + switch (iters) + { + case 0: + case 1: return LAUNCH(1); + case 2: return LAUNCH(2); + case 3: return LAUNCH(3); + case 4: return LAUNCH(4); + default: return cudaErrorInvalidValue; + } +#undef LAUNCH +} + +cudaError_t invokeCaptureNorm(__nv_bfloat16 const* h, long long hStride, __nv_bfloat16 const* res, long long resStride, + __nv_bfloat16 const* w, float eps, int H, int M, __nv_bfloat16* dst, long long dstStride, cudaStream_t stream) +{ + if (M == 0) + { + return cudaSuccess; + } + int const iters = norm_iters(H); + cudaLaunchAttribute attr[1]; + cudaLaunchConfig_t cfg = small_cfg(M, kThreads, true, attr, stream); +#define LAUNCH(I) \ + cudaLaunchKernelEx(&cfg, capture_norm_kernel, h, hStride, res, resStride, w, eps, H, dst, dstStride) + switch (iters) + { + case 0: + case 1: return LAUNCH(1); + case 2: return LAUNCH(2); + case 3: return LAUNCH(3); + case 4: return LAUNCH(4); + default: return cudaErrorInvalidValue; + } +#undef LAUNCH +} + +cudaError_t invokeStepGather(int const* newTok, long long ntS0, long long ntS1, int const* nextDraft, long long ndS0, + long long ndS1, int const* lens, int const* slots, int const* ppi, int P, int T, int D, int* ids, int ids0, + int* draft, int d0, int* posOff, int nPos, int* kvOff, int nKv, int b0, int work, cudaStream_t stream) +{ + cudaLaunchAttribute attr[1]; + cudaLaunchConfig_t cfg = small_cfg(std::max(1, std::min(64, (work + 255) / 256)), 256, true, attr, stream); + return cudaLaunchKernelEx(&cfg, step_gather_kernel, newTok, ntS0, ntS1, nextDraft, ndS0, ndS1, lens, slots, ppi, P, + T, D, ids, ids0, draft, d0, posOff, nPos, kvOff, nKv, b0); +} + +cudaError_t invokeSlotScatter(long long const* slots, int n, IntView sa, IntView da, IntView sb, IntView db, + int const* sl, int* dl, IntView sd, IntView dd, cudaStream_t stream) +{ + if (n == 0) + { + return cudaSuccess; + } + cudaLaunchAttribute attr[1]; + cudaLaunchConfig_t cfg = small_cfg(n, 32, true, attr, stream); + return cudaLaunchKernelEx(&cfg, slot_scatter_kernel, slots, n, sa.data, sa.s0, sa.s1, sa.width, da.data, da.s0, + da.s1, da.width, sb.data, sb.s0, sb.s1, sb.width, db.data, db.s0, db.s1, db.width, sl, dl, sd.data, sd.s0, + sd.s1, sd.width, dd.data, dd.s0, dd.s1, dd.width); +} + +} // namespace specTail +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/specTail.h b/cpp/tensorrt_llm/kernels/minimaxM3/specTail.h new file mode 100644 index 000000000000..c5bdf75ff482 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/specTail.h @@ -0,0 +1,80 @@ +/* + * 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 +{ +namespace minimaxM3 +{ +//! Small fused kernels of the MiniMax-M3 Eagle3 step: the drafter's fc_norm RMSNorms over the captured hidden states +//! and the step's input / output bookkeeping (GB300, decode). All launches use programmatic dependent launch where +//! noted. +namespace specTail +{ + +//! Threads per CTA of the norm kernels, and the hidden-size bound of their register tiles: +//! H % 8 == 0 and H <= kMaxIters * kThreads * 8. +constexpr int kThreads = 256; +constexpr int kMaxIters = 4; + +//! out[r] = [rmsnorm(table[ids[r]]) * wEmbed, rmsnorm(hidden[rows ? rows[r] : r]) * wHidden] ([M, 2H] BF16; ids +//! outside [0, vocab) embed to zero); residual[r] = hidden[rows[r]] when residual is set. +cudaError_t invokeDraftInput(int const* ids, int vocab, __nv_bfloat16 const* table, __nv_bfloat16 const* hidden, + long long const* rows, __nv_bfloat16 const* wEmbed, __nv_bfloat16 const* wHidden, float epsEmbed, float epsHidden, + int H, int M, __nv_bfloat16* out, __nv_bfloat16* residual, bool pdl, cudaStream_t stream); + +//! dst[r, c*H:(c+1)*H] = rmsnorm(src[r, c*H:(c+1)*H]) * w[c] for c < C (src row stride srcStride), dst [M, C*H]. +cudaError_t invokeChunkNorm(__nv_bfloat16 const* src, long long srcStride, __nv_bfloat16 const* w, float eps, int H, + int C, int M, __nv_bfloat16* dst, bool pdl, cudaStream_t stream); + +//! dst[r, :H] = rmsnorm(bf16(h[r] + res[r])) * w (res optional), with chunk-norm rounding. +cudaError_t invokeCaptureNorm(__nv_bfloat16 const* h, long long hStride, __nv_bfloat16 const* res, long long resStride, + __nv_bfloat16 const* w, float eps, int H, int M, __nv_bfloat16* dst, long long dstStride, cudaStream_t stream); + +//! The overlap scheduler's per-step input gathers (see step_gather_kernel); work = the largest of its four loops. +cudaError_t invokeStepGather(int const* newTok, long long ntS0, long long ntS1, int const* nextDraft, long long ndS0, + long long ndS1, int const* lens, int const* slots, int const* ppi, int P, int T, int D, int* ids, int ids0, + int* draft, int d0, int* posOff, int nPos, int* kvOff, int nKv, int b0, int work, cudaStream_t stream); + +//! A strided int32 2-D view: element (i, c) at data[i * s0 + c * s1], `width` columns. +struct IntView +{ + int* data; + long long s0; + long long s1; + int width; +}; + +//! The spec sampler's slot stores (see slot_scatter_kernel): for i < n, da / db / dd rows slots[i] get sa / sb / sd +//! row i, zero-padded to the destination widths, and dl[slots[i]] = sl[i]. +cudaError_t invokeSlotScatter(long long const* slots, int n, IntView sa, IntView da, IntView sb, IntView db, + int const* sl, int* dl, IntView sd, IntView dd, cudaStream_t stream); + +} // namespace specTail +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/stepInputsScatter.cu b/cpp/tensorrt_llm/kernels/minimaxM3/stepInputsScatter.cu new file mode 100644 index 000000000000..2e80de17bfd9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/stepInputsScatter.cu @@ -0,0 +1,69 @@ +/* + * 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/minimaxM3/stepInputsScatter.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kThreads = 256; + +__device__ __forceinline__ void copyBytes(uint8_t* dst, uint8_t const* src, int64_t bytes) +{ + for (int64_t b = threadIdx.x; b < bytes; b += blockDim.x) + { + dst[b] = src[b]; + } +} + +// One CTA: the batches this serves are a step's few hundred bytes of metadata, where a single launch is the point and +// the copy itself is negligible. +__global__ void __launch_bounds__(kThreads) stepInputsScatterKernel(uint8_t const* __restrict__ stage) +{ + auto const* header = reinterpret_cast(stage); + int64_t const numHostCopies = header[0]; + int64_t const numDeviceCopies = header[1]; + int64_t const* desc = header + 2; + // Copies complete in descriptor order: one may overwrite or read another's destination. + for (int64_t i = 0; i < numHostCopies; ++i) + { + copyBytes(reinterpret_cast(desc[3 * i]), stage + desc[3 * i + 1], desc[3 * i + 2]); + __syncthreads(); + } + for (int64_t i = numHostCopies; i < numHostCopies + numDeviceCopies; ++i) + { + copyBytes(reinterpret_cast(desc[3 * i]), reinterpret_cast(desc[3 * i + 1]), + desc[3 * i + 2]); + __syncthreads(); + } +} + +} // namespace + +void invokeStepInputsScatter(uint8_t const* stage, cudaStream_t stream) +{ + stepInputsScatterKernel<<<1, kThreads, 0, stream>>>(stage); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/stepInputsScatter.h b/cpp/tensorrt_llm/kernels/minimaxM3/stepInputsScatter.h new file mode 100644 index 000000000000..353197a479f9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/stepInputsScatter.h @@ -0,0 +1,40 @@ +/* + * 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 +{ + +//! Write a packed batch of byte copies described by a staging buffer. +//! +//! The staging buffer starts with an int64 header `[numHostCopies, numDeviceCopies]`, followed by one +//! `(dst, src, bytes)` int64 triple per copy. A host copy's `src` is a byte offset into the staging buffer (its +//! bytes were uploaded with the header); a device copy's `src` is a device address. Every host copy is written +//! before any device copy, so a device copy may read a destination a host copy of the same batch filled. +void invokeStepInputsScatter(uint8_t const* stage, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/targetHead.cu b/cpp/tensorrt_llm/kernels/minimaxM3/targetHead.cu new file mode 100644 index 000000000000..847c59779a6e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/targetHead.cu @@ -0,0 +1,372 @@ +/* + * 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/minimaxM3/targetHead.h" + +#include +#include +#include + +// The MiniMax-M3 target lm_head shard GEMV of a verify step of at most 8 rows (GB300, TP4), at stock precision: BF16 +// weight and activation, mma.m16n8k16 BF16 with FP32 accumulation over the full K, one rounding to BF16 per logit +// (the cuBLAS head differs only in the FP32 summation order). The full logit row is produced. +// +// * Lossless 12-bit weight coding: the BF16 exponents of a trained weight sit almost all in a window of 15 +// consecutive values [base, base + 14]. Each weight is one byte sign << 7 | mantissa plus a 4-bit code +// exponent - base; the decoded BF16 bits are the checkpoint's. Code 15 marks an exception: it decodes to +0 (it +// adds nothing to the MMA sums), and the segment epilogue adds the exact product x * w from a per-tile list (a +// product of two BF16 values is exact in FP32), so every product enters the sum exactly as in the BF16 head. +// * Layout: the weight is permuted into mma A-fragment order (per 16-row tile and k16 step, lane (g, j) holds +// {row g k 2j..2j+1, row g+8 k 2j.., row g k 2j+8.., row g+8 k 2j+8..}); per unit of 4 k16 steps a lane reads +// 48 bytes: the sign/mantissa bytes of steps 0-1 and of steps 2-3, then 4 code words (even values' nibbles in the +// low half, odd values' in the high half, so the codes of value pair p are (word >> 4p) & 0x000f000f). +// * A persistent grid (one 16-warp CTA per SM) splits the (tile, unit) space into equal contiguous ranges, one per +// warp, each streamed through a 3-unit cp.async ring in shared memory. The first units are issued before the PDL +// grid-dependency wait (the weight does not depend on the previous kernel); the activation is staged in shared +// memory in B-fragment order after it. +// * A tile split between warps writes per-warp FP32 partials; the last finisher (per-tile counter) adds them in warp +// order, so the result is deterministic run to run. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3 +{ +namespace targetHead +{ +namespace +{ + +constexpr int kUnit = 4; // k16 steps per coded unit + +struct Args +{ + Params p; + int tiles; + int ks; // K / 16 + long long units; // tiles * ks / kUnit +}; + +__device__ __forceinline__ void cp_async16(void* smem, void const* gmem) +{ + uint32_t const sa = static_cast(__cvta_generic_to_shared(smem)); + asm volatile("cp.async.cg.shared.global.L2::256B [%0], [%1], 16;\n" ::"r"(sa), "l"(gmem) : "memory"); +} + +__device__ __forceinline__ void cp_async_commit() +{ + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} + +template +__device__ __forceinline__ void cp_async_wait() +{ + asm volatile("cp.async.wait_group %0;\n" ::"n"(N) : "memory"); +} + +__device__ __forceinline__ void mma_bf16(float (&d)[4], uint4 const& a, uint2 const& b) +{ + 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.x), "r"(a.y), "r"(a.z), "r"(a.w), "r"(b.x), "r"(b.y)); +} + +// Accumulator layout of one tile: c[0..1] = (row g, tokens 2j, 2j + 1), c[2..3] = (row g + 8, ...). +__device__ __forceinline__ void store_tile(Args const& a, int tile, int lane, float const (&c)[4]) +{ + int const g = lane >> 2, j = lane & 3; +#pragma unroll + for (int e = 0; e < 4; ++e) + { + int const n = 2 * j + (e & 1); + int const row = tile * 16 + g + (e >> 1) * 8; + if (n < a.p.M && row < a.p.N) + a.p.out[n * a.p.out_stride + row] = __float2bfloat16_rn(c[e]); + } +} + +// Flush the warp's sums of `tile` over k16 steps [k0, k1). +__device__ __forceinline__ void flush(Args const& a, int warp, int tile, int k0, int k1, int lane, float const (&c)[4]) +{ + if (k0 == 0 && k1 == a.ks) + { + store_tile(a, tile, lane, c); + return; + } + // Warps owning parts of this tile: the owner of its first unit up to the owner of its last one, where + // owner(u) = the largest w with w * units / W <= u. The first owner stores into its slot 1 (the tile is its last + // segment), every later owner into slot 0 (its first segment), so a warp never reuses a slot. + long long const W = static_cast(gridDim.x) * (blockDim.x >> 5); + long long const first = static_cast(tile) * (a.ks / kUnit); + auto owner = [&](long long u) + { + long long w = (u * W) / a.units; + while (w + 1 < W && ((w + 1) * a.units) / W <= u) + ++w; + while (w > 0 && (w * a.units) / W > u) + --w; + return static_cast(w); + }; + int const wf = owner(first); + int const wl = owner(first + a.ks / kUnit - 1); + constexpr int kVals = 4 * 32; + float* mine = a.p.part + (static_cast(warp) * 2 + (warp == wf ? 1 : 0)) * kVals; +#pragma unroll + for (int e = 0; e < 4; ++e) + mine[e * 32 + lane] = c[e]; + __threadfence(); + __syncwarp(); + int done = 0; + if (lane == 0) + done = atomicAdd(&a.p.counter[tile], 1) == wl - wf; + done = __shfl_sync(0xffffffffu, done, 0); + if (!done) + return; + __threadfence(); + float s[4] = {0.f, 0.f, 0.f, 0.f}; + for (int w = wf; w <= wl; ++w) + { + float const* q = a.p.part + (static_cast(w) * 2 + (w == wf ? 1 : 0)) * kVals; +#pragma unroll + for (int e = 0; e < 4; ++e) + s[e] += __ldcg(q + e * 32 + lane); + } + store_tile(a, tile, lane, s); + if (lane == 0) + a.p.counter[tile] = 0; +} + +// BF16 pair p (values 2p, 2p + 1) of a lane's 8 coded values: sign/mantissa bytes sm (4 per word), codes word cw. +__device__ __forceinline__ uint32_t decode_pair(uint32_t sm_word, int byte0, uint32_t cw, int p, uint32_t base) +{ + uint32_t const b = __byte_perm(sm_word, 0u, byte0 == 0 ? 0x4140 : 0x4342); // {b0, 0, b1, 0} + uint32_t const q = (cw >> (4 * p)) & 0x000f000fu; + uint32_t const e = q * 128u + base; // exponent field of both halves + // Code 15 (an exception, sign/mantissa byte 0) decodes to +0. + return ((b & 0x007f007fu) | ((b << 8) & 0x80008000u) | e) & __vcmpne2(q, 0x000f000fu); +} + +// WARPS warps per CTA, each with a cp.async ring of S units (1536 bytes each). The range is walked tile segment by +// tile segment, so the inner loop has no boundary checks; k16 steps alternate between two accumulators (added at the +// segment end). x is staged as [ks / 4][32][4] B fragments, two 16-byte shared loads per unit. +template +__global__ void __launch_bounds__(WARPS * 32, 1) coded_kernel(Args const a) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ uint4 smem4[]; + int const lane = threadIdx.x & 31; + int const wid = threadIdx.x >> 5; + int const warp = blockIdx.x * WARPS + wid; + long long const W = static_cast(gridDim.x) * WARPS; + long long const u0 = warp * a.units / W; + long long const u1 = (warp + 1) * a.units / W; + int const n_units = static_cast(u1 - u0); + int const upt = a.ks / kUnit; // units per tile + uint4 const* wp = a.p.w + u0 * 96 + lane; + uint4* const xs4 = smem4; // [upt][32][2] uint4 + uint4* const ring = smem4 + upt * 64 + wid * (S * 96) + lane; + +#pragma unroll + for (int u = 0; u < S; ++u) + { + if (u < n_units) +#pragma unroll + for (int l = 0; l < 3; ++l) + cp_async16(ring + (u * 3 + l) * 32, wp + (u * 3 + l) * 32); + cp_async_commit(); + } + asm volatile("griddepcontrol.wait;" ::: "memory"); + { + // x[n][8q..8q+7] -> word h = q & 1 of entries (s = q / 2, lane = n * 4 + j), entry index (s / 4, lane, s % 4). + uint32_t* xs32 = reinterpret_cast(xs4); + int const qn = a.ks * 2; + int const items = kMaxRows * qn; + constexpr int kSU = 8; + for (int i0 = threadIdx.x; i0 < items; i0 += kSU * blockDim.x) + { + uint4 v[kSU]; +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + int const n = i / qn, q = i - n * qn; + v[u] = (i < items && n < a.p.M) ? *reinterpret_cast(a.p.x + n * a.p.x_stride + q * 8) + : make_uint4(0u, 0u, 0u, 0u); + } +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + if (i >= items) + continue; + int const n = i / qn, q = i - n * qn; + int const sstep = q >> 1; + int const base = (((sstep >> 2) * 32 + n * 4) * 4 + (sstep & 3)) * 2 + (q & 1); + xs32[base + 0] = v[u].x; + xs32[base + 8] = v[u].y; + xs32[base + 16] = v[u].z; + xs32[base + 24] = v[u].w; + } + } + } + __syncthreads(); + asm volatile("griddepcontrol.launch_dependents;"); + + uint4 const* xl = xs4 + lane * 2; + int tile = static_cast(u0 / upt); + int ku = static_cast(u0 % upt); // unit within the tile + int done = 0, slot = 0; + while (done < n_units) + { + int const seg = min(n_units - done, upt - ku); + float ca[4] = {0.f, 0.f, 0.f, 0.f}; + float cb[4] = {0.f, 0.f, 0.f, 0.f}; + int const e0 = __ldg(a.p.exc_off + tile), e1 = __ldg(a.p.exc_off + tile + 1); + for (int i = 0; i < seg; ++i) + { + cp_async_wait(); + uint4 const sm01 = ring[slot * 96]; + uint4 const sm23 = ring[slot * 96 + 32]; + uint4 const cw = ring[slot * 96 + 64]; + uint4 const xa = xl[(ku + i) * 64]; + uint4 const xb = xl[(ku + i) * 64 + 1]; + uint32_t const smw[8] = {sm01.x, sm01.y, sm01.z, sm01.w, sm23.x, sm23.y, sm23.z, sm23.w}; + uint32_t const cws[4] = {cw.x, cw.y, cw.z, cw.w}; + uint2 const xb2[4] + = {make_uint2(xa.x, xa.y), make_uint2(xa.z, xa.w), make_uint2(xb.x, xb.y), make_uint2(xb.z, xb.w)}; +#pragma unroll + for (int ci = 0; ci < 4; ++ci) + { + uint4 wv; + wv.x = decode_pair(smw[2 * ci], 0, cws[ci], 0, a.p.exp_base); + wv.y = decode_pair(smw[2 * ci], 2, cws[ci], 1, a.p.exp_base); + wv.z = decode_pair(smw[2 * ci + 1], 0, cws[ci], 2, a.p.exp_base); + wv.w = decode_pair(smw[2 * ci + 1], 2, cws[ci], 3, a.p.exp_base); + if (ci & 1) + mma_bf16(cb, wv, xb2[ci]); + else + mma_bf16(ca, wv, xb2[ci]); + } + // Refill the slot only after its values were consumed (no ordering between ld.shared and cp.async). + int const nxt = done + i + S; + if (nxt < n_units) +#pragma unroll + for (int l = 0; l < 3; ++l) + cp_async16(ring + (slot * 3 + l) * 32, wp + (nxt * 3 + l) * 32); + cp_async_commit(); + slot = slot + 1 == S ? 0 : slot + 1; + } +#pragma unroll + for (int e = 0; e < 4; ++e) + ca[e] += cb[e]; + // Exceptions of this tile with k in the segment: lanes load 32 at a time, each is broadcast; the owner lanes + // (row g or g + 8) add the exact product x[n][k] * w with x from shared memory. + { + int const g = lane >> 2, j = lane & 3; + int const kb = ku * 64, ke = (ku + seg) * 64; + uint32_t const* xs32 = reinterpret_cast(xs4); + for (int b0 = e0; b0 < e1; b0 += 32) + { + uint32_t pos = 0xffffffffu; + float2 wv = make_float2(0.f, 0.f); + if (b0 + lane < e1) + { + pos = __ldg(a.p.exc_pos + b0 + lane); + wv = __ldg(a.p.exc_w + b0 + lane); + } + int const cnt = min(32, e1 - b0); + for (int q = 0; q < cnt; ++q) + { + uint32_t const pq = __shfl_sync(0xffffffffu, pos, q); + float const w0 = __shfl_sync(0xffffffffu, wv.x, q); + int const r = static_cast(pq >> 16), k = static_cast(pq & 0xffffu); + // No early continue: every lane reaches each iteration's shuffles together. + if (k >= kb && k < ke && (r & 7) == g) + { + int const st = k >> 4, rr = k & 15; + bool const hi = r >= 8; +#pragma unroll + for (int q2 = 0; q2 < 2; ++q2) + { + int const n = 2 * j + q2; + uint32_t const word + = xs32[((((st >> 2) * 32 + n * 4 + ((rr & 7) >> 1)) * 4 + (st & 3)) * 2) + (rr >> 3)]; + float const xv = __uint_as_float((rr & 1) ? (word & 0xffff0000u) : (word << 16)); + float const d = __fmul_rn(xv, w0); + ca[q2] += hi ? 0.f : d; + ca[2 + q2] += hi ? d : 0.f; + } + } + __syncwarp(); + } + } + } + flush(a, warp, tile, ku * kUnit, (ku + seg) * kUnit, lane, ca); + done += seg; + ++tile; + ku = 0; + } +#endif +} + +} // namespace + +size_t targetHeadSmemBytes(int K) +{ + return static_cast(K / 64) * 64 * 16 + static_cast(kWarps) * kStages * 96 * 16; +} + +cudaError_t invokeTargetHead(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + if (p.M < 1 || p.M > kMaxRows || K % 64 != 0 || p.N < 1 || grid < 1) + return cudaErrorInvalidValue; + Args a; + a.p = p; + a.tiles = (p.N + 15) / 16; + a.ks = K / 16; + a.units = static_cast(a.tiles) * a.ks / kUnit; + // Every warp must own at least one unit: a split tile's finisher counts the warps between its first and last + // owner, so a warp with an empty range in between would leave the tile unstored. + if (a.units < kWarps) + return cudaErrorInvalidValue; + grid = static_cast(std::min(grid, a.units / kWarps)); + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(static_cast(grid)); + cfg.blockDim = dim3(kWarps * 32); + cfg.dynamicSmemBytes = targetHeadSmemBytes(K); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + auto kern = coded_kernel; + cudaError_t const err = cudaFuncSetAttribute( + kern, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(cfg.dynamicSmemBytes)); + if (err != cudaSuccess) + return err; + return cudaLaunchKernelEx(&cfg, kern, a); +} + +} // namespace targetHead +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/targetHead.h b/cpp/tensorrt_llm/kernels/minimaxM3/targetHead.h new file mode 100644 index 000000000000..84b68ee2152a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/targetHead.h @@ -0,0 +1,68 @@ +/* + * 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 +{ +namespace minimaxM3 +{ +//! The MiniMax-M3 target lm_head shard GEMV over a lossless 12-bit coding of its BF16 weight (verify rows <= 8, +//! sm_103a): BF16 x BF16 MMAs with FP32 accumulation over the full K, one BF16 rounding per logit. See targetHead.cu. +namespace targetHead +{ + +constexpr int kMaxRows = 8; // activation rows (verify rows) of one launch +constexpr int kWarps = 16; // warps per CTA +constexpr int kStages = 3; // coded units in flight per warp (cp.async ring) + +struct Params +{ + uint4 const* w; // [tiles][K / 64][3][32] coded weight (see targetHead.cu) + __nv_bfloat16 const* x; // [M, K] activation rows + long long x_stride; // elements between rows of x + __nv_bfloat16* out; // [M, N] logits + long long out_stride; // elements between rows of out + float* part; // [grid * kWarps][2][4 * 32] partial sums of tiles split between warps + int* counter; // [tiles] finishers per split tile, zero between calls (re-armed by the kernel) + int const* exc_off; // [tiles + 1] exception list offsets + uint32_t const* exc_pos; // row in tile << 16 | k, per exception + float2 const* exc_w; // {w, 0} per exception: the exact weight + uint32_t exp_base; // exponent base << 7 in both BF16 halves (base << 7 | base << 23) + int M; // rows, 1..kMaxRows + int N; // logits per row (shard rows) +}; + +//! Shared memory a launch over hidden size K takes (the staged activation and the cp.async rings). +size_t targetHeadSmemBytes(int K); + +//! Launch over `grid` 512-thread CTAs (K a multiple of 64). +cudaError_t invokeTargetHead(Params const& p, int K, int grid, bool pdl, cudaStream_t stream); + +} // namespace targetHead +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.cu b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.cu new file mode 100644 index 000000000000..dfc570239a5c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.cu @@ -0,0 +1,770 @@ +/* + * 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. + */ + +// MiniMax-M3 TP8 fused index score + top-16 block select in one launch. +// +// Scoring: grid (splits, batch), one CTA per SM. CTA (s, b) owns blocks s, s + splits, ... of request b and TMA-loads +// their index-K pages through a stage ring (one producer warp, 8 consumer warps; each consumer warp scores a 32-row +// slice of a page against all q columns with m16n8k16 FP16 MMAs, column = head * 4 + query row). Block scores go to a +// row-major scratch; the last CTA of the request (arrival counter) bulk-copies the request's rows back into its idle +// stage ring and selects the top 16 per (query row, KV head): a threshold bound from per-thread maxima, a candidate +// list, exact ranking by (key desc, block asc), and an exact radix select when ties overflow the candidate list. + +#include "tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.h" + +#include "tensorrt_llm/common/cudaDriverWrapper.h" + +#include +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kTopK = 16; +constexpr int kBlockK = 128; +constexpr int kHeadDim = 128; +constexpr int kStageBytes = kBlockK * kHeadDim; // fp8 +constexpr int kConsumerWarps = 8; +constexpr int kConsumerThreads = kConsumerWarps * 32; +constexpr int kThreads = kConsumerThreads + 32; // + producer warp +constexpr int kMaxDql = 4; // B-operand column = head * kMaxDql + query row +constexpr int kCols = 8; // one n8 tile: heads * kMaxDql <= 8 +constexpr int kMaxRows = 8; // kv_heads * dql +constexpr int kMaxLocal = 128; // blocks per CTA +constexpr int kCandCap = 256; +constexpr int kEpiBytes = kMaxLocal * 4 * kCols * 4; +constexpr float kInitScore = 1.0e30F; +constexpr float kLocalScore = 1.0e29F; +constexpr int kNStage = 12; // TMA stage ring depth (pages in flight per CTA) + +using Params = MinimaxM3Tp8FusedScoreSelectParams; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +__device__ __forceinline__ uint32_t twiddle(float x) +{ + uint32_t const bits = __float_as_uint(x); + return bits ^ ((bits & 0x80000000u) ? 0xFFFFFFFFu : 0x80000000u); +} + +__device__ __forceinline__ uint32_t cvt_e4m3x2(uint32_t pair16) +{ + uint32_t out; + unsigned short const in = static_cast(pair16 & 0xFFFFu); + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(out) : "h"(in)); + return out; +} + +__device__ __forceinline__ void mma16816( + float (&c)[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};" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void mbar_init(uint32_t bar, int count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(bar), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbar_expect_tx(uint32_t bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(bar), "r"(bytes) : "memory"); +} + +__device__ __forceinline__ void mbar_arrive(uint32_t bar) +{ + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(bar) : "memory"); +} + +__device__ __forceinline__ void mbar_wait(uint32_t bar, uint32_t parity) +{ + uint32_t done; + do + { + asm volatile( + "{\n\t.reg .pred p;\n\tmbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n\tselp.u32 %0, 1, 0, p;\n\t}" + : "=r"(done) + : "r"(bar), "r"(parity) + : "memory"); + } while (!done); +} + +__device__ __forceinline__ void tma_load_page( + uint32_t dst, CUtensorMap const* tmap, int page, uint32_t bar, uint64_t policy) +{ + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3, %4}], [%5], %6;" ::"r"(dst), + "l"(reinterpret_cast(tmap)), "r"(0), "r"(0), "r"(page), "r"(bar), "l"(policy) + : "memory"); +} + +// Exact radix select over n smem keys (fallback when the threshold bound leaves too many candidates). +__device__ void radix_emit(uint32_t const* kr, int n, uint32_t* hist, int32_t* dst, uint32_t neg_inf_key) +{ + int const lane = threadIdx.x & 31; + uint32_t prefix = 0, mask = 0; + int need = kTopK; + bool exact_bin = false; + for (int shift = 24; shift >= 0; shift -= 8) + { + for (int i = lane; i < 256; i += 32) + hist[i] = 0; + __syncwarp(); + for (int i = lane; i < n; i += 32) + { + uint32_t const k = kr[i]; + if ((k & mask) == prefix) + atomicAdd(&hist[(k >> shift) & 255u], 1u); + } + __syncwarp(); + uint32_t c[8]; + uint32_t sum = 0; +#pragma unroll + for (int j = 0; j < 8; ++j) + { + c[j] = hist[255 - 8 * lane - j]; + sum += c[j]; + } + uint32_t incl = sum; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) + { + uint32_t const v = __shfl_up_sync(0xFFFFFFFFu, incl, o); + if (lane >= o) + incl += v; + } + uint32_t above = incl - sum; + int found_bin = -1; + uint32_t found_above = 0, found_cnt = 0; +#pragma unroll + for (int j = 0; j < 8; ++j) + { + if (found_bin < 0 && above < static_cast(need) && above + c[j] >= static_cast(need)) + { + found_bin = 255 - 8 * lane - j; + found_above = above; + found_cnt = c[j]; + } + above += c[j]; + } + unsigned const bal = __ballot_sync(0xFFFFFFFFu, found_bin >= 0); + int const src = __ffs(bal) - 1; + int const bin = __shfl_sync(0xFFFFFFFFu, found_bin, src); + uint32_t const fa = __shfl_sync(0xFFFFFFFFu, found_above, src); + uint32_t const fc = __shfl_sync(0xFFFFFFFFu, found_cnt, src); + __syncwarp(); + need -= static_cast(fa); + prefix |= static_cast(bin) << shift; + mask |= 255u << shift; + if (fc == static_cast(need)) + { + exact_bin = true; + break; + } + } + int emitted = 0, ties_used = 0; + uint32_t const lt = (1u << lane) - 1u; + for (int c0 = 0; c0 < n; c0 += 32) + { + int const i = c0 + lane; + bool const valid = i < n; + uint32_t const k = valid ? kr[i] : 0u; + uint32_t const km = k & mask; + bool const neg = valid && k == neg_inf_key; + bool sel = valid && (km > prefix || (exact_bin && km == prefix)); + bool const tie = valid && !exact_bin && km == prefix; + unsigned const tie_bal = __ballot_sync(0xFFFFFFFFu, tie); + if (tie) + sel = ties_used + __popc(tie_bal & lt) < need; + ties_used += __popc(tie_bal); + bool const out_ok = sel && !neg; + unsigned const ob = __ballot_sync(0xFFFFFFFFu, out_ok); + if (out_ok) + dst[emitted + __popc(ob & lt)] = i; + emitted += __popc(ob); + } + if (lane >= emitted && lane < kTopK) + dst[lane] = -1; + __syncwarp(); +} + +// Descending bitonic sort of one value per lane (lane 0 ends with the largest). +__device__ __forceinline__ uint32_t warp_sort_desc(uint32_t v, int lane) +{ +#pragma unroll + for (int k = 2; k <= 32; k <<= 1) + { +#pragma unroll + for (int j = k >> 1; j > 0; j >>= 1) + { + uint32_t const o = __shfl_xor_sync(0xFFFFFFFFu, v, j); + bool const desc_block = (lane & k) == 0; + bool const low = (lane & j) == 0; + v = (low == desc_block) ? max(v, o) : min(v, o); + } + } + return v; +} + +__device__ __forceinline__ uint32_t forced_key(float sc, int i, int init_blocks, int local_start) +{ + if (i < init_blocks) + sc = kInitScore; + if (i >= local_start) + sc = kLocalScore; + return twiddle(sc); +} + +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +template +__global__ void __launch_bounds__(kThreads, 1) + m3_fused_score_select_kernel(__grid_constant__ CUtensorMap const tmap, Params p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ uint8_t smem_raw[]; + uint8_t* smem = smem_raw + ((1024u - (smem_u32(smem_raw) & 1023u)) & 1023u); + uint8_t* stages = smem; + float* epi = reinterpret_cast(smem + NSTAGE * kStageBytes); // [kMaxLocal][4 slices][kCols] + uint64_t* bars + = reinterpret_cast(smem + NSTAGE * kStageBytes + kEpiBytes); // full[NSTAGE], empty[NSTAGE] + __shared__ uint32_t s_wtop[kMaxRows][kConsumerWarps][16]; + __shared__ int s_cnt[kMaxRows]; + __shared__ int32_t s_sel_blk[kMaxRows][kTopK]; + __shared__ int s_sel_neg[kMaxRows][kTopK]; + __shared__ int s_last; + __shared__ int32_t s_sb[kMaxRows][32]; + __shared__ __align__(8) uint64_t s_selbar; + + int const s = blockIdx.x, b = blockIdx.y, splits = gridDim.x; + int const tid = threadIdx.x, warp = tid >> 5, lane = tid & 31; + int const g = lane >> 2, tq = lane & 3; + uint32_t const full0 = smem_u32(bars), empty0 = smem_u32(bars + NSTAGE); + + if (tid == 0) + { + for (int i = 0; i < NSTAGE; ++i) + { + mbar_init(full0 + 8 * i, 1); + mbar_init(empty0 + 8 * i, 4); + } + mbar_init(smem_u32(&s_selbar), 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + if (warp == kConsumerWarps && lane == 0) + asm volatile("prefetch.tensormap [%0];" ::"l"(reinterpret_cast(&tmap)) : "memory"); + __syncthreads(); + + // The producer lanes fetch their first page ids together with seq_lens (bounded by the table width), so a CTA + // pays one dependent round trip before its loads go out. That happens, and the loads of the pages that hold only + // tokens older than this step go out, before the grid dependency wait. + int pid_first = 0; + int seqlen = 0; + uint64_t policy = 0; + if (warp == kConsumerWarps) + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" : "=l"(policy)); + int issued = 0; // producer lane: its first page is already in flight + { + if (warp == kConsumerWarps && lane < NSTAGE) + { + int const blk = s + lane * splits; + if (blk < p.bt_width) + pid_first = __ldg(p.bt + static_cast(b) * p.bt_stride + blk); + } + seqlen = p.seq_lens[b]; + if (warp == kConsumerWarps) + { + int const nb0 = seqlen > 0 ? min((seqlen + kBlockK - 1) / kBlockK, p.max_tiles) : 0; + int const my0 = s < nb0 ? (nb0 - s + splits - 1) / splits : 0; + int const stable = max(seqlen - p.dql, 0) / kBlockK; // blocks holding only older tokens + int const blk = s + lane * splits; + if (lane < min(my0, NSTAGE) && blk < stable) + { + mbar_expect_tx(full0 + 8 * lane, kStageBytes); + tma_load_page(smem_u32(stages + lane * kStageBytes), &tmap, pid_first, full0 + 8 * lane, policy); + issued = 1; + } + } + } + + asm volatile("griddepcontrol.wait;" ::: "memory"); + + // The select's per-row valid-block count, loaded by every CTA alongside seq_lens (only the last one uses it). + int const sel_tpr = p.dql * p.kv_heads <= 4 ? kConsumerThreads / 4 : kConsumerThreads / 8; + int const nv_raw = tid < kConsumerThreads && tid / sel_tpr < p.dql * p.kv_heads + ? p.n_valid[b * p.dql + (tid / sel_tpr) % p.dql] + : 0; + int const nb = seqlen > 0 ? min((seqlen + kBlockK - 1) / kBlockK, p.max_tiles) : 0; + if (nb == 0) + { + // Nothing to score: every row of the request is all -1. + if (s == 0 && tid < p.dql * p.kv_heads * kTopK) + { + int const r = tid / kTopK, j = tid % kTopK; + int const t = b * p.dql + r % p.dql, kvh = r / p.dql; + p.out[static_cast(kvh) * p.out_h + static_cast(t) * p.out_t + j] = -1; + } + return; + } + if (s >= nb) + return; + int const parts = min(splits, nb); + int const my = (nb - s + splits - 1) / splits; + int const q_start = seqlen - p.dql; + int const group = p.heads / p.kv_heads; + int const rows = p.dql * p.kv_heads; + + if (warp == kConsumerWarps) + { + if (lane < min(my, NSTAGE) && !issued) + { + mbar_expect_tx(full0 + 8 * lane, kStageBytes); + tma_load_page(smem_u32(stages + lane * kStageBytes), &tmap, pid_first, full0 + 8 * lane, policy); + } + if (lane == 0) + { + for (int j = NSTAGE; j < my; ++j) + { + int const st = j % NSTAGE; + int const page = __ldg(p.bt + static_cast(b) * p.bt_stride + s + j * splits); + mbar_wait(empty0 + 8 * st, ((j / NSTAGE) - 1) & 1); + mbar_expect_tx(full0 + 8 * st, kStageBytes); + tma_load_page(smem_u32(stages + st * kStageBytes), &tmap, page, full0 + 8 * st, policy); + } + } + } + else + { + // Q fragments: column g is (head = g / kMaxDql, row = g % kMaxDql). + uint32_t qlo[4][2], qhi[4][2]; + { + int const head = g / kMaxDql, qpos = g % kMaxDql; + bool const valid = head < p.heads && qpos < p.dql; + uint8_t const* qrow = p.q + static_cast(b * p.dql + qpos) * p.q_st + head * p.q_sh; + uint32_t w0[4], w1[4]; +#pragma unroll + for (int k = 0; k < 4; ++k) + { + w0[k] = valid ? *reinterpret_cast(qrow + 32 * k + 4 * tq) : 0u; + w1[k] = valid ? *reinterpret_cast(qrow + 32 * k + 16 + 4 * tq) : 0u; + } +#pragma unroll + for (int k = 0; k < 4; ++k) + { + qlo[k][0] = cvt_e4m3x2(w0[k]); + qhi[k][0] = cvt_e4m3x2(w0[k] >> 16); + qlo[k][1] = cvt_e4m3x2(w1[k]); + qhi[k][1] = cvt_e4m3x2(w1[k] >> 16); + } + } + int const slice = warp & 3; +#pragma unroll 1 + for (int lb = warp >> 2; lb < my; lb += kConsumerWarps / 4) + { + int const st = lb % NSTAGE; + mbar_wait(full0 + 8 * st, (lb / NSTAGE) & 1); + uint8_t const* stg = stages + st * kStageBytes; + float acc[2][4]; +#pragma unroll + for (int m = 0; m < 2; ++m) +#pragma unroll + for (int e = 0; e < 4; ++e) + acc[m][e] = 0.f; +#pragma unroll + for (int k = 0; k < 4; ++k) + { +#pragma unroll + for (int m = 0; m < 2; ++m) + { + int const ra = 32 * slice + 16 * m + g; // rb = ra + 8 has the same (row & 7) + int const sw = ra & 7; + uint8_t const* pa = stg + ra * 128; + uint8_t const* pb = pa + 8 * 128; + uint32_t const d0 = *reinterpret_cast(pa + (((2 * k) ^ sw) << 4) + 4 * tq); + uint32_t const d1 = *reinterpret_cast(pb + (((2 * k) ^ sw) << 4) + 4 * tq); + uint32_t const d2 = *reinterpret_cast(pa + (((2 * k + 1) ^ sw) << 4) + 4 * tq); + uint32_t const d3 = *reinterpret_cast(pb + (((2 * k + 1) ^ sw) << 4) + 4 * tq); + uint32_t const l0 = cvt_e4m3x2(d0), l1 = cvt_e4m3x2(d1), l2 = cvt_e4m3x2(d2), l3 = cvt_e4m3x2(d3); + uint32_t const h0 = cvt_e4m3x2(d0 >> 16), h1 = cvt_e4m3x2(d1 >> 16), h2 = cvt_e4m3x2(d2 >> 16), + h3 = cvt_e4m3x2(d3 >> 16); + mma16816(acc[m], l0, l1, l2, l3, qlo[k][0], qlo[k][1]); + mma16816(acc[m], h0, h1, h2, h3, qhi[k][0], qhi[k][1]); + } + } + __syncwarp(); + if (lane == 0) + mbar_arrive(empty0 + 8 * st); + // Causal mask + max over this warp's 32 positions. + int const blk = s + lb * splits; + int const pos0 = blk * kBlockK + 32 * slice; +#pragma unroll + for (int j = 0; j < 2; ++j) + { + int const col = 2 * tq + j; + int const qabs = q_start + col % kMaxDql; + float v = -INFINITY; +#pragma unroll + for (int m = 0; m < 2; ++m) + { + float const a = qabs >= pos0 + 16 * m + g ? acc[m][j] : -INFINITY; + float const c = qabs >= pos0 + 16 * m + g + 8 ? acc[m][2 + j] : -INFINITY; + v = fmaxf(v, fmaxf(a, c)); + } +#pragma unroll + for (int o = 4; o < 32; o <<= 1) + v = fmaxf(v, __shfl_xor_sync(0xFFFFFFFFu, v, o)); + if (g == 0) + epi[(lb * 4 + slice) * kCols + col] = v; + } + } + } + __syncthreads(); + + // Per block: max over the 4 slices, amax over each KV head's index heads, row-major scratch stores. + for (int lb = tid; lb < my; lb += kThreads) + { + int const blk = s + lb * splits; + float v[kCols]; +#pragma unroll + for (int c = 0; c < kCols; ++c) + { + float x = epi[(lb * 4) * kCols + c]; +#pragma unroll + for (int sl = 1; sl < 4; ++sl) + x = fmaxf(x, epi[(lb * 4 + sl) * kCols + c]); + v[c] = x; + } + float* dst = p.scr + b * p.scr_b + blk; +#pragma unroll + for (int r = 0; r < kMaxRows; ++r) + { + if (r < rows) + { + int const kvh = r / p.dql, qpos = r % p.dql; + float x = -INFINITY; +#pragma unroll + for (int c = 0; c < kCols; ++c) + if (c / kMaxDql >= kvh * group && c / kMaxDql < (kvh + 1) * group && c % kMaxDql == qpos) + x = fmaxf(x, v[c]); + dst[r * p.scr_r] = x; + } + } + } + __syncthreads(); + + __syncthreads(); + if (tid == 0) + { + int old; + asm volatile("atom.add.acq_rel.gpu.s32 %0, [%1], 1;" : "=r"(old) : "l"(p.counters + b) : "memory"); + s_last = old == parts - 1; + } + __syncthreads(); + if (!s_last) + return; + asm volatile("griddepcontrol.launch_dependents;"); + + // ---- last CTA of request b: top-16 of every (query row, KV head) ---- + // Row r is owned by tpr threads (2 warps for <= 4 rows, 1 warp for <= 8); thread rt of the row takes blocks + // rt, rt + tpr, ... and stages their forced keys in the (now idle) stage ring. + if (tid == 0) + p.counters[b] = 0; + if (tid < kMaxRows) + s_cnt[tid] = 0; + int const tpr = sel_tpr; + int const r = tid / tpr, rt = tid % tpr; + bool const active = tid < kConsumerThreads && r < rows; + int const ldk = (p.max_tiles + 3) & ~3; + uint32_t* keys = reinterpret_cast(stages) + (active ? r : 0) * ldk; + // One bulk copy per row brings the request's block scores (blocks [0, nb), contiguous per row) into the stage + // ring, so the whole row arrives in a single L2 round trip. + uint32_t const selbar = smem_u32(&s_selbar); + if (tid == 0) + { + asm volatile("fence.proxy.async.global;" ::: "memory"); + uint32_t const bytes = static_cast((nb + 3) & ~3) * 4u; + mbar_expect_tx(selbar, bytes * rows); + for (int rr = 0; rr < rows; ++rr) + { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + smem_u32(stages) + static_cast(rr * ldk) * 4u), + "l"(p.scr + b * p.scr_b + rr * p.scr_r), "r"(bytes), "r"(selbar) + : "memory"); + } + } + int n = 0, lst = 0; + uint32_t mx = 0u; + if (active) + { + n = max(0, min(nv_raw, p.max_tiles)); + lst = max(nv_raw - p.local_blocks, 0); + mbar_wait(selbar, 0); + float const* rowf = reinterpret_cast(keys); +#pragma unroll 1 + for (int i0 = rt; i0 < n; i0 += 8 * tpr) + { + float x[8]; +#pragma unroll + for (int u = 0; u < 8; ++u) + x[u] = i0 + u * tpr < n ? rowf[i0 + u * tpr] : 0.f; +#pragma unroll + for (int u = 0; u < 8; ++u) + { + int const i = i0 + u * tpr; + if (i < n) + { + uint32_t const k = forced_key(x[u], i, p.init_blocks, lst); + keys[i] = k; + mx = max(mx, k); + } + } + } + } + // T0 = 16th largest thread maximum of the row: 16 threads each hold a key >= T0, so every top-16 key (ties + // included) is >= T0 (a row with fewer than 16 keys gets T0 = 0, i.e. all of them). + uint32_t const sorted = warp_sort_desc(mx, lane); + uint32_t t0 = __shfl_sync(0xFFFFFFFFu, sorted, kTopK - 1); + if (tpr == 64 && active && lane < kTopK) + s_wtop[r][rt >> 5][lane] = sorted; + __syncthreads(); + uint32_t* cand_key = reinterpret_cast(epi); // [rows][kCandCap] + int32_t* cand_blk = reinterpret_cast(epi) + kMaxRows * kCandCap; + if (active) + { + if (tpr == 64) + { + // 16th largest of two descending 16-lists A, B: max over i of min(A[i-1], B[15-i]), A[-1] = B[-1] = inf. + uint32_t const* A = s_wtop[r][0]; + uint32_t const* B = s_wtop[r][1]; + uint32_t x = 0u; + if (lane <= kTopK) + x = min(lane > 0 ? A[lane - 1] : 0xFFFFFFFFu, lane < kTopK ? B[kTopK - 1 - lane] : 0xFFFFFFFFu); + t0 = __reduce_max_sync(0xFFFFFFFFu, x); + } +#pragma unroll 1 + for (int i0 = rt; i0 < n; i0 += 8 * tpr) + { + uint32_t kk[8]; +#pragma unroll + for (int u = 0; u < 8; ++u) + kk[u] = i0 + u * tpr < n ? keys[i0 + u * tpr] : 0u; + uint32_t hit = 0; +#pragma unroll + for (int u = 0; u < 8; ++u) + hit |= static_cast(kk[u] >= t0 && kk[u] != 0u) << u; + if (hit) + { +#pragma unroll + for (int u = 0; u < 8; ++u) + { + if ((hit >> u) & 1u) + { + int const at = atomicAdd(&s_cnt[r], 1); + if (at < kCandCap) + { + cand_key[r * kCandCap + at] = kk[u]; + cand_blk[r * kCandCap + at] = i0 + u * tpr; + } + } + } + } + } + } + __syncthreads(); + // One warp per row: exact rank by (key desc, block asc), drop -inf, emit ascending ids with a -1 tail. + if (warp < rows) + { + int const rr = warp; + int const t = b * p.dql + rr % p.dql, kvh = rr / p.dql; + int32_t* dst = p.out + static_cast(kvh) * p.out_h + static_cast(t) * p.out_t; + int const c = s_cnt[rr]; + uint32_t const neg_key = twiddle(-INFINITY); + if (c <= 32) + { + // Lane j ranks candidate j against all of them (smem broadcasts), then places the selected ids by + // counting the smaller selected ids. + uint32_t const* ck = cand_key + rr * kCandCap; + int32_t const* cb = cand_blk + rr * kCandCap; + uint32_t const kj = lane < c ? ck[lane] : 0u; + int const bj = lane < c ? cb[lane] : 0x7FFFFFFF; + int rank = 0; +#pragma unroll 8 + for (int o = 0; o < 32; ++o) + { + uint32_t const ko = ck[o]; + int const bo = cb[o]; + rank += (o < c) & ((ko > kj) | ((ko == kj) & (bo < bj))); + } + bool const sel = lane < c && rank < kTopK && kj != neg_key; + int const sb = sel ? bj : 0x7FFFFFFF; + s_sb[rr][lane] = sb; + __syncwarp(); + int pos = 0; +#pragma unroll 8 + for (int o = 0; o < 32; ++o) + pos += s_sb[rr][o] < sb; + if (sel) + dst[pos] = bj; + int const cnt = __popc(__ballot_sync(0xFFFFFFFFu, sel)); + if (lane >= cnt && lane < kTopK) + dst[lane] = -1; + } + else if (c <= kCandCap) + { + for (int j = lane; j < c; j += 32) + { + uint32_t const kj = cand_key[rr * kCandCap + j]; + int const bj = cand_blk[rr * kCandCap + j]; + int rank = 0; + for (int o = 0; o < c; ++o) + { + uint32_t const ko = cand_key[rr * kCandCap + o]; + rank += (ko > kj) || (ko == kj && cand_blk[rr * kCandCap + o] < bj); + } + if (rank < kTopK) + { + s_sel_blk[rr][rank] = bj; + s_sel_neg[rr][rank] = kj == neg_key; + } + } + __syncwarp(); + bool const has = lane < kTopK && !s_sel_neg[rr][lane & 15]; + int const sk = has ? s_sel_blk[rr][lane & 15] : 0x7FFFFFFF; + int pos = 0; +#pragma unroll + for (int l = 0; l < kTopK; ++l) + { + int const o = __shfl_sync(0xFFFFFFFFu, sk, l); + pos += (o < sk) || (o == sk && l < lane); + } + if (lane < kTopK) + dst[pos] = sk == 0x7FFFFFFF ? -1 : sk; + } + else + { + // Massive exact ties: exact radix select over the staged keys. + int const nrr = max(0, min(__ldcg(p.n_valid + t), p.max_tiles)); + uint32_t const* kr = reinterpret_cast(stages) + rr * ldk; + uint32_t* hist = reinterpret_cast(stages) + rows * ldk + rr * 256; + radix_emit(kr, nrr, hist, dst, neg_key); + } + } +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +// TMA descriptor of the [pages, 128, 128] e4m3 index-K pool (one page per box, 128-byte swizzle), cached per +// (base, pages, strides). +bool make_page_map(CUtensorMap& map, void const* base, int64_t pages, int64_t stride_row, int64_t stride_page) +{ + using Key = std::tuple; + static std::mutex mu; + static std::map cache; + Key const key{reinterpret_cast(base), pages, stride_row, stride_page}; + std::lock_guard lock(mu); + auto it = cache.find(key); + if (it != cache.end()) + { + map = it->second; + return true; + } + cuuint64_t dims[3] + = {static_cast(kHeadDim), static_cast(kBlockK), static_cast(pages)}; + cuuint64_t strides[2] = {static_cast(stride_row), static_cast(stride_page)}; + cuuint32_t box[3] = {kHeadDim, kBlockK, 1}; + cuuint32_t estr[3] = {1, 1, 1}; + CUresult const res = tensorrt_llm::common::CUDADriverWrapper::getInstance()->cuTensorMapEncodeTiled(&map, + CU_TENSOR_MAP_DATA_TYPE_UINT8, 3, const_cast(base), dims, strides, box, estr, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_L2_256B, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + if (res != CUDA_SUCCESS) + { + return false; + } + if (cache.size() > 64) + { + cache.clear(); + } + cache.emplace(key, map); + return true; +} + +constexpr size_t kSmemBytes = 1024 + kNStage * kStageBytes + kEpiBytes + 2 * kNStage * 8; + +cudaError_t configure() +{ + return cudaFuncSetAttribute(m3_fused_score_select_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(kSmemBytes)); +} + +} // namespace + +cudaError_t invokeMinimaxM3Tp8FusedScoreSelect(MinimaxM3Tp8FusedScoreSelectParams const& params, void const* kPages, + int64_t numPages, int64_t kRowStride, int64_t kPageStride, int32_t batch, int32_t splits, bool usePdl, + cudaStream_t stream) +{ + if (batch < 1 || batch > kMinimaxM3Tp8FusedScoreSelectMaxBatch || splits < 1 + || (params.max_tiles + splits - 1) / splits > kMaxLocal) + { + return cudaErrorInvalidValue; + } + CUtensorMap map; + if (!make_page_map(map, kPages, numPages, kRowStride, kPageStride)) + { + return cudaErrorInvalidValue; + } + // Kernel attributes are set once per process, on the first launch (outside CUDA-graph capture: warmup runs the + // kernel eagerly first). + static cudaError_t const configured = configure(); + if (configured != cudaSuccess) + { + return configured; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(splits, batch); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = kSmemBytes; + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = usePdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, m3_fused_score_select_kernel, map, params); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.h b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.h new file mode 100644 index 000000000000..ad92b71bac21 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.h @@ -0,0 +1,77 @@ +/* + * 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 +{ + +//! Largest number of requests (decode spans) one fused score + select launch takes. +constexpr int32_t kMinimaxM3Tp8FusedScoreSelectMaxBatch = 2; +//! Largest decode query length (query rows per request). +constexpr int32_t kMinimaxM3Tp8FusedScoreSelectMaxDql = 4; +//! Largest number of selected rows per request (num_kv_heads * dql). +constexpr int32_t kMinimaxM3Tp8FusedScoreSelectMaxRows = 8; +//! Largest number of blocks one CTA scores (ceil(max_tiles / splits) must not exceed it). +constexpr int32_t kMinimaxM3Tp8FusedScoreSelectMaxLocalBlocks = 128; +//! Bytes of the CTA's TMA stage ring, reused by the select: rows * ceil4(max_tiles) * 4 + rows * 1024 must fit. +constexpr int64_t kMinimaxM3Tp8FusedScoreSelectRingBytes = 12 * 128 * 128; + +//! Operands of the MiniMax-M3 TP8 fused index score + top-16 block select. Strides are in elements. +struct MinimaxM3Tp8FusedScoreSelectParams +{ + uint8_t const* q; // e4m3 [total_q, heads, 128] + int64_t q_st, q_sh; + int32_t const* bt; // [batch, bt_width] + int64_t bt_stride; + int bt_width; + int32_t const* seq_lens; // [batch], including this step's dql query rows + int32_t const* n_valid; // [total_q] causal valid-block counts + float* scr; // [batch][rows][>= ceil4(max_tiles)] block scores + int64_t scr_b, scr_r; + int32_t* out; // logical [total_q, kv_heads, 16] + int64_t out_t, out_h; + int32_t* counters; // [batch], zero at rest + int heads, kv_heads, dql, max_tiles, init_blocks, local_blocks; +}; + +//! One launch of index scoring and top-16 selection for at most 2 decode requests. Grid (splits, batch), one CTA of +//! 8 consumer warps + 1 TMA producer warp per SM: CTA (s, b) TMA-loads the e4m3 [128, 128] index-K pages of blocks +//! s, s + splits, ... of request b (12-stage ring; the pages holding only tokens older than this step are requested +//! before the grid dependency wait), scores them against the request's q rows on FP16 tensor cores, and stores each +//! block's score (max over its causal positions, amax over each KV head's index heads) to the scratch. The last CTA +//! of a request (per-request arrival counter) releases PDL dependents, resets that counter to zero, and selects the +//! top 16 of every (query row, KV head) row exactly like minimax_m3_select_blocks (init / local forcing, ties to the +//! smaller block, -inf picks dropped, ascending ids with a -1 tail). +//! +//! kPages is the e4m3 index-K page pool [numPages, 128, 128] with the given byte strides (row stride and page stride +//! multiples of 16, base 16-byte aligned); its TMA descriptor is encoded on first use and cached per (base, pages, +//! strides). The device code is SM100-family only; on other architectures the kernel compiles to an empty body. +cudaError_t invokeMinimaxM3Tp8FusedScoreSelect(MinimaxM3Tp8FusedScoreSelectParams const& params, void const* kPages, + int64_t numPages, int64_t kRowStride, int64_t kPageStride, int32_t batch, int32_t splits, bool usePdl, + cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.cu b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.cu new file mode 100644 index 000000000000..5d5b0bf605d2 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.cu @@ -0,0 +1,610 @@ +/* + * 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. + */ + +// MiniMax-M3 TP8 long-context sparse decode attention in one launch (8 query heads on one KV head). +// +// Grid: one cluster of 16 CTAs per query token; CTA `r` attends the token's `r`-th selected block. The QK and PV +// MMAs are transposed (S^T = K q^T, O^T = V^T P^T) so the 8 heads fill the MMA n dimension. Each CTA pushes its +// normalized partial over DSMEM so that CTA `r` holds (head r / 2, dims (r % 2) * 64 + [0, 64)) of all 16 partials, +// merges them, and writes the BF16 row and, optionally, its MXFP8 copy and swizzled UE8M0 scales. +// +// Per-block math: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16, QK and PV on FP16 tensor cores with +// FP32 accumulation, log2-domain softmax with either the linear denominator or the log-domain recurrence. +// +// PREWAIT stages everything that does not depend on the top-16 table before the grid dependency wait: the request's +// seq_lens entry, the q rows and the request's block-table row (in dynamic shared memory after Smem8, when it has at +// most kBtCap entries). The arithmetic is the same in both variants. + +#include "tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kDim = 128; +constexpr int kPage = 128; +constexpr int kTopk = 16; +constexpr int kThreads = 128; +constexpr int kPitch = kDim + 8; // halves per smem row: 272 B, conflict-free ldmatrix +constexpr float kNegInf = -__builtin_huge_valf(); +constexpr int kHeads8 = 8; +constexpr int kHalfDim = kDim / 2; +constexpr int kBtCap = 8192; // block-table entries staged in shared memory (8192 x 128 = 1M tokens) + +using Params = MinimaxM3Tp8SparseDecodeLc8Params; + +struct Smem8 +{ + __half qs[kHeads8 * kPitch]; + __half vs[kPage * kPitch]; + __half ps[kHeads8 * kPitch]; // P^T: [head][token] + float recv[kTopk][kHalfDim]; // (head rank / 2, dims (rank % 2) * 64 + [0, 64)) partials of the 16 CTAs + float recv_lse[kTopk]; + float part[kHalfDim]; + float red[4][kHeads8]; + float rowm[kHeads8]; + float rowl[kHeads8]; +}; + +template +constexpr size_t smemBytes() +{ + return PREWAIT ? sizeof(Smem8) + kBtCap * sizeof(int) : sizeof(Smem8); +} + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ uint32_t e2m1x2_to_f16x2(uint32_t byte) +{ + uint32_t out; + asm("{\n .reg .b8 s;\n cvt.u8.u32 s, %1;\n cvt.rn.f16x2.e2m1x2 %0, s;\n}" : "=r"(out) : "r"(byte)); + return out; +} + +__device__ __forceinline__ uint32_t e4m3x2_to_f16x2(uint16_t pair) +{ + uint32_t out; + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(out) : "h"(pair)); + return out; +} + +__device__ __forceinline__ float ex2(float x) +{ + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ float lg2(float x) +{ + float y; + asm("lg2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ void ldsm_x4(uint32_t (&r)[4], uint32_t addr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x4_t(uint32_t (&r)[4], uint32_t addr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void mma(float (&c)[4], uint32_t const (&a)[4], 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};" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ void cluster_sync() +{ + asm volatile("barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;" ::: "memory"); +} + +__device__ __forceinline__ uint32_t map_cluster(void const* local, uint32_t rank) +{ + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(smem_u32(local)), "r"(rank)); + return remote; +} + +__device__ __forceinline__ void st_cluster_f32(uint32_t addr, float a) +{ + asm volatile("st.shared::cluster.f32 [%0], %1;" ::"r"(addr), "f"(a) : "memory"); +} + +// One packed NVFP4 row (64 bytes, 128 values) with its 8 E4M3 scales -> 128 fp16 values in smem. +__device__ __forceinline__ void dequant_row(__half* dst, uint8_t const* packed, uint8_t const (&sc)[8], bool valid) +{ + uint4 const* src = reinterpret_cast(packed); + uint32_t* d = reinterpret_cast(dst); +#pragma unroll + for (int c = 0; c < 4; ++c) + { + uint4 w = valid ? __ldg(src + c) : make_uint4(0, 0, 0, 0); + uint32_t words[4] = {w.x, w.y, w.z, w.w}; + // 16 bytes = 32 values = scale columns 2c and 2c + 1. +#pragma unroll + for (int h = 0; h < 2; ++h) + { + uint8_t const s = valid ? sc[2 * c + h] : 0; + uint32_t const sh = e4m3x2_to_f16x2(static_cast(s | (s << 8))); + __half2 const scale = *reinterpret_cast<__half2 const*>(&sh); +#pragma unroll + for (int j = 0; j < 2; ++j) + { + uint32_t const word = words[2 * h + j]; +#pragma unroll + for (int b = 0; b < 4; ++b) + { + uint32_t const v2 = e2m1x2_to_f16x2((word >> (8 * b)) & 0xFF); + __half2 const x = __hmul2(*reinterpret_cast<__half2 const*>(&v2), scale); + d[c * 16 + h * 8 + j * 4 + b] = *reinterpret_cast(&x); + } + } + } + } +} + +// q: 8 heads x 128 e4m3 -> fp16; threads 0..63 = (head, 16-value chunk). +__device__ __forceinline__ void load_q(Smem8& s, Params const& p, int tok, int tid) +{ + if (tid < kHeads8 * 8) + { + int const h = tid >> 3; + int const c = (tid & 7) * 16; + uint4 const w = *reinterpret_cast(p.q + tok * p.q_tok + h * p.q_head + c); + uint32_t const words[4] = {w.x, w.y, w.z, w.w}; + uint32_t* d = reinterpret_cast(s.qs + h * kPitch + c); +#pragma unroll + for (int i = 0; i < 4; ++i) + { + d[2 * i] = e4m3x2_to_f16x2(static_cast(words[i] & 0xFFFF)); + d[2 * i + 1] = e4m3x2_to_f16x2(static_cast(words[i] >> 16)); + } + } +} +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +template +__global__ void __launch_bounds__(kThreads) m3_tp8_lc8_attn_kernel(__grid_constant__ Params const p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ __align__(128) uint8_t smem_raw[]; + Smem8& s = *reinterpret_cast(smem_raw); + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const g = lane >> 2; + int const qd = lane & 3; + int const tok = blockIdx.x / kTopk; + uint32_t rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(rank)); + + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); + if constexpr (!PREWAIT) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); + } + + int const req = tok / p.dql; + int const qoff = tok - req * p.dql; + int const seq = p.seq_lens[req]; + int const kv_len = max(seq - p.dql + qoff + 1, 0); + int const nblocks = (kv_len + kPage - 1) / kPage; + int const real_topk = min(kTopk, nblocks); + + // The request's block-table row, staged before the grid dependency (written before the graph runs). + int* const s_bt = reinterpret_cast(smem_raw + sizeof(Smem8)); + bool bt_smem = false; + if constexpr (PREWAIT) + { + int const nbt = (max(seq, 0) + kPage - 1) / kPage; + bt_smem = nbt <= kBtCap; + if (bt_smem) + { + for (int i = tid; i < nbt; i += kThreads) + { + s_bt[i] = __ldg(p.block_table + req * p.bt_stride + i); + } + } + load_q(s, p, tok, tid); + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); + } + int blk = static_cast(rank) < real_topk ? p.topk[tok * p.t_tok + rank * p.t_k] : -1; + bool const active = blk >= 0; + if constexpr (PREWAIT) + { + __syncthreads(); // s_bt and the q rows are complete + } + else + { + load_q(s, p, tok, tid); + } + + // Lane holds heads 2 * qd + {0, 1}; acc[mt]: dims warp * 32 + mt * 16 + g + {0, 8}. + float m_row[2] = {kNegInf, kNegInf}; + float l_row[2] = {0.F, 0.F}; + float acc[2][4]; +#pragma unroll + for (int j = 0; j < 2; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + acc[j][i] = 0.F; + + // K rows n = warp * 32 + j * 8 + g stay packed in registers. As the MMA A operand, m-tile mt takes rows + // j = 2 mt (row g) and 2 mt + 1 (row g + 8); k-step ks takes packed bytes 8 ks + qd and 8 ks + 4 + qd. + uint32_t kw[4][16]; + uint2 ksc[4]; + if (active) + { + int64_t const page = bt_smem ? s_bt[blk] : p.block_table[req * p.bt_stride + blk]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + int const n = warp * 32 + j * 8 + g; + bool const valid = blk * kPage + n < kv_len; + uint4 const* src = reinterpret_cast(p.k + page * p.k_blk + n * 64); +#pragma unroll + for (int c = 0; c < 4; ++c) + { + uint4 const w = valid ? __ldg(src + c) : make_uint4(0, 0, 0, 0); + kw[j][4 * c] = w.x; + kw[j][4 * c + 1] = w.y; + kw[j][4 * c + 2] = w.z; + kw[j][4 * c + 3] = w.w; + } + ksc[j] = valid ? __ldg(reinterpret_cast(p.ksc + page * p.ksc_blk + n * 8)) : make_uint2(0, 0); + } + bool const valid = blk * kPage + tid < kv_len; + uint8_t vsc[8]; + if (valid) + { + uint8_t const* vb = p.vsc + page * p.vsc_blk + (tid >> 2) * 32 + (tid & 3); +#pragma unroll + for (int c = 0; c < 8; ++c) + vsc[c] = __ldg(vb + 4 * c); + } + dequant_row(s.vs + tid * kPitch, p.v + page * p.v_blk + tid * 64, vsc, valid); + } + __syncthreads(); + + if (active) + { + float const scale_log2 = p.sm_scale * 1.4426950409F * __ldg(p.kgs); + int const q4 = qd * 8; + // S^T = K q^T for this warp's 32 tokens: rows (tokens) g / g + 8 of each m-tile, columns (heads) 2 qd + {0, 1}. + float sacc[2][4]; +#pragma unroll + for (int j = 0; j < 2; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + sacc[j][i] = 0.F; +#pragma unroll + for (int ks2 = 0; ks2 < kDim / 16; ks2 += 2) + { + uint32_t bq[4]; + ldsm_x4(bq, smem_u32(s.qs + (lane & 7) * kPitch + ks2 * 16 + (lane >> 3) * 8)); +#pragma unroll + for (int u = 0; u < 2; ++u) + { + int const ks = ks2 + u; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) + { + uint32_t a[4]; +#pragma unroll + for (int hi = 0; hi < 2; ++hi) + { + int const j = 2 * mt + hi; + uint32_t const scw = ks < 4 ? ksc[j].x : ksc[j].y; + uint8_t const sc = static_cast(scw >> (8 * (ks & 3))); + uint32_t const sh = e4m3x2_to_f16x2(static_cast(sc | (sc << 8))); + __half2 const scale = *reinterpret_cast<__half2 const*>(&sh); + uint32_t const r0 = e2m1x2_to_f16x2((kw[j][2 * ks] >> q4) & 0xFF); + uint32_t const r1 = e2m1x2_to_f16x2((kw[j][2 * ks + 1] >> q4) & 0xFF); + __half2 const h0 = __hmul2(*reinterpret_cast<__half2 const*>(&r0), scale); + __half2 const h1 = __hmul2(*reinterpret_cast<__half2 const*>(&r1), scale); + a[hi] = *reinterpret_cast(&h0); + a[2 + hi] = *reinterpret_cast(&h1); + } + mma(sacc[mt], a, bq[2 * u], bq[2 * u + 1]); + } + } + } + float mloc[2] = {kNegInf, kNegInf}; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const t = warp * 32 + mt * 16 + g + 8 * (i >> 1); + float const v = blk * kPage + t < kv_len ? 0.F + sacc[mt][i] * scale_log2 : kNegInf; + sacc[mt][i] = v; + mloc[i & 1] = fmaxf(mloc[i & 1], v); + } +#pragma unroll + for (int e = 0; e < 2; ++e) + { + mloc[e] = fmaxf(mloc[e], __shfl_xor_sync(0xffffffff, mloc[e], 4)); + mloc[e] = fmaxf(mloc[e], __shfl_xor_sync(0xffffffff, mloc[e], 8)); + mloc[e] = fmaxf(mloc[e], __shfl_xor_sync(0xffffffff, mloc[e], 16)); + } + if (g == 0) + { + s.red[warp][2 * qd] = mloc[0]; + s.red[warp][2 * qd + 1] = mloc[1]; + } + __syncthreads(); + if (tid < kHeads8) + s.rowm[tid] = fmaxf(fmaxf(s.red[0][tid], s.red[1][tid]), fmaxf(s.red[2][tid], s.red[3][tid])); + __syncthreads(); + m_row[0] = s.rowm[2 * qd]; + m_row[1] = s.rowm[2 * qd + 1]; + float lloc[2] = {0.F, 0.F}; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int i = 0; i < 4; ++i) + { + float const pv = ex2(sacc[mt][i] - m_row[i & 1]); + lloc[i & 1] += pv; + int const t = warp * 32 + mt * 16 + g + 8 * (i >> 1); + s.ps[(2 * qd + (i & 1)) * kPitch + t] = __float2half_rn(pv); + } +#pragma unroll + for (int e = 0; e < 2; ++e) + { + lloc[e] += __shfl_xor_sync(0xffffffff, lloc[e], 4); + lloc[e] += __shfl_xor_sync(0xffffffff, lloc[e], 8); + lloc[e] += __shfl_xor_sync(0xffffffff, lloc[e], 16); + } + __syncthreads(); // rowm reads done before red is reused; ps complete + if (g == 0) + { + s.red[warp][2 * qd] = lloc[0]; + s.red[warp][2 * qd + 1] = lloc[1]; + } + // O^T[warp's 32 dims, 8 heads] = V^T P^T over all 128 tokens. +#pragma unroll + for (int kk = 0; kk < kPage; kk += 32) + { + uint32_t b[4]; + ldsm_x4(b, smem_u32(s.ps + (lane & 7) * kPitch + kk + (lane >> 3) * 8)); +#pragma unroll + for (int u = 0; u < 2; ++u) + { +#pragma unroll + for (int mt = 0; mt < 2; ++mt) + { + uint32_t a[4]; + int const t = kk + u * 16 + (lane & 7) + ((lane >> 4) & 1) * 8; + int const d = warp * 32 + mt * 16 + ((lane >> 3) & 1) * 8; + ldsm_x4_t(a, smem_u32(s.vs + t * kPitch + d)); + mma(acc[mt], a, b[2 * u], b[2 * u + 1]); + } + } + } + __syncthreads(); + if (tid < kHeads8) + s.rowl[tid] = (s.red[0][tid] + s.red[1][tid]) + (s.red[2][tid] + s.red[3][tid]); + __syncthreads(); + l_row[0] = s.rowl[2 * qd]; + l_row[1] = s.rowl[2 * qd + 1]; + } + + // Normalize; push (head, dim) to CTA 2 * head + dim / 64, slot recv[rank][dim % 64], plus each head's LSE (log2) + // to both of its CTAs. Posted DSMEM stores, made visible by the cluster barrier. + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); + { + float const vgs = __ldg(p.vgs); + float scale[2], lse[2]; +#pragma unroll + for (int e = 0; e < 2; ++e) + { + if constexpr (LINEAR) + { + scale[e] = l_row[e] > 0.F ? 1.0F / l_row[e] : 0.F; + lse[e] = l_row[e] > 0.F ? m_row[e] + lg2(l_row[e]) : kNegInf; + } + else + { + float const lse_i = active ? m_row[e] + lg2(ex2(kNegInf - m_row[e]) + l_row[e]) : kNegInf; + scale[e] = lse_i > kNegInf ? ex2(m_row[e] - lse_i) : 0.F; + lse[e] = lse_i; + } + scale[e] = scale[e] * vgs; + } + uint32_t dst[2]; +#pragma unroll + for (int e = 0; e < 2; ++e) + dst[e] = map_cluster(&s.recv[rank][0], (2 * qd + e) * 2 + (warp >> 1)); +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const dl = (warp & 1) * 32 + mt * 16 + g + 8 * (i >> 1); + st_cluster_f32(dst[i & 1] + dl * 4, acc[mt][i] * scale[i & 1]); + } + if (warp == 0 && g == 0) + { +#pragma unroll + for (int e = 0; e < 2; ++e) +#pragma unroll + for (int hf = 0; hf < 2; ++hf) + st_cluster_f32(map_cluster(&s.recv_lse[rank], (2 * qd + e) * 2 + hf), lse[e]); + } + } + cluster_sync(); + + // CTA `rank` merges head rank / 2, dims (rank % 2) * 64 + dl; thread = (chunk half, dl). + int const h = static_cast(rank) >> 1; + int const dh = static_cast(rank) & 1; + int const dl = tid & (kHalfDim - 1); + int const hc = tid >> 6; + float lse_c[kTopk]; + float lse_max = kNegInf; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + { + lse_c[c] = s.recv_lse[c]; + lse_max = fmaxf(lse_max, lse_c[c]); + } + lse_max = lse_max == kNegInf ? 0.F : lse_max; + float denom = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + { + lse_c[c] = ex2(lse_c[c] - lse_max); + denom += lse_c[c]; + } + denom = denom > 0.F ? denom : 1.F; + float o = 0.F; +#pragma unroll + for (int c = 0; c < kTopk / 2; ++c) + { + int const cc = hc * (kTopk / 2) + c; + o += s.recv[cc][dl] * __fdiv_rn(lse_c[cc], denom); + } + if (hc) + s.part[dl] = o; + __syncthreads(); + if (!hc) + { + o += s.part[dl]; + int const col = h * kDim + dh * kHalfDim + dl; + __nv_bfloat16 const ob = __float2bfloat16_rn(o); + p.out[tok * p.out_stride + col] = ob; + if constexpr (MXFP8) + { + float const x = __bfloat162float(ob); + float amax = fabsf(x); +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, off)); + float rcp; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(rcp) : "f"(448.0F)); + float const sfv = amax * rcp; + int const bits = __float_as_int(sfv); + int e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0); + e = min(e, 254); + float const qs = amax != 0.F ? __int_as_float((254 - e) << 23) : 0.F; + float const xq = x * qs; + uint16_t q2; + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(q2) : "f"(0.F), "f"(xq)); + p.q8[static_cast(tok) * (p.num_sf_k_tiles * 128) + col] = static_cast(q2 & 0xFF); + if (lane == 0) + { + int const k_sf = col / 32; + int64_t const off = static_cast(tok / 128) * p.num_sf_k_tiles * 512 + (k_sf / 4) * 512 + + (tok % 32) * 16 + ((tok % 128) / 32) * 4 + k_sf % 4; + p.sf[off] = static_cast(e); + } + } + } +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +template +cudaError_t configure() +{ + auto kernel = m3_tp8_lc8_attn_kernel; + cudaError_t err = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smemBytes())); + if (err != cudaSuccess) + { + return err; + } + return cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1); +} + +template +cudaError_t launch(Params const& prm, int num_tokens, bool use_pdl, cudaStream_t stream) +{ + // Kernel attributes are set once per process, on the first launch (outside CUDA-graph capture: warmup runs + // every variant eagerly first). + static cudaError_t const configured = configure(); + if (configured != cudaSuccess) + { + return configured; + } + auto kernel = m3_tp8_lc8_attn_kernel; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(num_tokens * kTopk); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = smemBytes(); + cfg.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kTopk; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; + cfg.attrs = attrs; + cfg.numAttrs = 2; + return cudaLaunchKernelEx(&cfg, kernel, prm); +} + +template +cudaError_t launchPrewait(Params const& prm, int num_tokens, bool prewait, bool use_pdl, cudaStream_t stream) +{ + return prewait ? launch(prm, num_tokens, use_pdl, stream) + : launch(prm, num_tokens, use_pdl, stream); +} + +} // namespace + +cudaError_t invokeMinimaxM3Tp8SparseDecodeLc8(MinimaxM3Tp8SparseDecodeLc8Params const& params, int32_t numTokens, + bool linearSoftmax, bool mxfp8, bool prewait, bool usePdl, cudaStream_t stream) +{ + if (linearSoftmax) + { + return mxfp8 ? launchPrewait(params, numTokens, prewait, usePdl, stream) + : launchPrewait(params, numTokens, prewait, usePdl, stream); + } + return mxfp8 ? launchPrewait(params, numTokens, prewait, usePdl, stream) + : launchPrewait(params, numTokens, prewait, usePdl, stream); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.h b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.h new file mode 100644 index 000000000000..ad65a577eab8 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.h @@ -0,0 +1,75 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +//! Operands of the MiniMax-M3 TP8 one-launch long-context sparse decode. Strides are in elements of the pointed-to +//! type. Geometry is fixed: 8 query heads on one KV head, head_dim 128, page 128, top-k 16. +struct MinimaxM3Tp8SparseDecodeLc8Params +{ + uint8_t const* q; // e4m3 [T, 8, 128] + int64_t q_tok, q_head; + uint8_t const* k; // packed e2m1 [pages, 1, 128, 64] + uint8_t const* v; + int64_t k_blk, v_blk; + uint8_t const* ksc; // e4m3 [pages, 1, 128, 8], token-major + uint8_t const* vsc; // e4m3, 4x4 token-quad order + int64_t ksc_blk, vsc_blk; + float const* kgs; + float const* vgs; + int const* topk; // [T, 16] view (kv head 0) + int64_t t_tok, t_k; + int const* block_table; + int64_t bt_stride; + int const* seq_lens; + int dql; + float sm_scale; + __nv_bfloat16* out; // [T, 1024] + int64_t out_stride; + uint8_t* q8; // e4m3 [rows, num_sf_k_tiles * 128] + uint8_t* sf; + int num_sf_k_tiles; +}; + +//! One cluster of 16 CTAs per query token: CTA r attends the token's r-th selected block with transposed MMAs (the 8 +//! heads fill the MMA n dimension), then CTA r merges (head r / 2, dims (r % 2) * 64 + [0, 64)) of the 16 partials +//! over DSMEM into the BF16 [T, 1024] row (and, when mxfp8, its MXFP8 copy with swizzled UE8M0 scales). Dependents +//! are released right after the grid dependency wait. +//! +//! With prewait, the request's seq_lens entry, its q rows (widened to FP16) and its block-table row (up to 8192 +//! entries, in dynamic shared memory; longer rows are read from global memory) are staged before the grid dependency +//! wait; outputs are bit-identical to prewait=false. +//! +//! Instantiated variants: linearSoftmax x mxfp8 x prewait, all in {false, true}. The device code is SM100-family +//! only; on other architectures the kernel compiles to an empty body, so callers check the device first. +cudaError_t invokeMinimaxM3Tp8SparseDecodeLc8(MinimaxM3Tp8SparseDecodeLc8Params const& params, int32_t numTokens, + bool linearSoftmax, bool mxfp8, bool prewait, bool usePdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index cd7525f45a1c..bd4faeea80f9 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -67,8 +67,6 @@ add_th_common_family( IndexerKCacheScatterOp.cpp IndexerTopKOp.cpp sparseKvCacheCompactOp.cpp - minimaxM3Fp8IndexerOp.cpp - minimaxM3SelectBlocksOp.cpp mlaRopeInplaceOp.cpp relativeAttentionBiasOp.cpp dsv3RopeOp.cpp @@ -76,6 +74,10 @@ add_th_common_family( nvfp4MlaKvCacheGatherOp.cpp trtllmGenQKVProcessOp.cpp) +# MiniMax-M3 torch ops live in minimaxM3/; that directory adds them to +# th_common_attention with target_sources() so the list stays next to the files. +add_subdirectory(minimaxM3) + target_sources( th_common_attention PRIVATE "${KIMI_K3_MODEL_CSRC}/k3AttnResOp.cpp" diff --git a/cpp/tensorrt_llm/thop/minimaxM3/CMakeLists.txt b/cpp/tensorrt_llm/thop/minimaxM3/CMakeLists.txt new file mode 100644 index 000000000000..0ea4ebdce2e0 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/CMakeLists.txt @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not +# use this file except in compliance with the License. You may obtain a copy of +# the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. + +# torch.ops.trtllm MiniMax-M3 ops (kernels in kernels/minimaxM3/). The parent +# defines th_common_attention. +file(GLOB SRC_CPP CONFIGURE_DEPENDS *.cpp) +target_sources(th_common_attention PRIVATE ${SRC_CPP}) diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3AttnC1Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3AttnC1Op.cpp new file mode 100644 index 000000000000..cfc915f939a9 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3AttnC1Op.cpp @@ -0,0 +1,191 @@ +/* + * 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/minimaxM3/attnC1.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace dims = tensorrt_llm::kernels::minimax_m3_attn_c1; +using tensorrt_llm::kernels::MinimaxM3AttnC1Params; + +constexpr int64_t kHeadDim = 128; + +void checkByteTensor(at::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.element_size() == 1, name, " must be a CUDA tensor of one-byte elements"); +} + +} // namespace + +//! One launch for a one-request MiniMax-M3 MSA verify layer of 6 query tokens: the fused norm/RoPE/KV-insert + +//! identity-selection sparse decode + merge of minimax_m3_fused_decode_attn on the packed QKV rows, and the MXFP8 +//! o_proj GEMV, 96 CTAs that are each an attention CTA and an o_proj unit (fed through the MXFP8 hand-off). +//! +//! @param packed bf16 [6, 2560] QKV+index projection rows. +//! @param flags int32 [>= 32 + 96]: the per-CTA call counters (zero-initialized once); written. +//! @param mxfp8Data / mxfp8Scale the two call-parity buffers of the attention's MXFP8 hand-off (0xFF at rest). +//! @param y bf16 [T, 6144] o_proj output. arWs / arFlags / arRank / arWorld / arMaxTokens / arMode: the next +//! multicast all-reduce's exchange (device addresses) to push y into when arMode is 2; 0 for none. +void minimaxM3AttnC1(at::Tensor const& packed, at::Tensor& flags, at::Tensor& kvData, at::Tensor& kvScale, + at::Tensor& idxK, at::Tensor const& outCacheLoc, at::Tensor const& positionIds, at::Tensor const& blockTable, + at::Tensor const& seqLens, at::Tensor const& kvQuantScale, at::Tensor const& kGlobalScale, + at::Tensor const& vGlobalScale, at::Tensor const& qNormWeight, at::Tensor const& kNormWeight, + at::Tensor const& indexKNormWeight, at::Tensor const& rotaryCosSin, double eps, double smScale, + int64_t decodeQueryLen, int64_t tokensPerCluster, int64_t blocksPerCta, at::Tensor& output, at::Tensor& mxfp8Data, + at::Tensor& mxfp8Scale, bool usePdl, at::Tensor const& oWeight, at::Tensor const& oWeightScale, at::Tensor& y, + int64_t arWs, int64_t arFlags, int64_t arRank, int64_t arWorld, int64_t arMaxTokens, int64_t arMode) +{ + TORCH_CHECK(packed.is_cuda() && packed.scalar_type() == at::kBFloat16 && packed.dim() == 2 + && packed.size(1) == 20 * kHeadDim && packed.stride(1) == 1, + "packed must be CUDA bf16 [T, 2560] with contiguous rows"); + checkByteTensor(oWeight, "o_weight"); + checkByteTensor(oWeightScale, "o_weight_scale"); + checkByteTensor(kvData, "kv_data"); + checkByteTensor(kvScale, "kv_scale"); + checkByteTensor(idxK, "idx_k"); + checkByteTensor(mxfp8Data, "mxfp8_data"); + checkByteTensor(mxfp8Scale, "mxfp8_scale"); + TORCH_CHECK( + oWeight.dim() == 2 && oWeight.size(0) == dims::kOpN && oWeight.size(1) == dims::kOpK && oWeight.is_contiguous(), + "o_weight must be contiguous [6144, 2048]"); + TORCH_CHECK( + kvData.dim() == 5 && kvScale.dim() == 5 && idxK.dim() == 4, "kv_data, kv_scale and idx_k must be paged"); + for (auto const* t : {&outCacheLoc, &positionIds, &blockTable, &seqLens}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kInt, + "slot, position, block-table and length tensors " + "must be CUDA int32"); + } + TORCH_CHECK(blockTable.dim() == 2 && blockTable.stride(1) == 1, "block_table must be row-major 2-D"); + for (auto const* t : {&kvQuantScale, &kGlobalScale, &vGlobalScale, &rotaryCosSin}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kFloat, "scales and rotary_cos_sin must be CUDA fp32"); + } + for (auto const* t : {&qNormWeight, &kNormWeight, &indexKNormWeight}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kBFloat16 && t->numel() == kHeadDim, + "norm weights must be CUDA bf16 [128]"); + } + TORCH_CHECK(flags.is_cuda() && flags.scalar_type() == at::kInt && flags.numel() >= 32 + dims::kCtas, + "bad c1 flags: int32 [>= 128]"); + TORCH_CHECK(mxfp8Data.numel() == 2 * dims::kMaxTokens * dims::kOpK && mxfp8Scale.numel() == 2 * dims::kOpSfBytes + && mxfp8Data.is_contiguous() && mxfp8Scale.is_contiguous(), + "c1 MXFP8 hand-off"); + int64_t const numTokens = packed.size(0); + int64_t const numReqs = seqLens.size(0); + TORCH_CHECK(numTokens == numReqs * decodeQueryLen && numTokens <= dims::kMaxTokens, "token count"); + TORCH_CHECK(output.is_cuda() && output.scalar_type() == at::kBFloat16 && output.dim() == 2 + && output.size(0) >= numTokens && output.size(1) == 16 * kHeadDim && output.stride(1) == 1, + "output must be CUDA bf16 [T, 2048] with contiguous rows"); + TORCH_CHECK(y.is_cuda() && y.scalar_type() == at::kBFloat16 && y.dim() == 2 && y.size(0) == numTokens + && y.size(1) == dims::kOpN && y.is_contiguous(), + "y must be [M, 6144]"); + TORCH_CHECK(decodeQueryLen == 6 && tokensPerCluster == 1 && blocksPerCta == 1 && numReqs == 1, + "minimax_m3_attn_c1 is compiled for one request with six query tokens (geometry (6, 1, 1)), got ", numReqs, + " x (", decodeQueryLen, ", ", tokensPerCluster, ", ", blocksPerCta, ")"); + TORCH_CHECK(arMode == 0 || arMode == 2, "ar_mode must be 0 (none) or 2 (push)"); + + c10::cuda::CUDAGuard const deviceGuard{packed.device()}; + auto const* prop = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK( + prop->major == 10, "minimax_m3_attn_c1 requires an sm_100-family GPU, got sm_", prop->major, prop->minor); + + MinimaxM3AttnC1Params prm = {}; + prm.kv_data = reinterpret_cast(kvData.data_ptr()); + prm.d_page = kvData.stride(0); + prm.d_plane = kvData.stride(1); + prm.kv_scale = reinterpret_cast(kvScale.data_ptr()); + prm.s_page = kvScale.stride(0); + prm.s_plane = kvScale.stride(1); + prm.idx_k = reinterpret_cast(idxK.data_ptr()); + prm.i_page = idxK.stride(0); + prm.i_tok = idxK.stride(2); + prm.num_pages = kvData.size(0); + prm.out_cache_loc = outCacheLoc.data_ptr(); + prm.position_ids = positionIds.data_ptr(); + prm.block_table = blockTable.data_ptr(); + prm.bt_stride = blockTable.stride(0); + prm.bt_cols = static_cast(blockTable.size(1)); + prm.seq_lens = seqLens.data_ptr(); + prm.kv_quant_scale = kvQuantScale.data_ptr(); + prm.k_gs = kGlobalScale.data_ptr(); + prm.v_gs = vGlobalScale.data_ptr(); + prm.q_w = reinterpret_cast<__nv_bfloat16 const*>(qNormWeight.data_ptr()); + prm.k_w = reinterpret_cast<__nv_bfloat16 const*>(kNormWeight.data_ptr()); + prm.ik_w = reinterpret_cast<__nv_bfloat16 const*>(indexKNormWeight.data_ptr()); + prm.rope = rotaryCosSin.data_ptr(); + prm.eps = static_cast(eps); + prm.sm_scale = static_cast(smScale); + prm.out = reinterpret_cast<__nv_bfloat16*>(output.data_ptr()); + prm.out_stride = output.stride(0); + prm.q8 = reinterpret_cast(mxfp8Data.data_ptr()); + prm.sf = reinterpret_cast(mxfp8Scale.data_ptr()); + prm.num_sf_k_tiles = static_cast(output.size(1) / 128); + prm.packed = reinterpret_cast<__nv_bfloat16 const*>(packed.data_ptr()); + prm.packed_stride = packed.stride(0); + prm.num_tokens = static_cast(numTokens); + prm.flags = reinterpret_cast(flags.data_ptr()); + prm.wsf_o = reinterpret_cast(oWeightScale.data_ptr()); + prm.y_out = reinterpret_cast<__nv_bfloat16*>(y.data_ptr()); + prm.ar_ws = reinterpret_cast(arWs); + prm.ar_flags = reinterpret_cast(arFlags); + prm.ar_rank = static_cast(arRank); + prm.ar_world = static_cast(arWorld); + prm.ar_maxt = static_cast(arMaxTokens); + prm.ar_mode = static_cast(arMode); + + cudaError_t const err = tensorrt_llm::kernels::invokeMinimaxM3AttnC1(prm, oWeight.data_ptr(), + static_cast(decodeQueryLen), static_cast(tokensPerCluster), static_cast(blocksPerCta), + static_cast(numReqs), usePdl, at::cuda::getCurrentCUDAStream(packed.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "minimax_m3_attn_c1 launch failed (qlen=", decodeQueryLen, " qg=", tokensPerCluster, + " bpc=", blocksPerCta, "): ", cudaGetErrorString(err)); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_attn_c1(Tensor packed, Tensor(a!) flags, Tensor(b!) kv_data, Tensor(c!) kv_scale, " + "Tensor(d!) idx_k, Tensor out_cache_loc, Tensor position_ids, Tensor block_table, Tensor seq_lens, " + "Tensor kv_quant_scale, Tensor k_global_scale, Tensor v_global_scale, Tensor q_norm_weight, " + "Tensor k_norm_weight, Tensor index_k_norm_weight, Tensor rotary_cos_sin, float eps, float sm_scale, " + "int decode_query_len, int tokens_per_cluster, int blocks_per_cta, Tensor(e!) output, Tensor(f!) mxfp8_data, " + "Tensor(g!) mxfp8_scale, bool use_pdl, Tensor o_weight, Tensor o_weight_scale, Tensor(h!) y, " + "int ar_workspace, int ar_flags, int ar_rank, int ar_world, int ar_max_tokens, int ar_mode) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_attn_c1", &tensorrt_llm::torch_ext::minimaxM3AttnC1); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3AttnMegaOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3AttnMegaOp.cpp new file mode 100644 index 000000000000..b94942e052f5 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3AttnMegaOp.cpp @@ -0,0 +1,207 @@ +/* + * 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/minimaxM3/attnMega.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace dims = tensorrt_llm::kernels::minimax_m3_attn_mega; +using tensorrt_llm::kernels::MinimaxM3AttnMegaParams; + +constexpr int64_t kHeadDim = 128; + +void checkByteTensor(at::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.element_size() == 1, name, " must be a CUDA tensor of one-byte elements"); +} + +} // namespace + +//! One launch for a one-request, one-token MiniMax-M3 MSA decode layer: the MXFP8 QKV GEMV of the layer input, the +//! fused norm/RoPE/KV-insert + identity-selection sparse decode + merge of minimax_m3_fused_decode_attn, and the MXFP8 +//! o_proj GEMV, joined by per-tile ready counters in `flags`. +//! +//! @param xq / xsf the layer input's E4M3 rows [T, 6144] and 128x4-swizzled UE8M0 scale atoms. +//! @param qkvWeight / qkvWeightScale E4M3 [2560, 6144] and its swizzled UE8M0 scales. +//! @param qkvOut bf16 [2, 8, 2560] call-parity buffers of the packed QKV rows (0xFFFF "not yet" at rest); written. +//! @param flags int32 [>= 32 + 112]: tile counters and the per-CTA call counters (zero-initialized once); written. +//! @param mxfp8Data / mxfp8Scale the two call-parity buffers of the attention's MXFP8 hand-off (0xFF at rest). +//! @param y bf16 [T, 6144] o_proj output. arWs / arFlags / arRank / arWorld / arMaxTokens / arMode: the next +//! multicast all-reduce's exchange (device addresses) to push y into when arMode is 2; 0 for none. +void minimaxM3AttnMega(at::Tensor const& xq, at::Tensor const& xsf, at::Tensor const& qkvWeight, + at::Tensor const& qkvWeightScale, at::Tensor& qkvOut, at::Tensor& flags, at::Tensor& kvData, at::Tensor& kvScale, + at::Tensor& idxK, at::Tensor const& outCacheLoc, at::Tensor const& positionIds, at::Tensor const& blockTable, + at::Tensor const& seqLens, at::Tensor const& kvQuantScale, at::Tensor const& kGlobalScale, + at::Tensor const& vGlobalScale, at::Tensor const& qNormWeight, at::Tensor const& kNormWeight, + at::Tensor const& indexKNormWeight, at::Tensor const& rotaryCosSin, double eps, double smScale, + int64_t decodeQueryLen, int64_t tokensPerCluster, int64_t blocksPerCta, at::Tensor& output, at::Tensor& mxfp8Data, + at::Tensor& mxfp8Scale, bool usePdl, at::Tensor const& oWeight, at::Tensor const& oWeightScale, at::Tensor& y, + int64_t arWs, int64_t arFlags, int64_t arRank, int64_t arWorld, int64_t arMaxTokens, int64_t arMode) +{ + checkByteTensor(xq, "xq"); + checkByteTensor(xsf, "xsf"); + checkByteTensor(qkvWeight, "qkv_weight"); + checkByteTensor(qkvWeightScale, "qkv_weight_scale"); + checkByteTensor(oWeight, "o_weight"); + checkByteTensor(oWeightScale, "o_weight_scale"); + checkByteTensor(kvData, "kv_data"); + checkByteTensor(kvScale, "kv_scale"); + checkByteTensor(idxK, "idx_k"); + checkByteTensor(mxfp8Data, "mxfp8_data"); + checkByteTensor(mxfp8Scale, "mxfp8_scale"); + TORCH_CHECK(xq.dim() == 2 && xq.size(1) == dims::kQkvK && xq.stride(1) == 1, "xq must be [T, 6144] row-major"); + TORCH_CHECK(qkvWeight.dim() == 2 && qkvWeight.size(0) == dims::kQkvN && qkvWeight.size(1) == dims::kQkvK + && qkvWeight.is_contiguous(), + "qkv_weight must be contiguous [2560, 6144]"); + TORCH_CHECK( + oWeight.dim() == 2 && oWeight.size(0) == dims::kOpN && oWeight.size(1) == dims::kOpK && oWeight.is_contiguous(), + "o_weight must be contiguous [6144, 2048]"); + TORCH_CHECK( + kvData.dim() == 5 && kvScale.dim() == 5 && idxK.dim() == 4, "kv_data, kv_scale and idx_k must be paged"); + for (auto const* t : {&outCacheLoc, &positionIds, &blockTable, &seqLens}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kInt, + "slot, position, block-table and length tensors " + "must be CUDA int32"); + } + TORCH_CHECK(blockTable.dim() == 2 && blockTable.stride(1) == 1, "block_table must be row-major 2-D"); + for (auto const* t : {&kvQuantScale, &kGlobalScale, &vGlobalScale, &rotaryCosSin}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kFloat, "scales and rotary_cos_sin must be CUDA fp32"); + } + for (auto const* t : {&qNormWeight, &kNormWeight, &indexKNormWeight}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kBFloat16 && t->numel() == kHeadDim, + "norm weights must be CUDA bf16 [128]"); + } + TORCH_CHECK(flags.is_cuda() && flags.scalar_type() == at::kInt && flags.numel() >= 32 + dims::kGrid, + "flags must be CUDA int32 [>= 144]"); + TORCH_CHECK(qkvOut.is_cuda() && qkvOut.element_size() == 2 && qkvOut.dim() == 3 && qkvOut.size(0) == 2 + && qkvOut.size(1) == dims::kMaxTokens && qkvOut.size(2) == dims::kQkvN && qkvOut.is_contiguous(), + "bad megakernel workspace: qkv_out must be contiguous [2, 8, 2560]"); + TORCH_CHECK(mxfp8Data.numel() == 2 * dims::kMaxTokens * dims::kOpK && mxfp8Scale.numel() == 2 * dims::kOpSfBytes + && mxfp8Data.is_contiguous() && mxfp8Scale.is_contiguous(), + "megakernel MXFP8 hand-off"); + int64_t const numTokens = xq.size(0); + int64_t const numReqs = seqLens.size(0); + TORCH_CHECK(numTokens == numReqs * decodeQueryLen && numTokens <= dims::kMaxTokens, "token count"); + TORCH_CHECK(output.is_cuda() && output.scalar_type() == at::kBFloat16 && output.dim() == 2 + && output.size(0) >= numTokens && output.size(1) == 16 * kHeadDim && output.stride(1) == 1, + "output must be CUDA bf16 [T, 2048] with contiguous rows"); + TORCH_CHECK(y.is_cuda() && y.scalar_type() == at::kBFloat16 && y.dim() == 2 && y.size(0) == numTokens + && y.size(1) == dims::kOpN && y.is_contiguous(), + "y must be [M, 6144]"); + TORCH_CHECK(decodeQueryLen == 1 && tokensPerCluster == 1 && blocksPerCta == 1 && numReqs == 1, + "minimax_m3_attn_mega is compiled for one request with one query token (geometry (1, 1, 1)), got ", numReqs, + " x (", decodeQueryLen, ", ", tokensPerCluster, ", ", blocksPerCta, ")"); + TORCH_CHECK(arMode == 0 || arMode == 2, "ar_mode must be 0 (none) or 2 (push)"); + + c10::cuda::CUDAGuard const deviceGuard{xq.device()}; + auto const* prop = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK( + prop->major == 10, "minimax_m3_attn_mega requires an sm_100-family GPU, got sm_", prop->major, prop->minor); + + MinimaxM3AttnMegaParams prm = {}; + prm.kv_data = reinterpret_cast(kvData.data_ptr()); + prm.d_page = kvData.stride(0); + prm.d_plane = kvData.stride(1); + prm.kv_scale = reinterpret_cast(kvScale.data_ptr()); + prm.s_page = kvScale.stride(0); + prm.s_plane = kvScale.stride(1); + prm.idx_k = reinterpret_cast(idxK.data_ptr()); + prm.i_page = idxK.stride(0); + prm.i_tok = idxK.stride(2); + prm.num_pages = kvData.size(0); + prm.out_cache_loc = outCacheLoc.data_ptr(); + prm.position_ids = positionIds.data_ptr(); + prm.block_table = blockTable.data_ptr(); + prm.bt_stride = blockTable.stride(0); + prm.bt_cols = static_cast(blockTable.size(1)); + prm.seq_lens = seqLens.data_ptr(); + prm.kv_quant_scale = kvQuantScale.data_ptr(); + prm.k_gs = kGlobalScale.data_ptr(); + prm.v_gs = vGlobalScale.data_ptr(); + prm.q_w = reinterpret_cast<__nv_bfloat16 const*>(qNormWeight.data_ptr()); + prm.k_w = reinterpret_cast<__nv_bfloat16 const*>(kNormWeight.data_ptr()); + prm.ik_w = reinterpret_cast<__nv_bfloat16 const*>(indexKNormWeight.data_ptr()); + prm.rope = rotaryCosSin.data_ptr(); + prm.eps = static_cast(eps); + prm.sm_scale = static_cast(smScale); + prm.out = reinterpret_cast<__nv_bfloat16*>(output.data_ptr()); + prm.out_stride = output.stride(0); + prm.q8 = reinterpret_cast(mxfp8Data.data_ptr()); + prm.sf = reinterpret_cast(mxfp8Scale.data_ptr()); + prm.num_sf_k_tiles = static_cast(output.size(1) / 128); + prm.xq = reinterpret_cast(xq.data_ptr()); + prm.ldxq = static_cast(xq.stride(0)); + prm.xsf_act = reinterpret_cast(xsf.data_ptr()); + prm.num_tokens = static_cast(numTokens); + prm.wsf = reinterpret_cast(qkvWeightScale.data_ptr()); + prm.qkv_out = reinterpret_cast<__nv_bfloat16*>(qkvOut.data_ptr()); + prm.flags = reinterpret_cast(flags.data_ptr()); + prm.wsf_o = reinterpret_cast(oWeightScale.data_ptr()); + prm.y_out = reinterpret_cast<__nv_bfloat16*>(y.data_ptr()); + prm.ar_ws = reinterpret_cast(arWs); + prm.ar_flags = reinterpret_cast(arFlags); + prm.ar_rank = static_cast(arRank); + prm.ar_world = static_cast(arWorld); + prm.ar_maxt = static_cast(arMaxTokens); + prm.ar_mode = static_cast(arMode); + + cudaError_t const err + = tensorrt_llm::kernels::invokeMinimaxM3AttnMega(prm, qkvWeight.data_ptr(), oWeight.data_ptr(), + static_cast(decodeQueryLen), static_cast(tokensPerCluster), static_cast(blocksPerCta), + static_cast(numReqs), usePdl, at::cuda::getCurrentCUDAStream(xq.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "minimax_m3_attn_mega launch failed (qlen=", decodeQueryLen, " qg=", tokensPerCluster, + " bpc=", blocksPerCta, "): ", cudaGetErrorString(err)); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_attn_mega(Tensor xq, Tensor xsf, Tensor qkv_weight, Tensor qkv_weight_scale, Tensor(a!) qkv_out, " + "Tensor(b!) flags, Tensor(c!) kv_data, Tensor(d!) kv_scale, Tensor(e!) idx_k, Tensor out_cache_loc, " + "Tensor position_ids, Tensor block_table, Tensor seq_lens, Tensor kv_quant_scale, Tensor k_global_scale, " + "Tensor v_global_scale, Tensor q_norm_weight, Tensor k_norm_weight, Tensor index_k_norm_weight, " + "Tensor rotary_cos_sin, float eps, float sm_scale, int decode_query_len, int tokens_per_cluster, " + "int blocks_per_cta, Tensor(f!) output, Tensor(g!) mxfp8_data, Tensor(h!) mxfp8_scale, bool use_pdl, " + "Tensor o_weight, Tensor o_weight_scale, Tensor(i!) y, int ar_workspace, int ar_flags, int ar_rank, " + "int ar_world, int ar_max_tokens, int ar_mode) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_attn_mega", &tensorrt_llm::torch_ext::minimaxM3AttnMega); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3DraftStepOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3DraftStepOp.cpp new file mode 100644 index 000000000000..6bfc88d79213 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3DraftStepOp.cpp @@ -0,0 +1,198 @@ +/* + * 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/minimaxM3/draftStep.h" + +#include +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace ds = tensorrt_llm::kernels::minimaxM3::draftStep; + +using OptTensor = std::optional; + +void checkArch(torch::Tensor const& t) +{ + TORCH_CHECK(t.is_cuda(), "minimax_m3_draft_step expects CUDA tensors"); + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_draft_step requires SM100 or newer"); +} + +void check(cudaError_t err) +{ + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, cudaGetErrorString(err)); + } +} + +void checkInt(torch::Tensor const& t, int64_t n, char const* name) +{ + TORCH_CHECK(t.scalar_type() == at::kInt && t.is_contiguous() && t.numel() >= n, "draft_step: ", name, + " must be contiguous int32 with at least ", n, " elements"); +} + +} // namespace + +//! One fused greedy draft pass over `rows` <= 8 requests: the BF16 head (this rank's fragment-ordered shard wf +//! [N / 16, K / 16, 256], or with excOff / excPos / excW / expBase its 12-bit coding, wf uint8 [N * K * 3 / 2]; +//! global ids vocabOffset..) over x[xRows[b]] (or x[b]), the global argmax across the TP ranks through the multicast +//! slots wsMcPtr / wsUc, then draftOut[b, step] = idsNext[b] = token, posNext[b] = posSrc[xRows[b] or b] + 1 and the +//! KV-length bookkeeping (see draftStep.h). +void minimaxM3DraftStep(torch::Tensor const& x, OptTensor const& xRows, int64_t rows, torch::Tensor const& wf, + int64_t vocabOffset, torch::Tensor& keys, torch::Tensor& counters, torch::Tensor& wsUc, int64_t wsMcPtr, + int64_t rank, int64_t world, torch::Tensor& draftOut, int64_t step, torch::Tensor& idsNext, + torch::Tensor const& posSrc, torch::Tensor& posNext, torch::Tensor& kvLens, OptTensor const& numAcc, + OptTensor const& seqLens, int64_t draftLen, int64_t numCtas, bool pdl, OptTensor const& outKeys, + OptTensor const& excOff, OptTensor const& excPos, OptTensor const& excW, int64_t expBase) +{ + checkArch(x); + for (torch::Tensor const* t : std::initializer_list{ + &wf, &keys, &counters, &wsUc, &draftOut, &idsNext, &posSrc, &posNext, &kvLens}) + { + TORCH_CHECK(t->device() == x.device(), "draft_step: every tensor must be on x's device"); + } + for (OptTensor const* t : + std::initializer_list{&xRows, &numAcc, &seqLens, &outKeys, &excOff, &excPos, &excW}) + { + TORCH_CHECK(!t->has_value() || (*t)->device() == x.device(), "draft_step: every tensor must be on x's device"); + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + TORCH_CHECK(x.scalar_type() == at::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "draft_step: x must be contiguous BF16 [rows, K]"); + int const K = static_cast(x.size(1)); + TORCH_CHECK(K % 16 == 0, "draft_step: K must be a multiple of 16"); + TORCH_CHECK(rows >= 1 && rows <= ds::kMaxTokens, "draft_step: 1..", ds::kMaxTokens, " rows"); + bool const coded = excOff.has_value(); + TORCH_CHECK(coded == excPos.has_value() && coded == excW.has_value(), + "draft_step: exc_off, exc_pos and exc_w come together"); + int nTiles = 0; + if (coded) + { + TORCH_CHECK(K % 64 == 0 && K / 16 <= 0xFFFF, "draft_step: a coded head needs K a multiple of 64"); + TORCH_CHECK(excOff->scalar_type() == at::kInt && excOff->is_contiguous() && excOff->numel() >= 2, + "draft_step: exc_off must be int32 [tiles + 1]"); + nTiles = static_cast(excOff->numel() - 1); + TORCH_CHECK(wf.scalar_type() == at::kByte && wf.is_contiguous() + && wf.numel() == static_cast(nTiles) * 16 * K * 3 / 2, + "draft_step: a coded wf must be the contiguous uint8 stream of [tiles * 16, K]"); + TORCH_CHECK(excPos->scalar_type() == at::kInt && excPos->is_contiguous(), "draft_step: exc_pos must be int32"); + TORCH_CHECK(excW->scalar_type() == at::kFloat && excW->is_contiguous() && excW->numel() == 2 * excPos->numel(), + "draft_step: exc_w must be fp32 [E, 2]"); + TORCH_CHECK(expBase >= 0 && expBase <= 240, "draft_step: exponent base out of range"); + } + else + { + TORCH_CHECK(wf.scalar_type() == at::kBFloat16 && wf.dim() == 3 && wf.is_contiguous() && wf.size(1) == K / 16 + && wf.size(2) == 256, + "draft_step: wf must be the fragment-ordered BF16 weight [N / 16, K / 16, 256]"); + nTiles = static_cast(wf.size(0)); + } + TORCH_CHECK(nTiles >= 1 && numCtas >= 1 && (nTiles + numCtas - 1) / numCtas <= ds::kMaxTilesPerCta, + "draft_step: too many weight tiles per CTA"); + TORCH_CHECK(vocabOffset >= 0 && vocabOffset + 16LL * nTiles <= 0xFFFFFFFFLL, "draft_step: vocab offset"); + if (xRows.has_value()) + { + TORCH_CHECK(xRows->scalar_type() == at::kLong && xRows->is_contiguous() && xRows->numel() >= rows, + "draft_step: x_rows must be contiguous int64 [rows]"); + } + else + { + TORCH_CHECK(x.size(0) >= rows, "draft_step: x must hold a row per request"); + } + TORCH_CHECK(world >= 1 && rank >= 0 && rank < world, "draft_step: rank / world"); + TORCH_CHECK(keys.scalar_type() == at::kLong && keys.numel() >= rows, "draft_step: keys must be int64 [>= rows]"); + TORCH_CHECK(counters.scalar_type() == at::kInt && counters.numel() >= 2, "draft_step: counters must be int32 [2]"); + TORCH_CHECK(wsUc.scalar_type() == at::kLong && wsUc.numel() >= 2 * ds::kMaxTokens * world, + "draft_step: ws_uc must hold [2][8][world] slots"); + TORCH_CHECK(draftOut.scalar_type() == at::kInt && draftOut.dim() == 2 && draftOut.stride(1) == 1 + && draftOut.size(0) >= rows && step >= 0 && step < draftOut.size(1), + "draft_step: draft_out must be int32 [rows, > step] with unit column stride"); + checkInt(idsNext, rows, "ids_next"); + checkInt(posNext, rows, "pos_next"); + checkInt(kvLens, rows, "kv_lens"); + TORCH_CHECK(posSrc.scalar_type() == at::kInt && posSrc.is_contiguous(), "draft_step: pos_src must be int32"); + TORCH_CHECK(numAcc.has_value() == seqLens.has_value(), "draft_step: num_acc and seq_lens come together"); + TORCH_CHECK(!numAcc.has_value() || step == 0, "draft_step: the KV rollback (num_acc) is pass 0's"); + if (numAcc.has_value()) + { + checkInt(*numAcc, rows, "num_acc"); + checkInt(*seqLens, rows, "seq_lens"); + } + if (outKeys.has_value()) + { + TORCH_CHECK(outKeys->scalar_type() == at::kLong && outKeys->is_contiguous() && outKeys->numel() >= rows, + "draft_step: out_keys must be contiguous int64 [rows]"); + } + ds::Params p{}; + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.x_rows = xRows.has_value() ? reinterpret_cast(xRows->data_ptr()) : nullptr; + p.M = static_cast(rows); + p.wf = reinterpret_cast(wf.data_ptr()); + p.n_tiles = nTiles; + p.vocab_offset = static_cast(vocabOffset); + p.keys = reinterpret_cast(keys.data_ptr()); + p.counters = reinterpret_cast(counters.data_ptr()); + p.ws_uc = reinterpret_cast(wsUc.data_ptr()); + p.ws_mc = reinterpret_cast(wsMcPtr); + p.rank = static_cast(rank); + p.world = static_cast(world); + p.out_keys = outKeys.has_value() ? reinterpret_cast(outKeys->data_ptr()) : nullptr; + p.st = ds::StepArgs{draftOut.data_ptr(), static_cast(draftOut.stride(0)), static_cast(step), + idsNext.data_ptr(), posSrc.data_ptr(), posNext.data_ptr(), kvLens.data_ptr(), + numAcc.has_value() ? numAcc->data_ptr() : nullptr, + seqLens.has_value() ? seqLens->data_ptr() : nullptr, static_cast(draftLen)}; + p.cd = ds::CodedArgs{coded ? excOff->data_ptr() : nullptr, + coded ? reinterpret_cast(excPos->data_ptr()) : nullptr, + coded ? reinterpret_cast(excW->data_ptr()) : nullptr, + static_cast((expBase << 7) | (expBase << 23))}; + check(ds::invokeDraftStep(p, K, static_cast(numCtas), pdl, at::cuda::getCurrentCUDAStream(x.get_device()))); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +// Mutations through the raw multicast address ws_mc_ptr are not visible to the schema: that buffer is owned by the +// Python workspace and never aliased by the op's tensor arguments. +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_draft_step(Tensor x, Tensor? x_rows, int rows, Tensor wf, int vocab_offset, Tensor(a!) keys, " + "Tensor(b!) counters, Tensor(c!) ws_uc, int ws_mc_ptr, int rank, int world, Tensor(d!) draft_out, int step, " + "Tensor(e!) ids_next, Tensor pos_src, Tensor(f!) pos_next, Tensor(g!) kv_lens, Tensor? num_acc, " + "Tensor(h!)? seq_lens, int draft_len, int num_ctas, bool pdl, Tensor(i!)? out_keys, Tensor? exc_off=None, " + "Tensor? exc_pos=None, Tensor? exc_w=None, int exp_base=0) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_draft_step", &tensorrt_llm::torch_ext::minimaxM3DraftStep); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3Fp8IndexerOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Fp8IndexerOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3Fp8IndexerOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Fp8IndexerOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3FusedDecodeAttnOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3FusedDecodeAttnOp.cpp new file mode 100644 index 000000000000..e431a62f0028 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3FusedDecodeAttnOp.cpp @@ -0,0 +1,175 @@ +/* + * 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/minimaxM3/fusedDecodeAttn.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +using tensorrt_llm::kernels::MinimaxM3FusedDecodeAttnParams; + +constexpr int64_t kNumHeads = 16; +constexpr int64_t kHeadDim = 128; +constexpr int64_t kPackedWidth = (kNumHeads + 4) * kHeadDim; + +void checkByteTensor(at::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.element_size() == 1, name, " must be a CUDA tensor of one-byte elements"); +} + +} // namespace + +//! Fused per-head Gemma RMSNorm + partial NeoX RoPE + NVFP4 K/V and FP8 index-K cache insert + identity-selection +//! sparse decode attention over the paged NVFP4 cache + split-K merge, for the MiniMax-M3 MSA sparse layers in +//! short-sequence decode (every request attends at most 16 KV blocks of 128 tokens). +//! +//! @param packed bf16 [T, 2560] QKV+index projection rows ([Q x16 | K | V | idxQ | idxK] x 128), T = R * qlen. +//! @param kvData uint8 [pages, 2, 1, 128, 64] HND NVFP4 pool (K plane 0, V plane 1); the new tokens are written. +//! @param kvScale uint8 [pages, 2, 1, 128, 8] E4M3 block scales (K linear, V 4x4 token quads); written. +//! @param idxK one-byte [pages, 1, 128, 128] FP8 index-K cache; written. +//! @param output bf16 [T, 2048] attention output; mxfp8Data / mxfp8Scale its MXFP8 copy (E4M3 rows and 128x4-swizzled +//! UE8M0 scales) for the o_proj GEMM. +void minimaxM3FusedDecodeAttn(at::Tensor const& packed, at::Tensor& kvData, at::Tensor& kvScale, at::Tensor& idxK, + at::Tensor const& outCacheLoc, at::Tensor const& positionIds, at::Tensor const& blockTable, + at::Tensor const& seqLens, at::Tensor const& kvQuantScale, at::Tensor const& kGlobalScale, + at::Tensor const& vGlobalScale, at::Tensor const& qNormWeight, at::Tensor const& kNormWeight, + at::Tensor const& indexKNormWeight, at::Tensor const& rotaryCosSin, double eps, double smScale, + int64_t decodeQueryLen, int64_t tokensPerCluster, int64_t blocksPerCta, at::Tensor& output, at::Tensor& mxfp8Data, + at::Tensor& mxfp8Scale, bool usePdl) +{ + TORCH_CHECK(packed.is_cuda() && packed.scalar_type() == at::kBFloat16 && packed.dim() == 2 + && packed.size(1) == kPackedWidth && packed.stride(1) == 1, + "packed must be CUDA bf16 [T, 2560] with contiguous rows"); + checkByteTensor(kvData, "kv_data"); + checkByteTensor(kvScale, "kv_scale"); + checkByteTensor(idxK, "idx_k"); + TORCH_CHECK( + kvData.dim() == 5 && kvScale.dim() == 5 && idxK.dim() == 4, "kv_data, kv_scale and idx_k must be paged"); + for (auto const* t : {&outCacheLoc, &positionIds, &blockTable, &seqLens}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kInt, + "slot, position, block-table and length tensors " + "must be CUDA int32"); + } + TORCH_CHECK(blockTable.dim() == 2 && blockTable.stride(1) == 1, "block_table must be row-major 2-D"); + for (auto const* t : {&kvQuantScale, &kGlobalScale, &vGlobalScale, &rotaryCosSin}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kFloat, "scales and rotary_cos_sin must be CUDA fp32"); + } + for (auto const* t : {&qNormWeight, &kNormWeight, &indexKNormWeight}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kBFloat16 && t->numel() == kHeadDim, + "norm weights must be CUDA bf16 [128]"); + } + TORCH_CHECK(output.is_cuda() && output.scalar_type() == at::kBFloat16 && output.dim() == 2 + && output.size(1) == kNumHeads * kHeadDim && output.stride(1) == 1, + "output must be CUDA bf16 [T, 2048] with contiguous rows"); + checkByteTensor(mxfp8Data, "mxfp8_data"); + checkByteTensor(mxfp8Scale, "mxfp8_scale"); + int64_t const numTokens = packed.size(0); + int64_t const numReqs = seqLens.size(0); + TORCH_CHECK(numTokens == numReqs * decodeQueryLen, "packed rows must equal num_requests * decode_query_len"); + TORCH_CHECK(output.size(0) >= numTokens && mxfp8Data.dim() == 2 && mxfp8Data.size(0) >= numTokens + && mxfp8Data.size(1) == kNumHeads * kHeadDim && mxfp8Data.is_contiguous() && mxfp8Scale.is_contiguous(), + "mxfp8_data must be contiguous [>= T, 2048]"); + TORCH_CHECK(tensorrt_llm::kernels::minimaxM3FusedDecodeAttnSupports(static_cast(decodeQueryLen), + static_cast(tokensPerCluster), static_cast(blocksPerCta)), + "minimax_m3_fused_decode_attn: geometry (", decodeQueryLen, ", ", tokensPerCluster, ", ", blocksPerCta, + ") is not compiled"); + if (numTokens == 0) + { + return; + } + + c10::cuda::CUDAGuard const deviceGuard{packed.device()}; + auto const* prop = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(prop->major == 10, "minimax_m3_fused_decode_attn requires an sm_100-family GPU, got sm_", prop->major, + prop->minor); + + MinimaxM3FusedDecodeAttnParams prm = {}; + prm.packed = reinterpret_cast<__nv_bfloat16 const*>(packed.data_ptr()); + prm.packed_stride = packed.stride(0); + prm.kv_data = reinterpret_cast(kvData.data_ptr()); + prm.d_page = kvData.stride(0); + prm.d_plane = kvData.stride(1); + prm.kv_scale = reinterpret_cast(kvScale.data_ptr()); + prm.s_page = kvScale.stride(0); + prm.s_plane = kvScale.stride(1); + prm.idx_k = reinterpret_cast(idxK.data_ptr()); + prm.i_page = idxK.stride(0); + prm.i_tok = idxK.stride(2); + prm.num_pages = kvData.size(0); + prm.out_cache_loc = outCacheLoc.data_ptr(); + prm.position_ids = positionIds.data_ptr(); + prm.block_table = blockTable.data_ptr(); + prm.bt_stride = blockTable.stride(0); + prm.bt_cols = static_cast(blockTable.size(1)); + prm.seq_lens = seqLens.data_ptr(); + prm.kv_quant_scale = kvQuantScale.data_ptr(); + prm.k_gs = kGlobalScale.data_ptr(); + prm.v_gs = vGlobalScale.data_ptr(); + prm.q_w = reinterpret_cast<__nv_bfloat16 const*>(qNormWeight.data_ptr()); + prm.k_w = reinterpret_cast<__nv_bfloat16 const*>(kNormWeight.data_ptr()); + prm.ik_w = reinterpret_cast<__nv_bfloat16 const*>(indexKNormWeight.data_ptr()); + prm.rope = rotaryCosSin.data_ptr(); + prm.eps = static_cast(eps); + prm.sm_scale = static_cast(smScale); + prm.out = reinterpret_cast<__nv_bfloat16*>(output.data_ptr()); + prm.out_stride = output.stride(0); + prm.q8 = reinterpret_cast(mxfp8Data.data_ptr()); + prm.sf = reinterpret_cast(mxfp8Scale.data_ptr()); + prm.num_sf_k_tiles = static_cast(output.size(1) / 128); + + cudaError_t const err = tensorrt_llm::kernels::invokeMinimaxM3FusedDecodeAttn(prm, static_cast(decodeQueryLen), + static_cast(tokensPerCluster), static_cast(blocksPerCta), static_cast(numReqs), usePdl, + at::cuda::getCurrentCUDAStream(packed.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "minimax_m3_fused_decode_attn launch failed (qlen=", decodeQueryLen, + " qg=", tokensPerCluster, " bpc=", blocksPerCta, "): ", cudaGetErrorString(err)); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_fused_decode_attn(Tensor packed, Tensor(a!) kv_data, Tensor(b!) kv_scale, Tensor(c!) idx_k, " + "Tensor out_cache_loc, Tensor position_ids, Tensor block_table, Tensor seq_lens, Tensor kv_quant_scale, " + "Tensor k_global_scale, Tensor v_global_scale, Tensor q_norm_weight, Tensor k_norm_weight, " + "Tensor index_k_norm_weight, Tensor rotary_cos_sin, float eps, float sm_scale, int decode_query_len, " + "int tokens_per_cluster, int blocks_per_cta, Tensor(d!) output, Tensor(e!) mxfp8_data, " + "Tensor(f!) mxfp8_scale, bool use_pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_fused_decode_attn", &tensorrt_llm::torch_ext::minimaxM3FusedDecodeAttn); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3McAllreduceNormOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3McAllreduceNormOp.cpp new file mode 100644 index 000000000000..3f4558f8a330 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3McAllreduceNormOp.cpp @@ -0,0 +1,159 @@ +/* + * 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/minimaxM3/mcAllreduceNorm.h" + +#include +#include +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kMaxTokens = 48; +constexpr int64_t kCluster = 8; +constexpr int64_t kMaxWorld = 8; + +} // namespace + +//! MiniMax-M3 multicast all-reduce + residual + RMSNorm (+ optional MoE finalize, + optional MXFP8 output). +//! +//! @param x BF16 [M, H] partial, or with `perm` the permuted FC2 output [rows, H]. +//! @param x2 Optional BF16 [M, H] second addend (the shared expert in finalize mode). +//! @param residual BF16 [M, H], 1 <= M <= 48. +//! @param normWeight BF16 [H]. +//! @param wsUc This rank's unicast view of the exchange workspace (int32 words); re-armed in place. +//! @param wsMcPtr Device address of the workspace's multicast mapping. +//! @param flags Int32 per-CTA call-parity counters, at least maxt * 8; updated in place. +//! @param cluster Cluster size; 8 is the only one built. +//! @param perm, scales Finalize mode: int32 [M, topk] permuted-row index and fp32 / bf16 [M, topk] routing weights. +//! @param maxt Token capacity of the workspace bucket (8, 16 or 48). +//! @param quant Also return the MXFP8 copy of the normed rows (E4M3 [M, H], swizzled UE8M0 scales). +//! @param prepushed 0, or 2 / 3 when producers already pushed this rank's rows (see the kernel header). +//! @param opts Epilogue variant; every variant is bit-identical. +//! @return [normed, updated residual] or, with `quant`, [normed, updated residual, E4M3 rows, scale bytes]. +std::vector minimaxM3McAllreduceNorm(torch::Tensor const& x, std::optional const& x2, + torch::Tensor const& residual, torch::Tensor const& normWeight, double eps, torch::Tensor& wsUc, int64_t wsMcPtr, + torch::Tensor& flags, int64_t rank, int64_t world, int64_t cluster, bool pdl, + std::optional const& perm, std::optional const& scales, int64_t maxt, bool quant, + int64_t prepushed, int64_t opts) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sm == 100 || sm == 103, "minimax_m3_mc_allreduce_norm requires SM100 or SM103, got SM", sm); + TORCH_CHECK(cluster == kCluster, "minimax_m3_mc_allreduce_norm: only cluster 8 is built, got ", cluster); + TORCH_CHECK(tensorrt_llm::kernels::isMinimaxM3McAllreduceNormSupported( + static_cast(maxt), static_cast(opts)), + "minimax_m3_mc_allreduce_norm: unsupported (bucket ", maxt, ", opts ", opts, ")"); + TORCH_CHECK( + world >= 1 && world <= kMaxWorld && rank >= 0 && rank < world, "invalid rank ", rank, " / world ", world); + + tensorrt_llm::kernels::MinimaxM3McAllreduceNormParams p; + if (x2.has_value()) + { + TORCH_CHECK(x2->scalar_type() == at::kBFloat16 && x2->is_contiguous() && x2->sizes() == residual.sizes(), + "x2 must be a contiguous bf16 tensor shaped like residual"); + p.x2 = reinterpret_cast<__nv_bfloat16 const*>(x2->data_ptr()); + } + if (perm.has_value()) + { + // Finalize mode: x is the permuted FC2 output [rows, H]. + TORCH_CHECK(scales.has_value() && perm->scalar_type() == at::kInt && perm->is_contiguous() && perm->dim() == 2 + && perm->size(0) == residual.size(0), + "perm must be a contiguous int32 [M, topk] tensor, with scales"); + TORCH_CHECK(scales->is_contiguous() && scales->numel() == perm->numel() + && (scales->scalar_type() == at::kFloat || scales->scalar_type() == at::kBFloat16), + "scales must be a contiguous fp32 or bf16 tensor with one weight per perm entry"); + TORCH_CHECK(x.dim() == 2 && x.size(1) == residual.size(1), "fc2 must be [rows, H]"); + p.fc2 = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.perm = perm->data_ptr(); + p.scales = scales->data_ptr(); + p.scaleFp32 = scales->scalar_type() == at::kFloat; + p.topk = static_cast(perm->size(1)); + } + TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.is_contiguous() && x.dim() == 2, + "x must be a contiguous 2-D bf16 CUDA tensor"); + TORCH_CHECK(residual.dim() == 2 && residual.size(0) >= 1 && residual.size(0) <= std::min(kMaxTokens, maxt), + "residual must hold 1..min(48, bucket) rows"); + TORCH_CHECK(residual.scalar_type() == at::kBFloat16 && residual.is_contiguous() + && (perm.has_value() || residual.sizes() == x.sizes()), + "residual must be a contiguous bf16 tensor shaped like x"); + TORCH_CHECK( + normWeight.scalar_type() == at::kBFloat16 && normWeight.numel() == residual.size(1), "norm_w must be bf16 [H]"); + TORCH_CHECK(wsUc.scalar_type() == at::kInt && wsUc.is_cuda(), "ws_uc must be an int32 CUDA tensor"); + TORCH_CHECK(flags.scalar_type() == at::kInt && flags.is_contiguous() && flags.numel() >= maxt * kCluster, + "flags must be a contiguous int32 tensor of at least bucket * 8 counters"); + + int64_t const m = residual.size(0), h = residual.size(1); + TORCH_CHECK(!quant || h % 128 == 0, "MXFP8 output needs H % 128 == 0"); + c10::cuda::CUDAGuard const deviceGuard{residual.device()}; + auto outNormed = torch::empty_like(residual); + auto outResidual = torch::empty_like(residual); + auto outQ = quant ? torch::empty({m, h}, residual.options().dtype(at::kFloat8_e4m3fn)) : torch::Tensor(); + auto outSf = quant ? torch::empty({(m + 127) / 128 * 128 * (h / 32)}, residual.options().dtype(at::kByte)) + : torch::Tensor(); + + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.numTokens = static_cast(m); + p.hidden = static_cast(h); + p.residual = reinterpret_cast<__nv_bfloat16 const*>(residual.data_ptr()); + p.normWeight = reinterpret_cast<__nv_bfloat16 const*>(normWeight.data_ptr()); + p.eps = static_cast(eps); + p.wsUc = wsUc.data_ptr(); + p.wsMc = reinterpret_cast(wsMcPtr); + p.flags = flags.data_ptr(); + p.rank = static_cast(rank); + p.world = static_cast(world); + p.outNormed = reinterpret_cast<__nv_bfloat16*>(outNormed.data_ptr()); + p.outResidual = reinterpret_cast<__nv_bfloat16*>(outResidual.data_ptr()); + p.outQuant = quant ? reinterpret_cast(outQ.data_ptr()) : nullptr; + p.outSf = quant ? outSf.data_ptr() : nullptr; + p.prepushed = static_cast(prepushed); + tensorrt_llm::kernels::invokeMinimaxM3McAllreduceNorm(p, static_cast(maxt), static_cast(opts), + pdl, at::cuda::getCurrentCUDAStream(residual.get_device())); + if (quant) + { + return {outNormed, outResidual, outQ, outSf}; + } + return {outNormed, outResidual}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_mc_allreduce_norm(Tensor x, Tensor? x2, Tensor residual, Tensor norm_w, float eps, " + "Tensor(a!) ws_uc, int ws_mc_ptr, Tensor(b!) flags, int rank, int world, int cluster, bool pdl, " + "Tensor? perm, Tensor? scales, int maxt, bool quant, int prepushed, int opts) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_mc_allreduce_norm", &tensorrt_llm::torch_ext::minimaxM3McAllreduceNorm); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3MsaKvUpdateOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3MsaKvUpdateOp.cpp new file mode 100644 index 000000000000..dfea9d3af430 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3MsaKvUpdateOp.cpp @@ -0,0 +1,110 @@ +/* + * 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/minimaxM3/msaKvUpdate.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +void checkInt32Cuda(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda(), "minimax_m3_msa_kv_update: ", name, " must be a CUDA tensor"); + TORCH_CHECK( + t.scalar_type() == torch::kInt32, "minimax_m3_msa_kv_update: ", name, " must be int32, got ", t.scalar_type()); +} + +} // namespace + +//! MiniMax-M3 MSA KV-length correction after the overlap scheduler moved kv_lens (see invokeMinimaxM3MsaKvUpdate). +//! +//! Every tensor is int32 on one CUDA device. block_table is [rows, ncols] with unit column stride; every other +//! operand is contiguous with at least `batch` (per request) or `total_q` (per query token) elements, and cu_kv_lens +//! at least `batch + 1`. Writes msa_seq_lens, out_cache_loc and, when given, cu_kv_lens and n_valid_blocks. +void minimaxM3MsaKvUpdate(torch::Tensor const& kvLens, torch::Tensor const& kvLensStaged, torch::Tensor const& seqLens, + torch::Tensor const& qBatchRow, torch::Tensor const& qIntra, torch::Tensor const& blockTable, + int64_t tokensPerBlock, int64_t batch, int64_t totalQ, torch::Tensor& msaSeqLens, + std::optional cuKvLens, torch::Tensor& outCacheLoc, std::optional nValidBlocks) +{ + checkInt32Cuda(kvLens, "kv_lens"); + checkInt32Cuda(kvLensStaged, "kv_lens_staged"); + checkInt32Cuda(seqLens, "seq_lens"); + checkInt32Cuda(qBatchRow, "q_batch_row"); + checkInt32Cuda(qIntra, "q_intra"); + checkInt32Cuda(blockTable, "block_table"); + checkInt32Cuda(msaSeqLens, "msa_seq_lens"); + checkInt32Cuda(outCacheLoc, "out_cache_loc"); + TORCH_CHECK(blockTable.dim() == 2 && blockTable.stride(1) == 1, + "minimax_m3_msa_kv_update: block_table must be 2-D with unit column stride"); + TORCH_CHECK(tokensPerBlock > 0, "minimax_m3_msa_kv_update: tokens_per_block must be positive"); + TORCH_CHECK(batch >= 0 && totalQ >= 0, "minimax_m3_msa_kv_update: batch and total_q must be non-negative"); + TORCH_CHECK(kvLens.numel() >= batch && kvLensStaged.numel() >= batch && seqLens.numel() >= batch + && msaSeqLens.numel() >= batch, + "minimax_m3_msa_kv_update: a per-request operand has fewer than batch elements"); + TORCH_CHECK(qBatchRow.numel() >= totalQ && qIntra.numel() >= totalQ && outCacheLoc.numel() >= totalQ, + "minimax_m3_msa_kv_update: a per-token operand has fewer than total_q elements"); + int32_t* cuKvPtr = nullptr; + if (cuKvLens.has_value()) + { + checkInt32Cuda(*cuKvLens, "cu_kv_lens"); + TORCH_CHECK(cuKvLens->numel() >= batch + 1, "minimax_m3_msa_kv_update: cu_kv_lens needs batch + 1 elements"); + cuKvPtr = cuKvLens->data_ptr(); + } + int32_t* nValidPtr = nullptr; + if (nValidBlocks.has_value()) + { + checkInt32Cuda(*nValidBlocks, "n_valid_blocks"); + TORCH_CHECK(nValidBlocks->numel() >= totalQ, "minimax_m3_msa_kv_update: n_valid_blocks needs total_q elements"); + nValidPtr = nValidBlocks->data_ptr(); + } + + c10::cuda::CUDAGuard const deviceGuard{kvLens.device()}; + auto const stream = at::cuda::getCurrentCUDAStream(kvLens.get_device()); + cudaError_t const err = tensorrt_llm::kernels::invokeMinimaxM3MsaKvUpdate(kvLens.data_ptr(), + kvLensStaged.data_ptr(), seqLens.data_ptr(), qBatchRow.data_ptr(), + qIntra.data_ptr(), blockTable.data_ptr(), blockTable.stride(0), + static_cast(blockTable.size(1)), static_cast(tokensPerBlock), static_cast(batch), + static_cast(totalQ), msaSeqLens.data_ptr(), cuKvPtr, outCacheLoc.data_ptr(), + nValidPtr, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_msa_kv_update launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_msa_kv_update(Tensor kv_lens, Tensor kv_lens_staged, Tensor seq_lens, Tensor q_batch_row, " + "Tensor q_intra, Tensor block_table, int tokens_per_block, int batch, int total_q, Tensor(a!) msa_seq_lens, " + "Tensor(b!)? cu_kv_lens, Tensor(c!) out_cache_loc, Tensor(d!)? n_valid_blocks) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_msa_kv_update", &tensorrt_llm::torch_ext::minimaxM3MsaKvUpdate); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Mxfp8LowMGemmOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Mxfp8LowMGemmOp.cpp new file mode 100644 index 000000000000..c8f6b4a0355d --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Mxfp8LowMGemmOp.cpp @@ -0,0 +1,83 @@ +/* + * 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/minimaxM3/mxfp8LowMGemm.h" + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! MiniMax-M3 fused MXFP8 activation quantization + low-M MXFP8 GEMM (decode). +//! +//! @param x BF16 [M, K] activation, unit column stride, 16-byte aligned rows (row stride a multiple of 8). +//! @param w E4M3 [N, K] weight, contiguous. +//! @param wsf UE8M0 weight scales in the 128x4 swizzled layout of `mxfp8_quantize(w, True)`. +//! @param warps, clusterK, rowTiles The kernel configuration; see `isMinimaxM3Mxfp8LowMGemmSupported`. +//! @return BF16 [M, N]. +torch::Tensor minimaxM3Mxfp8LowMGemm(torch::Tensor const& x, torch::Tensor const& w, torch::Tensor const& wsf, + int64_t warps, int64_t clusterK, int64_t rowTiles, bool pdl) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sm == 100 || sm == 103, "minimax_m3_mxfp8_low_m_gemm requires SM100 or SM103, got SM", sm); + TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.dim() == 2 && x.stride(1) == 1, + "x must be a 2-D bf16 CUDA tensor with unit column stride"); + TORCH_CHECK(w.scalar_type() == at::kFloat8_e4m3fn && w.is_contiguous() && w.dim() == 2 && w.size(1) == x.size(1), + "w must be a contiguous e4m3 [N, K] tensor"); + TORCH_CHECK(wsf.scalar_type() == at::kByte, "wsf must be uint8 UE8M0 scales"); + TORCH_CHECK( + reinterpret_cast(x.data_ptr()) % 16 == 0 && x.stride(0) % 8 == 0, "x rows must be 16-byte aligned"); + TORCH_CHECK( + reinterpret_cast(w.data_ptr()) % 16 == 0 && reinterpret_cast(wsf.data_ptr()) % 16 == 0, + "w and wsf must be 16-byte aligned"); + TORCH_CHECK(tensorrt_llm::kernels::isMinimaxM3Mxfp8LowMGemmSupported(static_cast(warps), + static_cast(clusterK), static_cast(rowTiles), static_cast(x.size(0))), + "minimax_m3_mxfp8_low_m_gemm: unsupported configuration (warps ", warps, ", cluster_k ", clusterK, + ", row_tiles ", rowTiles, ") at M=", x.size(0)); + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto out = torch::empty({x.size(0), w.size(0)}, x.options()); + tensorrt_llm::kernels::invokeMinimaxM3Mxfp8LowMGemm(reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()), + static_cast(x.stride(0)), static_cast(x.size(0)), static_cast(w.size(0)), + static_cast(x.size(1)), reinterpret_cast(w.data_ptr()), + reinterpret_cast(wsf.data_ptr()), reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), + static_cast(out.stride(0)), static_cast(warps), static_cast(clusterK), + static_cast(rowTiles), pdl, at::cuda::getCurrentCUDAStream(x.get_device())); + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_mxfp8_low_m_gemm(Tensor x, Tensor w, Tensor wsf, int warps, int cluster_k, int row_tiles, " + "bool pdl) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_mxfp8_low_m_gemm", &tensorrt_llm::torch_ext::minimaxM3Mxfp8LowMGemm); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3ProducerPwOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3ProducerPwOp.cpp new file mode 100644 index 000000000000..ec73b498e40e --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3ProducerPwOp.cpp @@ -0,0 +1,181 @@ +/* + * 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/minimaxM3/producerPw.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace pw = tensorrt_llm::kernels::minimaxM3::producerPw; + +constexpr int64_t kHeadDim = 128; +constexpr int64_t kPageSize = 128; + +void checkCuda(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda(), "producer_pw: ", name, " must be a CUDA tensor"); +} + +} // namespace + +//! The NVFP4 QKV + indexer norm/RoPE/KV-insert producer (same arguments and outputs as +//! minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert for 1 KV head, 1 index head, head_dim 128, rotary_dim 64): +//! returns (q FP8 [T, headsQ, 128], index_q FP8 [T, 1, 128]) and writes the NVFP4 K/V and FP8 index-K of each token's +//! cache slot. `early` triggers the dependents on entry. +std::vector minimaxM3ProducerPw(torch::Tensor const& packed, torch::Tensor& kvData, + torch::Tensor& kvScale, torch::Tensor& indexK, torch::Tensor const& outCacheLoc, torch::Tensor const& kvQuantScale, + double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, torch::Tensor const& indexQWeight, + torch::Tensor const& indexKWeight, torch::Tensor const& rotaryCosSin, torch::Tensor const& positionIds, + int64_t headsQ, bool pdl, bool early) +{ + checkCuda(packed, "packed"); + for (torch::Tensor const* t : std::initializer_list{&kvData, &kvScale, &indexK, &outCacheLoc, + &kvQuantScale, &qWeight, &kWeight, &indexQWeight, &indexKWeight, &rotaryCosSin, &positionIds}) + { + TORCH_CHECK(t->is_cuda() && t->get_device() == packed.get_device(), + "producer_pw: every operand must be on packed's CUDA device"); + } + TORCH_CHECK(reinterpret_cast(packed.data_ptr()) % 8 == 0 + && reinterpret_cast(kvData.data_ptr()) % 2 == 0 + && reinterpret_cast(indexK.data_ptr()) % 4 == 0, + "producer_pw: packed and the caches must start at the vector-load / store alignment"); + c10::cuda::CUDAGuard const deviceGuard{packed.device()}; + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_producer_pw requires SM100 or newer"); + TORCH_CHECK(headsQ == 8 || headsQ == 16, "producer_pw: 8 or 16 query heads"); + int64_t const totalHeads = headsQ + 4; + TORCH_CHECK(packed.scalar_type() == at::kBFloat16 && packed.dim() == 2 && packed.size(1) == totalHeads * kHeadDim + && packed.is_contiguous(), + "producer_pw: packed must be contiguous BF16 [T, (headsQ + 4) * 128]"); + // The kernel's layout (the native producer's checkMinimaxM3Nvfp4HndKVPools at one KV head): HND pools + // [pages, 2 planes, 1 head, 128 tokens, 64 data / 8 scale bytes], dense token rows and head blocks, planes and + // pages that do not overlap. + TORCH_CHECK(kvData.scalar_type() == at::kByte && kvData.dim() == 5 && kvData.size(0) > 0 && kvData.size(1) == 2 + && kvData.size(2) == 1 && kvData.size(3) == kPageSize && kvData.size(4) == kHeadDim / 2, + "producer_pw: kv_data must be the NVFP4 pool [pages, 2, 1, 128, 64] uint8"); + TORCH_CHECK(kvScale.scalar_type() == at::kByte && kvScale.dim() == 5 && kvScale.size(0) == kvData.size(0) + && kvScale.size(1) == 2 && kvScale.size(2) == 1 && kvScale.size(3) == kPageSize + && kvScale.size(4) == kHeadDim / 16, + "producer_pw: kv_scale must be the E4M3 scale pool [pages, 2, 1, 128, 8] uint8"); + TORCH_CHECK(kvData.stride(4) == 1 && kvData.stride(3) == kvData.size(4) + && kvData.stride(2) == kvData.size(3) * kvData.stride(3), + "producer_pw: kv_data must have dense token rows and head blocks"); + TORCH_CHECK(kvScale.stride(4) == 1 && kvScale.stride(3) == kvScale.size(4) + && kvScale.stride(2) == kvScale.size(3) * kvScale.stride(3), + "producer_pw: kv_scale must have dense token rows and head blocks"); + TORCH_CHECK(kvData.stride(1) >= kvData.size(2) * kvData.stride(2) && kvData.stride(0) >= 2 * kvData.stride(1) + && kvScale.stride(1) >= kvScale.size(2) * kvScale.stride(2) && kvScale.stride(0) >= 2 * kvScale.stride(1), + "producer_pw: pool planes or pages overlap"); + TORCH_CHECK(kvData.stride(0) % 2 == 0 && kvData.stride(1) % 2 == 0, + "producer_pw: kv_data page and plane strides must keep the uint16 stores aligned"); + TORCH_CHECK(indexK.scalar_type() == at::kFloat8_e4m3fn && indexK.dim() == 4 && indexK.size(0) == kvData.size(0) + && indexK.size(1) == 1 && indexK.size(2) == kPageSize && indexK.size(3) == kHeadDim && indexK.stride(3) == 1 + && indexK.stride(2) == kHeadDim && indexK.stride(0) >= kPageSize * kHeadDim && indexK.stride(0) % 4 == 0, + "producer_pw: index_k must be FP8 [pages, 1, 128, 128] with a page stride keeping the 32-bit stores aligned"); + int64_t const numTokens = packed.size(0); + TORCH_CHECK( + outCacheLoc.scalar_type() == at::kInt && outCacheLoc.is_contiguous() && outCacheLoc.numel() >= numTokens, + "producer_pw: out_cache_loc must be int32 [>= T]"); + TORCH_CHECK( + positionIds.scalar_type() == at::kInt && positionIds.is_contiguous() && positionIds.numel() >= numTokens, + "producer_pw: position_ids must be int32 [>= T]"); + TORCH_CHECK(kvQuantScale.scalar_type() == at::kFloat && kvQuantScale.is_contiguous() && kvQuantScale.numel() >= 3, + "producer_pw: kv_quant_scale must be fp32 [>= 3]"); + TORCH_CHECK(rotaryCosSin.scalar_type() == at::kFloat && rotaryCosSin.is_contiguous() && rotaryCosSin.dim() == 3 + && rotaryCosSin.size(1) == 2 && rotaryCosSin.size(2) == 32 + && reinterpret_cast(rotaryCosSin.data_ptr()) % 16 == 0, + "producer_pw: rotary_cos_sin must be fp32 [positions, 2, 32], 16-byte aligned (read as float4)"); + for (auto const* w : {&qWeight, &kWeight, &indexQWeight, &indexKWeight}) + { + TORCH_CHECK(w->scalar_type() == at::kBFloat16 && w->is_contiguous() && w->numel() == kHeadDim + && reinterpret_cast(w->data_ptr()) % 8 == 0, + "producer_pw: norm weights must be BF16 [128], 8-byte aligned (read as 4-value words)"); + } + auto const fp8 = packed.options().dtype(at::kFloat8_e4m3fn); + auto qOut = torch::empty({numTokens, headsQ, kHeadDim}, fp8); + auto indexQOut = torch::empty({numTokens, 1, kHeadDim}, fp8); + if (numTokens == 0) + { + return {qOut, indexQOut}; + } + pw::Params p{}; + p.packed = reinterpret_cast<__nv_bfloat16 const*>(packed.data_ptr()); + p.qOut = reinterpret_cast<__nv_fp8_e4m3*>(qOut.data_ptr()); + p.indexQOut = reinterpret_cast<__nv_fp8_e4m3*>(indexQOut.data_ptr()); + p.kvData = kvData.data_ptr(); + p.kvScale = kvScale.data_ptr(); + p.indexK = reinterpret_cast<__nv_fp8_e4m3*>(indexK.data_ptr()); + p.outCacheLoc = outCacheLoc.data_ptr(); + p.kvQuantScale = kvQuantScale.data_ptr(); + p.dataPageStride = kvData.stride(0); + p.dataPlaneStride = kvData.stride(1); + p.dataHeadStride = kvData.stride(2); + p.dataTokenStride = kvData.stride(3); + p.scalePageStride = kvScale.stride(0); + p.scalePlaneStride = kvScale.stride(1); + p.scaleHeadStride = kvScale.stride(2); + p.indexPageStride = indexK.stride(0); + p.indexTokenStride = indexK.stride(2); + p.numPages = kvData.size(0); + p.numTokens = static_cast(numTokens); + p.eps = static_cast(eps); + p.qWeight = reinterpret_cast<__nv_bfloat16 const*>(qWeight.data_ptr()); + p.kWeight = reinterpret_cast<__nv_bfloat16 const*>(kWeight.data_ptr()); + p.indexQWeight = reinterpret_cast<__nv_bfloat16 const*>(indexQWeight.data_ptr()); + p.indexKWeight = reinterpret_cast<__nv_bfloat16 const*>(indexKWeight.data_ptr()); + p.rotaryCosSin = rotaryCosSin.data_ptr(); + p.positionIds = positionIds.data_ptr(); + cudaError_t const err = pw::invokeProducerPw(p, static_cast(numTokens), static_cast(headsQ), pdl, early, + at::cuda::getCurrentCUDAStream(packed.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "producer_pw launch: ", cudaGetErrorString(err)); + } + return {qOut, indexQOut}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_producer_pw(Tensor packed, Tensor(a!) kv_data, Tensor(b!) kv_scale, Tensor(c!) index_k, " + "Tensor out_cache_loc, Tensor kv_quant_scale, float eps, Tensor q_weight, Tensor k_weight, " + "Tensor index_q_weight, Tensor index_k_weight, Tensor rotary_cos_sin, Tensor position_ids, int heads_q, " + "bool pdl, bool early) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_producer_pw", &tensorrt_llm::torch_ext::minimaxM3ProducerPw); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3RouterGemvOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3RouterGemvOp.cpp new file mode 100644 index 000000000000..059d29e43d85 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3RouterGemvOp.cpp @@ -0,0 +1,149 @@ +/* + * 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/minimaxM3/routerGemv.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kGemvMaxM = 16; +constexpr int64_t kTcMaxM = 48; +constexpr int64_t kTcExperts = 128; + +void checkSmVersion(char const* op) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sm == 100 || sm == 103, op, " requires SM100 or SM103, got SM", sm); +} + +void checkActivation(torch::Tensor const& x, torch::Tensor const& w, int64_t maxM) +{ + TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.dim() == 2 && x.stride(1) == 1, + "x must be a 2-D bf16 CUDA tensor with unit column stride"); + TORCH_CHECK(x.size(0) >= 1 && x.size(0) <= maxM, "x must hold 1..", maxM, " rows, got ", x.size(0)); + TORCH_CHECK(w.scalar_type() == at::kFloat && w.is_contiguous() && w.dim() == 2 && w.size(1) == x.size(1), + "w must be a contiguous fp32 [E, K] tensor"); + TORCH_CHECK(x.size(1) % 1024 == 0 && (x.stride(0) * 2) % 16 == 0, "K must be a multiple of 1024 with 16-byte rows"); +} + +} // namespace + +//! fp32 router logits `x @ w^T` for 1..16 bf16 rows (exact fp32 FMA, one CTA per expert). +torch::Tensor minimaxM3RouterGemv(torch::Tensor const& x, torch::Tensor const& w, int64_t rows, bool pdl) +{ + checkSmVersion("minimax_m3_router_gemv"); + checkActivation(x, w, kGemvMaxM); + TORCH_CHECK(rows == 1, "minimax_m3_router_gemv: only rows=1 is built, got ", rows); + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto out = torch::empty({x.size(0), w.size(0)}, x.options().dtype(at::kFloat)); + tensorrt_llm::kernels::invokeMinimaxM3RouterGemv(reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()), + static_cast(x.stride(0)), static_cast(x.size(0)), static_cast(x.size(1)), + w.data_ptr(), static_cast(w.size(0)), out.data_ptr(), nullptr, nullptr, nullptr, pdl, + at::cuda::getCurrentCUDAStream(x.get_device())); + return out; +} + +//! Router logits as `minimax_m3_router_gemv` plus the NVFP4 activation: [logits fp32 [M, E], packed e2m1 +//! [M, K / 2] uint8, linear E4M3 scales [M * K / 16] uint8]. +std::vector minimaxM3RouterGemvFp4( + torch::Tensor const& x, torch::Tensor const& w, torch::Tensor const& sfScale, bool pdl) +{ + checkSmVersion("minimax_m3_router_gemv_fp4"); + checkActivation(x, w, kGemvMaxM); + TORCH_CHECK(sfScale.scalar_type() == at::kFloat && sfScale.numel() == 1 && sfScale.is_cuda(), + "sf_scale must be a one-element fp32 CUDA tensor"); + int64_t const M = x.size(0), K = x.size(1); + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto out = torch::empty({M, w.size(0)}, x.options().dtype(at::kFloat)); + auto q = torch::empty({M, K / 2}, x.options().dtype(at::kByte)); + auto sf = torch::empty({M * (K / 16)}, x.options().dtype(at::kByte)); + tensorrt_llm::kernels::invokeMinimaxM3RouterGemv(reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()), + static_cast(x.stride(0)), static_cast(M), static_cast(K), w.data_ptr(), + static_cast(w.size(0)), out.data_ptr(), q.data_ptr(), sf.data_ptr(), + sfScale.data_ptr(), pdl, at::cuda::getCurrentCUDAStream(x.get_device())); + return {out, q, sf}; +} + +//! TF32 tensor-core router logits for 128 experts and 1..48 rows: [logits fp32 [M, 128]] or, with `sf_scale`, +//! [logits, packed e2m1 [M, K / 2], linear E4M3 scales [M * K / 16]]. +std::vector minimaxM3RouterTc( + torch::Tensor const& x, torch::Tensor const& w, std::optional const& sfScale, bool pdl) +{ + checkSmVersion("minimax_m3_router_tc"); + TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.dim() == 2 && x.stride(1) == 1, + "x must be a 2-D bf16 CUDA tensor with unit column stride"); + int64_t const M = x.size(0), K = x.size(1); + TORCH_CHECK(M >= 1 && M <= kTcMaxM && w.scalar_type() == at::kFloat && w.dim() == 2 && w.size(0) == kTcExperts + && w.size(1) == K && w.is_contiguous(), + "minimax_m3_router_tc takes 1..48 rows and a contiguous fp32 [128, K] weight"); + TORCH_CHECK(K % 1024 == 0 && (x.stride(0) * 2) % 16 == 0, "K must be a multiple of 1024 with 16-byte rows"); + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto out = torch::empty({M, w.size(0)}, x.options().dtype(at::kFloat)); + torch::Tensor q, sf; + uint8_t* qp = nullptr; + uint8_t* sfp = nullptr; + float const* scp = nullptr; + if (sfScale.has_value()) + { + TORCH_CHECK(sfScale->scalar_type() == at::kFloat && sfScale->numel() == 1 && sfScale->is_cuda(), + "sf_scale must be a one-element fp32 CUDA tensor"); + q = torch::empty({M, K / 2}, x.options().dtype(at::kByte)); + sf = torch::empty({M * (K / 16)}, x.options().dtype(at::kByte)); + qp = q.data_ptr(); + sfp = sf.data_ptr(); + scp = sfScale->data_ptr(); + } + tensorrt_llm::kernels::invokeMinimaxM3RouterTc(reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()), + static_cast(x.stride(0)), static_cast(M), static_cast(K), w.data_ptr(), + out.data_ptr(), qp, sfp, scp, pdl, at::cuda::getCurrentCUDAStream(x.get_device())); + if (sfScale.has_value()) + { + return {out, q, sf}; + } + return {out}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_router_gemv(Tensor x, Tensor w, int rows, bool pdl) -> Tensor"); + m.def("minimax_m3_router_gemv_fp4(Tensor x, Tensor w, Tensor sf_scale, bool pdl) -> Tensor[]"); + m.def("minimax_m3_router_tc(Tensor x, Tensor w, Tensor? sf_scale, bool pdl) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_router_gemv", &tensorrt_llm::torch_ext::minimaxM3RouterGemv); + m.impl("minimax_m3_router_gemv_fp4", &tensorrt_llm::torch_ext::minimaxM3RouterGemvFp4); + m.impl("minimax_m3_router_tc", &tensorrt_llm::torch_ext::minimaxM3RouterTc); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SampledAcceptOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SampledAcceptOp.cpp new file mode 100644 index 000000000000..1c2183647374 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SampledAcceptOp.cpp @@ -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. + */ + +#include "tensorrt_llm/kernels/minimaxM3/sampledAccept.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace sa = tensorrt_llm::kernels::minimaxM3::sampledAccept; + +//! uint64 words of the sampled-accept exchange workspace for `world` ranks. +int64_t minimaxM3SampledAcceptWsWords(int64_t world) +{ + return sa::sampledAcceptWsWords(static_cast(world)); +} + +//! Sampled (temperature + top-p) verification from this rank's BF16 logits shard [rows, Vs], rows = nreq * +//! (draft_len + 1) <= 48: the sampled token of every row (tokens), then the strict acceptance against draft (accTokens, +//! numAcc, gatherIds, optional nextNew column 0). Per-row Philox (seeds, offsets); the exchange goes through the +//! multicast workspace wsMcPtr / wsUc with the (finished rows, call) counters. +void minimaxM3SampledAccept(torch::Tensor const& logits, int64_t vocabOffset, torch::Tensor const& temperatures, + torch::Tensor const& topPs, torch::Tensor const& seeds, torch::Tensor const& offsets, torch::Tensor& wsUc, + int64_t wsMcPtr, torch::Tensor& counters, torch::Tensor& tokens, int64_t rank, int64_t world, + torch::Tensor const& draft, torch::Tensor const& batchIdx, torch::Tensor& accTokens, torch::Tensor& numAcc, + torch::Tensor& gatherIds, int64_t forceTotal, std::optional const& nextNew, + std::optional const& forceVec, std::optional const& extraAttempts, bool usePdl) +{ + c10::cuda::CUDAGuard const deviceGuard{logits.device()}; + sa::Args a; + a.logits = reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()); + a.ld = logits.stride(0); + a.rows = static_cast(logits.size(0)); + a.Vs = static_cast(logits.size(1)); + TORCH_CHECK(a.Vs <= sa::kMaxVs && a.rows <= sa::kMaxRows && world <= sa::kMaxWorld + && static_cast(vocabOffset) + a.Vs <= (1LL << sa::kIdBits)); + a.vocab_offset = static_cast(vocabOffset); + a.temperatures = temperatures.data_ptr(); + a.top_ps = topPs.data_ptr(); + a.seeds = reinterpret_cast(seeds.data_ptr()); + a.offsets = reinterpret_cast(offsets.data_ptr()); + a.ws_uc = reinterpret_cast(wsUc.data_ptr()); + a.ws_mc = reinterpret_cast(wsMcPtr); + a.counters = reinterpret_cast(counters.data_ptr()); + a.tokens = tokens.data_ptr(); + a.extra_attempts = extraAttempts.has_value() ? extraAttempts->data_ptr() : nullptr; + a.rank = static_cast(rank); + a.world = static_cast(world); + a.draft = draft.data_ptr(); + a.batch_idx = batchIdx.data_ptr(); + a.acc_tokens = accTokens.data_ptr(); + a.num_acc = numAcc.data_ptr(); + a.gather_ids = reinterpret_cast(gatherIds.data_ptr()); + a.draft_len = static_cast(draft.size(1)); + a.nreq = static_cast(draft.size(0)); + a.force_total = static_cast(forceTotal); + a.next_new = nextNew.has_value() ? nextNew->data_ptr() : nullptr; + a.next_stride = nextNew.has_value() ? static_cast(nextNew->stride(0)) : 0; + a.force_vec = forceVec.has_value() ? forceVec->data_ptr() : nullptr; + cudaError_t const err = sa::invokeSampledAccept(a, usePdl, at::cuda::getCurrentCUDAStream(logits.get_device())); + TORCH_CHECK(err == cudaSuccess, "m3_sampled_accept launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +// ws_mc_ptr is the multicast address of ws_uc's mapping: the op writes through it, which the schema cannot show. +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + // No tensor arguments, so no dispatch key: registered with the schema as a catch-all kernel. + m.def("minimax_m3_sampled_accept_ws_words(int world) -> int", + &tensorrt_llm::torch_ext::minimaxM3SampledAcceptWsWords); + m.def( + "minimax_m3_sampled_accept(Tensor logits, int vocab_offset, Tensor temperatures, Tensor top_ps, " + "Tensor seeds, Tensor offsets, Tensor(a!) ws_uc, int ws_mc_ptr, Tensor(b!) counters, Tensor(c!) tokens, " + "int rank, int world, Tensor draft, Tensor batch_idx, Tensor(d!) acc_tokens, Tensor(e!) num_acc, " + "Tensor(f!) gather_ids, int force_total, Tensor(g!)? next_new, Tensor? force_vec, " + "Tensor(h!)? extra_attempts, bool use_pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sampled_accept", &tensorrt_llm::torch_ext::minimaxM3SampledAccept); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SelectBlocksOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SelectBlocksOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SelectLongOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SelectLongOp.cpp new file mode 100644 index 000000000000..2def696ce875 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SelectLongOp.cpp @@ -0,0 +1,80 @@ +/* + * 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/minimaxM3/selectLong.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! Top-16 block select for MiniMax-M3 decode rows of up to 9216 blocks, into a caller-owned table. +//! +//! @param scores Float32 CUDA tensor [num_kv_heads, num_blocks, total_q], any strides. +//! @param nValidBlocks Contiguous int32 CUDA tensor with at least total_q elements. +//! @param out Int32 CUDA tensor of logical shape [total_q, num_kv_heads, 16] with unit stride in the last dim (the +//! head-major backing the target hands over is a permuted view). Every entry is written. +void minimaxM3SelectLong(torch::Tensor const& scores, torch::Tensor const& nValidBlocks, torch::Tensor& out, + int64_t initBlocks, int64_t localBlocks, bool usePdl, bool earlyTrigger) +{ + TORCH_CHECK(scores.is_cuda() && scores.scalar_type() == torch::kFloat32 && scores.dim() == 3, + "minimax_m3_select_long: scores must be a 3-D float32 CUDA tensor [num_kv_heads, num_blocks, total_q]"); + TORCH_CHECK(nValidBlocks.is_cuda() && nValidBlocks.scalar_type() == torch::kInt32 && nValidBlocks.is_contiguous() + && nValidBlocks.numel() >= scores.size(2), + "minimax_m3_select_long: n_valid_blocks must be contiguous int32 CUDA with at least total_q elements"); + int64_t const numKvHeads = scores.size(0); + int64_t const numBlocks = scores.size(1); + int64_t const totalQ = scores.size(2); + TORCH_CHECK(numBlocks <= tensorrt_llm::kernels::kMinimaxM3SelectLongMaxBlocks, + "minimax_m3_select_long supports at most ", tensorrt_llm::kernels::kMinimaxM3SelectLongMaxBlocks, + " blocks, got ", numBlocks); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kInt32 && out.dim() == 3 && out.size(0) == totalQ + && out.size(1) == numKvHeads && out.size(2) == 16 && out.stride(2) == 1, + "minimax_m3_select_long: out must be int32 [total_q, num_kv_heads, 16] with unit last-dim stride"); + + c10::cuda::CUDAGuard const deviceGuard{scores.device()}; + auto const stream = at::cuda::getCurrentCUDAStream(scores.get_device()); + cudaError_t const err = tensorrt_llm::kernels::invokeMinimaxM3SelectLong(scores.data_ptr(), scores.stride(0), + scores.stride(1), scores.stride(2), nValidBlocks.data_ptr(), out.data_ptr(), out.stride(0), + out.stride(1), static_cast(numKvHeads), static_cast(numBlocks), static_cast(totalQ), + static_cast(initBlocks), static_cast(localBlocks), usePdl, earlyTrigger, stream); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "minimax_m3_select_long launch failed: ", cudaGetErrorString(err)); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_select_long(Tensor scores, Tensor n_valid_blocks, Tensor(a!) out, int init_blocks, " + "int local_blocks, bool use_pdl=False, bool early_trigger=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_select_long", &tensorrt_llm::torch_ext::minimaxM3SelectLong); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3ShardAcceptOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3ShardAcceptOp.cpp new file mode 100644 index 000000000000..5b4ca3e6004f --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3ShardAcceptOp.cpp @@ -0,0 +1,118 @@ +/* + * 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/minimaxM3/shardAccept.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace sa = tensorrt_llm::kernels::minimaxM3::shardAccept; + +using OptTensor = std::optional; + +void checkArch(torch::Tensor const& t) +{ + TORCH_CHECK(t.is_cuda(), "minimax_m3_shard_accept expects CUDA tensors"); + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_shard_accept requires SM100 or newer"); +} + +void check(cudaError_t err) +{ + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, cudaGetErrorString(err)); + } +} + +cudaStream_t currentStream(torch::Tensor const& t) +{ + return at::cuda::getCurrentCUDAStream(t.get_device()); +} + +//! Per-request accepted counts for fractional forced acceptance (int32 [nreq]), or nullptr. +int const* forceVecPtr(OptTensor const& forceVec, int nreq) +{ + if (!forceVec.has_value()) + { + return nullptr; + } + TORCH_CHECK(forceVec->scalar_type() == at::kInt && forceVec->is_contiguous() && forceVec->numel() >= nreq, + "accept: force_vec must be contiguous int32 [nreq]"); + return forceVec->data_ptr(); +} + +} // namespace + +//! Greedy verification of nreq * (draft_len + 1) <= 48 verify rows from this rank's BF16 logits shard [rows, V / tp] +//! (global ids vocabOffset..), exchanged across the TP ranks through the multicast slots wsMcPtr / wsUc, plus the +//! strict acceptance (accepted tokens, counts, first-pass gather rows, optional nextNew[:, 0]). +void minimaxM3ShardAccept(torch::Tensor const& logits, int64_t vocabOffset, torch::Tensor& keys, + torch::Tensor& counters, torch::Tensor& wsUc, int64_t wsMcPtr, int64_t rank, int64_t world, + torch::Tensor const& draft, torch::Tensor const& batchIdx, torch::Tensor& accTokens, torch::Tensor& numAcc, + torch::Tensor& gatherIds, int64_t forceTotal, OptTensor const& nextNew, int64_t numCtas, OptTensor const& forceVec) +{ + checkArch(logits); + c10::cuda::CUDAGuard const deviceGuard{logits.device()}; + TORCH_CHECK(logits.scalar_type() == at::kBFloat16 && logits.dim() == 2 && logits.stride(1) == 1); + int const M = logits.size(0); + int const nreq = accTokens.size(0); + int const k1 = accTokens.size(1); + TORCH_CHECK(M >= 1 && M <= sa::kMaxRows && nreq * k1 == M && nreq <= sa::kThreads); + TORCH_CHECK(draft.scalar_type() == at::kInt && batchIdx.scalar_type() == at::kInt && accTokens.is_contiguous()); + sa::AcceptArgs acc{draft.data_ptr(), batchIdx.data_ptr(), accTokens.data_ptr(), + numAcc.data_ptr(), reinterpret_cast(gatherIds.data_ptr()), k1 - 1, nreq, + static_cast(forceTotal), nextNew.has_value() ? nextNew->data_ptr() : nullptr, + nextNew.has_value() ? static_cast(nextNew->stride(0)) : 0, forceVecPtr(forceVec, nreq)}; + check(sa::invokeShardAccept(reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()), M, + static_cast(logits.size(1)), static_cast(logits.stride(0)), static_cast(vocabOffset), + reinterpret_cast(keys.data_ptr()), reinterpret_cast(counters.data_ptr()), + reinterpret_cast(wsUc.data_ptr()), reinterpret_cast(wsMcPtr), + static_cast(rank), static_cast(world), acc, static_cast(numCtas), currentStream(logits))); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +// Mutations through the raw multicast address ws_mc_ptr are not visible to the schema: that buffer is owned by the +// Python workspace and never aliased by the op's tensor arguments. +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_shard_accept(Tensor logits, int vocab_offset, Tensor(a!) keys, Tensor(b!) counters, " + "Tensor(c!) ws_uc, int ws_mc_ptr, int rank, int world, Tensor draft, Tensor batch_idx, " + "Tensor(d!) acc_tokens, Tensor(e!) num_acc, Tensor(f!) gather_ids, int force_total, Tensor(g!)? next_new, " + "int num_ctas, Tensor? force_vec) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_shard_accept", &tensorrt_llm::torch_ext::minimaxM3ShardAccept); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SparseDecodeLcOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SparseDecodeLcOp.cpp new file mode 100644 index 000000000000..62a56779f88b --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SparseDecodeLcOp.cpp @@ -0,0 +1,139 @@ +/* + * 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/minimaxM3/sparseDecodeLc.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kHeads = 16; +constexpr int64_t kHeadDim = 128; +constexpr int64_t kTopK = 16; + +void checkBytes(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.element_size() == 1, "minimax_m3_sparse_decode_lc: ", name, + " must be a one-byte CUDA view (uint8 / e4m3 / packed e2m1)"); +} + +} // namespace + +//! One-launch MiniMax-M3 long-context sparse decode over the selected NVFP4 blocks (TP4 geometry: 16 query heads on +//! one KV head, head_dim 128, page 128, top-k 16). +//! +//! @param q E4M3 bytes [T, 16, 128], unit stride in the last dim. +//! @param k, v Packed E2M1 bytes [pages, 1, 128, 64] (the per-layer HND views). +//! @param kBlockScale, vBlockScale E4M3 bytes [pages, 1, 128, 8]: K token-major, V in 4x4 token-quad order. +//! @param kGlobalScale, vGlobalScale One-element float32 per-tensor dequant scales. +//! @param topk Int32 [T, 16] ascending selected block ids with -1 tail (kv head 0 of the head-major table). +//! @param out BF16 [T, 2048], written. q8 / sf, when given, receive the MXFP8 copy and its swizzled UE8M0 scales. +void minimaxM3SparseDecodeLc(torch::Tensor const& q, torch::Tensor const& k, torch::Tensor const& v, + torch::Tensor const& kBlockScale, torch::Tensor const& vBlockScale, torch::Tensor const& kGlobalScale, + torch::Tensor const& vGlobalScale, torch::Tensor const& topk, torch::Tensor const& blockTable, + torch::Tensor const& seqLens, torch::Tensor& out, std::optional q8, std::optional sf, + int64_t numSfKTiles, int64_t decodeQueryLen, double smScale, bool linearSoftmax, bool usePdl, bool earlyTrigger) +{ + TORCH_CHECK(earlyTrigger, "minimax_m3_sparse_decode_lc is instantiated with early_trigger=True only"); + TORCH_CHECK(q.is_cuda(), "minimax_m3_sparse_decode_lc: q must be a CUDA tensor"); + c10::cuda::CUDAGuard const deviceGuard{q.device()}; + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 10, + "minimax_m3_sparse_decode_lc requires an SM100-family GPU"); + checkBytes(q, "q"); + checkBytes(k, "k"); + checkBytes(v, "v"); + checkBytes(kBlockScale, "k_block_scale"); + checkBytes(vBlockScale, "v_block_scale"); + TORCH_CHECK(q.dim() == 3 && q.size(1) == kHeads && q.size(2) == kHeadDim && q.stride(2) == 1, + "minimax_m3_sparse_decode_lc: q must be [T, 16, 128] with unit last-dim stride"); + TORCH_CHECK(kGlobalScale.scalar_type() == torch::kFloat32 && vGlobalScale.scalar_type() == torch::kFloat32, + "minimax_m3_sparse_decode_lc: global scales must be float32"); + TORCH_CHECK(topk.scalar_type() == torch::kInt32 && topk.dim() == 2 && topk.size(1) == kTopK, + "minimax_m3_sparse_decode_lc: topk must be int32 [T, 16]"); + TORCH_CHECK(blockTable.scalar_type() == torch::kInt32 && seqLens.scalar_type() == torch::kInt32, + "minimax_m3_sparse_decode_lc: block_table and seq_lens must be int32"); + TORCH_CHECK(out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.size(0) == q.size(0) + && out.size(1) == kHeads * kHeadDim && out.stride(1) == 1, + "minimax_m3_sparse_decode_lc: out must be bfloat16 [T, 2048] with unit last-dim stride"); + TORCH_CHECK(q8.has_value() == sf.has_value(), "minimax_m3_sparse_decode_lc: pass both q8 and sf, or neither"); + bool const mxfp8 = q8.has_value(); + if (mxfp8) + { + checkBytes(*q8, "q8"); + checkBytes(*sf, "sf"); + } + + tensorrt_llm::kernels::MinimaxM3SparseDecodeLcParams prm{}; + prm.q = reinterpret_cast(q.data_ptr()); + prm.q_tok = q.stride(0); + prm.q_head = q.stride(1); + prm.k = reinterpret_cast(k.data_ptr()); + prm.v = reinterpret_cast(v.data_ptr()); + prm.k_blk = k.stride(0); + prm.v_blk = v.stride(0); + prm.ksc = reinterpret_cast(kBlockScale.data_ptr()); + prm.vsc = reinterpret_cast(vBlockScale.data_ptr()); + prm.ksc_blk = kBlockScale.stride(0); + prm.vsc_blk = vBlockScale.stride(0); + prm.kgs = kGlobalScale.data_ptr(); + prm.vgs = vGlobalScale.data_ptr(); + prm.topk = topk.data_ptr(); + prm.t_tok = topk.stride(0); + prm.t_k = topk.stride(1); + prm.block_table = blockTable.data_ptr(); + prm.bt_stride = blockTable.stride(0); + prm.seq_lens = seqLens.data_ptr(); + prm.dql = static_cast(decodeQueryLen); + prm.sm_scale = static_cast(smScale); + prm.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr()); + prm.out_stride = out.stride(0); + prm.q8 = mxfp8 ? reinterpret_cast(q8->data_ptr()) : nullptr; + prm.sf = mxfp8 ? reinterpret_cast(sf->data_ptr()) : nullptr; + prm.num_sf_k_tiles = static_cast(numSfKTiles); + + auto const stream = at::cuda::getCurrentCUDAStream(q.get_device()); + cudaError_t const err = tensorrt_llm::kernels::invokeMinimaxM3SparseDecodeLc( + prm, static_cast(q.size(0)), linearSoftmax, mxfp8, earlyTrigger, usePdl, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sparse_decode_lc launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sparse_decode_lc(Tensor q, Tensor k, Tensor v, Tensor k_block_scale, Tensor v_block_scale, " + "Tensor k_global_scale, Tensor v_global_scale, Tensor topk, Tensor block_table, Tensor seq_lens, " + "Tensor(a!) out, Tensor(b!)? q8, Tensor(c!)? sf, int num_sf_k_tiles, int decode_query_len, float sm_scale, " + "bool linear_softmax, bool use_pdl, bool early_trigger) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sparse_decode_lc", &tensorrt_llm::torch_ext::minimaxM3SparseDecodeLc); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SpecTailOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SpecTailOp.cpp new file mode 100644 index 000000000000..cea04a45cd86 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SpecTailOp.cpp @@ -0,0 +1,202 @@ +/* + * 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/minimaxM3/specTail.h" + +#include +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace st = tensorrt_llm::kernels::minimaxM3::specTail; + +void check(cudaError_t err) +{ + TORCH_CHECK(err == cudaSuccess, cudaGetErrorString(err)); +} + +cudaStream_t currentStream(torch::Tensor const& t) +{ + return at::cuda::getCurrentCUDAStream(t.get_device()); +} + +__nv_bfloat16 const* bf16In(torch::Tensor const& t) +{ + return reinterpret_cast<__nv_bfloat16 const*>(t.data_ptr()); +} + +__nv_bfloat16* bf16Out(torch::Tensor const& t) +{ + return reinterpret_cast<__nv_bfloat16*>(t.data_ptr()); +} + +//! Register-tile count of the norm kernels for hidden size H (H / 8 elements over kThreads threads). +int normIters(int H) +{ + return (H / 8 + st::kThreads - 1) / st::kThreads; +} + +st::IntView intView(torch::Tensor const& t) +{ + return st::IntView{t.data_ptr(), t.stride(0), t.stride(1), static_cast(t.size(1))}; +} + +} // namespace + +//! out = [rmsnorm(embed(ids)) * wEmbed, rmsnorm(hidden[rows]) * wHidden] [M, 2H]; residual = hidden[rows]. +void minimaxM3DraftInput(torch::Tensor const& ids, torch::Tensor const& table, torch::Tensor const& hidden, + std::optional const& rows, torch::Tensor const& wEmbed, torch::Tensor const& wHidden, + double epsEmbed, double epsHidden, torch::Tensor& out, std::optional const& residual, bool pdl) +{ + c10::cuda::CUDAGuard const deviceGuard{table.device()}; + int const M = ids.size(0); + int const H = table.size(1); + TORCH_CHECK(ids.scalar_type() == at::kInt && ids.is_contiguous() && table.scalar_type() == at::kBFloat16 + && table.is_contiguous() && H % 8 == 0); + TORCH_CHECK( + hidden.scalar_type() == at::kBFloat16 && hidden.dim() == 2 && hidden.is_contiguous() && hidden.size(1) == H); + TORCH_CHECK(out.scalar_type() == at::kBFloat16 && out.size(0) == M && out.size(1) == 2 * H && out.is_contiguous()); + TORCH_CHECK(wEmbed.scalar_type() == at::kBFloat16 && wEmbed.numel() == H && wHidden.scalar_type() == at::kBFloat16 + && wHidden.numel() == H); + TORCH_CHECK(!rows.has_value() || (rows->scalar_type() == at::kLong && rows->is_contiguous() && rows->size(0) == M)); + TORCH_CHECK(rows.has_value() || hidden.size(0) >= M, "draft_input: hidden must hold a row per id"); + TORCH_CHECK(!residual.has_value() + || (residual->scalar_type() == at::kBFloat16 && residual->is_contiguous() && residual->size(0) == M + && residual->size(1) == H), + "draft_input: residual must be contiguous [M, H]"); + TORCH_CHECK(normIters(H) <= st::kMaxIters, "draft_input: hidden size too large"); + check(st::invokeDraftInput(ids.data_ptr(), static_cast(table.size(0)), bf16In(table), bf16In(hidden), + rows.has_value() ? reinterpret_cast(rows->data_ptr()) : nullptr, bf16In(wEmbed), + bf16In(wHidden), static_cast(epsEmbed), static_cast(epsHidden), H, M, bf16Out(out), + residual.has_value() ? bf16Out(*residual) : nullptr, pdl, currentStream(table))); +} + +//! Per-chunk RMSNorm of src [M, >= C*H] (row stride may exceed C*H) with weights w [C, H] into dst [M, C*H]. +void minimaxM3ChunkNorm(torch::Tensor const& src, torch::Tensor const& w, double eps, torch::Tensor& dst, bool pdl) +{ + c10::cuda::CUDAGuard const deviceGuard{src.device()}; + int const M = src.size(0); + int const C = w.size(0); + int const H = w.size(1); + TORCH_CHECK(src.scalar_type() == at::kBFloat16 && w.scalar_type() == at::kBFloat16 && src.stride(1) == 1); + TORCH_CHECK(src.size(1) >= C * H && dst.is_contiguous() && dst.numel() == static_cast(M) * C * H); + TORCH_CHECK(H % 8 == 0 && src.stride(0) % 8 == 0); + TORCH_CHECK(normIters(H) <= st::kMaxIters, "chunk_norm: hidden size too large"); + check(st::invokeChunkNorm(bf16In(src), src.stride(0), bf16In(w), static_cast(eps), H, C, M, bf16Out(dst), + pdl, currentStream(src))); +} + +//! dst [M, H] (row stride may exceed H) = rmsnorm(bf16(h + r)) * w, r optional. +void minimaxM3CaptureNorm(torch::Tensor const& h, std::optional const& r, torch::Tensor const& w, + double eps, torch::Tensor& dst) +{ + c10::cuda::CUDAGuard const deviceGuard{h.device()}; + int const M = h.size(0); + int const H = w.size(0); + TORCH_CHECK( + h.scalar_type() == at::kBFloat16 && w.scalar_type() == at::kBFloat16 && dst.scalar_type() == at::kBFloat16); + TORCH_CHECK(h.dim() == 2 && h.size(1) == H && h.stride(1) == 1 && dst.dim() == 2 && dst.size(0) == M + && dst.size(1) == H && dst.stride(1) == 1 && H % 8 == 0 && h.stride(0) % 8 == 0 && dst.stride(0) % 8 == 0); + TORCH_CHECK(!r.has_value() + || (r->sizes() == h.sizes() && r->stride(1) == 1 && r->stride(0) % 8 == 0 + && r->scalar_type() == at::kBFloat16)); + TORCH_CHECK(normIters(H) <= st::kMaxIters, "capture_norm: hidden size too large"); + check(st::invokeCaptureNorm(bf16In(h), h.stride(0), r.has_value() ? bf16In(*r) : nullptr, + r.has_value() ? r->stride(0) : 0, bf16In(w), static_cast(eps), H, M, bf16Out(dst), dst.stride(0), + currentStream(h))); +} + +//! The overlap scheduler's step inputs from the previous step's device outputs (see step_gather_kernel). +void minimaxM3StepGather(torch::Tensor const& newTok, torch::Tensor const& nextDraft, torch::Tensor const& lens, + torch::Tensor const& slots, torch::Tensor const& ppi, int64_t P, int64_t T, int64_t D, torch::Tensor& ids, + int64_t ids0, torch::Tensor& draft, int64_t d0, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t b0) +{ + c10::cuda::CUDAGuard const deviceGuard{newTok.device()}; + for (auto const* t : std::initializer_list{ + &newTok, &nextDraft, &lens, &slots, &ppi, &ids, &draft, &posOff, &kvOff}) + { + TORCH_CHECK(t->scalar_type() == at::kInt, "step_gather: int32 tensors expected"); + } + TORCH_CHECK(newTok.dim() >= 2 && nextDraft.dim() == 2 && lens.is_contiguous() && posOff.is_contiguous() + && kvOff.is_contiguous() && ids.is_contiguous() && draft.is_contiguous()); + int const work = std::max(std::max(P * T, P * D), std::max(posOff.numel(), kvOff.numel())); + check(st::invokeStepGather(newTok.data_ptr(), newTok.stride(0), newTok.stride(1), nextDraft.data_ptr(), + nextDraft.stride(0), nextDraft.stride(1), lens.data_ptr(), slots.data_ptr(), ppi.data_ptr(), + static_cast(P), static_cast(T), static_cast(D), ids.data_ptr(), static_cast(ids0), + draft.data_ptr(), static_cast(d0), posOff.data_ptr(), static_cast(posOff.numel()), + kvOff.data_ptr(), static_cast(kvOff.numel()), static_cast(b0), work, currentStream(newTok))); +} + +//! The spec sampler's slot stores from one step's outputs (see slot_scatter_kernel). +void minimaxM3SlotScatter(torch::Tensor const& slots, torch::Tensor const& sa, torch::Tensor& da, + torch::Tensor const& sb, torch::Tensor& db, torch::Tensor const& sl, torch::Tensor& dl, torch::Tensor const& sd, + torch::Tensor& dd) +{ + c10::cuda::CUDAGuard const deviceGuard{slots.device()}; + for (auto const* t : std::initializer_list{&sa, &da, &sb, &db, &sl, &dl, &sd, &dd}) + { + TORCH_CHECK(t->scalar_type() == at::kInt && t->dim() <= 2, "slot_scatter: int32 tensors expected"); + } + TORCH_CHECK(slots.scalar_type() == at::kLong); + int const n = slots.numel(); + if (n == 0) + { + return; + } + check(st::invokeSlotScatter(reinterpret_cast(slots.data_ptr()), n, intView(sa), intView(da), + intView(sb), intView(db), sl.data_ptr(), dl.data_ptr(), intView(sd), intView(dd), + currentStream(slots))); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_draft_input(Tensor ids, Tensor table, Tensor hidden, Tensor? rows, Tensor w_embed, " + "Tensor w_hidden, float eps_embed, float eps_hidden, Tensor(a!) out, Tensor(b!)? residual, bool pdl) -> ()"); + m.def("minimax_m3_chunk_norm(Tensor src, Tensor w, float eps, Tensor(a!) dst, bool pdl) -> ()"); + m.def("minimax_m3_capture_norm(Tensor h, Tensor? r, Tensor w, float eps, Tensor(a!) dst) -> ()"); + m.def( + "minimax_m3_step_gather(Tensor new_tok, Tensor next_draft, Tensor lens, Tensor slots, Tensor ppi, int P, " + "int T, int D, Tensor(a!) ids, int ids0, Tensor(b!) draft, int d0, Tensor(c!) pos_off, Tensor(d!) kv_off, " + "int b0) -> ()"); + m.def( + "minimax_m3_slot_scatter(Tensor slots, Tensor sa, Tensor(a!) da, Tensor sb, Tensor(b!) db, Tensor sl, " + "Tensor(c!) dl, Tensor sd, Tensor(d!) dd) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_draft_input", &tensorrt_llm::torch_ext::minimaxM3DraftInput); + m.impl("minimax_m3_chunk_norm", &tensorrt_llm::torch_ext::minimaxM3ChunkNorm); + m.impl("minimax_m3_capture_norm", &tensorrt_llm::torch_ext::minimaxM3CaptureNorm); + m.impl("minimax_m3_step_gather", &tensorrt_llm::torch_ext::minimaxM3StepGather); + m.impl("minimax_m3_slot_scatter", &tensorrt_llm::torch_ext::minimaxM3SlotScatter); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3TargetHeadOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3TargetHeadOp.cpp new file mode 100644 index 000000000000..448989e08c3a --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3TargetHeadOp.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/minimaxM3/targetHead.h" + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace th = tensorrt_llm::kernels::minimaxM3::targetHead; + +void checkArch(torch::Tensor const& t) +{ + TORCH_CHECK(t.is_cuda(), "minimax_m3_target_head expects CUDA tensors"); + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_target_head requires SM100 or newer"); +} + +} // namespace + +//! out [M, n] (BF16) = x [M <= 8, K] (BF16) x W^T for the 12-bit coded BF16 shard W [n, K]: w the coded stream +//! (uint8 [tiles * K / 64 * 1536]), exc_off int32 [tiles + 1] / exc_pos int32 [E] / exc_w fp32 [E, 2] its exception +//! lists, exp_base its exponent window's base. part (fp32 [grid * 16 * 256]) and counter (int32 [tiles], zero) are +//! the split-tile scratch. +void minimaxM3TargetHead(torch::Tensor const& w, torch::Tensor const& x, torch::Tensor& out, torch::Tensor& part, + torch::Tensor& counter, torch::Tensor const& excOff, torch::Tensor const& excPos, torch::Tensor const& excW, + int64_t expBase, int64_t n, int64_t grid, bool pdl) +{ + checkArch(x); + for (torch::Tensor const* t : + std::initializer_list{&w, &out, &part, &counter, &excOff, &excPos, &excW}) + { + TORCH_CHECK(t->device() == x.device(), "target_head: every tensor must be on x's device"); + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + TORCH_CHECK(x.scalar_type() == at::kBFloat16 && x.dim() == 2, "target_head: x must be BF16 [M, K]"); + int const M = static_cast(x.size(0)); + int const K = static_cast(x.size(1)); + TORCH_CHECK(M >= 1 && M <= th::kMaxRows, "target_head: 1..", th::kMaxRows, " rows"); + TORCH_CHECK(K % 64 == 0 && K / 16 <= 0xFFFF, "target_head: K must be a multiple of 64"); + TORCH_CHECK(x.stride(1) == 1 && x.stride(0) % 8 == 0 && reinterpret_cast(x.data_ptr()) % 16 == 0, + "target_head: x rows must be unit-stride and 16-byte aligned"); + TORCH_CHECK(n >= 16 && n % 16 == 0, "target_head: n must be a positive multiple of 16 (whole tiles)"); + int64_t const tiles = (n + 15) / 16; + TORCH_CHECK(w.scalar_type() == at::kByte && w.is_contiguous() && w.numel() == tiles * 16 * K * 3 / 2, + "target_head: w must be the contiguous uint8 coded stream of [ceil(n / 16) * 16, K]"); + TORCH_CHECK(out.scalar_type() == at::kBFloat16 && out.dim() == 2 && out.size(0) >= M && out.size(1) >= n + && out.stride(1) == 1, + "target_head: out must be BF16 [M, >= n] with unit column stride"); + TORCH_CHECK(grid >= 1, "target_head: grid must be positive"); + TORCH_CHECK(part.scalar_type() == at::kFloat && part.is_contiguous() && part.numel() >= grid * th::kWarps * 2 * 128, + "target_head: part must be fp32 [grid * 16 * 256]"); + TORCH_CHECK(counter.scalar_type() == at::kInt && counter.is_contiguous() && counter.numel() >= tiles, + "target_head: counter must be int32 [tiles]"); + TORCH_CHECK(excOff.scalar_type() == at::kInt && excOff.is_contiguous() && excOff.numel() == tiles + 1, + "target_head: exc_off must be int32 [tiles + 1]"); + TORCH_CHECK(excPos.scalar_type() == at::kInt && excPos.is_contiguous(), "target_head: exc_pos must be int32"); + TORCH_CHECK(excW.scalar_type() == at::kFloat && excW.is_contiguous() && excW.numel() == 2 * excPos.numel(), + "target_head: exc_w must be fp32 [E, 2]"); + TORCH_CHECK(expBase >= 0 && expBase <= 240, "target_head: exponent base out of range"); + th::Params p{}; + p.w = reinterpret_cast(w.data_ptr()); + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.x_stride = x.stride(0); + p.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr()); + p.out_stride = out.stride(0); + p.part = part.data_ptr(); + p.counter = counter.data_ptr(); + p.exc_off = excOff.data_ptr(); + p.exc_pos = reinterpret_cast(excPos.data_ptr()); + p.exc_w = reinterpret_cast(excW.data_ptr()); + p.exp_base = static_cast((expBase << 7) | (expBase << 23)); + p.M = M; + p.N = static_cast(n); + cudaError_t const err + = th::invokeTargetHead(p, K, static_cast(grid), pdl, at::cuda::getCurrentCUDAStream(x.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, cudaGetErrorString(err)); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_target_head(Tensor w, Tensor x, Tensor(a!) out, Tensor(b!) part, Tensor(c!) counter, " + "Tensor exc_off, Tensor exc_pos, Tensor exc_w, int exp_base, int n, int grid, bool pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_target_head", &tensorrt_llm::torch_ext::minimaxM3TargetHead); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Tp8FusedScoreSelectOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Tp8FusedScoreSelectOp.cpp new file mode 100644 index 000000000000..1d28f41fdc7a --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Tp8FusedScoreSelectOp.cpp @@ -0,0 +1,173 @@ +/* + * 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/minimaxM3Tp8/fusedScoreSelect.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kHeadDim = 128; +constexpr int64_t kBlockK = 128; +constexpr int64_t kTopK = 16; + +} // namespace + +//! MiniMax-M3 TP8 fused index score + top-16 block select for decode spans of at most 2 requests. Equals +//! minimax_m3_select_blocks applied to the block scores of cute_dsl_minimax_m3_index_decode_score. Grid (splits, +//! batch) with splits = max(SM count / batch, ceil(max_tiles / 128), 1). +//! +//! @param idxQ E4M3 [total_q, heads, 128] (total_q = batch * dql, heads * 4 <= 8), unit last-dim stride, row and head +//! strides multiples of 4 bytes, 4-byte aligned. +//! @param idxKPages E4M3 [num_pages, 128, 128], unit last-dim stride, page and row strides multiples of 16 bytes, +//! 16-byte aligned. +//! @param blockTable Int32 [batch, >= blocks], unit last-dim stride. +//! @param seqLens Int32 [batch] (1 or 2 requests), including this step's dql query rows. +//! @param nValid Contiguous int32 [total_q] per-token causal valid-block counts. +//! @param out Int32, logical [total_q, num_kv_heads, 16] with unit last-dim stride (a permuted head-major backing is +//! fine). Every entry is written: ascending block ids with a -1 tail. +//! @param scratch Contiguous float32 [>= batch, >= num_kv_heads * dql, L] with L >= max_tiles and L % 4 == 0, 16-byte +//! aligned; [batch, 8, ceil4(max_tiles)] covers every supported geometry. Contents need no initialization. +//! @param counters Contiguous int32 with >= batch elements, all zero before the first call. The kernel returns each +//! request's counter to zero before it exits, so the buffer stays reusable across calls and graph replays without a +//! reset (as long as launches sharing it are stream-ordered). +//! @param maxTiles Worst-case block count n_valid is clamped to; num_kv_heads * dql * ceil4(max_tiles) * 4 + +//! num_kv_heads * dql * 1024 must not exceed 196608 bytes. +void minimaxM3Tp8FusedScoreSelect(torch::Tensor const& idxQ, torch::Tensor const& idxKPages, + torch::Tensor const& blockTable, torch::Tensor const& seqLens, torch::Tensor const& nValid, torch::Tensor& out, + torch::Tensor& scratch, torch::Tensor& counters, int64_t decodeQueryLen, int64_t numKvHeads, int64_t initBlocks, + int64_t localBlocks, int64_t maxTiles, bool usePdl) +{ + namespace tk = tensorrt_llm::kernels; + TORCH_CHECK(idxQ.is_cuda(), "minimax_m3_tp8_fused_score_select: idx_q must be a CUDA tensor"); + c10::cuda::CUDAGuard const deviceGuard{idxQ.device()}; + auto const* props = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(props->major >= 10, "minimax_m3_tp8_fused_score_select requires an SM100-family GPU"); + + TORCH_CHECK(idxQ.dim() == 3 && idxQ.size(2) == kHeadDim && idxQ.stride(2) == 1 && idxQ.element_size() == 1, + "minimax_m3_tp8_fused_score_select: idx_q must be one-byte [total_q, heads, 128] with unit last-dim stride"); + TORCH_CHECK( + idxQ.stride(0) % 4 == 0 && idxQ.stride(1) % 4 == 0 && (reinterpret_cast(idxQ.data_ptr()) & 3) == 0, + "minimax_m3_tp8_fused_score_select: idx_q strides and base must be 4-byte aligned"); + TORCH_CHECK(idxKPages.is_cuda() && idxKPages.element_size() == 1 && idxKPages.dim() == 3 + && idxKPages.size(1) == kBlockK && idxKPages.size(2) == kHeadDim && idxKPages.stride(2) == 1, + "minimax_m3_tp8_fused_score_select: idx_k_pages must be one-byte [num_pages, 128, 128] with unit last-dim " + "stride"); + TORCH_CHECK(idxKPages.stride(1) % 16 == 0 && idxKPages.stride(0) % 16 == 0 + && (reinterpret_cast(idxKPages.data_ptr()) & 15) == 0, + "minimax_m3_tp8_fused_score_select: idx_k_pages strides and base must be 16-byte aligned"); + TORCH_CHECK(blockTable.scalar_type() == torch::kInt32 && blockTable.dim() == 2 && blockTable.stride(1) == 1, + "minimax_m3_tp8_fused_score_select: block_table must be int32 [batch, blocks] with unit last-dim stride"); + TORCH_CHECK( + seqLens.scalar_type() == torch::kInt32 && nValid.scalar_type() == torch::kInt32 && nValid.is_contiguous(), + "minimax_m3_tp8_fused_score_select: seq_lens and n_valid must be int32, n_valid contiguous"); + TORCH_CHECK(out.scalar_type() == torch::kInt32 && out.dim() == 3 && out.size(2) == kTopK && out.stride(2) == 1, + "minimax_m3_tp8_fused_score_select: out must be int32 [total_q, num_kv_heads, 16] with unit last-dim stride"); + + int64_t const batch = seqLens.size(0); + int64_t const totalQ = idxQ.size(0); + int64_t const heads = idxQ.size(1); + int64_t const rows = decodeQueryLen * numKvHeads; + TORCH_CHECK(batch >= 1 && batch <= tk::kMinimaxM3Tp8FusedScoreSelectMaxBatch, + "minimax_m3_tp8_fused_score_select takes 1 to ", tk::kMinimaxM3Tp8FusedScoreSelectMaxBatch, " requests, got ", + batch); + TORCH_CHECK(decodeQueryLen >= 1 && decodeQueryLen <= tk::kMinimaxM3Tp8FusedScoreSelectMaxDql + && totalQ == batch * decodeQueryLen, + "minimax_m3_tp8_fused_score_select: decode_query_len must be in [1, 4] with total_q = batch * " + "decode_query_len"); + TORCH_CHECK(numKvHeads >= 1 && heads % numKvHeads == 0 && heads * tk::kMinimaxM3Tp8FusedScoreSelectMaxDql <= 8 + && rows <= tk::kMinimaxM3Tp8FusedScoreSelectMaxRows, + "minimax_m3_tp8_fused_score_select: unsupported head geometry"); + TORCH_CHECK(out.size(0) == totalQ && out.size(1) == numKvHeads, + "minimax_m3_tp8_fused_score_select: out must be [total_q, num_kv_heads, 16]"); + TORCH_CHECK(nValid.numel() >= totalQ, "minimax_m3_tp8_fused_score_select: n_valid needs total_q elements"); + TORCH_CHECK(maxTiles >= 1, "minimax_m3_tp8_fused_score_select: max_tiles must be positive"); + + TORCH_CHECK(scratch.is_cuda() && scratch.scalar_type() == torch::kFloat32 && scratch.dim() == 3 + && scratch.is_contiguous() && scratch.size(0) >= batch && scratch.size(1) >= rows + && scratch.size(2) >= maxTiles && scratch.stride(1) % 4 == 0 + && (reinterpret_cast(scratch.data_ptr()) & 15) == 0, + "minimax_m3_tp8_fused_score_select: scratch must be contiguous 16-byte aligned float32 [>= batch, >= rows, " + "L] with L >= max_tiles and L % 4 == 0"); + TORCH_CHECK(counters.is_cuda() && counters.scalar_type() == torch::kInt32 && counters.is_contiguous() + && counters.numel() >= batch, + "minimax_m3_tp8_fused_score_select: counters must be contiguous int32 with >= batch elements"); + TORCH_CHECK(rows * ((maxTiles + 3) & ~int64_t{3}) * 4 + rows * 1024 <= tk::kMinimaxM3Tp8FusedScoreSelectRingBytes, + "minimax_m3_tp8_fused_score_select: select scratch does not fit the stage ring"); + + int64_t const splits = std::max({props->multiProcessorCount / batch, + (maxTiles + tk::kMinimaxM3Tp8FusedScoreSelectMaxLocalBlocks - 1) + / tk::kMinimaxM3Tp8FusedScoreSelectMaxLocalBlocks, + int64_t{1}}); + TORCH_CHECK((maxTiles + splits - 1) / splits <= tk::kMinimaxM3Tp8FusedScoreSelectMaxLocalBlocks, + "minimax_m3_tp8_fused_score_select: too few splits"); + + tk::MinimaxM3Tp8FusedScoreSelectParams p{}; + p.q = reinterpret_cast(idxQ.data_ptr()); + p.q_st = idxQ.stride(0); + p.q_sh = idxQ.stride(1); + p.bt = blockTable.data_ptr(); + p.bt_stride = blockTable.stride(0); + p.bt_width = static_cast(blockTable.size(1)); + p.seq_lens = seqLens.data_ptr(); + p.n_valid = nValid.data_ptr(); + p.scr = scratch.data_ptr(); + p.scr_b = scratch.stride(0); + p.scr_r = scratch.stride(1); + p.out = out.data_ptr(); + p.out_t = out.stride(0); + p.out_h = out.stride(1); + p.counters = counters.data_ptr(); + p.heads = static_cast(heads); + p.kv_heads = static_cast(numKvHeads); + p.dql = static_cast(decodeQueryLen); + p.max_tiles = static_cast(maxTiles); + p.init_blocks = static_cast(initBlocks); + p.local_blocks = static_cast(localBlocks); + + auto const stream = at::cuda::getCurrentCUDAStream(idxQ.get_device()); + cudaError_t const err + = tk::invokeMinimaxM3Tp8FusedScoreSelect(p, idxKPages.data_ptr(), idxKPages.size(0), idxKPages.stride(1), + idxKPages.stride(0), static_cast(batch), static_cast(splits), usePdl, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_tp8_fused_score_select launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_tp8_fused_score_select(Tensor idx_q, Tensor idx_k_pages, Tensor block_table, Tensor seq_lens, " + "Tensor n_valid, Tensor(a!) out, Tensor(b!) scratch, Tensor(c!) counters, int decode_query_len, " + "int num_kv_heads, int init_blocks, int local_blocks, int max_tiles, bool use_pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_tp8_fused_score_select", &tensorrt_llm::torch_ext::minimaxM3Tp8FusedScoreSelect); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Tp8SparseDecodeLc8Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Tp8SparseDecodeLc8Op.cpp new file mode 100644 index 000000000000..3f40efdb9bb6 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Tp8SparseDecodeLc8Op.cpp @@ -0,0 +1,139 @@ +/* + * 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/minimaxM3Tp8/sparseDecodeLc8.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kHeads = 8; +constexpr int64_t kHeadDim = 128; +constexpr int64_t kTopK = 16; + +void checkBytes(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.element_size() == 1, "minimax_m3_tp8_sparse_decode_lc8: ", name, + " must be a one-byte CUDA view (uint8 / e4m3 / packed e2m1)"); +} + +} // namespace + +//! One-launch MiniMax-M3 TP8 long-context sparse decode over the selected NVFP4 blocks (8 query heads on one KV head, +//! head_dim 128, page 128, top-k 16). Dependents are released right after the grid dependency wait. +//! +//! @param q E4M3 bytes [T, 8, 128], unit stride in the last dim. +//! @param k, v Packed E2M1 bytes [pages, 1, 128, 64] (the per-layer HND views). +//! @param kBlockScale, vBlockScale E4M3 bytes [pages, 1, 128, 8]: K token-major, V in 4x4 token-quad order. +//! @param kGlobalScale, vGlobalScale One-element float32 per-tensor dequant scales. +//! @param topk Int32 [T, 16] ascending selected block ids with -1 tail (kv head 0 of the head-major table). +//! @param out BF16 [T, 1024], written. q8 / sf, when given, receive the MXFP8 copy and its swizzled UE8M0 scales. +//! @param prewait Stage seq_lens, q and the request's block-table row before the grid dependency wait (same output). +void minimaxM3Tp8SparseDecodeLc8(torch::Tensor const& q, torch::Tensor const& k, torch::Tensor const& v, + torch::Tensor const& kBlockScale, torch::Tensor const& vBlockScale, torch::Tensor const& kGlobalScale, + torch::Tensor const& vGlobalScale, torch::Tensor const& topk, torch::Tensor const& blockTable, + torch::Tensor const& seqLens, torch::Tensor& out, std::optional q8, std::optional sf, + int64_t numSfKTiles, int64_t decodeQueryLen, double smScale, bool linearSoftmax, bool usePdl, bool prewait) +{ + TORCH_CHECK(q.is_cuda(), "minimax_m3_tp8_sparse_decode_lc8: q must be a CUDA tensor"); + c10::cuda::CUDAGuard const deviceGuard{q.device()}; + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 10, + "minimax_m3_tp8_sparse_decode_lc8 requires an SM100-family GPU"); + checkBytes(q, "q"); + checkBytes(k, "k"); + checkBytes(v, "v"); + checkBytes(kBlockScale, "k_block_scale"); + checkBytes(vBlockScale, "v_block_scale"); + TORCH_CHECK(q.dim() == 3 && q.size(1) == kHeads && q.size(2) == kHeadDim && q.stride(2) == 1, + "minimax_m3_tp8_sparse_decode_lc8: q must be [T, 8, 128] with unit last-dim stride"); + TORCH_CHECK(kGlobalScale.scalar_type() == torch::kFloat32 && vGlobalScale.scalar_type() == torch::kFloat32, + "minimax_m3_tp8_sparse_decode_lc8: global scales must be float32"); + TORCH_CHECK(topk.scalar_type() == torch::kInt32 && topk.dim() == 2 && topk.size(1) == kTopK, + "minimax_m3_tp8_sparse_decode_lc8: topk must be int32 [T, 16]"); + TORCH_CHECK(blockTable.scalar_type() == torch::kInt32 && seqLens.scalar_type() == torch::kInt32, + "minimax_m3_tp8_sparse_decode_lc8: block_table and seq_lens must be int32"); + TORCH_CHECK(out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.size(0) == q.size(0) + && out.size(1) == kHeads * kHeadDim && out.stride(1) == 1, + "minimax_m3_tp8_sparse_decode_lc8: out must be bfloat16 [T, 1024] with unit last-dim stride"); + TORCH_CHECK(q8.has_value() == sf.has_value(), "minimax_m3_tp8_sparse_decode_lc8: pass both q8 and sf, or neither"); + bool const mxfp8 = q8.has_value(); + if (mxfp8) + { + checkBytes(*q8, "q8"); + checkBytes(*sf, "sf"); + } + + tensorrt_llm::kernels::MinimaxM3Tp8SparseDecodeLc8Params prm{}; + prm.q = reinterpret_cast(q.data_ptr()); + prm.q_tok = q.stride(0); + prm.q_head = q.stride(1); + prm.k = reinterpret_cast(k.data_ptr()); + prm.v = reinterpret_cast(v.data_ptr()); + prm.k_blk = k.stride(0); + prm.v_blk = v.stride(0); + prm.ksc = reinterpret_cast(kBlockScale.data_ptr()); + prm.vsc = reinterpret_cast(vBlockScale.data_ptr()); + prm.ksc_blk = kBlockScale.stride(0); + prm.vsc_blk = vBlockScale.stride(0); + prm.kgs = kGlobalScale.data_ptr(); + prm.vgs = vGlobalScale.data_ptr(); + prm.topk = topk.data_ptr(); + prm.t_tok = topk.stride(0); + prm.t_k = topk.stride(1); + prm.block_table = blockTable.data_ptr(); + prm.bt_stride = blockTable.stride(0); + prm.seq_lens = seqLens.data_ptr(); + prm.dql = static_cast(decodeQueryLen); + prm.sm_scale = static_cast(smScale); + prm.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr()); + prm.out_stride = out.stride(0); + prm.q8 = mxfp8 ? reinterpret_cast(q8->data_ptr()) : nullptr; + prm.sf = mxfp8 ? reinterpret_cast(sf->data_ptr()) : nullptr; + prm.num_sf_k_tiles = static_cast(numSfKTiles); + + auto const stream = at::cuda::getCurrentCUDAStream(q.get_device()); + cudaError_t const err = tensorrt_llm::kernels::invokeMinimaxM3Tp8SparseDecodeLc8( + prm, static_cast(q.size(0)), linearSoftmax, mxfp8, prewait, usePdl, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_tp8_sparse_decode_lc8 launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_tp8_sparse_decode_lc8(Tensor q, Tensor k, Tensor v, Tensor k_block_scale, Tensor v_block_scale, " + "Tensor k_global_scale, Tensor v_global_scale, Tensor topk, Tensor block_table, Tensor seq_lens, " + "Tensor(a!) out, Tensor(b!)? q8, Tensor(c!)? sf, int num_sf_k_tiles, int decode_query_len, float sm_scale, " + "bool linear_softmax, bool use_pdl, bool prewait) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_tp8_sparse_decode_lc8", &tensorrt_llm::torch_ext::minimaxM3Tp8SparseDecodeLc8); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/stepInputsScatterOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/stepInputsScatterOp.cpp new file mode 100644 index 000000000000..b78a779a29d1 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/stepInputsScatterOp.cpp @@ -0,0 +1,57 @@ +/* + * 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/minimaxM3/stepInputsScatter.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! Issue a batch of packed host-to-device and device-to-device copies in one launch. +//! +//! @param stage Contiguous uint8 CUDA tensor laid out as described in kernels/minimaxM3/stepInputsScatter.h. Its +//! destinations are raw device addresses: the caller guarantees each one is live and that no other work on the stream +//! writes it between staging and this call. +void stepInputsScatter(torch::Tensor const& stage) +{ + TORCH_CHECK(stage.is_cuda(), "step_inputs_scatter expects a CUDA staging buffer"); + TORCH_CHECK(stage.scalar_type() == torch::kUInt8, "step_inputs_scatter expects a uint8 staging buffer"); + TORCH_CHECK(stage.is_contiguous(), "step_inputs_scatter expects a contiguous staging buffer"); + TORCH_CHECK(stage.numel() >= 16, "step_inputs_scatter staging buffer is shorter than its header"); + c10::cuda::CUDAGuard const deviceGuard{stage.device()}; + tensorrt_llm::kernels::invokeStepInputsScatter( + stage.data_ptr(), at::cuda::getCurrentCUDAStream(stage.get_device())); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("step_inputs_scatter(Tensor stage) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("step_inputs_scatter", &tensorrt_llm::torch_ext::stepInputsScatter); +} diff --git a/tensorrt_llm/_device_copy.py b/tensorrt_llm/_device_copy.py index 0aef62d3fa05..b54667235c45 100644 --- a/tensorrt_llm/_device_copy.py +++ b/tensorrt_llm/_device_copy.py @@ -9,20 +9,63 @@ class DeviceCopyBuffer: - def __init__(self, tensor: torch.Tensor, *, pin_memory: bool = False): + def __init__(self, tensor: torch.Tensor, *, pin_memory: bool = False) -> None: self._tensor = tensor self._pin_memory = pin_memory self._last_host_values = None + self._last_host_shape: torch.Size | None = None + self._copy_key: object = None @property def tensor(self) -> torch.Tensor: return self._tensor + def invalidate(self) -> None: + """Forget a pending or externally overwritten copy.""" + self._last_host_values = None + self._last_host_shape = None + self._copy_key = None + + def matches(self, host: torch.Tensor, *, key: object = None, exact: bool = False) -> bool: + """Whether the recorded host inputs match, including any copy-specific layout key. + + ``exact`` requires the same host shape instead of accepting a matching prefix. + Callers that transform or defer their copies must record only successful submissions + and invalidate the record if the submission is discarded or another writer intervenes. + """ + last = self._last_host_values + flat = host.reshape(-1) + return ( + last is not None + and key == self._copy_key + and last.dtype == host.dtype + and (not exact or self._last_host_shape == host.shape) + and last.numel() >= flat.numel() + and torch.equal(last[: flat.numel()], flat) + ) + + def record(self, host: torch.Tensor, *, key: object = None) -> None: + """Remember the inputs of a copy issued by this buffer or a specialized copy path.""" + flat = host.reshape(-1) + last = self._last_host_values + if ( + last is None + or key != self._copy_key + or last.dtype != host.dtype + or last.numel() <= flat.numel() + ): + self._last_host_values = flat.clone() + else: + last = last.clone() + last[: flat.numel()] = flat + self._last_host_values = last + self._last_host_shape = host.shape + self._copy_key = key + def update(self, host: torch.Tensor, *, force: bool = False) -> None: flat = host.reshape(-1) n = flat.numel() - last = self._last_host_values - if not force and last is not None and last.numel() >= n and torch.equal(last[:n], flat): + if not force and self.matches(host): if _verify_skipped_device_copies: if self._tensor.is_cuda: if torch.cuda.is_current_stream_capturing(): @@ -37,9 +80,4 @@ def update(self, host: torch.Tensor, *, force: bool = False) -> None: staging = torch.empty_like(flat, device="cpu", pin_memory=self._pin_memory) staging.copy_(flat) self._tensor.view(-1)[:n].copy_(staging, non_blocking=True) - if last is None or last.numel() <= n: - self._last_host_values = flat.clone() - else: - last = last.clone() - last[:n] = flat - self._last_host_values = last + self.record(host) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py index 16d312d3e7b1..f08e3a6be5b8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py @@ -32,6 +32,7 @@ NoTarget, modeling_v2_resolve, modeling_v2_resolve_on_all_ranks, + resolved_mode, step_contract_enabled, ) @@ -40,6 +41,7 @@ "ModelingV2Mode", "NoTarget", "modeling_v2_resolve", + "resolved_mode", "modeling_v2_resolve_on_all_ranks", "step_contract_enabled", ] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index 3814683f7fd7..1d714f7dd885 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -113,15 +113,31 @@ def of(cls, llm_args: "TorchLlmArgs") -> "ModelingV2Mode": return cls(llm_args.modeling_v2) -class NoTarget(Exception): - """Raised by a family's ``route`` when no target fits. +class NoTarget(ValueError): + """Raised by a family's ``route`` when no target fits, or by a target at construction when the runtime + the engine arguments set is not the one it is written for. The message says, in words, which fact ruled the targets out. Under - ``auto`` the resolver logs it and falls back to the built-in implementation; + ``auto`` the resolver logs it and falls back to the built-in implementation + (for a refusal at construction, ``AutoModelForCausalLM.from_config`` does); under ``require`` it becomes the error the user sees. """ +# The mode each target class was last resolved under (the loader builds the class right after resolving it). +_RESOLVED_MODES: Dict[type, ModelingV2Mode] = {} + + +def resolved_mode(target: type) -> ModelingV2Mode: + """The mode ``modeling_v2_resolve`` returned ``target`` under; ``auto`` for a class built without it. + + A target reads it at construction: under ``require`` its in-target dispatch + raises where it would otherwise fall back to a built-in path inside the + shapes it claims. + """ + return _RESOLVED_MODES.get(target, ModelingV2Mode.AUTO) + + def tp_reaches_across_nodes(mapping) -> bool: """Whether an attention-TP allreduce over ``mapping`` can reach every rank through the workspaces the default targets build. @@ -182,6 +198,7 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti when ``llm_args.modeling_v2`` is ``off``, when no ``route`` is registered for the architecture, or when that ``route`` raises ``NoTarget``. In ``require`` every one of those but ``off`` raises ``ValueError`` instead. + """ mode = ModelingV2Mode.of(llm_args) if mode is ModelingV2Mode.OFF: @@ -197,7 +214,7 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti if route is None: known = ", ".join(sorted(registered_routes())) or "(none)" raise NoTarget(f"no route for this architecture; routed architectures: {known}") - return route(config, llm_args, _device_compute_capability()) + target = route(config, llm_args, _device_compute_capability()) except NoTarget as why: if mode is ModelingV2Mode.REQUIRE: raise ValueError(f"modeling_v2='require', but no target fits {arch!r}: {why}") from None @@ -205,6 +222,9 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti f"modeling_v2: no target fits {arch!r} ({why}); using the built-in implementation" ) return None + if target is not None: + _RESOLVED_MODES[target] = mode + return target def _qualified(target: Optional[type]) -> Optional[str]: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/__init__.py index 42b425503ddf..1607ad9fcf76 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/__init__.py @@ -15,4 +15,5 @@ from .glm_moe_dsa import routing as glm_moe_dsa_routing # noqa: F401 from .gpt_oss import routing as gpt_oss_routing # noqa: F401 from .kimi_k3_vl import routing as kimi_k3_vl_routing # noqa: F401 +from .minimax_m3 import routing as minimax_m3_routing # noqa: F401 from .qwen4_exp import routing as qwen4_exp_routing # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/__init__.py new file mode 100644 index 000000000000..a3dfe4a6ce39 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMaxM3SparseForConditionalGeneration targets.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/fake_ops.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/fake_ops.py new file mode 100644 index 000000000000..76ea86dfc52e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/fake_ops.py @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Register MiniMax-M3 fake ops through their shared kernel module.""" + +from .kernels import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/__init__.py new file mode 100644 index 000000000000..44365a74ba2b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared MiniMax-M3 kernel support.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/fake_ops.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/fake_ops.py new file mode 100644 index 000000000000..a031ea7f9e70 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/fake_ops.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fake implementations of the MiniMax-M3 decode ops for compile and fake-tensor tracing. + +Target packages register these implementations when imported after routing. +""" + +import torch + + +@torch.library.register_fake("trtllm::minimax_m3_mxfp8_low_m_gemm") +def _(x, w, wsf, warps, cluster_k, row_tiles, pdl): + return x.new_empty((x.shape[0], w.shape[0])) + + +@torch.library.register_fake("trtllm::minimax_m3_router_gemv") +def _(x, w, rows, pdl): + return x.new_empty((x.shape[0], w.shape[0]), dtype=torch.float32) + + +def _minimax_m3_router_fp4_outputs(x, w): + m, k = x.shape + return [ + x.new_empty((m, w.shape[0]), dtype=torch.float32), + x.new_empty((m, k // 2), dtype=torch.uint8), + x.new_empty((m * (k // 16),), dtype=torch.uint8), + ] + + +@torch.library.register_fake("trtllm::minimax_m3_router_gemv_fp4") +def _(x, w, sf_scale, pdl): + return _minimax_m3_router_fp4_outputs(x, w) + + +@torch.library.register_fake("trtllm::minimax_m3_router_tc") +def _(x, w, sf_scale, pdl): + outputs = _minimax_m3_router_fp4_outputs(x, w) + return outputs if sf_scale is not None else outputs[:1] + + +@torch.library.register_fake("trtllm::minimax_m3_producer_pw") +def _( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + heads_q, + pdl, + early, +): + num_tokens = packed.shape[0] + return [ + packed.new_empty((num_tokens, heads_q, 128), dtype=torch.float8_e4m3fn), + packed.new_empty((num_tokens, 1, 128), dtype=torch.float8_e4m3fn), + ] + + +@torch.library.register_fake("trtllm::minimax_m3_mc_allreduce_norm") +def _( + x, + x2, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + perm, + scales, + maxt, + quant, + prepushed, + opts, +): + outputs = [torch.empty_like(residual), torch.empty_like(residual)] + if quant: + m, h = residual.shape + outputs += [ + residual.new_empty((m, h), dtype=torch.float8_e4m3fn), + residual.new_empty(((m + 127) // 128 * 128 * (h // 32),), dtype=torch.uint8), + ] + return outputs diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/__init__.py new file mode 100644 index 000000000000..4550fe8e4de7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/__init__.py @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 NVFP4 / sm_103 / tp4.""" + +from ..kernels import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py new file mode 100644 index 000000000000..c0bf5cd12b17 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py @@ -0,0 +1,4643 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ModelingV2 target: MiniMax-M3 NVFP4 / sm_103 / tp4 -- self-contained modeling code. + +MiniMax-M3's language model: 60 layers, hidden 6144, 64 query heads and 4 KV heads of dim 128. Layers 0-2 are dense +attention with a dense SwiGLU-OAI MLP; layers 3-59 are MSA block-sparse attention (an index branch scores 128-token +blocks and the top 16 are attended) with a MoE of 128 routed experts (top 4, sigmoid scores with a correction bias, +renormalized and scaled by 2) plus one shared expert. The checkpoint is a modelopt mixed-precision export: NVFP4 routed +experts, MXFP8 projections, fp32 router, bf16 embedding and head; the KV cache is NVFP4 for the sparse layers and FP8 +for the dense ones. + +``tp4`` is four GB300s: head-split attention (16 query heads, one KV head and one index head per rank) and +tensor-parallel experts (every expert on every rank, intermediate 768 per rank), no attention DP. Speculative +decoding, when configured, is one-model Eagle3 (draft length 3): the BF16 drafter (``M3Eagle3Draft``) and the BF16 +target head, both at checkpoint precision; decode steps of up to 8 requests verify and draft through fused kernels +at that precision (``M3Eagle3Worker``). + +**The core** (``MiniMaxM3Tp4ModelingV2``) holds the weights (``weights.MODEL_WEIGHTS``) and the runtime resources the +steps use -- one MSA attention backend per layer, the TP all-reduce communicator, the RoPE table, the MXFP8 GEMM +dispatch state the engine tunes -- and dispatches each step, on the host, from its context count alone +(``_select_target``): + +* ``DecodeTarget`` takes every pure generation step. Up to ``DECODE_MAX_TOKENS`` tokens (up to 8 requests, each one + token or a draft verify of 1 + 3 rows) with 1-D input ids it runs the decode-step kernels and picks each layer + boundary's collective by the step's token bucket (``DECODE_PLANS``); it runs its own generic layer for the first + layer when the fused embedding does not take the step, for every layer whose output Eagle3 captures, and for every + layer of a wider step. +* ``PrefillTarget`` takes every step with context rows: prefill chunks and mixed steps, every layer in its generic + form over the MSA backend. + +The two targets share no body: each owns its layer, attention and MoE compositions and its ``step_args``, the one +projection of the step metadata its forward reads. What both call are the module-level op adapters, each picking the +kernel for one call from shapes and layouts (``_linear``, ``_qk_norm_rope``, ``_sparse_qkv_producer``, the router and +MoE-runner calls, the all-reduce + norm calls), which read no step metadata themselves. + +Every fast path the decode step can take, the rows each covers (``FAST_PATHS``, each paired with the entry whose +``CELLS`` certify it) and what a step that cannot take one does are below: **uncovered** (outside every covered +range: log once, run the generic path) or **missed** (inside a covered range but a precondition is missing: raise +``FastPathUnavailable`` under ``modeling_v2='require'``, else log once and fall back). + +**Hand-offs between kernels.** A producer that already holds an activation in registers writes its MXFP8 form for the +next MXFP8 projection: the decode boundary all-reduce for the next layer's QKV, the sparse-decode merge and the fused +decode attention for the o_proj. It publishes it in ``ops/m3_prequant.SLOT`` keyed by the BF16 tensor; the next +MXFP8 projection on a decode graph takes it (``_linear``), and every take clears the slot. The decode boundaries' +push targets (the o_proj / shared-expert GEMVs and the MoE engine writing straight into the boundary all-reduce's +exchange) are passed explicitly from the layer to the producer and back. + +**What this target asserts rather than adapts**: SM 10.3, the topology and checkpoint (``routing.py``); the +MiniMax-M3 MSA sparse attention runtime with the fused five-way QKV/index projection, an FP8 indexer and an FP8 or +NVFP4 KV cache, no torch.compile and no piecewise prefill graphs (``_validate_runtime``); text-only input (the vision +tower is not built and a step carrying image input raises). +""" + +from __future__ import annotations + +import atexit +import copy +import dataclasses +import os +import threading +import time +from dataclasses import dataclass +from typing import ( + Any, + Callable, + Dict, + FrozenSet, + Hashable, + Iterable, + List, + Mapping, + Optional, + Tuple, +) + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + ModelingV2Mode, + NoTarget, + resolved_mode, + step_contract_enabled, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( + fp4_block_scale_moe_runner, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 + flashinfer_fused_add_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import _aligned_nvfp4_dequant_scales +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + MSA_REQUIRED_TOPK, + msa_paged_kv, +) +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.distributed import AllReduce, 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_speculative import ( + _ensure_draft_vocab_size, + _get_requested_draft_moe_backend, + _set_draft_kv_cache_quant_algo, + _slice_spec_position_ids, +) +from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import ( + MXFP8LinearMethod, + TensorParallelMode, + flashinfer_mxfp8_autotune_active, + flashinfer_mxfp8_decode_graph_capture_active, +) +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.moe.fused_moe.routing import RoutingMethodType +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy, RuntimeHooks +from tensorrt_llm._torch.speculative import SpecMetadata +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + should_use_separate_draft_kv_cache, +) +from tensorrt_llm._torch.utils import ActType_TrtllmGen +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantAlgo + +from . import weights as _weights +from .ops import m3_attn_mega as _attn_mega_kernel +from .ops import m3_attn_mega_c1 as _attn_c1_kernel +from .ops import m3_draft_step as m3_draft_step_kernel +from .ops import m3_fused_attn as _fused_attn_kernel +from .ops import m3_index_score as _index_score_kernel +from .ops import m3_mc_allreduce, m3_prequant, m3_spec_tail +from .ops import m3_sampled_accept as m3_sampled_accept_kernel +from .ops import m3_shard_accept as m3_shard_accept_kernel +from .ops import m3_target_head as m3_target_head_kernel +from .ops import msa_kv_update as _msa_kv_update_kernel +from .ops.activation.m3_silu_and_mul import m3_silu_and_mul +from .ops.attention.m3_attn_c1 import m3_attn_c1 +from .ops.attention.m3_attn_mega import m3_attn_mega +from .ops.attention.m3_draft_attention import m3_draft_attention +from .ops.attention.m3_fp8_indexer_qk_norm_rope import m3_fp8_indexer_qk_norm_rope +from .ops.attention.m3_fused_decode_attn import m3_fused_decode_attn +from .ops.attention.m3_index_score import SPLIT_K as INDEX_SCORE_SPLIT_K +from .ops.attention.m3_index_score import m3_index_score +from .ops.attention.m3_msa_attention import m3_msa_attention +from .ops.attention.m3_producer import m3_producer +from .ops.attention.m3_qk_norm_rope_fp8 import m3_qk_norm_rope_fp8 +from .ops.attention.m3_qkv_insert import m3_qkv_insert +from .ops.attention.m3_select_blocks import m3_select_blocks +from .ops.attention.m3_select_long import MAX_BLOCKS as SELECT_LONG_MAX_BLOCKS +from .ops.attention.m3_select_long import m3_select_long +from .ops.attention.m3_sparse_decode import m3_sparse_decode +from .ops.comm.m3_allreduce import m3_allreduce +from .ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm +from .ops.comm.m3_mc_finalize_allreduce_norm import m3_mc_finalize_allreduce_norm +from .ops.embedding.m3_embedding import m3_embedding +from .ops.gemm.m3_bf16_linear import m3_bf16_linear +from .ops.gemm.m3_fp32_linear import m3_fp32_linear +from .ops.gemm.m3_mx_gemv import ( + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + SHARED_DOWN, + SHARED_GATE_UP, + m3_mx_gemv, +) +from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 +from .ops.gemm.m3_mxfp8_gemm import m3_mxfp8_gemm +from .ops.gemm.m3_mxfp8_linear import m3_mxfp8_linear +from .ops.gemm.m3_target_head import MAX_ROWS as TARGET_HEAD_KERNEL_MAX_ROWS +from .ops.gemm.m3_target_head import m3_target_head +from .ops.gemm.mxfp8_low_m_gemm import mxfp8_low_m_gemm +from .ops.long_select import long_select_blocks +from .ops.m3_router_gemv import m3_router_gemv as _router_gemv +from .ops.m3_router_gemv import m3_router_gemv_fp4 as _router_gemv_fp4 +from .ops.m3_router_gemv import m3_router_tc as _router_tc +from .ops.moe.m3_moe16_engine import M3Moe16State, m3_moe16_engine +from .ops.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine +from .ops.moe.m3_router_gemv import m3_router_gemv +from .ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 +from .ops.moe.m3_router_tc import m3_router_tc +from .ops.mx_gemv import op as mx_gemv +from .ops.mxfp8_low_m_gemm import mxfp8_low_m_gemm as _low_m_gemm +from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm +from .ops.spec.m3_draft_input import m3_draft_input +from .ops.spec.m3_draft_step import MAX_TOKENS as DRAFT_STEP_KERNEL_MAX_TOKENS +from .ops.spec.m3_draft_step import m3_draft_step +from .ops.spec.m3_force_vec import m3_force_vec +from .ops.spec.m3_shard_accept import m3_shard_accept +from .ops.triton_sparse_decode import MERGE_MXFP8, MERGE_MXFP8_MIN_TOKENS + +_SM = (10, 3) +# -------------------------------------------------------------------------------------------------------------------- +# Row ranges +# -------------------------------------------------------------------------------------------------------------------- + +# The largest generation step the decode kernels take: 8 requests verifying 1 + 3 drafts and their padding rows. A +# wider pure generation step stays on the decode target and runs its generic layers. +DECODE_MAX_TOKENS = 48 # capacity: fixed: the decode kernels' step bound + + +@dataclass(frozen=True) +class DecodePlan: + """The layer-boundary collectives of one decode token bucket (GB300 tp4 measurements in the comments).""" + + name: str + max_tokens: int + # The attention boundary's multicast all-reduce + residual + RMSNorm (5.05 vs 5.27 us native at 1 token). + attn_mc_ar: bool + # The same kernel at the MoE boundary with the shared + routed add folded in: wins up to 48 tokens. + moe_mc_ar: bool + # The MoE-boundary all-reduce also emits the next QKV activation in MXFP8 for the decode GEMV / GEMM. + qkv_prequant: bool + + +# tiny: every step up to 8 tokens (c1..c8 without speculation, the c1 verify, every draft pass); small: the c2 +# verify; verify: the c4 / c8 verifies. Each is one multicast all-reduce bucket (ops/m3_mc_allreduce.BUCKETS). +DECODE_PLANS = ( + DecodePlan("tiny", 8, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), + DecodePlan("small", 16, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), + DecodePlan("verify", DECODE_MAX_TOKENS, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), +) + +# The sparse layers' NVFP4 QKV + indexer producer (ops/attention/m3_producer, every step its geometry fits) triggers its +# dependent kernel on entry, so the index scorer launches and stages its pre-wait loads while the producer runs (its +# outputs are read only after the scorer's own grid wait). +PRODUCER_EARLY_TRIGGER = True + +# The fused embedding + first RMSNorm kernel: one owner per row count and index dtype. +EMBED_FIRSTNORM_MAX_TOKENS = 8 # capacity: fixed: built rows; wider steps embed generically +# The o_proj pushes its rows into the attention boundary's multicast all-reduce (prepushed mode 2) up to this many +# tokens; wider steps run the GEMV and a plain multicast all-reduce. +ATTN_PUSH_MAX_TOKENS = 16 # capacity: fixed: measured push crossover +AR_PUSH_MODE = 2 +# The 1..8-token MoE engine (router, routed experts and the shared expert in one launch) pushes its sum into the MoE +# boundary's all-reduce (prepushed mode 2). +MOE_ENGINE_MAX_TOKENS = 8 # capacity: fixed: the engine's built rows +MOE_PUSH_MAX_TOKENS = MOE_ENGINE_MAX_TOKENS # capacity: fixed: the engine's rows +# The 9..13-row engine (router, routed and shared experts in one kernel); its bf16 sum goes to the boundary. +MOE16_MIN_TOKENS = 9 # capacity: fixed: the 9..13-row engine's lower bound +MOE16_MAX_TOKENS = 13 # capacity: fixed: the 9..13-row engine's built rows +# From this many tokens the MoE runs unfinalized through the TRTLLM-Gen runner and its finalize (top-k gather, +# routing weights, shared expert add) runs inside the boundary all-reduce. 9..13 take it only when their engine is +# missing (outside require mode). +MC_FINALIZE_MIN_TOKENS = 9 # capacity: fixed: the rows above the engines +# The router: the fp32 GEMV at one row, the tensor-core router (fused with the routed experts' NVFP4 activation +# quantization) at 2..48 rows. GB300 tp4, behind a PDL producer, with the NVFP4 activation: 2.90 / 2.93 / 3.20 / +# 3.82 / 7.36 us at 1 / 4 / 8 / 12 / 48 tokens against the fp32 GEMV's 2.85 / 3.96 / 4.88 / 6.01 / - us. +ROUTER_GEMV_MAX_TOKENS = 1 # capacity: fixed: measured router crossover +ROUTER_TC_MAX_TOKENS = 48 # capacity: fixed: the tensor-core router's built rows +# The fused norm/RoPE/KV-insert/sparse-decode/merge kernel: up to 8 query rows per request and 16 query rows per +# step (above, its 16-CTA clusters run in waves and the three-kernel path is faster: GB300 -13% fAL spec c8). +FUSED_ATTN_MAX_QUERY_LEN = 8 # capacity: fixed: compiled query rows +FUSED_ATTN_MAX_TOKENS = 16 # capacity: fixed: measured wave crossover +# One request verifying a draft of 6 query rows runs attention + o_proj as one launch (96 CTAs: 6 clusters of 16). +ATTN_C1_QUERY_LEN = 6 # capacity: fixed: the verify kernel's query rows +# A greedy Eagle3 verify step of up to this many requests verifies with the fused shard argmax + acceptance kernel +# over the target lm_head's BF16 shard logits (ops/m3_shard_accept), at the configured draft length. +SHARD_ACCEPT_DRAFT_LEN = 3 # capacity: fixed: the built Eagle3 draft length +SHARD_ACCEPT_MAX_REQUESTS = DECODE_MAX_TOKENS // ( + SHARD_ACCEPT_DRAFT_LEN + 1 +) # capacity: fixed: decode rows / (draft length + 1) +# A fused-verify step of up to this many verify rows (one or two requests at draft length 3) computes the target +# logits shard from the 12-bit coded copy of the BF16 head (ops/m3_target_head; the same BF16 x BF16 -> fp32 +# arithmetic, 75% of the bytes) instead of the in-tree head GEMM; wider verifies keep the in-tree GEMM. +TARGET_HEAD_MAX_ROWS = TARGET_HEAD_KERNEL_MAX_ROWS +# On a fused-verify step of up to this many verify rows the drafter's first fc is forked onto the aux stream right +# after the last captured target layer, overlapping the remaining target layers and the head (spec_worker). +FC_EARLY_MAX_ROWS = 8 # capacity: fixed: measured fork crossover +# A decode-only Eagle3 step of up to this many requests drafts with the fused draft loop (spec_worker): per pass the +# draft input in one launch (ops/m3_spec_tail.draft_input), the drafter layer's in-tree attention, MLP and +# all-reduces, and one kernel for the draft head (its 12-bit coded copy, or the BF16 fragments when the coding does +# not decode bit for bit), the cross-rank argmax and the pass's bookkeeping (ops/m3_draft_step), at the +# configured draft length. The head kernel takes one row per request, at most 8. +# The fused draft loop streams the 12-bit coded copy of the draft head (when it decodes bit for bit); False keeps the +# BF16 fragments. +DRAFT_HEAD_CODED = True +DRAFT_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN # capacity: fixed: the built draft length +DRAFT_STEP_MAX_REQUESTS = DRAFT_STEP_KERNEL_MAX_TOKENS +# The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows +# (ops/triton_sparse_decode.MQ_MIN_TOKENS). +SPARSE_DECODE_MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query crossover +# The long-row top-16 selector takes engines of 129..LONG_SELECT_MAX_BLOCKS KV blocks of 128 tokens; it holds 2, 4 or +# 9 keys per thread by the row's valid-block count. At most 128 blocks the built-in selector runs. +LONG_SELECT_MIN_BLOCKS = 129 # capacity: fixed: past the built-in selector's 128 +LONG_SELECT_MAX_BLOCKS = SELECT_LONG_MAX_BLOCKS +LONG_SELECT_KEY_WIDTHS = (2, 4, 9) +# The long-row select releases the sparse decode after it right after its own grid wait (the decode stages q and the +# request length before its wait) up to this many query rows; on wider steps the early-launched attention CTAs crowd +# the select's CTAs off the SMs (8 requests x 4 rows: +3.1 us per layer), so they keep the trigger at the select's exit. +SELECT_EARLY_MAX_ROWS = 8 # capacity: fixed: measured early-release crossover +# The MSA index scorer (ops/attention/m3_index_score) prefetches the older index-K pages into L2 before its grid wait +# at every batch; up to this many requests it also issues Q and the first K tile right after the wait (larger batches +# are bandwidth bound, where the extra early tiles only crowd the prefetch). +INDEX_SCORE_EARLY_MAX_REQUESTS = 1 # capacity: fixed: measured early-issue crossover + +# The decode MXFP8 GEMVs, per projection (N, K): the rows each serves. On an activation its producer quantized +# (the MoE boundary for the QKV, the attention merge for the o_proj) the GEMV serves every row up to 48 but the +# unsplit K=768 shared down projection at 17..32 rows, where FlashInfer's GEMM is faster. +_ALL_ROWS = frozenset(range(1, DECODE_MAX_TOKENS + 1)) +MX_GEMV_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { + QKV_SPARSE: _ALL_ROWS, + QKV: _ALL_ROWS, + O_PROJ: _ALL_ROWS, + SHARED_GATE_UP: _ALL_ROWS, + SHARED_DOWN: _ALL_ROWS - frozenset(range(17, 33)), + DENSE_GATE_UP: _ALL_ROWS, + DENSE_DOWN: _ALL_ROWS, +} +# On a bf16 activation the GEMV quantizes it in its prologue, which pays only up to 8 32-element blocks per thread: +# the most rows each projection takes. +MX_GEMV_BF16_MAX_TOKENS: Dict[Tuple[int, int], int] = { + QKV_SPARSE: 16, + QKV: 16, + O_PROJ: 32, + SHARED_GATE_UP: 16, + SHARED_DOWN: 16, + DENSE_GATE_UP: 8, + DENSE_DOWN: 16, +} +# The fused low-M MXFP8 GEMM (projections marked for it), per projection: its built configurations' row limit +# (ops/mxfp8_low_m_gemm). The bf16 GEMV serves every one of these rows first; the GEMM takes them only if it +# declines. +LOW_M_GEMM_MAX_TOKENS: Dict[Tuple[int, int], int] = { + QKV_SPARSE: 8, + O_PROJ: 8, + DENSE_GATE_UP: 1, + DENSE_DOWN: 8, +} + +Key = Hashable + + +@dataclass(frozen=True) +class RowRange: + """The shapes one fast path takes, and the catalog entry whose cells certify each of them. + + ``keys`` are the covered shapes: row counts for most paths, (query rows per request, requests) for the attention + kernels, a variant name for thresholds. ``certifies`` maps a cell's spec of ``entry`` to the keys that cell + drives (empty for a cell of another variant). + """ + + path: str + keys: FrozenSet[Key] + entry: OpWrapper + certifies: Callable[[Mapping[str, Any]], Iterable[Key]] + + def covers(self, key: Key) -> bool: + return key in self.keys + + def certified(self) -> FrozenSet[Key]: + return frozenset(k for cell in self.entry.CELLS for k in self.certifies(cell.spec)) + + +def _rows(lo: int, hi: int) -> FrozenSet[int]: + return frozenset(range(lo, hi + 1)) + + +def _spec_rows(field: str, **match: Any) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + """A cell's ``spec[field]`` when the cell's other fields match ``match``.""" + + def certifies(spec: Mapping[str, Any]) -> Iterable[Key]: + if any(spec.get(name) != value for name, value in match.items()): + return () + return (spec[field],) + + return certifies + + +def _const(*keys: Key) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + return lambda spec: keys + + +def _fused_attn_keys() -> FrozenSet[Tuple[int, int]]: + return frozenset( + (qlen, reqs) + for qlen in range(1, FUSED_ATTN_MAX_QUERY_LEN + 1) + for reqs in range(1, FUSED_ATTN_MAX_TOKENS // qlen + 1) + ) + + +def _sparse_decode_producer(spec: Mapping[str, Any]) -> Iterable[Key]: + if not spec["identity"]: + return () + rows = spec["decode_query_len"] * len(spec["seq_lens"]) + return ("multi_query" if rows >= SPARSE_DECODE_MQ_MIN_TOKENS else "per_row",) + + +def _long_select_widths(spec: Mapping[str, Any]) -> Iterable[Key]: + """The key widths a cell's rows reach: every width up to its blocks' when rows of mixed lengths are drawn.""" + blocks = spec["num_blocks"] + widest = 2 if blocks <= 2048 else 4 if blocks <= 4096 else 9 + if spec["valid"] == "mixed": + return tuple(w for w in LONG_SELECT_KEY_WIDTHS if w <= widest) + return (widest,) + + +EMBED_FIRSTNORM = RowRange( + "embed_firstnorm", + _rows(1, EMBED_FIRSTNORM_MAX_TOKENS), + m3_embedding_firstnorm, + _spec_rows("rows"), +) +# One request decoding one token on an MXFP8 layer input: QKV GEMV + attention + o_proj in one launch. +ATTN_MEGA = RowRange("attn_mega", frozenset({(1, 1)}), m3_attn_mega, _const((1, 1))) +ATTN_C1 = RowRange( + "attn_c1", frozenset({(ATTN_C1_QUERY_LEN, 1)}), m3_attn_c1, _const((ATTN_C1_QUERY_LEN, 1)) +) +FUSED_DECODE_ATTN = RowRange( + "fused_decode_attn", + _fused_attn_keys(), + m3_fused_decode_attn, + lambda spec: ((spec["qlen"], len(spec["seq_lens"])),), +) +SPARSE_DECODE_PRODUCER = RowRange( + "sparse_decode_producer", + frozenset({"per_row", "multi_query"}), + m3_sparse_decode, + _sparse_decode_producer, +) +# Every token count of a sparse layer on an NVFP4 KV cache at the TP4 heads (prefill and decode): the producer takes the +# geometry, not a row range. +PRODUCER = RowRange("m3_producer", frozenset({"nvfp4_tp4"}), m3_producer, _const("nvfp4_tp4")) +INDEX_SCORE = RowRange( + "m3_index_score", + frozenset({"early", "prefetch"}), + m3_index_score, + lambda spec: ("early" if spec["early"] else "prefetch",), +) +LONG_SELECT = RowRange( + "long_select", frozenset(LONG_SELECT_KEY_WIDTHS), m3_select_long, _long_select_widths +) +MC_ALLREDUCE = RowRange( + "mc_allreduce_norm", + _rows(1, DECODE_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=0), +) +ATTN_PUSH = RowRange( + "attn_push", + _rows(1, ATTN_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +MOE_PUSH = RowRange( + "moe_push", + _rows(1, MOE_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +MC_FINALIZE = RowRange( + "mc_finalize_allreduce_norm", + _rows(MC_FINALIZE_MIN_TOKENS, DECODE_MAX_TOKENS), + m3_mc_finalize_allreduce_norm, + _spec_rows("m"), +) +MOE_ENGINE = RowRange( + "m3_moe_engine", _rows(1, MOE_ENGINE_MAX_TOKENS), m3_moe_engine, _spec_rows("tokens") +) +MOE16_ENGINE = RowRange( + "m3_moe16_engine", + _rows(MOE16_MIN_TOKENS, MOE16_MAX_TOKENS), + m3_moe16_engine, + _spec_rows("tokens"), +) +ROUTER_GEMV = RowRange( + "m3_router_gemv", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv, _spec_rows("m") +) +ROUTER_GEMV_FP4 = RowRange( + "m3_router_gemv_fp4", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv_fp4, _spec_rows("m") +) +ROUTER_TC = RowRange( + "m3_router_tc", + _rows(ROUTER_GEMV_MAX_TOKENS + 1, ROUTER_TC_MAX_TOKENS), + m3_router_tc, + _spec_rows("m"), +) +# The GEMM paths are keyed (projection (N, K), rows): each projection runs its own split-K / ring / MMA +# configuration, so a row count is certified per projection. +MX_GEMV = RowRange( + "m3_mx_gemv", + frozenset((nk, m) for nk, rows in MX_GEMV_ROWS.items() for m in rows), + m3_mx_gemv, + lambda spec: ((spec["nk"], spec["m"]),), +) +MX_GEMV_BF16 = RowRange( + "m3_mx_gemv_bf16", + frozenset((nk, m) for nk, hi in MX_GEMV_BF16_MAX_TOKENS.items() for m in _rows(1, hi)), + m3_mx_gemv_bf16, + lambda spec: ((spec["nk"], spec["m"]),), +) +LOW_M_GEMM = RowRange( + "mxfp8_low_m_gemm", + frozenset((nk, m) for nk, hi in LOW_M_GEMM_MAX_TOKENS.items() for m in _rows(1, hi)), + mxfp8_low_m_gemm, + lambda spec: (((spec["n"], spec["k"]), spec["m"]),), +) +# (draft length, requests) of a greedy verify step. +SHARD_ACCEPT = RowRange( + "m3_shard_accept", + frozenset((SHARD_ACCEPT_DRAFT_LEN, n) for n in range(1, SHARD_ACCEPT_MAX_REQUESTS + 1)), + m3_shard_accept, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) + +# Verify rows of a fused-verify step whose target logits come from the coded head. +TARGET_HEAD = RowRange( + "m3_target_head", + _rows(1, TARGET_HEAD_MAX_ROWS), + m3_target_head, + _spec_rows("m"), +) + +# (draft length, requests) of a fused draft loop, and the rows of its passes' draft input: the first pass's +# requests x (draft length + 1), one per request on the later passes. +DRAFT_STEP = RowRange( + "m3_draft_step", + frozenset((DRAFT_STEP_DRAFT_LEN, n) for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1)), + m3_draft_step, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) +DRAFT_INPUT = RowRange( + "m3_draft_input", + frozenset( + rows + for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1) + for rows in (n, n * (DRAFT_STEP_DRAFT_LEN + 1)) + ), + m3_draft_input, + _spec_rows("m"), +) + +FAST_PATHS: Tuple[RowRange, ...] = ( + EMBED_FIRSTNORM, + ATTN_MEGA, + ATTN_C1, + FUSED_DECODE_ATTN, + SPARSE_DECODE_PRODUCER, + PRODUCER, + INDEX_SCORE, + LONG_SELECT, + MC_ALLREDUCE, + ATTN_PUSH, + MOE_PUSH, + MC_FINALIZE, + MOE_ENGINE, + MOE16_ENGINE, + ROUTER_GEMV, + ROUTER_GEMV_FP4, + ROUTER_TC, + MX_GEMV, + MX_GEMV_BF16, + LOW_M_GEMM, + SHARD_ACCEPT, + TARGET_HEAD, + DRAFT_STEP, + DRAFT_INPUT, +) + +# -------------------------------------------------------------------------------------------------------------------- +# Selection: pure functions of the step shape +# -------------------------------------------------------------------------------------------------------------------- + + +def plan_for(num_tokens: int) -> DecodePlan: + for plan in DECODE_PLANS: + if num_tokens <= plan.max_tokens: + return plan + raise AssertionError(f"no decode plan covers {num_tokens} tokens") + + +def is_decode_step(num_tokens: int) -> bool: + """Whether a pure generation step of ``num_tokens`` rows runs the decode kernels (else generic layers).""" + return num_tokens <= DECODE_MAX_TOKENS + + +def moe_path(num_tokens: int) -> Optional[str]: + """The MoE kernel of a decode step: an engine, the unfinalized MoE + finalize all-reduce, or None (generic).""" + if MOE_ENGINE.covers(num_tokens): + return MOE_ENGINE.path + if MOE16_ENGINE.covers(num_tokens): + return MOE16_ENGINE.path + if MC_FINALIZE.covers(num_tokens): + return MC_FINALIZE.path + return None + + +def router_path(num_tokens: int) -> Optional[str]: + """The router kernel when the MoE runs unfused: the fp32 GEMV, the tensor-core router, or None (built-in gate).""" + if ROUTER_GEMV.covers(num_tokens): + return ROUTER_GEMV.path + if ROUTER_TC.covers(num_tokens): + return ROUTER_TC.path + return None + + +MX_GEMV_PATH = MX_GEMV.path +MX_GEMV_BF16_PATH = MX_GEMV_BF16.path +LOW_M_GEMM_PATH = LOW_M_GEMM.path +# FlashInfer's CuTe DSL MXFP8 GEMM on the producer's MXFP8 activation (rows the GEMV does not serve). +PREQUANT_GEMM_PATH = "flashinfer_mxfp8_gemm" +# The MXFP8 quantize (tuned backend) and FlashInfer's CuTe DSL GEMM: every row the decode kernels do not serve. +GENERIC_GEMM_PATH = "mxfp8_gemm" + + +def linear_path(n: int, k: int, num_tokens: int, prequant: bool, low_m_gemm: bool) -> str: + """The decode-graph MXFP8 GEMM of an [N, K] projection on ``num_tokens`` rows. + + ``prequant``: the producer published the activation in MXFP8. ``low_m_gemm``: the projection is marked for the + fused low-M GEMM. + """ + if prequant: + if num_tokens in MX_GEMV_ROWS.get((n, k), ()): + return MX_GEMV_PATH + return PREQUANT_GEMM_PATH + if num_tokens <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0): + return MX_GEMV_BF16_PATH + if low_m_gemm and num_tokens <= LOW_M_GEMM_MAX_TOKENS.get((n, k), 0): + return LOW_M_GEMM_PATH + return GENERIC_GEMM_PATH + + +# -------------------------------------------------------------------------------------------------------------------- +# Fallback policy +# -------------------------------------------------------------------------------------------------------------------- + + +class FastPathUnavailable(RuntimeError): + """A covered decode shape could not take its fast path under ``modeling_v2='require'``.""" + + +_lock = threading.Lock() +_require = [False] +_logged: set = set() + + +def set_require(require: bool) -> None: + """Whether a miss raises (``modeling_v2='require'``); set by the target at construction.""" + _require[0] = bool(require) + + +def require_mode() -> bool: + return _require[0] + + +def _log_once(level: Callable[[str], None], path: str, reason: str, message: str) -> None: + key = (path, reason) + with _lock: + if key in _logged: + return + _logged.add(key) + level(message) + + +def missed(path: str, rows: Optional[int], reason: str) -> None: + """``path`` covers this step but cannot run: raise under require, else log once; the caller falls back.""" + at = "" if rows is None else f" at {rows} rows" + if _require[0]: + raise FastPathUnavailable( + f"modeling_v2='require': the MiniMax-M3 decode path {path!r} covers this step{at} but cannot run: " + f"{reason}" + ) + _log_once( + logger.warning, + path, + reason, + f"[M3 dispatch] {path}{at}: {reason}; falling back (logged once per reason)", + ) + + +def capture_missed(path: str, rows: Optional[int], what: str) -> None: + """``what`` is built on eager steps but a CUDA-graph capture found it missing (see ``missed``).""" + missed(path, rows, f"{what} is not built and a CUDA graph is being captured") + + +def uncovered(path: str, reason: str) -> None: + """The step is outside ``path``'s covered shapes: log once per (path, reason); the generic path runs.""" + _log_once(logger.info, path, reason, f"[M3 dispatch] {path} not taken: {reason} (logged once)") + + +def reset_logged() -> None: + """Forget which (path, reason) pairs were logged.""" + with _lock: + _logged.clear() + + +# -------------------------------------------------------------------------------------------------------------------- +# Target constants +# -------------------------------------------------------------------------------------------------------------------- + +# The generic boundary takes the low-latency multicast all-reduce + residual + RMSNorm up to this many tokens; with +# the shared + routed add folded in, up to its own maximum (m3_mc_allreduce.MAX_TOKENS). +MC_AR_MAX_TOKENS = 8 # capacity: fixed: measured crossover vs the one-shot all-reduce +# The shared expert folds into the routed engine from one token (dedicated shared-expert CTAs). +MOE_FOLD_MIN_TOKENS = 1 # capacity: fixed: a row lower bound +# SMs left to the concurrent kernels while a persistent routed engine runs. +MOE_ENGINE_RESERVED_SMS = 24 # capacity: fixed: an SM split, not a size +# Id of the target's final hidden state in Eagle3's ``layers_to_capture``. +FINAL_HIDDEN_LAYER_ID = -1 # capacity: fixed: an Eagle3 layer id, not a size +# The trtllm-gen MoE runner's routing kind for MiniMax-M3 (sigmoid + correction bias, renormalized, scaled) and its +# gated activation (SwiGLU with the per-expert alpha / beta / clamp operands). +ROUTING_MINIMAX2 = int(RoutingMethodType.MiniMax2) +ACT_SWIGLU = int(ActType_TrtllmGen.SwiGlu) + +# -------------------------------------------------------------------------------------------------------------------- +# The runtime-value dump (TRTLLM_MODELING_V2_VALUE_DUMP=1): the metadata of every fast-path and entry decision a step +# takes, recorded on eager steps and captures (replay runs no Python), printed once at exit. +# -------------------------------------------------------------------------------------------------------------------- + +_VALUE_DUMP = os.environ.get("TRTLLM_MODELING_V2_VALUE_DUMP") == "1" +_VALUES: set = set() + + +def _record(*meta: Hashable) -> None: + if _VALUE_DUMP: + _VALUES.add(meta) + + +def _print_values() -> None: + if _VALUES: + rank = os.environ.get("OMPI_COMM_WORLD_RANK", os.environ.get("RANK", "0")) + for meta in sorted(_VALUES, key=repr): + print(f"[modeling_v2 value dump] rank {rank}: {meta}", flush=True) + + +if _VALUE_DUMP: + atexit.register(_print_values) + +# -------------------------------------------------------------------------------------------------------------------- +# Configuration: the text tower and the runtime the target assumes +# -------------------------------------------------------------------------------------------------------------------- + + +def _is_vl_config(pretrained_config: PretrainedConfig) -> bool: + model_type = getattr(pretrained_config, "model_type", None) + if model_type == "minimax_m3_vl": + return True + if model_type == "minimax_m3": + return False + return any( + "ForConditionalGeneration" in a + for a in (getattr(pretrained_config, "architectures", None) or []) + ) + + +def _text_model_config(model_config: ModelConfig) -> ModelConfig: + """The model config with the language model's ``text_config`` as its pretrained config. + + The checkpoint declares the vision-language wrapper; its ``text_config`` arrives as a plain dict (no typed + sub-config is registered), so it is materialized into a ``PretrainedConfig`` with the wrapper's dtype and + embedding tie carried over. A text config is returned unchanged. + """ + if not _is_vl_config(model_config.pretrained_config): + return model_config + cfg = copy.deepcopy(model_config) + text = getattr(cfg.pretrained_config, "text_config", None) + if text is None: + raise ValueError("MiniMax-M3 VL config has no `text_config`") + if isinstance(text, dict): + out = PretrainedConfig() + for k, v in text.items(): + setattr(out, k, v) + text = out + if ( + getattr(text, "torch_dtype", None) is None + and getattr(cfg.pretrained_config, "torch_dtype", None) is not None + ): + text.torch_dtype = cfg.pretrained_config.torch_dtype + if ( + getattr(text, "tie_word_embeddings", None) is None + and getattr(cfg.pretrained_config, "tie_word_embeddings", None) is not None + ): + text.tie_word_embeddings = cfg.pretrained_config.tie_word_embeddings + return dataclasses.replace(cfg, pretrained_config=text) + + +def _validate_runtime(model_config: ModelConfig) -> None: + """The runtime this target is written for. It is set by engine arguments rather than the checkpoint, so routing + sees only its KV cache algorithm; this check at construction refuses every other runtime. + + The MiniMax-M3 MSA backend (it owns the step metadata and the index-K cache), the fused five-way QKV/index + projection or the separate QKV and index projections, an FP8 indexer and an FP8 or NVFP4 main KV cache. + """ + sparse = model_config.sparse_attention_config + if ( + sparse is None + or getattr(sparse, "algorithm", None) != "minimax_m3" + or sparse.implementation != "msa" + ): + raise NoTarget( + "this MiniMax-M3 target requires sparse_attention_config.algorithm='minimax_m3' with implementation='msa'" + ) + if getattr(sparse, "indexer_kv_dtype", None) != "fp8": + raise NoTarget( + "this MiniMax-M3 target requires sparse_attention_config.indexer_kv_dtype='fp8'" + ) + mode = model_config.quant_config.quant_mode + if mode is None or not (mode.has_fp8_kv_cache() or mode.has_fp4_kv_cache()): + raise NoTarget("this MiniMax-M3 target requires an FP8 or NVFP4 KV cache") + + +def _sparse_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = cfg.sparse_attention_config["sparse_attention_freq"] + return [i for i, f in enumerate(freq) if int(f) != 0] + + +def _moe_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = getattr(cfg, "moe_layer_freq", None) + if freq is None: + return list(range(cfg.num_hidden_layers)) + return [i for i, f in enumerate(freq) if int(f) != 0] + + +class _Mxfp8GemmState(nn.Module): + """The MXFP8 GEMM backend state the engine's warmup drives. + + The engine finds MXFP8 projections by a module's ``quant_method`` and, for every one, opts it into native GEMM + tuning and (on decode CUDA graphs) the auto-dispatched FlashInfer GEMM, then marks each tuned. Every projection of + this target dispatches identically, so one method object carries that state for all of them; the forward reads it + at each call (``_mx_mode``) and passes the resolved backend to the GEMM entries. + """ + + def __init__(self) -> None: + super().__init__() + self.quant_method = MXFP8LinearMethod() + + +@dataclass(frozen=True) +class MxMode: + """The MXFP8 GEMM dispatch of one call, resolved from ``_Mxfp8GemmState`` and the engine's contexts.""" + + decode_graph: bool # inside a decode CUDA-graph warmup or capture + autotune: bool # inside the engine's FlashInfer MXFP8 tuning pass + tuned_quantize_gemm: ( + bool # per-bucket tuned quantize + GEMM (decode graphs, tune_decode_graph_backends) + ) + flashinfer: bool # FlashInfer's MXFP8 GEMM on the quantized activation + native_tuning: ( + bool # the native GEMM through the autotuner (warmup, before its tactics are registered) + ) + + +def _mx_mode(state: _Mxfp8GemmState) -> MxMode: + method = state.quant_method + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + autotune = flashinfer_mxfp8_autotune_active() + return MxMode( + decode_graph=decode_graph, + autotune=autotune, + tuned_quantize_gemm=method.tune_decode_graph_backends and decode_graph, + flashinfer=method.backend == "flashinfer" + or ( + method.backend == "auto" + and (autotune or (method._flashinfer_autotuned and decode_graph)) + ), + native_tuning=method.needs_native_autotune, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# Projections: the MXFP8 GEMM dispatch (``linear_path``) and its hand-offs +# -------------------------------------------------------------------------------------------------------------------- + + +def _mxfp8_linear( + core, x: torch.Tensor, w: torch.Tensor, s: torch.Tensor, mode: MxMode +) -> torch.Tensor: + """One MXFP8 projection on the engine-tuned path: MXFP8 quantize + GEMM, backend per ``mode``.""" + _record( + "m3_mxfp8_linear", + tuple(w.shape), + mode.tuned_quantize_gemm, + mode.flashinfer, + mode.native_tuning, + ) + return m3_mxfp8_linear( + x.contiguous(), + w, + s, + core.mx_alpha, + torch.bfloat16, + mode.tuned_quantize_gemm, + mode.autotune, + mode.flashinfer, + mode.native_tuning, + ) + + +def _gemv_variants(n: int, k: int, push: Optional[str], world: int): + """Every (rows, bf16 activation, push (world, bucket tokens, mode)) call the decode dispatch routes to an [N, K] + projection's GEMV: an MXFP8 activation on the rows ``MX_GEMV_ROWS`` gives it, a bf16 one up to + ``MX_GEMV_BF16_MAX_TOKENS``; each pushing when the projection's boundary push covers the rows.""" + push_rows = {"attn": ATTN_PUSH.keys, "moe": MOE_PUSH.keys}.get(push, frozenset()) + variants = [] + for rows in range(1, DECODE_MAX_TOKENS + 1): + for x_bf16, served in ( + (False, rows in MX_GEMV_ROWS.get((n, k), ())), + (True, rows <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0)), + ): + if not served: + continue + variants.append((rows, x_bf16, None)) + if rows in push_rows: + bucket = next(b for b in m3_mc_allreduce.BUCKETS if rows <= b) + variants.append((rows, x_bf16, (world, bucket, AR_PUSH_MODE))) + return variants + + +def _low_m_gemm_supports(x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor) -> bool: + """Whether the fused MXFP8 low-M kernel handles these operands.""" + return ( + x.dim() == 2 + and x.dtype in (torch.bfloat16, torch.float16) + and x.stride(-1) == 1 + and weight.dtype == torch.float8_e4m3fn + and weight.is_contiguous() + and weight_scale.dtype == torch.uint8 + and x.shape[1] == weight.shape[1] + and weight.shape[1] % 256 == 0 + ) + + +def _uncompiled_note() -> str: + key = mx_gemv.last_uncompiled + return "" if key is None else f" (its kernel {key} was not compiled before warmup)" + + +def _linear( + core, x: torch.Tensor, wk: str, sk: str, low_m_gemm: bool, push: Optional[tuple] = None +) -> Tuple[torch.Tensor, bool]: + """One MXFP8 projection ``x @ w[wk]^T``; returns (output, whether its GEMV pushed into ``push``). + + Outside decode CUDA graphs: the engine-tuned MXFP8 quantize + GEMM. Inside them the projection's (N, K) and the + step's rows pick one path (``linear_path``): + + 1. an activation its producer already quantized to MXFP8 (``m3_prequant.SLOT``, taken here): the tcgen05 decode + GEMV on the rows it serves, else the decode-graph MXFP8 GEMM on that activation (AutoTuner-selected backend + per token bucket); + 2. a bf16 activation on the rows the GEMV serves quantizing in its prologue; + 3. ``low_m_gemm`` projections, on its rows: the fused low-M MXFP8 GEMM; + 4. otherwise the engine-tuned MXFP8 quantize + GEMM. + + A kernel declining rows its path covers is a miss; the chain then continues down the list. ``push`` (an + ``m3_mx_gemv`` ``ar`` tuple) is the boundary all-reduce exchange the GEMV writes its rows into. + """ + w, s = core.w[wk], core.w[sk] + x = x.contiguous() + mode = _mx_mode(core.mxfp8) + if not mode.decode_graph: + return _mxfp8_linear(core, x, w, s, mode), False + use_pdl = get_env_enable_pdl() + rows = x.shape[0] + n, k = w.shape + prequant = m3_prequant.SLOT.take(x) + if prequant is None and x.dtype != torch.bfloat16: + path = GENERIC_GEMM_PATH + else: + path = linear_path(n, k, rows, prequant is not None, low_m_gemm) + _record("linear", (n, k), rows, path, push is not None) + output = None + if prequant is not None: + act, act_scale = prequant + if path == MX_GEMV_PATH: + output = m3_mx_gemv(act, act_scale, w, s, True, None, use_pdl, push) + if output is not None: + return output, push is not None and mx_gemv.last_pushed + missed( + path, + rows, + f"the GEMV declined the [{n}, {k}] projection's MXFP8 activation" + + _uncompiled_note(), + ) + if path == PREQUANT_GEMM_PATH: + uncovered(MX_GEMV_PATH, f"the [{n}, {k}] projection at rows the GEMV does not serve") + return m3_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16, mode.autotune), False + if path == MX_GEMV_BF16_PATH: + output = m3_mx_gemv_bf16(x, w, s, True, None, use_pdl, push) + if output is not None: + return output, push is not None and mx_gemv.last_pushed + missed( + path, rows, f"the GEMV declined the [{n}, {k}] projection's input" + _uncompiled_note() + ) + if ( + path in (MX_GEMV_BF16_PATH, LOW_M_GEMM_PATH) + and low_m_gemm + and rows <= LOW_M_GEMM_MAX_TOKENS.get((n, k), 0) + and _low_m_gemm_supports(x, w, s) + ): + output = _low_m_gemm(x, w, s, torch.bfloat16, use_pdl=use_pdl) + if output is not None: + return output, False + if path == LOW_M_GEMM_PATH: + missed(path, rows, f"the GEMM declined the [{n}, {k}] projection's input") + if path == GENERIC_GEMM_PATH: + uncovered( + MX_GEMV_BF16_PATH, + f"the [{n}, {k}] projection's bf16 input beyond the decode GEMV's rows", + ) + return _mxfp8_linear(core, x, w, s, mode), False + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer boundaries: all-reduce + residual + RMSNorm +# -------------------------------------------------------------------------------------------------------------------- + + +def _can_mc(x, folded: bool = False) -> bool: + """Whether the generic boundary takes the multicast all-reduce + residual + RMSNorm.""" + return ( + isinstance(x, torch.Tensor) + and x.dim() == 2 + and x.shape[0] <= (m3_mc_allreduce.MAX_TOKENS if folded else MC_AR_MAX_TOKENS) + and x.dtype == torch.bfloat16 + and x.is_contiguous() + ) + + +def _mc_allreduce_norm(core, x, x2, residual, norm_w, *, quant=False, prepushed=0): + """The multicast all-reduce + optional second addend + residual + RMSNorm (``m3_mc_allreduce_norm``), or None + without a workspace or for operands the kernel does not take (the same on every rank). With ``quant`` the normed + rows are also published, MXFP8-quantized, to ``m3_prequant.SLOT`` for the next MXFP8 projection.""" + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + if workspace is None: + return None + _record("m3_mc_allreduce_norm", x.shape[0], x2 is not None, quant, prepushed) + return m3_mc_allreduce.m3_mc_allreduce_norm( + x, + residual.contiguous(), + norm_w, + core.rms_eps, + workspace, + use_pdl=get_env_enable_pdl(), + x2=x2, + quant=quant, + prepushed=prepushed, + ) + + +def _native_allreduce_norm(core, x, residual, norm_w): + """The built-in fused all-reduce + residual + RMSNorm.""" + _record("m3_allreduce", x.shape[0]) + return m3_allreduce( + x, + core.allreduce, + AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=norm_w, + eps=core.rms_eps, + trigger_completion_at_end=False, + ), + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The MoE: router, routed experts (trtllm-gen NVFP4), shared expert +# -------------------------------------------------------------------------------------------------------------------- + + +@dataclass +class _MoeScalars: + """The routed experts' kernel scalars of one layer, derived from the checkpoint's per-tensor scales once the + weights are real (``MiniMaxM3Tp4ModelingV2._derive_moe``).""" + + fc31_input_scale: torch.Tensor + fc2_input_scale: torch.Tensor + fc31_alpha: torch.Tensor + fc2_alpha: torch.Tensor + fc31_scale_c: torch.Tensor + act_alpha: torch.Tensor + act_beta: torch.Tensor + act_clamp: torch.Tensor + + +def _router_logits(core, i: int, h: torch.Tensor) -> torch.Tensor: + """The fp32 router logits of MoE layer ``i``: the fp32 GEMV at one row, the tensor-core router up to 48, else the + fp32 linear.""" + rows = h.shape[0] if h.dim() == 2 else 0 + path = router_path(rows) + w = core.w[f"l{i}_router"] + if path == ROUTER_GEMV.path: + logits = _router_gemv(h, w, use_pdl=get_env_enable_pdl()) + if logits is not None: + return logits + missed(path, rows, "the router GEMV declined the operands") + elif path == ROUTER_TC.path: + fused = _router_tc(h, w, use_pdl=get_env_enable_pdl()) + if fused is not None: + return fused[0] + missed(path, rows, "the tensor-core router declined the operands") + else: + uncovered(ROUTER_TC.path, f"more than {ROUTER_TC_MAX_TOKENS} rows") + return m3_fp32_linear(h, w) + + +def _router_fp4(core, i: int, h: torch.Tensor): + """Router logits and the routed experts' NVFP4 activation from one kernel, or None (then ``_router_logits``).""" + if h.dim() != 2 or h.shape[0] > ROUTER_TC_MAX_TOKENS: + return None + rows = h.shape[0] + path = router_path(rows) + if path == ROUTER_GEMV.path: + path = ROUTER_GEMV_FP4.path + scale = core.moe[i].fc31_input_scale + if path == ROUTER_TC.path: + fused = _router_tc(h, core.w[f"l{i}_router"], scale, use_pdl=get_env_enable_pdl()) + else: + fused = _router_gemv_fp4(h, core.w[f"l{i}_router"], scale, use_pdl=get_env_enable_pdl()) + if fused is None: + missed(path, rows, "the router declined the operands") + return fused + + +def _routed_experts(core, i: int, x, x_sf, router_logits, do_finalize: bool): + """The trtllm-gen NVFP4 MoE runner over this rank's tensor-parallel slice of every expert, routing in-kernel. + + ``x_sf`` None means ``x`` is the bf16 activation, quantized here with the experts' global input scale. Returns + the finalized rows, or the unfinalized ``[gemm2 out, expert scales, expanded->permuted index]``. + """ + moe = core.moe[i] + if x_sf is None: + x, x_sf = fp4_quantize(x, moe.fc31_input_scale, 16, False, False) + rows = x.shape[0] + x_sf = x_sf.view(rows, -1).flatten() + _record("fp4_block_scale_moe_runner", rows, do_finalize) + outs = fp4_block_scale_moe_runner( + router_logits, + core.w[f"l{i}_router_bias"], + x, + x_sf.view(torch.float8_e4m3fn), + core.w[f"l{i}_e_w31"], + core.w[f"l{i}_e_s31"].view(torch.float8_e4m3fn), + None, + moe.act_alpha, + moe.act_beta, + moe.act_clamp, + core.w[f"l{i}_e_w2"], + core.w[f"l{i}_e_s2"].view(torch.float8_e4m3fn), + None, + moe.fc31_scale_c, + moe.fc31_alpha, + moe.fc2_alpha, + core.num_experts, + core.top_k, + None, + None, + core.moe_inter, + 0, + core.num_experts, + core.routed_scaling_factor, + ROUTING_MINIMAX2, + do_finalize, + act_type=ACT_SWIGLU, + topk_weights=None, + topk_ids=None, + output=None, + tune_max_num_tokens=core.max_num_tokens, + use_dp=False, + ) + return outs[0] if do_finalize else outs + + +def _split_chunk(rows: int, num_chunks: int) -> List[int]: + """The MoE scheduler's even split: the first ``rows % num_chunks`` chunks take one more row.""" + base, extra = divmod(rows, num_chunks) + return [base + 1] * extra + [base] * (num_chunks - extra) + + +# -------------------------------------------------------------------------------------------------------------------- +# Attention +# -------------------------------------------------------------------------------------------------------------------- + + +def _nvfp4_layer(manager, i: int) -> bool: + return bool(getattr(manager, "is_nvfp4_layer", lambda _layer_idx: False)(i)) + + +def _producer_covered(core) -> bool: + """Whether ``PRODUCER`` covers this target's sparse layers: an NVFP4 KV cache at the TP4 heads (16 query, one KV + and one index head). A covered layer that cannot take ``m3_producer`` is a miss.""" + return core.main_kv_is_nvfp4 and (core.heads, core.kv_heads, core.index_heads) == (16, 1, 1) + + +def _producer_layout(kv_data, kv_scale) -> bool: + """Whether a sparse layer's NVFP4 pools are the layout ``m3_producer`` is built for: [pages, 2, 1, 128, 64] with + its scale pool [pages, 2, 1, 128, 8].""" + return ( + kv_data is not None + and kv_scale is not None + and kv_data.dim() == 5 + and tuple(kv_data.shape[1:]) == (2, 1, 128, 64) + and kv_data.stride()[2:] == (128 * 64, 64, 1) + and kv_scale.dim() == 5 + and tuple(kv_scale.shape[1:]) == (2, 1, 128, 8) + and kv_scale.stride()[2:] == (128 * 8, 8, 1) + and kv_scale.shape[0] == kv_data.shape[0] + ) + + +def _nvfp4_fused_write_layout(data: torch.Tensor, scales: torch.Tensor, num_heads: int) -> bool: + """Reject invalid HND pools; select the fused producer only for its packed layout.""" + if ( + data.dim() != 5 + or scales.dim() != 5 + or data.shape[0] <= 0 + or data.shape[0] != scales.shape[0] + or tuple(data.shape[1:]) != (2, num_heads, 128, 64) + or tuple(scales.shape[1:]) != (2, num_heads, 128, 8) + or data.dtype not in (torch.int8, torch.uint8) + or scales.dtype != torch.uint8 + ): + raise ValueError("MiniMax-M3 NVFP4 data and scale caches must have matching HND geometry") + for cache in (data, scales): + if ( + cache.stride(4) != 1 + or cache.stride(3) != cache.shape[4] + or cache.stride(2) < 128 * cache.stride(3) + or cache.stride(1) < num_heads * cache.stride(2) + or cache.stride(0) < 2 * cache.stride(1) + ): + raise ValueError( + "MiniMax-M3 NVFP4 cache rows must be packed and heads/planes/pages must not overlap" + ) + if scales.stride(2) != 128 * 8: + raise ValueError("MiniMax-M3 NVFP4 MSA scale cache must have packed head blocks") + return data.stride(2) == 128 * 64 and data.stride(0) % 2 == 0 and data.stride(1) % 2 == 0 + + +def _producer_miss(covered: bool, rows: Optional[int], reason: str) -> None: + if covered: + missed(PRODUCER.path, rows, reason) + return None + + +def _sparse_qkv_producer( + core, i: int, packed, position_ids, kv_cache_manager, msa_idx_k_cache, msa_out_cache_loc +): + """The horizontal producer of a sparse layer: Q / index-Q (FP8) out, K / V / index-K written into the paged + caches (NVFP4 or FP8 main cache), for every step type (token-major). An NVFP4 cache at the TP4 heads takes + ``m3_producer`` (``PRODUCER``); a covered layer that cannot is a miss. Returns (q, index_q), or None for caches + neither kernel takes. + + ``kv_cache_manager``, ``msa_idx_k_cache`` (the per-layer index-K cache getter) and ``msa_out_cache_loc`` are the + calling target's ``step_args`` entries.""" + covered = _producer_covered(core) + rows = int(packed.shape[0]) + if packed.dtype != torch.bfloat16 or position_ids is None: + return _producer_miss( + covered, + rows, + "the packed QKV dtype, position ids, head dims, norm or RoPE configuration", + ) + manager = kv_cache_manager + if manager is None: + return _producer_miss(covered, rows, "the step has no KV cache manager") + buffers = manager.get_buffers(i, kv_layout="HND") + index_k_cache = msa_idx_k_cache(i) + out_cache_loc = msa_out_cache_loc + layer_uses_nvfp4 = _nvfp4_layer(manager, i) + scale_buffers = ( + manager.get_block_scale_buffers(i, kv_layout="HND") if layer_uses_nvfp4 else None + ) + kv_quant_scale = core.inv_kv_scales[i] if layer_uses_nvfp4 else None + if layer_uses_nvfp4: + supported_main = ( + buffers is not None + and scale_buffers is not None + and _nvfp4_fused_write_layout(buffers, scale_buffers, core.kv_heads) + and buffers.is_cuda + and scale_buffers.is_cuda + and scale_buffers.device == buffers.device + and kv_quant_scale.is_cuda + and kv_quant_scale.dtype == torch.float32 + and kv_quant_scale.numel() >= 3 + and kv_quant_scale.is_contiguous() + ) + else: + supported_main = ( + buffers is not None + and buffers.is_cuda + and buffers.dtype == torch.float8_e4m3fn + and buffers.dim() == 5 + and tuple(buffers.shape[1:]) == (2, core.kv_heads, 128, 128) + and buffers.stride(4) == 1 + and buffers.stride(3) == 128 + and buffers.stride(2) == 128 * 128 + and buffers.stride(1) >= core.kv_heads * buffers.stride(2) + and buffers.stride(0) >= 2 * buffers.stride(1) + and buffers.stride(0) % 4 == 0 + and buffers.stride(1) % 4 == 0 + ) + supported_index = ( + index_k_cache.is_cuda + and index_k_cache.dtype == torch.float8_e4m3fn + and index_k_cache.dim() == 4 + and tuple(index_k_cache.shape[1:]) == (1, 128, 128) + and index_k_cache.stride(3) == 1 + and index_k_cache.stride(2) == 128 + and index_k_cache.stride(1) >= 128 * 128 + and index_k_cache.stride(0) >= index_k_cache.stride(1) + and index_k_cache.stride(0) % 4 == 0 + and index_k_cache.stride(1) % 4 == 0 + and buffers is not None + and index_k_cache.shape[0] == buffers.shape[0] + ) + cos_sin = core.rotary_cos_sin + supported_rope = ( + cos_sin.is_cuda + and cos_sin.dtype == torch.float32 + and cos_sin.is_contiguous() + and cos_sin.dim() == 3 + and tuple(cos_sin.shape[1:]) == (2, 32) + ) + if ( + not supported_main + or not supported_index + or not supported_rope + or out_cache_loc is None + or not out_cache_loc.is_cuda + or out_cache_loc.dtype != torch.int32 + or not out_cache_loc.is_contiguous() + or out_cache_loc.numel() < rows + ): + return _producer_miss( + covered and layer_uses_nvfp4, + rows, + "the layer's NVFP4 / index caches, RoPE table, KV quant scale or cache slots are not the layout the " + "producer is built for", + ) + w = core.w + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + if ( + layer_uses_nvfp4 + and (core.heads, core.kv_heads, core.index_heads) == (16, 1, 1) + and _producer_layout(buffers, scale_buffers) + ): + _record("m3_producer", rows) + q, index_q = m3_producer( + packed.contiguous(), + buffers.view(torch.uint8), + scale_buffers.view(torch.uint8), + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + core.heads, + get_env_enable_pdl(), + PRODUCER_EARLY_TRIGGER, + ) + else: + if layer_uses_nvfp4 and covered: + missed( + PRODUCER.path, + rows, + f"the NVFP4 pool {tuple(buffers.shape)} / scale pool {tuple(scale_buffers.shape)} is not the " + "producer's layout", + ) + _record("m3_qkv_insert", rows, layer_uses_nvfp4) + q, index_q = m3_qkv_insert( + packed.contiguous(), + buffers.view(torch.uint8) if layer_uses_nvfp4 else buffers, + scale_buffers, + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.heads, + core.kv_heads, + core.index_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + ) + return q.flatten(1), index_q.flatten(1) + + +def _qk_norm_rope(core, i: int, qkv: torch.Tensor, position_ids, out_fp8: bool) -> torch.Tensor: + """The per-head Gemma RMSNorm of Q and K and the partial NeoX RoPE of a dense layer's packed [Q | K | V], in one + kernel: in place, or an FP8 copy (Q / K normed and roped, V cast) when ``out_fp8``.""" + qkv = qkv.contiguous() + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + args = ( + qkv, + core.heads, + core.kv_heads, + core.kv_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_q_norm"], + core.w[f"l{i}_k_norm"], + core.rope_theta, + True, + pos, + 1.0, # factor: no YARN + 0.0, # low + 0.0, # high + 1.0, # attention_factor + True, # is_qk_norm + True, # use_gemma + False, # use_mrope + 0, + 0, + ) + _record("qk_norm_rope", qkv.shape[0], out_fp8) + if out_fp8: + return m3_qk_norm_rope_fp8(*args) + return fused_qk_norm_rope(*args) + + +def _attn_mega(core, i: int, hidden: torch.Tensor, args: dict, push: Optional[tuple]): + """(o_proj output, pushed) of the one-launch QKV + attention + o_proj kernel for one request decoding one token + whose layer input the preceding all-reduce handed over in MXFP8; None, with nothing consumed, otherwise.""" + if not ATTN_MEGA.covers((args["decode_query_len"], int(args["seq_lens"].shape[0]))): + return None + if not flashinfer_mxfp8_decode_graph_capture_active(): + return None + pending = m3_prequant.SLOT.pending + if ( + pending is None + or hidden.dim() != 2 + or pending[0] != hidden.data_ptr() + or pending[1] != 1 + or pending[2] != hidden.shape[1] + ): + # After a layer whose boundary does not quantize (a dense layer, an Eagle3 capture layer). + uncovered(ATTN_MEGA.path, "the layer input has no MXFP8 hand-off") + return None + w = core.w + if w[f"l{i}_qkv_w"].shape != (2560, 6144) or w[f"l{i}_o_w"].shape != (6144, 2048): + uncovered(ATTN_MEGA.path, "not the tp4 QKV / o_proj shapes") + return None + workspace = core.mega_ws.get(i) + if workspace is None: + # The first call JIT-builds the module; both only outside capture (the eager warmup steps). + if torch.cuda.is_current_stream_capturing(): + capture_missed(ATTN_MEGA.path, 1, "the attn_mega module and workspace") + return None + _attn_mega_kernel._load_module() + workspace = core.mega_ws[i] = _attn_mega_kernel.MegaWorkspace(hidden.device) + xq, xsf = m3_prequant.SLOT.take(hidden) + output = hidden.new_empty((1, core.heads * core.head_dim), dtype=torch.bfloat16) + y = hidden.new_empty((1, w[f"l{i}_o_w"].shape[0]), dtype=torch.bfloat16) + _record("m3_attn_mega", push is not None) + _attn_mega_kernel.m3_attn_mega( + xq, + xsf, + w[f"l{i}_qkv_w"], + w[f"l{i}_qkv_s"], + workspace, + args["kv_data"], + args["kv_scale"], + args["index_k_cache"], + args["out_cache_loc"], + args["position_ids"], + args["block_table"], + args["seq_lens"], + args["kv_quant_scale"], + args["k_global_scale"], + args["v_global_scale"], + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_k_norm"], + core.rotary_cos_sin, + w[f"l{i}_o_w"], + w[f"l{i}_o_s"], + eps=core.rms_eps, + sm_scale=core.sm_scale(i), + decode_query_len=1, + output=output, + y=y, + use_pdl=get_env_enable_pdl(), + ar=push, + ) + return y, push is not None + + +def _attn_c1(core, i: int, packed: torch.Tensor, args: dict, push: Optional[tuple]): + """(o_proj output, True) of the one-launch attention + o_proj kernel for one request verifying a draft, pushed + into the attention boundary's all-reduce; None, with nothing consumed, when the step or the layer does not + qualify.""" + qlen, num_tokens = args["decode_query_len"], args["num_tokens"] + if not ATTN_C1.covers((qlen, int(args["seq_lens"].shape[0]))): + return None + if not flashinfer_mxfp8_decode_graph_capture_active(): + return None + if push is None: + missed(ATTN_C1.path, num_tokens, "no armed attention-boundary push to write into") + return None + w = core.w + if w[f"l{i}_o_w"].shape != (6144, 2048): + uncovered(ATTN_C1.path, "not the tp4 o_proj shape") + return None + workspace = core.c1_ws.get(i) + if workspace is None: + # The first call JIT-builds the module; both only outside capture (the eager warmup steps). + if torch.cuda.is_current_stream_capturing(): + capture_missed(ATTN_C1.path, num_tokens, "the attn_c1 module and workspace") + return None + _attn_c1_kernel._load_module() + workspace = core.c1_ws[i] = _attn_c1_kernel.MegaWorkspace(packed.device) + output = packed.new_empty((num_tokens, core.heads * core.head_dim), dtype=torch.bfloat16) + y = packed.new_empty((num_tokens, w[f"l{i}_o_w"].shape[0]), dtype=torch.bfloat16) + _record("m3_attn_c1", num_tokens) + _attn_c1_kernel.m3_attn_c1( + packed.contiguous(), + workspace, + args["kv_data"], + args["kv_scale"], + args["index_k_cache"], + args["out_cache_loc"], + args["position_ids"], + args["block_table"], + args["seq_lens"], + args["kv_quant_scale"], + args["k_global_scale"], + args["v_global_scale"], + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_k_norm"], + core.rotary_cos_sin, + w[f"l{i}_o_w"], + w[f"l{i}_o_s"], + eps=core.rms_eps, + sm_scale=core.sm_scale(i), + decode_query_len=qlen, + output=output, + y=y, + use_pdl=get_env_enable_pdl(), + ar=push[:5], + ) + return y, True + + +def _fused_attn(core, i: int, packed: torch.Tensor, args: dict) -> torch.Tensor: + """The fused norm/RoPE/KV-insert/sparse-decode/merge kernel after the packed QKV projection; its MXFP8 copy of + the output goes to the o_proj through ``m3_prequant.SLOT``.""" + num_tokens = args["num_tokens"] + width = core.heads * core.head_dim + output = packed.new_empty((num_tokens, width), dtype=torch.bfloat16) + mxfp8 = None + if MERGE_MXFP8_MIN_TOKENS <= num_tokens <= MERGE_MXFP8.MAX_ROWS: + mxfp8 = MERGE_MXFP8.buffers(width, output.device) + w = core.w + _record("m3_fused_decode_attn", num_tokens, args["decode_query_len"]) + _fused_attn_kernel.m3_fused_decode_attn( + packed.contiguous(), + args["kv_data"], + args["kv_scale"], + args["index_k_cache"], + args["out_cache_loc"], + args["position_ids"], + args["block_table"], + args["seq_lens"], + args["kv_quant_scale"], + args["k_global_scale"], + args["v_global_scale"], + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_k_norm"], + core.rotary_cos_sin, + eps=core.rms_eps, + sm_scale=core.sm_scale(i), + decode_query_len=args["decode_query_len"], + output=output, + mxfp8_out=mxfp8, + use_pdl=get_env_enable_pdl(), + ) + if mxfp8 is not None: + m3_prequant.SLOT.publish(output, mxfp8[0][:num_tokens], mxfp8[1]) + return output + + +def _fp8_index_qk_norm_rope( + core, i: int, idx_qk: torch.Tensor, position_ids, msa_idx_k_cache, msa_out_cache_loc +) -> torch.Tensor: + """The index branch of an unfused sparse layer: index-Q in FP8 out, index-K into the FP8 index-K cache, normed and + roped in one kernel (the built-in ``_fused_fp8_index_qk_norm_rope``). ``msa_idx_k_cache`` / ``msa_out_cache_loc`` + are the calling target's ``step_args`` entries.""" + if position_ids is None or idx_qk.dtype != torch.bfloat16: + raise NotImplementedError( + "MiniMax-M3 fused FP8 indexer requires BF16 activations and position_ids." + ) + index_k_cache = msa_idx_k_cache(i) + if index_k_cache.dtype != torch.float8_e4m3fn: + raise ValueError( + "MiniMax-M3 fused FP8 indexer requires an E4M3 index-K cache, " + f"got {index_k_cache.dtype}." + ) + rows = int(idx_qk.shape[0]) + _record("m3_fp8_indexer_qk_norm_rope", rows) + return m3_fp8_indexer_qk_norm_rope( + idx_qk.contiguous(), + index_k_cache, + msa_out_cache_loc[:rows], + core.index_heads, + core.index_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_idx_q_norm"], + core.w[f"l{i}_idx_k_norm"], + core.rope_theta, + position_ids.reshape(-1).contiguous().to(torch.int32), + ).flatten(1) + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer helpers both targets' bodies call: the Eagle3 capture hook and the norm after a layer +# -------------------------------------------------------------------------------------------------------------------- + + +def _plain_capture(spec_metadata, layer_id: int, hidden_states, residual) -> None: + """The in-tree Eagle3 hidden-state capture.""" + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + + +def _next_norm(core, i: int) -> torch.Tensor: + """The norm the boundary after layer ``i`` applies: the next layer's input norm, the final norm after the last.""" + return core.w[f"l{i + 1}_in_norm"] if i + 1 < core.num_layers else core.w["final_norm"] + + +# -------------------------------------------------------------------------------------------------------------------- +# The core +# -------------------------------------------------------------------------------------------------------------------- + +# Engine states (compiled kernel + workspace) per (device, token count[, shared fold]), shared by every layer: a step +# runs the layers one after another on one stream. +_ENGINE_STATES: Dict[tuple, object] = {} +_MOE16_STATES: Dict[tuple, object] = {} + + +class MiniMaxM3Tp4ModelingV2(DecoderModel): + """The MiniMax-M3 text model of this target: its weights, its runtime resources, and the per-step dispatch.""" + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + _validate_runtime(model_config) + super().__init__(model_config) + cfg = model_config.pretrained_config + mapping = model_config.mapping + self.mapping = mapping + self.vocab_size = cfg.vocab_size + d = _weights.MODEL_WEIGHTS.dims(self) + self.num_layers = d.num_layers + self.hidden = d.hidden + self.heads = d.heads + self.kv_heads = d.kv_heads + self.head_dim = d.head_dim + self.index_heads = d.index_heads + self.index_dim = d.index_dim + self.q_size = d.heads * d.head_dim + self.kv_size = d.kv_heads * d.head_dim + self.moe_inter = d.moe_inter + self.num_experts = d.num_experts + self.top_k = int(cfg.num_experts_per_tok) + self.routed_scaling_factor = float(getattr(cfg, "routed_scaling_factor", 1.0)) + self.rms_eps = float(cfg.rms_norm_eps) + self.swiglu_alpha = float(getattr(cfg, "swiglu_alpha", 1.702)) + self.swiglu_beta = 1.0 # the (up + 1) offset of SwiGLU-OAI + self.swiglu_limit = float(getattr(cfg, "swiglu_limit", 7.0)) + self.max_num_tokens = model_config.max_num_tokens + self.moe_max_num_tokens = model_config.moe_max_num_tokens + self.sparse_layers = frozenset(_sparse_layer_ids(cfg)) + self.moe_layers = frozenset(_moe_layer_ids(cfg)) + mode = model_config.quant_config.quant_mode + self.main_kv_is_fp8 = bool(mode is not None and mode.has_fp8_kv_cache()) + self.main_kv_is_nvfp4 = bool(mode is not None and mode.has_fp4_kv_cache()) + # The sparse layers' QKV and index projections: one five-way packed projection, or the QKV projection and a + # separate BF16 index projection (fuse_qkv_index_projection=False) run on two streams. + self.fused_projection = bool( + getattr(model_config.sparse_attention_config, "fuse_qkv_index_projection", False) + ) + + self.w = _weights.MODEL_WEIGHTS.declare(d) + + # The engine reads these two off the module tree: the MXFP8 GEMM state it tunes, and the opt-in to the + # automatic FlashInfer MXFP8 path on decode graphs. + self.mxfp8 = _Mxfp8GemmState() + self._use_flashinfer_mxfp8_decode_graph_default = True + + # The built-in fused all-reduce + residual + RMSNorm of the layer boundaries the multicast kernel does not take. + self.allreduce = AllReduce( + mapping=mapping, strategy=model_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + + # Streams: the shared / routed expert overlap and, when the MoE chunks its input, the chunking overlap. + self.aux_stream_moe = torch.cuda.Stream() + chunked = self.moe_max_num_tokens < model_config.max_num_tokens * mapping.dp_size + self.aux_stream_moe_chunk = torch.cuda.Stream() if chunked else None + self.moe_events = {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers} + self.moe_chunk_events = { + i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers + } + # The unfused sparse layers' index branch overlaps the QKV branch on its own stream. + self.aux_stream_attn = None if self.fused_projection else torch.cuda.Stream() + self.attn_events = ( + {} + if self.fused_projection + else {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.sparse_layers} + ) + + # Partial NeoX RoPE (rotary_dim of head_dim): the cos/sin table every attention kernel reads. + rope = RopeParams.from_config(cfg) + self.rope_theta = rope.theta + self.rotary_dim = int(rope.dim) + self.rotary_cos_sin = rope.create_rope_const_params(interleave=False)[1].reshape( + rope.max_positions, 2, -1 + ) + + # One MSA attention backend per layer (dense and sparse; the step metadata is the backend's). + quant_config = model_config.get_quant_config() + self.attn: List[MiniMaxM3MsaSparseAttention] = [] + for i in range(self.num_layers): + backend = create_attention( + model_config.attn_backend, + i, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=None, + quant_config=quant_config, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + q_scaling=1.0, + sparse_params=model_config.sparse_attention_config.to_sparse_params( + pretrained_config=cfg, layer_idx=i + ), + ) + backend.update_quant_config(quant_config) + if not isinstance(backend, MiniMaxM3MsaSparseAttention): + raise ValueError( + f"layer {i}: expected the MiniMax-M3 MSA backend, got {type(backend).__name__}" + ) + self.attn.append(backend) + + # Built once the weights are real (``post_load_weights``): the KV (Q, K, V) scales, the MoE scalars, the + # decode engines and workspaces, the two targets. + self.kv_scales: Dict[int, torch.Tensor] = {} + self.inv_kv_scales: Dict[int, torch.Tensor] = {} + self.moe: Dict[int, _MoeScalars] = {} + self.mx_alpha: Optional[torch.Tensor] = None + self.moe_engines: Dict[int, Dict[Tuple[int, bool, bool], Optional[object]]] = {} + self.moe16_engines: Dict[int, Dict[int, Optional[object]]] = {} + self.embed_owners: Dict[tuple, DecodeEmbedding] = {} + self.mega_ws: Dict[int, object] = {} + self.c1_ws: Dict[int, object] = {} + self._built_for: Optional[tuple] = None + self._prefill: Optional[PrefillTarget] = None + self._decode: Optional[DecodeTarget] = None + + def sm_scale(self, i: int) -> float: + return (self.head_dim**-0.5) / float(self.attn[i].q_scaling) + + # ---------------------------------------------------------------------------------------------- after the load + + def post_load_weights(self) -> None: + """Derive the kernels' state from the loaded weights, then build both targets and every decode kernel. + + Idempotent per set of weight storages: the loader's post-load walk reaches this module directly and through + the shell, and a second call on the same storages does nothing. + """ + identity = ( + self.w["embed"].data_ptr(), + self.w["l0_in_norm"].data_ptr(), + self.w["final_norm"].data_ptr(), + ) + if self._built_for == identity: + return + device = self.w["embed"].device + derived = nn.Module() + derived.register_buffer( + "mx_alpha", torch.ones(1, dtype=torch.float32, device=device), persistent=False + ) + for i in range(self.num_layers): + # No KV calibration in the checkpoint: unit (Q, K, V) scales and their reciprocals. + derived.register_buffer( + f"l{i}_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + derived.register_buffer( + f"l{i}_inv_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + for i in sorted(self.moe_layers): + for name, value in dataclasses.asdict(self._derive_moe(i)).items(): + derived.register_buffer(f"l{i}_{name}", value, persistent=False) + self.derived = derived + self.mx_alpha = derived.mx_alpha + self.kv_scales = {i: getattr(derived, f"l{i}_kv_scales") for i in range(self.num_layers)} + self.inv_kv_scales = { + i: getattr(derived, f"l{i}_inv_kv_scales") for i in range(self.num_layers) + } + self.moe = { + i: _MoeScalars( + **{ + f.name: getattr(derived, f"l{i}_{f.name}") + for f in dataclasses.fields(_MoeScalars) + } + ) + for i in sorted(self.moe_layers) + } + self._build_embed_owners() + self.prebuild_kernels() + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + self._built_for = identity + + def cache_derived_state(self) -> None: + """The staged loaders' entry to the same build (``post_load_weights``).""" + self.post_load_weights() + + def _derive_moe(self, i: int) -> _MoeScalars: + """The routed experts' kernel scalars from the checkpoint's per-tensor scales (trtllm-gen NVFP4 layout). + + Global input scales: the reciprocal of the largest per-expert input scale (gate and up must agree). Per expert: + ``alpha = 1 / (input_scale * (1 / weight_scale_2))`` for the gate (fc31_alpha), the up half (its GEMM2-input + rescale ``fc31_scale_c = fc2_input_scale * up_alpha``) and the down projection; SwiGLU-OAI's activation + constants per expert, with beta and the clamp in the pre-dequant GEMM output domain (divided by fc31_alpha). + """ + w = self.w + isc1, isc3, isc2 = w[f"l{i}_e_isc1"].data, w[f"l{i}_e_isc3"].data, w[f"l{i}_e_isc2"].data + if not torch.allclose(isc1, isc3): + raise ValueError(f"layer {i}: the routed experts' w1 and w3 input scales differ") + fc31_input_scale = isc1.max().reciprocal() + fc2_input_scale = isc2.max().reciprocal() + ws2_1, ws2_3, ws2_2 = ( + w[f"l{i}_e_ws2_1"].data, + w[f"l{i}_e_ws2_3"].data, + w[f"l{i}_e_ws2_2"].data, + ) + fc31_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_1)) + up_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_3)) + fc2_alpha = 1.0 / (fc2_input_scale * (1.0 / ws2_2)) + fc31_scale_c = fc2_input_scale * up_alpha + n = self.num_experts + device = fc31_alpha.device + act_alpha = torch.full((n,), self.swiglu_alpha, dtype=torch.float32, device=device) + act_beta = torch.full((n,), self.swiglu_beta, dtype=torch.float32, device=device) + act_clamp = torch.full((n,), self.swiglu_limit, dtype=torch.float32, device=device) + act_beta.div_(fc31_alpha) + act_clamp.div_(fc31_alpha) + return _MoeScalars( + fc31_input_scale=fc31_input_scale, + fc2_input_scale=fc2_input_scale, + fc31_alpha=fc31_alpha, + fc2_alpha=fc2_alpha, + fc31_scale_c=fc31_scale_c, + act_alpha=act_alpha, + act_beta=act_beta, + act_clamp=act_clamp, + ) + + def _build_embed_owners(self) -> None: + """The fused embedding + first RMSNorm's caller-owned buffers, one per (rows, id dtype) of a decode step.""" + table, weight = self.w["embed"], self.w["l0_in_norm"] + self.embed_owners = { + (rows, dtype): DecodeEmbedding(table, weight, self.rms_eps, rows, dtype) + for rows in range(1, EMBED_FIRSTNORM_MAX_TOKENS + 1) + for dtype in (torch.int32, torch.int64) + } + + def projections(self) -> List[Tuple[str, str, Optional[str]]]: + """Every MXFP8 projection in layer order: (weight key, scale key, the boundary its GEMV pushes into).""" + out = [] + for i in range(self.num_layers): + out.append((f"l{i}_qkv_w", f"l{i}_qkv_s", None)) + out.append((f"l{i}_o_w", f"l{i}_o_s", "attn")) + if i in self.moe_layers: + out.append((f"l{i}_sh_gu_w", f"l{i}_sh_gu_s", None)) + out.append((f"l{i}_sh_dn_w", f"l{i}_sh_dn_s", "moe")) + else: + out.append((f"l{i}_mlp_gu_w", f"l{i}_mlp_gu_s", None)) + out.append((f"l{i}_mlp_dn_w", f"l{i}_mlp_dn_s", None)) + return out + + def decode_query_lens(self) -> List[int]: + """The query rows per request a decode step can have: 1, and 1 + each draft length speculation runs.""" + lens = {1} + spec = getattr(self.model_config, "spec_config", None) + if spec is not None: + for drafts in ( + getattr(spec, "max_draft_len", None), + getattr(spec, "max_total_draft_tokens", None), + *(getattr(spec, "draft_len_schedule", None) or {}).values(), + ): + if drafts: + lens.add(1 + int(drafts)) + return sorted(lens) + + def prebuild_kernels(self) -> Dict[str, object]: + """Compile (and load) every kernel a decode forward dispatches to that would otherwise compile on its first + call: each projection's decode GEMVs (every row count, activation and boundary push), the MoE engines of + ``MOE_ENGINE`` / ``MOE16_ENGINE``'s token counts, and the index scorer at the decode query lengths. Runs on the + current (main) stream before the executor's warmup, in the same order on every rank, then synchronizes; after + it, GEMV compiles are frozen and the forward only looks kernels up (a missing one is a miss). Returns the + seconds each part took and the scorer keys.""" + times: Dict[str, object] = {} + if not self.mxfp8.quant_method.use_cutlass: + raise ValueError("the MiniMax-M3 decode dispatch needs the CUTLASS MXFP8 GEMM") + mx_gemv.freeze(False) + use_pdl = get_env_enable_pdl() + world = int(self.mapping.tp_size) + start = time.perf_counter() + for wk, sk, push in self.projections(): + weight, scale = self.w[wk], self.w[sk] + mx_gemv.ensure_compiled(weight, scale, use_pdl) + n, k = weight.shape + for rows, x_bf16, ar in _gemv_variants(n, k, push, world): + mx_gemv.precompile(weight, scale, rows, x_bf16, use_pdl=use_pdl, ar=ar) + torch.cuda.synchronize() + times["gemv_s"] = time.perf_counter() - start + start = time.perf_counter() + for i in sorted(self.moe_layers): + self._prebuild_engines(i) + torch.cuda.synchronize() + times["moe_s"] = time.perf_counter() - start + start = time.perf_counter() + times["scorer_keys"] = self._prebuild_index_scorer( + self.decode_query_lens(), self.w["embed"].device + ) + torch.cuda.synchronize() + times["scorer_s"] = time.perf_counter() - start + mx_gemv.freeze(True) + logger.info( + "[M3 prebuild] decode kernels built before warmup: GEMVs %.1f s, MoE engines %.1f s, index scorer " + "%.1f s (%d variants)" + % (times["gemv_s"], times["moe_s"], times["scorer_s"], len(times["scorer_keys"])) + ) + return times + + # ------------------------------------------------------------------------------------------ the MoE engines + + def _expert_operands(self, i: int): + """The routed experts' operands the engines read (trtllm-gen NVFP4 SwiGLU-OAI layout), and the input scale.""" + w, m = self.w, self.moe[i] + tensors = [ + w[f"l{i}_e_w31"].data, + w[f"l{i}_e_s31"].data, + w[f"l{i}_e_w2"].data, + w[f"l{i}_e_s2"].data, + m.fc31_alpha, + m.fc31_scale_c, + m.act_alpha, + m.act_beta, + m.act_clamp, + m.fc2_alpha, + ] + return tensors, m.fc31_input_scale + + def _shared_weights(self, i: int): + w = self.w + return ( + w[f"l{i}_sh_gu_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_gu_s"].data, + w[f"l{i}_sh_dn_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_dn_s"].data, + ) + + def engine(self, i: int, num_tokens: int, shared: bool) -> Optional[object]: + """Layer ``i``'s prebuilt routed-expert engine for ``num_tokens`` (1..8), router folded in (and the shared + expert with ``shared``). None after a miss: the build failed or was never made (a forward never builds one).""" + return self._lookup( + i, self.moe_engines[i], (num_tokens, True, shared), MOE_ENGINE.path, num_tokens + ) + + def moe16_engine(self, i: int, num_tokens: int) -> Optional[object]: + """Layer ``i``'s prebuilt 9..13-row engine (router, routed and shared experts in one kernel), or None.""" + if not MOE16_MIN_TOKENS <= num_tokens <= MOE16_MAX_TOKENS: + return None + return self._lookup(i, self.moe16_engines[i], num_tokens, MOE16_ENGINE.path, num_tokens) + + def _lookup(self, i: int, engines: dict, key, path: str, num_tokens: int) -> Optional[object]: + if key not in engines: + what = f"layer {i}'s engine" + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, what) + else: + missed(path, num_tokens, f"{what} was not built before warmup") + return None + engine = engines[key] + if engine is None: + missed(path, num_tokens, "its build failed (logged)") + return engine + + def _prebuild_engines(self, i: int) -> None: + """Build (compile, on first use of each shared state, and warm) every engine a decode step of layer ``i`` can + take: ``MOE_ENGINE``'s token counts with the shared expert folded in, and ``MOE16_ENGINE``'s. Whether an + engine is kept is agreed across the TP group (one host collective per engine): a build that fails on one rank + drops that engine on every rank, so the ranks never mix an engine (and its boundary push) with the generic + MoE.""" + engines = self.moe_engines.setdefault(i, {}) + engines16 = self.moe16_engines.setdefault(i, {}) + for num_tokens in sorted(MOE_ENGINE.keys): + key = (num_tokens, True, True) + if key not in engines: + engines[key] = m3_mc_allreduce.tp_agree_built( + self.mapping, + self._try_build( + i, MOE_ENGINE.path, num_tokens, self._build_engine, shared=True + ), + f"[M3 dispatch] {MOE_ENGINE.path} at {num_tokens} rows: layer {i}'s engine", + ) + for num_tokens in sorted(MOE16_ENGINE.keys): + if num_tokens not in engines16: + engines16[num_tokens] = m3_mc_allreduce.tp_agree_built( + self.mapping, + self._try_build(i, MOE16_ENGINE.path, num_tokens, self._build_moe16_engine), + f"[M3 dispatch] {MOE16_ENGINE.path} at {num_tokens} rows: layer {i}'s engine", + ) + + def _try_build(self, i: int, path: str, num_tokens: int, build, **kwargs) -> Optional[object]: + try: + return build(i, num_tokens, **kwargs) + except (ImportError, RuntimeError, ValueError) as err: + logger.warning( + f"[M3 dispatch] {path} at {num_tokens} rows: layer {i}'s engine failed to build: {err}" + ) + return None + + def _build_engine(self, i: int, num_tokens: int, shared: bool) -> object: + tensors, scale = self._expert_operands(i) + device = tensors[0].device + sms = torch.cuda.get_device_properties(device).multi_processor_count + kwargs = {} + if shared: + kwargs["shared_weights"] = self._shared_weights(i) + state = _ENGINE_STATES.get((str(device), num_tokens, shared)) + if state is None: + state = _ENGINE_STATES[(str(device), num_tokens, shared)] = M3MoeState( + device, num_tokens=num_tokens, num_ctas=sms - MOE_ENGINE_RESERVED_SMS, shared=shared + ) + engine = state.layer( + *tensors, + self.w[f"l{i}_router_bias"].data, + router_weight=self.w[f"l{i}_router"].data, + fc31_input_scale=scale.reshape(1).float(), + **kwargs, + ) + # Compile (first use of the state) and warm. + m3_moe_engine( + torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device), + engine, + ) + return engine + + def _build_moe16_engine(self, i: int, num_tokens: int) -> object: + tensors, scale = self._expert_operands(i) + device = tensors[0].device + state = _MOE16_STATES.get((str(device), num_tokens)) + if state is None: + state = _MOE16_STATES[(str(device), num_tokens)] = M3Moe16State( + device, num_tokens=num_tokens + ) + engine = state.layer( + *tensors, + self.w[f"l{i}_router_bias"].data, + router_weight=self.w[f"l{i}_router"].data, + fc31_input_scale=scale.reshape(1).float(), + shared_weights=self._shared_weights(i), + ) + # Compile (first use of the state) and warm. + m3_moe16_engine( + torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device), + engine, + ) + return engine + + # ------------------------------------------------------------------------------------------ the index scorer + + def _scorer_variants(self, query_lens) -> list: + """The scorer keys a decode step can take: every sparse layer's (index dtype, index heads per rank, index dim) + at each decode query length, with the early-issue schedule and without it, where the kernel serves it.""" + runner = _index_score_kernel.IndexScoreRunner + keys = [] + for i in sorted(self.sparse_layers): + backend = self.attn[i] + config = backend.m3_config + dtype = torch.float8_e4m3fn if backend.indexer_kv_dtype == "fp8" else torch.bfloat16 + for qlen in sorted(set(query_lens)): + if not runner.is_supported( + q_dtype=dtype, + num_heads=config.num_index_heads, + head_dim=config.sparse_index_dim, + page_size=runner.SUPPORTED_PAGE_SIZE, + max_decode_query_len=qlen, + ): + continue + for early in (True, False): + key = ( + dtype, + int(config.num_index_heads), + int(qlen), + int(config.sparse_index_dim), + int(INDEX_SCORE_SPLIT_K), + bool(early), + ) + if key not in keys: + keys.append(key) + return keys + + def _prebuild_index_scorer(self, query_lens, device) -> list: + """Compile every scorer key and launch each once on real-shaped dummy operands (loading its module), on the + current stream. The kernel's batch, token, page and block counts are symbolic: its compile needs no step + metadata. Returns the keys.""" + runner = _index_score_kernel.IndexScoreRunner + keys = self._scorer_variants(query_lens) + page = runner.SUPPORTED_PAGE_SIZE + for key in keys: + dtype, heads, qlen, dim, split_k, early = key + runner._compile(dtype, heads, qlen, dim, split_k, early) + rows = 1 if early else INDEX_SCORE_EARLY_MAX_REQUESTS + 1 + blocks = 2 + tokens = rows * qlen + idx_q = torch.zeros(tokens, heads, dim, dtype=torch.bfloat16, device=device).to(dtype) + idx_k = torch.zeros( + rows * blocks, 1, page, dim, dtype=torch.bfloat16, device=device + ).to(dtype) + block_table = torch.arange(rows * blocks, dtype=torch.int32, device=device).view( + rows, blocks + ) + seq_lens = torch.full((rows,), blocks * page, dtype=torch.int32, device=device) + score = torch.empty(heads, blocks, tokens, dtype=torch.float32, device=device) + runner.kernel_cache[key]( + idx_q, idx_k.squeeze(1), block_table, score.transpose(1, 2), seq_lens + ) + return keys + + # --------------------------------------------------------------------------------------------- the dispatch + + def _select_target(self, attn_metadata: AttentionMetadata): + """Which target runs this step: one host-side integer, fixed per CUDA-graph capture. A step with context rows + (a mixed step included) is prefill's; every pure generation step is decode's, whatever its width.""" + return self._decode if attn_metadata.num_contexts == 0 else self._prefill + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + capture: Callable[..., None] = _plain_capture, + **kwargs, + ) -> torch.Tensor: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + if kwargs.get("lora_params") is not None: + raise ValueError("this MiniMax-M3 target serves no LoRA adapters") + return self._select_target(attn_metadata).forward( + attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, capture + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The targets +# -------------------------------------------------------------------------------------------------------------------- + + +class PrefillTarget: + """A step with context rows: prefill chunks, and the mixed steps that carry generation rows beside them. Every + layer runs in its generic form; the attention core is the MSA backend's for every row.""" + + def __init__(self, core: MiniMaxM3Tp4ModelingV2) -> None: + self.core = core + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + Counts and per-step host scalars are read here; device tensors are the metadata's persistent buffers, held + by reference; ``msa_idx_k_cache`` is the metadata's per-layer index-K cache getter. ``metadata`` is the + prepared object itself: the MSA backend's step (``m3_msa_attention``) takes it as its argument. + """ + return dict( + num_contexts=md.num_contexts, + num_ctx_tokens=md.num_ctx_tokens, + num_tokens=md.num_tokens, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + metadata=md, + ) + + def forward( + self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, capture + ): + core = self.core + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target does not run under piecewise (breakable) prefill CUDA graphs" + ) + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + _record("prefill", int(step["num_tokens"]), int(step["num_contexts"])) + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + hidden, residual = self._layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + return hidden + + def _swiglu_mlp(self, h: torch.Tensor, prefix: str, low_m_gemm: bool) -> torch.Tensor: + """A gated MXFP8 MLP (the dense MLP, the shared expert): gate_up, SwiGLU-OAI, down, no reduction.""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s", low_m_gemm) + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + out, _ = _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s", low_m_gemm) + return out + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: + the MoE scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor): + """Router + routed experts, finalized: the fused router with the NVFP4 activation up to 48 rows, else the + router logits and the experts quantizing their bf16 input (chunked).""" + core = self.core + fused = _router_fp4(core, i, h) + if fused is not None: + router_logits, x_fp4, x_sf = fused + return _routed_experts(core, i, x_fp4, x_sf, router_logits, True) + router_logits = _router_logits(core, i, h) + return self._routed_chunks(i, h, router_logits) + + def _moe(self, i: int, h: torch.Tensor, *, return_parts: bool): + """The MoE block of layer ``i`` without its reduction. Returns ``(out, routed)``: with ``return_parts`` the + shared and routed parts (routed None when the engine folded the shared expert in), else their sum in + ``out``.""" + core = self.core + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, False) + folded = m3_moe_engine(h.contiguous(), engine, None, None) + return folded, None + routed_out, shared_out = maybe_execute_in_parallel( + lambda: self._routed(i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the + # same collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if return_parts: + return shared_out, routed_out + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + if _can_mc(hidden): + fused = _mc_allreduce_norm(core, hidden, None, residual, core.w[f"l{i}_post_norm"]) + hidden, residual = ( + fused + if fused is not None + else _native_allreduce_norm(core, hidden, residual, core.w[f"l{i}_post_norm"]) + ) + else: + hidden, residual = _native_allreduce_norm( + core, hidden, residual, core.w[f"l{i}_post_norm"] + ) + if i in core.moe_layers: + out, routed = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out = self._swiglu_mlp(hidden, f"l{i}_mlp", True) + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step): + """Attention of layer ``i`` up to its o_proj (no reduction). Returns (o_proj output, pushed): a prefill step + arms no boundary push, so pushed is always False.""" + core = self.core + if i not in core.sparse_layers: + # Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v for an FP8 cache, else in place in bf16), the + # MSA backend (it writes the layer's cache), o_proj. + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + if not core.fused_projection: + return self._unfused_sparse_attention(i, position_ids, hidden, step) + packed, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection + and its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and + the FP8 indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the + o_proj. A layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): the MSA + backend (``m3_msa_attention``: the KV write unless a producer wrote the caches, the sparse layers' top-k block + selection, the FMHA over the context rows and any generation rows beside them).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + live = int(step["num_tokens"]) + q_live, out_live = q[:live], output[:live] + k_live = k[:live] if k is not None else None + v_live = v[:live] if v is not None else None + idx_q_live = idx_q[:live] if idx_q is not None else None + sparse = i in core.sparse_layers + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k_live is not None) + m3_msa_attention( + core.attn[i], + q_live, + k_live, + v_live, + idx_q_live, + step["metadata"], + out_live, + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + return output + + +class DecodeTarget: + """A pure generation step (no context rows), of any width. + + Up to ``DECODE_MAX_TOKENS`` tokens with 1-D input ids, per layer: attention (o_proj all-reduce deferred, its rows + pushed into the boundary all-reduce up to 16 tokens) -> attention boundary (all-reduce + residual + post-attention + RMSNorm) -> MoE (the 1..8 / 9..13-row engines, the first pushing its sum into the boundary all-reduce; the + unfinalized MoE with the finalize in the all-reduce from 14 rows) or dense MLP -> MoE boundary (all-reduce + + shared/routed add + residual + the next layer's RMSNorm, emitting the next QKV's MXFP8 activation). A layer whose + output Eagle3 captures runs the generic layer, as does the first layer when the fused embedding does not take the + step. A wider step (or one carrying embeddings) runs every layer in its generic form, with this target's attention + (the decode selection and sparse decode where the producer wrote the caches). + """ + + def __init__(self, core: MiniMaxM3Tp4ModelingV2) -> None: + self.core = core + self._moe_push_descs: Dict[tuple, Any] = {} + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + States the two context counts instead of reading them back: this target runs only when ``num_contexts`` is + zero, so reading them would ask the metadata to confirm what the routing already decided. + + Graph classes: counts, the generation span, the identity-selection flag and the worst-case score extent are + host values fixed per capture (prepare() stages them before the forward); the block table, lengths, cache + slots, identity table and valid-block counts are the metadata's persistent buffers, held by reference and + refreshed in place every step; ``msa_idx_k_cache`` / ``msa_proxy_max_score_view`` are the metadata's getters + of the per-layer index-K cache and of the score scratch view. ``metadata`` is the prepared object itself: the + MSA backend's step (``m3_msa_attention``) takes it as its argument. + """ + return dict( + num_contexts=0, + num_ctx_tokens=0, + num_tokens=md.num_tokens, + num_seqs=md.num_seqs, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + msa_decode_span=md.msa_decode_span, + msa_block_table=md.msa_block_table, + msa_seq_lens_cuda=md.msa_seq_lens_cuda, + msa_short_ok=md._msa_short_ok, + msa_short_table=md.msa_short_table, + msa_proxy_max_score_view=md.msa_proxy_max_score_view, + msa_worst_case_max_k_tiles=md.msa_worst_case_max_k_tiles, + msa_n_valid_blocks=md.msa_n_valid_blocks, + metadata=md, + ) + + def forward( + self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, capture + ): + core = self.core + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + # The decode kernels' rows (a per-capture constant): up to DECODE_MAX_TOKENS tokens of 1-D input ids. + kernels = ( + inputs_embeds is None and input_ids.ndim == 1 and is_decode_step(input_ids.numel()) + ) + if kernels: + num_tokens = int(input_ids.numel()) + plan = plan_for(num_tokens) + _record("decode", num_tokens, plan.name) + hidden, residual = self._embed(input_ids, num_tokens) + else: + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target does not run under piecewise (breakable) CUDA graphs" + ) + plan = None + _record("decode", int(step["num_tokens"]), "generic") + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + if ( + not kernels + or residual is None + or (spec_metadata is not None and spec_metadata.is_layer_capture(i)) + ): + hidden, residual = self._generic_layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + else: + hidden, residual = self._layer(i, plan, position_ids, hidden, residual, step) + return hidden + + def _embed(self, input_ids: torch.Tensor, num_tokens: int): + """(hidden, residual) of the fused embedding + first RMSNorm; (embedding, None) when it does not run, so the + first layer runs its own input norm.""" + core = self.core + if not EMBED_FIRSTNORM.covers(num_tokens): + uncovered(EMBED_FIRSTNORM.path, f"more than {EMBED_FIRSTNORM_MAX_TOKENS} tokens") + return m3_embedding(input_ids, core.w["embed"]), None + owner = core.embed_owners.get((num_tokens, input_ids.dtype)) + if owner is None or not input_ids.is_contiguous(): + missed( + EMBED_FIRSTNORM.path, + num_tokens, + f"no embedding owner for {input_ids.dtype} ids (built for int32 / int64 after the load)" + if owner is None + else "the input ids are not contiguous", + ) + return m3_embedding(input_ids, core.w["embed"]), None + _record("m3_embedding_firstnorm", num_tokens, input_ids.dtype) + return m3_embedding_firstnorm( + input_ids, owner.table, owner.weight, owner.eps, owner.raw, owner.output, owner.use_pdl + ) + + def _bucket_flags(self, num_tokens: int): + """(bucket workspace, its 8-CTA call counters) of the boundary all-reduce for ``num_tokens``, or None.""" + workspace = m3_mc_allreduce.get_workspace(self.core.mapping, self.core.hidden) + workspace = None if workspace is None else m3_mc_allreduce._bucket(workspace, num_tokens) + flags = None if workspace is None else workspace.flags_for(8) + return None if flags is None else (workspace, flags) + + @staticmethod + def _no_bucket(path: str, num_tokens: int) -> None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, "the multicast all-reduce workspace") + else: + missed(path, num_tokens, "no multicast all-reduce bucket holds the step") + + def _attn_push_target(self, num_tokens: int): + """The o_proj's push tuple into the attention boundary's all-reduce, or None (after a miss). + + None comes only from the token count or from workspace state created collectively in the same eager call on + every rank, so all ranks agree. (A lone non-pushing rank would still be safe: its all-reduce pushes the same + bf16 rows into the same set A.)""" + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(ATTN_PUSH.path, num_tokens) + return None + workspace, flags = found + return ( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + AR_PUSH_MODE, + ) + + def _moe_push_targets(self, num_tokens: int): + """(engine push descriptor, shared down-projection push tuple) for the MoE boundary's exchange sets, or None + (after a miss). + + This decision must match across ranks: prepushed 3 reads set B of every rank, which a non-pushing rank never + fills. None comes only from the token count or from state created in the same eager call on every rank + (workspace, flags, the descriptor cache), and the engine it is armed for is agreed at build time + (``_prebuild_engines``), so it does.""" + hidden = self.core.hidden + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(MOE_PUSH.path, num_tokens) + return None + workspace, flags = found + key = (workspace.mc.data_ptr(), flags.data_ptr()) + desc = self._moe_push_descs.get(key) + if desc is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(MOE_PUSH.path, num_tokens, "the MoE engine's push descriptor") + return None + desc = self._moe_push_descs[key] = ar_push_desc( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + flags.device, + ) + set_b = workspace.mc.data_ptr() + 2 * 4 * workspace.max_tokens * workspace.world_size * ( + hidden // 2 + ) + return desc, ( + set_b, + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + 2, + ) + + def _boundary(self, use_mc: bool, x, residual, norm_w, x2=None, quant=False): + """All-reduce + (second addend) + residual + RMSNorm: the multicast kernel, else the built-in all-reduce.""" + core = self.core + if use_mc: + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w, quant=quant) + if fused is not None: + return fused + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + missed( + MC_ALLREDUCE.path, + x.shape[0], + m3_mc_allreduce.decline_reason(x, residual, workspace, x2=x2), + ) + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _prepushed(self, i: int, x, residual, norm_w, quant, prepushed, x2=None): + fused = _mc_allreduce_norm( + self.core, x, x2, residual, norm_w, quant=quant, prepushed=prepushed + ) + if fused is None: + raise RuntimeError(f"layer {i}: rows were pushed into an all-reduce that did not run") + return fused + + def _layer(self, i: int, plan: DecodePlan, position_ids, hidden, residual, step): + core = self.core + num_tokens = hidden.shape[0] + # Padded rows: the boundary all-reduce polls and re-arms rows t < M = residual.shape[0], CUDA-graph padding + # rows included. Every producer pushes all rows of the same tensor: the o_proj / shared down GEMV t < m = + # x.shape[0] (its row count comes from the host shape, not a live count), the MoE engine every row of the + # num_tokens it was built for (= num_tokens here). The all-reduce checks x.sizes() == residual.sizes(), so + # no padded row is left with a stale flag. + # The GEMV dispatch (and with it every push into a boundary all-reduce) is a decode-graph one. + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + if not decode_graph: + uncovered( + "decode_graph", + "an eager decode step: the decode-graph GEMV dispatch and the boundary pushes are off", + ) + push = plan.attn_mc_ar and decode_graph and ATTN_PUSH.covers(num_tokens) + armed = self._attn_push_target(num_tokens) if push else None + hidden, attn_pushed = self._attention(i, position_ids, hidden, step, armed) + post_norm = core.w[f"l{i}_post_norm"] + if armed is not None and attn_pushed: + hidden, residual = self._prepushed(i, hidden, residual, post_norm, False, AR_PUSH_MODE) + else: + if armed is not None: + missed( + ATTN_PUSH.path, + num_tokens, + "the o_proj did not push its rows into the armed boundary all-reduce", + ) + hidden, residual = self._boundary(plan.attn_mc_ar, hidden, residual, post_norm) + + norm_w = _next_norm(core, i) + if i not in core.moe_layers: + out, _ = self._swiglu_mlp(hidden, f"l{i}_mlp", True) + return self._boundary(plan.attn_mc_ar, out, residual, norm_w) + + quant = plan.qkv_prequant and i + 1 < core.num_layers + path = moe_path(num_tokens) + moe_push = None + if path == MOE_ENGINE.path: + # The engine runs inside the MoE; looked up (or found missing) here first so a push is armed only for an + # engine that will take it. + engine = core.engine(i, num_tokens, shared=True) + if engine is not None and plan.moe_mc_ar and decode_graph: + moe_push = self._moe_push_targets(num_tokens) + elif path == MOE16_ENGINE.path and decode_graph: + # Router, routed and shared experts in one kernel; its bf16 sum goes to the boundary all-reduce. + engine = core.moe16_engine(i, num_tokens) + if engine is not None: + _record("m3_moe16_engine", num_tokens) + out = m3_moe16_engine(hidden.contiguous(), engine) + return self._boundary(plan.moe_mc_ar, out, residual, norm_w, quant=quant) + if path in (MOE16_ENGINE.path, MC_FINALIZE.path) and plan.moe_mc_ar: + if self._bucket_flags(num_tokens) is None: + self._no_bucket(MC_FINALIZE.path, num_tokens) + else: + (shared, fc2, scales, perm), _, _, _ = self._moe( + i, hidden, return_parts=False, do_finalize=False + ) + _record("m3_mc_finalize_allreduce_norm", num_tokens, quant) + fused = m3_mc_allreduce.m3_mc_finalize_allreduce_norm( + fc2, + perm.view(num_tokens, -1).to(torch.int32), + scales.reshape(num_tokens, -1), + shared, + residual, + norm_w, + core.rms_eps, + m3_mc_allreduce.get_workspace(core.mapping, core.hidden), + use_pdl=get_env_enable_pdl(), + quant=quant, + ) + if fused is None: + raise RuntimeError( + "the MC finalize all-reduce rejected the unfinalized MoE outputs" + ) + return fused + shared, routed, engine_pushed, shared_pushed = self._moe( + i, hidden, return_parts=True, push=moe_push + ) + if moe_push is not None: + if routed is None: + # The engine folded the shared expert in and pushed routed + shared into set A. + if engine_pushed: + return self._prepushed(i, shared, residual, norm_w, quant, 2) + else: + if engine_pushed and shared_pushed: + return self._prepushed(i, shared, residual, norm_w, quant, 3, x2=routed) + if engine_pushed or shared_pushed: + raise RuntimeError("only one MoE part pushed into the boundary all-reduce") + missed( + MOE_PUSH.path, + num_tokens, + "the MoE did not push its output into the armed boundary all-reduce", + ) + return self._boundary(plan.moe_mc_ar, shared, residual, norm_w, x2=routed, quant=quant) + + def _swiglu_mlp( + self, h: torch.Tensor, prefix: str, low_m_gemm: bool, push: Optional[tuple] = None + ): + """A gated MXFP8 MLP (the dense MLP, the shared expert): gate_up, SwiGLU-OAI, down, no reduction. Returns + (output, whether the down projection's GEMV pushed into ``push``).""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s", low_m_gemm) + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + return _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s", low_m_gemm, push) + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: the MoE + scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor, do_finalize: bool): + """Router + routed experts: the fused router with the NVFP4 activation up to 48 rows, else the router logits and + the experts quantizing their bf16 input.""" + core = self.core + fused = _router_fp4(core, i, h) + if fused is not None: + router_logits, x_fp4, x_sf = fused + return _routed_experts(core, i, x_fp4, x_sf, router_logits, do_finalize) + router_logits = _router_logits(core, i, h) + if not do_finalize: + return _routed_experts(core, i, h, None, router_logits, False) + return self._routed_chunks(i, h, router_logits) + + def _moe( + self, i: int, h: torch.Tensor, *, return_parts: bool, do_finalize: bool = True, push=None + ): + """The MoE block of layer ``i`` without its reduction. + + ``push`` = (the engine's ``arp`` descriptor, the shared down projection's ``ar`` tuple) the decode layer armed + for the MoE boundary's exchange sets, or None. Returns ``(out, routed, engine_pushed, shared_pushed)``: with + ``return_parts`` the shared and routed parts (routed None when the engine folded the shared expert in), else + their sum in ``out``; with ``do_finalize=False`` ``out`` is ``[shared, fc2 out, expert scales, permuted + index]``. + """ + core = self.core + engine_arp, shared_ar = push if push is not None else (None, None) + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + do_finalize + and h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, engine_arp is not None) + folded = m3_moe_engine(h.contiguous(), engine, None, engine_arp) + return folded, None, engine_arp is not None, False + routed_out, (shared_out, shared_pushed) = maybe_execute_in_parallel( + lambda: self._routed(i, h, do_finalize), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False, shared_ar), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the same + # collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if not do_finalize: + return [shared_out, *routed_out], None, False, shared_pushed + if return_parts: + return shared_out, routed_out, False, shared_pushed + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None, False, shared_pushed + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _generic_layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + if _can_mc(hidden): + fused = _mc_allreduce_norm(core, hidden, None, residual, core.w[f"l{i}_post_norm"]) + hidden, residual = ( + fused + if fused is not None + else _native_allreduce_norm(core, hidden, residual, core.w[f"l{i}_post_norm"]) + ) + else: + hidden, residual = _native_allreduce_norm( + core, hidden, residual, core.w[f"l{i}_post_norm"] + ) + if i in core.moe_layers: + out, routed, _, _ = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out, _ = self._swiglu_mlp(hidden, f"l{i}_mlp", True) + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step, push: Optional[tuple] = None): + """Attention of layer ``i`` up to its o_proj (no reduction). ``push`` is the attention boundary's exchange the + o_proj may write into. Returns (o_proj output, pushed).""" + core = self.core + if i not in core.sparse_layers: + # Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v: an FP8 cache, or the FP8 pool the dense layers + # sit in within a hybrid NVFP4 cache on a decode step), the MSA backend, o_proj. + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + out_fp8 = core.main_kv_is_fp8 or core.main_kv_is_nvfp4 + fused = _qk_norm_rope(core, i, qkv, position_ids, out_fp8) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + if not core.fused_projection: + return self._unfused_sparse_attention(i, position_ids, hidden, step, push) + # Gathered before the QKV projection so nothing runs between it and the fused attention (PDL hand-off). + fused = self._fused_attn_args(i, position_ids, step) + if fused is not None: + mega = _attn_mega(core, i, hidden, fused, push) + if mega is not None: + return mega + packed, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + if fused is not None and ( + packed.shape[0] != fused["num_tokens"] or packed.shape[1] != 2560 + ): + missed( + FUSED_DECODE_ATTN.path, + fused["num_tokens"], + f"the packed QKV is {tuple(packed.shape)}, not [rows, 2560]", + ) + fused = None + if fused is not None: + c1 = _attn_c1(core, i, packed, fused, push) + if c1 is not None: + return c1 + o = _fused_attn(core, i, packed, fused) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step, push): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection and + its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and the FP8 + indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the o_proj. A + layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _fused_attn_args(self, i: int, position_ids, step) -> Optional[dict]: + """Arguments of the fused decode attention for this step, or None when it does not apply. + + The layer must be an NVFP4 sparse MSA layer and the step a pure generation step whose selection the metadata + staged as the identity (every request within 16 blocks of 128 tokens), with at most 8 query tokens per request. + """ + core = self.core + path = FUSED_DECODE_ATTN.path + if i not in core.sparse_layers or not core.main_kv_is_nvfp4: + uncovered(path, "not an NVFP4 sparse MSA layer with an FP8 indexer") + return None + if position_ids is None: + uncovered(path, "not a pure generation step with positions") + return None + if not step["msa_short_ok"]: + uncovered( + path, "a request beyond 16 KV blocks: the block selection is not the identity" + ) + return None + span = step["msa_decode_span"] + if span is None or span.row_first != 0: + uncovered(path, "no decode span starting at row 0") + return None + num_reqs = int(step["num_seqs"]) + qlen = int(span.query_len) + if not FUSED_DECODE_ATTN.covers((qlen, num_reqs)): + uncovered( + path, + f"more than {FUSED_ATTN_MAX_QUERY_LEN} query rows per request or {FUSED_ATTN_MAX_TOKENS} per step", + ) + return None + manager = step["kv_cache_manager"] + if manager is not None and not manager.is_nvfp4_layer(i): + uncovered(path, "the layer's KV cache is not NVFP4") + return None + if _fused_attn_kernel._module is None: + # Building compiles every geometry (minutes): only outside capture (the eager warmup steps). + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_reqs * qlen, "the fused attention module") + return None + _fused_attn_kernel._load_module() + if manager is None: + missed(path, num_reqs * qlen, "no KV-cache manager") + return None + k_scale, v_scale = _aligned_nvfp4_dequant_scales(core.attn[i], core.kv_scales[i]) + return dict( + num_tokens=num_reqs * qlen, + decode_query_len=qlen, + kv_data=manager.get_buffers(i, kv_layout="HND").view(torch.uint8), + kv_scale=manager.get_block_scale_buffers(i, kv_layout="HND"), + index_k_cache=step["msa_idx_k_cache"](i), + out_cache_loc=step["msa_out_cache_loc"], + position_ids=position_ids.reshape(-1).contiguous().to(torch.int32), + block_table=step["msa_block_table"][:num_reqs], + seq_lens=step["msa_seq_lens_cuda"][:num_reqs], + kv_quant_scale=core.inv_kv_scales[i], + k_global_scale=k_scale, + v_global_scale=v_scale, + ) + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): a sparse + layer whose caches the producer wrote runs this target's selection and sparse decode when the generation span + starts at row 0; everything else the MSA backend (``m3_msa_attention``: the KV write unless a producer wrote the + caches, the sparse layers' top-k block selection, the FMHA).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + live = int(step["num_tokens"]) + q_live, out_live = q[:live], output[:live] + k_live = k[:live] if k is not None else None + v_live = v[:live] if v is not None else None + idx_q_live = idx_q[:live] if idx_q is not None else None + span = step["msa_decode_span"] + sparse = i in core.sparse_layers + if sparse and k_live is None and span is not None and span.row_first == 0: + self._decode_sparse_core(i, q_live, idx_q_live, step, out_live) + return output + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k_live is not None) + m3_msa_attention( + core.attn[i], + q_live, + k_live, + v_live, + idx_q_live, + step["metadata"], + out_live, + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + return output + + def _decode_sparse_core(self, i: int, q, idx_q, step, output) -> None: + """Top-16 block selection and the block-sparse decode of one pure generation step of a sparse layer whose caches + the producer already wrote. + + The selection is the identity when the metadata staged it (every request within 16 blocks). Otherwise the CuTe + DSL scorer writes the per-index-head block maxima, reduced to the KV head, and the long-row selector takes the + top 16 (the built-in selector beyond its shapes). + """ + core = self.core + backend = core.attn[i] + config = backend.m3_config + num_tokens = int(q.shape[0]) + rows = int(step["num_seqs"]) + qlen = int(step["msa_decode_span"].query_len) + block_table = step["msa_block_table"][:rows] + seq_lens = step["msa_seq_lens_cuda"][:rows] + num_kv_heads = int(backend.indexer.config.num_kv_heads) + identity = step["msa_short_ok"] and step["msa_short_table"] is not None + if identity: + topk = step["msa_short_table"][: num_kv_heads * num_tokens * MSA_REQUIRED_TOPK].view( + num_kv_heads, num_tokens, MSA_REQUIRED_TOPK + ) + else: + idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) + idx_k_cache = step["msa_idx_k_cache"](i) + max_score = step["msa_proxy_max_score_view"]( + config.num_index_heads, step["msa_worst_case_max_k_tiles"], num_tokens + ) + early = rows <= INDEX_SCORE_EARLY_MAX_REQUESTS + key = ( + idx_q_view.dtype, + int(config.num_index_heads), + int(qlen), + int(idx_q_view.shape[-1]), + int(INDEX_SCORE_SPLIT_K), + bool(early), + ) + if key not in _index_score_kernel.IndexScoreRunner.kernel_cache: + # Outside require mode the op below compiles it (the only scorer this step has). + missed( + INDEX_SCORE.path, + num_tokens, + f"the index scorer {key} was not compiled before warmup", + ) + _record("m3_index_score", key) + m3_index_score( + idx_q_view, + idx_k_cache.squeeze(1), + block_table, + seq_lens, + max_score.transpose(1, 2), + qlen, + INDEX_SCORE_SPLIT_K, + early, + ) + n_valid_blocks = step["msa_n_valid_blocks"][:num_tokens] + group = config.num_index_heads // config.num_kv_heads + max_score_kv = ( + max_score.view( + config.num_kv_heads, group, max_score.shape[1], max_score.shape[2] + ).amax(dim=1) + if group > 1 + else max_score + ) + table = long_select_blocks( + max_score_kv, + n_valid_blocks=n_valid_blocks, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + use_pdl=get_env_enable_pdl(), + early_trigger=num_tokens <= SELECT_EARLY_MAX_ROWS, + ) + if table is None: + uncovered( + LONG_SELECT.path, + f"{max_score_kv.shape[1]} KV blocks, outside {LONG_SELECT_MIN_BLOCKS}..{LONG_SELECT_MAX_BLOCKS}: " + "the built-in selector", + ) + table = m3_select_blocks( + max_score_kv, + MSA_REQUIRED_TOPK, + n_valid_blocks, + config.init_blocks, + config.local_blocks, + True, + ) + topk = table.permute(1, 0, 2) + manager = step["kv_cache_manager"] + k_paged, v_paged = msa_paged_kv(manager, i) + nvfp4 = {} + if manager.is_nvfp4_layer(i): + k_scale, v_scale = _aligned_nvfp4_dequant_scales(backend, core.kv_scales[i]) + scales = manager.get_block_scale_buffers(i, "HND") + nvfp4 = dict( + k_block_scale=scales[:, 0], + v_block_scale=scales[:, 1], + k_global_scale=k_scale, + v_global_scale=v_scale, + ) + _record("m3_sparse_decode", num_tokens, qlen, identity) + m3_sparse_decode( + q.view(num_tokens, core.heads, core.head_dim), + k_paged, + v_paged, + topk, + block_table, + seq_lens, + sm_scale=core.sm_scale(i), + output=output.view(num_tokens, core.heads, core.head_dim), + decode_query_len=qlen, + identity_blocks=identity, + **nvfp4, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# Speculative decoding: one-model Eagle3 at checkpoint precision +# -------------------------------------------------------------------------------------------------------------------- + + +class M3SpecStepFold: + """The target's speculative step-input staging (the in-tree ``SpecStepFold`` protocol). + + Between two decode graphs of a one-model Eagle3 step, the spec sampler stores the step's outputs into its per-slot + store and the engine gathers them back into its flat input buffers for the next step. This object does both in + one launch each (``m3_spec_tail.slot_scatter`` / ``step_gather``): integer copies only, the same values the + engine's and the sampler's own copies write. The forward never stores the sampler's outputs itself, so + ``begin_step`` and ``sampler_rows_stored`` always report that nothing was staged. + """ + + def __init__(self) -> None: + self.store: Optional[object] = None + self.step_inputs: Optional[tuple] = None + + def bind_engine_buffers( + self, input_ids, draft_tokens, position_id_offsets, kv_lens_offsets, max_batch_size + ): + self.step_inputs = (input_ids, draft_tokens, position_id_offsets, kv_lens_offsets) + + def bind_sampler_store(self, store: object) -> None: + self.store = store + + def begin_step(self, seq_slots: Optional[list], num_rows: int) -> bool: + return False + + def gather( + self, + new_tokens, + next_draft_tokens, + new_tokens_lens, + previous_slots, + previous_pos_indices, + previous_batch_len: int, + tokens_per_gen_step: int, + draft_token_buffer_width: int, + num_tokens: int, + num_draft_tokens: int, + first_previous_row: int, + ) -> bool: + if self.step_inputs is None: + return False + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets = self.step_inputs + return m3_spec_tail.step_gather( + new_tokens, + next_draft_tokens, + new_tokens_lens, + previous_slots, + previous_pos_indices, + previous_batch_len, + tokens_per_gen_step, + draft_token_buffer_width, + input_ids, + num_tokens, + draft_tokens, + num_draft_tokens, + position_id_offsets, + kv_lens_offsets, + first_previous_row, + ) + + def sampler_rows_stored(self, outputs: dict, num_skip: int, seq_slots: list) -> bool: + return False + + def scatter( + self, seq_slots, new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens + ) -> bool: + store = self.store + if store is None: + return False + return m3_spec_tail.slot_scatter( + seq_slots, + new_tokens, + store.new_tokens.squeeze(-1).T, + next_new_tokens, + store.next_new_tokens.squeeze(-1).T, + new_tokens_lens, + store.new_tokens_lens, + next_draft_tokens, + store.next_draft_tokens, + ) + + +class _DraftForward: + """``draft_model.model`` for the in-tree worker: the drafter layer's forward, called as + ``model(input_ids=, position_ids=, hidden_states=, attn_metadata=, spec_metadata=, all_rank_num_tokens=)`` and + returning (hidden, hidden to carry to the next pass). ``d2t`` None: the draft vocabulary is the target's.""" + + d2t = None + + def __init__(self, draft: "M3Eagle3Draft") -> None: + self._draft = draft + + def __call__(self, **kwargs): + return self._draft.forward_layer(**kwargs) + + +class M3Eagle3Draft(nn.Module): + """The one-layer BF16 Eagle3 drafter (``speculative_model``): its weights (``weights.DRAFT_WEIGHTS``), its + vocab-parallel BF16 head and logits processor, the drafter layer's TRTLLM attention backend (layer index = the + target's layer count, FP8 KV) and its TP all-reduce. + + What the worker reads: ``model`` (the layer forward), ``apply_eagle3_fc`` (the per-layer fc_norms + fc over the + captured target hidden states), ``lm_head`` and ``logits_processor``; the fused draft loop reads the layer's parts + (``draft_input_weights``, ``attention``, ``mlp``, ``ar_norm``). The drafter embeds with the target's table. + """ + + def __init__( + self, + core: MiniMaxM3Tp4ModelingV2, + draft_config: ModelConfig, + start_layer_idx: int, + capture_layers: int, + ) -> None: + super().__init__() + cfg = draft_config.pretrained_config + _ensure_draft_vocab_size(cfg) + eagle = getattr(cfg, "eagle_config", None) or {} + mapping = draft_config.mapping + self.draft_config = draft_config + # The engine's draft weight mapper binds to these two. + self.model_config = draft_config + self.config = cfg + self.capture_layers = int(capture_layers) + self.layer_idx = int(start_layer_idx) + self.dtype = cfg.torch_dtype + self.hidden = int(cfg.hidden_size) + self.hidden_size_in = int(getattr(cfg, "target_hidden_size", cfg.hidden_size)) + self.rms_eps = float(cfg.rms_norm_eps) + self.return_hidden_post_norm = bool( + eagle.get("return_hidden_post_norm", False) or getattr(cfg, "norm_output", False) + ) + if ( + int(cfg.num_hidden_layers) != 1 + or eagle.get("eh_proj_before_attn", False) + or eagle.get("norm_before_fc", False) + or not getattr(cfg, "fc_norm", False) + or getattr(cfg, "kv_lora_rank", None) + or self.capture_layers <= 1 + or self.hidden_size_in != self.hidden + or cfg.draft_vocab_size != cfg.vocab_size + or getattr(cfg, "attention_bias", False) + or getattr(cfg, "mlp_bias", False) + or getattr(cfg, "bias", False) + or getattr(cfg, "tie_word_embeddings", False) + ): + raise NoTarget( + "this MiniMax-M3 target drafts with a one-layer Llama-style Eagle3 drafter with per-layer " + "fc_norms over " + "several captured layers, the target's hidden size and vocabulary and no biases" + ) + tp = mapping.tp_size + self.heads = cfg.num_attention_heads // tp + self.kv_heads = (cfg.num_key_value_heads + tp - 1) // tp + self.head_dim = int( + getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads + ) + self.q_size = self.heads * self.head_dim + self.kv_size = self.kv_heads * self.head_dim + self.w = _weights.DRAFT_WEIGHTS.declare(_weights.DRAFT_WEIGHTS.dims(self)) + self.lm_head = LMHead( + cfg.draft_vocab_size, + cfg.hidden_size, + dtype=cfg.torch_dtype, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=draft_config.lm_head_gather_output, + reduce_output=False, + use_custom_cublas_mm=False, + quant_config=None, + ) + self.logits_processor = LogitsProcessor() + self.allreduce = AllReduce( + mapping=mapping, strategy=draft_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + quant_config = draft_config.get_quant_config() + self.attn = create_attention( + draft_config.attn_backend, + self.layer_idx, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, rope=RopeParams.from_config(cfg) + ), + quant_config=quant_config, + skip_create_weights_in_init=draft_config.skip_create_weights_in_init, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + if not self.attn.support_fused_rope() or not self.attn.support_fused_qkv(): + raise NoTarget( + f"the drafter's attention backend {type(self.attn).__name__} must fuse RoPE and QKV" + ) + self.attn.update_quant_config(quant_config) + # The drafter embeds with the target's table, read off the core on every call (the engine materializes the + # core's weights after construction). + self._core = [core] + self.model = _DraftForward(self) + + @property + def embed_table(self) -> torch.Tensor: + return self._core[0].w["embed"] + + def load_weights_from_target_model(self, target_model: nn.Module) -> None: + """Nothing to alias: the drafter reads the target's embedding table through the core and owns its head.""" + + # ------------------------------------------------------------------------------------------------- the layer + + def fc(self, hs: torch.Tensor) -> torch.Tensor: + return m3_bf16_linear(hs, self.w["fc"]) + + def apply_eagle3_fc(self, hidden_states: torch.Tensor) -> torch.Tensor: + """The captured target hidden states [T, layers x H] to the drafter's input width: each layer's fc_norm, + then the fc (the in-tree ``Eagle3ForCausalLM.apply_eagle3_fc``).""" + hidden_states = hidden_states.to(self.dtype) + if hidden_states.shape[-1] != self.hidden: + chunks = hidden_states.chunk(self.capture_layers, dim=-1) + hidden_states = torch.cat( + [ + flashinfer_rmsnorm(c, self.w[f"fc_norm{i}"], self.rms_eps) + for i, c in enumerate(chunks) + ], + dim=-1, + ) + hidden_states = self.fc(hidden_states) + return hidden_states + + def attention(self, x: torch.Tensor, md, reduce: bool) -> torch.Tensor: + """QKV, the backend's fused RoPE + KV append + causal attention, o_proj (all-reduced unless ``reduce`` is + False, where the caller fuses the reduction).""" + qkv = m3_bf16_linear(x, self.w["qkv"]) + _record("m3_draft_attention", int(md.num_tokens)) + o = m3_draft_attention(self.attn, qkv, md) + out = m3_bf16_linear(o, self.w["o"]) + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def mlp(self, h: torch.Tensor, reduce: bool) -> torch.Tensor: + gu = m3_bf16_linear(h, self.w["gu"]) + act = m3_silu_and_mul(gu, None, None, None) + out = m3_bf16_linear(act, self.w["dn"]) + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def ar_norm(self, x: torch.Tensor, residual: torch.Tensor, norm_key: str): + """One draft-layer boundary in the in-tree structure: the projection's plain all-reduce (BF16 out), then the + fused residual add + RMSNorm (FP32 add, norm of the FP32 sum). Returns (normed, residual sum).""" + return flashinfer_fused_add_rmsnorm( + m3_allreduce(x, self.allreduce, None), residual, self.w[norm_key], self.rms_eps + ) + + def forward_layer( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata, + spec_metadata, + all_rank_num_tokens: Optional[List[int]] = None, + ): + """The drafter layer (the in-tree ``Eagle3DraftModel.forward`` for one layer): the embedded tokens and the + fc'd hidden states normed and concatenated, attention, the post-attention norm with the residual, the MLP and + the final norm. Returns (normed output, the hidden state the next pass carries).""" + previous = attn_metadata.all_rank_num_tokens + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = all_rank_num_tokens + try: + embeds = m3_embedding(input_ids, self.embed_table).to(self.dtype) + residual = hidden_states + h = flashinfer_rmsnorm(hidden_states, self.w["hidden_norm"], self.rms_eps) + e = flashinfer_rmsnorm(embeds, self.w["in_norm"], self.rms_eps) + h = torch.cat([e, h], dim=-1) + h = self.attention(h, attn_metadata, reduce=True) + h, residual = flashinfer_fused_add_rmsnorm( + h, residual, self.w["post_norm"], self.rms_eps + ) + h = self.mlp(h, reduce=True) + spec_metadata.maybe_capture_hidden_states(self.layer_idx, h, residual) + h, to_save = flashinfer_fused_add_rmsnorm( + h, residual, self.w["final_norm"], self.rms_eps + ) + return (h, h) if self.return_hidden_post_norm else (h, to_save) + finally: + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = previous + + +def _eagle3_capture(spec_metadata, layer_id: int, hidden_states, residual, capture_norm) -> None: + """The Eagle3 hidden-state capture of one target layer. + + Mirrors ``Eagle3OneModelSpecMetadata.maybe_capture_hidden_states``. With ``capture_norm`` = (the drafter's stacked + fc_norm weights [layers, H], eps), the capture writes rmsnorm(bf16(h + residual)) * fc_norm[i] in one kernel (the + in-tree capture's bf16 add, then the in-tree fc_norm arithmetic), which the fused first drafter inputs read as is; + otherwise it is the in-tree capture. + """ + if capture_norm is None: + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + return + for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): + if captured_layer_id == layer_id: + assert spec_metadata.hidden_states is not None + num_tokens = spec_metadata.num_capture_tokens + if num_tokens > spec_metadata.hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds buffer capacity: " + f"num_tokens={num_tokens}, capacity={spec_metadata.hidden_states.shape[0]}" + ) + if num_tokens > hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds available hidden states: " + f"num_tokens={num_tokens}, hidden_states={hidden_states.shape[0]}" + ) + hidden_size = spec_metadata.hidden_size + m3_spec_tail.capture_norm( + hidden_states[:num_tokens], + residual[:num_tokens] if residual is not None else None, + capture_norm[0][i], + capture_norm[1], + spec_metadata.hidden_states[:num_tokens, i * hidden_size : (i + 1) * hidden_size], + ) + break + + +def _sharded_head(lm_head: nn.Module) -> bool: + """A vocab-sharded TP head without padding (the fused kernels read the local shard).""" + return ( + getattr(lm_head, "tp_mode", None) == TensorParallelMode.COLUMN + and getattr(lm_head, "padding_size", 0) == 0 + and getattr(lm_head, "tp_size", 1) > 1 + ) + + +class M3Eagle3Worker(Eagle3OneModelWorker): + """``Eagle3OneModelWorker`` with the MiniMax-M3 fused verify and fused draft loop, at the in-tree numerics. + + * Verify: on decode-only batches whose verify rows the fused kernels cover (``tail_eligible``), the target's BF16 + lm_head computes this rank's vocabulary shard of the logits (the in-tree head's own local GEMM; at most + ``TARGET_HEAD_MAX_ROWS`` verify rows stream a lossless 12-bit coding of the same BF16 weight) and one kernel + replaces the logits all-gather, the argmax (or the temperature + top-p draw) and the strict acceptance. + * Draft: the drafter's first pass may take the per-layer fc_norms in one launch (or folded into the hidden-state + capture) and run its fc beside the target head GEMM (forked right after the last captured target layer at most + ``FC_EARLY_MAX_ROWS`` verify rows). On a decode-only step ``DRAFT_STEP`` covers, the draft loop is fused: each + pass's draft input in one launch, the drafter layer's attention and MLP with their all-reduces fused with the + residual add and the next RMSNorm, and one kernel for the BF16 head, the cross-rank argmax and the pass's + bookkeeping. + + Every other batch takes the in-tree paths. Methods whose docstring says "Mirrors" copy the named in-tree method + with the fused branches inserted. + """ + + def __init__(self, spec_config, mapping, use_separate_draft_kv_cache: bool = False) -> None: + super().__init__( + spec_config, mapping=mapping, use_separate_draft_kv_cache=use_separate_draft_kv_cache + ) + # Per-step state, reset at the top of every _forward_impl: the first draft pass's gather rows when this step + # takes the fused first drafter inputs, and the fc output a fork computed beside the target head. + self._m3_first_gather_ids: Optional[torch.Tensor] = None + self._m3_fc_hidden: Optional[torch.Tensor] = None + self._m3_target_lm_head: Optional[nn.Module] = None + # (fc_norm0 weight, stacked fc_norm weights [layers, H]), built eagerly on the first fused first pass. + self._m3_fc_norm_w: Optional[tuple] = None + # Armed per step by arm_capture_norm: (stacked fc_norm weights, eps) or None. + self._m3_capture_norm: Optional[tuple] = None + self._m3_aux_stream: Optional[torch.cuda.Stream] = None + # Armed per step by arm_capture_norm: the drafter whose first fc is forked at the last captured target layer, + # and the fork that did it, joined by _m3_target_accept. + self._m3_fc_early_drafter: Optional[M3Eagle3Draft] = None + self._m3_fc_early_stream: Optional[torch.cuda.Stream] = None + # The 12-bit coded copy of the target head shard the fused verify streams, built before the warmup. + self._m3_target_head = None + self._m3_force_vec_warm = False + # The streamed copy of the drafter's BF16 head shard the fused draft loop reads, built before the warmup. + self._m3_draft_head = None + # Per step: the next step's input tokens [nreq, D + 1], column 0 written by the fused verify (or the seed) and + # columns 1.. by the fused draft passes; set as _m3_next_new_tokens once the fused draft loop filled it. + self._m3_next_new: Optional[torch.Tensor] = None + self._m3_next_new_tokens: Optional[torch.Tensor] = None + # Per step: the first draft inputs are the fused ones (accepted tokens as they lie, fc output of every row). + self._m3_fused_first_inputs = False + + def capture_hidden_states( + self, spec_metadata, layer_id: int, hidden_states, residual=None + ) -> None: + """The Eagle3 capture with this step's capture norm; on a step armed for the early fc fork + (``arm_capture_norm``), the last captured layer then forks the drafter's first fc onto the aux stream, where + it overlaps the remaining target layers and the head (joined by ``_m3_target_accept``).""" + _eagle3_capture(spec_metadata, layer_id, hidden_states, residual, self._m3_capture_norm) + drafter = self._m3_fc_early_drafter + layers = spec_metadata.layers_to_capture + if drafter is not None and layers and layer_id == layers[-1]: + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = self._m3_fc_fork( + spec_metadata, spec_metadata.num_capture_tokens, drafter + ) + + def prebuild(self, lm_head: nn.Module, draft: Optional[M3Eagle3Draft]) -> None: + """The fused verify's coded target head and the fused draft loop's head copy, from the loaded weights, + before the executor's warmup, on every rank.""" + start = time.perf_counter() + if self.prebuild_target_head(lm_head): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] coded target head built before warmup: %.1f s" + % (time.perf_counter() - start) + ) + start = time.perf_counter() + if self.prebuild_draft_head(draft): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] fused draft head built before warmup: %.1f s" + % (time.perf_counter() - start) + ) + + def _forward_impl( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + """Mirrors Eagle3OneModelWorker._forward_impl, with the fused verify branch.""" + runtime_draft_len = spec_metadata.runtime_draft_len + # skip the draft forward if the runtime draft length is 0 + if runtime_draft_len == 0: + return self.skip_drafting( + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + ) + + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + + raw_logits = logits + + self._m3_first_gather_ids = None + self._m3_next_new = None + self._m3_next_new_tokens = None + self._m3_fused_first_inputs = False + if logits is None: + # tail_eligible() skipped the gathered target logits: verify from this rank's logits shard. + accepted_tokens, num_accepted_tokens, self._m3_first_gather_ids = ( + self._m3_target_accept( + hidden_states, spec_metadata, batch_size, runtime_draft_len, draft_model + ) + ) + else: + self._execute_guided_decoder_if_present(logits) + + # Sample and accept tokens. ``input_ids`` is required by the relaxed-acceptance path (scans for + # thinking-phase tokens); ignored otherwise. + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + input_ids, logits, attn_metadata, spec_metadata + ) + if self.draft_only_eligible(attn_metadata, spec_metadata, draft_model): + self._m3_seed_first_pass(num_accepted_tokens, spec_metadata, batch_size) + + # Mamba hybrid models need state updates after token acceptance (a no-op on this target's cache manager). + if self._is_mamba_hybrid_cache is None: + self._is_mamba_hybrid_cache = isinstance( + attn_metadata.kv_cache_manager, MambaHybridCacheManager + ) + if num_gens > 0 and self._is_mamba_hybrid_cache: + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices, + ) + + sa_manager = getattr(spec_metadata.spec_resource_manager, "sa_manager", None) + if self.sa_enhancer is not None and sa_manager is not None: + self.sa_enhancer.extend_and_prepare( + sa_manager=sa_manager, + request_ids=spec_metadata.request_ids, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_gens=num_gens, + num_contexts=num_contexts, + max_draft_len=runtime_draft_len, + ) + + if self._uses_external_shared_target_kv: + next_draft_tokens = self._forward_external_shared_target_kv_draft_loop( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + ) + else: + # Save the old attn_metadata and spec_metadata + self._prepare_attn_metadata_for_spec_dec(attn_metadata) + + # Prepare inputs for the 1st draft model forward + position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + # Predict draft tokens. ``original_all_rank_num_tokens`` is saved here so the post-loop restore (below) + # can put attn_metadata back into a state the target model expects. + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + + # Get the draft KV cache manager if using separate layouts + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + + next_draft_tokens = self._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + # restore attn_metadata to support cuda graph + self._restore_attn_metadata_from_spec_dec(attn_metadata) + # restore all_rank_num_tokens for attention DP + if original_all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + + # prepare next new tokens to support overlap scheduler + if self._m3_next_new_tokens is not None: + # Filled in place by the fused verify or the seed (column 0) and the fused draft passes (columns 1..). + next_new_tokens = self._m3_next_new_tokens + else: + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + attn_metadata.use_spec_decoding = True + + # Commit model-owned target state only after draft preparation has succeeded; SpecWorkerBase aborts the + # still-pending snapshots on any exception above this point. + if self._auxiliary_state_handlers: + self.commit_auxiliary_speculative_states( + num_accepted_tokens, + attn_metadata.mamba_metadata.state_indices[:batch_size], + num_contexts, + ) + + return { + "logits": raw_logits, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + } + + def prepare_1st_drafter_inputs( + self, + input_ids, + position_ids, + hidden_states, + accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + ): + """The fused first draft inputs on a step that verified with the fused kernels (or was seeded for them), + else the in-tree ones.""" + fused = self._m3_first_drafter_inputs( + position_ids, + accepted_tokens, + input_ids.shape[0], + attn_metadata.num_contexts, + spec_metadata, + draft_model, + ) + self._m3_fused_first_inputs = fused is not None + if fused is not None: + fused.update(attn_metadata=attn_metadata, spec_metadata=spec_metadata) + return fused + return super().prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + def _m3_structural_ok(self, attn_metadata, spec_metadata) -> bool: + """The batch conditions shared by the fused verify and the fused first drafter inputs: a decode-only plain + Eagle3 step without attention DP, guided decoding, relaxed acceptance or a suffix-automaton enhancer, of at + most DECODE_MAX_TOKENS verify rows. Host state only, so the TP ranks agree.""" + if self.is_mtp_eagle or self.use_dynamic_tree: + return False + if self._uses_external_shared_target_kv or self.sa_enhancer is not None: + return False + if ( + self.guided_decoder is not None + or self.mapping is None + or self.mapping.enable_attention_dp + ): + return False + if getattr(self.spec_config, "use_relaxed_acceptance_for_thinking", False): + return False + if attn_metadata.num_contexts != 0 or attn_metadata.num_seqs == 0: + return False + draft_len = spec_metadata.runtime_draft_len + if draft_len == 0 or attn_metadata.num_seqs * (draft_len + 1) > DECODE_MAX_TOKENS: + return False + return not spec_metadata.wants_advanced_draft_sampling + + def tail_eligible(self, attn_metadata, spec_metadata, lm_head: nn.Module) -> bool: + """Whether this step verifies with the fused kernels instead of the gathered target logits. + + Decode-only Eagle3 batches (``_m3_structural_ok``) over an unpadded vocab-sharded bias-free target head, all + greedy at a request count ``SHARD_ACCEPT`` covers, or sampled under the temperature + top-p contract (see + _m3_sampled_verify_ok), with nothing else reading the target logits. Every input is host state shared by all + TP ranks, so the ranks agree. + """ + if not self._m3_structural_ok(attn_metadata, spec_metadata): + return False + if getattr(spec_metadata, "enable_penalty", False): + return False + if not _sharded_head(lm_head) or getattr(lm_head, "bias", None) is not None: + return False + draft_len = spec_metadata.runtime_draft_len + if spec_metadata.is_all_greedy_sample: + key = (draft_len, attn_metadata.num_seqs) + if not SHARD_ACCEPT.covers(key): + uncovered( + SHARD_ACCEPT.path, + f"draft length {draft_len} at {attn_metadata.num_seqs} requests", + ) + return False + elif not self._m3_sampled_verify_ok(spec_metadata): + return False + self._m3_target_lm_head = lm_head + return True + + def draft_only_eligible(self, attn_metadata, spec_metadata, draft_model) -> bool: + """Whether a step verified by the in-tree path can still take the fused first drafter inputs: the structural + conditions of tail_eligible, without the target-sampling ones.""" + return self._m3_structural_ok(attn_metadata, spec_metadata) and draft_model is not None + + def _m3_seed_first_pass(self, num_accepted_tokens, spec_metadata, batch_size: int) -> None: + """The first draft pass's gather rows (the last accepted row of each request) from the in-tree acceptance, as + _m3_target_accept leaves them.""" + k1 = spec_metadata.runtime_draft_len + 1 + n = num_accepted_tokens[:batch_size].to(torch.int64) + rows = spec_metadata.batch_indices_cuda[:batch_size].to(torch.int64) + self._m3_first_gather_ids = rows * k1 + n - 1 + + def _m3_first_drafter_inputs( + self, + position_ids, + accepted_tokens, + num_tokens: int, + num_contexts: int, + spec_metadata, + draft_model: Optional[M3Eagle3Draft], + ): + """The in-tree first draft inputs of a decode-only step, with the per-layer fc_norm RMSNorms of + ``apply_eagle3_fc`` in one kernel (unless the capture already applied them) before the fc; the input ids are + the accepted tokens as they lie. None to use the in-tree path.""" + if ( + self._m3_first_gather_ids is None + or self.is_mtp_eagle + or num_contexts != 0 + or draft_model is None + ): + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + ): + return None + cache = self._m3_fc_norm_w + if cache is None or cache[0] is not draft_model.w["fc_norm0"]: + if torch.cuda.is_current_stream_capturing(): + capture_missed("spec_tail", None, "the stacked Eagle3 fc_norm weights") + return None + cache = ( + draft_model.w["fc_norm0"], + torch.stack( + [draft_model.w[f"fc_norm{i}"].detach() for i in range(layers)] + ).contiguous(), + ) + self._m3_fc_norm_w = cache + hidden = self._m3_fc_hidden + self._m3_fc_hidden = None + if hidden is None or hidden.shape[0] != num_tokens: + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + hs = m3_spec_tail.chunk_norm(hs, cache[1], draft_model.rms_eps) + hidden = draft_model.fc(hs) + input_ids = accepted_tokens[:, : spec_metadata.runtime_draft_len + 1].reshape(-1) + return {"input_ids": input_ids, "position_ids": position_ids, "hidden_states": hidden} + + def _m3_fc_fork( + self, + spec_metadata, + num_tokens: int, + draft_model: Optional[M3Eagle3Draft], + after: Optional[torch.cuda.Event] = None, + ) -> Optional[torch.cuda.Stream]: + """Launch the first draft pass's fc_norm + fc on a forked stream beside the target head GEMM (after the event + ``after`` when given, else after everything queued on the current stream); the result is picked up by + _m3_first_drafter_inputs. Returns the stream to join, or None when the fc path does not apply (the first + drafter inputs then compute it).""" + self._m3_fc_hidden = None + if draft_model is None: + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + or self._m3_fc_norm_w is None + ): + return None + aux = self._m3_aux_stream + if aux is None: + aux = self._m3_aux_stream = torch.cuda.Stream() + if after is not None: + aux.wait_event(after) + else: + aux.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(aux): + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + hs = m3_spec_tail.chunk_norm( + hs, self._m3_fc_norm_w[1], draft_model.rms_eps, use_pdl=False + ) + self._m3_fc_hidden = draft_model.fc(hs) + return aux + + def _m3_fc_join(self, fork: torch.cuda.Stream) -> None: + """Join the fc fork onto the current stream and keep its output's memory alive for that stream. + + _m3_fc_hidden is allocated in fork (aux) stream order and read on the current stream. Without the record, + dropping it returns its block to the aux stream's pool, where later aux work can take it while the current + stream's reads are still queued (eager warm-up / mixed steps). Nothing is recorded under graph capture, as in + maybe_execute_in_parallel: a captured graph relies on the next fork's aux.wait_stream(main) to order the + reuse.""" + main = torch.cuda.current_stream() + main.wait_stream(fork) + hidden = self._m3_fc_hidden + if ( + hidden is not None + and hidden.is_cuda + and not (torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing()) + ): + hidden.record_stream(main) + + def arm_capture_norm( + self, attn_metadata, spec_metadata, lm_head: nn.Module, draft_model: Optional[M3Eagle3Draft] + ) -> None: + """Before every target forward: whether this step's Eagle3 captures apply the drafter's fc_norm themselves. + + On a fused-verify step whose fc_norm weights are already stacked (by an earlier fused first pass), the + capture writes rmsnorm(h + r) * fc_norm[i] (one kernel per captured layer instead of add + copy) and the fused + first drafter inputs read it without their own chunk_norm. A fused-verify step of at most + ``FC_EARLY_MAX_ROWS`` verify rows is also armed to fork the drafter's first fc at the last captured layer + (``capture_hidden_states``). + """ + self._m3_capture_norm = None + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = None + if draft_model is None: + return + if not self.tail_eligible(attn_metadata, spec_metadata, lm_head): + return + # The fork reads every captured column at the last captured layer, so a final-state capture, which is written + # after the model returns, keeps the fc at its stock place. + if ( + attn_metadata.num_seqs * (spec_metadata.runtime_draft_len + 1) <= FC_EARLY_MAX_ROWS + and FINAL_HIDDEN_LAYER_ID not in spec_metadata.layers_to_capture + ): + self._m3_fc_early_drafter = draft_model + fc_norm_w = self._m3_fc_norm_w + if ( + fc_norm_w is not None + and fc_norm_w[0] is draft_model.w["fc_norm0"] + and len(spec_metadata.layers_to_capture) == fc_norm_w[1].shape[0] + ): + self._m3_capture_norm = (fc_norm_w[1], draft_model.rms_eps) + + @staticmethod + def _m3_sampled_verify_ok(spec_metadata) -> bool: + """Whether sampled verify rows can use the fused sampled op: only under advanced_sampling_mode=no_topk, the + deploy-time guarantee that no row uses top-k or min-p (that op applies neither).""" + mode = getattr(spec_metadata, "advanced_sampling_mode", None) + return mode is not None and mode.skips_top_k and not mode.skips_top_p + + def _m3_target_accept( + self, + hidden_states, + spec_metadata, + batch_size: int, + draft_len: int, + draft_model: Optional[M3Eagle3Draft] = None, + ): + """Fused verification (see tail_eligible); same outputs as the strict acceptance path.""" + force_total = 0 + force_vec = None + if self.force_num_accepted_tokens != 0.0: + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if not torch.cuda.is_current_stream_capturing(): + # The draw pool is built by an H2D copy and the draw kernel JIT-compiles per block width, so both + # happen here, before the first captured step (a scratch counter keeps the real draw sequence). + self._ensure_force_accept_rng_state(hidden_states.device) + if not self._m3_force_vec_warm: + scratch = torch.zeros(1, dtype=torch.int64, device=hidden_states.device) + dummy = torch.empty(64, dtype=torch.int32, device=hidden_states.device) + for block in (16, 32, 64): + m3_force_vec( + self._force_accept_rng_pool, + scratch, + dummy, + block, + 1, + 2, + 0.5, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + self._m3_force_vec_warm = True + if not is_warmup: + force_total = min(int(self.force_num_accepted_tokens) + 1, draft_len + 1) + force_vec = self._m3_force_vec( + batch_size, force_total, draft_len + 1, hidden_states.device + ) + draft_tokens = spec_metadata.draft_tokens.view(batch_size, draft_len) + if draft_tokens.dtype != torch.int32: + draft_tokens = draft_tokens.to(torch.int32) + # Column 0 receives the last accepted token, the next step's first input (the fused draft loop adds the rest). + self._m3_next_new = torch.empty( + (batch_size, draft_len + 1), dtype=torch.int32, device=hidden_states.device + ) + rows = hidden_states[: batch_size * (draft_len + 1)] + lm_head = self._m3_target_lm_head + batch_idx = spec_metadata.batch_indices_cuda[:batch_size] + head = self._m3_coded_head(lm_head, rows.shape[0]) + # The Eagle3 fc on the aux stream: the fork made at the last captured layer, else one beside the head. A + # sampled verify's coded head holds every SM while it streams: its fc is forked after it instead (overlapping + # the sampled acceptance), where it does not delay the head's start. + fork = self._m3_fc_early_stream + self._m3_fc_early_stream = None + head_first = fork is None and head is not None and not spec_metadata.is_all_greedy_sample + if fork is None and not head_first: + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model) + # This rank's vocabulary shard of the target logits: the head's own local GEMM (or its coded copy), whose + # all-gather the fused kernels replace (they exchange per-row records across ranks, not logits). + if head is not None: + logits = m3_target_head_kernel.logits(head, rows) + else: + logits = lm_head.apply_linear(rows, None) + head_done = None + if head_first: + head_done = torch.cuda.Event() + head_done.record() + vocab_offset = self.mapping.tp_rank * lm_head.weight.shape[0] + if spec_metadata.is_all_greedy_sample: + result = m3_shard_accept_kernel.shard_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + path = SHARD_ACCEPT.path + else: + n_rows = rows.shape[0] + seeds, offsets = self._rng_state_per_token(spec_metadata, n_rows) + result = m3_sampled_accept_kernel.sampled_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + spec_metadata.temperatures[:n_rows], + spec_metadata.top_ks[:n_rows], + spec_metadata.top_ps[:n_rows], + None, + seeds, + offsets, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + path = "spec_tail" + if head_done is not None: + # Enqueued after the sampled acceptance (which starts first) but waits on the head only. + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model, after=head_done) + if fork is not None: + self._m3_fc_join(fork) + if result is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, None, "the fused verify's exchange workspace") + raise RuntimeError( + "M3 fused verify unavailable after tail_eligible() accepted the step" + ) + return result + + def prebuild_target_head(self, lm_head: nn.Module) -> bool: + """Build the 12-bit coded copy of the target head shard the fused verify streams at <= TARGET_HEAD_MAX_ROWS + rows (rebuilt from the current weight values on every call), on every rank. True when the coded copy decodes + to the BF16 weight bit for bit; otherwise (or for a head the fused verify does not read) nothing is kept and + the head's own GEMM runs.""" + self._m3_target_head = None + if not _sharded_head(lm_head) or not m3_target_head_kernel.head_supported(lm_head): + return False + head = m3_target_head_kernel.prepare_head(lm_head) + self._m3_target_head = head + if not head.coded: + logger.warning( + f"[M3 target head] the coded shard {tuple(lm_head.weight.shape)} does not decode to the BF16 weight " + "bit for bit; the fused verify keeps the head's BF16 GEMM" + ) + return head.coded + + def _m3_coded_head(self, lm_head: nn.Module, rows: int): + """The coded target head for a fused verify of ``rows`` rows, or None for the head's own GEMM: above + ``TARGET_HEAD_MAX_ROWS`` rows (uncovered), on a shard whose coding failed its round trip (uncovered: the GEMM + is the exact path), and when the copy was not built for this weight (a miss).""" + path = TARGET_HEAD.path + if not TARGET_HEAD.covers(rows): + uncovered(path, f"{rows} verify rows (the head's own GEMM runs)") + return None + head = self._m3_target_head + if head is None or head.weight is not lm_head.weight: + missed( + path, + rows, + "the coded target head is not built for the head's current weight (prebuild)", + ) + return None + if not head.coded: + uncovered( + path, "the coded shard failed its bitwise round trip (the head's own GEMM runs)" + ) + return None + return head + + def _m3_force_vec( + self, batch_size: int, base_total: int, max_total: int, device + ) -> Optional[torch.Tensor]: + """Per-request forced accepted counts for a fractional forced acceptance, else None. + + Draws exactly as _apply_force_accepted_tokens does for a decode-only batch (same fixed-seed pool, one counter + bump per step, the same (counter, slot) hash), so every TP rank and the generic path agree on the counts. + """ + frac = self.force_num_accepted_tokens - int(self.force_num_accepted_tokens) + if frac <= 0.0 or base_total >= max_total: + return None + self._ensure_force_accept_rng_state(device) + out = torch.empty(batch_size, dtype=torch.int32, device=device) + return m3_force_vec( + self._force_accept_rng_pool, + self._force_accept_rng_counter, + out, + batch_size, + base_total, + max_total, + float(frac), + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + max(16, 1 << (batch_size - 1).bit_length()), + ) + + # ---------------------------------------------------------------------------------------------- fused draft loop + + def prebuild_draft_head(self, draft_model: Optional[M3Eagle3Draft]) -> bool: + """Build the streamed copy of the drafter's BF16 head shard that the fused draft loop reads (rebuilt from the + current weight values on every call), on every rank. False when the drafter is not one the loop serves.""" + self._m3_draft_head = None + if draft_model is None or not self._m3_drafter_supported(draft_model): + return False + self._m3_draft_head = m3_draft_step_kernel.prepare_head( + draft_model.lm_head, self.mapping.tp_rank, DRAFT_HEAD_CODED + ) + return True + + def _m3_drafter_supported(self, draft_model: M3Eagle3Draft) -> bool: + """The drafter the fused draft loop runs (``M3Eagle3Draft`` validates its shape at construction) over this + worker's tensor-parallel group with a bias-free vocab-sharded BF16 head without padding.""" + if self.mapping is None or self.mapping.tp_size <= 1 or self.mapping.enable_attention_dp: + return False + if self._d2t is not None: + return False + lm_head = draft_model.lm_head + return _sharded_head(lm_head) and m3_draft_step_kernel.head_supported(lm_head) + + def _m3_fused_loop_ok( + self, attn_metadata, spec_metadata, batch_size: int, draft_model: Optional[M3Eagle3Draft] + ) -> bool: + """Whether this step drafts with the fused loop: a decode-only step the fused verify or the seed prepared + (``_m3_first_gather_ids``) with the fused first draft inputs, a supported drafter, a (draft length, requests) + ``DRAFT_STEP`` covers, the prebuilt head and the exchange workspace. Host state only, so the TP ranks + agree.""" + if ( + self._m3_first_gather_ids is None + or not self._m3_fused_first_inputs + or draft_model is None + ): + return False + path = DRAFT_STEP.path + if not self._m3_drafter_supported(draft_model) or not hasattr( + attn_metadata, "kv_lens_cuda" + ): + uncovered( + path, "the drafter or its attention metadata is not one the fused loop serves" + ) + return False + draft_len = spec_metadata.runtime_draft_len + if not DRAFT_STEP.covers((draft_len, batch_size)): + uncovered(path, f"draft length {draft_len} at {batch_size} requests") + return False + head = self._m3_draft_head + if head is None or head.weight is not draft_model.lm_head.weight: + missed( + path, + batch_size, + "the drafter head's streamed copy is not built for its current weight (prebuild)", + ) + return False + if m3_draft_step_kernel.workspace(self.mapping, head.wf.device) is None: + capture_missed(path, None, "the draft-step exchange workspace") + return False + return True + + def _forward_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ): + """The fused draft loop on a step it covers (``_m3_fused_loop_ok``), else the in-tree loop.""" + if self._m3_fused_loop_ok(attn_metadata, spec_metadata, batch_size, draft_model): + return self._m3_fused_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ) + return super()._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + + def _m3_fused_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model: M3Eagle3Draft, + draft_kv_cache_manager, + batch_size: int, + num_accepted_tokens, + ): + """Mirrors Eagle3OneModelWorker._forward_linear_draft_loop for a decode-only Eagle3 step, with each pass's + glue fused (see ``_m3_fused_loop_ok``): + + * the draft input -- the embedding lookup, the layer's input norm and hidden norm and their concat -- in one + launch (``m3_spec_tail.draft_input``; the second pass gathers the last accepted rows there); + * the drafter layer's attention and MLP, each followed by the all-reduce with the residual add and the next + RMSNorm (the post-attention norm, then the drafter's final norm); + * one kernel for the BF16 head over the head rows (pass 0: the last accepted row of each request), the + cross-rank argmax and the bookkeeping the in-tree loop does between passes: the draft token, the next pass's + input id and position, the draft KV lengths (and pass 0's sequence lengths) (``m3_draft_step``). + + The host metadata transitions are the in-tree loop's. The drafts land in columns 1.. of the step's next-input + tokens (column 0 is the last accepted token), which ``_forward_impl`` returns as they lie. + """ + draft_len = spec_metadata.runtime_draft_len + table = draft_model.embed_table + dev = table.device + head = self._m3_draft_head + gather_ids = self._m3_first_gather_ids + ids = inputs["input_ids"] + hidden = inputs["hidden_states"] + pos = inputs["position_ids"] + if ids.dtype != torch.int32: + ids = ids.to(torch.int32) + if pos.dtype != torch.int32: + pos = pos.to(torch.int32) + num_acc = num_accepted_tokens[:batch_size] + if num_acc.dtype != torch.int32: + num_acc = num_acc.to(torch.int32) + next_new = self._m3_next_new + if next_new is None: + # Seeded from the in-tree acceptance: the last accepted token is the first-pass input at the gather row. + next_new = torch.empty((batch_size, draft_len + 1), dtype=torch.int32, device=dev) + next_new[:, 0] = ids[gather_ids] + draft_tokens = next_new[:, 1:] + w = draft_model.w + eps = draft_model.rms_eps + + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager) as md: + for i in range(draft_len): + x, residual = m3_spec_tail.draft_input( + ids, + table, + hidden, + gather_ids if i == 1 else None, + w["in_norm"], + eps, + w["hidden_norm"], + eps, + ) + h = draft_model.attention(x, md, reduce=False) + h, residual = draft_model.ar_norm(h, residual, "post_norm") + h = draft_model.mlp(h, reduce=False) + h, to_save = draft_model.ar_norm(h, residual, "final_norm") + if draft_model.return_hidden_post_norm: + to_save = h + ids_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + pos_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + if not m3_draft_step_kernel.draft_step( + head, + h, + batch_size, + gather_ids if i == 0 else None, + self.mapping, + draft_tokens, + i, + ids_next, + pos, + pos_next, + md.kv_lens_cuda, + num_acc if i == 0 else None, + md._seq_lens_cuda if i == 0 else None, + draft_len, + ): + raise RuntimeError( + "M3 fused draft step unavailable after _m3_fused_loop_ok() accepted the step" + ) + ids, pos, hidden = ids_next, pos_next, to_save + if i == 0: + # The draft step set the device sequence lengths to 1 and rolled the KV lengths back. + md._seq_lens[:batch_size].fill_(1) + md.on_update() + has_kv_cache = md.kv_cache_manager is not None + if has_kv_cache and hasattr(md, "host_request_types"): + md.host_request_types[: md.num_contexts].fill_(1) + if has_kv_cache: + md.num_contexts = 0 + self._prepare_flash_mla_generation_layout(md, 0, batch_size) + md.update_for_spec_dec() + md.use_spec_decoding = False + else: + md.update_for_spec_dec() + self._m3_next_new_tokens = next_new + return draft_tokens + + +def _build_speculation(shell, model_config: ModelConfig): + """The drafter config, the drafter and the worker of one-model Eagle3 (the in-tree one-engine construction for + the llama3 drafter architecture).""" + spec = model_config.spec_config + if ( + not spec.spec_dec_mode.use_one_engine() + or not spec.spec_dec_mode.is_eagle3_one_model() + or spec.eagle3_model_arch != "llama3" + or getattr(spec, "use_dynamic_tree", False) + ): + raise NoTarget( + "this MiniMax-M3 target speculates with linear one-model Eagle3 (llama3 drafter) only" + ) + draft_config = ModelConfig.from_pretrained( + spec.speculative_model, + trust_remote_code=True, + attn_backend=model_config.attn_backend, + moe_backend=_get_requested_draft_moe_backend(model_config, spec), + mapping=model_config.mapping, + spec_config=model_config.spec_config, + max_num_tokens=model_config.max_num_tokens, + moe_max_num_tokens=model_config.moe_max_num_tokens, + ) + _set_draft_kv_cache_quant_algo(draft_config, model_config) + draft_config.extra_attrs = model_config.extra_attrs + use_separate = should_use_separate_draft_kv_cache(spec) + draft = M3Eagle3Draft( + shell.model, + draft_config, + model_config.pretrained_config.num_hidden_layers, + spec.num_capture_layers, + ) + worker = M3Eagle3Worker(spec, model_config.mapping, use_separate_draft_kv_cache=use_separate) + worker.set_draft_model(draft) + shell.epilogue.append(draft) + shell.epilogue.append(worker) + return draft_config, draft, worker, use_separate + + +# -------------------------------------------------------------------------------------------------------------------- +# Runtime hooks and the shell +# -------------------------------------------------------------------------------------------------------------------- + + +def _runtime_hooks(spec_fold: Optional[M3SpecStepFold]) -> RuntimeHooks: + """The step-loop hooks this target opts into, all measured on GB300 tp4 AgentX decode. + + * ``pack_step_inputs``: a decode step's small input copies go out as one upload and one scatter at the graph + replay instead of a chain of copies the next graph waits on (~30 us of idle GPU at one request). + * ``skip_unchanged_h2d``: steady decode re-uploads slot indices, query lengths and page tables that rarely change. + * ``loop_on_execution_stream``: the overlap loop's own stream work and the forward on one stream, so step N+1's + inputs wait on graph N in program order rather than through two cross-stream event hops. + * ``merge_tuned_tactics_across_ranks``: tensor-parallel peers capture the same GEMM tactics (rank-local profiling + noise otherwise splits them, and the slowest rank sets every step). + * ``skip_greedy_rng_upload``: all-greedy speculative batches never read the Philox seeds. + * ``spec_step_fold``: the spec sampler's slot stores and the engine's next-step input gathers in one launch each. + * MSA: identity selections for short sequences (and short/long graph families), the one-launch KV-length update, + unchanged block-table copies skipped, pool views built once. + """ + return RuntimeHooks( + pack_step_inputs=True, + skip_unchanged_h2d=True, + loop_on_execution_stream=True, + merge_tuned_tactics_across_ranks=True, + skip_greedy_rng_upload=True, + spec_step_fold=spec_fold, + msa=MsaRuntimePolicy( + short_sequence_identity=True, + kv_update=_msa_kv_update_kernel.msa_kv_update, + skip_unchanged_block_offsets=True, + reuse_pool_views=True, + ), + ) + + +class ModelingV2MiniMaxM3Nvfp4Sm103Tp4( + DecoderModelForCausalLM[MiniMaxM3Tp4ModelingV2, PretrainedConfig] +): + """MiniMax-M3 NVFP4 / sm_103 / tp4, text only, with one-model Eagle3 or no speculation. + + The shell owns the vocab-parallel BF16 ``lm_head`` and the logits processor (``DecoderModelForCausalLM``); with + speculation, the drafter (``M3Eagle3Draft``, its weights from the speculative checkpoint through + ``load_draft_weights``) and the worker (``M3Eagle3Worker``) that verifies and drafts each step. + """ + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + assert torch.cuda.get_device_capability() == _SM, ( + f"this target is built for sm_{_SM[0]}{_SM[1]}, the device is sm_" + f"{''.join(map(str, torch.cuda.get_device_capability()))}" + ) + # Under modeling_v2='require' a covered decode shape that cannot take its fast path raises. + set_require(resolved_mode(type(self)) is ModelingV2Mode.REQUIRE) + model_config = _text_model_config(model_config) + cfg = model_config.pretrained_config + spec_config = getattr(model_config, "spec_config", None) + if model_config.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4: + # The one-model Eagle3 layer has no NVFP4 trtllm-gen decode kernel: keep its modules and the shared + # draft cache on FP8 while the target stays NVFP4. + model_config.extra_attrs["draft_kv_cache_quant_algo_override"] = QuantAlgo.FP8 + super().__init__( + MiniMaxM3Tp4ModelingV2(model_config), + config=model_config, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + self.spec_config = spec_config + self.draft_model: Optional[M3Eagle3Draft] = None + self.draft_config: Optional[ModelConfig] = None + self.spec_worker: Optional[M3Eagle3Worker] = None + self.use_separate_draft_kv_cache = False + self.layer_idx = FINAL_HIDDEN_LAYER_ID + spec_fold = None + if spec_config is not None: + ( + self.draft_config, + self.draft_model, + self.spec_worker, + self.use_separate_draft_kv_cache, + ) = _build_speculation(self, model_config) + spec_fold = M3SpecStepFold() + self.runtime_hooks = _runtime_hooks(spec_fold) + + def load_weights(self, weights: Dict, *args, **kwargs) -> None: + _weights.MODEL_WEIGHTS.load(self, weights) + + def load_draft_weights(self, weights: Dict, weight_mapper=None) -> None: + """The drafter's checkpoint (``speculative_model``), loaded by the engine after the target's.""" + _weights.DRAFT_WEIGHTS.load(self.draft_model, weights) + + def set_guided_decoder(self, guided_decoder) -> bool: + """Mirrors ``SpecDecOneEngineForCausalLM.set_guided_decoder``: guided decoding with speculation is installed + on the worker.""" + if hasattr(self.spec_worker, "set_guided_decoder"): + return self.spec_worker.set_guided_decoder(guided_decoder) + return False + + def cache_derived_state(self) -> None: + """Build the targets and every decode kernel once the weights are real, before the executor's warmup; with + speculation also the fused verify's coded target head and the fused draft loop's head copy.""" + super().cache_derived_state() + self.model.post_load_weights() + if self.spec_worker is not None: + self.spec_worker.prebuild(self.lm_head, self.draft_model) + + def post_load_weights(self) -> None: + super().post_load_weights() + self.model.post_load_weights() + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + return_context_logits: bool = False, + spec_metadata: Optional[SpecMetadata] = None, + resource_manager: Optional[Any] = None, + **kwargs, + ) -> Any: + """One engine step: the target forward, then the logits (or, with speculation, the worker's verify and + draft). Text only.""" + multimodal_params = kwargs.pop("multimodal_params", None) or [] + if any(p is not None and getattr(p, "multimodal_data", None) for p in multimodal_params): + raise ValueError( + "this MiniMax-M3 target is text only; the request carries multimodal input" + ) + worker = self.spec_worker + capture = _plain_capture + if worker is not None: + capture = worker.capture_hidden_states + if spec_metadata is not None: + worker.arm_capture_norm( + attn_metadata, spec_metadata, self.lm_head, self.draft_model + ) + hidden_states = self.model( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + capture=capture, + **kwargs, + ) + if getattr(attn_metadata, "csa2_remote_tail_mode", None) == "source": + # A context worker stopped at the encoder/decoder boundary: a disposable token for the context-only + # completion, and no draft worker. + return self.logits_processor.forward(hidden_states, self.lm_head, attn_metadata, False) + if spec_metadata is not None and spec_metadata.is_layer_capture(FINAL_HIDDEN_LAYER_ID): + capture(spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None) + if attn_metadata.padded_num_tokens is not None: + hidden_states = hidden_states[: attn_metadata.num_tokens] + if worker is None: + return self.logits_processor.forward( + hidden_states, self.lm_head, attn_metadata, return_context_logits + ) + # The fused verify computes this rank's logits shard itself and needs no gathered logits. + if worker.tail_eligible(attn_metadata, spec_metadata, self.lm_head): + logits = None + else: + logits = self.logits_processor.forward( + hidden_states[spec_metadata.gather_ids], self.lm_head, attn_metadata, True + ) + spec_input_ids = input_ids if input_ids is not None else kwargs.get("orig_input_ids") + spec_position_ids = position_ids + if attn_metadata.padded_num_tokens is not None: + if spec_input_ids is not None: + spec_input_ids = spec_input_ids[: attn_metadata.num_tokens] + if position_ids is not None: + spec_position_ids = _slice_spec_position_ids(position_ids, attn_metadata.num_tokens) + return worker( + input_ids=spec_input_ids, + position_ids=spec_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, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/__init__.py new file mode 100644 index 000000000000..1059445198db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this target. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/activation/__init__.py new file mode 100644 index 000000000000..b8c7426e3933 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/activation/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Activation ops of this target: the gated MLPs' fused SwiGLU.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/activation/m3_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/activation/m3_silu_and_mul.py new file mode 100644 index 000000000000..64cb00558760 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/activation/m3_silu_and_mul.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The gated MLPs' fused SwiGLU: ``silu_and_mul`` with MiniMax-M3's SwiGLU-OAI gain, offset and clamp.""" + +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 M3SiluAndMul(OpWrapper): + """``out = g * sigmoid(alpha * g) * (u + beta)`` over ``x = [g | u]`` (bf16 [M, 2I] -> [M, I]), with + ``g = min(gate, limit)`` and ``u = clamp(up, -limit, limit)``; ``limit``, ``alpha`` and ``beta`` None give plain + SwiGLU ``silu(gate) * up``. + + The MXFP8 dense MLPs (layers 0-2, intermediate 3072 per rank) and the MoE's shared expert (intermediate 768 per + rank) run it at MiniMax-M3's (limit 7, alpha 1.702, beta 1); the BF16 Eagle3 drafter's MLP (intermediate 4608 per + rank) runs it plain. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert, one decode token", spec=dict(m=1, inter=768, oai=True)), + Cell(why="shared expert, a 48-row decode verify", spec=dict(m=48, inter=768, oai=True)), + Cell(why="dense MLP, a prefill chunk", spec=dict(m=2048, inter=3072, oai=True)), + Cell(why="dense MLP, one decode token", spec=dict(m=1, inter=3072, oai=True)), + Cell( + why="drafter MLP (plain SwiGLU), eight decode rows", + spec=dict(m=8, inter=4608, oai=False), + ), + Cell( + why="drafter MLP (plain SwiGLU), a prefill chunk", + spec=dict(m=1024, inter=4608, oai=False), + ), + ) + + note = """ + The kernel computes in fp32 and rounds once to bf16; `reference` does the same arithmetic in fp32, so the two + differ only by the fast sigmoid's last bits (the test's band). + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + + def reference( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + gate, up = x.float().chunk(2, dim=-1) + if limit is not None: + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + a = 1.0 if alpha is None else alpha + b = 0.0 if beta is None else beta + return (gate * torch.sigmoid(a * gate) * (up + b)).to(x.dtype) + + def is_valid( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> None: + assert x.stride(-1) == 1, "the gate/up halves are read from contiguous rows" + + +m3_silu_and_mul = M3SiluAndMul() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/__init__.py new file mode 100644 index 000000000000..a7f7fee7a234 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Attention ops of this target: MSA block selection, the sparse decode and the KV-length update.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_attn_c1.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_attn_c1.py new file mode 100644 index 000000000000..625a7d5d4edc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_attn_c1.py @@ -0,0 +1,280 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The c1 kernel: fused MSA decode attention + o_proj GEMV in one 96-CTA launch, for a one-request, 6-token +verify layer.""" + +from __future__ import annotations + +from typing import NamedTuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_attn_mega import HIDDEN, check_handoff_at_rest, workspace_parity +from .m3_fused_decode_attn import ( + HEAD_DIM, + MXFP8_BLOCK, + NUM_HEADS, + check_attention_inputs, + mx_gemm, + mxfp8_rows, + reference_attention, + unswizzle_scales, +) + +#: 6 clusters of 16 CTAs, each an attention CTA and an o_proj unit: ``flags[32 + CTA]`` are the call counters. +GRID = 96 + + +class C1State(NamedTuple): + """The c1 kernel's results: the attention's bf16 output, the caches after the insert, and the o_proj output.""" + + output: torch.Tensor + kv_data: torch.Tensor + kv_scale: torch.Tensor + idx_k: torch.Tensor + y: torch.Tensor + + +class M3AttnC1(OpWrapper): + """The attention of a one-request verify step of 6 query tokens and its o_proj, in one 96-CTA launch. + + The fused attention of ``m3_fused_decode_attn`` on the packed QKV rows (bf16 ``output``, MXFP8 hand-off into the + call's parity of ``mxfp8_data`` / ``mxfp8_scale``), and the MXFP8 o_proj GEMV into ``y``: every attention CTA + is also an o_proj unit (48 row tiles x 2 K splits). ``flags[32 + CTA]`` are the per-CTA call counters; the + workspace is allocated once and never reset. With ``ar_mode`` 2, ``y`` is also pushed into the next multicast + all-reduce's exchange at ``ar_workspace`` / ``ar_flags``. + + Returns None; the cells read ``output``, ``y``, the hand-off and the caches. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="Eagle3 verify of 5 drafts (qlen 6) a few blocks in; both call parities (two calls, one workspace)", + spec=dict(seq_len=300, calls=2), + ), + Cell( + why="the full 16-block window: the six tokens fill the last rows of the last page", + spec=dict(seq_len=2048, calls=1), + ), + Cell( + why="the six tokens straddle a page boundary (positions 126..131)", + spec=dict(seq_len=132, calls=1), + ), + Cell(why="the first six tokens of a request", spec=dict(seq_len=6, calls=1)), + ) + + note = """ + Compiled for the one shape the target launches: one request of 6 query tokens (geometry (6, 1, 1), 96 CTAs), + the MXFP8 hand-off. The all-reduce push (`ar_mode` 2) needs a multicast workspace and is not driven here: the + cells run `ar_mode` 0, where `y` is the o_proj output and nothing else is written. + """ + + def __call__( + self, + packed: torch.Tensor, + flags: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + y: torch.Tensor, + ar_workspace: int, + ar_flags: int, + ar_rank: int, + ar_world: int, + ar_max_tokens: int, + ar_mode: int, + ) -> None: + torch.ops.trtllm.minimax_m3_attn_c1( + packed, + flags, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + eps, + sm_scale, + decode_query_len, + tokens_per_cluster, + blocks_per_cta, + output, + mxfp8_data, + mxfp8_scale, + use_pdl, + o_weight, + o_weight_scale, + y, + ar_workspace, + ar_flags, + ar_rank, + ar_world, + ar_max_tokens, + ar_mode, + ) + + def reference( + self, + packed: torch.Tensor, + flags: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + y: torch.Tensor, + ar_workspace: int, + ar_flags: int, + ar_rank: int, + ar_world: int, + ar_max_tokens: int, + ar_mode: int, + ) -> C1State: + """The composition: the fused attention's reference on ``packed``, then + ``y = bf16(mxfp8(output) @ dequant(W_o)^T)`` (fp64 accumulation, one rounding). + + The workspace, the output buffers, the geometry, ``use_pdl`` and the all-reduce arguments are ignored. + """ + attn = reference_attention( + packed, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + eps, + sm_scale, + decode_query_len, + ) + q8, ex = mxfp8_rows(attn.output) + width = NUM_HEADS * HEAD_DIM + out_y = mx_gemm( + q8, + ex, + o_weight.view(torch.uint8), + unswizzle_scales(o_weight_scale, HIDDEN, width // MXFP8_BLOCK), + ) + return C1State(attn.output, attn.kv_data, attn.kv_scale, attn.idx_k, out_y) + + def is_valid( + self, + packed: torch.Tensor, + flags: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + y: torch.Tensor, + ar_workspace: int, + ar_flags: int, + ar_rank: int, + ar_world: int, + ar_max_tokens: int, + ar_mode: int, + ) -> None: + check_attention_inputs( + packed.shape[0], + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + decode_query_len, + ) + check_handoff_at_rest(mxfp8_data, mxfp8_scale, workspace_parity(flags, GRID)) + assert ar_mode != 2 or (ar_workspace != 0 and ar_flags != 0), ( + "ar_mode 2 needs the exchange's addresses" + ) + + +m3_attn_c1 = M3AttnC1() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_attn_mega.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_attn_mega.py new file mode 100644 index 000000000000..3445c8c65cbd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_attn_mega.py @@ -0,0 +1,340 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The attention megakernel: QKV GEMV + fused MSA decode attention + o_proj GEMV in one launch, for a one-request, +one-token decode layer.""" + +from __future__ import annotations + +from typing import NamedTuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_fused_decode_attn import ( + HEAD_DIM, + MXFP8_BLOCK, + NUM_HEADS, + PACKED_WIDTH, + check_attention_inputs, + mx_gemm, + mxfp8_rows, + reference_attention, + unswizzle_scales, +) + +#: The layer input (hidden size) and the o_proj output width at tp4. +HIDDEN = 6144 +#: Persistent grid of the megakernel (7 clusters of 16): ``flags[32 + CTA]`` are its per-CTA call counters. +GRID = 112 +#: Rows of each call-parity buffer. +MAX_TOKENS = 8 + + +class MegaState(NamedTuple): + """The megakernel's results: the packed QKV rows of the call, the attention's bf16 output, the caches after the + insert, and the o_proj output.""" + + packed: torch.Tensor + output: torch.Tensor + kv_data: torch.Tensor + kv_scale: torch.Tensor + idx_k: torch.Tensor + y: torch.Tensor + + +def workspace_parity(flags: torch.Tensor, grid: int) -> int: + """The call parity every CTA of the grid will use (their counters move together), checked.""" + counters = flags[32 : 32 + grid] + assert bool((counters == counters[0]).all()), ( + "the per-CTA call counters disagree: the workspace is corrupt" + ) + return int(counters[0]) & 1 + + +def check_handoff_at_rest(mxfp8_data: torch.Tensor, mxfp8_scale: torch.Tensor, parity: int) -> None: + """The MXFP8 hand-off's parity buffer this call writes must hold the 0xFF "not yet" marker (Lamport).""" + assert bool((mxfp8_data.view(torch.uint8).reshape(2, -1)[parity] == 0xFF).all()) and bool( + (mxfp8_scale.view(torch.uint8).reshape(2, -1)[parity] == 0xFF).all() + ), ( + "the MXFP8 hand-off's buffer for this call is not at rest (0xFF): the o_proj units would read stale codes" + ) + + +class M3AttnMega(OpWrapper): + """The whole attention half of a decode layer for one request and one query token, in one launch. + + The MXFP8 QKV GEMV of the layer input (as the preceding all-reduce hands it over: E4M3 rows + 128x4-swizzled + UE8M0 scales) into the call's parity buffer of ``qkv_out``, the fused attention of ``m3_fused_decode_attn`` on + those rows (bf16 ``output``, MXFP8 hand-off into the call's parity of ``mxfp8_data`` / ``mxfp8_scale``), and the + MXFP8 o_proj GEMV into ``y``. ``flags`` holds the per-tile ready counters and the per-CTA call counters; the + workspace is allocated once and never reset. With ``ar_mode`` 2, ``y`` is also pushed into the next multicast + all-reduce's exchange at ``ar_workspace`` / ``ar_flags``. + + Returns None; the cells read ``qkv_out``, ``output``, ``y``, the hand-off and the caches. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode of one request a few blocks in; both call parities (two calls on one workspace)", + spec=dict(seq_len=300, calls=2), + ), + Cell( + why="the full 16-block window, the new token on the last row of the last page", + spec=dict(seq_len=2048, calls=1), + ), + Cell( + why="the new token opens a page (position 1024): its own page's other rows are past the length", + spec=dict(seq_len=1025, calls=1), + ), + Cell(why="the first token of a request: one attended row", spec=dict(seq_len=1, calls=1)), + ) + + note = """ + Compiled for the one shape the target launches: one request, one query token (geometry (1, 1, 1)), the MXFP8 + hand-off. The all-reduce push (`ar_mode` 2) needs a multicast workspace and is not driven here: the cells run + `ar_mode` 0, where `y` is the o_proj output and nothing else is written. + + The QKV rows are handed from the GEMV CTAs to the attention CTAs by polling for the 0xFFFF marker the + previous call left: a GEMV output whose bf16 bits are 0xFFFF (a NaN) would never be seen as written. Finite + inputs cannot produce it. + """ + + def __call__( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + qkv_weight: torch.Tensor, + qkv_weight_scale: torch.Tensor, + qkv_out: torch.Tensor, + flags: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + y: torch.Tensor, + ar_workspace: int, + ar_flags: int, + ar_rank: int, + ar_world: int, + ar_max_tokens: int, + ar_mode: int, + ) -> None: + torch.ops.trtllm.minimax_m3_attn_mega( + xq, + xsf, + qkv_weight, + qkv_weight_scale, + qkv_out, + flags, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + eps, + sm_scale, + decode_query_len, + tokens_per_cluster, + blocks_per_cta, + output, + mxfp8_data, + mxfp8_scale, + use_pdl, + o_weight, + o_weight_scale, + y, + ar_workspace, + ar_flags, + ar_rank, + ar_world, + ar_max_tokens, + ar_mode, + ) + + def reference( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + qkv_weight: torch.Tensor, + qkv_weight_scale: torch.Tensor, + qkv_out: torch.Tensor, + flags: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + y: torch.Tensor, + ar_workspace: int, + ar_flags: int, + ar_rank: int, + ar_world: int, + ar_max_tokens: int, + ar_mode: int, + ) -> MegaState: + """The composition: ``packed = bf16(dequant(xq) @ dequant(W_qkv)^T)``, the fused attention's reference on + those rows, then ``y = bf16(mxfp8(output) @ dequant(W_o)^T)`` (fp64 accumulation, one rounding per GEMV). + + The workspace, the output buffers, the geometry, ``use_pdl`` and the all-reduce arguments are ignored. + """ + t = xq.shape[0] + xq_b, w_b = xq.view(torch.uint8), qkv_weight.view(torch.uint8) + packed = mx_gemm( + xq_b, + unswizzle_scales(xsf, t, HIDDEN // MXFP8_BLOCK), + w_b, + unswizzle_scales(qkv_weight_scale, PACKED_WIDTH, HIDDEN // MXFP8_BLOCK), + ) + attn = reference_attention( + packed, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + eps, + sm_scale, + decode_query_len, + ) + q8, ex = mxfp8_rows(attn.output) + width = NUM_HEADS * HEAD_DIM + out_y = mx_gemm( + q8, + ex, + o_weight.view(torch.uint8), + unswizzle_scales(o_weight_scale, HIDDEN, width // MXFP8_BLOCK), + ) + return MegaState(packed, attn.output, attn.kv_data, attn.kv_scale, attn.idx_k, out_y) + + def is_valid( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + qkv_weight: torch.Tensor, + qkv_weight_scale: torch.Tensor, + qkv_out: torch.Tensor, + flags: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + y: torch.Tensor, + ar_workspace: int, + ar_flags: int, + ar_rank: int, + ar_world: int, + ar_max_tokens: int, + ar_mode: int, + ) -> None: + check_attention_inputs( + xq.shape[0], + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + decode_query_len, + ) + assert xsf.numel() >= 128 * (HIDDEN // MXFP8_BLOCK), ( + "xsf must hold one 128-row tile of swizzled scales" + ) + parity = workspace_parity(flags, GRID) + assert bool((qkv_out.view(torch.int16)[parity] == -1).all()), ( + "qkv_out's buffer for this call is not at rest (0xFFFF): the attention would read stale QKV rows" + ) + check_handoff_at_rest(mxfp8_data, mxfp8_scale, parity) + assert ar_mode != 2 or (ar_workspace != 0 and ar_flags != 0), ( + "ar_mode 2 needs the exchange's addresses" + ) + + +m3_attn_mega = M3AttnMega() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_dense_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_dense_decode.py new file mode 100644 index 000000000000..bf18485cb406 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_dense_decode.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' decode attention over the full page table with valid sequence-length bounds.""" + +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.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( # noqa: E501 + minimax_m3_trtllm_gen_dense_decode, +) + +__all__ = ["m3_dense_decode"] + +#: The TP4 dense layers' geometry: 16 query heads and one KV head of dim 128 per rank, an FP8 KV pool. +HEADS = 16 +KV_HEADS = 1 +HEAD_DIM = 128 +#: Valid upper bounds used to compare short-context and model-maximum dispatch. +SHORT_BOUND = 16384 +MODEL_MAX = 1 << 20 + + +def dense_attention_reference( + q: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, +) -> torch.Tensor: + """fp32 causal decode attention: query row ``b * qlen + i`` of request b attends its first + ``seq_lens[b] - qlen + i + 1`` keys. ``keys`` / ``values``: [batch, kv_heads, max_len, head_dim] (the request's + tokens in order). Returns fp32 [total_q, heads, head_dim].""" + batch, kv_heads, max_len, _ = keys.shape + heads = q.shape[1] + group = heads // kv_heads + out = torch.zeros(q.shape, dtype=torch.float32, device=q.device) + pos = torch.arange(max_len, device=q.device) + with true_fp32_matmul(): + for b in range(batch): + full = int(seq_lens[b]) + rows = q[b * decode_query_len : (b + 1) * decode_query_len].float() # [qlen, heads, d] + limit = full - decode_query_len + 1 + torch.arange(decode_query_len, device=q.device) + valid = pos[None, :] < limit[:, None] # [qlen, max_len] + for h in range(kv_heads): + qh = rows[:, h * group : (h + 1) * group] # [qlen, group, d] + logits = torch.einsum("qgd,kd->qgk", qh, keys[b, h].float()) * sm_scale + logits = logits.masked_fill(~valid[:, None, :], float("-inf")) + probs = torch.softmax(logits, dim=-1) + out[ + b * decode_query_len : (b + 1) * decode_query_len, h * group : (h + 1) * group + ] = torch.einsum("qgk,kd->qgd", probs, values[b, h].float()) + return out + + +def gather_paged_kv( + kv_cache_manager, layer_idx: int, block_table: torch.Tensor, seq_lens: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """The K and V of each request in token order, [batch, kv_heads, max_len, head_dim], from the flat sub-page pool + (``get_kv_subpage_pool``: K of slot s at sub-page s * subpages_per_slot, V one sub-page later).""" + pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + page = pool.shape[2] + pages = (int(seq_lens.max()) + page - 1) // page + slots = block_table[:, :pages].long() + k = pool[slots * subpages] # [batch, pages, kv_heads, page, d] + v = pool[slots * subpages + 1] + batch, _, kv_heads, _, d = k.shape + + def order(t): + return t.permute(0, 2, 1, 3, 4).reshape(batch, kv_heads, pages * page, d) + + return order(k), order(v) + + +class M3DenseDecode(OpWrapper): + """Decode over each request's full causal context in the paged FP8 pool. + + The sequence-length bound covers every request and the BF16 output is written in place. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over a 20000-token context, its full sequence-length bound", + spec=dict(kv_lens=(20000,), qlen=4, max_seq_len=20000, tight_bound=True), + ), + Cell( + why="one request, one row, over a 65536-token context with its full sequence-length bound", + spec=dict(kv_lens=(65536,), qlen=1, max_seq_len=65536, tight_bound=True), + ), + Cell( + why="one request verifying 1 + 3 rows over a 3000-token context: an upper bound above the length", + spec=dict(kv_lens=(3000,), qlen=4, max_seq_len=SHORT_BOUND, tight_bound=True), + ), + Cell( + why="two requests with the model maximum as the upper bound", + spec=dict(kv_lens=(20000, 5000), qlen=4, max_seq_len=MODEL_MAX, tight_bound=False), + ), + ) + + note = """ + Outputs at two valid sequence-length bounds agree to a few BF16 ulps and both match the fp32 + reference within the FP8-KV tolerance of the in-tree dense decode test. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + minimax_m3_trtllm_gen_dense_decode( + q, + kv_cache_manager, + layer_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + max_seq_len=max_seq_len, + max_num_requests=max_num_requests, + ) + return output + + def reference( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + """fp32 attention over the pool's K / V (q rounded to the pool's dtype first, as the kernel runs it). + + ``output``, ``max_seq_len`` and ``max_num_requests`` do not change the result.""" + keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) + if keys.dtype == torch.float8_e4m3fn: + q = q.to(torch.float8_e4m3fn) + return dense_attention_reference(q, keys, values, seq_lens, decode_query_len, sm_scale) + + def is_valid( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> None: + # The multi-CTA KV counters are sized for max_num_requests: more requests overrun them. + assert seq_lens.shape[0] <= max_num_requests, "more requests than the counters cover" + assert bool((seq_lens <= max_seq_len).all()), "sequence length exceeds its bound" + # Each verify row attends a causal prefix of at least one key. + assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" + + +m3_dense_decode = M3DenseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_draft_attention.py new file mode 100644 index 000000000000..6347348d1681 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_draft_attention.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's attention on the TRTLLM backend: RoPE, KV-cache append and causal attention.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + PredefinedAttentionMask, +) +from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType + + +class M3DraftAttention(OpWrapper): + """The drafter layer's attention over its fused QKV rows ``qkv`` [T, (q + 2 kv) heads x 128] (bf16): the + backend applies the drafter's RoPE (fused), appends K / V to the drafter layer's paged cache (FP8, the target's + NVFP4 cache keeps its draft layer in FP8) and runs causal attention over the step's live rows (``metadata`` + ``num_tokens``); returns [num_tokens, q heads x 128] bf16. + + ``backend`` (``TrtllmAttention``) and ``metadata`` are the runtime's objects for the drafter layer and step. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="first draft pass of a 4-request verify (4 x 4 rows)", spec=dict(reqs=4, qlen=4)), + Cell(why="a later draft pass, one row per request", spec=dict(reqs=8, qlen=1)), + Cell(why="a context chunk's draft rows", spec=dict(reqs=1, context=512)), + ) + + note = """ + `reference` is exact (float64) causal softmax attention: Q rotated by the backend's NeoX RoPE at each token's + position, over the drafter layer's K / V rows read back from its FP8 cache (dequantized by the backend's + ``kv_scale_quant_orig``). What remains between the two is the FMHA's own arithmetic (Q's rounding to the + kernel's input precision, its P rounding and accumulation order), which the test bounds. + """ + + def __call__(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + num_tokens = metadata.num_tokens + out = backend.forward( + qkv[:num_tokens, :], + None, + None, + metadata, + forward_args=AttentionForwardArgs(attention_mask=PredefinedAttentionMask.CAUSAL), + ) + return out[0] if isinstance(out, tuple) else out + + def reference(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + """The call's [num_tokens, q heads x 128] output, in ``qkv``'s dtype. + + Query token ``j`` of request ``r`` sits at position ``cached_r + j`` (the metadata's cached lengths) and + attends positions up to it. Its Q is rotated here; every K / V row, the step's own included, is read from + the drafter layer's paged cache at the slots the metadata's KV cache manager holds, so this is evaluated + after the call (whose roped, quantized append it then reads back; the test checks that append itself). + Mirrors one KV head per rank and un-scaled NeoX RoPE, the drafter's configuration. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + assert backend.num_kv_heads == 1, "mirrors the one-KV-head-per-rank layout only" + rope = backend.rope_params + assert backend.position_embedding_type == int(PositionEmbeddingType.rope_gpt_neox) + assert rope.scale_type == RotaryScalingType.none and rope.alpha == 1.0 + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + kv_scale = float(backend.kv_scale_quant_orig.reshape(-1)[0]) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = int(metadata.num_tokens) + assert sum(qlens) == total + device = qkv.device + + rot = int(rope.dim) + inv_freq = float(rope.theta) ** ( + -torch.arange(0, rot, 2, dtype=torch.float64, device=device) / rot + ) + pool = manager.get_buffers(layer, "HND") + q = qkv[:total, : heads * dim].to(torch.float64).reshape(total, heads, dim) + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][ + : (length + page - 1) // page + ], + dtype=torch.int64, + device=device, + ) + keys = pool[pages, 0, 0].to(torch.float64).reshape(-1, dim)[:length] * kv_scale + values = pool[pages, 1, 0].to(torch.float64).reshape(-1, dim)[:length] * kv_scale + qpos = cached[row] + torch.arange(n, device=device) + angles = qpos.to(torch.float64)[:, None] * inv_freq[None, :] + cos, sin = angles.cos()[:, None, :], angles.sin()[:, None, :] + rows = q[first : first + n] + x1, x2 = rows[..., : rot // 2], rows[..., rot // 2 : rot] + rows = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, rows[..., rot:]], dim=-1) + allowed = torch.arange(length, device=device)[None, :] <= qpos[:, None] + logits = torch.einsum("thd,kd->thk", rows, keys) * sm_scale + logits = logits.masked_fill(~allowed[:, None, :], float("-inf")) + out[first : first + n] = torch.softmax(logits, dim=-1) @ values + first += n + return out.reshape(total, heads * dim).to(qkv.dtype) + + +m3_draft_attention = M3DraftAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..bcd928fccf5b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The unfused sparse layers' index branch: per-head Gemma RMSNorm + partial NeoX RoPE of [index-Q | index-K], index-Q +out in FP8 E4M3 and index-K inserted into the FP8 index-K cache, in one kernel.""" + +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 +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3Fp8IndexerQkNormRope(OpWrapper): + """``qk`` [T, (heads + 1) * 128] BF16 (the index projection's output): returns index-Q [T, heads, 128] E4M3 and + writes index-K (E4M3) at ``out_cache_loc`` of ``index_k_cache`` [pages, 1, 128, 128]. Gemma RMSNorm of each head, + then RoPE on its leading ``rotary_dim`` channels (NeoX). + + The unfused sparse layers (``fuse_qkv_index_projection=False``) run it on the index branch with an FP8 indexer; + the fused layers' producer writes the index cache itself. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, one index head per rank", spec=dict(m=1, heads=1)), + Cell(why="a 48-row decode verify", spec=dict(m=48, heads=1)), + Cell(why="a prefill chunk", spec=dict(m=2048, heads=1)), + ) + + note = """ + `reference` is the in-place fused norm + RoPE in BF16 followed by the E4M3 cast (the kernel rounds through BF16 + before the cast); its index-K is what the kernel writes to the cache. The two kernels may round a handful of + values differently: the test allows fewer than 0.1% of the bytes to differ, as the in-tree test does. + """ + + def __call__( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_fp8_indexer_qk_norm_rope( + qk, + index_k_cache, + out_cache_loc, + num_heads_q, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + position_ids, + ) + + def reference( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """(index-Q [T, heads, head_dim] E4M3, index-K [T, head_dim] E4M3): the catalog's in-place kernel on a copy, + then the E4M3 cast.""" + out = fused_qk_norm_rope( + qk.clone(), + num_heads_q, + 1, + 0, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + True, + position_ids, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + q, k = out.split([num_heads_q * head_dim, head_dim], dim=-1) + return ( + q.reshape(q.shape[0], num_heads_q, head_dim).to(torch.float8_e4m3fn), + k.to(torch.float8_e4m3fn), + ) + + +m3_fp8_indexer_qk_norm_rope = M3Fp8IndexerQkNormRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_fused_decode_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_fused_decode_attn.py new file mode 100644 index 000000000000..1e1084aeccdf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_fused_decode_attn.py @@ -0,0 +1,643 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fused MSA short-sequence decode attention of this target's sparse layers, and the plain-torch +composition it fuses (shared with the megakernel and c1 entries, which wrap the same attention).""" + +from __future__ import annotations + +import math +from typing import NamedTuple, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Per-rank attention geometry at tp4: 16 query heads, one KV head, one index head, head_dim 128. +NUM_HEADS = 16 +HEAD_DIM = 128 +#: The packed QKV(+index) projection row: [Q x16 | K | V | idxQ | idxK] x 128. +PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM +#: Partial NeoX RoPE over the first 64 dims of every head. +ROTARY_DIM = 64 +#: One MSA block is one 128-token KV page; identity selection covers at most 16 of them. +PAGE = 128 +MAX_BLOCKS = 16 +MAX_SEQ_LEN = MAX_BLOCKS * PAGE +#: NVFP4 block: 16 elements per E4M3 scale. +NVFP4_BLOCK = 16 +#: MXFP8 block of the o_proj hand-off: 32 elements per UE8M0 scale. +MXFP8_BLOCK = 32 + +_E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +_E4M3_MAX = 448.0 + + +def e2m1_values(device) -> torch.Tensor: + """The 16 E2M1 codes' values (sign in bit 3).""" + pos = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=device) + return torch.cat([pos, -pos]) + + +def e4m3_bytes(x: torch.Tensor) -> torch.Tensor: + """Round fp32 to E4M3 the way ``__nv_fp8_e4m3(float)`` does (RN-even, saturating to +-448), as raw bytes.""" + return x.clamp(-_E4M3_MAX, _E4M3_MAX).to(torch.float8_e4m3fn).view(torch.uint8) + + +def e4m3_value(b: torch.Tensor) -> torch.Tensor: + return b.view(torch.float8_e4m3fn).float() + + +def e2m1_codes(x: torch.Tensor) -> torch.Tensor: + """``cvt.rn.satfinite.e2m1x2.f32`` per element: round-to-nearest-even onto {0, .5, 1, 1.5, 2, 3, 4, 6}, + saturating at 6, sign in bit 3 (a negative value that rounds to zero keeps its sign).""" + a = x.abs() + # Ties go to the even code: 0.25 -> 0, 0.75 -> 1.0, 1.25 -> 1.0, 1.75 -> 2.0, 2.5 -> 2.0, 3.5 -> 4.0, 5 -> 4. + code = ( + (a > 0.25).int() + + (a >= 0.75).int() + + (a > 1.25).int() + + (a >= 1.75).int() + + (a > 2.5).int() + + (a >= 3.5).int() + + (a > 5.0).int() + ) + return (code | (torch.signbit(x).int() << 3)).to(torch.uint8) + + +def nvfp4_quantize( + x: torch.Tensor, second_level_scale: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """The MiniMax-M3 producer's NVFP4 recipe for rows of 128: (packed data [N, 64] uint8, E4M3 scales [N, 8]). + + Elements are rounded to bf16 first; each 16-element block's scale is ``e4m3(sls * amax / 6)`` and the codes are + ``e2m1(x / (scale / sls))``. Element ``2i`` is the low nibble of byte ``i``. + """ + n = x.shape[0] + xb = x.to(torch.bfloat16).float().view(n, HEAD_DIM // NVFP4_BLOCK, NVFP4_BLOCK) + sls = second_level_scale.float().reshape(()) + amax = xb.abs().amax(-1) + scale = e4m3_bytes(sls * (amax * (1.0 / 6.0))) + out_scale = torch.where( + amax != 0, 1.0 / (e4m3_value(scale) * (1.0 / sls)), torch.zeros_like(amax) + ) + codes = e2m1_codes(xb * out_scale[..., None]).view(n, HEAD_DIM) + data = codes[:, 0::2] | (codes[:, 1::2] << 4) + return data, scale + + +def nvfp4_dequantize(data: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """[N, 64] packed E2M1 + [N, 8] E4M3 block scales -> fp32 [N, 128] (the per-tensor scale not applied).""" + table = e2m1_values(data.device) + lo = table[(data & 0xF).long()] + hi = table[(data >> 4).long()] + values = torch.stack([lo, hi], dim=-1).reshape(data.shape[0], HEAD_DIM) + return values * e4m3_value(scale).repeat_interleave(NVFP4_BLOCK, dim=-1) + + +def v_scale_rows(scale_plane: torch.Tensor) -> torch.Tensor: + """A page's V scale plane in the 4x4 token-quad order [token // 4][column][token % 4] -> [128, 8] row-major.""" + return ( + scale_plane.reshape(PAGE // 4, HEAD_DIM // NVFP4_BLOCK, 4) + .permute(0, 2, 1) + .reshape(PAGE, -1) + ) + + +def v_scale_plane(rows: torch.Tensor) -> torch.Tensor: + """Inverse of :func:`v_scale_rows`: [128, 8] row-major scales -> the plane's [128, 8] storage.""" + return rows.reshape(PAGE // 4, 4, HEAD_DIM // NVFP4_BLOCK).permute(0, 2, 1).reshape(PAGE, -1) + + +def norm_rope( + x: torch.Tensor, weight: torch.Tensor, eps: float, cos: torch.Tensor, sin: torch.Tensor +) -> torch.Tensor: + """Per-head Gemma RMSNorm (``x * rsqrt(mean(x^2) + eps) * (1 + w)``) then NeoX RoPE over the first 64 dims. + + x: fp32 [N, 128]; cos / sin: fp32 [N, 32] for each row's position.""" + y = x * (torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + weight.float())) + x1, x2 = y[:, : ROTARY_DIM // 2], y[:, ROTARY_DIM // 2 : ROTARY_DIM] + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, y[:, ROTARY_DIM:]], dim=-1) + + +def swizzled_scale_offsets(rows: int, k_blocks: int, device) -> torch.Tensor: + """Byte offset of scale (row, k-block) in the 128x4-swizzled UE8M0 layout: [rows, k_blocks] int64.""" + r = torch.arange(rows, device=device)[:, None] + k = torch.arange(k_blocks, device=device)[None, :] + k_tiles = (k_blocks + 3) // 4 + return ( + (r // 128) * k_tiles * 512 + (k // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + k % 4 + ) + + +def mxfp8_rows(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """The merge's MXFP8 hand-off recipe on bf16 rows [T, K]: (E4M3 bytes [T, K], UE8M0 exponents [T, K / 32]). + + Per 32-element block: ``e = ceil_exp(amax / 448)`` (biased, at most 254), values ``e4m3(x * 2^(127 - e))``.""" + t, k = x.shape + xf = x.float().view(t, k // MXFP8_BLOCK, MXFP8_BLOCK) + amax = xf.abs().amax(-1) + bits = (amax * (1.0 / _E4M3_MAX)).view(torch.int32) + ex = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).int() + ex = ex.clamp(max=254) + scale = torch.where(amax != 0, ((254 - ex) << 23).view(torch.float32), torch.zeros_like(amax)) + q8 = e4m3_bytes(xf * scale[..., None]).view(t, k) + return q8, ex.to(torch.uint8) + + +def mxfp8_dequantize(q8: torch.Tensor, ex: torch.Tensor) -> torch.Tensor: + """E4M3 bytes [T, K] with UE8M0 exponents [T, K / 32] -> fp64 values.""" + scale = torch.exp2(ex.double() - 127.0).repeat_interleave(MXFP8_BLOCK, dim=-1) + return e4m3_value(q8).double() * scale + + +def unswizzle_scales(swizzled: torch.Tensor, rows: int, k_blocks: int) -> torch.Tensor: + """The [rows, k_blocks] UE8M0 bytes of a 128x4-swizzled scale buffer.""" + flat = swizzled.reshape(-1).view(torch.uint8) + return flat[swizzled_scale_offsets(rows, k_blocks, flat.device)] + + +def mx_gemm( + a: torch.Tensor, a_exp: torch.Tensor, w: torch.Tensor, w_exp: torch.Tensor +) -> torch.Tensor: + """bf16 ``dequant(a) @ dequant(w)^T`` of MXFP8 operands (E4M3 bytes, row-major UE8M0 exponents), accumulated in + fp64 and rounded once.""" + return (mxfp8_dequantize(a, a_exp) @ mxfp8_dequantize(w, w_exp).T).to(torch.bfloat16) + + +class AttnState(NamedTuple): + """What the fused attention leaves behind: its bf16 output and the three caches after the new tokens' insert.""" + + output: torch.Tensor + kv_data: torch.Tensor + kv_scale: torch.Tensor + idx_k: torch.Tensor + + +def query_rows(packed, position_ids, q_norm_weight, eps, rotary_cos_sin) -> torch.Tensor: + """fp32 [T, 16, 128]: the query heads of packed rows after Gemma RMSNorm + NeoX RoPE (before E4M3).""" + rows = packed.float().view(packed.shape[0], NUM_HEADS + 4, HEAD_DIM) + pos_ids = position_ids.long() + cos = rotary_cos_sin[pos_ids, 0].float() + sin = rotary_cos_sin[pos_ids, 1].float() + return torch.stack( + [norm_rope(rows[:, h], q_norm_weight, eps, cos, sin) for h in range(NUM_HEADS)], dim=1 + ) + + +def attend_cached( + q, + kv_data, + kv_scale, + block_table, + seq_lens, + k_global_scale, + v_global_scale, + sm_scale, + decode_query_len, +) -> torch.Tensor: + """bf16 [T, 2048]: the identity-selection decode of fp32 queries ``q`` [T, 16, 128] over caches that already + hold the step's tokens. q is rounded to E4M3 (no scale), the operand the kernel's tensor core reads; request + ``r``'s token ``t`` sits at ``pos = seq_lens[r] - qlen + t`` and attends positions ``[0, pos]`` through + ``block_table``: scores ``q . dequant(K) * sm_scale * k_global_scale``, softmax, values + ``dequant(V) * v_global_scale``. fp64 throughout, one bf16 rounding at the end (the kernel's split-K partials + and their LSE merge are exact algebra).""" + q8 = e4m3_value(e4m3_bytes(q)).double() + t_total = q.shape[0] + qlen = int(decode_query_len) + out = torch.empty(t_total, NUM_HEADS, HEAD_DIM, dtype=torch.float64, device=q.device) + k_gs = float(k_global_scale.reshape(-1)[0]) + v_gs = float(v_global_scale.reshape(-1)[0]) + for r, seq_len in enumerate(seq_lens.tolist()): + pages = block_table[r, : math.ceil(seq_len / PAGE)].tolist() + kd = torch.cat([nvfp4_dequantize(kv_data[p, 0, 0], kv_scale[p, 0, 0]) for p in pages]) + vd = torch.cat( + [nvfp4_dequantize(kv_data[p, 1, 0], v_scale_rows(kv_scale[p, 1, 0])) for p in pages] + ) + kd, vd = kd.double() * k_gs, vd.double() * v_gs + for t in range(qlen): + row = r * qlen + t + n = seq_len - qlen + t + 1 + p = torch.softmax(q8[row] @ kd[:n].T * sm_scale, dim=-1) + out[row] = p @ vd[:n] + return out.to(torch.bfloat16).view(t_total, NUM_HEADS * HEAD_DIM) + + +def reference_attention( + packed, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + eps, + sm_scale, + decode_query_len, +) -> AttnState: + """The composition the fused kernel replaces, in plain torch, on copies of the caches. + + 1. Per token, q heads / k / index-k: Gemma RMSNorm + NeoX RoPE at ``position_ids``; v as is. + 2. k / v -> NVFP4 (``kv_quant_scale[1]`` / ``[2]``) into the token's ``out_cache_loc`` slot (K scales linear, V + scales in 4x4 token quads); index-k -> E4M3 into the index cache. + 3. :func:`attend_cached` over the caches after the insert. + """ + kv_data, kv_scale, idx_k = kv_data.clone(), kv_scale.clone(), idx_k.clone() + idx_bytes = idx_k.view(torch.uint8) + t_total = packed.shape[0] + rows = packed.float().view(t_total, NUM_HEADS + 4, HEAD_DIM) + pos_ids = position_ids.long() + cos = rotary_cos_sin[pos_ids, 0].float() + sin = rotary_cos_sin[pos_ids, 1].float() + q = query_rows(packed, position_ids, q_norm_weight, eps, rotary_cos_sin) + k = norm_rope(rows[:, NUM_HEADS], k_norm_weight, eps, cos, sin) + v = rows[:, NUM_HEADS + 1] + ik = norm_rope(rows[:, NUM_HEADS + 3], index_k_norm_weight, eps, cos, sin) + k_data, k_sf = nvfp4_quantize(k, kv_quant_scale[1]) + v_data, v_sf = nvfp4_quantize(v, kv_quant_scale[2]) + ik_bytes = e4m3_bytes(ik) + for t, slot in enumerate(out_cache_loc[:t_total].tolist()): + page, row = divmod(slot, PAGE) + kv_data[page, 0, 0, row] = k_data[t] + kv_data[page, 1, 0, row] = v_data[t] + kv_scale[page, 0, 0, row] = k_sf[t] + vrows = v_scale_rows(kv_scale[page, 1, 0]) + vrows[row] = v_sf[t] + kv_scale[page, 1, 0] = v_scale_plane(vrows) + idx_bytes[page, 0, row] = ik_bytes[t] + output = attend_cached( + q, + kv_data, + kv_scale, + block_table, + seq_lens, + k_global_scale, + v_global_scale, + sm_scale, + decode_query_len, + ) + return AttnState(output, kv_data, kv_scale, idx_k) + + +def check_attention_inputs( + packed_rows, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + decode_query_len, +) -> None: + """The inputs the attention kernels accept and answer wrongly (shared by the three entries' ``is_valid``).""" + assert ( + kv_data.dim() == 5 + and tuple(kv_data.shape[1:]) == (2, 1, PAGE, HEAD_DIM // 2) + and kv_data.stride(4) == 1 + and kv_data.stride(3) == HEAD_DIM // 2 + ), "kv_data must be [pages, 2, 1, 128, 64] with contiguous head blocks" + assert ( + kv_scale.dim() == 5 + and tuple(kv_scale.shape[1:]) == (2, 1, PAGE, HEAD_DIM // NVFP4_BLOCK) + and kv_scale.stride(4) == 1 + and kv_scale.stride(3) == HEAD_DIM // NVFP4_BLOCK + ), "kv_scale must be [pages, 2, 1, 128, 8] with contiguous head blocks" + assert ( + idx_k.dim() == 4 and tuple(idx_k.shape[1:]) == (1, PAGE, HEAD_DIM) and idx_k.stride(3) == 1 + ), "idx_k must be [pages, 1, 128, 128] with contiguous rows" + for name, w in ( + ("q_norm_weight", q_norm_weight), + ("k_norm_weight", k_norm_weight), + ("index_k_norm_weight", index_k_norm_weight), + ): + assert w.is_contiguous(), f"{name} must be contiguous (read as 4-element vectors)" + assert rotary_cos_sin.is_contiguous() and tuple(rotary_cos_sin.shape[1:]) == ( + 2, + ROTARY_DIM // 2, + ), "rotary_cos_sin must be contiguous [max_pos, 2, 32]" + for name, t in ( + ("out_cache_loc", out_cache_loc), + ("position_ids", position_ids), + ("seq_lens", seq_lens), + ): + assert t.is_contiguous(), f"{name} must be contiguous" + qlen = int(decode_query_len) + num_reqs = seq_lens.shape[0] + assert packed_rows == num_reqs * qlen, "one row per (request, query token)" + assert position_ids.numel() == packed_rows, "one position per token" + lens = seq_lens.tolist() + assert all(qlen <= n <= MAX_SEQ_LEN for n in lens), ( + f"every seq_len must be in [decode_query_len, {MAX_SEQ_LEN}]: identity selection covers 16 blocks of 128, " + "a longer request is silently cut to them" + ) + assert block_table.shape[1] >= math.ceil(max(lens) / PAGE), ( + "block_table must hold every attended block (missing columns read page 0)" + ) + assert int(position_ids.max()) < rotary_cos_sin.shape[0], "position past the rotary table" + table = block_table.tolist() + slots = out_cache_loc.tolist() + for r, n in enumerate(lens): + for t in range(qlen): + pos = n - qlen + t + want = table[r][pos // PAGE] * PAGE + pos % PAGE + assert slots[r * qlen + t] == want, ( + "out_cache_loc must be the slot block_table gives the token's position seq_len - qlen + t: the " + "attention reads that page's other rows from the cache it loaded before the insert" + ) + + +class M3FusedDecodeAttn(OpWrapper): + """One launch after the packed QKV GEMM of a sparse layer in short-sequence decode. + + Fuses q / k / index-k Gemma RMSNorm + NeoX partial RoPE, the NVFP4 K/V and FP8 index-K cache insert, the + block-sparse decode over every attended block (selection is the identity: at most 16 blocks of 128 tokens) and + the split-K LSE merge, writing the bf16 output and its MXFP8 copy for the o_proj GEMM. + + ``(tokens_per_cluster, blocks_per_cta)`` is the launch geometry; the target passes + ``ops.m3_fused_attn.default_geometry(qlen, num_requests)``. Returns None: the cells read ``output``, + ``mxfp8_data`` / ``mxfp8_scale`` and the three caches. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # Every compiled geometry, each at a step shape the target routes to it (decode_query_len x requests, at most + # 16 tokens). seq_lens is a per-request list; positions are seq_len - qlen + t. + CELLS: tuple[Cell, ...] = ( + Cell( + why="plain decode, one request, a few blocks: geometry (1, 1, 1)", + spec=dict(qlen=1, seq_lens=(300,), geometry=(1, 1)), + ), + Cell( + why="plain decode, 6 requests: the widest one-block-per-CTA batch; full 16 blocks and a page's 1st row", + spec=dict(qlen=1, seq_lens=(2048, 129, 1, 700, 1500, 256), geometry=(1, 1)), + ), + Cell( + why="plain decode, 8 requests: KV blocks paired per CTA (8-CTA clusters), geometry (1, 1, 2)", + spec=dict(qlen=1, seq_lens=(2048, 33, 128, 129, 1024, 1900, 640, 7), geometry=(1, 2)), + ), + Cell( + why="plain decode at the 16-token cap: 16 paired-block clusters", + spec=dict(qlen=1, seq_lens=tuple(range(97, 2048, 125))[:16], geometry=(1, 2)), + ), + Cell( + why="Eagle3 verify of 1 draft (qlen 2), 2 requests: geometry (2, 1, 1); a draft crossing a page", + spec=dict(qlen=2, seq_lens=(129, 900), geometry=(1, 1)), + ), + Cell( + why="qlen 2 at 8 requests: geometry (2, 1, 2)", + spec=dict(qlen=2, seq_lens=(2, 2048, 130, 257, 640, 1000, 1500, 64), geometry=(1, 2)), + ), + Cell( + why="qlen 3 at 5 requests (15 tokens): geometry (3, 1, 1)", + spec=dict(qlen=3, seq_lens=(3, 2048, 130, 385, 777), geometry=(1, 1)), + ), + Cell( + why="qlen 4 at 4 requests: geometry (4, 1, 1)", + spec=dict(qlen=4, seq_lens=(4, 2048, 131, 1030), geometry=(1, 1)), + ), + Cell( + why="qlen 5 at 3 requests: geometry (5, 1, 1)", + spec=dict(qlen=5, seq_lens=(5, 1100, 2048), geometry=(1, 1)), + ), + Cell( + why="qlen 6, one request (the c1 shape when c1 does not apply): geometry (6, 1, 1)", + spec=dict(qlen=6, seq_lens=(1283,), geometry=(1, 1)), + ), + Cell( + why="qlen 6 at 2 requests: geometry (6, 1, 2)", + spec=dict(qlen=6, seq_lens=(6, 2048), geometry=(1, 2)), + ), + Cell( + why="qlen 7, two requests: geometry (7, 1, 1)", + spec=dict(qlen=7, seq_lens=(132, 2000), geometry=(1, 1)), + ), + Cell( + why="qlen 8 at 2 requests (the 16-token cap): geometry (8, 1, 1)", + spec=dict(qlen=8, seq_lens=(8, 1029), geometry=(1, 1)), + ), + # Every other (query rows per request, requests) the decode dispatch sends here (dispatch.FUSED_DECODE_ATTN): + # lengths spread over the 16-block window, the geometry the target picks for the shape. + *( + Cell( + why=f"qlen {qlen} at {reqs} requests: a decode shape between the ones above", + spec=dict( + qlen=qlen, + seq_lens=tuple( + qlen + (r * 389 + 97 * qlen) % (2049 - qlen) for r in range(reqs) + ), + geometry=geometry, + ), + ) + for qlen, reqs, geometry in ( + (1, 2, (1, 1)), + (1, 3, (1, 1)), + (1, 4, (1, 1)), + (1, 5, (1, 1)), + (1, 7, (1, 2)), + (1, 9, (1, 2)), + (1, 10, (1, 2)), + (1, 11, (1, 2)), + (1, 12, (1, 2)), + (1, 13, (1, 2)), + (1, 14, (1, 2)), + (1, 15, (1, 2)), + (2, 1, (1, 1)), + (2, 3, (1, 1)), + (2, 4, (1, 2)), + (2, 5, (1, 2)), + (2, 6, (1, 2)), + (2, 7, (1, 2)), + (3, 1, (1, 1)), + (3, 2, (1, 1)), + (3, 3, (1, 1)), + (3, 4, (1, 1)), + (4, 1, (1, 1)), + (4, 2, (1, 1)), + (4, 3, (1, 1)), + (5, 1, (1, 1)), + (5, 2, (1, 1)), + (7, 1, (1, 1)), + (8, 1, (1, 1)), + ) + ), + ) + + note = """ + Only the variants the target launches are compiled: the geometries above, each with the MXFP8 hand-off (the + target always passes it). Q is quantized to E4M3 with no scale (`kv_quant_scale[0]` is not read): its operand + on the tensor core is the E4M3 value, which `reference` reproduces. + + Not certified: a KV pool whose pages interleave other layers (any `stride(0)` / `stride(1)` that is a multiple + of 16 is accepted; the cells use a contiguous pool), `position_ids` other than `seq_len - qlen + t` (RoPE reads + `position_ids`, placement reads `seq_lens`; the target passes equal values), and `use_pdl=False` timing. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + ) -> None: + torch.ops.trtllm.minimax_m3_fused_decode_attn( + packed, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + eps, + sm_scale, + decode_query_len, + tokens_per_cluster, + blocks_per_cta, + output, + mxfp8_data, + mxfp8_scale, + use_pdl, + ) + + def reference( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + ) -> AttnState: + """The bf16 output and the caches after the insert, as :func:`reference_attention` computes them. + + The geometry, the output buffers and ``use_pdl`` do not change what is computed and are ignored; the + MXFP8 hand-off is :func:`mxfp8_rows` of the bf16 output, which the test checks on the op's own output. + """ + return reference_attention( + packed, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + eps, + sm_scale, + decode_query_len, + ) + + def is_valid( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + ) -> None: + check_attention_inputs( + packed.shape[0], + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + decode_query_len, + ) + k_blocks = output.shape[1] // MXFP8_BLOCK + needed = math.ceil(packed.shape[0] / 128) * ((k_blocks + 3) // 4) * 512 + assert mxfp8_scale.numel() >= needed, ( + "mxfp8_scale must hold the swizzled scales of every output row tile" + ) + + +m3_fused_decode_attn = M3FusedDecodeAttn() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_index_score.py new file mode 100644 index 000000000000..c8a9b54d2a3a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_index_score.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# PyTorch oracle ported from vLLM (Apache-2.0), _reference_decode_index_score: +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/tests/kernels/attention/test_minimax_m3.py#L188 +"""The MSA index scorer of a sparse decode step: per (index head, query token, KV block) the causal max of Q . K over +the block's 128 index-K positions, with the pre-wait schedule of ``ops/m3_index_score``.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +__all__ = ["m3_index_score"] + +PAGE = 128 +HEAD_DIM = 128 +#: CTAs per request (each scores blocks split, split + SPLIT_K, ...; the scores do not depend on it). +SPLIT_K = 512 + + +def compiled_scorer(key: tuple): + """The compiled kernel of ``key`` (dtype, heads, query length, head dim, split_k, early), compiling it on a miss + (the target compiles every key it can take before warmup: ``layers.prebuild_index_scorer``).""" + from ..m3_index_score import IndexScoreRunner + + if key not in IndexScoreRunner.kernel_cache: + IndexScoreRunner._compile(*key) + return IndexScoreRunner.kernel_cache[key] + + +class M3IndexScore(OpWrapper): + """``score[h, t, j]`` (fp32, a transposed view of the [heads, blocks, tokens] selector buffer) = the max over the + positions ``p`` of block ``j`` of request ``b = t // qlen`` with ``p <= seq_lens[b] - qlen + t % qlen`` of + ``idx_q[t, h] . index_k[block_table[b, j], p % 128]``, for ``j < ceil(seq_lens[b] / 128)`` (no other entry is + written). FP8 or BF16 q and index-K, fp32 products and sums. + + ``early`` (the target's one-request kernel) issues the Q and first K tiles right after the grid wait from a + pre-wait length hint and prefetches one older block per CTA into L2 before the wait; otherwise every older block of + the CTA is prefetched. Neither changes a score. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over 20000 tokens: the early-issue kernel", + spec=dict(seq_lens=(20000,), qlen=4, early=True), + ), + Cell( + why="one request, one row, 65536 tokens: early issue, 512 blocks on 512 CTAs", + spec=dict(seq_lens=(65536,), qlen=1, early=True), + ), + Cell( + why="one request of 129 tokens: the early tiles of most CTAs lie past the length and are drained", + spec=dict(seq_lens=(129,), qlen=4, early=True), + ), + Cell( + why="one request longer than 512 blocks (the CTAs loop): early issue", + spec=dict(seq_lens=(70000,), qlen=4, early=True), + ), + Cell( + why="two requests: the prefetching kernel", + spec=dict(seq_lens=(20000, 3000), qlen=4, early=False), + ), + Cell( + why="eight requests of mixed lengths: the prefetching kernel", + spec=dict( + seq_lens=(300, 1500, 2049, 4096, 5000, 7777, 8192, 33001), qlen=4, early=False + ), + ), + ) + + note = """ + The products of FP8 (or BF16) values are exact in fp32 and the kernel sums 128 of them in its k-split order, so it + differs from the reference matmul by a few fp32 ulps; it is held to the in-tree scorer bit for bit. + """ + + def __call__( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + _, num_heads, head_dim = idx_q.shape + key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + compiled_scorer(key)(idx_q, index_k, block_table, score, seq_lens) + return score + + def reference( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + """The fp32 [heads, total_q, blocks] scores (``-inf`` where nothing is written). ``score``, ``split_k`` and + ``early`` do not change them.""" + total_q, heads, _ = idx_q.shape + qlen = max_decode_query_len + out = torch.full( + (heads, total_q, score.shape[2]), + -float("inf"), + dtype=torch.float32, + device=idx_q.device, + ) + with true_fp32_matmul(): + for b, seq in enumerate(seq_lens.tolist()): + nblocks = (seq + PAGE - 1) // PAGE + if nblocks == 0: + continue + t0 = b * qlen + q = idx_q[t0 : t0 + qlen].float() + k = index_k[block_table[b, :nblocks].long()].reshape(nblocks * PAGE, -1).float() + s = torch.einsum("qhd,kd->hqk", q, k) + q_pos = seq - qlen + torch.arange(qlen, device=idx_q.device) + k_pos = torch.arange(k.shape[0], device=idx_q.device) + s.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) + out[:, t0 : t0 + qlen, :nblocks] = s.reshape(heads, qlen, nblocks, PAGE).amax(dim=3) + return out + + def is_valid( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> None: + # The block-table row and the score row are read and written up to ceil(seq_len / 128) with no bound. + nblocks = (seq_lens + PAGE - 1) // PAGE + assert bool((nblocks <= block_table.shape[1]).all()), ( + "a request runs past its block-table row" + ) + assert bool((nblocks <= score.shape[2]).all()), "a request runs past the score buffer" + assert idx_q.shape[0] == seq_lens.shape[0] * max_decode_query_len, ( + "idx_q must hold max_decode_query_len rows per request" + ) + + +m3_index_score = M3IndexScore() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_msa_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_msa_attention.py new file mode 100644 index 000000000000..4043e306696d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_msa_attention.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MiniMax-M3 MSA attention backend's step: KV write, sparse top-k block selection, FMHA.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + SparseBackendForwardArgs, +) + +from .m3_sparse_decode_lc import dequant_nvfp4_pages + + +class M3MsaAttention(OpWrapper): + """One layer's attention core on the MSA backend (``MiniMaxM3MsaSparseAttention``), written into ``output`` + [T, heads * 128] (bf16, the step's live rows). + + * ``k`` / ``v`` given (the dense layers): written into the layer's paged main cache first (NVFP4 with + ``kv_scale_orig_quant``, else FP8 / BF16 as the pool is). None: a producer already wrote the caches. + * ``sparse``: the indexer scores the layer's index-K cache (already written) with ``idx_q`` and publishes the + top-k blocks; the FMHA attends those blocks. Otherwise the FMHA attends the whole page table. + + The FMHA dispatches on the step's metadata (context rows: the MSA prefill FMHA; generation rows: the MSA decode + FMHA, the dense layers' trtllm-gen decode over their full context). ``backend`` and ``metadata`` are the runtime's + objects for this layer and step. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="dense layer, a prefill chunk, FP8 pool", + spec=dict(layer="dense", phase="context", tokens=1024), + ), + Cell( + why="dense layer, a decode step of 4 requests", + spec=dict(layer="dense", phase="generation", reqs=4), + ), + Cell( + why="sparse layer, a prefill chunk beyond 16 blocks, NVFP4 pool", + spec=dict(layer="sparse", phase="context", tokens=4096), + ), + Cell( + why="sparse layer, a mixed step (context + generation rows)", + spec=dict(layer="sparse", phase="mixed", tokens=512, reqs=2), + ), + ) + + note = """ + `reference` is exact (float64) softmax attention over the cache rows the step attends, read back from the paged + cache the call wrote (or a producer wrote before it) and dequantized as the FMHA dequantizes them: FP8 rows as + E4M3, NVFP4 rows as E2M1 x E4M3 block scale x the per-tensor ``kv_scale_quant_orig``. A sparse layer's top-k blocks + are re-selected from ``idx_q`` and the index-K cache with the backend's rule (per-block max over the causal + keys, the query's own block forced in, ties to the lower block). On an FP8 pool the FMHA takes Q as a plain E4M3 + cast, which the reference applies too. What remains between the two is the FMHA's own arithmetic (its P + rounding and accumulation order), which the test bounds. + """ + + def __call__( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + if k is not None: + backend.write_layer_caches( + k, v, None, metadata, kv_scale_orig_quant=kv_scale_orig_quant + ) + if sparse: + blocks = backend.run_indexer(idx_q, None, metadata, idx_k_prewritten=True) + args = AttentionForwardArgs( + output=output, sparse_backend_args=SparseBackendForwardArgs(topk_indices=blocks) + ) + else: + args = AttentionForwardArgs(output=output) + args.kv_scale_orig_quant = kv_scale_orig_quant + args.kv_scale_quant_orig = kv_scale_quant_orig + backend.forward(q, None, None, metadata, forward_args=args) + return output + + def reference( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + """``output`` as the call is supposed to write it, [T, heads * 128] in ``output``'s dtype. + + Every attended row -- the step's own included -- is read from the layer's paged cache, at the slots the + metadata's KV cache manager holds for each request, so this is evaluated after the call (whose write of + ``k`` / ``v`` it then reads back). ``k``, ``v`` and ``kv_scale_orig_quant`` are therefore not read; the + test checks the write itself. Query token ``j`` of request ``r`` sits at position ``cached_r + j`` and + attends positions up to it; on a sparse layer only those in its top-k blocks. Mirrors one KV head per rank, + the target's layout. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + assert backend.num_kv_heads == 1, "mirrors the one-KV-head-per-rank layout only" + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = sum(qlens) + assert q.shape[0] == total and output.shape[0] == total + + pool = manager.get_buffers(layer, "HND") + nvfp4 = manager.is_nvfp4_layer(layer) + if nvfp4: + scale_pool = manager.get_block_scale_buffers(layer, "HND") + k_global = float(kv_scale_quant_orig.reshape(-1)[1]) + v_global = float(kv_scale_quant_orig.reshape(-1)[2]) + q_rows = q.reshape(total, heads, dim) + if q_rows.dtype != torch.float8_e4m3fn and pool.dtype == torch.float8_e4m3fn: + # The FP8 FMHA shares one dtype across Q / K / V: Q is cast to E4M3, unscaled. + q_rows = q_rows.to(torch.float8_e4m3fn) + q64 = q_rows.to(torch.float64) + if sparse: + config = backend.m3_config + idx_cache = metadata.msa_idx_k_cache(layer) + idx64 = idx_q.reshape(total, config.num_index_heads, config.sparse_index_dim).to( + torch.float64 + ) + + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=q.device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + blocks = (length + page - 1) // page + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][:blocks], + dtype=torch.int64, + device=q.device, + ) + if nvfp4: + keys = dequant_nvfp4_pages( + pool[:, 0].view(torch.uint8), scale_pool[:, 0], pages, quad_scales=False + ) + values = dequant_nvfp4_pages( + pool[:, 1].view(torch.uint8), scale_pool[:, 1], pages, quad_scales=True + ) + keys, values = keys * k_global, values * v_global + else: + keys = pool[pages, 0, 0].to(torch.float64) + values = pool[pages, 1, 0].to(torch.float64) + keys = keys.reshape(-1, dim)[:length] + values = values.reshape(-1, dim)[:length] + qpos = cached[row] + torch.arange(n, device=q.device) + allowed = torch.arange(length, device=q.device)[None, :] <= qpos[:, None] + if sparse: + idx_keys = idx_cache[pages, 0].to(torch.float64).reshape(-1, dim)[:length] + chosen = _selected_blocks( + idx64[first : first + n], idx_keys, allowed, qpos, page, config + ) + allowed &= chosen.repeat_interleave(page, dim=1)[:, :length] + for lo in range(0, n, 256): + hi = min(n, lo + 256) + logits = torch.einsum("thd,kd->thk", q64[first + lo : first + hi], keys) * sm_scale + logits = logits.masked_fill(~allowed[lo:hi, None, :], float("-inf")) + out[first + lo : first + hi] = torch.softmax(logits, dim=-1) @ values + first += n + return out.reshape(total, heads * dim).to(output.dtype) + + +def _selected_blocks( + idx_q: torch.Tensor, + idx_keys: torch.Tensor, + allowed: torch.Tensor, + qpos: torch.Tensor, + page: int, + config, +) -> torch.Tensor: + """The blocks each query token of one request attends, [n, blocks] bool. + + A block's score is the maximum, over the KV head's index heads and the block's causally visible keys, of + ``idx_q . idx_k``. Of the token's ``ceil((pos + 1) / page)`` valid blocks the first ``init_blocks`` and the last + ``local_blocks`` are forced in, and the ``topk`` highest-scoring are kept, ties to the lower block (the + backend's ``select_blocks_from_maxscore`` rule). The positive score scale the scorers apply does not change the + ranking, so none is applied. + """ + n, length = allowed.shape + blocks = (length + page - 1) // page + scores = torch.einsum("thd,kd->thk", idx_q, idx_keys).amax(dim=1) + scores = scores.masked_fill(~allowed, float("-inf")) + padded = torch.full((n, blocks * page), float("-inf"), dtype=scores.dtype, device=scores.device) + padded[:, :length] = scores + block_scores = padded.view(n, blocks, page).amax(dim=-1) + ids = torch.arange(blocks, device=scores.device)[None, :] + n_valid = torch.div(qpos + page, page, rounding_mode="floor")[:, None] + block_scores = block_scores.masked_fill(ids >= n_valid, float("-inf")) + forced = (ids < torch.clamp(n_valid, max=config.init_blocks)) | ( + (ids >= n_valid - config.local_blocks) & (ids < n_valid) + ) + block_scores = block_scores.masked_fill(forced, float("inf")) + order = torch.sort(block_scores, dim=1, descending=True, stable=True).indices[:, : config.topk] + keep = torch.arange(order.shape[1], device=scores.device)[None, :] < n_valid + chosen = torch.zeros(n, blocks, dtype=torch.bool, device=scores.device) + rows = torch.arange(n, device=scores.device)[:, None].expand_as(order) + chosen[rows[keep], order[keep]] = True + return chosen + + +m3_msa_attention = M3MsaAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_producer.py new file mode 100644 index 000000000000..957a1b82eef0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_producer.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer, with its constant loads staged +before the PDL grid-dependency wait and an early dependents trigger (``kernels/minimaxM3/producerPw.cu``).""" + +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 + +__all__ = ["m3_producer"] + +HEAD_DIM = 128 +ROTARY_DIM = 64 +PAGE = 128 +#: Query heads per rank at TP4 (MiniMax-M3: 64 query heads, 1 KV head and 1 index head per rank). +HEADS_Q = 16 +#: NVFP4 block: 16 values share one E4M3 scale. +FP4_BLOCK = 16 +E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def _norm_rope( + x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, eps: float +): + """Gemma RMSNorm (``x * rsqrt(mean(x^2) + eps) * (1 + w)``) then NeoX RoPE on the first 64 dims, in fp32. + x [T, 128] fp32; cos / sin [T, 32].""" + y = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) * (1.0 + weight.float()) + lo, hi, rest = y[:, :32], y[:, 32:64], y[:, 64:] + return torch.cat([lo * cos - hi * sin, hi * cos + lo * sin, rest], dim=-1) + + +def decode_nvfp4_token( + data: torch.Tensor, scales: torch.Tensor, quant_scale: float +) -> torch.Tensor: + """fp32 [128] of one token's packed E2M1 row (64 bytes, low nibble first) and its 8 E4M3 block scales.""" + table = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=data.device) + nib = torch.stack([data & 0xF, data >> 4], dim=-1).reshape(-1).long() + vals = table[nib & 7] * torch.where(nib >= 8, -1.0, 1.0) + s = scales.view(torch.float8_e4m3fn).float() / quant_scale + return vals * s.repeat_interleave(FP4_BLOCK) + + +def scale_index(within_page: int, plane: int) -> torch.Tensor: + """Positions of a token's 8 block scales in its page's [128 * 8] scale plane: row-major for K, the 4-token + swizzle for V.""" + cols = torch.arange(8) + if plane == 0: + return within_page * 8 + cols + return (within_page // 4) * 32 + cols * 4 + within_page % 4 + + +class M3Producer(OpWrapper): + """The NVFP4 QKV + indexer producer of a sparse layer: ``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert`` for one + KV head, one index head, head dim 128 and rotary dim 64, with the same arguments and outputs bit for bit. + + ``packed`` [T, (heads_q + 4) * 128] BF16 holds each token's q heads, k, v, index-q and index-k. Every head but v + gets the Gemma RMSNorm with its weight and NeoX RoPE on its first 64 dims (``rotary_cos_sin[position]``). Returns + (q FP8 [T, heads_q, 128], index-q FP8 [T, 1, 128]); writes the token's slot ``out_cache_loc[t]`` (skipped when + negative or past the pool): k and v as NVFP4 (16-value blocks, E4M3 scales against ``kv_quant_scale[1]`` / + ``[2]``, the V scales in the 4-token swizzle) and index-k as FP8. ``early`` triggers the dependent kernel on + entry (its outputs are read only after the dependent's own grid wait). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token (c1 without speculation)", spec=dict(tokens=1, scale=1.0)), + Cell(why="a c1 verify: 4 tokens", spec=dict(tokens=4, scale=3.0)), + Cell(why="a c4 verify: 16 tokens, small activations", spec=dict(tokens=16, scale=1e-3)), + Cell( + why="a c8 verify: 32 tokens, large activations, a skipped and an out-of-pool slot", + spec=dict(tokens=32, scale=40.0, bad_slots=True), + ), + Cell(why="the widest decode step: 48 tokens", spec=dict(tokens=48, scale=1.0)), + Cell( + why="a prefill chunk of 1000 tokens, early trigger off", + spec=dict(tokens=1000, scale=1.0, early=False), + ), + ) + + note = """ + The kernel is the native producer's arithmetic, compiled with the library's --use_fast_math, so it is held to that + kernel bit for bit on every output and cache byte. The torch reference states the math in fp32: the FP8 outputs are + held to it within one E4M3 step, and the NVFP4 K/V, decoded, within one E2M1 step of their block scale. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + q, index_q = torch.ops.trtllm.minimax_m3_producer_pw( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + heads_q, + pdl, + early, + ) + return q, index_q + + def reference( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + ) -> dict[str, torch.Tensor]: + """fp32 values before their quantization: "q" [T, heads_q, 128], "index_q" [T, 128], and per token "k", "v", + "index_k" [T, 128] (the values written to the token's slot). The caches, ``pdl`` and ``early`` do not change + them.""" + t = packed.shape[0] + x = packed.float().view(t, heads_q + 4, HEAD_DIM) + rope = rotary_cos_sin[position_ids[:t].long()] + cos, sin = rope[:, 0], rope[:, 1] + q = torch.stack( + [_norm_rope(x[:, h], q_weight, cos, sin, eps) for h in range(heads_q)], dim=1 + ) + return { + "q": q, + "k": _norm_rope(x[:, heads_q], k_weight, cos, sin, eps), + "v": x[:, heads_q + 1], + "index_q": _norm_rope(x[:, heads_q + 2], index_q_weight, cos, sin, eps), + "index_k": _norm_rope(x[:, heads_q + 3], index_k_weight, cos, sin, eps), + } + + def is_valid( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + ) -> None: + t = packed.shape[0] + # The RoPE row is read at the position with no bound. + pos = position_ids[:t] + assert bool(((pos >= 0) & (pos < rotary_cos_sin.shape[0])).all()), ( + "position_ids must index rotary_cos_sin" + ) + # Two tokens writing one slot race. + slots = out_cache_loc[:t] + live = slots[(slots >= 0) & (slots < kv_data.shape[0] * PAGE)] + assert live.unique().numel() == live.numel(), "out_cache_loc must not repeat a slot" + + +m3_producer = M3Producer() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..6493c4a45025 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_qk_norm_rope_fp8.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' per-head Gemma RMSNorm + partial NeoX RoPE of the packed [Q | K | V], into a fresh FP8 copy.""" + +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 +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3QkNormRopeFp8(OpWrapper): + """``fused_qk_norm_rope`` written out of place in FP8 E4M3: Q and K heads normed and roped in fp32 and rounded + once to E4M3 (saturating at +/-448), V cast, ``qkv`` left untouched. Same arguments as the catalog's in-place + entry. + + The dense layers (0-2) feed the FP8 main KV cache with it: on every step of an FP8 cache, and on the decode steps + of a hybrid NVFP4 cache (whose dense layers sit in its FP8 pool). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, tp4 heads (16 q, 1 kv)", spec=dict(m=1, q=16, kv=1)), + Cell(why="a 48-row decode verify", spec=dict(m=48, q=16, kv=1)), + Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=16, kv=1)), + ) + + note = """ + `reference` is the in-place kernel followed by a cast to E4M3. V matches it bit for bit; Q and K round once from + fp32 where the composition rounds through bf16 first, so they are within one E4M3 step of it. + """ + + def __call__( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + return torch.ops.trtllm.fused_qk_norm_rope_to_fp8( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + + def reference( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + """The catalog's in-place kernel on a copy, then the E4M3 cast (V exact; Q and K within one E4M3 step).""" + out = fused_qk_norm_rope( + qkv.clone(), + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + return out.to(torch.float8_e4m3fn) + + +m3_qk_norm_rope_fp8 = M3QkNormRopeFp8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_qkv_insert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_qkv_insert.py new file mode 100644 index 000000000000..91927ab0bc39 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_qkv_insert.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' generic horizontal producer: packed QKV/index rows -> FP8 Q / index-Q, K / V / index-K into +the paged caches.""" + +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 M3QkvInsert(OpWrapper): + """The packed ``[Q | K | V | index-Q | index-K]`` bf16 rows of a sparse layer: per-head Gemma RMSNorm and the + partial NeoX RoPE of Q, K, index-Q and index-K, K / V written to the main paged cache at ``out_cache_loc`` and + index-K to the FP8 index cache; returns (Q, index-Q) in FP8 E4M3 [M, heads, 128]. + + ``kv_scale`` given: the main cache is NVFP4 (packed E2M1 data [pages, 2, kv, 128, 64] uint8 + E4M3 scale pool + [pages, 2, kv, 128, 8], ``kv_quant_scale`` the fp32 (Q, K, V) quantization scales); None: FP8 E4M3 + [pages, 2, kv, 128, 128]. Token-major and step-type agnostic. ``m3_producer`` takes the NVFP4 cache at the tp4 + heads; this entry takes the FP8 cache and NVFP4 pools outside that layout. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="FP8 main cache, a prefill chunk", spec=dict(m=1024, cache="fp8")), + Cell(why="FP8 main cache, one decode token", spec=dict(m=1, cache="fp8")), + ) + + note = """ + The arithmetic is the CUDA producer's (fp32 norm and RoPE with bf16 roundings, then the E4M3 casts); `reference` + returns the fp32 values, and the test holds the FP8 outputs and cache rows within one E4M3 step of them. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + common = ( + num_heads, + num_kv_heads, + num_index_heads, + head_dim, + rotary_dim, + eps, + q_norm, + k_norm, + index_q_norm, + index_k_norm, + rotary_cos_sin, + position_ids, + ) + if kv_scale is not None: + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, kv_scale, index_k_cache, out_cache_loc, kv_quant_scale, *common + ) + return torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, index_k_cache, out_cache_loc, *common + ) + + def reference( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + """The fp32 values the producer rounds (a new dict, nothing written): Gemma RMSNorm + NeoX partial RoPE of + Q [M, heads, D], index-Q [M, index heads, D], K [M, kv heads, D] and index-K [M, D]; V [M, kv heads, D] as + packed. The cache rows at ``out_cache_loc`` and the returned FP8 Q / index-Q are these values cast.""" + m = packed.shape[0] + sizes = [ + num_heads * head_dim, + num_kv_heads * head_dim, + num_kv_heads * head_dim, + num_index_heads * head_dim, + head_dim, + ] + q, k, v, iq, ik = packed.float().split(sizes, dim=-1) + cos, sin = rotary_cos_sin[position_ids.long()].unbind(1) + half = rotary_dim // 2 + + def norm_rope(x, w, heads): + x = x.reshape(m, heads, head_dim) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w.float()) + x1, x2, rest = x[..., :half], x[..., half:rotary_dim], x[..., rotary_dim:] + c, s = cos.unsqueeze(1), sin.unsqueeze(1) + return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s, rest], dim=-1) + + return { + "q": norm_rope(q, q_norm, num_heads), + "index_q": norm_rope(iq, index_q_norm, num_index_heads), + "k": norm_rope(k, k_norm, num_kv_heads), + "v": v.reshape(m, num_kv_heads, head_dim), + "index_k": norm_rope(ik, index_k_norm, 1)[:, 0], + } + + +m3_qkv_insert = M3QkvInsert() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_select_blocks.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_select_blocks.py new file mode 100644 index 000000000000..6bf1764bfc0d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_select_blocks.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MSA top-k block selection from per-KV-head block scores (the built-in selector).""" + +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 + + +class M3SelectBlocks(OpWrapper): + """Per query row the ``topk`` highest-scoring KV blocks of ``max_score_kv`` [kv_heads, n_blocks, total_q] (fp32), + with the first ``init_blocks`` and the last ``local_blocks`` valid blocks forced in and blocks at or beyond the + row's ``n_valid_blocks`` excluded; [total_q, kv_heads, topk] int32 ascending ids, -1 padded. ``head_major_output`` + backs the result head-major so ``permute(1, 0, 2)`` is contiguous. + + The decode dispatch runs it where the long-row selector does not apply: requests of 17..128 KV blocks. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="eight decode rows of 17..128 blocks", + spec=dict(total_q=8, n_blocks=128, mixed=True), + ), + Cell( + why="one decode row at the 128-block edge", + spec=dict(total_q=1, n_blocks=128, mixed=False), + ), + Cell(why="a 4-row verify at 40 blocks", spec=dict(total_q=4, n_blocks=40, mixed=False)), + ) + + note = """ + Exact: ties between equal scores break toward the lower block id; `reference` sorts with the same rule. + """ + + def __call__( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous() + return torch.ops.trtllm.minimax_m3_select_blocks( + max_score_kv, nvb, topk, init_blocks, local_blocks, head_major_output + ) + + def reference( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + """Plain torch selection; ``head_major_output`` changes the backing only.""" + heads, n_blocks, total_q = max_score_kv.shape + out = torch.full((total_q, heads, topk), -1, dtype=torch.int32, device=max_score_kv.device) + blocks = torch.arange(n_blocks, device=max_score_kv.device) + for t in range(total_q): + nv = int(n_valid_blocks[t]) + for h in range(heads): + s = max_score_kv[h, :, t].clone().double() + s[blocks >= nv] = float("-inf") + forced = (blocks < min(init_blocks, nv)) | ( + (blocks >= nv - local_blocks) & (blocks < nv) + ) + s[forced] = float("inf") + order = sorted(range(nv), key=lambda b: (-float(s[b]), b))[: min(topk, nv)] + ids = torch.tensor(sorted(order), dtype=torch.int32, device=out.device) + out[t, h, : ids.numel()] = ids + return out + + +m3_select_blocks = M3SelectBlocks() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_select_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_select_long.py new file mode 100644 index 000000000000..91c1817f1610 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_select_long.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 MSA block selection for decode rows of up to 9216 blocks.""" + +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 + +#: Blocks per decode row the op covers (9 keys per thread of one 1024-thread CTA). +MAX_BLOCKS = 9216 +#: Blocks the selection keeps per (query, KV head). +TOPK = 16 +#: Scores the forcing substitutes: initial blocks rank above local ones, and both above any finite score. +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 + +# M3's checkpoint: no forced initial blocks, the one block holding the query forced in. +_M3 = dict(num_kv_heads=1, init_blocks=0, local_blocks=1) + + +class M3SelectLong(OpWrapper): + """Fill ``out`` [total_q, num_kv_heads, 16] with each query's top-16 blocks. + + ``scores`` is the per-KV-head block maxima [num_kv_heads, num_blocks, + total_q] the CuTe DSL index scorer wrote (amax over the KV head's index + heads). Per (query, KV head) row the op keeps the ``n_valid_blocks[query]`` + leading blocks (clamped to ``[0, num_blocks]``), forces the last + ``local_blocks`` of them (counted from the *unclamped* valid count) to + ``LOCAL_SCORE`` and the first ``init_blocks`` to ``INIT_SCORE`` -- local wins + where both apply -- and selects the 16 largest, ties to the smaller block. + A pick whose forced score is -inf is emitted as -1. The ids are written + ascending with a -1 tail. + + Same contract as ``torch.ops.trtllm.minimax_m3_select_blocks`` (the in-tree + ``select_blocks_from_maxscore``) and bit-identical to it, except on a row + where more than 1024 valid blocks tie at the 16th-largest score: this op + breaks that tie to the smallest blocks as the contract says, the in-tree op + does not. The test checks both against ``reference``. Returns ``out``. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="32K-context engine (257 worst-case k-tiles), Eagle3 verify at one request: 4 query rows", + spec=dict(num_blocks=257, total_q=4, valid="near_full", scores="random", **_M3), + ), + Cell( + why="decode without speculation at batch 8 on a 2048-block engine: the 2-key-per-thread path at its edge", + spec=dict(num_blocks=2048, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="1M-context engine at batch 8 x qlen 4: rows on all three key widths, short rows and padding rows", + spec=dict(num_blocks=8193, total_q=32, valid="mixed", scores="random", **_M3), + ), + Cell( + why="duplicated scores: ties broken to the smaller block, and > 1024 keys at the bound (re-read path)", + spec=dict(num_blocks=4500, total_q=4, valid="near_full", scores="ties", **_M3), + ), + Cell( + why="non-finite scores: +inf, NaN ranked above +inf, and rows whose picks run into -inf (emitted -1)", + spec=dict(num_blocks=600, total_q=8, valid="mixed", scores="nonfinite", **_M3), + ), + ) + + note = """ + Certified at one KV head per rank (TP4) with M3's forcing (`init_blocks=0`, + `local_blocks=1`). The kernel takes any head count and forcing; neither is + driven here. + + `scores` may be strided (the op reads it by stride); the target hands over + the scorer's contiguous view, which is what the cells drive. Rows of at + most 128 blocks are within the kernel's range, but the target keeps those on + the in-tree op, so they are not certified here. + + Ordering is by the descending-score key `orderedKey`: a positive NaN ranks + above +inf, +0 above -0, and a negative NaN below -inf -- which, unlike + -inf, is emitted as a block id rather than -1. The cells drive positive + NaN only. + """ + + def __call__( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_select_long( + scores, n_valid_blocks, out, init_blocks, local_blocks, use_pdl, early_trigger + ) + return out + + def reference( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + ) -> torch.Tensor: + """The [total_q, num_kv_heads, 16] int32 table the op is supposed to write. + + Exact integer ranking, no floating point beyond the forcing: each + score's bit pattern maps to an unsigned key that is smaller for a larger + score (NaN above +inf), packed with the block id so one sort ranks by + (score descending, block ascending). ``out``, ``use_pdl`` and + ``early_trigger`` only say where and how the op writes; they do not enter the result. + """ + num_kv_heads, num_blocks, total_q = scores.shape + device = scores.device + blocks = torch.arange(num_blocks, device=device, dtype=torch.int64) + raw = n_valid_blocks[:total_q].to(torch.int64) + valid = raw.clamp(0, num_blocks) + local_start = (raw - local_blocks).clamp_min(0) + result = torch.full((total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device=device) + for head in range(num_kv_heads): + forced = scores[head].float().t().clone() # [total_q, num_blocks] + forced = torch.where(blocks[None, :] < init_blocks, INIT_SCORE, forced) + forced = torch.where(blocks[None, :] >= local_start[:, None], LOCAL_SCORE, forced) + bits = forced.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + key = torch.where(bits >= 0x80000000, bits, (~bits) & 0x7FFFFFFF) + in_row = blocks[None, :] < valid[:, None] + packed = torch.where( + in_row, (key << 16) | blocks[None, :], torch.iinfo(torch.int64).max + ) + k = min(TOPK, num_blocks) + picked = torch.topk(packed, k, dim=1, largest=False).values + picked_key = picked >> 16 + keep = (picked != torch.iinfo(torch.int64).max) & (picked_key != 0xFF800000) + ids = torch.where(keep, picked & 0xFFFF, torch.iinfo(torch.int64).max) + ids = ids.sort(dim=1).values + ids = torch.where(ids == torch.iinfo(torch.int64).max, -1, ids) + result[:, head, :k] = ids.to(torch.int32) + return result + + def is_valid( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + ) -> None: + """The inputs this op takes and answers wrongly. + + Negative forcing counts are accepted and push the forced ranges past the + row instead of failing; the in-tree op this one stands in for refuses + them, so a caller relying on this one would silently diverge from it. + An ``out`` overlapping ``scores`` is overwritten while it is still read. + """ + assert init_blocks >= 0 and local_blocks >= 0, ( + f"init_blocks ({init_blocks}) and local_blocks ({local_blocks}) must be non-negative; the " + "op accepts negative counts and selects as if the forced ranges were elsewhere" + ) + storage = scores.untyped_storage() + assert not (storage.data_ptr() <= out.data_ptr() < storage.data_ptr() + storage.nbytes()), ( + "out must not alias scores: the op writes it while scores are still read" + ) + + +m3_select_long = M3SelectLong() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode.py new file mode 100644 index 000000000000..b82d733c263b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This target's block-sparse decode attention (Triton, with the one-launch long-context route).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..triton_sparse_decode import minimax_m3_sparse_attn_decode +from .m3_sparse_decode_lc import PAGE, sparse_attention_reference + + +class M3SparseDecode(OpWrapper): + """``minimax_m3_sparse_attn_decode`` as ``M3Attention._decode_sparse_core`` calls it. + + Attention of each query token over the blocks ``topk_idx`` selected for it + (``[num_kv_heads, T, topk]``, ascending ids with a -1 tail), each block + truncated at the token's causal extent, written into ``output`` in place. + The target passes FP8 q, the packed NVFP4 sparse-layer cache with its block + and per-tensor scales, and BF16 ``output`` [T, 16, 128]. Returns ``output``. + + Which kernel runs is decided inside, per step: a selected (non-identity) + table on the TP4 geometry goes to the one-launch ``m3_sparse_decode_lc`` + kernel; an identity table (every request within 16 blocks, staged by the + metadata) runs the Triton producer, or from 17 query rows the multi-query + producer, then the LSE merge. Every route also publishes the MXFP8 o_proj + hand-off to ``m3_prequant.SLOT``; this entry certifies the attention output. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="identity selection, decode without speculation at batch 8: the Triton producer + merge", + spec=dict( + seq_lens=(1, 128, 129, 700, 2048, 1500, 64, 0), decode_query_len=1, identity=True + ), + ), + Cell( + why="identity selection, Eagle3 verify at one request (T=4): one block per producer program", + spec=dict(seq_lens=(1900,), decode_query_len=4, identity=True), + ), + Cell( + why="identity selection, Eagle3 verify at batch 4 (T=16): the last row count below the MQ producer", + spec=dict(seq_lens=(300, 2048, 4, 1000), decode_query_len=4, identity=True), + ), + Cell( + why="identity selection, Eagle3 verify at batch 8 (T=32): the multi-query producer, padding rows", + spec=dict( + seq_lens=(2048, 130, 5, 0, 1777, 900, 2000, 0), decode_query_len=4, identity=True + ), + ), + Cell( + why="selected blocks, Eagle3 verify at batch 2: routed to the one-launch long-context kernel", + spec=dict(seq_lens=(9000, 40000), decode_query_len=4, identity=False), + ), + Cell( + why="selected blocks, decode without speculation at batch 8: the long-context route at T=8", + spec=dict( + seq_lens=(3000, 2100, 9999, 2048, 2049, 4500, 16000, 2200), + decode_query_len=1, + identity=False, + ), + ), + ) + + note = """ + Certified only as the target calls it: one KV head of 16 query heads, + head_dim 128, FP8 q, packed NVFP4 cache, BF16 output, `num_topk_chunks` + left to the dispatcher. The function also takes BF16 / FP8 caches, other + head geometries and an explicit split-K factor (the in-tree tests drive + those); none is certified here. + + The MXFP8 hand-off each route publishes is certified on the long-context + kernel's own entry (`m3_sparse_decode_lc`); the Triton merge writes it with + the same recipe (`_store_mxfp8_head`). + """ + + def __call__( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + minimax_m3_sparse_attn_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + kv_scale=kv_scale, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + num_topk_chunks=num_topk_chunks, + identity_blocks=identity_blocks, + ) + return output + + def reference( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + """``output`` as the call is supposed to write it: float64 block-sparse attention over the + dequantized NVFP4 blocks (``sparse_attention_reference``), rounded once to ``output``'s dtype. + + Only the packed NVFP4 cache with one KV head is mirrored -- the target's + only configuration. ``kv_scale``, ``num_topk_chunks`` and + ``identity_blocks`` do not change what is computed. + """ + assert k_block_scale is not None and k_paged.shape[1] == 1, ( + "mirrors the one-KV-head NVFP4 cache only" + ) + attn = sparse_attention_reference( + q, + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx[0], + block_table, + seq_lens, + decode_query_len, + sm_scale, + ) + return attn.to(output.dtype) + + def is_valid( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> None: + """The inputs this call takes and answers wrongly. + + Every producer walks a token's first ``min(topk, ceil(kv_len / 128))`` + table entries and indexes the block table with them unchecked, so they + must be real blocks of the request. ``identity_blocks`` is a promise the + multi-query producer relies on: every token of a request reads the same + blocks, all of them. Reads ``topk_idx`` and ``seq_lens`` back to the + host: a test-only cost. + """ + total_q = q.shape[0] + seq = seq_lens.tolist() + max_topk = topk_idx.shape[-1] + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk_idx[:, tok, : min(max_topk, nblocks)] + assert bool(((ids >= 0) & (ids < min(nblocks, block_table.shape[1]))).all()), ( + f"token {tok}: the first min(topk, block count) selected ids must be blocks of the request; " + "the block table is indexed with them unchecked" + ) + if identity_blocks: + assert nblocks <= max_topk and bool( + (ids == torch.arange(nblocks, device=ids.device, dtype=ids.dtype)).all() + ), ( + f"token {tok}: identity_blocks promises the table holds every block of the request in order" + ) + + +m3_sparse_decode = M3SparseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode_lc.py new file mode 100644 index 000000000000..2dfeba5cbd0a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode_lc.py @@ -0,0 +1,395 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-launch long-context sparse decode over the selected NVFP4 blocks (16-CTA clusters). + +Also the plain-torch block-sparse attention and MXFP8 hand-off mirrors the +Triton sparse decode entry (``m3_sparse_decode``) reuses: both compute the same +thing over the same cache. +""" + +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 + +#: The TP4 geometry the kernel is compiled for. +NUM_HEADS = 16 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +#: NVFP4 elements per E4M3 block scale. +SF_VEC = 16 +#: MXFP8 elements per UE8M0 scale, and the o_proj width the hand-off is laid out for. +MX_BLOCK = 32 +WIDTH = NUM_HEADS * HEAD_DIM + +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + + +def dequant_nvfp4_pages( + packed: torch.Tensor, block_scale: torch.Tensor, pages: torch.Tensor, quad_scales: bool +) -> torch.Tensor: + """Dequantize KV head 0 of the given pages of a packed NVFP4 cache to float64 [n, PAGE, HEAD_DIM]. + + ``packed`` is the [pages, kv_heads, PAGE, HEAD_DIM // 2] byte view (low + nibble = even element), ``block_scale`` the [pages, kv_heads, PAGE, + HEAD_DIM // 16] E4M3 bytes ``fused_write_layer_caches_nvfp4`` writes: token + ``t``'s column ``c`` at byte ``t * 8 + c`` of the page's region for K, and at + ``(t // 4) * 32 + c * 4 + t % 4`` for V (``quad_scales``). The per-tensor + global scale is not applied. Every product is exact in float64. + """ + table = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) + rows = packed[pages, 0].to(torch.int64) # [n, PAGE, HEAD_DIM // 2] + nibbles = torch.stack([rows & 0x0F, rows >> 4], dim=-1).reshape(rows.shape[0], PAGE, HEAD_DIM) + cols = HEAD_DIM // SF_VEC + region = block_scale[pages, 0].reshape(rows.shape[0], PAGE * cols) + t = torch.arange(PAGE, device=packed.device)[:, None] + c = torch.arange(cols, device=packed.device)[None, :] + index = (t // 4) * 32 + c * 4 + t % 4 if quad_scales else t * cols + c + scales = region[:, index.reshape(-1)].reshape(rows.shape[0], PAGE, cols) + scales = scales.view(torch.float8_e4m3fn).to(torch.float64) + return table[nibbles] * scales.repeat_interleave(SF_VEC, dim=-1) + + +def sparse_attention_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, +) -> torch.Tensor: + """Block-sparse decode attention over one KV head of a packed NVFP4 cache, float64 [T, heads, HEAD_DIM]. + + Query token ``tok`` belongs to request ``tok // decode_query_len`` at offset + ``tok % decode_query_len`` and attends the first ``min(16, ceil(kv_len / + PAGE))`` entries of ``topk[tok]`` that are not -1, each block truncated at + ``kv_len = max(seq_len - decode_query_len + offset + 1, 0)``. A token that + attends nothing (a CUDA-graph padding row) is all zeros. ``q`` is E4M3 (or + its byte view); softmax and both products are float64. + """ + q64 = ( + q.view(torch.float8_e4m3fn).to(torch.float64) + if q.element_size() == 1 + else q.to(torch.float64) + ) + total_q, num_heads, _ = q64.shape + kgs = float(k_global_scale.reshape(-1)[0]) + vgs = float(v_global_scale.reshape(-1)[0]) + out = torch.zeros(total_q, num_heads, HEAD_DIM, dtype=torch.float64, device=q.device) + seq = seq_lens.tolist() + positions = torch.arange(PAGE, device=q.device) + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + real_topk = min(TOPK, (kv_len + PAGE - 1) // PAGE) + blocks = topk[tok, :real_topk] + blocks = blocks[blocks >= 0].to(torch.int64) + if blocks.numel() == 0: + continue + pages = block_table[req, blocks].to(torch.int64) + keys = ( + dequant_nvfp4_pages(k, k_block_scale, pages, quad_scales=False).reshape(-1, HEAD_DIM) + * kgs + ) + values = ( + dequant_nvfp4_pages(v, v_block_scale, pages, quad_scales=True).reshape(-1, HEAD_DIM) + * vgs + ) + valid = ((blocks[:, None] * PAGE + positions[None, :]) < kv_len).reshape(-1) + logits = (q64[tok] @ keys.T) * sm_scale + logits = logits.masked_fill(~valid[None, :], float("-inf")) + out[tok] = torch.softmax(logits, dim=-1) @ values + return out + + +def mxfp8_handoff_reference( + out: torch.Tensor, num_rows_alloc: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """The MXFP8 o_proj hand-off the merge writes for BF16 rows ``out`` [T, WIDTH]: ``(q8 [T, WIDTH] uint8, + sf [num_rows_alloc * WIDTH // 32] uint8)``. + + Per 32-element block: ``e`` is the biased exponent of ``amax * rcp(448)`` + rounded up when its mantissa is non-zero (capped at 254), the elements are + scaled by ``2^(127 - e)`` and rounded to E4M3 (saturating), and ``e`` lands + in the 128x4-swizzled scale layout the MXFP8 GEMM reads. Entries of ``sf`` + for rows past ``T`` are zero. Computed from the op's own BF16 rows, so a + comparison against it is exact. + """ + total_q = out.shape[0] + x = out.float().reshape(total_q, WIDTH // MX_BLOCK, MX_BLOCK) + amax = x.abs().amax(dim=-1) + sfv = amax * torch.tensor(1.0 / 448.0, dtype=torch.float32, device=out.device) + bits = sfv.view(torch.int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) + e = e.clamp(max=254) + qs = torch.where(amax != 0, ((254 - e) << 23).view(torch.float32), torch.zeros_like(amax)) + xq = (x * qs[..., None]).clamp(-448.0, 448.0) + q8 = xq.to(torch.float8_e4m3fn).view(torch.uint8).reshape(total_q, WIDTH) + num_sf_k_tiles = WIDTH // 128 + sf = torch.zeros(num_rows_alloc * (WIDTH // MX_BLOCK), dtype=torch.uint8, device=out.device) + tok = torch.arange(total_q, device=out.device)[:, None] + k_sf = torch.arange(WIDTH // MX_BLOCK, device=out.device)[None, :] + offset = ( + (tok // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (tok % 32) * 16 + + ((tok % 128) // 32) * 4 + + k_sf % 4 + ) + sf[offset.reshape(-1)] = e.to(torch.uint8).reshape(-1) + return q8, sf + + +class M3SparseDecodeLc(OpWrapper): + """Sparse decode attention, split-K merge and MXFP8 hand-off of one TP4 decode step in one launch. + + One cluster of 16 CTAs per query token: CTA ``r`` attends the token's + ``r``-th selected block (``topk[tok, r]``), and the cluster merges the 16 + partials over distributed shared memory into the BF16 row ``out[tok]``. With + ``q8`` / ``sf`` it also writes the row's MXFP8 copy and swizzled UE8M0 + scales for the o_proj GEMM. All operands are the byte views + ``m3_attn_lc.m3_lc_attn`` hands over (E4M3 ``q``, packed E2M1 pages, E4M3 + block scales). Returns ``out``. + + Per block: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16 + (exact), QK and PV on FP16 tensor cores with FP32 accumulation, the softmax + probabilities rounded to FP16 for PV, log2-domain softmax with ``ex2.approx``. + ``linear_softmax`` picks the denominator recurrence (the two agree to + rounding); the target picks it per batch shape. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode without speculation, one request: T=1, log-domain recurrence, MXFP8 hand-off", + spec=dict(seq_lens=(9000,), decode_query_len=1, linear_softmax=False, mxfp8=True), + ), + Cell( + why="Eagle3 verify at one request: T=4, the log-domain recurrence the target keeps for local batch 1", + spec=dict(seq_lens=(20000,), decode_query_len=4, linear_softmax=False, mxfp8=True), + ), + Cell( + why="Eagle3 verify at two requests: T=8 on the linear denominator", + spec=dict(seq_lens=(5000, 12345), decode_query_len=4, linear_softmax=True, mxfp8=True), + ), + Cell( + why="Eagle3 verify at batch 8 (T=32): short requests with -1 tails and zero-length padding rows", + spec=dict( + seq_lens=(16384, 2049, 130, 0, 7777, 4, 30000, 0), + decode_query_len=4, + linear_softmax=True, + mxfp8=True, + ), + ), + Cell( + why="decode without speculation at batch 8: T=8 rows of distinct requests", + spec=dict( + seq_lens=(3000, 2100, 9999, 2048, 2049, 4500, 16000, 2200), + decode_query_len=1, + linear_softmax=False, + mxfp8=True, + ), + ), + Cell( + why="no hand-off buffers (the variant the call site takes past 128 rows), driven at T=8", + spec=dict(seq_lens=(5000, 12345), decode_query_len=4, linear_softmax=True, mxfp8=False), + ), + ) + + note = """ + Certified: the four instantiated variants (`linear_softmax` x hand-off on / + off), all with `early_trigger=True` -- the op refuses `False`. The + hand-off cells are certified bit-exactly against `mxfp8_handoff_reference` + applied to the op's own BF16 rows; the BF16 rows themselves within the band + the test derives. + + `topk` is read for the first `min(16, ceil(kv_len / 128))` entries of each + token; a -1 among them skips that CTA's block, entries past them are never + read. `use_pdl` only changes the launch attribute. The hand-off assumes + `num_sf_k_tiles = 16` (a 2048-wide o_proj input) and buffers with at least + as many rows as the step; `MERGE_MXFP8` sizes them at 128 rows. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + early_trigger: bool, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_sparse_decode_lc( + q, + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + out, + q8, + sf, + num_sf_k_tiles, + decode_query_len, + sm_scale, + linear_softmax, + use_pdl, + early_trigger, + ) + return out + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + early_trigger: bool, + ) -> torch.Tensor: + """The BF16 ``out`` [T, 2048] the op is supposed to write: ``sparse_attention_reference`` in float64, + rounded once. ``linear_softmax``, ``use_pdl`` and ``early_trigger`` do not change what is computed; the + hand-off buffers are checked separately with ``mxfp8_handoff_reference``. + """ + attn = sparse_attention_reference( + q.view(-1, NUM_HEADS, HEAD_DIM), + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + decode_query_len, + sm_scale, + ) + return attn.reshape(out.shape).to(out.dtype) + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + early_trigger: bool, + ) -> None: + """The inputs this op takes and answers wrongly (``m3_attn_lc.lc_applies`` is the target's form of the + layout half). + + The kernel addresses a page's token rows at a fixed 64-byte (data) and + 8-byte (scale) pitch and loads them 16 / 8 bytes at a time, so only the + page stride is read from the views. It reads the selected block ids, + the block table and the hand-off buffers unchecked. Reads ``topk`` and + ``seq_lens`` back to the host: a test-only cost. + """ + for name, t in (("k", k), ("v", v)): + assert t.dim() == 4 and t.shape[1] == 1 and t.shape[2:] == (PAGE, HEAD_DIM // 2), ( + f"{name} must be the [pages, 1, 128, 64] byte view of one KV head" + ) + assert t.stride(3) == 1 and t.stride(2) == HEAD_DIM // 2, ( + f"{name} rows must be dense 64-byte rows: the kernel ignores the row stride" + ) + assert t.data_ptr() % 16 == 0 and t.stride(0) % 16 == 0, ( + f"{name} pages must be 16-byte aligned" + ) + for name, t in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + assert t.dim() == 4 and t.stride(3) == 1 and t.stride(2) == HEAD_DIM // SF_VEC, ( + f"{name} rows must be dense 8-byte rows: the kernel ignores the row stride" + ) + assert t.data_ptr() % 8 == 0 and t.stride(0) % 8 == 0, ( + f"{name} pages must be 8-byte aligned" + ) + assert q.data_ptr() % 16 == 0 and q.stride(0) % 16 == 0 and q.stride(1) % 16 == 0, ( + "q rows are loaded 16 bytes at a time and must be 16-byte aligned" + ) + assert k_global_scale.numel() == 1 and v_global_scale.numel() == 1, ( + "global scales are single values" + ) + total_q = q.shape[0] + assert ( + total_q % decode_query_len == 0 and seq_lens.numel() >= total_q // decode_query_len + ), ( + "every query token must belong to a request: seq_lens is read at tok // decode_query_len unchecked" + ) + seq = seq_lens.tolist() + num_cols = block_table.shape[1] + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk[tok, : min(TOPK, nblocks)] + assert bool(((ids >= -1) & (ids < min(nblocks, num_cols))).all()), ( + f"token {tok}: selected block ids must be -1 or below the request's block count; the block " + "table is indexed with them unchecked" + ) + if q8 is not None: + assert num_sf_k_tiles == WIDTH // 128, ( + f"num_sf_k_tiles must be {WIDTH // 128}: the hand-off rows are {WIDTH} wide and a smaller " + "value overlaps them silently" + ) + assert ( + q8.numel() >= total_q * WIDTH + and sf is not None + and sf.numel() >= ((total_q + 127) // 128 * 128 * (WIDTH // MX_BLOCK)) + ), "the hand-off buffers must hold every row of the step: they are written unchecked" + + +m3_sparse_decode_lc = M3SparseDecodeLc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/msa_kv_update.py new file mode 100644 index 000000000000..ad05b934d4a9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/msa_kv_update.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MSA KV-length correction in one launch (this target's ``MsaKvUpdate`` hook).""" + +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 + +#: The MSA pool's page: one sparse block is one 128-token KV page. +TOKENS_PER_BLOCK = 128 + + +class MsaKvUpdate(OpWrapper): + """Re-derive the MSA decode state after the overlap scheduler moved ``kv_lens``. + + ``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` calls this with its + own buffers at the head of every captured speculative decode step. One + single-CTA launch writes, for ``batch`` requests and ``total_q`` query + tokens (all int32): + + * ``msa_seq_lens[b] = min(kv_lens[b], kv_lens_staged[b])`` and, when given, + ``cu_kv_lens`` as its exclusive prefix sum (``batch + 1`` entries); + * ``out_cache_loc[t]``, the KV write slot of the position token ``t`` + attends up to, through ``block_table``; + * ``n_valid_blocks[t]``, when given, the block count top-k selection is + bounded by. + + It returns None; the cells read the four buffers. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode without speculation at one request: qo = 1, every buffer written", + spec=dict(batch=1, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="decode without speculation at the captured batch 8", + spec=dict(batch=8, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="Eagle3 verify (draft 3): qo = 4 per request, partial acceptance shortens kv_lens below staged", + spec=dict(batch=8, query_len=4, accept="partial", cu_kv=True, n_valid=True), + ), + Cell( + why="the metadata without cu_kv_lens (no extend_ctx promotion): the optional prefix sum is skipped", + spec=dict(batch=4, query_len=4, accept="partial", cu_kv=False, n_valid=True), + ), + Cell( + why="CUDA-graph padding rows (kv_lens 0, pos -1) and no n_valid_blocks buffer", + spec=dict(batch=8, query_len=4, accept="padded", cu_kv=True, n_valid=False), + ), + Cell( + why="positions crossing a page boundary: the slot moves to the previous page of the request", + spec=dict(batch=3, query_len=4, accept="page_edge", cu_kv=True, n_valid=True), + ), + ) + + note = """ + The kernel covers any batch up to the block size's loop (it strides over + 256 threads), and any page size; only `tokens_per_block = 128`, the MSA + pool's page, is certified. + + The op does not bounds-check the request index a token reads + (`q_batch_row`) or the block-table column it derives: both come from the + metadata's own staging, which keeps them in range by construction, and an + out-of-range value reads whatever memory sits there. `is_valid` states both. + """ + + def __call__( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + torch.ops.trtllm.minimax_m3_msa_kv_update( + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + tokens_per_block, + batch, + total_q, + msa_seq_lens, + cu_kv_lens, + out_cache_loc, + n_valid_blocks, + ) + + def reference( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]: + """``(msa_seq_lens[:batch], cu_kv_lens[:batch + 1], out_cache_loc[:total_q], + n_valid_blocks[:total_q])`` as the op is supposed to write them. + + The in-tree metadata's own torch re-derivation + (``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` past its hook + call), statement for statement, returning new tensors rather than + writing the buffers. The buffer arguments only say which outputs exist: + a ``None`` buffer gives a ``None`` result. + """ + page = int(tokens_per_block) + kv_true = torch.minimum(kv_lens[:batch], kv_lens_staged[:batch]) + cu_kv = None + if cu_kv_lens is not None: + cu_kv = torch.zeros(batch + 1, dtype=torch.int32, device=kv_lens.device) + cu_kv[1:] = torch.cumsum(kv_true, 0) + qbr = q_batch_row[:total_q].to(torch.long) + qo_dev = seq_lens[:batch] + pos = kv_true[qbr] - qo_dev[qbr] + q_intra[:total_q] + pos_slot = pos.clamp_min(0) + block_col = torch.div(pos_slot, page, rounding_mode="floor").clamp( + max=int(block_table.shape[1]) - 1 + ) + slots = block_table[qbr, block_col.to(torch.long)] * page + torch.remainder(pos_slot, page) + n_valid = None + if n_valid_blocks is not None: + n_valid = torch.div( + (pos + 1).clamp_min(0) + (page - 1), page, rounding_mode="floor" + ).to(torch.int32) + return kv_true.to(torch.int32), cu_kv, slots.to(torch.int32), n_valid + + def is_valid( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + """The inputs this op takes and answers wrongly. + + Every operand but the block table is read as a dense run from its first + element, so a strided view reads the wrong elements. The request index of + each token and the table row it selects are not bounds-checked. Reads + ``q_batch_row`` back to the host: a test-only cost. + """ + for name, t in ( + ("kv_lens", kv_lens), + ("kv_lens_staged", kv_lens_staged), + ("seq_lens", seq_lens), + ("q_batch_row", q_batch_row), + ("q_intra", q_intra), + ("msa_seq_lens", msa_seq_lens), + ("out_cache_loc", out_cache_loc), + ("cu_kv_lens", cu_kv_lens), + ("n_valid_blocks", n_valid_blocks), + ): + assert t is None or t.is_contiguous(), ( + f"{name} must be contiguous: the kernel indexes it densely from data_ptr and ignores strides" + ) + rows = q_batch_row[:total_q] + assert total_q == 0 or ( + int(rows.min()) >= 0 and int(rows.max()) < min(batch, block_table.shape[0]) + ), ( + "q_batch_row must index a request below batch (and a block-table row): the op reads the " + "request's lengths and table row unchecked" + ) + + +msa_kv_update = MsaKvUpdate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/__init__.py new file mode 100644 index 000000000000..5da86fe24bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collective-communication entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_allreduce.py new file mode 100644 index 000000000000..1a275c4ad863 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_allreduce.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The built-in TP all-reduce, fused with the residual add and the next RMSNorm.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.distributed import AllReduce, AllReduceParams + + +class M3AllReduce(OpWrapper): + """``comm(x, all_reduce_params=params)``: the runtime's ``AllReduce`` over the TP group. ``params`` None: the + plain sum over the ranks (the drafter's o_proj and down projection). With ``RESIDUAL_RMS_NORM`` params: + ``updated = residual + sum_ranks(x)``, ``normed = rmsnorm(updated) * norm_weight``; + returns (normed, updated). The strategy (``AllReduceStrategy.AUTO``: one-shot / two-shot Lamport, NCCL) is the + communicator's per call. + + The layer boundaries the multicast all-reduce does not take run it: prefill and mixed steps, and generation steps + beyond the multicast kernel's rows. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="tp4, a prefill chunk", spec=dict(world=4, m=2048, hidden=6144)), + Cell(why="tp4, a 64-row generation step", spec=dict(world=4, m=64, hidden=6144)), + Cell( + why="tp4, plain sum of eight drafter rows", + spec=dict(world=4, m=8, hidden=6144, plain=True), + ), + ) + + note = """ + The reduction order across ranks is the strategy's; `reference` sums the ranks' rows in fp32 in rank order, then + adds the residual and normalizes (one bf16 rounding of the sum and one of the norm). A strategy may round the sum + to bf16 before the residual add, so the residual sum agrees within one bf16 ulp of its operands' magnitude. + """ + + def __call__(self, x: torch.Tensor, comm: AllReduce, params: Optional[AllReduceParams]): + return comm(x, all_reduce_params=params) + + def reference(self, xs, comm: AllReduce, params: Optional[AllReduceParams]): + """``xs``: every rank's ``x`` in rank order (a collective's reference takes the group's inputs).""" + total = torch.stack([x.float() for x in xs]).sum(0) + if params is None: + return total.to(xs[0].dtype) + updated = (params.residual.float() + total).to(xs[0].dtype) + h = updated.float() + normed = ( + h + * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + params.eps) + * params.norm_weight.float() + ) + return normed.to(xs[0].dtype), updated + + +m3_allreduce = M3AllReduce() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_mc_allreduce_norm.py new file mode 100644 index 000000000000..0da06e59b14e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_mc_allreduce_norm.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode layer boundary, 4 ranks.""" + +from __future__ import annotations + +from typing import 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 + +_E4M3_MAX = 448.0 + + +def mxfp8_rows(normed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """MXFP8 of BF16 rows as ``trtllm::mxfp8_quantize(normed, True)``: E4M3 [M, H] and 128x4-swizzled UE8M0 bytes. + + One scale per 32 columns, the smallest power of two at or above ``amax / 448`` (biased exponent, round up), + data = row / scale rounded to nearest E4M3. The scale buffer is padded to 128 rows, padding bytes left as the + op leaves them (the comparison masks them out). + """ + m, h = normed.shape + blocks = normed.float().view(m, h // 32, 32) + amax = blocks.abs().amax(-1) + bits = (amax / _E4M3_MAX).view(torch.int32) + byte = (((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) > 0).to(torch.int32)).clamp(0, 254) + scale = torch.exp2(byte.float() - 127.0).unsqueeze(-1) + data = torch.where(amax.unsqueeze(-1) != 0, blocks / scale, torch.zeros_like(blocks)) + data = data.view(m, h).to(torch.float8_e4m3fn) + return data, swizzle_sf(byte.to(torch.uint8)) + + +def swizzle_sf(byte: torch.Tensor) -> torch.Tensor: + """Scatter [rows, H / 32] scale bytes into the 128x4-swizzled buffer the op writes (zero padding).""" + rows, cols = byte.shape + padded_rows = (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (r // 128) * (128 * cols) + ) + out = torch.zeros(padded_rows * cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def sf_written(rows: int, cols: int, device) -> torch.Tensor: + """Mask of the swizzled scale bytes the op writes for ``rows`` rows (the padding rows are left unwritten).""" + return swizzle_sf(torch.ones(rows, cols, dtype=torch.uint8, device=device)).bool() + + +def boundary( + partials: Sequence[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, +) -> tuple[torch.Tensor, ...]: + """The boundary arithmetic once each rank's BF16 partial is known (shared with the finalize entry). + + Ranks are summed in FP32 in rank order, the residual added and rounded to BF16 once; the RMS statistic and the + normalization are done in float64 and rounded to BF16 once. + """ + total = torch.zeros(partials[0].shape, dtype=torch.float32, device=residual.device) + for p in partials: + total = total + p.float() + updated = (total + residual.float()).to(torch.bfloat16) + u = updated.double() + inv = torch.rsqrt(u.pow(2).mean(-1, keepdim=True) + eps) + normed = (u * inv * norm_w.double()).to(torch.bfloat16) + if not quant: + return normed, updated + return (normed, updated, *mxfp8_rows(normed)) + + +class M3McAllreduceNorm(OpWrapper): + """One TP group's boundary: ``updated = bf16(residual + sum_r bf16(x_r + x2_r))``, ``normed = RMSNorm(updated)``. + + Every rank pushes its partial through the multicast mapping ``ws_mc_ptr`` into every rank's exchange, then polls + its own view ``ws_uc`` (re-arming the words it read) and rotates the per-CTA call parities in ``flags``; the + workspace is the ``McAllReduceWorkspace`` bucket of ``maxt`` tokens. All ranks must make the same calls in the + same stream order. ``prepushed`` 2 means a producer already pushed ``x`` (and launched this grid after its + pushes); 3 that two producers pushed ``x`` into exchange set A and ``x2`` into set B. Returns + ``[normed, updated]`` or, with ``quant``, also the MXFP8 copy of ``normed`` (E4M3 rows, swizzled UE8M0). + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The tp4 decode boundaries (hidden 6144, 4 ranks, eps 1e-6). Buckets 8 / 16 / 48 are separate workspaces; the + # 8-token bucket runs epilogue option 7 at 1, 2, 7, 8 tokens and 15 at 3..6, the others 15. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token at the attention boundary, x only, bucket 8 option 7", + spec=dict(m=1, maxt=8, opts=7, x2=False, quant=False, prepushed=0), + ), + Cell( + why="MoE boundary at 4 tokens: shared + routed summed in-kernel, MXFP8 copy for the next QKV, option 15", + spec=dict(m=4, maxt=8, opts=15, x2=True, quant=True, prepushed=0), + ), + Cell( + why="a full 8-token bucket with the second addend", + spec=dict(m=8, maxt=8, opts=7, x2=True, quant=False, prepushed=0), + ), + Cell( + why="attention boundary whose rows the o_proj GEMV already pushed (prepushed 2)", + spec=dict(m=8, maxt=8, opts=7, x2=False, quant=False, prepushed=2), + ), + Cell( + why="routed engine pushed set A, shared down projection set B (prepushed 3), with the MXFP8 copy", + spec=dict(m=6, maxt=8, opts=15, x2=True, quant=True, prepushed=3), + ), + Cell( + why="12 tokens pushed by the o_proj GEMV into the 16-token bucket", + spec=dict(m=12, maxt=16, opts=15, x2=False, quant=False, prepushed=2), + ), + Cell( + why="the 9..13-token MoE kernel's sum into the 16-token bucket, MXFP8 copy", + spec=dict(m=16, maxt=16, opts=15, x2=False, quant=True, prepushed=0), + ), + Cell( + why="the 48-token bucket at a partial fill", + spec=dict(m=30, maxt=48, opts=15, x2=False, quant=False, prepushed=0), + ), + Cell( + why="the 48-token bucket full, second addend and MXFP8 copy", + spec=dict(m=48, maxt=48, opts=15, x2=True, quant=True, prepushed=0), + ), + # Every other row count of the decode boundaries (dispatch.MC_ALLREDUCE): the bucket and epilogue options the + # target picks for it (ops/m3_mc_allreduce), the second addend and the MXFP8 copy alternating. + *( + Cell( + why=f"{m} tokens: a boundary row count between the ones above", + spec=dict( + m=m, + maxt=8 if m <= 8 else 16 if m <= 16 else 48, + opts=(7, 7, 15, 15, 15, 15, 7, 7)[m - 1] if m <= 8 else 15, + x2=m % 2 == 0, + quant=m % 3 == 0, + prepushed=0, + ), + ) + for m in range(1, 49) + if m not in (1, 4, 8, 16, 30, 48) + ), + # Every other row count pushed by the o_proj (1..16) or the 1..8-token MoE engine (with the MXFP8 copy): + # dispatch.ATTN_PUSH / dispatch.MOE_PUSH. + *( + Cell( + why=f"{m} tokens prepushed: a pushed row count between the ones above", + spec=dict( + m=m, + maxt=8 if m <= 8 else 16, + opts=(7, 7, 15, 15, 15, 15, 7, 7)[m - 1] if m <= 8 else 15, + x2=False, + quant=m <= 8, + prepushed=2, + ), + ) + for m in range(1, 17) + if m not in (8, 12) + ), + ) + + note = """ + The epilogue option bits (1 early residual / weight load, 2 st.async cluster exchange, 4 early parity rotation, + 8 max carveout) are bit-identical by construction; only the four (bucket, option) pairs the target selects are + built. `prepushed` 1 (pushed, grid wait kept) is accepted by the op but no caller passes it. `pdl` changes only + launch overlap and is driven off in the cells. + + `normed` is gated to one BF16 rounding step: the kernel's RMS statistic is a warp / cluster tree of FP32 sums + and its `rsqrtf` is approximate, so a row whose scaled value lands next to a BF16 midpoint can round either + way. `updated` and the MXFP8 copy of the op's own `normed` are bit-exact. + """ + + def __call__( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_mc_allreduce_norm( + x, + x2, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + None, + None, + maxt, + quant, + prepushed, + opts, + ) + + def reference( + self, + per_rank_x: Sequence[torch.Tensor], + per_rank_x2: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute. + + The documented collective exception: it takes every rank's ``x`` (and ``x2``) in rank order, not the calling + rank's. Each rank's partial is ``bf16(x + x2)`` (``x`` alone without ``x2``), whatever pushed it. + """ + partials = [] + for r, x in enumerate(per_rank_x): + if per_rank_x2 is None: + partials.append(x) + else: + partials.append((x.float() + per_rank_x2[r].float()).to(torch.bfloat16)) + return boundary(partials, residual, norm_w, eps, quant) + + def is_valid( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + ) -> None: + h = residual.shape[-1] + # The op checks norm_w's element count, not its layout: a strided view is read as packed memory. + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + # The workspace is addressed as [set][parity][maxt][world][H / 2] words and never bounds-checked: a smaller + # buffer (another bucket's, or one built for fewer ranks) is polled and re-armed out of range. + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + # Flags are per (token, cluster CTA) of the bucket; the op checks the count, not that they belong to the + # workspace. Flags of another bucket hold another call history, and the parity they select is not the one the + # peers push into: the poll then waits on the other half forever or reads last call's rows. + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + + +m3_mc_allreduce_norm = M3McAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..7e6ebeef7d31 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MoE finalize + multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode MoE boundary, 4 ranks.""" + +from __future__ import annotations + +from typing import 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 .m3_mc_allreduce_norm import boundary + + +def finalize_partial( + fc2: torch.Tensor, perm: torch.Tensor, scales: torch.Tensor, shared: Optional[torch.Tensor] +) -> torch.Tensor: + """One rank's BF16 partial: ``bf16(bf16(sum_k scale_k * fc2[perm_k]) + shared)``, the routed sum in FP32 in k + order and rounded to BF16, then the shared expert added as stock adds it (one more BF16 rounding). + + The kernel's per-k step is a fused multiply-add; it is reproduced exactly by the float64 product (exact for two + FP32 operands) plus the FP32 accumulator, rounded back to FP32. Dropped entries (perm -1) are skipped. + """ + m, topk = perm.shape + acc = torch.zeros(m, fc2.shape[1], dtype=torch.float32, device=fc2.device) + sc = scales.reshape(m, topk).float() + for k in range(topk): + rows = perm[:, k].long() + keep = (rows >= 0).unsqueeze(-1) + gathered = fc2[rows.clamp_min(0)].double() + step = (gathered * sc[:, k : k + 1].double() + acc.double()).float() + acc = torch.where(keep, step, acc) + if shared is not None: + acc = acc.to(torch.bfloat16).float() + shared.float() + return acc.to(torch.bfloat16) + + +class M3McFinalizeAllreduceNorm(OpWrapper): + """The MoE boundary with the routed experts' finalize folded in. + + Per rank the partial is ``bf16(bf16(sum_k scales[t, k] * fc2[perm[t, k]]) + shared[t])`` (``perm`` -1 = dropped, + ``scales`` fp32 or bf16), then the group's all-reduce, residual add and RMSNorm exactly as + ``m3_mc_allreduce_norm``: same workspace, flags and call-order contract. Returns ``[normed, updated]`` or, with + ``quant``, also the MXFP8 copy of ``normed``. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The tp4 target takes this path from 9 tokens (top-4 of 128 experts, fp32 routing weights, the shared expert + # always present): the 16- and 48-token buckets, epilogue option 15. + CELLS: tuple[Cell, ...] = ( + Cell( + why="the smallest finalize call (9 tokens, 16-token bucket) with the next QKV's MXFP8 copy", + spec=dict(m=9, maxt=16, opts=15, topk=4, quant=True, dropped=0), + ), + Cell( + why="a full 16-token bucket with dropped expert slots", + spec=dict(m=16, maxt=16, opts=15, topk=4, quant=False, dropped=5), + ), + Cell( + why="the 48-token bucket at a partial fill", + spec=dict(m=24, maxt=48, opts=15, topk=4, quant=False, dropped=0), + ), + Cell( + why="the 48-token bucket full, MXFP8 copy", + spec=dict(m=48, maxt=48, opts=15, topk=4, quant=True, dropped=3), + ), + # Every other row count of the unfinalized MoE (dispatch.MC_FINALIZE): the bucket and options the target picks. + *( + Cell( + why=f"{m} tokens: a finalize row count between the ones above", + spec=dict( + m=m, + maxt=16 if m <= 16 else 48, + opts=15, + topk=4, + quant=m % 2 == 0, + dropped=m % 4, + ), + ) + for m in range(9, 49) + if m not in (9, 16, 24, 48) + ), + ) + + note = """ + BF16 routing weights and the 8-token bucket are accepted by the op (option 7 there) but the target only enters + this path at 9 tokens and up, with the built-in MoE's fp32 weights; neither is a cell. The rest is as in + `m3_mc_allreduce_norm`'s note. + """ + + def __call__( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_mc_allreduce_norm( + fc2, + shared, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + perm, + scales, + maxt, + quant, + 0, + opts, + ) + + def reference( + self, + per_rank_fc2: Sequence[torch.Tensor], + per_rank_perm: Sequence[torch.Tensor], + per_rank_scales: Sequence[torch.Tensor], + per_rank_shared: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute, from every rank's MoE outputs in rank order (the + collective exception, as in ``m3_mc_allreduce_norm.reference``).""" + partials = [ + finalize_partial( + fc2, + per_rank_perm[r], + per_rank_scales[r], + None if per_rank_shared is None else per_rank_shared[r], + ) + for r, fc2 in enumerate(per_rank_fc2) + ] + return boundary(partials, residual, norm_w, eps, quant) + + def is_valid( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + ) -> None: + h = residual.shape[-1] + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + # The kernel gathers fc2[perm] without a bound: an index past the FC2 rows reads whatever follows them. + lo, hi = int(perm.min()), int(perm.max()) + assert lo >= -1 and hi < fc2.shape[0], ( + f"perm spans [{lo}, {hi}]; entries must be -1 (dropped) or a row of fc2 (< {fc2.shape[0]})" + ) + + +m3_mc_finalize_allreduce_norm = M3McFinalizeAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding/__init__.py new file mode 100644 index 000000000000..e9f4d8e20b94 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Embedding ops of this target: the generic token-embedding gather.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding/m3_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding/m3_embedding.py new file mode 100644 index 000000000000..f979fdecf376 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding/m3_embedding.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The generic token-embedding gather over the replicated table (an id outside the table gives a zero row).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.embedding import pre_comm_embedding_ops + + +class M3Embedding(OpWrapper): + """``table[ids]`` with a zero row for an id outside ``[0, V)`` (the speculative sampler pads with -1). + + The table is the full (replicated) bf16 embedding [V, 6144]; ids int32 / int64 of any count. The fused decode + embedding (``m3_embedding_firstnorm``) takes decode steps of up to 8 rows; this entry takes every other step. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="a prefill chunk, int32 ids", spec=dict(rows=2048, id_dtype=torch.int32, pad=False) + ), + Cell( + why="a 48-row decode verify, int64 ids", + spec=dict(rows=48, id_dtype=torch.int64, pad=False), + ), + Cell( + why="draft rows padded with -1 ids", spec=dict(rows=32, id_dtype=torch.int32, pad=True) + ), + ) + + note = """ + An exact gather: `reference` must match bit for bit. + """ + + def __call__(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + return pre_comm_embedding_ops(ids, table, 1, 0, None, 0, table.shape[0], False, 0) + + def reference(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] + return rows * valid.unsqueeze(-1).to(table.dtype) + + +m3_embedding = M3Embedding() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding_firstnorm/__init__.py new file mode 100644 index 000000000000..4c6fc92cd49c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding_firstnorm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 decode embedding gather + first-layer RMSNorm as one CuTe DSL kernel (see op.py).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py new file mode 100644 index 000000000000..78563b4db40d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fixed M3 H6144 gather/raw residual/RMS, reusing pinned K3 norm kernel. + +128 threads x 8 BF16 x 6 iterations covers the row. Unlike the original +K3 wrapper this target does not promise FlashInfer's bit-identical reduction. +""" + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +from cutlass.experimental import primitives as prims + +NORM_THREADS = 128 +NORM_VEC = 8 +NORM_WARPS = 4 # capacity: fixed: the CTA's warps + + +@cute.kernel +def m3_embed_norm_kernel( + ids: cute.Tensor, # int32 or int64 [N] + table: cute.Tensor, # bf16 [V, H] + weight: cute.Tensor, # bf16 [H] + raw: cute.Tensor, # bf16 [N, H], out: the rows + out: cute.Tensor, # bf16 [N, H], out: the normed rows + vocab: cutlass.Int64, + eps: cutlass.Float32, + tv_layout: cute.Layout, + tiler_mn: cute.Shape, + hidden: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + prims.griddepcontrol(prims.GridDepAction.WAIT) + + smem = cutlass.utils.SmemAllocator() + reduction_buffer = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((1, NORM_WARPS)), byte_alignment=4 + ) + + tok = cutlass.Int64(ids[bidx]) + valid = (tok >= cutlass.Int64(0)) & (tok < vocab) + row = cutlass.Int64(cutlass.select_(valid, tok, cutlass.Int64(0))) + + gX = cute.local_tile(table, tiler_mn, (row, 0)) + gR = cute.local_tile(raw, tiler_mn, (bidx, 0)) + gY = cute.local_tile(out, tiler_mn, (bidx, 0)) + w_layout = cute.prepend(weight.layout, cute.make_layout((tiler_mn[0],), stride=(0,))) + gW = cute.local_tile(cute.make_tensor(weight.iterator, w_layout), tiler_mn, (0, 0)) + + copy_atom_load = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), table.element_type, num_bits_per_copy=128 + ) + copy_atom_store = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), out.element_type, num_bits_per_copy=128 + ) + thr_copy_X = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_W = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_O = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_mn).get_slice(tidx) + + tXgX = thr_copy_X.partition_S(gX) + tXrX = cute.make_fragment_like(tXgX) + tWgW = thr_copy_W.partition_S(gW) + tWrW = cute.make_fragment_like(tWgW) + tXrW = thr_copy_X.retile(tWrW) + tXgR = thr_copy_O.partition_D(gR) + tXgO = thr_copy_O.partition_D(gY) + tXrO = cute.make_fragment_like(tXgO) + + # The row (zero for an id outside [0, V)), straight to the snapshot slot; the weight. + tXrX.store(cute.zeros_like(tXrX, dtype=table.element_type)) + if valid: + cute.copy(copy_atom_load, tXgX, tXrX) + cute.copy(copy_atom_load, tWgW, tWrW) + cute.copy(copy_atom_store, tXrX, tXgR) + + # flashinfer's RMSNormKernel arithmetic, in its order (see the header). + x = tXrX.load().to(cutlass.Float32) + x_sq = x * x + sum_sq = x_sq.reduce(cute.ReductionOp.ADD, init_val=cutlass.Float32(0.0), reduction_profile=0) + for i in cutlass.range_constexpr(5): + sum_sq = sum_sq + cute.arch.shuffle_sync_bfly(sum_sq, offset=1 << i) + lane = cute.arch.lane_idx() + warp = cute.arch.warp_idx() + if lane == 0: + reduction_buffer[0, warp] = sum_sq + cute.arch.barrier() + total = cutlass.Float32(0.0) + if lane < NORM_WARPS: + total = reduction_buffer[0, lane] + for i in cutlass.range_constexpr(5): + total = total + cute.arch.shuffle_sync_bfly(total, offset=1 << i) + mean_sq = total / cutlass.Float32(hidden) + rstd = cute.math.rsqrt(mean_sq + eps, fastmath=True) + w = tXrW.load().to(cutlass.Float32) + y = x * rstd * (w + cutlass.Float32(0.0)) + tXrO.store(y.to(out.element_type)) + cute.copy(copy_atom_store, tXrO, tXgO) + + +@cute.jit +def m3_embed_norm( + ids: cute.Tensor, + table: cute.Tensor, + weight: cute.Tensor, + raw: cute.Tensor, + out: cute.Tensor, + vocab: cutlass.Int64, + eps: cutlass.Float32, + n_tokens: cutlass.Constexpr[int], + hidden: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + # flashinfer's thread-value layout for one row: thread t, value (v, k) -> column 8 t + v + 1024 k. + blocks = hidden // (NORM_VEC * NORM_THREADS) + tv_layout = cute.make_layout(((NORM_THREADS, 1), (NORM_VEC, blocks)), + stride=((NORM_VEC, 1), (1, NORM_VEC * NORM_THREADS))) # fmt: skip + m3_embed_norm_kernel( + ids, table, weight, raw, out, vocab, eps, tv_layout, (1, hidden), hidden + ).launch( + grid=[n_tokens, 1, 1], + block=[NORM_THREADS, 1, 1], + smem=NORM_WARPS * 4, + stream=stream, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding_firstnorm/op.py new file mode 100644 index 000000000000..276f8c32023b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding_firstnorm/op.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host side of the MiniMax-M3 decode embedding + first-layer RMSNorm (m3_embedding_firstnorm_kernel.py): H 6144, +1..8 rows, into caller-owned buffers. + +Compiled on the first call per (rows, id dtype, table shape, PDL), which must happen outside CUDA-graph capture +(:class:`DecodeEmbedding` compiles at construction). +""" + +import threading +from typing import Dict, Tuple + +import cuda.bindings.driver as cuda_driver +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +from .m3_embedding_firstnorm_kernel import m3_embed_norm as _m3_embed_norm_jit + +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size +MAX_ROWS = 8 # capacity: fixed: the kernel's built rows (decode steps) + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def arg(tensor): + return from_dlpack( + tensor.detach(), + assumed_align=tensor.element_size() + if tensor.ndim == 1 and tensor.dtype in (torch.int32, torch.int64) + else 16, + ) + + +def _check(ids, table, weight, raw, out) -> None: + rows = ids.shape[0] if ids.dim() == 1 else 0 + if not ( + 1 <= rows <= MAX_ROWS + and ids.dtype in (torch.int32, torch.int64) + and ids.is_contiguous() + and table.dtype == torch.bfloat16 + and table.dim() == 2 + and table.shape[1] == HIDDEN + and table.is_contiguous() + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (HIDDEN,) + and all( + t.dtype == torch.bfloat16 and tuple(t.shape) == (rows, HIDDEN) and t.is_contiguous() + for t in (raw, out) + ) + ): + raise ValueError( + "m3_embed_norm: expects ids int32/int64 [1..8], a contiguous bf16 table [V, 6144], bf16 weight [6144] " + "and contiguous bf16 raw / out [rows, 6144]" + ) + + +def _compile(ids, table, weight, eps, raw, out, use_pdl: bool): + key = (ids.shape[0], ids.dtype, tuple(table.shape), table.device, bool(use_pdl)) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_embed_norm compiles on its first call per shape: call it once before capture" + ) + with _lock: + fn = _compiled.get(key) + if fn is None: + stream = cuda_driver.CUstream(torch.cuda.current_stream(table.device).cuda_stream) + fn = _compiled[key] = cute.compile( + _m3_embed_norm_jit, + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + ids.shape[0], + HIDDEN, + bool(use_pdl), + stream, + ) + return fn + + +def m3_embed_norm( + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Write ``table[ids]`` into ``raw`` (a zero row for an id outside ``[0, V)``) and its RMSNorm with ``weight`` / + ``eps`` into ``out``; returns ``(out, raw)``.""" + _check(ids, table, weight, raw, out) + fn = _compile(ids, table, weight, eps, raw, out, use_pdl) + stream = cuda_driver.CUstream(torch.cuda.current_stream(ids.device).cuda_stream) + fn( + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + stream, + ) + return out, raw + + +class DecodeEmbedding: + """Caller-owned buffers of one (rows, id dtype) decode embedding: ``raw`` (the gathered rows, the first residual) + and ``output`` (their RMSNorm), over ``table`` and the first layer's plain (non-Gemma) input norm ``weight`` / + ``eps``. Built (and compiled) before CUDA-graph capture.""" + + def __init__(self, table, weight, eps, rows, id_dtype, use_pdl=True): + assert not torch.cuda.is_current_stream_capturing() + assert table.dtype == torch.bfloat16 and table.is_contiguous() and table.shape[1] == HIDDEN + assert weight.dtype == torch.bfloat16 and weight.shape == (HIDDEN,) + assert 1 <= rows <= MAX_ROWS and id_dtype in (torch.int32, torch.int64) + self.table, self.weight, self.rows, self.id_dtype, self.use_pdl = ( + table, + weight, + rows, + id_dtype, + use_pdl, + ) + self.eps = float(eps) + self.raw = torch.empty(rows, HIDDEN, device=table.device, dtype=torch.bfloat16) + self.output = torch.empty_like(self.raw) + ids = torch.empty(rows, device=table.device, dtype=id_dtype) + _compile(ids, table, weight, self.eps, self.raw, self.output, use_pdl) + + def __call__(self, ids): + assert ids.shape == (self.rows,) and ids.dtype == self.id_dtype and ids.is_contiguous() + return m3_embed_norm( + ids, self.table, self.weight, self.eps, self.raw, self.output, self.use_pdl + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/__init__.py new file mode 100644 index 000000000000..9cc374bdf1cc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GEMM entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_bf16_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_bf16_linear.py new file mode 100644 index 000000000000..d74fc709dbdf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_bf16_linear.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The BF16 Eagle3 drafter's projections: ``x @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3Bf16Linear(OpWrapper): + """``F.linear(x, w)`` for bf16 ``x`` [M, K] and a bf16 weight [N, K] (cuBLAS, fp32 accumulation), bf16 out. + + The drafter's fc over the captured target layers (K = 3 x 6144), its fused QKV (K = 2 x 6144, this rank's heads), + o_proj (this rank's input columns), gate_up and down projections run it. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="drafter fc, 32 verify rows", spec=dict(m=32, n=6144, k=18432)), + Cell(why="drafter QKV, eight draft rows", spec=dict(m=8, n=2304, k=12288)), + Cell(why="drafter o_proj, a prefill chunk", spec=dict(m=1024, n=6144, k=2048)), + Cell(why="drafter gate_up, one row", spec=dict(m=1, n=9216, k=6144)), + Cell(why="drafter down, eight rows", spec=dict(m=8, n=6144, k=4608)), + ) + + note = """ + cuBLAS's split and accumulation order: `reference` multiplies in fp32 (one bf16 rounding), within the test's band. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x, w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return (x.float() @ w.float().t()).to(x.dtype) + + +m3_bf16_linear = M3Bf16Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_fp32_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_fp32_linear.py new file mode 100644 index 000000000000..8d7374810675 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_fp32_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The router's fp32 projection beyond the decode routers' rows: ``float(x) @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + + +class M3Fp32Linear(OpWrapper): + """fp32 router logits ``F.linear(float(x[M, 6144]), w[128, 6144])`` for prefill and steps above 48 rows (the + fp32 GEMV and the tensor-core router take 1..48). Whatever fp32 matmul precision the process runs with applies + (the engine leaves torch's default).""" + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="a prefill chunk", spec=dict(m=2048, e=128, k=6144)), + Cell( + why="a generation step above the tensor-core router's 48 rows", + spec=dict(m=64, e=128, k=6144), + ), + ) + + note = """ + Under torch's default fp32 matmul precision the product may be TF32: the test's band covers that rounding against + the true fp32 reference. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x.to(torch.float32), w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + with true_fp32_matmul(): + return x.float() @ w.t() + + +m3_fp32_linear = M3Fp32Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv.py new file mode 100644 index 000000000000..a4cbb12d8085 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 x MXFP8 GEMV (tcgen05, block-scaled, CuTe DSL) on an activation its producer already quantized.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..moe.m3_moe_engine import mxfp8_dequant +from ..mx_gemv import op as _gemv + +__all__ = ["PROJECTIONS", "gemv_is_valid", "m3_mx_gemv", "served_cells"] + +# (N, K) of this target's MXFP8 projections at tp4, the GEMV's tuned shapes. +QKV_SPARSE = (2560, 6144) +QKV = (2304, 6144) +O_PROJ = (6144, 2048) +SHARED_GATE_UP = (1536, 6144) +SHARED_DOWN = (6144, 768) +DENSE_GATE_UP = (6144, 6144) +DENSE_DOWN = (6144, 3072) +PROJECTIONS = { + "sparse-layer QKV": QKV_SPARSE, + "dense-attention QKV": QKV, + "o_proj": O_PROJ, + "shared gate_up": SHARED_GATE_UP, + "shared down": SHARED_DOWN, + "dense gate_up": DENSE_GATE_UP, + "dense down": DENSE_DOWN, +} + + +def served_cells(named: tuple[Cell, ...], x_bf16: bool, max_m: int) -> tuple[Cell, ...]: + """A cell for every (projection, rows) the GEMV serves up to ``max_m`` rows that ``named`` does not list. + + Certification is per projection: a projection's split-K, ring and MMA width differ, so each (N, K, M) the + decode dispatch can send here is its own cell. + """ + listed = {(cell.spec["nk"], cell.spec["m"]) for cell in named} + return tuple( + Cell(why=f"{name}, {m} rows", spec=dict(nk=nk, m=m)) + for name, nk in PROJECTIONS.items() + for m in range(1, max_m + 1) + if (nk, m) not in listed and _gemv.serves(nk[0], nk[1], m, x_bf16) + ) + + +def gemv_reference( + x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor +) -> torch.Tensor: + """bf16(x @ dequant(weight)^T) for fp32 activation values x [M, K], summed in fp64 and rounded once.""" + return (x.double() @ mxfp8_dequant(weight, weight_scale).double().t()).to(torch.bfloat16) + + +def gemv_is_valid(weight: torch.Tensor, weight_scale: torch.Tensor) -> None: + n, k = weight.shape + # The kernel reads both from data_ptr() as dense buffers at the layout it was compiled for. + assert weight.dtype == torch.float8_e4m3fn and weight.is_contiguous(), ( + "weight must be contiguous e4m3 [N, K]" + ) + assert weight_scale.numel() == n * k // 32, ( + "weight_scale must be the 128x4-swizzled UE8M0 buffer [N * K / 32]" + ) + + +class M3MxGemv(OpWrapper): + """bf16 [M, N] = xq [M, K] (e4m3, 128x4-swizzled UE8M0 scales ``xsf``) x weight [N, K]^T (e4m3, swizzled + ``weight_scale``), M <= 48; None when the shape is not served (the caller falls back). + + ``trigger_early``: release the PDL dependents before the epilogue. ``cfg``: a (split, ring[, contig, nacc, a_tmem, + st_async]) tuning override; default the table's. ``use_pdl``: default the environment's. ``ar`` = (multicast + workspace address, flags address, rank, world, bucket tokens, mode): also push the rows into the next boundary + all-reduce's exchange (then ``op.last_pushed`` is True). + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The projections' (N, K) at the token counts each build differs in: the MMA width (8 / 16 / 32 / 48 columns) a + # call uses is the smallest that holds M, and the shared gate_up takes another split at 2 tokens and above 8. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token (split 4)", spec=dict(nk=QKV_SPARSE, m=1)), + Cell(why="sparse-layer QKV, 8 tokens", spec=dict(nk=QKV_SPARSE, m=8)), + Cell( + why="sparse-layer QKV, 48 verify rows: the widest MMA", spec=dict(nk=QKV_SPARSE, m=48) + ), + Cell(why="dense-attention QKV, one token", spec=dict(nk=QKV, m=1)), + Cell(why="dense-attention QKV, 13 rows: the 16-column MMA", spec=dict(nk=QKV, m=13)), + Cell(why="o_proj, one token (split 2)", spec=dict(nk=O_PROJ, m=1)), + Cell(why="o_proj, 24 rows: the 32-column MMA", spec=dict(nk=O_PROJ, m=24)), + Cell( + why="shared gate_up, one token (A from shared memory)", + spec=dict(nk=SHARED_GATE_UP, m=1), + ), + Cell(why="shared gate_up, 2 tokens: its split-4 config", spec=dict(nk=SHARED_GATE_UP, m=2)), + Cell( + why="shared gate_up, 12 rows: the wide split-4 config", + spec=dict(nk=SHARED_GATE_UP, m=12), + ), + Cell(why="shared down, one token (unsplit, K 768)", spec=dict(nk=SHARED_DOWN, m=1)), + Cell(why="shared down, 8 tokens", spec=dict(nk=SHARED_DOWN, m=8)), + Cell(why="dense gate_up, one token", spec=dict(nk=DENSE_GATE_UP, m=1)), + Cell(why="dense down, 8 tokens", spec=dict(nk=DENSE_DOWN, m=8)), + ) + # Every other (projection, rows) the decode dispatch can send here (dispatch.MX_GEMV_ROWS). + CELLS = CELLS + served_cells(CELLS, x_bf16=False, max_m=48) + + note = """ + The kernel sums in fp32 in split-K order where `reference` sums in fp64: the bf16 outputs agree within the band the + test derives. Not covered: the all-reduce push (`ar`), which needs a multicast exchange; `cfg` overrides beyond the + table (the target passes none). Shapes outside the table, M > 48, and 17..32 rows at K < 1024 return None. + """ + + def __call__( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv(xq, xsf, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The dequantized product; the launch options are ignored.""" + return gemv_reference(mxfp8_dequant(xq, xsf), weight, weight_scale) + + def is_valid( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + m, k = xq.shape + assert xq.dtype == torch.float8_e4m3fn and xq.is_contiguous(), ( + "xq must be contiguous e4m3 [M, K]" + ) + # A linear-layout [M, K / 32] scale buffer is read as the swizzled one: wrong scales, no error. + assert xsf.numel() >= (m + 127) // 128 * 128 * (k // 32), ( + "xsf must be the 128x4-swizzled scale buffer (rows padded to 128), not the linear one" + ) + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv = M3MxGemv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..1c6bbec4c10a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv_bf16.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 GEMV (CuTe DSL) on a bf16 activation, quantized to MXFP8 in the kernel's prologue.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..moe.m3_moe_engine import mxfp8_quant_dequant +from ..mx_gemv import op as _gemv +from .m3_mx_gemv import ( + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + SHARED_DOWN, + SHARED_GATE_UP, + gemv_is_valid, + gemv_reference, + served_cells, +) + +__all__ = ["m3_mx_gemv_bf16"] + + +class M3MxGemvBf16(OpWrapper): + """bf16 [M, N] = mxfp8_quantize(x) x weight [N, K]^T for a contiguous, 16-byte aligned bf16 ``x`` [M, K]; the + quantization is mxfp8_quantize's (UE8M0 scale per 32 = the smallest power of two >= amax / 448). None when the + shape is not served. The other arguments as ``m3_mx_gemv``'s. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The prologue quantizes one 32-element block per thread and round: the projections at the token counts the + # target sends here (no producer-quantized activation), up to the widest each serves. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token", spec=dict(nk=QKV_SPARSE, m=1)), + Cell( + why="sparse-layer QKV, 16 rows: the most its split-4 prologue serves", + spec=dict(nk=QKV_SPARSE, m=16), + ), + Cell(why="dense-attention QKV, 8 tokens", spec=dict(nk=QKV, m=8)), + Cell(why="o_proj, one token", spec=dict(nk=O_PROJ, m=1)), + Cell(why="o_proj, 32 rows: the 32-column MMA", spec=dict(nk=O_PROJ, m=32)), + Cell(why="shared gate_up, one token", spec=dict(nk=SHARED_GATE_UP, m=1)), + Cell(why="shared gate_up, 2 tokens: its split-4 config", spec=dict(nk=SHARED_GATE_UP, m=2)), + Cell(why="shared down, 8 tokens (K 768)", spec=dict(nk=SHARED_DOWN, m=8)), + Cell( + why="dense gate_up, 8 tokens: the most split 2 at K 6144 serves", + spec=dict(nk=DENSE_GATE_UP, m=8), + ), + Cell(why="dense down, one token", spec=dict(nk=DENSE_DOWN, m=1)), + ) + # Every other (projection, rows) the decode dispatch can send here (dispatch.MX_GEMV_BF16_MAX_TOKENS). + CELLS = CELLS + served_cells(CELLS, x_bf16=True, max_m=32) + + note = """ + As `m3_mx_gemv`. The in-kernel quantization is mxfp8_quantize's bit for bit except where the kernel's + rcp.approx(448) and an exact 1/448 put a block's amax on different sides of a power of two (not observed on the + cells). Not covered: `ar`, `cfg` overrides. Returns None for rows beyond 8 rounds of the prologue per CTA. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv_bf16(x, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The MXFP8 round trip of x times the dequantized weight; the launch options are ignored.""" + return gemv_reference(mxfp8_quant_dequant(x), weight, weight_scale) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv_bf16 = M3MxGemvBf16() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_gemm.py new file mode 100644 index 000000000000..b3d54e14ef2e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_gemm.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-graph MXFP8 GEMM on an activation its producer already quantized (rows the GEMV does not serve).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_gemm_autotuned + +from .m3_mxfp8_linear import mxfp8_reference + + +class M3Mxfp8Gemm(OpWrapper): + """``act @ weight^T`` for an MXFP8 activation (e4m3 [M, K] + 128x4-swizzled UE8M0 scales) and an MXFP8 weight + [N, K] (swizzled scales), bf16 out. The GEMM backend per token bucket is the AutoTuner's (the same runner, op key + and tactic cache as the engine-tuned quantize + GEMM); ``tune`` profiles in the engine's tuning pass.""" + + ARCHS = frozenset({Arch.SM_103}) + + # The decode dispatch hands it the boundary / merge MXFP8 activation on rows the GEMV does not serve: the + # unsplit shared down projection at 17..32 rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert down projection, 17 rows", spec=dict(m=17, n=6144, k=768)), + Cell(why="shared expert down projection, 32 rows", spec=dict(m=32, n=6144, k=768)), + ) + + note = """ + Exactly the MXFP8 product (fp32 accumulation, one bf16 rounding): `reference` dequantizes both operands and + multiplies in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + return mxfp8_gemm_autotuned(act, act_scale, weight, weight_scale, output_dtype, tune=tune) + + def reference( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + """The dequantized product; ``tune`` is ignored.""" + return mxfp8_reference(act, act_scale, weight, weight_scale).to(output_dtype) + + +m3_mxfp8_gemm = M3Mxfp8Gemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_linear.py new file mode 100644 index 000000000000..799cda007837 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_linear.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One MXFP8 projection on the engine-tuned path: dynamic MXFP8 quantize of a bf16 activation + block-scaled GEMM.""" + +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, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_quantize_gemm_autotuned + +_BLOCK = 32 + + +def _unswizzle(scale: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """[rows, cols] UE8M0 bytes from the 128x4-swizzled layout (rows padded to 128, columns to 4).""" + pad_r = (rows + 127) // 128 * 128 + pad_c = (cols + 3) // 4 * 4 + s = scale.view(torch.uint8).reshape(pad_r // 128, pad_c // 4, 32, 4, 4) + return s.permute(0, 3, 2, 1, 4).reshape(pad_r, pad_c)[:rows, :cols] + + +def _dequant(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + rows, k = x.shape + s = _unswizzle(scale, rows, k // _BLOCK).to(torch.int32) + factor = torch.pow(2.0, (s - 127).to(torch.float32)) + return ( + x.view(torch.float8_e4m3fn).float().view(rows, k // _BLOCK, _BLOCK) * factor.unsqueeze(-1) + ).view(rows, k) + + +def mxfp8_reference(act, act_scale, weight, weight_scale) -> torch.Tensor: + """fp32 product of two dequantized MXFP8 operands (swizzled UE8M0 scales).""" + with true_fp32_matmul(): + return _dequant(act, act_scale) @ _dequant(weight, weight_scale).t() + + +class M3Mxfp8Linear(OpWrapper): + """``x @ weight^T`` for bf16 ``x`` [M, K] and an MXFP8 weight [N, K] (e4m3 + 128x4-swizzled UE8M0 scales), bf16 + out, on the backend the engine selected: + + * ``tuned_quantize_gemm`` (decode CUDA graphs once the engine tuned per-bucket backends): the AutoTuner's quantize + and GEMM per token bucket (``tune`` profiles them in the engine's tuning pass); + * otherwise the native ``mxfp8_quantize`` (swizzled scales) and then FlashInfer's MXFP8 GEMM (``flashinfer``) or + the CUTLASS block-scaled GEMM, through the autotuner while the engine tunes it (``native_tuning``), scaled by + ``alpha`` (fp32 [1], 1.0 for these checkpoints). + + Every MXFP8 projection of the target (the attention QKV and o_proj, the dense MLPs and the shared expert) runs it + outside the decode-step kernels. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="sparse-layer QKV, a prefill chunk, native GEMM", + spec=dict(m=2048, n=2560, k=6144, path="native"), + ), + Cell( + why="o_proj, a prefill chunk, native GEMM", + spec=dict(m=2048, n=6144, k=2048, path="native"), + ), + Cell( + why="dense down projection, eager decode, native GEMM", + spec=dict(m=4, n=6144, k=3072, path="native"), + ), + Cell( + why="shared gate_up, FlashInfer GEMM", + spec=dict(m=64, n=1536, k=6144, path="flashinfer"), + ), + Cell(why="o_proj, tuned quantize + GEMM", spec=dict(m=64, n=6144, k=2048, path="tuned")), + ) + + note = """ + Exactly the MXFP8 product of the MXFP8 round trip of x (fp32 accumulation, one bf16 rounding): `reference` + quantizes x with the same block rule (UE8M0 scale = the block's power-of-two ceiling over e4m3's 448) and + multiplies dequantized operands in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + if tuned_quantize_gemm: + return mxfp8_quantize_gemm_autotuned(x, weight, weight_scale, output_dtype, tune=tune) + act, act_scale = torch.ops.trtllm.mxfp8_quantize(x, True) + if flashinfer: + return torch.ops.trtllm.flashinfer_mm_mxfp8( + act, act_scale, weight, weight_scale, output_dtype + ) + gemm = ( + torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned + if native_tuning + else torch.ops.trtllm.mxfp8_mxfp8_gemm + ) + return gemm(act, act_scale, weight, weight_scale, alpha, output_dtype) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + """The MXFP8 round trip of ``x`` times the dequantized weight, times ``alpha``; the backend flags are + ignored.""" + m, k = x.shape + blocks = x.float().view(m, k // _BLOCK, _BLOCK) + amax = blocks.abs().amax(-1, keepdim=True) + exp = torch.ceil(torch.log2(torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)))) + scale = torch.pow(2.0, exp) + q = (blocks / scale).to(torch.float8_e4m3fn).float() * scale + with true_fp32_matmul(): + out = q.view(m, k) @ _dequant(weight, weight_scale).t() + return (out * alpha.float()).to(output_dtype) + + +m3_mxfp8_linear = M3Mxfp8Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_target_head.py new file mode 100644 index 000000000000..5f3b4ad1e2b0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_target_head.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target lm_head shard GEMV of a verify step of at most 8 rows over a lossless 12-bit coding of its BF16 weight, +and that coding (``encode_chunked`` / ``decode_reference`` / ``round_trip_ok``). + +The coding: the BF16 exponents of a trained weight sit almost all in a window of 15 consecutive values +``[base, base + 14]``. Each weight is one byte ``sign << 7 | mantissa`` plus a 4-bit code ``exponent - base``; code 15 +marks an exception, which decodes to +0 (it adds nothing to the MMA sums); its exact product ``x * w`` is added from a +per-tile list. The bytes are laid out in ``mma.m16n8k16`` A-fragment order: +per 16-row tile and unit of 4 k16 steps, each lane reads the sign/mantissa bytes of steps 0-1, of steps 2-3, then 4 +code words (``[tiles][K / 64][3][32 lanes][16 bytes]``). A coded weight is used only after ``round_trip_ok`` proved +it decodes to the BF16 weight bit for bit. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +__all__ = ["m3_target_head"] + +#: Verify rows per call (requests x (draft length + 1)). +MAX_ROWS = 8 +#: Warps per CTA (the split-tile scratch holds 2 x 4 x 32 fp32 per warp). +CTA_WARPS = 16 +#: The target lm_head's TP4 shard (200064-id vocabulary) and the hidden size. +TARGET_SHARD = 50016 +HIDDEN = 6144 +#: Codes per window: 0..14 decode, 15 marks an exception. +WINDOW = 15 +#: The largest exponent base whose window top plus one (base + 15) is still a finite BF16 exponent. +MAX_BASE = 254 - WINDOW +#: Rows per block of encode_chunked / round_trip_ok (bounded temporaries). +ROWS_PER_CHUNK = 4096 + + +def _pad_rows(weight: torch.Tensor) -> torch.Tensor: + n = weight.shape[0] + tiles = (n + 15) // 16 + if tiles * 16 == n: + return weight + return torch.cat([weight, weight.new_zeros(tiles * 16 - n, weight.shape[1])]) + + +def _fragment_view(w: torch.Tensor) -> torch.Tensor: + """[tiles * 16, K] -> [tiles, K / 16, 32 lanes, 8 values] in mma A-fragment order (a permutation of the values).""" + n, k = w.shape + # rows: (tile, rh, g); cols: (ks, ch, j, e) -> (tile, ks, g, j, ch, rh, e) + return ( + w.view(n // 16, 2, 8, k // 16, 2, 4, 2) + .permute(0, 3, 2, 5, 4, 1, 6) + .reshape(n // 16, k // 16, 32, 8) + ) + + +def _exponents(weight: torch.Tensor) -> torch.Tensor: + return (weight.contiguous().view(torch.int16).to(torch.int32) >> 7) & 0xFF + + +def _window_base(hist: torch.Tensor) -> int: + """The base of the 15-exponent window holding the most weights (the lowest such base).""" + window = torch.stack([hist[b : b + WINDOW].sum() for b in range(0, MAX_BASE + 1)]) + return int(window.argmax()) + + +def encode(weight: torch.Tensor, base: Optional[int] = None): + """Lossless 12-bit coding of a BF16 [N, K] weight (see the module docstring); ``base`` defaults to this weight's + own window. + + Returns (packed uint8 [tiles * K / 64 * 1536], base, exc_off int32 [tiles + 1], exc_pos int32 [E] (row in tile + << 16 | k), exc_w fp32 [E, 2] ({w, 0})). + """ + w = _pad_rows(weight) + n, k = w.shape + tiles, ks = n // 16, k // 16 + assert ks % 4 == 0, "K must be a multiple of 64" + bits = w.contiguous().view(torch.int16).to(torch.int32) & 0xFFFF + expo = (bits >> 7) & 0xFF + if base is None: + base = _window_base(torch.bincount(expo.flatten(), minlength=256)) + code = expo - base + exc = (code < 0) | (code >= WINDOW) + code = torch.where(exc, torch.full_like(code, WINDOW), code) + sm = ((bits >> 8) & 0x80) | (bits & 0x7F) + sm = torch.where(exc, torch.zeros_like(sm), sm) + # Exceptions: (tile, row in tile, k) in tile-major order; code 15 decodes to +0. + rows, cols = exc.nonzero(as_tuple=True) + order = torch.argsort((rows // 16).long() * (16 * k) + (rows % 16).long() * k + cols.long()) + rows, cols = rows[order], cols[order] + exc_w = torch.stack( + [ + w[rows, cols].float(), + torch.zeros((rows.numel(),), device=w.device), + ], + dim=1, + ) + exc_pos = (((rows % 16) << 16) | cols).to(torch.int32) + exc_off = torch.zeros(tiles + 1, dtype=torch.int32, device=w.device) + exc_off[1:] = torch.cumsum(torch.bincount((rows // 16).long(), minlength=tiles), 0).to( + torch.int32 + ) + # Fragment order, then per (tile, group of 4 k16 steps, lane): 32 sign/mantissa bytes and 4 code words. + smf = _fragment_view(sm.to(torch.uint8).view(n, k)).view(tiles, ks // 4, 4, 32, 8) + cf = ( + _fragment_view(code.to(torch.uint8).view(n, k)) + .view(tiles, ks // 4, 4, 32, 8) + .to(torch.int32) + ) + # Code word of a k16 step: value 2p -> bits 4p..4p+3, value 2p+1 -> bits 16+4p.. + cw = torch.zeros(tiles, ks // 4, 4, 32, dtype=torch.int64, device=w.device) + for p in range(4): + cw |= (cf[..., 2 * p].long() << (4 * p)) | (cf[..., 2 * p + 1].long() << (16 + 4 * p)) + cw = cw.to(torch.int32) # wraps for bit 31; the kernel reads it as uint32 + # Word 0 of a lane = the bytes of steps 0 and 1; word 1 = steps 2 and 3; word 2 = the code words of steps 0..3. + seg01 = ( + smf.permute(0, 1, 3, 2, 4).reshape(tiles, ks // 4, 32, 2, 16).permute(0, 1, 3, 2, 4) + ) # [.., 2, 32, 16] + seg2 = cw.permute(0, 1, 3, 2).contiguous().view(torch.uint8).view(tiles, ks // 4, 1, 32, 16) + packed = torch.cat([seg01.contiguous(), seg2], dim=2) # [tiles, ks / 4, 3, 32, 16] + return packed.contiguous().view(-1), base, exc_off, exc_pos, exc_w.contiguous() + + +def encode_chunked(weight: torch.Tensor, rows_per_chunk: int = ROWS_PER_CHUNK): + """``encode`` over row blocks (bounded temporaries), with the exponent window chosen on the whole weight.""" + n, _ = weight.shape + assert rows_per_chunk % 16 == 0 + hist = torch.zeros(256, dtype=torch.int64, device=weight.device) + for r in range(0, n, rows_per_chunk): + hist += torch.bincount(_exponents(weight[r : r + rows_per_chunk]).flatten(), minlength=256) + base = _window_base(hist) + packed, offs, poss, ws = [], [], [], [] + n_exc = 0 + for r in range(0, n, rows_per_chunk): + p, _, off, pos, ew = encode(weight[r : r + rows_per_chunk], base=base) + packed.append(p) + offs.append(off[:-1] + n_exc) + poss.append(pos) + ws.append(ew) + n_exc += pos.numel() + exc_off = torch.cat(offs + [torch.tensor([n_exc], dtype=torch.int32, device=weight.device)]) + return torch.cat(packed), base, exc_off, torch.cat(poss), torch.cat(ws) + + +def decode_reference(packed: torch.Tensor, base: int, n: int, k: int) -> torch.Tensor: + """BF16 [N, K] from the coded stream with the in-kernel rule (exceptions, code 15, decode to +0).""" + tiles, ks = (n + 15) // 16, k // 16 + p = packed.view(tiles, ks // 4, 3, 32, 16) + sm = ( + p[:, :, :2] + .permute(0, 1, 3, 2, 4) + .reshape(tiles, ks // 4, 32, 4, 8) + .permute(0, 1, 3, 2, 4) + .to(torch.int32) + ) + cw = ( + p[:, :, 2] + .contiguous() + .view(torch.int32) + .view(tiles, ks // 4, 32, 4) + .permute(0, 1, 3, 2) + .long() + & 0xFFFFFFFF + ) + codes = torch.stack([(cw >> (4 * (i // 2) + 16 * (i % 2))) & 0xF for i in range(8)], dim=-1).to( + torch.int32 + ) + bits = ((sm & 0x80) << 8) | ((codes + base) << 7) | (sm & 0x7F) + bits = torch.where(codes == WINDOW, torch.zeros_like(bits), bits) + frag = bits.reshape(tiles, ks, 32, 8) + # Inverse of _fragment_view: (tile, ks, g, j, ch, rh, e) -> (tile, rh, g, ks, ch, j, e). + w = frag.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, k) + return w.to(torch.int16).view(torch.bfloat16)[:n] + + +def decode( + packed: torch.Tensor, + base: int, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + n: int, + k: int, + row0: int = 0, + row1: Optional[int] = None, +) -> torch.Tensor: + """BF16 rows [row0, row1) of the coded weight [N, K] with the exceptions applied (what the kernel multiplies by). + ``row0`` must be a multiple of 16.""" + row1 = n if row1 is None else row1 + assert row0 % 16 == 0 + t0, t1 = row0 // 16, (row1 + 15) // 16 + unit_bytes = (k // 64) * 3 * 32 * 16 + dec = decode_reference(packed[t0 * unit_bytes : t1 * unit_bytes], base, row1 - row0, k).clone() + counts = (exc_off[1:] - exc_off[:-1]).long() + tile_of = torch.repeat_interleave(torch.arange(counts.numel(), device=packed.device), counts) + rows = tile_of * 16 + (exc_pos >> 16).long() + cols = (exc_pos & 0xFFFF).long() + sel = (rows >= row0) & (rows < row1) + dec[rows[sel] - row0, cols[sel]] = exc_w.view(-1, 2)[sel, 0].to(torch.bfloat16) + return dec + + +def round_trip_ok( + packed: torch.Tensor, + base: int, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + weight: torch.Tensor, + rows_per_chunk: int = ROWS_PER_CHUNK, +) -> bool: + """Whether the coded weight (code values + exception list) decodes to ``weight`` bit for bit, with its exponent + window below the BF16 infinity exponent.""" + n, k = weight.shape + if not 0 <= base <= MAX_BASE or k % 64: + return False + if ( + packed.numel() != ((n + 15) // 16) * 16 * k * 3 // 2 + or exc_off.numel() != (n + 15) // 16 + 1 + ): + return False + for r0 in range(0, n, rows_per_chunk): + r1 = min(n, r0 + rows_per_chunk) + dec = decode(packed, base, exc_off, exc_pos, exc_w, n, k, r0, r1) + if not torch.equal(dec.view(torch.int16), weight[r0:r1].contiguous().view(torch.int16)): + return False + return True + + +class M3TargetHead(OpWrapper): + """``out[:M, :n]`` (BF16) = x [M, K] (BF16, M <= 8) times the coded BF16 shard W [n, K] transposed: each logit + the fp32 sum of the exact BF16 products, rounded once to BF16 (the in-tree BF16 head GEMM's arithmetic, in + another summation order). + + ``w`` / ``exp_base`` / ``exc_off`` / ``exc_pos`` / ``exc_w`` are ``encode_chunked``'s outputs for W; ``part`` + (fp32 [grid * 16 * 256]) and ``counter`` (int32 [ceil(n / 16)], zero between calls, re-armed by the kernel) are + the split-tile scratch; ``grid`` is the number of 16-warp CTAs (one per SM). + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The verify of a fused-verify step at most 8 rows wide: 1 + 3 drafts for one or two requests (4 / 8 rows), and + # every other row count the kernel's one n8 MMA tile serves; one cell with an exception in many tiles, one with + # a shard whose rows are not a multiple of 16. + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"{m} verify row(s) of the TP4 target head shard", + spec=dict(m=m, n=TARGET_SHARD, k=HIDDEN), + ) + for m in range(1, MAX_ROWS + 1) + ), + Cell( + why="4 rows with outliers and exact zeros in many tiles (the exception lists)", + spec=dict(m=4, n=TARGET_SHARD, k=HIDDEN, exceptions=True), + ), + Cell( + why="2 rows over a shard of 4000 rows (250 tiles), K 1024", + spec=dict(m=2, n=4000, k=1024, exceptions=True), + ), + ) + + note = """ + The weight the kernel multiplies by is the checkpoint's BF16 weight bit for bit (round_trip_ok); an exception + adds nothing to the MMA sums and its exact product x * w is added at its tile's flush. A logit therefore differs + from the in-tree head's only by the fp32 summation order before its single BF16 rounding (ulp-level); + integer-valued data, whose sums are exact in any order, is equal bit for bit. Split tiles are summed in warp + order: the result is deterministic. + """ + + def __call__( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_target_head( + w, x, out, part, counter, exc_off, exc_pos, exc_w, exp_base, n, grid, pdl + ) + return out + + def reference( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> torch.Tensor: + """The BF16 logits [M, n] of x against the decoded weight: the fp32 product, rounded once to BF16. + + ``out``, ``part``, ``counter``, ``grid`` and ``pdl`` do not change the result.""" + weight = decode(w, exp_base, exc_off, exc_pos, exc_w, n, x.shape[1]) + with true_fp32_matmul(): + logits = x.float() @ weight.float().t() + return logits.to(torch.bfloat16) + + def is_valid( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> None: + # A split tile's last finisher is the CTA whose arrival makes the count reach its owners: a stale count + # stores a partial sum. + assert not bool(counter[: (n + 15) // 16].any()), "counter must be zero between calls" + # The exceptions are read per tile in list order with no bound against the coded stream. + assert int(exc_off[-1].item()) == exc_pos.numel(), "exc_off must end at the exception count" + assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" + assert out.data_ptr() != x.data_ptr(), "out must not alias x" + + +m3_target_head = M3TargetHead() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/mxfp8_low_m_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/mxfp8_low_m_gemm.py new file mode 100644 index 000000000000..915c401a3ca1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/mxfp8_low_m_gemm.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused MXFP8 activation quantization + low-M MXFP8 GEMM of the decode projections.""" + +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, +) + +_E4M3_MAX = 448.0 + + +def _ue8m0_unswizzle(wsf: torch.Tensor, n: int, k: int) -> torch.Tensor: + """[N, K / 32] scale bytes out of the 128x4-swizzled buffer of ``mxfp8_quantize(w, True)``.""" + cols = k // 32 + padded_cols = (cols + 3) // 4 * 4 + r = torch.arange(n, device=wsf.device).view(-1, 1) + c = torch.arange(cols, device=wsf.device).view(1, -1) + index = ( + (c % 4) + + (c // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + return wsf.reshape(-1)[index] + + +def _mxfp8_dequant_rows(x: torch.Tensor) -> torch.Tensor: + """The activation as the kernel quantizes it (``trtllm::mxfp8_quantize``), dequantized back to fp32. + + Per 32 columns: scale = the smallest power of two at or above ``amax / 448`` (UE8M0, round up), data = x / scale + rounded to nearest E4M3; an all-zero block stays zero. + """ + m, k = x.shape + blocks = x.float().view(m, k // 32, 32) + amax = blocks.abs().amax(-1, keepdim=True) + bits = (amax / _E4M3_MAX).view(torch.int32) + byte = (((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) > 0).to(torch.int32)).clamp(0, 254) + scale = torch.exp2(byte.float() - 127.0) + data = torch.where(amax != 0, blocks / scale, torch.zeros_like(blocks)).to(torch.float8_e4m3fn) + return (data.float() * scale).view(m, k) + + +# The target's configuration per (N, K) projection marked for this kernel, and the most rows it routes here +# (dispatch.LOW_M_GEMM_MAX_TOKENS): (warps, cluster_k, row_tiles, max rows). +CONFIGS = { + (2560, 6144): (4, 4, 2, 8), + (6144, 2048): (8, 2, 4, 8), + (6144, 6144): (4, 4, 1, 1), + (6144, 3072): (4, 2, 2, 8), +} + + +def _config_cells(named: tuple[Cell, ...]) -> tuple[Cell, ...]: + """A cell for every (projection, rows) of ``CONFIGS`` that ``named`` does not list: certification is per + projection, since each runs its own configuration.""" + listed = {(cell.spec["n"], cell.spec["k"], cell.spec["m"]) for cell in named} + return tuple( + Cell( + why=f"{n} x {k}, {m} rows", + spec=dict(m=m, n=n, k=k, warps=warps, cluster_k=cluster_k, row_tiles=row_tiles), + ) + for (n, k), (warps, cluster_k, row_tiles, max_m) in CONFIGS.items() + for m in range(1, max_m + 1) + if (n, k, m) not in listed + ) + + +class Mxfp8LowMGemm(OpWrapper): + """``out[M, N] = mxfp8(x)[M, K] @ dequant(w, wsf)[N, K]^T`` in BF16, quantizing ``x`` in the kernel. + + ``w`` is E4M3 [N, K] with UE8M0 block-32 scales ``wsf`` in the 128x4 swizzled layout; ``x`` BF16 [M, K] with unit + column stride. ``(warps, cluster_k, row_tiles)`` is the kernel configuration the target's table picks per + weight shape; each CTA owns ``16 * row_tiles`` weight rows and ``K / cluster_k`` columns, split-K partials meet + in the cluster root and are summed in rank order, one BF16 rounding at the end. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The tp4 projections marked for this kernel (the decode GEMV runs them first and hands over when it declines), + # each at the configuration the target's table assigns and up to the row count it routes here. + CELLS: tuple[Cell, ...] = ( + Cell( + why="QKV with the index q / k projections (2560 x 6144), one token", + spec=dict(m=1, n=2560, k=6144, warps=4, cluster_k=4, row_tiles=2), + ), + Cell( + why="QKV at the configured row limit, 8 tokens", + spec=dict(m=8, n=2560, k=6144, warps=4, cluster_k=4, row_tiles=2), + ), + Cell( + why="o_proj (6144 x 2048), 8 warps x 4 row tiles, odd row count", + spec=dict(m=5, n=6144, k=2048, warps=8, cluster_k=2, row_tiles=4), + ), + Cell( + why="dense MLP gate_up (6144 x 6144), configured for one token only", + spec=dict(m=1, n=6144, k=6144, warps=4, cluster_k=4, row_tiles=1), + ), + Cell( + why="dense MLP down (6144 x 3072), 8 tokens", + spec=dict(m=8, n=6144, k=3072, warps=4, cluster_k=2, row_tiles=2), + ), + ) + # Every other (projection, rows) of its configurations. + CELLS = CELLS + _config_cells(CELLS) + + note = """ + Only the four configurations the target's table assigns are built, each for 1..8 rows (one 8-row MMA tile). + `pdl` changes launch overlap only and is driven off. Activation blocks whose amax is below `448 * 2^-127` take a + flush-to-zero path the reference does not model; no decode activation is that small. + """ + + def __call__( + self, + x: torch.Tensor, + w: torch.Tensor, + wsf: torch.Tensor, + warps: int, + cluster_k: int, + row_tiles: int, + pdl: bool, + ) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_mxfp8_low_m_gemm( + x, w, wsf, warps, cluster_k, row_tiles, pdl + ) + + def reference( + self, + x: torch.Tensor, + w: torch.Tensor, + wsf: torch.Tensor, + warps: int, + cluster_k: int, + row_tiles: int, + pdl: bool, + ) -> torch.Tensor: + """Dequantized fp32 product of the MXFP8 activation and weight, rounded to BF16 once. + + ``warps``, ``cluster_k``, ``row_tiles`` and ``pdl`` only pick the kernel's split; they do not enter the + arithmetic. + """ + n, k = w.shape + w_scale = torch.exp2(_ue8m0_unswizzle(wsf, n, k).float() - 127.0) + w_deq = (w.float().view(n, k // 32, 32) * w_scale.unsqueeze(-1)).view(n, k) + with true_fp32_matmul(): + out = _mxfp8_dequant_rows(x) @ w_deq.t() + return out.to(torch.bfloat16) + + def is_valid( + self, + x: torch.Tensor, + w: torch.Tensor, + wsf: torch.Tensor, + warps: int, + cluster_k: int, + row_tiles: int, + pdl: bool, + ) -> None: + n, k = w.shape + # The kernel bulk-copies 256-byte scale chunks by computed offset; the op checks the dtype only, so a short + # or linear-layout scale buffer is read past its end or in the wrong order. + need = (n + 127) // 128 * 128 * ((k // 32 + 3) // 4 * 4) + assert wsf.is_contiguous() and wsf.numel() >= need, ( + f"wsf must be the contiguous 128x4-swizzled UE8M0 buffer of {need} bytes; got {wsf.numel()}" + ) + + +mxfp8_low_m_gemm = Mxfp8LowMGemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/long_select.py new file mode 100644 index 000000000000..c01a82f0da41 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/long_select.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 block select for MiniMax-M3 decode rows of 129 to 9216 blocks. + +Same contract and bit-identical output as ``torch.ops.trtllm.minimax_m3_select_blocks`` on per-KV-head scores +[num_kv_heads, n_blocks, total_q]: init / local forcing, per-query valid-block count, ties to the smaller block index, +forced -inf picks emitted as -1, ascending block ids with -1 tail padding. + +One 1024-thread CTA per (query, KV head) row holds the row's ordered score keys in registers (2, 4 or 9 per +thread, by the row's valid count). The weakest of the 16 maxima of the interleaved groups ``block % 16`` is a lower +bound on the 16th-largest key, so the top 16 is among the keys at or above it; those are compacted to shared memory +and ranked exactly by (key, block). More than 1024 keys tied at the bound (degenerate rows) rank every candidate +against the whole row (re-read). + +The kernel is ``torch.ops.trtllm.minimax_m3_select_long``, called through this target's ``m3_select_long`` entry. +""" + +from typing import Optional + +import torch + +from ..ops.attention.m3_select_long import MAX_BLOCKS, m3_select_long + + +def long_select_blocks( + max_score_kv: torch.Tensor, + *, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, +) -> Optional[torch.Tensor]: + """Top-16 select over per-KV-head scores [num_kv_heads, n_blocks, total_q] with 129..9216 blocks. + + Returns the head-major backed [total_q, num_kv_heads, 16] int32 table, or None outside this kernel's shapes + (callers keep the reference op). ``early_trigger`` releases the dependent sparse decode right after the select's + grid wait (``SELECT_EARLY_MAX_ROWS`` in modeling.py). + """ + if ( + max_score_kv.dtype != torch.float32 + or not max_score_kv.is_cuda + or max_score_kv.dim() != 3 + or not 128 < max_score_kv.shape[1] <= MAX_BLOCKS + or n_valid_blocks.dtype != torch.int32 + or not n_valid_blocks.is_cuda + or n_valid_blocks.shape[0] < max_score_kv.shape[2] + ): + return None + num_kv_heads, _, total_q = max_score_kv.shape + out = torch.empty( + (num_kv_heads, total_q, 16), dtype=torch.int32, device=max_score_kv.device + ).permute(1, 0, 2) + return m3_select_long( + max_score_kv, + n_valid_blocks.contiguous(), + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_attn_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_attn_lc.py new file mode 100644 index 000000000000..88e64f35ceb4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_attn_lc.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 long-context sparse decode attention in one launch. + +Replaces the generic chain's split-K producer (``_gqa_sparse_decode_kernel``, one block per chunk) and the LSE +merge (``_merge_topk_attn_out_kernel``, with its MXFP8 copy for the o_proj) once the top-k block table exists. + +Grid: one cluster of 16 CTAs per query token; CTA ``r`` attends the token's ``r``-th selected block (the Triton +producer's chunk ``r`` at 16 chunks), keeps the partial in shared memory, and after a cluster barrier merges head +``r`` from all 16 partials over DSMEM, writing the BF16 row and, optionally, its MXFP8 copy and swizzled UE8M0 +scales exactly as ``_store_mxfp8_head`` lays them out. + +Per-block math follows the Triton producer: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16, QK and +PV on FP16 tensor cores with FP32 accumulation, log2-domain softmax with either the linear denominator or the +log-domain recurrence. Accumulation order inside the MMAs and the reductions differs, so results match to rounding, +not bit for bit. + +Schedule: the request length and the query rows (FP8 widened to FP16 in shared memory) are read before the PDL grid +wait, which overlaps them with the select ahead when it releases this grid early (modeling's +``SELECT_EARLY_MAX_ROWS``); +the split-K partials travel by ``st.async`` into each receiving CTA's own mbarrier instead of a cluster-wide barrier. +Neither changes a value or the merge order. + +Geometry: 16 query heads per KV head (TP4), one KV head, head_dim 128, page 128, top-k 16. The kernel is +``torch.ops.trtllm.minimax_m3_sparse_decode_lc``, called through this target's ``m3_sparse_decode_lc`` entry. +""" + +import torch + +from ..ops.attention.m3_sparse_decode_lc import m3_sparse_decode_lc + + +def lc_applies(q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output) -> bool: + """Whether the step's operands fit the kernel: TP4 geometry, packed NVFP4 pages, FP8 q, BF16 [T, 16, 128] out.""" + if q.dim() != 3 or tuple(q.shape[1:]) != (16, 128) or q.dtype != torch.float8_e4m3fn: + return False + if q.stride(2) != 1 or q.stride(1) % 16 or q.stride(0) % 16 or q.data_ptr() % 16: + return False + if int(k_paged.shape[1]) != 1 or int(topk_idx.shape[0]) != 1 or int(topk_idx.shape[-1]) != 16: + return False + for t in (k_paged, v_paged): + if ( + t.stride(3) != 1 + or t.stride(2) != 64 + or int(t.shape[3]) != 64 + or t.data_ptr() % 16 + or t.stride(0) % 16 + ): + return False + for t in (k_block_scale, v_block_scale): + if t is None or t.stride(3) != 1 or t.stride(2) != 8 or t.data_ptr() % 8 or t.stride(0) % 8: + return False + return ( + output.dtype == torch.bfloat16 + and output.dim() == 3 + and output.is_contiguous() + and tuple(output.shape) == tuple(q.shape) + ) + + +def m3_lc_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + linear_softmax=True, + mxfp8=None, + use_pdl=True, + early_trigger=True, +): + """Drop-in for ``minimax_m3_sparse_attn_decode`` on the TP4 NVFP4 shape; callers check ``lc_applies`` first. + + ``topk_idx`` is the head-major [1, total_q, 16] view the Triton producer takes. ``mxfp8`` is ``(q8, sf)`` from + ``MERGE_MXFP8.buffers(2048, dev)`` to also emit the o_proj hand-off. The op is built with the early PDL + trigger only (``early_trigger=True``). + """ + total_q = int(q.shape[0]) + q8, sf = ( + (mxfp8[0].view(torch.uint8), mxfp8[1].view(torch.uint8)) + if mxfp8 is not None + else (None, None) + ) + m3_sparse_decode_lc( + q.view(torch.uint8), + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx[0], + block_table, + seq_lens, + output.view(total_q, 2048), + q8, + sf, + 2048 // 128, + int(decode_query_len), + float(sm_scale), + bool(linear_softmax), + bool(use_pdl), + bool(early_trigger), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_attn_mega.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_attn_mega.py new file mode 100644 index 000000000000..7c5550387f7e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_attn_mega.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 attention megakernel: ``m3_attn_mega``, one launch for a one-request, one-token decode layer. + +The QKV GEMV (bit-identical to ``m3_mx_gemv``: split 4 with contiguous k-tiles, one accumulator, the MXFP8 +activation as produced by the all-reduce, the split-K partials reduced in the owner CTA over DSMEM in rank order +and rounded to bf16), the fused decode attention and the o_proj GEMV run in one persistent grid, joined by per-tile +ready counters instead of a PDL boundary. + +The kernel is ``torch.ops.trtllm.minimax_m3_attn_mega`` (``cpp/tensorrt_llm/kernels/minimaxM3/attnMega.cu``), +compiled for geometry (1, 1, 1) with the MXFP8 hand-off only. + +Grid: 7 clusters of 16 CTAs. The first attention cluster(s) run the attention; the last 80 CTAs are the QKV units +(row tile u / 4, split u % 4; the four splits of a tile are consecutive CTAs of one cluster). Each unit's owner CTA +of a 32-row slice stores it to ``packed_out`` and release-increments the tile's counter; every attention CTA waits +for 4 more increments on all 20 tiles than at its last call (a per-CTA target it alone keeps), then reads the rows. + +One kernel replaces the per-layer chain that follows the packed QKV GEMM in short-sequence decode mode (every +request attends at most 16 KV blocks of 128 tokens, so block selection is the identity): + +* q/k/index-k RMSNorm + NeoX partial RoPE, NVFP4 K/V cache insert and FP8 index-K insert + (``minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel``), +* block-sparse GQA decode over the 16 blocks (``_gqa_sparse_decode_kernel``), +* the split-K LSE merge, optionally with the MXFP8 copy of the output (``_merge_topk_attn_out_kernel``). + +Geometry: 16 query heads, one KV head, one index head, head_dim 128, rotary_dim 64, page 128. + +Launch: clusters of ``16 / BPC`` CTAs; CTA ``r`` owns KV blocks ``[r * BPC, r * BPC + BPC)`` for ``QG`` query +tokens of one request (``decode_query_len / QG`` clusters per request). Each block keeps its own softmax partial, +exactly the Triton producer's chunk, so the merge sees the same 16 partials. + +* Before ``griddepcontrol.wait``: request metadata, RoPE rows, norm weights, the packed K chunks (registers) and the + packed V tile (``cp.async``). These are staged before the CUDA graph or written by earlier steps; nothing the QKV + GEMM writes is read, and nothing is written, before the wait. +* After the wait: K is dequantized to fp16 (exact) while the packed-QKV loads are in flight; every CTA + normalizes/ropes the query rows it needs (recomputed rather than exchanged); the CTA whose block holds a new + token's position also normalizes, ropes and quantizes that token's K, V and index-K exactly as the C++ producer + does, and writes the dequantized row into its fp16 tiles. The PDL dependents are released once every K/V byte + has been read: an earlier release lets their weight streams delay the KV reads. +* S^T = K Q^T and O^T = V^T P^T run on ``tcgen05.mma`` (FP16 operands from shared memory, FP32 accumulators in + TMEM). The softmax runs in the TMEM-reading warps with the producer's per-block math (block max, exp2, fp16 P, + log-domain LSE), so each partial equals the Triton producer's chunk partial up to accumulation order. The + producer warps dequantize V meanwhile. +* Partials go from registers to the merging CTA with ``st.async`` + mbarrier transaction counts (DSMEM); CTA ``r`` + merges its heads with the merge kernel's LSE weighting and writes the BF16 output (and the MXFP8 hand-off). + +CUDA C++ rather than Triton or CuTe DSL: the design needs cluster DSMEM pushes, loads issued before the PDL wait, +tcgen05/TMEM, and bit-exact reuse of the C++ producer's norm/RoPE/NVFP4 arithmetic. +""" + +from typing import Optional + +import torch + +from ..ops.attention.m3_attn_mega import m3_attn_mega as _m3_attn_mega_op + +NUM_HEADS = 16 +HEAD_DIM = 128 +PAGE_SIZE = 128 +MAX_BLOCKS = 16 # capacity: fixed: the MSA top-16 blocks +PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM # capacity: fixed: the tp4 packed QKV width + +_module = None + + +def _load_module(): + """The compiled ops' namespace (``torch.ops.trtllm``), resolved once; ``_module`` stays None until then.""" + global _module + if _module is None: + try: + torch.ops.trtllm.minimax_m3_attn_mega + except AttributeError as exc: + raise RuntimeError( + "torch.ops.trtllm.minimax_m3_attn_mega is not registered: this TensorRT-LLM build lacks " + "cpp/tensorrt_llm/kernels/minimaxM3/attnMega.cu" + ) from exc + _module = torch.ops.trtllm + return _module + + +def _check(cond: bool, msg: str) -> None: + if not cond: + raise ValueError(f"m3_fused_decode_attn: {msg}") + + +class MegaWorkspace: + """Per-layer persistent state of the megakernel (allocate before CUDA-graph capture; nothing is ever reset): the + two call-parity buffers of the packed QKV rows and of the attention's MXFP8 hand-off (marker-filled), and + ``flags[32 + CTA]``, each CTA's call counter (its low bit selects the parity).""" + + def __init__(self, device, max_tokens: int = 8): + # Two call-parity buffers of the packed QKV rows, all set to the 0xFFFF "not yet" marker. + self.qkv_out = torch.full( + (2, max_tokens, PACKED_WIDTH), -1, dtype=torch.int16, device=device + ).view(torch.bfloat16) + self.flags = torch.zeros(256, dtype=torch.int32, device=device) + # Two call-parity buffers of the attention's MXFP8 hand-off to the o_proj units, all 0xFF ("not yet"). + self.q8 = torch.full( + (2 * max_tokens * 2048,), 255, dtype=torch.uint8, device=device + ) # capacity: fixed: a workspace of the kernel's fixed geometry + self.sf = torch.full( + (2 * 16 * 512,), 255, dtype=torch.uint8, device=device + ) # capacity: fixed: a workspace of the kernel's fixed geometry + self.calls = 0 # host mirror of the call count (eager use only; graph replays advance the device copy) + + +def m3_attn_mega( + xq: torch.Tensor, + xsf: torch.Tensor, + qkv_weight: torch.Tensor, + qkv_weight_scale: torch.Tensor, + ws: MegaWorkspace, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + *, + eps: float, + sm_scale: float, + decode_query_len: int, + output: torch.Tensor, + y: torch.Tensor, + use_pdl: bool = True, + ar: Optional[tuple] = None, +) -> torch.Tensor: + """QKV projection (bit-identical to m3_mx_gemv) of the MXFP8 layer input ``xq`` / ``xsf``, the fused decode + attention (bit-identical to m3_fused_decode_attn) and the o_proj projection (bit-identical to m3_mx_gemv) into + ``y`` [T, 6144], in one launch. + ``output`` receives the attention rows (bf16). ``ar`` = m3_mx_gemv's (multicast workspace address, flags address, + rank, world, bucket tokens, mode 2): also push y into the next m3_mc_allreduce's exchange. Returns the packed QKV + rows of this call.""" + num_tokens = int(xq.shape[0]) + num_reqs = int(seq_lens.shape[0]) + _check( + num_tokens == num_reqs * decode_query_len and num_tokens <= ws.qkv_out.shape[0], + "token count", + ) + _check( + num_reqs == 1 and decode_query_len == 1, + "the megakernel runs one request with one query token (one " + "attention cluster); use the m3_mx_gemv + m3_fused_decode_attn + m3_mx_gemv path otherwise", + ) + _check( + xq.dtype == torch.float8_e4m3fn and xq.shape[1] == 6144 and xq.stride(1) == 1, + "xq must be e4m3 [T, 6144]", + ) + _check(qkv_weight.shape == (PACKED_WIDTH, 6144), "QKV weight must be [2560, 6144]") + out2d = output.view(num_tokens, NUM_HEADS * HEAD_DIM) + q8, sf = ws.q8, ws.sf + _check( + y.dtype == torch.bfloat16 and y.shape == (num_tokens, 6144) and y.is_contiguous(), + "y must be bf16 [T, 6144]", + ) + qg, bpc = 1, 1 # one 16-CTA attention cluster, one KV block per CTA + _load_module() + _m3_attn_mega_op( + xq.view(torch.uint8), + xsf.view(torch.uint8), + qkv_weight.view(torch.uint8), + qkv_weight_scale.view(torch.uint8), + ws.qkv_out, + ws.flags, + kv_data, + kv_scale, + index_k_cache.view(torch.uint8), + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + float(eps), + float(sm_scale), + int(decode_query_len), + qg, + bpc, + out2d, + q8, + sf, + bool(use_pdl), + o_weight.view(torch.uint8), + o_weight_scale.view(torch.uint8), + y, + *((ar[0], ar[1], ar[2], ar[3], ar[4], 2) if ar is not None else (0, 0, 0, 1, 8, 0)), + ) + out_rows = ws.qkv_out[ws.calls & 1, :num_tokens] + ws.calls += 1 + return out_rows + + +__all__ = ["MegaWorkspace", "m3_attn_mega"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_attn_mega_c1.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_attn_mega_c1.py new file mode 100644 index 000000000000..2315fbe12afd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_attn_mega_c1.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 attention + o_proj kernel: ``m3_attn_c1``, one launch for a one-request, 6-token verify layer. + +The fused decode attention on the packed QKV rows and the o_proj GEMV run in one 96-CTA launch (6 clusters of 16): +every attention CTA is also an o_proj unit, fed by the attention's MXFP8 hand-off through the workspace (Lamport, +call parity in its flags), and the o_proj rows are pushed into the next multicast all-reduce's exchange. + +The kernel is ``torch.ops.trtllm.minimax_m3_attn_c1`` (``cpp/tensorrt_llm/kernels/minimaxM3/attnC1.cu``), compiled +for ``C1_GEOMETRIES`` only. + +One kernel replaces the per-layer chain that follows the packed QKV GEMM in short-sequence decode mode (every +request attends at most 16 KV blocks of 128 tokens, so block selection is the identity): + +* q/k/index-k RMSNorm + NeoX partial RoPE, NVFP4 K/V cache insert and FP8 index-K insert + (``minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel``), +* block-sparse GQA decode over the 16 blocks (``_gqa_sparse_decode_kernel``), +* the split-K LSE merge, optionally with the MXFP8 copy of the output (``_merge_topk_attn_out_kernel``). + +Geometry: 16 query heads, one KV head, one index head, head_dim 128, rotary_dim 64, page 128. + +Launch: clusters of ``16 / BPC`` CTAs; CTA ``r`` owns KV blocks ``[r * BPC, r * BPC + BPC)`` for ``QG`` query +tokens of one request (``decode_query_len / QG`` clusters per request). Each block keeps its own softmax partial, +exactly the Triton producer's chunk, so the merge sees the same 16 partials. + +* Before ``griddepcontrol.wait``: request metadata, RoPE rows, norm weights, the packed K chunks (registers) and the + packed V tile (``cp.async``). These are staged before the CUDA graph or written by earlier steps; nothing the QKV + GEMM writes is read, and nothing is written, before the wait. +* After the wait: K is dequantized to fp16 (exact) while the packed-QKV loads are in flight; every CTA + normalizes/ropes the query rows it needs (recomputed rather than exchanged); the CTA whose block holds a new + token's position also normalizes, ropes and quantizes that token's K, V and index-K exactly as the C++ producer + does, and writes the dequantized row into its fp16 tiles. The PDL dependents are released once every K/V byte + has been read: an earlier release lets their weight streams delay the KV reads. +* S^T = K Q^T and O^T = V^T P^T run on ``tcgen05.mma`` (FP16 operands from shared memory, FP32 accumulators in + TMEM). The softmax runs in the TMEM-reading warps with the producer's per-block math (block max, exp2, fp16 P, + log-domain LSE), so each partial equals the Triton producer's chunk partial up to accumulation order. The + producer warps dequantize V meanwhile. +* Partials go from registers to the merging CTA with ``st.async`` + mbarrier transaction counts (DSMEM); CTA ``r`` + merges its heads with the merge kernel's LSE weighting and writes the BF16 output (and the MXFP8 hand-off). + +CUDA C++ rather than Triton or CuTe DSL: the design needs cluster DSMEM pushes, loads issued before the PDL wait, +tcgen05/TMEM, and bit-exact reuse of the C++ producer's norm/RoPE/NVFP4 arithmetic. +""" + +from typing import Optional + +import torch + +from ..ops.attention.m3_attn_c1 import m3_attn_c1 as _m3_attn_c1_op + +NUM_HEADS = 16 +HEAD_DIM = 128 # capacity: fixed: the checkpoint's head dim +PAGE_SIZE = 128 +MAX_BLOCKS = 16 # capacity: fixed: the MSA top-16 blocks +PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM # capacity: fixed: the tp4 packed QKV width + +# The geometries of m3_fused_attn, which default_geometry() reads (only C1_GEOMETRIES is compiled into this op). +GEOMETRIES = frozenset( + [ + (1, 1, 1), + (1, 1, 2), + (2, 1, 1), + (2, 1, 2), + (3, 1, 1), + (4, 1, 1), + (5, 1, 1), + (6, 1, 1), + (6, 1, 2), + (7, 1, 1), + (8, 1, 1), + ] +) + +_module = None + + +def _load_module(): + """The compiled ops' namespace (``torch.ops.trtllm``), resolved once; ``_module`` stays None until then.""" + global _module + if _module is None: + try: + torch.ops.trtllm.minimax_m3_attn_c1 + except AttributeError as exc: + raise RuntimeError( + "torch.ops.trtllm.minimax_m3_attn_c1 is not registered: this TensorRT-LLM build lacks " + "cpp/tensorrt_llm/kernels/minimaxM3/attnC1.cu" + ) from exc + _module = torch.ops.trtllm + return _module + + +def default_geometry(decode_query_len: int, num_requests: int) -> tuple: + """(tokens per cluster, KV blocks per CTA) for a step shape, as m3_fused_attn.default_geometry.""" + clusters = decode_query_len * num_requests + if clusters <= 6: + return (1, 1) + if (decode_query_len, 1, 2) in GEOMETRIES and clusters <= 16: + return (1, 2) + return (1, 1) + + +def _check(cond: bool, msg: str) -> None: + if not cond: + raise ValueError(f"m3_fused_decode_attn: {msg}") + + +def m3_fused_decode_attn_checks( + num_tokens, + num_reqs, + decode_query_len, + kv_data, + kv_scale, + index_k_cache, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + output, +) -> None: + """Argument checks shared by the attention entry points (the packed or fused-projection QKV source).""" + _check( + 1 <= decode_query_len <= 8, f"decode_query_len must be in [1, 8], got {decode_query_len}" + ) + _check( + num_tokens == num_reqs * decode_query_len, + "token rows must equal num_requests * decode_query_len", + ) + _check( + kv_data.dtype == torch.uint8 + and kv_data.dim() == 5 + and tuple(kv_data.shape[1:]) == (2, 1, 128, 64) + and kv_data.stride(4) == 1 + and kv_data.stride(3) == 64 + and kv_data.stride(0) % 16 == 0 + and kv_data.stride(1) % 16 == 0 + and kv_data.data_ptr() % 16 == 0, + "kv_data must be uint8 [pages, 2, 1, 128, 64] with contiguous 16-byte aligned head blocks", + ) + _check( + kv_scale.dtype == torch.uint8 + and kv_scale.dim() == 5 + and tuple(kv_scale.shape[1:]) == (2, 1, 128, 8) + and kv_scale.stride(4) == 1 + and kv_scale.stride(3) == 8 + and kv_scale.stride(0) % 16 == 0 + and kv_scale.stride(1) % 16 == 0 + and kv_scale.data_ptr() % 16 == 0 + and kv_scale.shape[0] == kv_data.shape[0], + "kv_scale must be uint8 [pages, 2, 1, 128, 8] with contiguous 16-byte aligned head blocks", + ) + _check( + index_k_cache.dtype == torch.float8_e4m3fn + and index_k_cache.dim() == 4 + and tuple(index_k_cache.shape[1:]) == (1, 128, 128) + and index_k_cache.stride(3) == 1 + and index_k_cache.stride(2) == 128 + and index_k_cache.stride(0) % 4 == 0 + and index_k_cache.shape[0] == kv_data.shape[0], + "index_k_cache must be fp8 [pages, 1, 128, 128]", + ) + for name, t in ( + ("out_cache_loc", out_cache_loc), + ("position_ids", position_ids), + ("block_table", block_table), + ("seq_lens", seq_lens), + ): + _check(t.dtype == torch.int32 and t.is_cuda, f"{name} must be CUDA int32") + _check( + out_cache_loc.is_contiguous() and out_cache_loc.numel() >= num_tokens, + "out_cache_loc too short", + ) + _check( + position_ids.is_contiguous() and position_ids.numel() == num_tokens, + "position_ids must have T entries", + ) + _check( + seq_lens.is_contiguous() + and block_table.dim() == 2 + and block_table.stride(1) == 1 + and block_table.shape[0] >= num_reqs, + "block_table must be [num_requests, max_blocks] row-major", + ) + for name, t in ( + ("kv_quant_scale", kv_quant_scale), + ("k_global_scale", k_global_scale), + ("v_global_scale", v_global_scale), + ): + _check( + t.dtype == torch.float32 and t.is_cuda and t.is_contiguous(), + f"{name} must be CUDA fp32", + ) + _check(kv_quant_scale.numel() >= 3, "kv_quant_scale must hold [Q, K, V]") + for name, t in ( + ("q_norm_weight", q_norm_weight), + ("k_norm_weight", k_norm_weight), + ("index_k_norm_weight", index_k_norm_weight), + ): + _check( + t.dtype == torch.bfloat16 + and t.numel() == HEAD_DIM + and t.is_contiguous() + and t.data_ptr() % 8 == 0, + f"{name} must be contiguous bf16 [128]", + ) + _check( + rotary_cos_sin.dtype == torch.float32 + and rotary_cos_sin.is_contiguous() + and rotary_cos_sin.dim() == 3 + and tuple(rotary_cos_sin.shape[1:]) == (2, 32), + "rotary_cos_sin must be fp32 [max_pos, 2, 32]", + ) + out2d = output.view(num_tokens, NUM_HEADS * HEAD_DIM) + _check( + output.dtype == torch.bfloat16 and out2d.stride(1) == 1 and out2d.data_ptr() % 8 == 0, + "output must be bf16 [T, 2048] with contiguous rows", + ) + + +class MegaWorkspace: + """Per-layer persistent state of the megakernel (allocate before CUDA-graph capture; nothing is ever reset): the + two call-parity buffers of the packed QKV rows and of the attention's MXFP8 hand-off (marker-filled), and + ``flags[32 + CTA]``, each CTA's call counter (its low bit selects the parity).""" + + def __init__(self, device, max_tokens: int = 8): + # Two call-parity buffers of the packed QKV rows, all set to the 0xFFFF "not yet" marker. + self.qkv_out = torch.full( + (2, max_tokens, PACKED_WIDTH), -1, dtype=torch.int16, device=device + ).view(torch.bfloat16) + self.flags = torch.zeros(256, dtype=torch.int32, device=device) + # Two call-parity buffers of the attention's MXFP8 hand-off to the o_proj units, all 0xFF ("not yet"). + self.q8 = torch.full( + (2 * max_tokens * 2048,), 255, dtype=torch.uint8, device=device + ) # capacity: fixed: a workspace of the kernel's fixed geometry + self.sf = torch.full( + (2 * 16 * 512,), 255, dtype=torch.uint8, device=device + ) # capacity: fixed: a workspace of the kernel's fixed geometry + self.calls = 0 # host mirror of the call count (eager use only; graph replays advance the device copy) + + +C1_GEOMETRIES = frozenset([(6, 1, 1)]) + + +def c1_applies(decode_query_len: int, num_requests: int) -> bool: + """Whether the step shape has a c1 kernel: 96 attention CTAs (6 clusters of 16), one o_proj unit each.""" + qg, bpc = default_geometry(decode_query_len, num_requests) + return ( + num_requests * (decode_query_len // qg) * (16 // bpc) == 96 + and (decode_query_len, qg, bpc) in C1_GEOMETRIES + ) + + +def m3_attn_c1( + packed: torch.Tensor, + ws: MegaWorkspace, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + *, + eps: float, + sm_scale: float, + decode_query_len: int, + output: torch.Tensor, + y: torch.Tensor, + use_pdl: bool = True, + ar: Optional[tuple] = None, +) -> None: + """The fused decode attention on the packed QKV rows (bit-identical to m3_fused_decode_attn) and the o_proj + projection (bit-identical to m3_mx_gemv) into ``y`` [T, 6144], in one 96-CTA launch: every attention CTA is also + an o_proj unit, fed by the attention's MXFP8 hand-off through ``ws`` (Lamport, call parity in ``ws.flags``). + ``ar`` = m3_mx_gemv's (multicast workspace address, flags address, rank, world, bucket tokens): also push y into + the next m3_mc_allreduce's exchange.""" + num_tokens = int(packed.shape[0]) + num_reqs = int(seq_lens.shape[0]) + _check( + packed.dtype == torch.bfloat16 + and packed.dim() == 2 + and packed.shape[1] == PACKED_WIDTH + and packed.stride(1) == 1 + and packed.stride(0) % 4 == 0 + and packed.data_ptr() % 8 == 0, + "packed must be bf16 [T, 2560] with 8-byte aligned contiguous rows", + ) + m3_fused_decode_attn_checks( + num_tokens, + num_reqs, + decode_query_len, + kv_data, + kv_scale, + index_k_cache, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + output, + ) + qg, bpc = default_geometry(decode_query_len, num_reqs) + _check( + c1_applies(decode_query_len, num_reqs), f"no c1 kernel for {num_reqs} x {decode_query_len}" + ) + _check( + y.dtype == torch.bfloat16 and y.shape == (num_tokens, 6144) and y.is_contiguous(), + "y must be bf16 [T, 6144]", + ) + out2d = output.view(num_tokens, NUM_HEADS * HEAD_DIM) + _load_module() + _m3_attn_c1_op( + packed, + ws.flags, + kv_data, + kv_scale, + index_k_cache.view(torch.uint8), + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + float(eps), + float(sm_scale), + int(decode_query_len), + qg, + bpc, + out2d, + ws.q8, + ws.sf, + bool(use_pdl), + o_weight.view(torch.uint8), + o_weight_scale.view(torch.uint8), + y, + *((ar[0], ar[1], ar[2], ar[3], ar[4], 2) if ar is not None else (0, 0, 0, 1, 8, 0)), + ) + ws.calls += 1 + + +__all__ = ["MegaWorkspace", "c1_applies", "m3_attn_c1"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_draft_step.py new file mode 100644 index 000000000000..31768ebcd2ee --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_draft_step.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused greedy Eagle3 draft pass for the MiniMax-M3 drafter's vocab-sharded BF16 head (TP decode, <= 8 requests). + +One kernel replaces a draft pass's head and the draft loop's per-pass bookkeeping, at the in-tree precision: + +* The drafter's BF16 lm_head shard is kept a second time, built by ``prepare_head`` before the executor's warmup: in + the lossless 12-bit coding of the target head (``ops.gemm.m3_target_head``; 75% of the bytes, kept only when it + decodes to the BF16 shard bit for bit), else in ``mma.m16n8k16`` A-fragment order (``bf16_fragments``, the same + values). The kernel streams it with BF16 MMAs and FP32 accumulation over the BF16 head input rows; a coded + exception decodes to +0 and its exact product x * w is added at its tile's flush. +* Each logit is rounded to BF16 (the in-tree logits are the BF16 head GEMM's output) and keyed + ``ordered(value) << 32 | ~global_index`` with -0 as +0, so the larger value wins and ties take the lowest index, as + the in-tree ``torch.max`` / ``torch.argmax`` over the gathered logits (``greedy_sample_draft_with_tp_gather``). +* The last CTA exchanges each row's key across the TP ranks through a multicast mapping (Lamport slots, see + ``m3_shard_accept``), then stores the draft token, the next pass's input id and position, and advances the draft KV + lengths as the in-tree draft loop does between passes. + +The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/draftStep.cu) and called through the +catalog entry ``ops.spec.m3_draft_step``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from ..ops.gemm.m3_target_head import encode_chunked, round_trip_ok +from ..ops.spec.m3_draft_step import MAX_TILES_PER_CTA, MAX_TOKENS, bf16_fragments +from ..ops.spec.m3_draft_step import m3_draft_step as _draft_step_op +from .m3_shard_accept import AcceptWorkspace, create_workspace + + +@dataclass(eq=False) +class DraftHead: + """The streamed copy of one rank's draft lm_head shard and its first global id: ``wf`` is the 12-bit coded stream + with its exception lists (``coded``), or the fragment-ordered BF16 copy (``bf16_fragments``).""" + + weight: torch.Tensor + wf: torch.Tensor + vocab_offset: int + coded: bool = False + exc_off: Optional[torch.Tensor] = None + exc_pos: Optional[torch.Tensor] = None + exc_w: Optional[torch.Tensor] = None + exp_base: int = 0 + + +def head_supported(lm_head) -> bool: + """A bias-free BF16 head whose local shard the kernel can own: whole 16-row tiles, K a multiple of 16, at most + MAX_TILES_PER_CTA tiles per CTA of a one-CTA-per-SM grid.""" + weight = getattr(lm_head, "weight", None) + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: + return False + if getattr(lm_head, "bias", None) is not None: + return False + n, k = weight.shape + if n % 16 or k % 16: + return False + return (n // 16 + _num_sms() - 1) // _num_sms() <= MAX_TILES_PER_CTA + + +def prepare_head(lm_head, tp_rank: int, coded: bool) -> DraftHead: + """The streamed copy of ``lm_head``'s local shard (global ids from tp_rank * rows): its 12-bit coding when + ``coded``, K is a multiple of 64 and the coding decodes to the shard bit for bit, else (fail + closed) the fragment-ordered BF16 copy. + Built once, eagerly, before the executor's warmup (``M3Eagle3Worker.prebuild_draft_head``).""" + weight = lm_head.weight + offset = tp_rank * weight.shape[0] + with torch.no_grad(): + w = weight.detach() + if coded and w.shape[1] % 64 == 0: + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + if round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w): + return DraftHead( + weight=weight, + wf=packed, + vocab_offset=offset, + coded=True, + exc_off=exc_off, + exc_pos=exc_pos, + exc_w=exc_w.view(-1), + exp_base=base, + ) + del packed, exc_off, exc_pos, exc_w + wf = bf16_fragments(w) + return DraftHead(weight=weight, wf=wf, vocab_offset=offset) + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace ([2][MAX_TOKENS][world] slots), created collectively on the first eager call; None + during a CUDA-graph capture that finds it missing.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + # A capture that finds it missing is the caller's miss (it is built on the eager steps). + return None + _WS = create_workspace(mapping, device, MAX_TOKENS) + return _WS + + +def draft_step( + head: DraftHead, + x: torch.Tensor, + rows: int, + x_rows: Optional[torch.Tensor], + mapping, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, +) -> bool: + """One fused greedy draft pass over ``rows`` requests (see the module docstring); False when the exchange + workspace is missing during a capture. Collective over the TP group. + + Args: + head: the prepared head of this rank (``prepare_head``), coded or BF16 fragments. + x: [*, K] BF16 head input rows (the drafter's final-norm output); row ``x_rows[b]`` (int64) or ``b`` is + request b's. + draft_out: int32 [rows, >= step + 1], unit column stride: column ``step`` receives the tokens. + ids_next / pos_next: int32 [rows]: the tokens, and ``pos_src[x_rows[b] or b] + 1`` (pos_src int32). + kv_lens: int32 draft KV lengths: ``-= draft_len - num_acc`` with ``seq_lens = 1`` on pass 0 when ``num_acc`` + is given, ``+= 1`` on later passes. + """ + ws = workspace(mapping, x.device) + if ws is None: + return False + _draft_step_op( + x, + x_rows, + rows, + head.wf, + head.vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + _num_sms(), + True, + None, + head.exc_off, + head.exc_pos, + head.exc_w, + head.exp_base, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_force_vec.py new file mode 100644 index 000000000000..b4de37794ca9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_force_vec.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance (the in-tree speculative worker's +``force_num_accepted_tokens``), drawn exactly as its ``_apply_force_accepted_tokens`` does, in one launch. + +Triton: one CTA bumps the shared step counter and draws each request's count from the fixed-seed pool with the +in-tree (counter, slot) hash. Reached through the ``ops/spec`` entry ``m3_force_vec``. +""" + +import triton +import triton.language as tl + + +@triton.jit +def force_vec_kernel( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK: tl.constexpr, +): + c = tl.load(counter) + 1 + slots = tl.arange(0, BLOCK) + tl.store(counter, c) + idx = (c * counter_stride + slots.to(tl.int64) * slot_stride) & pool_mask + live = slots < batch_size + r = tl.load(pool + idx, mask=live, other=1.0) + v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) + tl.store(out + slots, v, mask=live) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_fused_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_fused_attn.py new file mode 100644 index 000000000000..e4bd922a1587 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_fused_attn.py @@ -0,0 +1,354 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 fused short-sequence decode attention: ``m3_fused_decode_attn``. + +The kernel is ``torch.ops.trtllm.minimax_m3_fused_decode_attn`` (``cpp/tensorrt_llm/kernels/minimaxM3/ +fusedDecodeAttn.cu``); it is compiled for the geometries in ``GEOMETRIES``, each with the MXFP8 hand-off. + +One kernel replaces the per-layer chain that follows the packed QKV GEMM in short-sequence decode mode (every +request attends at most 16 KV blocks of 128 tokens, so block selection is the identity): + +* q/k/index-k RMSNorm + NeoX partial RoPE, NVFP4 K/V cache insert and FP8 index-K insert + (``minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel``), +* block-sparse GQA decode over the 16 blocks (``_gqa_sparse_decode_kernel``), +* the split-K LSE merge, optionally with the MXFP8 copy of the output (``_merge_topk_attn_out_kernel``). + +Geometry: 16 query heads, one KV head, one index head, head_dim 128, rotary_dim 64, page 128. + +Launch: clusters of ``16 / BPC`` CTAs; CTA ``r`` owns KV blocks ``[r * BPC, r * BPC + BPC)`` for ``QG`` query +tokens of one request (``decode_query_len / QG`` clusters per request). Each block keeps its own softmax partial, +exactly the Triton producer's chunk, so the merge sees the same 16 partials. + +* Before ``griddepcontrol.wait``: request metadata, RoPE rows, norm weights, the packed K chunks (registers) and the + packed V tile (``cp.async``). These are staged before the CUDA graph or written by earlier steps; nothing the QKV + GEMM writes is read, and nothing is written, before the wait. +* After the wait: K is dequantized to fp16 (exact) while the packed-QKV loads are in flight; every CTA + normalizes/ropes the query rows it needs (recomputed rather than exchanged); the CTA whose block holds a new + token's position also normalizes, ropes and quantizes that token's K, V and index-K exactly as the C++ producer + does, and writes the dequantized row into its fp16 tiles. The PDL dependents are released once every K/V byte + has been read: an earlier release lets their weight streams delay the KV reads. +* S^T = K Q^T and O^T = V^T P^T run on ``tcgen05.mma`` (FP16 operands from shared memory, FP32 accumulators in + TMEM). The softmax runs in the TMEM-reading warps with the producer's per-block math (block max, exp2, fp16 P, + log-domain LSE), so each partial equals the Triton producer's chunk partial up to accumulation order. The + producer warps dequantize V meanwhile. +* Partials go from registers to the merging CTA with ``st.async`` + mbarrier transaction counts (DSMEM); CTA ``r`` + merges its heads with the merge kernel's LSE weighting and writes the BF16 output (and the MXFP8 hand-off). + +CUDA C++ rather than Triton or CuTe DSL: the design needs cluster DSMEM pushes, loads issued before the PDL wait, +tcgen05/TMEM, and bit-exact reuse of the C++ producer's norm/RoPE/NVFP4 arithmetic. +""" + +from typing import Optional + +import torch + +from ..ops.attention.m3_fused_decode_attn import m3_fused_decode_attn as _m3_fused_decode_attn_op + +NUM_HEADS = 16 # capacity: fixed: the tp4 query heads +HEAD_DIM = 128 # capacity: fixed: the checkpoint's head dim +PAGE_SIZE = 128 +MAX_BLOCKS = 16 # capacity: fixed: the MSA top-16 blocks +PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM # capacity: fixed: the tp4 packed QKV width + +# (decode_query_len, tokens per cluster, KV blocks per CTA) geometries compiled into the op: what +# default_geometry() returns for every step shape the target routes here (at most 16 tokens). +GEOMETRIES = frozenset( + [ + (1, 1, 1), + (1, 1, 2), + (2, 1, 1), + (2, 1, 2), + (3, 1, 1), + (4, 1, 1), + (5, 1, 1), + (6, 1, 1), + (6, 1, 2), + (7, 1, 1), + (8, 1, 1), + ] +) + +_module = None + + +def _load_module(): + """The compiled ops' namespace (``torch.ops.trtllm``), resolved once; ``_module`` stays None until then.""" + global _module + if _module is None: + try: + torch.ops.trtllm.minimax_m3_fused_decode_attn + except AttributeError as exc: + raise RuntimeError( + "torch.ops.trtllm.minimax_m3_fused_decode_attn is not registered: this TensorRT-LLM " + "build lacks cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.cu" + ) from exc + _module = torch.ops.trtllm + return _module + + +def default_geometry(decode_query_len: int, num_requests: int) -> tuple: + """(tokens per cluster, KV blocks per CTA) for a step shape: one token per 16-CTA cluster up to 6 clusters, + then KV blocks paired into 8-CTA clusters up to 16 clusters where that geometry is compiled.""" + clusters = decode_query_len * num_requests + if clusters <= 6: + return (1, 1) + if (decode_query_len, 1, 2) in GEOMETRIES and clusters <= 16: + return (1, 2) + return (1, 1) + + +def _check(cond: bool, msg: str) -> None: + if not cond: + raise ValueError(f"m3_fused_decode_attn: {msg}") + + +def m3_fused_decode_attn_checks( + num_tokens, + num_reqs, + decode_query_len, + kv_data, + kv_scale, + index_k_cache, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + output, +) -> None: + """Argument checks shared by the attention entry points (the packed or fused-projection QKV source).""" + _check( + 1 <= decode_query_len <= 8, f"decode_query_len must be in [1, 8], got {decode_query_len}" + ) + _check( + num_tokens == num_reqs * decode_query_len, + "token rows must equal num_requests * decode_query_len", + ) + _check( + kv_data.dtype == torch.uint8 + and kv_data.dim() == 5 + and tuple(kv_data.shape[1:]) == (2, 1, 128, 64) + and kv_data.stride(4) == 1 + and kv_data.stride(3) == 64 + and kv_data.stride(0) % 16 == 0 + and kv_data.stride(1) % 16 == 0 + and kv_data.data_ptr() % 16 == 0, + "kv_data must be uint8 [pages, 2, 1, 128, 64] with contiguous 16-byte aligned head blocks", + ) + _check( + kv_scale.dtype == torch.uint8 + and kv_scale.dim() == 5 + and tuple(kv_scale.shape[1:]) == (2, 1, 128, 8) + and kv_scale.stride(4) == 1 + and kv_scale.stride(3) == 8 + and kv_scale.stride(0) % 16 == 0 + and kv_scale.stride(1) % 16 == 0 + and kv_scale.data_ptr() % 16 == 0 + and kv_scale.shape[0] == kv_data.shape[0], + "kv_scale must be uint8 [pages, 2, 1, 128, 8] with contiguous 16-byte aligned head blocks", + ) + _check( + index_k_cache.dtype == torch.float8_e4m3fn + and index_k_cache.dim() == 4 + and tuple(index_k_cache.shape[1:]) == (1, 128, 128) + and index_k_cache.stride(3) == 1 + and index_k_cache.stride(2) == 128 + and index_k_cache.stride(0) % 4 == 0 + and index_k_cache.shape[0] == kv_data.shape[0], + "index_k_cache must be fp8 [pages, 1, 128, 128]", + ) + for name, t in ( + ("out_cache_loc", out_cache_loc), + ("position_ids", position_ids), + ("block_table", block_table), + ("seq_lens", seq_lens), + ): + _check(t.dtype == torch.int32 and t.is_cuda, f"{name} must be CUDA int32") + _check( + out_cache_loc.is_contiguous() and out_cache_loc.numel() >= num_tokens, + "out_cache_loc too short", + ) + _check( + position_ids.is_contiguous() and position_ids.numel() == num_tokens, + "position_ids must have T entries", + ) + _check( + seq_lens.is_contiguous() + and block_table.dim() == 2 + and block_table.stride(1) == 1 + and block_table.shape[0] >= num_reqs, + "block_table must be [num_requests, max_blocks] row-major", + ) + for name, t in ( + ("kv_quant_scale", kv_quant_scale), + ("k_global_scale", k_global_scale), + ("v_global_scale", v_global_scale), + ): + _check( + t.dtype == torch.float32 and t.is_cuda and t.is_contiguous(), + f"{name} must be CUDA fp32", + ) + _check(kv_quant_scale.numel() >= 3, "kv_quant_scale must hold [Q, K, V]") + for name, t in ( + ("q_norm_weight", q_norm_weight), + ("k_norm_weight", k_norm_weight), + ("index_k_norm_weight", index_k_norm_weight), + ): + _check( + t.dtype == torch.bfloat16 + and t.numel() == HEAD_DIM + and t.is_contiguous() + and t.data_ptr() % 8 == 0, + f"{name} must be contiguous bf16 [128]", + ) + _check( + rotary_cos_sin.dtype == torch.float32 + and rotary_cos_sin.is_contiguous() + and rotary_cos_sin.dim() == 3 + and tuple(rotary_cos_sin.shape[1:]) == (2, 32), + "rotary_cos_sin must be fp32 [max_pos, 2, 32]", + ) + out2d = output.view(num_tokens, NUM_HEADS * HEAD_DIM) + _check( + output.dtype == torch.bfloat16 and out2d.stride(1) == 1 and out2d.data_ptr() % 8 == 0, + "output must be bf16 [T, 2048] with contiguous rows", + ) + + +def m3_fused_decode_attn( + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + *, + eps: float, + sm_scale: float, + decode_query_len: int, + output: torch.Tensor, + mxfp8_out: Optional[tuple] = None, + use_pdl: bool = True, + geometry: Optional[tuple] = None, +) -> None: + """Fused norm/RoPE/KV-insert + short-sequence sparse decode + split-K merge, written into ``output``. + + packed: [T, 2560] bf16 QKV+index GEMM output, T = num_requests * decode_query_len. + kv_data: uint8 [num_pages, 2, 1, 128, 64] HND NVFP4 pool (K plane 0, V plane 1); written. + kv_scale: uint8 [num_pages, 2, 1, 128, 8] E4M3 block scales (K linear, V 4x4 token quads); written. + index_k_cache: float8_e4m3fn [num_pages, 1, 128, 128]; written. + out_cache_loc, position_ids: int32 [>= T] / [T]; block_table int32 [num_requests, >= 1]; seq_lens int32 + [num_requests] (attended length including this step's tokens, at most 16 * 128). + kv_quant_scale: fp32 [3] (Q, K, V) quantization scales; k/v_global_scale: fp32 [1] dequantization scales. + output: bf16 [T, 2048] (or [T, 16, 128], contiguous rows). + mxfp8_out: (e4m3 [>= T, 2048], uint8 swizzled UE8M0 scales) hand-off for the o_proj GEMM (required: only + the MXFP8 variants are compiled), + same layout as the merge kernel's MERGE_MXFP8 output. + geometry: (query tokens per cluster, KV blocks per CTA); default_geometry() when None. + """ + num_tokens = int(packed.shape[0]) + num_reqs = int(seq_lens.shape[0]) + _check( + 1 <= decode_query_len <= 8, f"decode_query_len must be in [1, 8], got {decode_query_len}" + ) + _check( + num_tokens == num_reqs * decode_query_len, + "packed rows must equal num_requests * decode_query_len", + ) + _check( + packed.dtype == torch.bfloat16 + and packed.dim() == 2 + and packed.shape[1] == PACKED_WIDTH + and packed.stride(1) == 1 + and packed.stride(0) % 4 == 0 + and packed.data_ptr() % 8 == 0, + "packed must be bf16 [T, 2560] with 8-byte aligned contiguous rows", + ) + m3_fused_decode_attn_checks( + num_tokens, + num_reqs, + decode_query_len, + kv_data, + kv_scale, + index_k_cache, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + output, + ) + out2d = output.view(num_tokens, NUM_HEADS * HEAD_DIM) + _check( + mxfp8_out is not None, + "only the variants with the MXFP8 hand-off are compiled; pass mxfp8_out", + ) + q8, sf = mxfp8_out + _check( + q8.dtype == torch.float8_e4m3fn + and q8.dim() == 2 + and q8.shape[1] == NUM_HEADS * HEAD_DIM + and q8.shape[0] >= num_tokens + and q8.is_contiguous(), + "mxfp8 data must be e4m3 [>= T, 2048]", + ) + _check(sf.dtype == torch.uint8 and sf.is_contiguous(), "mxfp8 scales must be contiguous uint8") + q8 = q8.view(torch.uint8) + if num_tokens == 0: + return + qg, bpc = geometry if geometry is not None else default_geometry(decode_query_len, num_reqs) + _check( + (decode_query_len, qg, bpc) in GEOMETRIES, + f"geometry {(decode_query_len, qg, bpc)} is not compiled", + ) + _load_module() + _m3_fused_decode_attn_op( + packed, + kv_data, + kv_scale, + index_k_cache.view(torch.uint8), + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + float(eps), + float(sm_scale), + int(decode_query_len), + qg, + bpc, + out2d, + q8, + sf, + bool(use_pdl), + ) + + +__all__ = ["GEOMETRIES", "default_geometry", "m3_fused_decode_attn"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_index_score.py new file mode 100644 index 000000000000..fad75d44694a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_index_score.py @@ -0,0 +1,638 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py +"""MiniMax-M3 index decode block scorer (CuTe DSL, Blackwell) with its pre-wait schedule. + +The kernel of ``tensorrt_llm/_torch/attention/kernels/blackwell/minimax_m3_index_decode_score.py`` (the same scores, bit +for bit) with two schedule options this target compiles in: + +* ``prefetch_old``: before the grid-dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that hold only + tokens older than this step (the producer ahead writes only this step's tokens), up to ``prefetch_depth`` blocks per + CTA (0 = all of them). +* ``early_issue``: the TMA warp reads a ``seq_lens`` hint and its first page id before the wait and issues the Q tile + and the first K tile right after it. The tiles are read after the wait, so their contents are final; the block table + is only written between graph replays, so the early page id is the right one whenever the block is valid. A CTA + whose early tiles lie past the request's true length drains them before it exits. + +Every kernel is compiled before the executor's warmup (``layers.prebuild_index_scorer``) through ``IndexScoreRunner`` +and called through this target's ``m3_index_score`` entry. +""" + +import cutlass +import torch +from cuda.bindings.driver import CUstream +from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import ( + EVICT_FIRST, + fp8x4_to_fp16x4, + mma_sync, + simple_tma_copy, +) +from tensorrt_llm._torch.kernels.blackwell.utils import TRTLLM_ENABLE_PDL, griddepcontrol_wait + +__all__ = ["IndexDecodeScoreKernel", "IndexScoreRunner"] + + +@dsl_user_op +def _l2_prefetch_bulk(addr: Int64, size: Int32, *, loc=None, ip=None) -> None: + """Prefetch ``size`` contiguous bytes (multiple of 16) at global ``addr`` into L2.""" + llvm.inline_asm( + res=None, + operands_=[Int64(addr).ir_value(loc=loc, ip=ip), Int32(size).ir_value(loc=loc, ip=ip)], + asm_string="cp.async.bulk.prefetch.L2.global [$0], $1;", + constraints="l,r", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _fp8_to_f16_mma_fragments(src: cute.Tensor): + """Split an FP8 ldmatrix fragment into the two FP16 k-fragments MMA wants.""" + src_elems = cute.size(src) + src_u32 = cute.recast_tensor(src, Uint32) + src_f16 = cute.make_rmem_tensor(src_elems, Float16) + src_f16_u32 = cute.recast_tensor(src_f16, Uint32) + # Packed conversion; faster and fewer SASS instructions than + # src.load().to(Float16). + for i in cutlass.range_constexpr(src_elems // 4): + converted = fp8x4_to_fp16x4(src_u32[i]) + src_f16_u32[i * 2] = converted[0] + src_f16_u32[i * 2 + 1] = converted[1] + lower = cute.make_rmem_tensor(src_elems // 2, Float16) + upper = cute.make_rmem_tensor(src_elems // 2, Float16) + + # FP8 ldmatrix yields four consecutive values along K. Split each group + # into the lower two and upper two for the two FP16 MMA k-fragments. + for i in cutlass.range_constexpr(src_elems // 2): + lower[i] = src_f16[(i // 2) * 4 + i % 2] + upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] + return lower, upper + + +class IndexDecodeScoreKernel: + """Per-block max index score for one decode step. + + Grid is (batch, split_k): CTA (b, s) walks KV blocks + s, s + split_k, s + 2 * split_k, ... of request b and writes each block's + score directly, so no cross-CTA reduction is needed. CTAs whose split_id + exceeds the request's block count exit immediately. + """ + + BLOCK_K = 128 + BAR_MMA = 1 + num_stages = 2 + + def __init__( + self, + dtype: type[cutlass.Numeric], + num_heads: int, + max_decode_query_len: int, + split_k: int, + head_dim: int = 128, + prefetch_old: bool = False, + early_issue: bool = False, + prefetch_depth: int = 0, + ): + # prefetch_old: before the grid dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that + # hold only tokens older than this step (the predecessor writes only this step's tokens). + self.prefetch_old = prefetch_old + # early_issue: the TMA warp reads a seq_lens hint and this CTA's first page id before the grid dependency + # wait, then issues the Q tile and the first K tile right after it, overlapping the post-wait seq_lens load. + # The tiles are read after the wait, so their contents are final; the block table is only written between + # graph replays, so the early page id is the right one whenever the block is valid. A CTA whose early tiles + # turn out to lie past the request's true length drains them before exiting. + self.early_issue = early_issue + # prefetch_depth: blocks per CTA that prefetch_old covers (0 = all of them). Prefetching every older page at + # large batches floods L2; the first few per CTA are the ones whose latency the pipeline cannot hide. + self.prefetch_depth = prefetch_depth + self.dtype = dtype + self.num_heads = num_heads + self.max_decode_query_len = max_decode_query_len + self.split_k = split_k + self.head_dim = head_dim + + def _swizzle_elems(self) -> int: + """Elements spanned by one 128-byte swizzle atom.""" + return 128 * 8 // self.dtype.width + + def _sq_layout(self): + """Composed SMEM layout for the Q tile, shared by the descriptor and the kernel.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_q = self.num_heads * self.max_decode_query_len + layout = cute.make_layout( + (self.max_decode_query_len, self.num_heads, (elems, head_dim // elems)), + stride=(elems, self.max_decode_query_len * elems, (1, block_q * elems)), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + def _sk_layout(self): + """Composed SMEM layout for the pipelined K tiles.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_k = self.BLOCK_K + layout = cute.make_layout( + (1, block_k, (elems, head_dim // elems), self.num_stages), + stride=(0, elems, (1, block_k * elems), block_k * head_dim), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + @cute.jit + def __call__( + self, + gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] + gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] + block_table: cute.Tensor, # [bs, max_pages] + score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] + seq_lens: cute.Tensor, # [bs] + stream: CUstream, + ): + num_heads = self.num_heads + head_dim = self.head_dim + MAX_DQL = self.max_decode_query_len + assert num_heads * MAX_DQL <= 32 + + batch = seq_lens.shape[0] + decode_query_len = gQ.shape[0] // batch + grid = (batch, self.split_k, 1) + block = (32 * 5, 1, 1) + + tma_g2s = cpasync.CopyBulkTensorTileG2SOp() + elems = self._swizzle_elems() + + q_tma_atom, q_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gQ, (None, None, elems)), + self._sq_layout(), + (MAX_DQL, num_heads, head_dim), + ) + k_tma_atom, k_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gK_cache, (None, None, elems)), + self._sk_layout(), + (1, self.BLOCK_K, head_dim), + ) + + self.kernel( + q_tma_atom, + q_tma_tensor, + k_tma_atom, + k_tma_tensor, + gK_cache, + block_table, + score, + seq_lens, + decode_query_len, + ).launch(grid=grid, block=block, stream=stream, use_pdl=TRTLLM_ENABLE_PDL) + + @cute.kernel + def kernel( + self, + q_tma_atom: cute.CopyAtom, + q_tma_tensor: cute.Tensor, + k_tma_atom: cute.CopyAtom, + k_tma_tensor: cute.Tensor, + gK_cache: cute.Tensor, + block_table: cute.Tensor, + score: cute.Tensor, + seq_lens: cute.Tensor, + decode_query_len, + ): + batch_id, split_id, _ = cute.arch.block_idx() + _, split_k, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = cute.arch.lane_idx() + + NUM_HEADS = self.num_heads + MAX_DQL = self.max_decode_query_len + BLOCK_Q = NUM_HEADS * MAX_DQL + BLOCK_K = self.BLOCK_K + head_dim = self.head_dim + dtype = self.dtype + MMA_N = 8 + num_stages = self.num_stages + Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) + EPI_Q = Q_TILES * MMA_N + + sq_layout = self._sq_layout() + sk_layout = self._sk_layout() + + smem = SmemAllocator() + sK = smem.allocate_tensor( + dtype, + sk_layout.outer, + byte_alignment=128, + swizzle=sk_layout.inner, + )[0, None, None, None] + # sQ aliases the first K stage: Q is consumed into registers before the + # first K tile is needed, so the two never overlap in time. + sQ_tma = cute.make_tensor(sK[None, None, 0].iterator, layout=sq_layout.outer) + # TMA sees Q as (query, head, dim) while ldmatrix consumes a flattened Q + # column mode. The target profile keeps the rank-2 view even for + # degenerate shapes such as MAX_DQL == 1. + q_tma_elems = self._swizzle_elems() + sQ = cute.coalesce( + cute.group_modes(sQ_tma, 0, 2), + target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), + ) + epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + + # Barrier init and descriptor prefetch do not depend on the predecessor + # grid, so they are the work PDL can overlap. Every CTA runs them, even + # the ones that own no blocks, since ownership is not known until + # seq_lens is read below. + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, 1) + cute.arch.mbarrier_init(tma_empty_mbar + i, 128) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(q_tma_atom) + cpasync.prefetch_descriptor(k_tma_atom) + cute.arch.sync_threads() + + if cutlass.const_expr(self.prefetch_old): + if warp_id == 4: + with cute.arch.elect_one(): + seqlen0 = seq_lens[batch_id] + num_blocks0 = cute.ceil_div(seqlen0, BLOCK_K) + stable = (seqlen0 - decode_query_len) // BLOCK_K + if stable > num_blocks0: + stable = num_blocks0 + if cutlass.const_expr(self.prefetch_depth > 0): + if stable > split_id + self.prefetch_depth * split_k: + stable = split_id + self.prefetch_depth * split_k + for block_id in range(split_id, stable, split_k): + page = block_table[batch_id, block_id] + _l2_prefetch_bulk( + gK_cache[page, None, None].iterator.toint().to(Int64), + Int32(BLOCK_K * head_dim * (dtype.width // 8)), + ) + + # No explicit launch_dependents: the only consumer is the block selector, + # which reduces over every block this grid writes, so triggering it early + # would leave it resident and spinning in its own wait. The implicit + # trigger as the CTAs exit still gives it the launch overlap. + early_q = False + early_page = Int32(0) + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + hint_blocks = cute.ceil_div(seq_lens[batch_id], BLOCK_K) + if hint_blocks > block_table.shape[1]: + hint_blocks = block_table.shape[1] + if split_id < hint_blocks: + early_q = True + early_page = block_table[batch_id, split_id] + + griddepcontrol_wait() + + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + gQ_early = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar, BLOCK_Q * head_dim * (dtype.width // 8) + ) + simple_tma_copy(q_tma_atom, gQ_early, sQ_tma, tma_full_mbar) + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, BLOCK_K * head_dim * (dtype.width // 8) + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + # seq_lens must be read after the wait. It is a graph-stable buffer that + # on_update_kv_lens patches every step, and num_blocks bounds an unmasked + # block_table load and an unmasked score store, so a stale length is an + # out-of-bounds write rather than just a wrong answer. + seqlen = seq_lens[batch_id] + num_blocks = cute.ceil_div(seqlen, BLOCK_K) + + if split_id < num_blocks: + if warp_id == 4: + # TMA warp + tma_stage = 0 + tma_parity = 1 + first_block = split_id + + if not early_q: + gQ_tile = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) + with cute.arch.elect_one(): + Q_size = BLOCK_Q * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) + # TMA bounds-checks rows when the runtime decode_query_len is + # below MAX_DQL; padded Q columns are masked before the stores. + simple_tma_copy(q_tma_atom, gQ_tile, sQ_tma, tma_full_mbar) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + if early_q: + # The first K tile was issued right after the wait. + first_block = split_id + split_k + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + for block_id in range(first_block, num_blocks, split_k): + page_id = block_table[batch_id, block_id] + gK_tile = k_tma_tensor[page_id, None, None] + k_mbar = tma_full_mbar + tma_stage + + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) + with cute.arch.elect_one(): + K_size = BLOCK_K * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) + simple_tma_copy( + k_tma_atom, + gK_tile, + sK[None, None, tma_stage], + k_mbar, + cache_policy=EVICT_FIRST, + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + else: + # MMA warps; each handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T + sK_warp = cute.local_tile(sK, (32, head_dim, num_stages), (warp_id, 0, 0)) + q_start = seqlen - decode_query_len + + elems = 128 // dtype.width # 16B + MMA_K = 32 * 8 // dtype.width # 32B + + # Pre-compute ldmatrix addresses. + # sK loads a [16 x 16B] tile: + # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) + # sQ loads an [8 x 32B] tile: + # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = cute.zipped_divide(sK_warp, (16, cute.make_layout((elems, 2)), 1)) + sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) + + # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) + # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] + sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] + + ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) + ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) + + rQ = cute.make_rmem_tensor( + ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype + ) + rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) + rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + for q in cutlass.range_constexpr(Q_TILES): + cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) + cute.arch.mbarrier_arrive(tma_empty_mbar) + + tma_stage = 1 % self.num_stages + tma_parity = 0 + if tma_stage == 0: + tma_parity ^= 1 + + if cutlass.const_expr(dtype is Float8E4M3FN): + rQ_f16 = cute.make_rmem_tensor((4, head_dim // MMA_K, Q_TILES, 2), Float16) + q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) + rQ_f16[None, None, None, 0].store(q_lower.load()) + rQ_f16[None, None, None, 1].store(q_upper.load()) + + for block_id in range(split_id, num_blocks, split_k): + rC.fill(0.0) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + for k in cutlass.range_constexpr(head_dim // MMA_K): + cute.copy( + ldsm_atom, + sK_ldsm[None, (None, k, tma_stage)], + rK[None, None, k], + ) + for m in cutlass.range_constexpr(2): + if cutlass.const_expr(dtype is Float8E4M3FN): + rK_lower, rK_upper = _fp8_to_f16_mma_fragments(rK[None, m, k]) + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK_lower, + rQ_f16[None, k, n, 0], + rC[None, m, n], + ) + rC[None, m, n] = mma_sync( + rK_upper, + rQ_f16[None, k, n, 1], + rC[None, m, n], + ) + else: + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK[None, m, k], + rQ[(None, k % 2), k // 2, n], + rC[None, m, n], + ) + + cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) + + k_start = block_id * BLOCK_K + warp_id * 32 + + for q in cutlass.range_constexpr(Q_TILES): + for i in cutlass.range_constexpr(4): + for j in cutlass.range_constexpr(2): + col = q * 8 + (lane_id % 4) * 2 + j + q_local_pos = col % MAX_DQL + q_pos = q_start + q_local_pos + k_pos = k_start + i * 8 + lane_id // 4 + rC[q * 8 + i * 2 + j] = ( + rC[q * 8 + i * 2 + j] if q_pos >= k_pos else float("-inf") + ) + + for q in cutlass.range_constexpr(Q_TILES): + # Thread-local reduction along the BLOCK_K dim. + rScore = cute.make_rmem_tensor(2, Float32) + rScore.fill(float("-inf")) + for i in cutlass.range_constexpr(4): + rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) + rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) + + # Warp reduction among lanes 0, 4, 8, 12, ... + for i in cutlass.range_constexpr(3): + offset = 4 << i + other0 = cute.arch.shuffle_sync_bfly( + rScore[0], offset=offset, mask=-1, mask_and_clamp=31 + ) + other1 = cute.arch.shuffle_sync_bfly( + rScore[1], offset=offset, mask=-1, mask_and_clamp=31 + ) + rScore[0] = cute.arch.fmax(rScore[0], other0) + rScore[1] = cute.arch.fmax(rScore[1], other1) + + if lane_id * 2 < MMA_N: + epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] + epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + head_id = lane_id // MAX_DQL + q_local_pos = lane_id - head_id * MAX_DQL + valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len + if lane_id < BLOCK_Q and valid_q: + final_score = epi_buffer[lane_id, 0] + for i in cutlass.range_constexpr(1, 4): + final_score = cute.arch.fmax(final_score, epi_buffer[lane_id, i]) + + t = batch_id * decode_query_len + q_local_pos + score[head_id, t, block_id] = final_score + + tma_stage = (tma_stage + 1) % self.num_stages + if tma_stage == 0: + tma_parity ^= 1 + else: + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + # Early tiles past the request's true length: let them land before the CTA exits. + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.mbarrier_wait(tma_full_mbar + 1 % num_stages, 0) + + +class IndexScoreRunner: + """Compiled ``IndexDecodeScoreKernel`` variants, keyed by (index dtype, index heads, decode query length, head dim, + split_k, early issue). Batch, query-token count, page count and block-table width stay symbolic, so one compile + serves every decode step of its key. ``prefetch_old`` is always on; ``early`` also caps the prefetch at one block + per CTA.""" + + kernel_cache: dict = {} + SUPPORTED_HEAD_DIM = 128 # capacity: fixed: the checkpoint's index head dim + SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K + # BLOCK_Q must fit one warp's worth of epilogue lanes. + MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile + EARLY_PREFETCH_DEPTH = 1 + + @classmethod + def is_supported( + cls, + *, + q_dtype: torch.dtype, + num_heads: int, + head_dim: int, + page_size: int, + max_decode_query_len: int, + ) -> bool: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + from tensorrt_llm._utils import is_sm_100f + + return ( + is_sm_100f() + and q_dtype in TORCH_TO_CUTE_DTYPE + and head_dim == cls.SUPPORTED_HEAD_DIM + and page_size == cls.SUPPORTED_PAGE_SIZE + and num_heads * max_decode_query_len <= cls.MAX_BLOCK_Q + and max_decode_query_len >= 1 + ) + + @classmethod + def _compile( + cls, + q_dtype: torch.dtype, + num_heads: int, + max_decode_query_len: int, + head_dim: int, + split_k: int, + early: bool, + ) -> None: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + + key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + if key in cls.kernel_cache: + return + cute_dtype = TORCH_TO_CUTE_DTYPE[q_dtype] + page_size = cls.SUPPORTED_PAGE_SIZE + sym_total_tokens = cute.sym_int() + sym_batch = cute.sym_int() + + # 16-element divisibility on every non-innermost stride is what the TMA descriptors for Q and K assume. + def _sym_stride(): + return cute.sym_int64(divisibility=16) + + q_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (sym_total_tokens, num_heads, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + # The index-K pool may be coalesced with the main K/V cache (per-page stride above page_size * head_dim). + k_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (cute.sym_int(), page_size, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + bt_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch, cute.sym_int()), stride_order=(1, 0) + ) + # The target passes a transposed view of the [heads, blocks, tokens] selector buffer. + score_fake = cute.runtime.make_fake_tensor( + cutlass.Float32, + (num_heads, sym_total_tokens, cute.sym_int()), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + sl_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch,), stride_order=(0,) + ) + fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = IndexDecodeScoreKernel( + cute_dtype, + num_heads, + max_decode_query_len, + split_k, + head_dim, + prefetch_old=True, + early_issue=early, + prefetch_depth=cls.EARLY_PREFETCH_DEPTH if early else 0, + ) + cls.kernel_cache[key] = cute.compile( + kernel, + q_fake, + k_fake, + bt_fake, + score_fake, + sl_fake, + fake_stream, + options="--enable-tvm-ffi", + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_mc_allreduce.py new file mode 100644 index 000000000000..d922b2aadbd2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_mc_allreduce.py @@ -0,0 +1,365 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-latency TP all-reduce + residual + RMSNorm for decode (M <= 8 tokens) over a multicast mapping. + + partial = bf16(x + x2) (optional second input, e.g. shared + routed experts) + updated = bf16(residual + sum_r partial_r) (ranks summed in FP32, rank order) + normed = bf16(updated * rsqrt(mean(updated^2) + eps) * w) + +One cluster of CLUSTER CTAs per token; CTA c owns columns [c * H / CLUSTER, ...), one 16-byte chunk (8 BF16) +per thread. Each CTA pushes its slice of this rank's partial through the multicast mapping (one store lands in +every rank's buffer), releases dependents, polls its own buffer until all ranks' slices arrived, sums them, adds +the residual and reduces the RMS statistic across the cluster through DSMEM. + +Exchange protocol as in the K3 sandwich: int32 words, 0x80000000 = empty (pushed -0.0 halves become +0.0), each +CTA keeps a call counter whose parity selects the buffer half, readers re-arm the words they read. All ranks of +the TP group must issue the same calls in the same stream order. + +The workspace holds two exchange sets. With prepushed 3 two producers pushed this rank's x (set A) and x2 (set B), +e.g. the routed experts' finalize and the shared expert's down projection, and the kernel forms bf16(x + x2) per +rank from the two slots, exactly as the unpushed path does. + +Producers that push (the o_proj / shared-expert GEMV ``ar`` tuple, the routed-expert engine's ``arp`` descriptor) are +pointed at a bucket's exchange by its ``mc`` address, ``flags_for(8)`` counters, ``rank``, ``world_size`` and +``max_tokens``; ``get_workspace`` and ``_bucket`` select it. + +The kernel is ``torch.ops.trtllm.minimax_m3_mc_allreduce_norm``, reached through the ``ops/comm`` entries +``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``; this module owns the workspaces, the bucket and +epilogue-option choice and the MXFP8 publication. +""" + +from dataclasses import dataclass +from typing import Any + +import torch + +from ..ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm as _allreduce_norm_op +from ..ops.comm.m3_mc_finalize_allreduce_norm import ( + m3_mc_finalize_allreduce_norm as _finalize_allreduce_norm_op, +) + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +BUCKETS = (8, 16, 48) # token capacity of each workspace (grid = bucket * cluster CTAs) +EMPTY_WORD = -2147483648 + +# Attention / MoE boundary epilogue options (each bit-identical): 1 residual + norm weight loaded at entry (before +# the poll); 2 cluster RMS exchange by st.async + mbarrier complete_tx instead of cluster.sync; 4 call parity rotated +# right after the poll (no trailing __syncthreads); 8 max-smem carveout so the next kernel's CTAs co-reside early. +# The 8-token bucket, by token count 1..8 (GB300 TP4 tuning). +AR_OPTS_BY_M = (7, 7, 15, 15, 15, 15, 7, 7) +# The 16- and 48-token buckets. +AR_OPTS_BY_BUCKET = {16: 15, 48: 15} +# The MoE-finalize all-reduce, per bucket. +AR_FIN_OPTS_BY_BUCKET = {8: 7, 16: 15, 48: 15} + + +# Option 2 is also the safe ordering of the cluster RMS exchange: its stores into peer CTAs' shared memory follow a +# cluster barrier arrive / wait pair. Without it they run before any cluster barrier, when a peer need not have +# started, so every call keeps it. +_OPT_CLUSTER_BARRIER = 2 + + +def _ar_opts(m: int, bucket: int, finalize: bool = False) -> int: + """Epilogue options for an m-row call on the given bucket; every option is bit-identical. A bucket missing from + the tables gets 0, and every result keeps _OPT_CLUSTER_BARRIER.""" + if finalize: + opts = AR_FIN_OPTS_BY_BUCKET.get(bucket, 0) + elif bucket == 8: + opts = AR_OPTS_BY_M[m - 1] + else: + opts = AR_OPTS_BY_BUCKET.get(bucket, 0) + return opts | _OPT_CLUSTER_BARRIER + + +@dataclass(eq=False) +class McAllReduceWorkspace: + """Two halves of [max_tokens][world][H] BF16 (int32 words) per rank behind one multicast mapping, plus one call + counter per CTA per cluster size. Create collectively and eagerly on every TP rank.""" + + uc: torch.Tensor + mc: torch.Tensor + flags: dict + rank: int + world_size: int + hidden: int + handle: Any + comm: Any + max_tokens: int = 8 + + @classmethod + def create(cls, mapping, hidden: int, max_tokens: int = 8) -> "McAllReduceWorkspace": + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_disabled + + comm = _get_mnnvl_workspace_comm(mapping) + words = 4 * max_tokens * mapping.tp_size * (hidden // 2) # two exchange sets of two halves + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + if not mpi_disabled(): + comm.Free() + raise RuntimeError( + "McAllReduceWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 4, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("McAllReduceWorkspace.create failed on another rank") + return cls( + uc=uc, + mc=mc, + flags={}, + rank=mapping.tp_rank, + world_size=mapping.tp_size, + hidden=hidden, + handle=handle, + comm=comm, + max_tokens=max_tokens, + ) + + def flags_for(self, cluster: int): + """This cluster size's per-CTA call counters, or None when they do not exist yet during a capture. + get_workspace creates the cluster-8 counters with the workspace, so for cluster 8 (the only size built) + that None never happens and cannot differ between ranks.""" + flags = self.flags.get(cluster) + if flags is None: + if torch.cuda.is_current_stream_capturing(): + return None + flags = torch.zeros(self.max_tokens * cluster, dtype=torch.int32, device=self.uc.device) + self.flags[cluster] = flags + return flags + + +class McBuckets: + """One workspace per token bucket (8 / 16 / 48): calls of different sizes never share exchange slots.""" + + def __init__(self, workspaces): + self.workspaces = workspaces + self.hidden = workspaces[0].hidden + + def for_tokens(self, m: int): + for ws in self.workspaces: + if m <= ws.max_tokens: + return ws + return None + + +def _bucket(workspace, m): + return workspace.for_tokens(m) if isinstance(workspace, McBuckets) else workspace + + +def decline_reason(x, residual, workspace, cluster=8, x2=None) -> str: + """Why ``m3_mc_allreduce_norm`` returned None for these operands (for the dispatch's miss message).""" + if workspace is None: + return "no multicast workspace: it is created collectively on the first eager decode step" + bucket = _bucket(workspace, x.shape[0]) + if bucket is None: + return f"no workspace bucket holds {x.shape[0]} tokens (buckets {BUCKETS})" + if bucket.flags.get(cluster) is None: + return f"no {cluster}-CTA call counters for the {bucket.max_tokens}-token bucket" + return ( + f"operands outside the kernel: x {tuple(x.shape)} {x.dtype} contiguous={x.is_contiguous()}, " + f"residual contiguous={residual.is_contiguous()}, hidden {bucket.hidden}" + + ("" if x2 is None else f", x2 {tuple(x2.shape)} contiguous={x2.is_contiguous()}") + ) + + +def m3_mc_allreduce_norm( + x, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + x2=None, + quant=False, + prepushed=0, +): + """Return (normed, updated_residual) or None when unsupported. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM. + + A None makes the caller take another collective, so it must be the same on every TP rank. It is: it depends only + on the shapes, dtypes and layouts of the operands (the same code builds them on every rank) and on whether the + workspace's flags exist, which ``get_workspace`` creates collectively, before any capture, for cluster 8.""" + workspace = _bucket(workspace, x.shape[0]) + if workspace is None: + return None + if ( + x.dim() != 2 + or x.shape[0] > MAX_TOKENS + or x.dtype != torch.bfloat16 + or x.shape[1] != workspace.hidden + or not x.is_contiguous() + or not residual.is_contiguous() + or (x2 is not None and (not x2.is_contiguous() or x2.shape != x.shape)) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + outs = _allreduce_norm_op( + x, + x2, + residual, + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and x.shape[1] % 128 == 0), + int(prepushed), + _ar_opts(x.shape[0], workspace.max_tokens), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +def m3_mc_finalize_allreduce_norm( + fc2, + perm, + scales, + shared, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + quant=False, +): + """MoE finalize (sum_k scale * fc2[perm] + shared) + all-reduce + residual + RMSNorm in one kernel. + Returns (normed, updated_residual) or None when unsupported. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM.""" + m = residual.shape[0] + workspace = _bucket(workspace, m) + if workspace is None: + return None + if ( + residual.dim() != 2 + or m > MAX_TOKENS + or residual.shape[1] != workspace.hidden + or fc2.dtype != torch.bfloat16 + or not fc2.is_contiguous() + or fc2.dim() != 2 + or perm.dtype != torch.int32 + or perm.dim() != 2 + or perm.shape[0] != m + or not perm.is_contiguous() + or scales.dtype not in (torch.float32, torch.bfloat16) + or not scales.is_contiguous() + or (shared is not None and (shared.shape != residual.shape or not shared.is_contiguous())) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + outs = _finalize_allreduce_norm_op( + fc2, + perm, + scales, + shared, + residual.contiguous(), + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and residual.shape[1] % 128 == 0), + _ar_opts(m, workspace.max_tokens, finalize=True), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +_WORKSPACE: dict = {} +_AGREE_COMMS: dict = {} + + +def _tp_comm(mapping): + """A communicator over mapping's TP group, created once (collective over that group only).""" + key = tuple(mapping.tp_group) + comm = _AGREE_COMMS.get(key) + if comm is None: + from tensorrt_llm._utils import mpi_comm, mpi_disabled + + if mpi_disabled(): + comm = mapping.tp_group_pg + else: + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + _AGREE_COMMS[key] = comm + return comm + + +def tp_agree(mapping, ok: bool) -> bool: + """Whether ``ok`` holds on every rank of mapping's TP group (one blocking host collective). + + For a fast path whose availability is decided from rank-local state (a kernel build that can fail on one rank): + every rank must take the same path, so that they issue the same collectives with the same exchange protocol and + the same kernels. Call it once per decision, at build time, at the same point on every rank, outside CUDA-graph + capture; never per step. With one rank, or attention DP (whose ranks run different batches), it returns ``ok``. + """ + if mapping.tp_size <= 1 or mapping.enable_attention_dp: + return ok + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("tp_agree must run outside CUDA-graph capture") + from tensorrt_llm._torch.distributed.ops import _mnnvl_workspace_all_succeeded + + return _mnnvl_workspace_all_succeeded(_tp_comm(mapping), bool(ok)) + + +def tp_agree_built(mapping, built, what: str): + """``built`` when it is not None on every TP rank, else None on every rank (see ``tp_agree``).""" + if tp_agree(mapping, built is not None) or built is None: + return built + from tensorrt_llm.logger import logger + + logger.warning(f"{what}: built here but not on another TP rank; every rank falls back") + return None + + +def get_workspace(mapping, hidden: int): + """Process-wide workspace (created collectively on first eager use); None while capturing before creation. + + The first eager call for a hidden size (the warmup forward, before any CUDA-graph capture) creates one + McAllReduceWorkspace per bucket in BUCKETS on every TP rank (MPI collective over the TP group) plus each + bucket's 8-CTA-cluster call counters; later calls, captured or not, return the same McBuckets. The workspaces + live until process exit. + """ + ws = _WORKSPACE.get(hidden) + if ws is None: + if torch.cuda.is_current_stream_capturing(): + return None + wss = [McAllReduceWorkspace.create(mapping, hidden, b) for b in BUCKETS] + for w in wss: + w.flags_for(8) + ws = McBuckets(wss) + _WORKSPACE[hidden] = ws + return ws diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/__init__.py new file mode 100644 index 000000000000..0aa75c0bf82c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/__init__.py @@ -0,0 +1,16 @@ +# 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. +"""MiniMax-M3 MoE layers for low-batch decode as CuTe DSL kernels: ``m3_moe`` for 1..8 tokens (op.py) and ``m3_moe16`` +for 9..16 tokens (op16.py).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/m3_moe16_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/m3_moe16_kernel.py new file mode 100644 index 000000000000..ff609e78ef29 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/m3_moe16_kernel.py @@ -0,0 +1,3871 @@ +# 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. +"""MiniMax-M3 routed + shared experts for 9..16 decode tokens on one TP rank: ``m3_moe16`` (v2). + +Same math as the 1..8-token engine (m3_moe_kernel.py), operation for operation for the routed part (bitwise equal to +TRTLLM-Gen's fused MoE on the fp32 router-GEMV logits); a different schedule for 9..16 rows: + +- Front: the hidden states land in smem by bulk copies in K chunks (the idle ring); all 16 warps run the router + GEMV (m3_router_gemv's per-thread order) and the NVFP4 quantization of B1 from smem; the logits cross CTAs as + tagged words; top-k one warp per token. +- Ring: FC1 / FC2 weight stages by TMA; their scale words go by 4-byte cp.async straight into the instruction atom + layout (SFA_CP; no raw scale box, no splicer warp), completing on the stage's full barrier. +- FC1 -> FC2 hand-off (B2IMG): the FC1 epilogue writes each unit's FP4 values and scale byte into a per-group global + image laid out as the FC2 B tile + its scale atoms; one counter per group (release add per finished unit, every + counter gains exactly 192 per call); warp 5 loads a completed image with two bulk copies. +- Shared expert on NSH more CTAs (MXFP8 tcgen05), as the 1..8 engine. +Configuration via M3_CONFIG (m_max, num_ctas, stages, nb2_cap, lag, shared, sfa_cp, b2buf, ...). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +_CFG = globals().get("M3_CONFIG") or {} + + +def _cfg(key: str, default): + return type(default)(_CFG.get(key, default)) + + +H = 6144 # capacity: fixed: the checkpoint's hidden size +I_TP = 768 +TWO_I = 2 * I_TP +E_LOCAL = 128 +TOP_K = 4 # capacity: fixed: the checkpoint's top-k +ROUTE_SCALE = 2.0 +SUM_EPS = 1e-20 +M_MAX = _cfg("m_max", 12) +assert 9 <= M_MAX <= 16, M_MAX +NCTA = _cfg("num_ctas", 128) +FUSED_ROUTER = True +NROUTE = E_LOCAL +NGEMV = (NROUTE + NCTA - 1) // NCTA +assert NGEMV == 1, "m3_moe16: one router expert per CTA (num_ctas >= 128)" +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +LAG = _cfg("lag", 3) +G_CAP = TOP_K * M_MAX +NG_CAP = M_MAX +THREADS = 512 # capacity: fixed: the CTA size +NWARPS = THREADS // 32 +GEMV_THREADS = 256 +SHARED = bool(_cfg("shared", 0)) +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +SFA_CP = bool( + _cfg("sfa_cp", 1) +) # A scales by 4-byte cp.async into the atom layout (else raw box + splicer warp) +B2IMG = True +PERM = bool(_cfg("perm", 1)) +PERM_MIN_NG = max( + 3, _cfg("perm_min_ng", 5) +) # (measured: helps at random routing G >= 23, hurts at spec G 13-14) +# a partial last slot group becomes group 1 (the plan must match: op16) +DBG_DELAY = bool( + _cfg("dbg_delay", 0) +) # (stress builds) pseudo-random producer / epilogue / release stalls +DBG_NOSTORE = bool(_cfg("dbg_nostore", 0)) # (timing experiment only: wrong results) +EPI_FENCE = bool( + _cfg("epi_fence", 0) +) # (rel_warp 0) writer lanes' fence.acq_rel.gpu before the release add +EPI_PFENCE = bool( + _cfg("epi_pfence", 1) +) # producer-side fence.proxy.async.global before the release add +REL_WARP = bool( + _cfg("rel_warp", 1) +) # warp 8 does the per-group release adds (off the epilogue's path) +NSH = 24 if SHARED else 0 +SH_I = 768 +SH_KB = H // 32 # capacity: fixed: a kernel tile constant +SH_DKB = SH_I // 32 # capacity: fixed: a kernel tile constant +SH_SPLIT = 2 +SH_GT = SH_KB // 4 // SH_SPLIT # capacity: fixed: a kernel tile constant +SH_STAGES_GU = SH_GT +SH_STAGES = SH_GT + 2 * (SH_I // 128) # capacity: fixed: pipeline stages +SHX_BAR = 6 +SH_PRETILE = False +SH_STAGE_BYTES = 16384 + 512 +MMA_M = 128 +KT_BYTES = 128 +INC_K = 2 +NUM_KB = 4 +# FC1 +N1 = _cfg("n1", 48 if _cfg("m_max", 12) <= 12 else 64) # B rows QS q + t +N1_LD = 1 << (N1 - 1).bit_length() # tcgen05.ld .num (a power of two) covering N1 columns +QS = N1 // 4 +assert QS >= M_MAX +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 +K1_TILES = H // 256 +S1 = K1_TILES // C1 # capacity: fixed: pipeline stages +NUM_BYTES_A = MMA_M * KT_BYTES # capacity: fixed: a stage's operand bytes +NUM_BYTES_BOX1 = R1 * KT_BYTES +SF1_KA = H // 64 +SF1_RB = TWO_I // 128 +NUM_BYTES_SFA_RAW = 16 * 512 +NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 +NUM_BYTES_SFA = NUM_KB * 512 +NUM_WORDS_SFA = NUM_BYTES_SFA // 4 +NUM_BYTES_B1 = N1 * KT_BYTES +NUM_SFB1_ATOMS = S1 * NUM_KB +# FC2 +R2 = 32 +NB2 = H // R2 +KT2 = I_TP // 256 +SF2_KA = I_TP // 64 +SF2_RB = H // 128 +N2 = _cfg("n2", (4 * M_MAX + 15) // 16 * 16) # B rows j M + t +NUM_BYTES_B2 = N2 * KT_BYTES +IMG_B2 = KT2 * NUM_BYTES_B2 # one group's B2 tile image (global and smem) +IMG_SF = KT2 * NUM_KB * 512 # its scale atoms +IMG_BYTES = IMG_B2 + IMG_SF +CNT_PER_CALL = 4 * (TWO_I // R1) # every group counter gains 192 per call +B1IMG_OFF = None # set below (after N1): the activation image after the NG_CAP group images +TK_WARPS = [0, 1, 2, 3, 6, 8, 9, 10, 11, 12, 13, 14, 5, 7, 15, 4][ + : _cfg("m_max", 12) +] # top-k warp of token t (4 last: early0) +B1_WARP = 15 if _cfg("m_max", 12) <= 14 else -1 # (-1: warp 0 after the routing) +NB2_CAP = _cfg("nb2_cap", 2) +B1IMG_OFF = M_MAX * IMG_BYTES +B1IMG_BYTES = (S1 * N1 * 128 + M_MAX * (H // 16) + 15) // 16 * 16 +IMG_TOTAL = B1IMG_OFF + B1IMG_BYTES +UNITS_ROW = TWO_I // R1 +PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +TOPK_BAR = 4 +ALL_BAR = 5 +POLL_NS = _cfg("poll_ns", 150) +B2_POLL_NS = _cfg("b2_poll_ns", 16) +LG_REPL = _cfg("lg_repl", 64) +LG_STRIDE = M_MAX * E_LOCAL + 32 +W_WARP = 4 # capacity: fixed: a warp role index +X_WARP = 5 # B2 image loader; slots + plan # capacity: fixed: a warp role index +S_WARP = 6 # (SFA_CP = 0) scale splicer # capacity: fixed: a warp role index +MMA_WARP = 7 # capacity: fixed: a warp role index +EARLY0 = bool(_cfg("early0", 1)) +EVICT_FIRST = 0x12F0000000000000 +EVICT_LAST = 0x14F0000000000000 +_LOG2E = 1.4426950408889634 +NB2BUF = _cfg("b2buf", 2) +SPB = 1 +# x staging: K chunks of XCK elements, NXB smem buffers (rotating if fewer than the chunks) +XCK = _cfg("xck", 2048) +NXC = H // XCK +XCB = M_MAX * XCK * 2 # bytes of one chunk (all tokens) +STAGE_A = NUM_BYTES_A +STAGE_SF = NUM_BYTES_SFA if SFA_CP else NUM_BYTES_SFA_RAW + NUM_BYTES_SFA +_FIXED_SMEM = ( + S1 * NUM_BYTES_B1 + + NB2BUF * IMG_B2 + + IMG_SF + + NB2_CAP * M_MAX * TOP_K * R2 * 2 + + SPB * M_MAX * C1 * R1 * 4 + + (G_CAP + 1) * (NG_CAP + 1) * 4 + + 6 * 1024 +) +_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // (STAGE_A + STAGE_SF) +ACC1_COL = 0 +ACC2_COL = 2 * N1 +ACC2_STRIDE = 1 << (N2 - 1).bit_length() +SFB1_COL = ACC2_COL + 2 * ACC2_STRIDE +SFB1_WSTEP = 1 +SFB2_COL = ( + SFB1_COL + NUM_SFB1_ATOMS * 4 +) # one TMEM slot (loaded after the previous group's FC2 MMAs) +SFA1_COL = SFB2_COL + KT2 * NUM_KB * 4 +STAGES = min(_cfg("stages", 8), _STAGES_FIT, (512 - SFA1_COL) // 16) +TMEM_COLS = 512 +assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) +# smem: [ring A | ring SF | B2 bufs | B2 SF staging | B1]; x chunks alias [0, _OFF_B1) during the front +_OFF_SFA = STAGES * NUM_BYTES_A +_OFF_B2 = (_OFF_SFA + STAGES * STAGE_SF + 1023) // 1024 * 1024 +_OFF_SFB2 = _OFF_B2 + NB2BUF * IMG_B2 +_OFF_B1 = (_OFF_SFB2 + IMG_SF + 1023) // 1024 * 1024 +BIG = _OFF_B1 + S1 * NUM_BYTES_B1 +NXB = min(NXC, _OFF_B1 // XCB) +assert NXB >= 2, (NXB, _OFF_B1, XCB) +# shared CTA layout in BIG: one pass over all M tokens; B rows SHN q + t (gate_up, 2 K slices), t (down) +SHN = 16 # capacity: fixed: a kernel tile constant +SH_NPASS = 1 +SH_HINT = EVICT_FIRST +SH_PF_POST = bool( + _cfg("sh_pf_post", 1) +) # L2 prefetch of the shared weights after the grid dependency (else before) +_SB_A = 0 +_SH_FIX = ( + SH_GT * 2 * SHN * 128 + + 6 * SHN * 128 + + M_MAX * SH_KB + + M_MAX * SH_DKB + + 64 * SHN * 4 + + 2 * 32 * SHN * 4 + + 2048 +) +SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 +_SB_RAW = SH_NST * NUM_BYTES_A +_SB_SPL = _SB_RAW + SH_NST * 2048 +_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 +_SB_BH = _SB_BX + SH_GT * 2 * SHN * 128 +_SB_XC = _SB_BH + 6 * SHN * 128 +_SB_HC = _SB_XC + M_MAX * SH_KB +_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 +_SB_GU = _SB_XCH + 64 * SHN * 4 +_SB_END = _SB_GU + 2 * 32 * SHN * 4 +SH_ACC_GU = 0 +SH_ACC_DN = 2 * SHN +SH_SFA = SH_ACC_DN + 2 * SHN +SH_SFB_GU = SH_SFA + 4 * 8 +SH_SFB_DN = SH_SFB_GU + 4 * SH_GT # capacity: fixed: a kernel tile constant +SH_TMEM = 256 # capacity: fixed: TMEM columns +if SHARED: + assert SH_STAGES - SH_NST <= 64, SH_STAGES + assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 24 <= SH_TMEM, (SH_NST, _SB_END, BIG) +TD_STRIDE = 64 +PFD = _cfg("pfd", 0) # L2 prefetch distance in ring stages (0: off) +PF_WARP = 12 # capacity: fixed: a warp role index +TICKET = bool(_cfg("ticket", 0)) # CTA roles by entry order +# priority tokens: token 0 and PT1 (M / 2: the second request's first row at spec c2) get their logits published first +NPT = _cfg("npt", 0) +PT1 = M_MAX // 2 +PF_PRIO = bool( + _cfg("pf_prio", 0) +) # L2 prefetch of the priority tokens' experts (this CTA's 1/NCTA of each) +PF_W1 = TWO_I * (H // 2) // NCTA +PF_S1 = TWO_I * (H // 16) // NCTA +PF_W2 = H * (I_TP // 2) // NCTA +PF_S2 = H * (I_TP // 16) // NCTA +assert PF_W1 % 16 == 0 and PF_S1 % 16 == 0 and PF_W2 % 16 == 0 and PF_S2 % 16 == 0 + + +# ============================================================================= +# Inline-PTX helpers. +# ============================================================================= +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- all-reduce push (the MoE boundary's multicast exchange, set A): +# arp = int64 [enable, mc, flags, rank, world, maxt] +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): + """st.global.b32 when pred != 0.""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _ar_push(arp, v, low, t, cl): + """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on + lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call + parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + world = arp.load(idx=4) + parity = cutlass.Int64( + _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) + ) + & cutlass.Int32(1) + ) + off = ( + parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" + _llvm.inline_asm( + None, + [], + f"bar.arrive {bar_id}, {threads};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "st.relaxed.gpu.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int64(val)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _nanosleep(ns: int, *, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _gtime(*, loc=None, ip=None): + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): + if cutlass.const_expr(TIMING): + tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _sub_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_bits(x, *, loc=None, ip=None): + """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3_to_f32(b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(b)], + "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" + "cvt.f32.f16 $0, l;\n}", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x2(hi, lo, *, loc=None, ip=None): + """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """x rounded to bf16 (nearest even), as f32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(v, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(v)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16 (nearest even) bits of x in the low 16 bits.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ue8m0_rp(x, *, loc=None, ip=None): + """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3x2(hi, lo, *, loc=None, ip=None): + """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _min_xorsign_abs(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_full(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2_full(x, *, loc=None, ip=None): + """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): + """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" + _llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _mx_osc(amax): + """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" + code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) + cf = cutlass.Int32( + cutlass.select_( + code == cutlass.Int32(255), + cutlass.Int32(0x7FFFFFFF), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) + ), + ) + ).bitcast(cutlass.Float32) + return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) + + +def _mx_words(v, osc): + """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" + words = [] + for t in range(len(v) // 4): + lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) + hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) + words.append(lo | (hi << cutlass.Int32(16))) + return words + + +def _twiddle(score): + """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" + bits = ( + score.bitcast(cutlass.Int32) + if hasattr(score, "bitcast") + else cutlass.Float32(score).bitcast(cutlass.Int32) + ) + return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) + + +@cute.jit +def _wait(bar: cutlass.Array, parity): + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +def _tptr(base_row, col, dtype=cutlass.Int32): + return cutlass.inttoptr((base_row << 16) | col, 6, dtype) + + +@dsl_user_op +def _match_any(value, *, loc=None, ip=None): + """The mask of the warp's lanes holding the same 32-bit value.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(value)], + "match.any.sync.b32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@cute.jit +def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): + """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" + if cutlass.const_expr(HANG_DBG): + n = cutlass.Int32(0) + while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( + n < cutlass.Int32(1 << 22) + ): + n = n + cutlass.Int32(1) + if n >= cutlass.Int32(1 << 22): + bx, _, _ = cute.arch.block_idx() + w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tdbg.store( + cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w + ) + else: + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +@cute.jit +def _adv(stage, phase): + """The next ring position's (stage, phase).""" + stage = stage + cutlass.Int32(1) + wrap = stage == cutlass.Int32(STAGES) + return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( + wrap, phase ^ cutlass.Int32(1), phase + ) + + +@cute.jit +def _seg(i0, i1, k): + """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: + the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" + uf = i0 // cutlass.Int32(S1) + ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) + unit = uf + k - cutlass.Int32(1) + unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) + unit = cutlass.select_(ul == uf, uf, unit) + j0 = i0 - unit * cutlass.Int32(S1) + j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) + j1 = i1 - unit * cutlass.Int32(S1) + j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) + return unit, j0, j1 + + +@cute.jit +def _nsegs(i0, i1): + n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) + return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) + + +@cute.jit +def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): + """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile + 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" + # B rows n >= 32 (N1 = 64) go to the atom's next 32-bit column (row n: lane n % 32, column n // 32) + for w in cutlass.range_constexpr((N1 + 31) // 32): + n = lane + cutlass.Int32(32 * w) + q = n // cutlass.Int32(QS) + tl = n - q * cutlass.Int32(QS) + sf_ok = (n < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) + sfv = [] + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + j = a // NUM_KB + kb = a % NUM_KB + wsf = src.load( + idx=cutlass.select_( + sf_ok, + tl * cutlass.Int32(H // 64) + + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) + + cutlass.Int32(kb), + cutlass.Int32(0), + ) + ) + sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + prims.tcgen05_st( + "32x32b", + _tptr( + base_row + warp * cutlass.Int32(32), + base_col + SFB1_COL + a * 4 + w * SFB1_WSTEP, + ), + sfv[a], + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(b1_ready) + + +@cute.jit +def _b1_addr(t, kb16): + """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" + kt = kb16 >> cutlass.Int32(4) + j = kt >> cutlass.Int32(2) + n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t + cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) + return ( + j * cutlass.Int32(NUM_BYTES_B1) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) + ) // cutlass.Int32(4) + + +@cute.jit +def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): + """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, + gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 + and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" + neg_max = cutlass.Float32(-3.40282346638528859812e38) + gc = cutlass.select_(gate_raw > lim, lim, gate_raw) + gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) + uc = cutlass.select_(xv > lim, lim, xv) + uc = cutlass.select_(xv < -lim, -lim, uc) + c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) + sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) + gs = _mul_rn(gc, sgm) + h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) + h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) + amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) + for off in (16, 8, 4, 2, 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) + sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) + sff = _e4m3_to_f32(sfb) + osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) + return _mul_rn(_mul_rn(h, osc), c), sfb + + +@cute.jit +def _plan_range(s_plan: cutlass.Array, p): + w = s_plan.load(idx=p) + return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) + + +@cute.jit +def _sh_issue( + lane, + sid, + g, + s_big: cutlass.Array, + sh_full: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + gpos, +): + """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both + slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its + raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile + (g - SH_GT) % 6.""" + st = gpos % cutlass.Int32(SH_NST) # ring position (stage g of this pass) + bar = sh_full.subview(st) + if cutlass.const_expr(SH_PRETILE): + if lane == 0: + src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) + prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + gusf.subview(src + cutlass.Int32(16384)), + bar, + 512, + ) + cute.arch.sync_warp() + else: + if g < cutlass.Int32(SH_GT): + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) + cute.arch.sync_warp() + if lane < 4: + q = lane >> cutlass.Int32(1) + u = lane & cutlass.Int32(1) + row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) + kt = q * cutlass.Int32(SH_GT) + g + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), + tma_gu, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=SH_HINT, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + cutlass.Int32(_SB_RAW) + + st * cutlass.Int32(2048) + + lane * cutlass.Int32(512) + ), + gusf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), + tma_dn, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=SH_HINT, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), + dnsf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + cute.arch.sync_warp() + + +@cute.jit +def _shared_cta( + tidx, + warp, + lane, + sid, + s_big: cutlass.Array, + sh_bar: cutlass.Array, + epochs: cutlass.Array, + tdbg: cutlass.Array, + bx, + xb32: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + tmem_ptr_i32: cutlass.Array, + gu8: cutlass.Array, + dn8: cutlass.Array, +): + """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one + 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles + of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" + big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) + bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) + big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) + sh_full = sh_bar # [0, SH_NST): stage full + sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) + b_gu_rdy = sh_bar.subview(16) + b_dn_rdy = sh_bar.subview(17) + acc_gu = sh_bar.subview(18) + acc_dn = sh_bar.subview(19) + if warp == 0: + prims.prefetch_tensormap(tma_gu) + prims.prefetch_tensormap(tma_dn) + if lane < SH_NST: + prims.mbarrier_init(sh_full.subview(lane), 1) + prims.mbarrier_init( + sh_empty.subview(lane), 33 + ) # the MMAs' commit + each splicer lane's reads done + if lane == 0: + prims.mbarrier_init(b_gu_rdy, 1) + prims.mbarrier_init(b_dn_rdy, 1) + prims.mbarrier_init(acc_gu, 1) + prims.mbarrier_init(acc_dn, 1) + if warp == 1: + prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency + _sh_issue( + lane, + sid, + cutlass.Int32(g), + s_big, + sh_full, + tma_gu, + tma_dn, + gusf, + dnsf, + cutlass.Int32(g), + ) + # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) + if cutlass.const_expr(SH_PRETILE): + # this CTA's stages beyond the ring, 4 per lane + for r in cutlass.range_constexpr(2): + pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(2) + cutlass.Int32(r) + if pg < cutlass.Int32(SH_STAGES): + _l2_prefetch( + gusf.subview( + (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) + ) + .data_ptr() + .toint(), + SH_STAGE_BYTES, + ) + elif cutlass.const_expr(not SH_PF_POST): + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(SH_PF_POST): + if warp == 0: # after the grid dependency: DRAM is idle during the routed front + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + NQ = THREADS - 64 # activation quant threads (warps 2 ..) + for ps in cutlass.range_constexpr(SH_NPASS): + TT0 = 0 # one pass: B rows SHN q + t (gate_up), t (down) + MS = M_MAX + GB = ps * SH_STAGES # ring positions before this pass + PAR = ps & 1 # parity of the per-pass barriers (B ready, accumulators) + if warp == 0: + for g in cutlass.range(SH_NST if ps == 0 else 0, SH_STAGES, unroll=1): + gg = g + cutlass.Int32(GB) # ring position over the passes + _wait( + sh_empty.subview(gg % cutlass.Int32(SH_NST)), + ((gg // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), + ) + _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf, gg) + elif warp == 1: + # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down + idesc_gu = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=2 * SHN, + m_dim=128, + ) + idesc_dn = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=SHN, + m_dim=128, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + desc_a0 = prims.Tcgen05SmemDesc.build( + s_big.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bx = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BX), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bh = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BH), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + sq = sid & cutlass.Int32(3) + _wait(b_gu_rdy, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for g in cutlass.range(SH_STAGES, unroll=1): + gg = g + cutlass.Int32(GB) + st = gg % cutlass.Int32(SH_NST) + if g == cutlass.Int32(SH_GT): + _wait(b_dn_rdy, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _wait(sh_full.subview(st), (gg // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) + raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) + if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): + # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 + w4 = [ + big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) + for qq in range(4) + ] + big32.store( + ( + cutlass.Int32(w4[0]), + cutlass.Int32(w4[1]), + cutlass.Int32(w4[2]), + cutlass.Int32(w4[3]), + ), + idx=cutlass.Int32(_SB_SPL // 4) + + st * cutlass.Int32(128) + + lane * cutlass.Int32(4), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + prims.mbarrier_arrive( + sh_empty.subview(st) + ) # this lane's generic reads of the slot are done + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + if g < cutlass.Int32(SH_GT): + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bx + cutlass.Int64(g * cutlass.Int32((2 * SHN * 128) >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), + da + 2 * kb, + db + 2 * kb, + idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_GT - 1): + prims.tcgen05_commit(acc_gu) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview( + cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) + + st * cutlass.Int32(512 if SH_PRETILE else 2048) + ), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bh + cutlass.Int64(kt * cutlass.Int32((SHN * 128) >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr( + tcol + SH_ACC_DN + jt * cutlass.Int32(SHN), 6, cutlass.Float32 + ), + da + 2 * kb, + db + 2 * kb, + idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_STAGES - 1): + prims.tcgen05_commit(acc_dn) + cute.arch.sync_warp() + if warp >= 2: + if cutlass.const_expr(ps > 0): + _wait( + acc_dn, PAR ^ 1 + ) # the previous pass's MMAs are done: its B tiles and TMEM scales are free + # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = + # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes + qt = tidx - cutlass.Int32(64) + for r in cutlass.range_constexpr((MS * SH_KB + NQ - 1) // NQ): + q_ = qt + cutlass.Int32(r * NQ) + if q_ < cutlass.Int32(MS * SH_KB): + t = q_ // cutlass.Int32(SH_KB) + b = q_ - t * cutlass.Int32(SH_KB) + wv = [ + xb32.load( + idx=(t + cutlass.Int32(TT0)) * cutlass.Int32(H // 2) + + b * cutlass.Int32(16) + + cutlass.Int32(4 * z), + vector_size=4, + alignment=16, + ) + for z in range(4) + ] + v = [] + for z in cutlass.range_constexpr(4): + for zz in cutlass.range_constexpr(4): + w = cutlass.Int32(wv[z][zz]) + v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) + v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(v[z])) + code, osc = _mx_osc(amax) + words = _mx_words(v, osc) + kt = b >> cutlass.Int32(2) + sl = kt // cutlass.Int32(SH_GT) + it = kt - sl * cutlass.Int32(SH_GT) + n = sl * cutlass.Int32(SHN) + t + c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BX) + + it * cutlass.Int32(2 * SHN * 128) + + n * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) + if warp >= 4: + if warp < 8: + qd = warp - cutlass.Int32(4) + # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of + # k-tile SH_GT q + i + n = lane + t = n & cutlass.Int32(SHN - 1) + ok = (n < cutlass.Int32(2 * SHN)) & (t < cutlass.Int32(MS)) + for i in cutlass.range_constexpr(SH_GT): + kt = (n // cutlass.Int32(SHN)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) + wsf = big32.load( + idx=( + cutlass.Int32(_SB_XC) + + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) + + kt * cutlass.Int32(4) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), + cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + et = tidx - cutlass.Int32(128) + if et == 0: + prims.mbarrier_arrive(b_gu_rdy) + _stamp(tdbg, bx, 54) + # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t + # at column 8 q + t + _wait(acc_gu, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=2 * SHN, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(2): + for tt in cutlass.range_constexpr(MS): + bigf.store( + cutlass.Float32(av[SHN + tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) + * cutlass.Int32(SHN) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd < cutlass.Int32(2): + # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 + for tt in cutlass.range_constexpr(MS): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(SHN) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(SHN) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag + for r in cutlass.range_constexpr((32 * MS + 127) // 128): + pq = et + cutlass.Int32(r * 128) + if pq < cutlass.Int32(32 * MS): + p = pq // cutlass.Int32(MS) + tt = pq - p * cutlass.Int32(MS) + ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(SHN) + tt) + up = bigf.load( + idx=cutlass.Int32(_SB_GU // 4) + + (cutlass.Int32(32) + p) * cutlass.Int32(SHN) + + tt + ) + q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) + q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) + q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) + q_e2 = _ex2_full( + _mul_rn( + _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), + cutlass.Float32(1.4426950216293334961), + ) + ) + q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) + q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) + _st_relaxed_u32( + hx.subview( + (tt + cutlass.Int32(TT0)) * cutlass.Int32(SH_I) + + sid * cutlass.Int32(32) + + p + ) + .data_ptr() + .toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), + ) + if et == 0: + _stamp(tdbg, bx, 50) + # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes + for rr in cutlass.range_constexpr((MS * SH_DKB + 127) // 128): + q_ = et + cutlass.Int32(rr * 128) + if q_ < cutlass.Int32(MS * SH_DKB): + t = q_ // cutlass.Int32(SH_DKB) + kb = q_ - t * cutlass.Int32(SH_DKB) + base = (t + cutlass.Int32(TT0)) * cutlass.Int32(SH_I) + kb * cutlass.Int32( + 32 + ) + hv = [cutlass.Int32(0) for _ in range(32)] + pend = cutlass.Int32(1) + while pend != cutlass.Int32(0): + for q4 in cutlass.range_constexpr(8): + v4 = prims.load_ext( + hx.subview(base + cutlass.Int32(4 * q4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for z in cutlass.range_constexpr(4): + hv[4 * q4 + z] = cutlass.Int32(v4[z]) + okw = cutlass.Boolean(True) + for z in cutlass.range_constexpr(32): + okw = okw & ( + ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 + ) + pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) + vals = [ + (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) + ] + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) + code, osc = _mx_osc(amax) + words = _mx_words(vals, osc) + kt = kb >> cutlass.Int32(2) + c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BH) + + kt * cutlass.Int32(SHN * 128) + + t * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (t & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + ( + ((c0 + cutlass.Int32(1)) ^ (t & cutlass.Int32(7))) + << cutlass.Int32(2) + ), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + okh = lane < cutlass.Int32(MS) + for kt in cutlass.range_constexpr(6): + wsf = big32.load( + idx=( + cutlass.Int32(_SB_HC) + + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) + + cutlass.Int32(4 * kt) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), + cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if et == 0: + prims.mbarrier_arrive(b_dn_rdy) + _stamp(tdbg, bx, 51) + # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l + _wait(acc_dn, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jt in cutlass.range_constexpr(2): + dv = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_DN + jt * SHN, + cutlass.Float32, + ), + num=SHN, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + rowo = ( + sid * cutlass.Int32(256) + + cutlass.Int32(128 * jt) + + qd * cutlass.Int32(32) + + lane + ) + for tt in cutlass.range_constexpr(MS): + _st_relaxed_u32( + dnx.subview(cutlass.Int32((tt + TT0) * H) + rowo).data_ptr().toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), + ) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if et == 0: + _stamp(tdbg, bx, 52) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if warp == 1: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) + if tidx == 0: + epochs.store(ep + cutlass.Int32(1), idx=bx) + + +# ============================================================================= +# v2 helpers: global release / acquire, proxy fences. +# ============================================================================= +@dsl_user_op +def _red_release_add(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _atom_add_u64(addr, val, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr), cutlass.Int64(val)], + "atom.relaxed.gpu.global.add.u64 $0, [$1], $2;", + "=l,l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + "fence.proxy.async.global;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _fence_acq_rel_gpu(*, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + "fence.acq_rel.gpu;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _dbg_delay(salt): + """(DBG_DELAY stress builds) a pseudo-random ~0..4 us stall: protocol robustness under skew.""" + if cutlass.const_expr(DBG_DELAY): + r = cutlass.Int32(_gtime() >> cutlass.Int64(6)) * cutlass.Int32(1103515245) + salt + if ((r >> cutlass.Int32(12)) & cutlass.Int32(7)) == cutlass.Int32(0): + for _ in cutlass.range((r >> cutlass.Int32(16)) & cutlass.Int32(31), unroll=1): + _nanosleep(128) + + +@cute.jit +def _waitt(bar: cutlass.Array, parity, acc): + """_wait that (TIMING) adds the waited ns to acc.""" + t0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + if cutlass.const_expr(TIMING): + acc = acc + (_gtime() - t0) + return acc + + +@cute.jit +def _tstore(tdbg: cutlass.Array, bx, i: cutlass.Constexpr, v): + if cutlass.const_expr(TIMING): + tdbg.store(cutlass.Int64(v), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@cute.jit +def _issue_fc1( + lane, + stage, + e, + c, + j, + sA: cutlass.Array, + sSF: cutlass.Array, + ab_full: cutlass.Array, + tma_a1, + tma_s1, + s1w: cutlass.Array, +): + """FC1 stage j of unit c of expert e into ring stage `stage` (the caller waited its empty barrier): 4 TMA boxes + (rows 32 c .., k-tiles 4 j + q) and the stage's scale atoms (SFA_CP: word [r][q] of atom kb <- raw atom + (k-tile 4 j + q, kb) word [r][c % 4], by this warp's 32 lanes).""" + row = c * cutlass.Int32(R1) + kx = j * cutlass.Int32(C1 * KT_BYTES) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview(stage * cutlass.Int32(NUM_BYTES_A) + cutlass.Int32(q * NUM_BYTES_BOX1)), + tma_a1, + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) * cutlass.Int32(SF1_KA) + + j * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + prims.cp_async_bulk_tensor_shared_cta_global( + sSF.subview(stage * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1, + (cutlass.Int32(0), sf_line), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + atom0 = (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) * cutlass.Int32( + SF1_KA + ) + j * cutlass.Int32(16) + src0 = atom0 * cutlass.Int32(128) + lane * cutlass.Int32(4) + (c & cutlass.Int32(3)) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for q in cutlass.range_constexpr(C1): + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSF.subview(dst0 + cutlass.Int32(kb * 128 + q)), + s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + # one group per ring stage (racecheck: the issuer waits it before the slot's reuse) + prims.cp_async_commit_group() + cute.arch.sync_warp() + + +@cute.jit +def _issue_fc2( + lane, + stage, + g, + blk, + t2: cutlass.Constexpr, + n_here, + s_el: cutlass.Array, + sA: cutlass.Array, + sSF: cutlass.Array, + ab_full: cutlass.Array, + tma_a2, + tma_s2, + s2w: cutlass.Array, +): + """FC2 k-tile t2 of output block blk for the group's n_here slots: box jj = rows 32 blk .. of slot jj's w2 at + k-tile t2; scale word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) word [r][blk % 4].""" + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + n_here + * cutlass.Int32(NUM_BYTES_BOX1 if SFA_CP else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4), + ) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + cutlass.Int32(jj * NUM_BYTES_BOX1) + ), + tma_a2, + (cutlass.Int32(t2 * KT_BYTES), blk * cutlass.Int32(R2), e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSF.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) + ), + tma_s2, + ( + cutlass.Int32(0), + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(4), + ), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + src0 = ( + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (blk & cutlass.Int32(3)) + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSF.subview(dst0 + cutlass.Int32(kb * 128 + jj)), + s2w.subview(src0 + cutlass.Int32(kb * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + # one group per ring stage (racecheck: the issuer waits it before the slot's reuse) + prims.cp_async_commit_group() + cute.arch.sync_warp() + + +@cute.jit +def _quant_words(w, sc, rsc, rsix): + """NVFP4 of one 16-value block given as 8 bf16-pair words: (lo word, hi word, scale e4m3 bits).""" + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), cutlass.Float32(0.0) + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + return lo, hi, sfb + + +@cute.jit +def _quant_block(xs32: cutlass.Array, widx, sc, rsc, rsix): + """NVFP4 of one 16-value block (8 bf16 words at widx): (lo word, hi word, scale e4m3 bits), fp4_quantize's ops.""" + w = [] + for h in cutlass.range_constexpr(2): + v4 = xs32.load(idx=widx + cutlass.Int32(4 * h), vector_size=4, alignment=16) + for z in cutlass.range_constexpr(4): + w.append(cutlass.Int32(v4[z])) + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), cutlass.Float32(0.0) + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + return lo, hi, sfb + + +@cute.kernel +def m3_moe_kernel( + tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_gu_desc: cutlass.GridConstant[cuda.TensorMap], + tma_dn_desc: cutlass.GridConstant[cuda.TensorMap], + bias: cutlass.Array, # f32 [128] + escal: cutlass.Array, # f32 [6, 128] + plan: cutlass.Array, # int32 [4 M + 1, NCTA, M + 1] + partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums + img32: cutlass.Array, # int32 [NG_CAP, IMG_BYTES / 4]: per-group FC2 B image (tile + scale atoms) + cnt: cutlass.Array, # int32 [NG_CAP * 8]: per-group unit counters (index 8 g) + s1w: cutlass.Array, # int32 words of w3_w1_weight_scale + s2w: cutlass.Array, # int32 words of w2_weight_scale + w1b: cutlass.Array, # int8 bytes of w3_w1_weight (L2 prefetch) + w2b: cutlass.Array, # int8 bytes of w2_weight (L2 prefetch) + epochs: cutlass.Array, # int32 [NCTA + NSH] + out: cutlass.Array, # bf16 [M, 6144] + route_dbg: cutlass.Array, # int32 [M, 8] + tdbg: cutlass.Array, # int64 [NCTA + NSH, 64] + xb32: cutlass.Array, # int32 words (bf16 pairs) of the hidden states [M, 6144] + rw: cutlass.Array, # f32 router weight [128, 6144] + sfsc: cutlass.Array, # f32 [1] fc31_input_scale + lgbuf: cutlass.Array, # int64 [LG_REPL, LG_STRIDE] + guw: cutlass.Array, + gusf: cutlass.Array, + dnw: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, +): + tidx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + if cutlass.const_expr(TICKET): + # role by entry order: the CTAs that become resident last (behind the predecessor's CTAs) take the shared + # expert roles (most slack), the first NCTA the routed ones (the router GEMV / logits exchange) + s_tk = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + if tidx == 0: + old = _atom_add_u64( + cnt.subview(cutlass.Int32(NG_CAP * 8 + 8)).data_ptr().toint(), cutlass.Int64(1) + ) + s_tk.store(cutlass.Int32(old % cutlass.Int64(NCTA + NSH)), idx=0) + prims.barrier_cta_sync(0) + bx = s_tk.load(idx=0) + if tidx == 0: + _stamp(tdbg, bx, 0) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) + ) + + s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) + sA = cutlass.Array( + s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 + ) + # ring scale words: SFA_CP: the instruction atoms [STAGES][4 x 512 B]; else raw boxes [STAGES][8 KB] + spliced + sSF = cutlass.Array( + s_big.data_ptr(_OFF_SFA), + shape=(STAGES * STAGE_SF // 4,), + dtype=cutlass.Int32, + alignment=1024, + ) + sSFA = cutlass.Array( + s_big.data_ptr(_OFF_SFA + (0 if SFA_CP else STAGES * NUM_BYTES_SFA_RAW)), + shape=(NUM_WORDS_SFA * STAGES,), + dtype=cutlass.Int32, + alignment=512, + ) + sB1 = cutlass.Array( + s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 + ) + sB2 = cutlass.Array( + s_big.data_ptr(_OFF_B2), shape=(NB2BUF * IMG_B2,), dtype=cutlass.Int8, alignment=1024 + ) + sSFB2 = cutlass.Array( + s_big.data_ptr(_OFF_SFB2), shape=(IMG_SF // 4,), dtype=cutlass.Int32, alignment=1024 + ) + ys = cutlass.Array( + cutlass.BFloat16, + NB2_CAP * M_MAX * TOP_K * R2, + space=cutlass.AddressSpace.smem, + alignment=16, + ) # bf16 values + s_part = cutlass.Array( + cutlass.Float32, SPB * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_red = cutlass.Array( + cutlass.Float32, 9 * M_MAX + 16, space=cutlass.AddressSpace.smem, alignment=16 + ) + # the activation's scale codes [M, 384] until they are in TMEM: aliases ys (first written by the FC2 epilogue) + s_xsf = cutlass.Array( + ys.data_ptr(0), shape=(M_MAX * (H // 64),), dtype=cutlass.Int32, alignment=16 + ) + ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, NB2BUF, space=cutlass.AddressSpace.smem, alignment=8) + b2_free = cutlass.Array(cutlass.Int64, NB2BUF, space=cutlass.AddressSpace.smem, alignment=8) + img_bar = cutlass.Array(cutlass.Int64, NB2BUF, space=cutlass.AddressSpace.smem, alignment=8) + sfb2_free = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + x_full = cutlass.Array(cutlass.Int64, NXB, space=cutlass.AddressSpace.smem, alignment=8) + x_empty = cutlass.Array(cutlass.Int64, NXB, space=cutlass.AddressSpace.smem, alignment=8) + route_bar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + ys_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + grp_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) + b1_load = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + topk_bar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tok0_bar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + px_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + s_prog = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + s_el = cutlass.Array(cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + s_kpos = cutlass.Array( + cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_w = cutlass.Array( + cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_ids = cutlass.Array( + cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + s_gn = cutlass.Array( + cutlass.Int32, NG_CAP, space=cutlass.AddressSpace.smem, alignment=16 + ) # slots per group + s_emap = cutlass.Array(cutlass.Int32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_plan_all = cutlass.Array( + cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 + ) + s_esc = cutlass.Array(cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32) + + is_routed = cutlass.Boolean(True) + if cutlass.const_expr(SHARED): + sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) + is_routed = bx < cutlass.Int32(NCTA) + if bx >= cutlass.Int32(NCTA): + _shared_cta( + tidx, + warp, + lane, + bx - cutlass.Int32(NCTA), + s_big, + sh_full, + epochs, + tdbg, + bx, + xb32, + tma_gu_desc.get_ptr(), + tma_dn_desc.get_ptr(), + gusf, + dnsf, + hx, + dnx, + arp, + tmem_ptr_i32, + guw, + dnw, + ) + if is_routed: + # ---- before the grid dependency + if warp == 0: + if tidx < STAGES: + prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) + prims.mbarrier_init(ab_empty.subview(tidx), 1) + prims.mbarrier_init(sfa_ready.subview(tidx), 1) + if tidx < 2: + prims.mbarrier_init(acc_full.subview(tidx), 1) + prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(acc2_full.subview(tidx), 1) + prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) + if tidx < NB2BUF: + prims.mbarrier_init(b2_ready.subview(tidx), 1) + prims.mbarrier_init(b2_free.subview(tidx), 1) + prims.mbarrier_init(img_bar.subview(tidx), 1) + if tidx < NXB: + prims.mbarrier_init(x_full.subview(tidx), 1) + prims.mbarrier_init(x_empty.subview(tidx), NWARPS) + if tidx == 0: + prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) + prims.mbarrier_init(sfb2_free.subview(0), 1) + prims.mbarrier_init(route_bar.subview(0), 32) + prims.mbarrier_init(ys_done.subview(0), EPI_THREADS) + prims.mbarrier_init(b1_load.subview(0), 1) + prims.mbarrier_init(topk_bar.subview(0), M_MAX * TOP_K) + prims.mbarrier_init(tok0_bar.subview(0), TOP_K) + prims.mbarrier_init(px_full.subview(0), 1) + if tidx < 4: + prims.mbarrier_init(grp_done.subview(tidx), EPI_THREADS) + if tidx == 0: + s_prog.store(cutlass.Int32(0), idx=0) + if warp == X_WARP: # slot tables: empty (the slot warp fills them after the top-k) + for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): + s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + for z in cutlass.range_constexpr(E_LOCAL // 32): + s_emap.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + if warp == MMA_WARP: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == W_WARP: + prims.prefetch_tensormap(tma_a1_desc.get_ptr()) + prims.prefetch_tensormap(tma_a2_desc.get_ptr()) + if cutlass.const_expr(not SFA_CP): + prims.prefetch_tensormap(tma_s1_desc.get_ptr()) + prims.prefetch_tensormap(tma_s2_desc.get_ptr()) + for i in cutlass.range_constexpr(E_LOCAL // 32): + s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) + for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): + q = lane + cutlass.Int32(32 * i) + if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): + gq = q // cutlass.Int32(NG_CAP + 1) + s_plan_all.store( + plan.load( + idx=(gq * cutlass.Int32(NCTA) + bx) * cutlass.Int32(NG_CAP + 1) + + q + - gq * cutlass.Int32(NG_CAP + 1) + ), + idx=q, + ) + # router weights (expert bx) in registers: thread tg = tid % 256 owns elements 4 tg + 1024 j + tg = tidx % cutlass.Int32(GEMV_THREADS) + half = tidx // cutlass.Int32(GEMV_THREADS) + rwv = [ + rw.load( + idx=bx * cutlass.Int32(H) + tg * cutlass.Int32(4) + cutlass.Int32(1024 * j), + vector_size=4, + alignment=16, + ) + for j in range(H // 1024) + ] + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + + cute.arch.griddepcontrol_wait() + # ---- x chunks into the idle ring (warp 6): chunk c = K [XCK c, XCK (c + 1)) of every token + if warp == S_WARP: + if cutlass.const_expr(NPT > 0): + # the priority tokens' rows first (their logits are published ahead of the others) + if lane == 0: + prims.mbarrier_arrive_expect_tx(px_full.subview(0), NPT * H * 2) + cute.arch.sync_warp() + if lane < NPT: + prims.cp_async_bulk_shared_cluster_global( + sB1.subview(lane * cutlass.Int32(H * 2)), + xb32.subview( + cutlass.select_( + lane == cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(PT1) + ) + * cutlass.Int32(H // 2) + ), + px_full.subview(0), + H * 2, + ) + for c in cutlass.range_constexpr(NXB): + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_full.subview(c), XCB) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(c * XCB) + lane * cutlass.Int32(XCK * 2)), + xb32.subview(lane * cutlass.Int32(H // 2) + cutlass.Int32(c * XCK // 2)), + x_full.subview(c), + XCK * 2, + ) + ar_on = arp.load(idx=0) != cutlass.Int64(0) + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) + + sc = cutlass.Float32(sfsc.load(idx=0)) + rsc = _rcp(sc) + rsix = _rcp(cutlass.Float32(6.0)) + img8f = cutlass.Array( + img32.data_ptr(0), shape=(IMG_TOTAL,), dtype=cutlass.Int8, alignment=16 + ) + if cutlass.const_expr(NPT > 0): + # ---- priority tokens (0 and PT1): their logits first (m3_router_gemv's per-thread / warp order) + _wait(px_full.subview(0), 0) + ph = cutlass.select_(half < cutlass.Int32(NPT), half, cutlass.Int32(0)) + pacc = cutlass.Float32(0.0) + px32 = cutlass.Array( + sB1.data_ptr(0), shape=(NPT * H // 2,), dtype=cutlass.Int32, alignment=16 + ) + for j in cutlass.range_constexpr(H // 1024): + pxw = px32.load( + idx=ph * cutlass.Int32(H // 2) + cutlass.Int32(512 * j) + tg * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + px0 = cutlass.Int32(pxw[0]) + px1 = cutlass.Int32(pxw[1]) + pacc = _fma_rn( + (px0 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][0]), + pacc, + ) + pacc = _fma_rn( + (px0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][1]), + pacc, + ) + pacc = _fma_rn( + (px1 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][2]), + pacc, + ) + pacc = _fma_rn( + (px1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][3]), + pacc, + ) + for off in (16, 8, 4, 2, 1): + pacc = _add_rn(pacc, cute.arch.shuffle_sync_bfly(pacc, off)) + if lane == 0: + s_red.store(pacc, idx=cutlass.Int32(9 * M_MAX) + warp) # [half][8 warps] + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ALL_BAR, number_of_threads=THREADS) + if tidx < cutlass.Int32(NPT * LG_REPL): + hh = tidx // cutlass.Int32(LG_REPL) + rr = tidx - hh * cutlass.Int32(LG_REPL) + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn( + sacc, s_red.load(idx=cutlass.Int32(9 * M_MAX + w) + hh * cutlass.Int32(8)) + ) + tp = cutlass.select_(hh == cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(PT1)) + _st_relaxed_u64( + lgbuf.subview(rr * cutlass.Int32(LG_STRIDE) + tp * cutlass.Int32(E_LOCAL) + bx) + .data_ptr() + .toint(), + ltag | (cutlass.Int64(sacc.bitcast(cutlass.Int32)) & cutlass.Int64(0xFFFFFFFF)), + ) + if tidx == 0: + _stamp(tdbg, bx, 43) + # ---- this CTA's share of the NVFP4 activation (QB blocks, from global) for the global B1 image: loads first + QB = (M_MAX * (H // 16) + NCTA - 1) // NCTA + qb = bx * cutlass.Int32(QB) + tidx + q_ok = (tidx < cutlass.Int32(QB)) & (qb < cutlass.Int32(M_MAX * (H // 16))) + qb = cutlass.select_(q_ok, qb, cutlass.Int32(0)) + qtok = qb // cutlass.Int32(H // 16) + qkb = qb - qtok * cutlass.Int32(H // 16) + xq_v = [ + xb32.load( + idx=qtok * cutlass.Int32(H // 2) + qkb * cutlass.Int32(8) + cutlass.Int32(4 * h), + vector_size=4, + alignment=16, + ) + for h in range(2) + ] + # ---- router GEMV from the staged chunks, all 512 threads (rolled: the front runs once per call) + MH = (M_MAX + 1) // 2 # tokens per GEMV half + acc_r = cute.make_rmem_tensor((MH,), cutlass.Float32) + for i in cutlass.range_constexpr(MH): + acc_r[i] = cutlass.Float32(0.0) + big32 = cutlass.Array( + s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16 + ) + JPC = XCK // 1024 # router j-steps per chunk + for xc in cutlass.range(NXC, unroll=1): + buf = xc % cutlass.Int32(NXB) + _wait(x_full.subview(buf), (xc // cutlass.Int32(NXB)) & cutlass.Int32(1)) + if (tidx == 0) & (xc == cutlass.Int32(0)): + _stamp(tdbg, bx, 41) + xw0 = buf * cutlass.Int32(XCB // 4) # chunk word base in s_big + for jj in cutlass.range_constexpr(JPC): + jg = xc * cutlass.Int32(JPC) + cutlass.Int32(jj) + wsel = [cutlass.Float32(rwv[0][z]) for z in range(4)] + for jx in cutlass.range_constexpr(1, H // 1024): + for z in cutlass.range_constexpr(4): + wsel[z] = cutlass.select_( + jg == cutlass.Int32(jx), cutlass.Float32(rwv[jx][z]), wsel[z] + ) + for i in cutlass.range_constexpr(MH): + t = half * cutlass.Int32(MH) + cutlass.Int32(i) + t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) + xw = big32.load( + idx=xw0 + + t * cutlass.Int32(XCK // 2) + + cutlass.Int32(512 * jj) + + tg * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + x0 = cutlass.Int32(xw[0]) + x1 = cutlass.Int32(xw[1]) + acc_g = acc_r[i] + acc_g = _fma_rn( + (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), wsel[0], acc_g + ) + acc_g = _fma_rn( + (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), wsel[1], acc_g + ) + acc_g = _fma_rn( + (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), wsel[2], acc_g + ) + acc_g = _fma_rn( + (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), wsel[3], acc_g + ) + acc_r[i] = acc_g + if cutlass.const_expr(NXB < NXC): + if xc + cutlass.Int32(NXB) < cutlass.Int32(NXC): + # every thread's reads of this chunk are done (CTA barrier, then a generic -> async proxy fence by + # the refilling warp: the bulk copy overwrites what the generic proxy read); racecheck-clean + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ALL_BAR, number_of_threads=THREADS) + if warp == S_WARP: + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_full.subview(buf), XCB) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + buf * cutlass.Int32(XCB) + lane * cutlass.Int32(XCK * 2) + ), + xb32.subview( + lane * cutlass.Int32(H // 2) + + (xc + cutlass.Int32(NXB)) * cutlass.Int32(XCK // 2) + ), + x_full.subview(buf), + XCK * 2, + ) + if tidx == 0: + _stamp(tdbg, bx, 42) + for i in cutlass.range_constexpr(MH): + v = acc_r[i] + for off in (16, 8, 4, 2, 1): + v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) + t = half * cutlass.Int32(MH) + cutlass.Int32(i) + if (lane == 0) & (t < cutlass.Int32(M_MAX)): + s_red.store(v, idx=(warp % cutlass.Int32(8)) * cutlass.Int32(M_MAX) + t) + # this CTA's quantization blocks into the global B1 image (B1 tile bytes, then the [M, 384] scale codes) + if q_ok: + qw = [] + for h in cutlass.range_constexpr(2): + for z in cutlass.range_constexpr(4): + qw.append(cutlass.Int32(xq_v[h][z])) + qlo, qhi, qsf = _quant_words(qw, sc, rsc, rsix) + img32.store( + (qlo, qhi), idx=cutlass.Int32(B1IMG_OFF // 4) + _b1_addr(qtok, qkb), alignment=8 + ) + img8f.store(cutlass.Int8(qsf), idx=cutlass.Int32(B1IMG_OFF + S1 * NUM_BYTES_B1) + qb) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ALL_BAR, number_of_threads=THREADS) + if warp == cutlass.Int32(0): + _dbg_delay(bx * cutlass.Int32(3)) + if tidx < LG_REPL: + for t in cutlass.range_constexpr(M_MAX): + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX + t))) + lw = ltag | (cutlass.Int64(sacc.bitcast(cutlass.Int32)) & cutlass.Int64(0xFFFFFFFF)) + _st_relaxed_u64( + lgbuf.subview(tidx * cutlass.Int32(LG_STRIDE) + cutlass.Int32(t * E_LOCAL) + bx) + .data_ptr() + .toint(), + lw, + ) + if tidx == 0: + _stamp(tdbg, bx, 21) + if tidx == cutlass.Int32( + THREADS - 32 + ): # B1 share released after the logits (its gpu fence stalls warp 15 only) + if cutlass.const_expr(EPI_PFENCE): + _fence_proxy_async_global() # once, after the join (ALL_BAR) of the writer threads + _red_release_add( + cnt.subview(cutlass.Int32(NG_CAP * 8)).data_ptr().toint(), cutlass.Int32(1) + ) + # ---- the B1 image (every CTA's share) into smem (warp 13; the top-k warps are 0 .. M - 1) + if cutlass.const_expr(B1_WARP >= 0): + if warp == cutlass.Int32(B1_WARP): + if lane == 0: + b1addr = cnt.subview(cutlass.Int32(NG_CAP * 8)).data_ptr().toint() + b1tgt = (ep + cutlass.Int32(1)) * cutlass.Int32(NCTA) + cv = _ld_acquire_u32(b1addr) + while (cv - b1tgt) < cutlass.Int32(0): + _nanosleep(64) + cv = _ld_acquire_u32(b1addr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx( + b1_load.subview(0), S1 * NUM_BYTES_B1 + M_MAX * (H // 16) + ) + prims.cp_async_bulk_shared_cluster_global( + sB1.subview(0), + img32.subview(cutlass.Int32(B1IMG_OFF // 4)), + b1_load.subview(0), + S1 * NUM_BYTES_B1, + ) + prims.cp_async_bulk_shared_cluster_global( + s_xsf.subview(0), + img32.subview(cutlass.Int32((B1IMG_OFF + S1 * NUM_BYTES_B1) // 4)), + b1_load.subview(0), + M_MAX * (H // 16), + ) + _stamp(tdbg, bx, 23) + cute.arch.sync_warp() + + # ---- top-k: token t on warp t (keys sigmoid + bias, ties to the lower id), TRTLLM-Gen MiniMax2 ops + tk_t = cutlass.Int32(-1) # this warp's top-k token (TK_WARPS[t] = warp) + for i_tk in cutlass.range_constexpr(M_MAX): + tk_t = cutlass.select_(warp == cutlass.Int32(TK_WARPS[i_tk]), cutlass.Int32(i_tk), tk_t) + if tk_t >= cutlass.Int32(0): + t = tk_t + want = ltag >> cutlass.Int64(32) + rbase = ( + (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) + + t * cutlass.Int32(E_LOCAL) + + lane + ) + done = cutlass.Int32(0) + res = [cutlass.Int64(0) for _ in range(4)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(4): + res[z] = _ld_relaxed_u64( + lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() + ) + ok = ok & ((res[z] >> cutlass.Int64(32)) == want) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): + done = cutlass.Int32(1) + else: + _nanosleep(POLL_NS) + if (lane == 0) & (t == cutlass.Int32(0)): + _stamp(tdbg, bx, 22) + lgv = [cutlass.Int32(res[z]).bitcast(cutlass.Float32) for z in range(4)] + keys = [] + for i in cutlass.range_constexpr(4): + sg = _fma_rn( + cutlass.Float32(0.5), + cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), + cutlass.Float32(0.5), + ) + keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) + hk = [_twiddle(keys[i]) for i in range(4)] + taken = [cutlass.Boolean(False) for _ in range(4)] + ids_t = [] + kt_t = [] + for k in cutlass.range_constexpr(TOP_K): + bh = cutlass.Int32(0) + bl = cutlass.Int32(0) + for i in cutlass.range_constexpr(4): + lo_i = cutlass.Int32(65535 - i * 32) - lane + hv = cutlass.select_(taken[i], cutlass.Int32(0), hk[i]) + lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) + better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( + (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) + ) + bh = cutlass.select_(better, hv, bh) + bl = cutlass.select_(better, lv, bl) + mh = _redux_max_u32(bh) + ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) + e_k = cutlass.Int32(65535) - ml + for i in cutlass.range_constexpr(4): + taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) + ids_t.append(e_k) + kv = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(4): + kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) + kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) + sig = cutlass.Float32(0.0) + e_l = cutlass.Int32(0) + for k in cutlass.range_constexpr(TOP_K): + sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) + sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) + e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) + total = sig + for off in (16, 8, 4, 2, 1): + total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) + wv = _bf16_round( + _div_rn( + _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), + _add_rn(total, cutlass.Float32(SUM_EPS)), + ) + ) + if lane < TOP_K: + s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) + s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) + if bx == 0: + route_dbg.store( + wv.bitcast(cutlass.Int32), + idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, + ) + if (lane == 0) & (t == cutlass.Int32(0)): + _stamp(tdbg, bx, 20) + if cutlass.const_expr(PF_PRIO): + if (t == cutlass.Int32(0)) | (t == cutlass.Int32(PT1)): + if lane < TOP_K: + # this CTA's 1/NCTA of each of the token's experts into L2 (DRAM is idle until the routing is + # complete; every byte is streamed later anyway) + _l2_prefetch( + w1b.subview( + e_l * cutlass.Int32(TWO_I * H // 2) + bx * cutlass.Int32(PF_W1) + ) + .data_ptr() + .toint(), + PF_W1, + ) + _l2_prefetch( + s1w.subview( + (e_l * cutlass.Int32(TWO_I * H // 16) + bx * cutlass.Int32(PF_S1)) + // cutlass.Int32(4) + ) + .data_ptr() + .toint(), + PF_S1, + ) + _l2_prefetch( + w2b.subview( + e_l * cutlass.Int32(H * I_TP // 2) + bx * cutlass.Int32(PF_W2) + ) + .data_ptr() + .toint(), + PF_W2, + ) + _l2_prefetch( + s2w.subview( + (e_l * cutlass.Int32(H * I_TP // 16) + bx * cutlass.Int32(PF_S2)) + // cutlass.Int32(4) + ) + .data_ptr() + .toint(), + PF_S2, + ) + cute.arch.sync_warp() + if lane < TOP_K: # every writer lane arrives (its own release of s_ids / s_w) + prims.mbarrier_arrive(topk_bar.subview(0)) # this token's ids / weights are in smem + if t == cutlass.Int32(0): + prims.mbarrier_arrive( + tok0_bar.subview(0) + ) # token 0's experts: the early group-0 stages + cute.arch.sync_warp() + + if cutlass.const_expr(EARLY0): + if warp == W_WARP: + # slots 0..3 = token 0's four distinct experts; group 0's stage ranges are the same for every G + _wait(tok0_bar.subview(0), 0) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(STAGES): + _issue_fc1( + lane, + n0, + e, + c, + j, + sA, + sSF, + ab_full, + tma_a1_desc.get_ptr(), + tma_s1_desc.get_ptr(), + s1w, + ) + n0 = n0 + cutlass.Int32(1) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + if warp == X_WARP: + # slots: entry i < 4 M is (token i / 4, top-k i % 4); distinct experts in entry order + NCH = (TOP_K * M_MAX + 31) // 32 + _wait(topk_bar.subview(0), 0) + gcnt = cutlass.Int32(0) + rec = [] # per chunk: (first, rank, slot rank of the entry, expert, entry valid) + for ch in cutlass.range_constexpr(NCH): + ei = lane + cutlass.Int32(32 * ch) + valid = ei < cutlass.Int32(TOP_K * M_MAX) + my_e = cutlass.select_( + valid, + s_ids.load(idx=cutlass.select_(valid, ei, cutlass.Int32(0))), + cutlass.Int32(-1) - lane, + ) + known = cutlass.Boolean(False) + prev = cutlass.Int32(-1) + key = my_e + if cutlass.const_expr(ch > 0): + prev = cutlass.select_( + valid, + s_emap.load(idx=cutlass.select_(valid, my_e, cutlass.Int32(0))), + cutlass.Int32(-1), + ) + known = prev >= cutlass.Int32(0) + key = cutlass.select_(known, cutlass.Int32(-1000) - lane, my_e) + cute.arch.sync_warp() # every lane's s_emap read before any lane's write + same = _match_any(key) + leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) + first = valid & (known == cutlass.Boolean(False)) & (lane == leader) + bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) + rank = gcnt + cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) + slot_l = cute.arch.shuffle_sync(rank, leader) + slot_l = cutlass.select_(known, prev, slot_l) + if first: + s_emap.store(rank, idx=my_e) + if valid: + if bx == 0: + route_dbg.store( + my_e, + idx=(ei // cutlass.Int32(TOP_K)) * cutlass.Int32(8) + + ei % cutlass.Int32(TOP_K), + ) + rec.append((first, rank, slot_l, my_e, valid)) + cute.arch.sync_warp() + gcnt = gcnt + cute.arch.popc(bal) + # slot of rank r: with a partial last group (tail < 4) and >= 3 groups, that group becomes group 1 (its + # FC2 then is not the last hand-off behind a short FC1 phase); the full groups 1.. shift up by one + ngr = (gcnt + cutlass.Int32(3)) >> cutlass.Int32(2) + tail = gcnt - (ngr - cutlass.Int32(1)) * cutlass.Int32(4) + do_perm = ( + cutlass.Boolean(PERM) + & (tail < cutlass.Int32(4)) + & (ngr >= cutlass.Int32(PERM_MIN_NG)) + ) + for z in cutlass.range_constexpr((G_CAP + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP): + s_el.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + cute.arch.sync_warp() + for ch in cutlass.range_constexpr(NCH): + first, rank, slot_r, my_e, valid = rec[ch] + ei = lane + cutlass.Int32(32 * ch) + for which in cutlass.range_constexpr(2): + r = rank if which == 0 else slot_r + sl_m = cutlass.select_( + r < cutlass.Int32(4), + r, + cutlass.select_( + r >= (ngr - cutlass.Int32(1)) * cutlass.Int32(4), + r - (ngr - cutlass.Int32(2)) * cutlass.Int32(4), + r + cutlass.Int32(4), + ), + ) + sl_m = cutlass.select_(do_perm, sl_m, r) + if cutlass.const_expr(which == 0): + if first: + s_el.store(my_e, idx=sl_m) + else: + if valid: + s_kpos.store( + ei % cutlass.Int32(TOP_K), + idx=sl_m * cutlass.Int32(M_MAX) + ei // cutlass.Int32(TOP_K), + ) + # slots per group + if lane < cutlass.Int32(NG_CAP): + nh = gcnt - lane * cutlass.Int32(4) + nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) + nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) + nh = cutlass.select_(do_perm & (lane == cutlass.Int32(1)), tail, nh) + nh = cutlass.select_( + do_perm & (lane == ngr - cutlass.Int32(1)), cutlass.Int32(4), nh + ) + s_gn.store(nh, idx=lane) + cute.arch.sync_warp() + if lane == 0: + s_el.store(gcnt, idx=G_CAP) + if lane < cutlass.Int32(NG_CAP + 1): + s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 24) + prims.mbarrier_arrive( + route_bar.subview(0) + ) # every lane: its slots / kpos / plan-row writes are in smem + cute.arch.sync_warp() + # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) + if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): + nh = s_gn.load(idx=lane) + pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) + if pad > cutlass.Int32(0): + _red_release_add(cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad) + else: + _wait(route_bar.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 2) + if cutlass.const_expr(B1_WARP < 0): + if warp == cutlass.Int32(0): + if lane == 0: + b1addr = cnt.subview(cutlass.Int32(NG_CAP * 8)).data_ptr().toint() + b1tgt = (ep + cutlass.Int32(1)) * cutlass.Int32(NCTA) + cv = _ld_acquire_u32(b1addr) + while (cv - b1tgt) < cutlass.Int32(0): + _nanosleep(64) + cv = _ld_acquire_u32(b1addr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx( + b1_load.subview(0), S1 * NUM_BYTES_B1 + M_MAX * (H // 16) + ) + prims.cp_async_bulk_shared_cluster_global( + sB1.subview(0), + img32.subview(cutlass.Int32(B1IMG_OFF // 4)), + b1_load.subview(0), + S1 * NUM_BYTES_B1, + ) + prims.cp_async_bulk_shared_cluster_global( + s_xsf.subview(0), + img32.subview(cutlass.Int32((B1IMG_OFF + S1 * NUM_BYTES_B1) // 4)), + b1_load.subview(0), + M_MAX * (H // 16), + ) + _stamp(tdbg, bx, 23) + cute.arch.sync_warp() + if warp < 4: # B1 scales to TMEM (the MMA needs them only with the first ring stage) + _wait(b1_load.subview(0), 0) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + + g_n = s_el.load(idx=G_CAP) + ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) + pr2 = s_plan.load(idx=NG_CAP) + blk0 = pr2 & cutlass.Int32(0xFFFF) + nb = pr2 >> cutlass.Int32(16) + + # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG + if warp == W_WARP: + stage = cutlass.Int32(0) + phase = cutlass.Int32(1) + skip = cutlass.Int32(0) + n_iss = cutlass.Int32( + 0 + ) # stages issued (early ones included): the L2-prefetch warp's throttle + w_emp = cutlass.Int64(0) + if cutlass.const_expr(EARLY0): + skip = s_misc.load(idx=0) + stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) + phase = cutlass.select_( + skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) + ) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + for j in cutlass.range(j0, j1, unroll=1): + if skip > cutlass.Int32(0): + skip = skip - cutlass.Int32(1) + else: + w_emp = _waitt(ab_empty.subview(stage), phase, w_emp) + prims.cp_async_wait_group(STAGES - 1) + _dbg_delay(bx) + _issue_fc1( + lane, + stage, + e, + c, + j, + sA, + sSF, + ab_full, + tma_a1_desc.get_ptr(), + tma_s1_desc.get_ptr(), + s1w, + ) + stage, phase = _adv(stage, phase) + if cutlass.const_expr(PFD > 0): + n_iss = n_iss + cutlass.Int32(1) + if lane == 0: + prims.store_ext( + n_iss.ir_value(), + s_prog.subview(0), + order="relaxed", + scope="cta", + ) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = s_gn.load(idx=g) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + w_emp = _waitt(ab_empty.subview(stage), phase, w_emp) + prims.cp_async_wait_group(STAGES - 1) + _issue_fc2( + lane, + stage, + g, + blk, + t2, + n_here, + s_el, + sA, + sSF, + ab_full, + tma_a2_desc.get_ptr(), + tma_s2_desc.get_ptr(), + s2w, + ) + stage, phase = _adv(stage, phase) + if cutlass.const_expr(PFD > 0): + n_iss = n_iss + cutlass.Int32(1) + if lane == 0: + prims.store_ext( + n_iss.ir_value(), + s_prog.subview(0), + order="relaxed", + scope="cta", + ) + if lane == 0: + _stamp(tdbg, bx, 3) + _tstore(tdbg, bx, 30, w_emp) + + # ---- (SFA_CP = 0) scale splicer (warp 6): atom kb word [r][q] = raw atom 4 q + kb word [r][m1] + if cutlass.const_expr(not SFA_CP): + if warp == S_WARP: + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + i0 = cutlass.Int32(0) + i1 = cutlass.Int32(0) + for part in cutlass.range_constexpr(2): + cnt_ = cutlass.Int32(0) + m1 = cutlass.Int32(0) + nk = cutlass.Int32(0) + if cutlass.const_expr(part == 0): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + nk = _nsegs(i0, i1) + else: + if p >= cutlass.Int32(LAG): + nk = nb + for k in range(nk): + if cutlass.const_expr(part == 0): + ug, j0, j1 = _seg(i0, i1, k) + m1 = ( + ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + ) & cutlass.Int32(3) + cnt_ = j1 - j0 + else: + m1 = (blk0 + k) & cutlass.Int32(3) + cnt_ = cutlass.Int32(KT2) + for _j in cutlass.range(cnt_, unroll=1): + _waitd(ab_full.subview(stage), phase, tdbg, 3) + raw0 = ( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + lane * cutlass.Int32(4) + + m1 + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32( + 4 + ) + for kb in cutlass.range_constexpr(NUM_KB): + w0 = sSF.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) + w1 = sSF.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) + w2 = sSF.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) + w3 = sSF.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) + sSFA.store( + ( + cutlass.Int32(w0), + cutlass.Int32(w1), + cutlass.Int32(w2), + cutlass.Int32(w3), + ), + idx=dst0 + cutlass.Int32(kb * 128), + alignment=16, + ) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + cute.arch.sync_warp() + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + prims.tcgen05_commit(sfa_ready.subview(stage)) + stage, phase = _adv(stage, phase) + + # ---- FC2 B loader (warp 5): group g's image once its counter is complete + if warp == X_WARP: + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + for g in range(ng): + buf = g % cutlass.Int32(NB2BUF) + if g >= cutlass.Int32(NB2BUF): + _waitd( + b2_free.subview(buf), + ((g // cutlass.Int32(NB2BUF)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 29) + caddr = cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() + if lane == 0: + cv = _ld_acquire_u32(caddr) + while (cv - cnt_target) < cutlass.Int32(0): + _nanosleep(B2_POLL_NS) + cv = _ld_acquire_u32(caddr) + cute.arch.sync_warp() + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 51) + _fence_proxy_async_global() + if g >= cutlass.Int32(1): + _waitd( + sfb2_free.subview(0), (g - cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 107 + ) + if lane == 0: + prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) + prims.cp_async_bulk_shared_cluster_global( + sB2.subview(buf * cutlass.Int32(IMG_B2)), + img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), + img_bar.subview(buf), + IMG_B2, + ) + prims.cp_async_bulk_shared_cluster_global( + sSFB2.subview(0), + img32.subview( + g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) + ), + img_bar.subview(buf), + IMG_SF, + ) + cute.arch.sync_warp() + _waitd( + img_bar.subview(buf), (g // cutlass.Int32(NB2BUF)) & cutlass.Int32(1), tdbg, 108 + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 52) + if g >= cutlass.Int32(1): + # the single TMEM scale slot: the previous group's FC2 MMAs are done + gp = g - cutlass.Int32(1) + _waitd( + b2_free.subview(gp % cutlass.Int32(NB2BUF)), + (gp // cutlass.Int32(NB2BUF)) & cutlass.Int32(1), + tdbg, + 109, + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 53) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for aa in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SFB2_COL + aa * 4), + sd0 + cutlass.Int64(aa * (512 >> 4)), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + prims.tcgen05_commit(sfb2_free.subview(0)) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 25) + + # ---- (PFD) L2 prefetch warp (12): the same stage sequence as the producer, PFD stages ahead of it + if cutlass.const_expr(PFD > 0): + if warp == cutlass.Int32(PF_WARP): + pf = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + for j in cutlass.range(j0, j1, unroll=1): + if lane == 0: + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + while prg + cutlass.Int32(PFD) <= pf: + _nanosleep(100) + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + if pf >= prg: + kx = j * cutlass.Int32(C1 * KT_BYTES) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_prefetch( + tma_a1_desc.get_ptr(), + [ + kx + cutlass.Int32(q * KT_BYTES), + c * cutlass.Int32(R1), + e, + ], + [], + ) + atom0 = ( + e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2)) + ) * cutlass.Int32(SF1_KA) + j * cutlass.Int32(16) + _l2_prefetch( + s1w.subview(atom0 * cutlass.Int32(128)) + .data_ptr() + .toint(), + 8192, + ) + cute.arch.sync_warp() + pf = pf + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = s_gn.load(idx=g) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + if lane == 0: + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + while prg + cutlass.Int32(PFD) <= pf: + _nanosleep(100) + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + if pf >= prg: + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load( + idx=g * cutlass.Int32(4) + cutlass.Int32(jj) + ) + prims.cp_async_bulk_tensor_prefetch( + tma_a2_desc.get_ptr(), + [ + cutlass.Int32(t2 * KT_BYTES), + blk * cutlass.Int32(R2), + e, + ], + [], + ) + a2 = ( + e * cutlass.Int32(SF2_RB) + + (blk >> cutlass.Int32(2)) + ) * cutlass.Int32(SF2_KA) + cutlass.Int32(4 * t2) + _l2_prefetch( + s2w.subview(a2 * cutlass.Int32(128)) + .data_ptr() + .toint(), + 2048, + ) + cute.arch.sync_warp() + pf = pf + cutlass.Int32(1) + + # ---- (REL_WARP) group releases (warp 8): this CTA's finished units of group p, one release add per phase + if cutlass.const_expr(REL_WARP): + if warp == cutlass.Int32(8): + for p in range(ng): + i0, i1 = _plan_range(s_plan, p) + n_done = cutlass.Int32(0) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + n_done = n_done + cutlass.select_( + (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), + cutlass.Int32(1), + cutlass.Int32(0), + ) + _wait( + grp_done.subview(p & cutlass.Int32(3)), + (p >> cutlass.Int32(2)) & cutlass.Int32(1), + ) + _dbg_delay(bx * cutlass.Int32(13) + p) + if (lane == 0) & (n_done > cutlass.Int32(0)): + if cutlass.const_expr(EPI_PFENCE): + _fence_proxy_async_global() + _red_release_add( + cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done + ) + if (lane == 0) & (p == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 54) + cute.arch.sync_warp() + + # ---- MMA (warp 7) + if warp == MMA_WARP: + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N1, + m_dim=MMA_M, + sparsity_version=0, + ) + idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N2, + m_dim=MMA_M, + sparsity_version=0, + ) + desc_a0 = prims.Tcgen05SmemDesc.build( + sA.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b10 = prims.Tcgen05SmemDesc.build( + sB1.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b20 = prims.Tcgen05SmemDesc.build( + sB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + _waitd(b1_ready, 0, tdbg, 5) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + m_b2 = cutlass.Int64(0) + m_full = cutlass.Int64(0) + m_acc = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + m_acc = _waitt( + acc_empty.subview(ab), + ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + m_acc, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_ptr = cutlass.inttoptr( + tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 + ) + for j in cutlass.range(j0, j1, unroll=1): + if cutlass.const_expr(SFA_CP): + m_full = _waitt(ab_full.subview(stage), phase, m_full) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 6) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc1, + (j > j0) | cutlass.Boolean(kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + cute.arch.sync_warp() + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(ab)) + cute.arch.sync_warp() + k1 = k1 + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + buf = g % cutlass.Int32(NB2BUF) + m_b2_0 = m_b2 + m_b2 = _waitt( + b2_ready.subview(buf), (g // cutlass.Int32(NB2BUF)) & cutlass.Int32(1), m_b2 + ) + if cutlass.const_expr(TIMING): + if (lane == 0) & (g < cutlass.Int32(4)): + tdbg.store( + m_b2 - m_b2_0, + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(44) + g, + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 4) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + m_acc = _waitt( + acc2_empty.subview(a2), + ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + m_acc, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc2_ptr = cutlass.inttoptr( + tcol + ACC2_COL + a2 * cutlass.Int32(ACC2_STRIDE), 6, cutlass.Float32 + ) + for t2 in cutlass.range_constexpr(KT2): + if cutlass.const_expr(SFA_CP): + m_full = _waitt(ab_full.subview(stage), phase, m_full) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 7) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b20 + cutlass.Int64( + (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) + * cutlass.Int32(NUM_BYTES_B2 >> 4) + ) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB2_COL + cutlass.Int32(t2 * NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc2_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc2, + cutlass.Boolean(t2 > 0 or kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + cute.arch.sync_warp() + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc2_full.subview(a2)) + cute.arch.sync_warp() + k2 = k2 + cutlass.Int32(1) + if prims.elect_sync(): + prims.tcgen05_commit( + b2_free.subview(buf) + ) # this group's B2 tile and TMEM scales are free + cute.arch.sync_warp() + if lane == 0: + _tstore(tdbg, bx, 26, m_b2) + _tstore(tdbg, bx, 27, m_full) + _tstore(tdbg, bx, 28, m_acc) + + # ---- epilogue warps (0-3) + if warp < 4: + row_w = base_row + warp * 32 + img8 = cutlass.Array( + img32.data_ptr(0), shape=(NG_CAP * IMG_BYTES,), dtype=cutlass.Int8, alignment=16 + ) + for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): + qi = tidx + cutlass.Int32(r * EPI_THREADS) + sl_e = qi // cutlass.Int32(6) + if qi < cutlass.Int32(6 * G_CAP): + e_sl = s_el.load(idx=sl_e) + if e_sl >= cutlass.Int32(0): + fld = qi - sl_e * cutlass.Int32(6) + s_esc.store( + cutlass.Float32(escal.load(idx=fld * cutlass.Int32(E_LOCAL) + e_sl)), + idx=sl_e * cutlass.Int32(8) + fld, + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + e_wait = cutlass.Int64(0) + e_busy = cutlass.Int64(0) + e_a = cutlass.Int64(0) + e_pw = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + unit = slot * cutlass.Int32(U1) + cblk + unit_done = (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)) + e_wait = _waitt( + acc_full.subview(ab), + (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), + e_wait, + ) + tb0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 + ), + num=N1_LD, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc_empty.subview(ab)) + for t in cutlass.range_constexpr(M_MAX): + part = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + part = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part + ) + s_part.store(part, idx=(t * C1 + warp) * R1 + lane) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + tc0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(TIMING): + e_a = e_a + (tc0 - tb0) + g_u = slot >> cutlass.Int32(2) + jrow = slot & cutlass.Int32(3) + # this warp's tokens t = warp + 4 tr, all in flight together (latency-bound otherwise) + NT = (M_MAX + 3) // 4 + xr = cute.make_rmem_tensor((NT,), cutlass.Float32) + for tr in cutlass.range_constexpr(NT): + tq = cutlass.select_( + warp + cutlass.Int32(4 * tr) < cutlass.Int32(M_MAX), + warp + cutlass.Int32(4 * tr), + cutlass.Int32(0), + ) + xr[tr] = _add_rn( + _add_rn( + _add_rn( + s_part.load(idx=(tq * C1 + 0) * R1 + lane), + s_part.load(idx=(tq * C1 + 1) * R1 + lane), + ), + s_part.load(idx=(tq * C1 + 2) * R1 + lane), + ), + s_part.load(idx=(tq * C1 + 3) * R1 + lane), + ) + if (j0 == cutlass.Int32(0)) & (j1 != cutlass.Int32(S1)): + # head of a unit split with the next CTA: its partial sums + for tr in cutlass.range_constexpr(NT): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + _st_relaxed_u32( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint(), + xr[tr].bitcast(cutlass.Int32), + ) + else: + if j0 != cutlass.Int32(0): + # tail: the head's partial sums (head first, as TRTLLM-Gen's K order) + for tr in cutlass.range_constexpr(NT): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + paddr = ( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint() + ) + pbits = _ld_relaxed_u32(paddr) + while pbits == cutlass.Int32(PARTIAL_EMPTY): + pbits = _ld_relaxed_u32(paddr) + _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) + xr[tr] = _add_rn(pbits.bitcast(cutlass.Float32), xr[tr]) + cute.arch.sync_warp() + _dbg_delay(bx * cutlass.Int32(7) + warp) + es = slot * cutlass.Int32(8) + ea = s_esc.load(idx=es) + ec = s_esc.load(idx=es + cutlass.Int32(1)) + eal = s_esc.load(idx=es + cutlass.Int32(2)) + ebe = s_esc.load(idx=es + cutlass.Int32(3)) + eli = s_esc.load(idx=es + cutlass.Int32(4)) + is_up = ( + (lane >> cutlass.Int32(3)) & cutlass.Int32(1) + ) == cutlass.Int32(0) + n0 = jrow * cutlass.Int32(M_MAX) + t2 = cblk >> cutlass.Int32(4) + cc = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) + gbase = g_u * cutlass.Int32(IMG_BYTES) + sfoff = ( + gbase + + cutlass.Int32(IMG_B2) + + ( + ( + t2 * cutlass.Int32(NUM_KB) + + ((cblk & cutlass.Int32(15)) >> cutlass.Int32(2)) + ) + * cutlass.Int32(512) + ) + + (cblk & cutlass.Int32(3)) + ) + for tr in cutlass.range_constexpr(NT): + # (skip) only tokens routed to this slot (warp-uniform); other B2 rows are never used + if s_kpos.load( + idx=slot * cutlass.Int32(M_MAX) + + cutlass.select_( + warp + cutlass.Int32(4 * tr) < cutlass.Int32(M_MAX), + warp + cutlass.Int32(4 * tr), + cutlass.Int32(0), + ) + ) >= cutlass.Int32(0): + xv = xr[tr] + gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) + qv, sfb = _swiglu_quant( + xv, gate_raw, is_up, ea, ec, eal, ebe, eli + ) + byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) + bt = [ + cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) + for i in range(8) + ] + lo = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | (bt[3] << cutlass.Int32(24)) + ) + hi = ( + bt[4] + | (bt[5] << cutlass.Int32(8)) + | (bt[6] << cutlass.Int32(16)) + | (bt[7] << cutlass.Int32(24)) + ) + t = warp + cutlass.Int32(4 * tr) + if ( + (lane == 0) + & (t < cutlass.Int32(M_MAX)) + & cutlass.Boolean(not DBG_NOSTORE) + ): + # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 + # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), row n + # at byte 16 (n % 32) + 4 (n / 32), byte c % 4 + n = n0 + t + img32.store( + (lo, hi), + idx=( + gbase + + t2 * cutlass.Int32(NUM_BYTES_B2) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + img8.store( + cutlass.Int8(sfb & cutlass.Int32(0xFF)), + idx=sfoff + + (n & cutlass.Int32(31)) * cutlass.Int32(16) + + (n >> cutlass.Int32(5)) * cutlass.Int32(4), + ) + tc1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(TIMING): + e_pw = e_pw + (tc1 - tc0) + if cutlass.const_expr(not REL_WARP): + if unit_done: + if lane == 0: + if cutlass.const_expr(EPI_PFENCE): + _fence_proxy_async_global() + if cutlass.const_expr(EPI_FENCE): + _fence_acq_rel_gpu() + cute.arch.sync_warp() + cute.arch.barrier( + barrier_id=EPI_BAR, number_of_threads=EPI_THREADS + ) # s_part reads done + if cutlass.const_expr(not REL_WARP): + if unit_done & (tidx == 0): + _red_release_add( + cnt.subview(g_u * cutlass.Int32(8)).data_ptr().toint(), + cutlass.Int32(1), + ) + if cutlass.const_expr(TIMING): + e_busy = e_busy + (_gtime() - tb0) + k1 = k1 + cutlass.Int32(1) + if cutlass.const_expr(REL_WARP): + prims.mbarrier_arrive( + grp_done.subview(p & cutlass.Int32(3)) + ) # this CTA's group-p units are out + if tidx == 0: + _stamp(tdbg, bx, 10) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + e_wait = _waitt( + acc2_full.subview(a2), + (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), + e_wait, + ) + if cutlass.const_expr(TIMING): + if ( + (tidx == 0) + & (g == ng - cutlass.Int32(1)) + & (bb == cutlass.Int32(0)) + ): + _stamp(tdbg, bx, 8) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = [] + ya = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, + base_col + ACC2_COL + a2 * cutlass.Int32(ACC2_STRIDE), + cutlass.Float32, + ), + num=ACC2_STRIDE, + ) + for t in cutlass.range_constexpr(M_MAX): + v = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + v = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(ya[w * M_MAX + t]), v + ) + yv.append(v) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc2_empty.subview(a2)) + sl = g * cutlass.Int32(4) + warp + if warp < s_gn.load(idx=g): + a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) + for t in cutlass.range_constexpr(M_MAX): + kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + if kp >= cutlass.Int32(0): + ys.store( + cutlass.BFloat16( + _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)) + ), + idx=( + (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + * cutlass.Int32(TOP_K) + + kp + ) + * cutlass.Int32(R2) + + lane, + ) + k2 = k2 + cutlass.Int32(1) + if tidx == 0: + _stamp(tdbg, bx, 6) + _tstore(tdbg, bx, 31, e_wait) + _tstore(tdbg, bx, 34, e_busy) + _tstore(tdbg, bx, 36, e_a) + _tstore(tdbg, bx, 38, e_pw) + prims.mbarrier_arrive(ys_done.subview(0)) # this thread's ys values are in smem + + # ---- finalize on all 16 warps: (block bb, token t) pairs, the fma chain over t's top-4 in order, + y_sh + NPAIR = NB2_CAP * M_MAX + NPW = (NPAIR + NWARPS - 1) // NWARPS + wyl = [] + for bt in cutlass.range_constexpr(NPW): + pair = warp + cutlass.Int32(bt * NWARPS) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + ok = (pair < cutlass.Int32(NPAIR)) & (bb < nb) + ch = cutlass.select_( + ok, + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + cutlass.Int32(0), + ) + wyl.append( + _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + if cutlass.const_expr(SHARED) + else cutlass.Int32(0) + ) + _wait(ys_done.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 19) + for bt in cutlass.range_constexpr(NPW): + pair = warp + cutlass.Int32(bt * NWARPS) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + if (pair < cutlass.Int32(NPAIR)) & (bb < nb): + acc_o = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(TOP_K): + acc_o = _fma_rn( + s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), + cutlass.Float32( + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ) + ), + acc_o, + ) + ch = ( + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(SHARED): + wy = wyl[bt] + while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: + wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + acc_o = _add_rn( + (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) + ) + out.store(cutlass.BFloat16(acc_o), idx=ch) + if ar_on: + cl = ( + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + _ar_push( + arp, + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + + # ---- teardown + if tidx == 0: + _stamp(tdbg, bx, 7) + epochs.store(ep + cutlass.Int32(1), idx=bx) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if ar_on | cutlass.Boolean(LATE_TRIGGER): + cute.arch.griddepcontrol_launch_dependents() + if warp == MMA_WARP: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS + ) + + +@cute.jit +def m3_moe( + a1_tensor: cute.Tensor, + sfa1_tensor: cute.Tensor, + a2_tensor: cute.Tensor, + sfa2_tensor: cute.Tensor, + bias: cute.Tensor, + escal: cute.Tensor, + plan: cute.Tensor, + partials: cute.Tensor, + img32: cute.Tensor, + cnt: cute.Tensor, + s1w: cute.Tensor, + s2w: cute.Tensor, + w1b: cute.Tensor, + w2b: cute.Tensor, + epochs: cute.Tensor, + out: cute.Tensor, + route_dbg: cute.Tensor, + tdbg: cute.Tensor, + xb32: cute.Tensor, + rw: cute.Tensor, + sfsc: cute.Tensor, + lgbuf: cute.Tensor, + guw: cute.Tensor, + gusf: cute.Tensor, + dnw: cute.Tensor, + dnsf: cute.Tensor, + hx: cute.Tensor, + dnx: cute.Tensor, + arp: cute.Tensor, + stream: cuda_driver.CUstream, +): + tma_a1_desc = cuda.create_tensor_map_tiled( + global_address=a1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H // 2, TWO_I, E_LOCAL], + global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], + box_dims=(KT_BYTES, R1, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_a2_desc = cuda.create_tensor_map_tiled( + global_address=a2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[I_TP // 2, H, E_LOCAL], + global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], + box_dims=(KT_BYTES, R2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_s1_desc = cuda.create_tensor_map_tiled( + global_address=sfa1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_s2_desc = cuda.create_tensor_map_tiled( + global_address=sfa2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_gu_desc = tma_a1_desc + tma_dn_desc = tma_a2_desc + if cutlass.const_expr(SHARED): + tma_gu_desc = cuda.create_tensor_map_tiled( + global_address=guw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H, 2 * SH_I], + global_strides=[H // 16], + box_dims=(KT_BYTES, 32), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_dn_desc = cuda.create_tensor_map_tiled( + global_address=dnw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[SH_I, H], + global_strides=[SH_I // 16], + box_dims=(KT_BYTES, 128), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + m3_moe_kernel( + tma_a1_desc, + tma_a2_desc, + tma_s1_desc, + tma_s2_desc, + tma_gu_desc, + tma_dn_desc, + bias, + escal, + plan, + partials, + img32, + cnt, + s1w, + s2w, + w1b, + w2b, + epochs, + out, + route_dbg, + tdbg, + xb32, + rw, + sfsc, + lgbuf, + guw, + gusf, + dnw, + dnsf, + hx, + dnx, + arp, + ).launch( + grid=(NCTA + NSH, 1, 1), + block=(THREADS, 1, 1), + cluster=(1, 1, 1), + stream=stream, + use_pdl=True, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/m3_moe_kernel.py new file mode 100644 index 000000000000..1ad72e733c64 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/m3_moe_kernel.py @@ -0,0 +1,3719 @@ +# 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. +"""MiniMax-M3 routed experts for up to 8 decode tokens on one TP rank: ``m3_moe`` (group-phased weight stream). + +One persistent kernel replaces TRTLLM-Gen's MiniMax2 routing, FC1 (NVFP4 x NVFP4 + SwiGLU-OAI + NVFP4 requantization), +FC2 and finalize, optionally with the router GEMV and the activation's NVFP4 quantization folded in (FUSED_ROUTER). It +writes the rank's routed partial [M, 6144] bf16, bit for bit what the TRTLLM-Gen op returns. + +- Routing (warps 4..7, token t on warp 4 + t % 4): TRTLLM-Gen's MiniMax2 top-4, operation for operation; warp 4 then + makes the distinct experts slots in (token, top-k) order (match.any), slot s in group s / 4. +- Work: the groups' FC1 stages and FC2 tiles stream through one ring per CTA in phases: phase p = FC1 of group p (a + stream-K range of the group's unit-major stages), then FC2 of group p - LAG (this CTA's 32-row output blocks, 4 + slots' slices stacked in one 128-row MMA, B row j M + t). LAG phases of streaming hide the hand-off of group + p - LAG's intermediate between CTAs. +- FC1 stage: a 32-row unit's rows at 4 consecutive 256-K tiles (rows 32 q + r) against resident B rows QS q + t. + FC2 stage: 4 slots' 32 rows at one 256-K tile. Both carry 8 KB of raw scale atoms, spliced into the 4 instruction + atoms by warp 6 with one formula (word [r][q] <- raw atom 4 q + kb, word [r][m1]). +- FC1 epilogue: k-quarters summed in order, SwiGLU-OAI, NVFP4 requant into tagged 16-byte records (3 words of 3 data + bytes + the call's tag); a unit split between two CTAs is summed by its tail's holder (Lamport slot). +- Warps 5 and 8.. (XW warps) build each group's FC2 B (tiles + scale atoms) from the records once they carry the + call's tag; warp 5 copies the scales to TMEM. +- Finalize: bf16(fc2_alpha * acc) per (slot, token) into its top-k position, then per token the fma chain over its + top-4 in order, bf16 out. With an all-reduce descriptor, every output word is also pushed to the exchange + (AR_PRE: call parities and descriptor loaded during the last group's FC2). + +Configuration (trace-time constants) via M3_CONFIG = {"m_max", "num_ctas", "stages", "nb2_cap", "fused_router", "lag", +"xw", ...}; the plan (op.make_phase_plan) must be built for the same lag. +Launched with PDL: barrier setup, TMEM allocation, tensormap prefetch, bias / plan / router weight loads before +griddepcontrol.wait; every producer-output read and every global write after it. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +_CFG = globals().get("M3_CONFIG") or {} + + +def _cfg(key: str, default): + return type(default)(_CFG.get(key, default)) + + +# ============================================================================= +# Problem shape (trace-time constants). +# ============================================================================= +H = 6144 # capacity: fixed: the checkpoint's hidden size +I_TP = 768 +TWO_I = 2 * I_TP +E_LOCAL = 128 +TOP_K = 4 # capacity: fixed: the checkpoint's top-k +ROUTE_SCALE = 2.0 +SUM_EPS = 1e-20 +M_MAX = _cfg("m_max", 1) +NCTA = _cfg("num_ctas", 152) +FUSED_ROUTER = bool(_cfg("fused_router", 0)) +NROUTE = E_LOCAL if FUSED_ROUTER else 0 # experts whose logits the router GEMV computes +NGEMV = (NROUTE + NCTA - 1) // NCTA # experts per CTA: bx, bx + NCTA, .. +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +XVDUMP = bool( + _cfg("xvdump", 0) +) # debug only (needs a plan without split units): FC1 sums into `partials` +YDUMP = bool( + _cfg("ydump", 0) +) # debug only: FC2 outputs per (token, k) into `partials` (breaks its protocol) +LAG = _cfg("lag", 2) # phase p runs FC1 of group p and FC2 of group p - LAG +G_CAP = TOP_K * M_MAX # distinct experts per call +NG_CAP = M_MAX # groups of 4 slots +XW = _cfg("xw", 5 if M_MAX > 4 else (4 if M_MAX > 2 else 1)) # FC2 B builder warps: 5, then 8 .. +THREADS = 256 + 32 * (XW - 1) +GEMV_THREADS = 256 # capacity: fixed: the CTA size +# SHARED: the MoE's shared expert (MXFP8 gate_up [1536, 6144] -> SwiGLU-OAI -> MXFP8 down [6144, 768]) on NSH more +# CTAs (warps 0-3), its output added in finalize: out = bf16(shared + routed) (shared_output.add_(routed_output)). +# The math is mxfp8_low_m_gemm's (native MXFP8 activation quant, mma.sync m16n8k32 e4m3 per 32-K block, block scales +# as one fma, its partial-sum order, bf16) and the Triton silu_and_mul's, operation for operation. Shared CTA s owns +# intermediate positions 32 s .. 32 s + 31 (gate and up rows: MMA tiles = warps 0-3) and down rows 256 s .. + 255; +# 16-row x 768-byte weight chunks (+ their scale atoms) stream through NBW buffers per warp, the first NBW before +# the grid dependency. h crosses the shared CTAs, y_sh reaches the routed CTAs' finalize (both tagged words). +SHARED = bool(_cfg("shared", 0)) +# LATE_TRIGGER: release the dependents at teardown, not after the grid dependency. An early release lets the next +# kernels' CTAs (the all-reduce, then the attention it releases at entry) become resident during this kernel, packed +# onto the SMs it leaves free, where the attention then runs on a fraction of the GPU. +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +# AR_PRE: the finalize's all-reduce push uses the call parities and descriptor fields loaded while the last group's +# FC2 runs (not two dependent global loads per push on the kernel's tail). +AR_PRE = bool(_cfg("ar_pre", 1)) +# XSTAGE (fused router): the hidden states land in the (still idle) ring by one bulk copy per token right after the +# grid dependency; the NVFP4 quantization of B1 reads them from shared memory (one HBM / L2 round trip, not ~8). +XSTAGE = bool(_cfg("xstage", 1)) and bool(_cfg("fused_router", 0)) +NSH = 24 if SHARED else 0 +SH_I = 768 +SH_KB = H // 32 # gate_up K blocks # capacity: fixed: a kernel tile constant +SH_DKB = SH_I // 32 # down K blocks # capacity: fixed: a kernel tile constant +SH_SPLIT = 2 # gate_up K slices (= m3_mx_gemv's cluster ranks, contiguous K ranges): MMA rows 64 q + .. = slice q +SH_GT = ( + SH_KB // 4 // SH_SPLIT +) # gate_up k-tiles (128 K) per slice # capacity: fixed: a kernel tile constant +SH_STAGES_GU = SH_GT +SH_STAGES = SH_GT + 2 * ( + SH_I // 128 +) # + 2 down tiles x 6 k-tiles # capacity: fixed: pipeline stages +SHX_BAR = 6 +# SH_PRETILE: the shared weights come pre-tiled (op._pretile_shared): per shared CTA its 36 stages, each the 16 KB +# shared-memory image of its 128-row MMA tile (128-byte swizzle) followed by the stage's 512-byte instruction scale +# atom, so a stage is two contiguous bulk copies (not 128 row segments of 128 bytes) and needs no splicing. +SH_PRETILE = bool( + _cfg("sh_pretile", 1 if M_MAX <= 2 else 0) +) # (in-chain: a win at M <= 2, not at 4..8) +SH_STAGE_BYTES = 16384 + 512 +MMA_M = 128 +KT_BYTES = 128 +INC_K = 2 +NUM_KB = 4 +# FC1 +N1 = 16 if M_MAX <= 4 else 32 # B rows: QS q + t (token t at the stage's k-tile 4 j + q) +QS = N1 // 4 +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 # 48 units per slot +K1_TILES = H // 256 +S1 = K1_TILES // C1 # 6 stages per unit # capacity: fixed: pipeline stages +NUM_BYTES_A = MMA_M * KT_BYTES # capacity: fixed: a stage's operand bytes +NUM_BYTES_BOX1 = R1 * KT_BYTES +SF1_KA = H // 64 +SF1_RB = TWO_I // 128 +NUM_BYTES_SFA_RAW = 16 * 512 +NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 +NUM_BYTES_SFA = NUM_KB * 512 # capacity: fixed: a stage's scale bytes +NUM_WORDS_SFA = NUM_BYTES_SFA // 4 +NUM_BYTES_B1 = N1 * KT_BYTES +NUM_SFB1_ATOMS = S1 * NUM_KB +# FC2 +R2 = 32 +NB2 = H // R2 # 192 +KT2 = I_TP // 256 # 3 +SF2_KA = I_TP // 64 +SF2_RB = H // 128 +N2 = max(8, (4 * M_MAX + 7) // 8 * 8) # B rows j M + t +NUM_BYTES_B2 = N2 * KT_BYTES +NB2_CAP = _cfg("nb2_cap", 2) +UNITS_ROW = TWO_I // R1 # 48 records per (slot, token) row +H_ROW = UNITS_ROW * 16 +PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +B2_BAR = 5 +TD_STRIDE = 64 # tdbg columns per CTA +TOPK_BAR = 4 +B1_AFTER_TOPK = bool( + _cfg("b1_after_topk", 0) +) # (experiment) the B1 build waits for the routing warps' top-k +TOPK2_BAR = 7 +EARLY0 = bool( + _cfg("early0", 1) +) # warp 4 streams group 0 (token 0's experts) right after its own top-k +EARLY0_N = _cfg("early0_n", 99) # stages issued early (capped at STAGES) +# routing warps: 4..7, then the FC2 B builder warps 8.. (idle until the routing is done) up to one warp per token +RW_EXTRA = max(0, min(M_MAX - 4, (THREADS // 32) - 8)) if bool(_cfg("rw_extra", 1)) else 0 +NRW0 = ( + min(4, M_MAX) if bool(_cfg("rw_trim", 1)) else 4 +) # warps 4.. with a token (no idle pollers in the barrier) +NRW = NRW0 + RW_EXTRA +LG_REPL = _cfg("lg_repl", 32 if M_MAX > 4 else 128) +GEMV_SMEM = bool( + _cfg("gemv_smem", 1) +) # (XSTAGE) the router GEMV reads the staged hidden states from smem +LG_STRIDE = M_MAX * E_LOCAL + 32 +POLL_NS = _cfg("poll_ns", 150) +B2_POLL_NS = _cfg("b2_poll_ns", 200) +# B2IMG (needs SFA_CP: warp 6 is free): the FC1 epilogue writes each finished unit's FP4 bytes and scale byte straight +# into a per-group global image laid out as the FC2 B tile + its scale atoms (exactly what the builder warps assembled +# from the tagged records); per group, warp 6 joins the epilogue (fc1_done), fences the async proxy and does one +# release add of this CTA's finished units to the group counter; warp 5 acquires the counter (every group gains +# CNT_PER_CALL per call, CTA 0 pads absent slots), fences, and bulk-copies the image into the B2 buffer. +B2IMG = bool(_cfg("b2img", 0)) and bool(_cfg("sfa_cp", 0)) +IMG_B2 = KT2 * NUM_BYTES_B2 +IMG_SF = KT2 * NUM_KB * 512 +IMG_BYTES = IMG_B2 + IMG_SF +CNT_PER_CALL = 4 * (TWO_I // R1) +B2IMG_NS = _cfg("b2img_ns", 16) # (B2IMG) counter poll back-off +B2_DLY = _cfg( + "b2_dly", 0 +) # (stress) the release warp spins up to this many ns (CTA- and group-dependent) first +W_WARP = 4 # capacity: fixed: a warp role index +X_WARP = 5 # capacity: fixed: a warp role index +S_WARP = 6 # capacity: fixed: a warp role index +MMA_WARP = 7 # capacity: fixed: a warp role index +# (M <= 3) the early stages are issued by the idle MMA warp, not by warp 4 (routing's critical path) +E0_WARP = MMA_WARP # capacity: fixed: a warp role index +E0_BAR = 8 +E1_BAR = 9 +E0_OFFLOAD = EARLY0 and NRW0 <= 3 and bool(_cfg("e0_off", 1)) +EVICT_FIRST = 0x12F0000000000000 +_LOG2E = 1.4426950408889634 +LD_POW2 = ( + 1 if M_MAX == 1 else 2 if M_MAX == 2 else 4 if M_MAX <= 4 else 8 +) # tcgen05.ld .num covering M columns + +_FIXED_SMEM = ( + S1 * NUM_BYTES_B1 + + 2 * KT2 * NUM_BYTES_B2 + + 2 * KT2 * NUM_KB * 512 + + NB2_CAP * M_MAX * TOP_K * R2 * 4 + + 2 * M_MAX * C1 * R1 * 4 + + (M_MAX * (H // 16) if FUSED_ROUTER else 0) + + 6 * 1024 +) +# SFA_CP: the A scale words go by 4-byte cp.async straight into the instruction atom layout (no raw scale box, no +# splicer warp; the MMA warp tcgen05.cp's them), completing on the stage's full barrier: 18 KB stages instead of 26 KB +SFA_CP = bool(_cfg("sfa_cp", 0)) +_RAW_STAGE = 0 if SFA_CP else NUM_BYTES_SFA_RAW +_STAGE_SMEM = NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA +_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // _STAGE_SMEM +ACC1_COL = 0 +ACC2_COL = 2 * N1 +SFB1_COL = ACC2_COL + 2 * N2 +SFB2_COL = SFB1_COL + NUM_SFB1_ATOMS * 4 +SFA1_COL = SFB2_COL + 2 * KT2 * NUM_KB * 4 +STAGES = min(_cfg("stages", 12 if SFA_CP else 8), _STAGES_FIT, (512 - SFA1_COL) // 16) +TMEM_COLS = 512 +assert 1 <= M_MAX <= 8, M_MAX +assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) +assert M_MAX * H * 2 <= STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA), ( + "x staging must fit the ring" +) +# one shared-memory region for the ring and the FC1 / FC2 B operands (a shared CTA reuses all of it) +_OFF_SFAR = STAGES * NUM_BYTES_A +_OFF_SFA = _OFF_SFAR + STAGES * _RAW_STAGE +_OFF_B1 = _OFF_SFA + STAGES * NUM_BYTES_SFA +_OFF_B2 = _OFF_B1 + S1 * NUM_BYTES_B1 +_OFF_SFB2 = _OFF_B2 + 2 * KT2 * NUM_BYTES_B2 +BIG = _OFF_SFB2 + 2 * KT2 * NUM_KB * 512 +# shared CTA layout in BIG: weight buffers, x quarters (2), x scales [M, 192], gate / up [4, 16, M], h, h scales +# shared CTA layout in BIG: A ring (16 KB stages), raw scale atoms (4 per stage), spliced atoms, B (x: [stage][16 rows] +# [128 B]; h: [6][8][128 B]), x / h scale codes, gate_up partial exchange +_SB_A = 0 +_SB_RAW = None +_SH_FIX = ( + SH_GT * 16 * 128 + + 6 * 8 * 128 + + M_MAX * SH_KB + + M_MAX * SH_DKB + + 64 * 8 * 4 + + 2 * 32 * 8 * 4 + + 1024 +) +SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 +_SB_RAW = SH_NST * NUM_BYTES_A +_SB_SPL = _SB_RAW + SH_NST * 2048 +_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 +_SB_BH = _SB_BX + SH_GT * 16 * 128 +_SB_XC = _SB_BH + 6 * 8 * 128 +_SB_HC = _SB_XC + M_MAX * SH_KB +_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 +_SB_GU = _SB_XCH + 64 * 8 * 4 +_SB_END = _SB_GU + 2 * 32 * 8 * 4 +# shared CTA TMEM columns: gate_up accumulator (N 16), down accumulators (2 x N 8), A scales per stage, B scales +SH_ACC_GU = 0 +SH_ACC_DN = 16 +SH_SFA = 32 +SH_SFB_GU = SH_SFA + 4 * 8 +SH_SFB_DN = SH_SFB_GU + 4 * SH_GT # capacity: fixed: a kernel tile constant +SH_TMEM = 256 # capacity: fixed: TMEM columns +if SHARED: + assert FUSED_ROUTER, "shared expert: fused router only" + assert SH_STAGES - SH_NST <= 64, SH_STAGES + assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 24 <= SH_TMEM, (SH_NST, _SB_END, BIG) + + +# ============================================================================= +# Inline-PTX helpers. +# ============================================================================= +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- all-reduce push (the MoE boundary's multicast exchange, set A): +# arp = int64 [enable, mc, flags, rank, world, maxt] +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): + """st.global.b32 when pred != 0.""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _ar_push(arp, v, low, t, cl): + """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on + lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call + parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + world = arp.load(idx=4) + parity = cutlass.Int64( + _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) + ) + & cutlass.Int32(1) + ) + off = ( + parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +def _ar_parity(arp, t, cl): + """The call parity _ar_push reads for token t's channel cl (fixed between this kernel's grid dependency and the + next all-reduce, which this kernel releases at teardown): loaded ahead of the finalize (AR_PRE).""" + return _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cl // cutlass.Int32(H // 8)) + ) + + +def _ar_push_pre(ard, parity32, v, low, t, cl): + """_ar_push with the parity (_ar_parity) and the descriptor fields ard = (mc, rank, world, maxt) preloaded.""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + mc, rk, world, maxt = ard + parity = cutlass.Int64(parity32 & cutlass.Int32(1)) + off = ( + parity * maxt * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + rk) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(mc + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" + _llvm.inline_asm( + None, + [], + f"bar.arrive {bar_id}, {threads};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "st.relaxed.gpu.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int64(val)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _nanosleep(ns: int, *, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _gtime(*, loc=None, ip=None): + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): + if cutlass.const_expr(TIMING): + tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _sub_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_bits(x, *, loc=None, ip=None): + """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3_to_f32(b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(b)], + "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" + "cvt.f32.f16 $0, l;\n}", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x2(hi, lo, *, loc=None, ip=None): + """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """x rounded to bf16 (nearest even), as f32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(v, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(v)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16 (nearest even) bits of x in the low 16 bits.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ue8m0_rp(x, *, loc=None, ip=None): + """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3x2(hi, lo, *, loc=None, ip=None): + """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _min_xorsign_abs(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_full(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2_full(x, *, loc=None, ip=None): + """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): + """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" + _llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _mx_osc(amax): + """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" + code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) + cf = cutlass.Int32( + cutlass.select_( + code == cutlass.Int32(255), + cutlass.Int32(0x7FFFFFFF), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) + ), + ) + ).bitcast(cutlass.Float32) + return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) + + +def _mx_words(v, osc): + """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" + words = [] + for t in range(len(v) // 4): + lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) + hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) + words.append(lo | (hi << cutlass.Int32(16))) + return words + + +def _twiddle(score): + """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" + bits = ( + score.bitcast(cutlass.Int32) + if hasattr(score, "bitcast") + else cutlass.Float32(score).bitcast(cutlass.Int32) + ) + return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) + + +@cute.jit +def _wait(bar: cutlass.Array, parity): + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +def _tptr(base_row, col, dtype=cutlass.Int32): + return cutlass.inttoptr((base_row << 16) | col, 6, dtype) + + +@dsl_user_op +def _match_any(value, *, loc=None, ip=None): + """The mask of the warp's lanes holding the same 32-bit value.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(value)], + "match.any.sync.b32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@cute.jit +def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): + """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" + if cutlass.const_expr(HANG_DBG): + n = cutlass.Int32(0) + while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( + n < cutlass.Int32(1 << 22) + ): + n = n + cutlass.Int32(1) + if n >= cutlass.Int32(1 << 22): + bx, _, _ = cute.arch.block_idx() + w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tdbg.store( + cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w + ) + else: + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +@cute.jit +def _adv(stage, phase): + """The next ring position's (stage, phase).""" + stage = stage + cutlass.Int32(1) + wrap = stage == cutlass.Int32(STAGES) + return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( + wrap, phase ^ cutlass.Int32(1), phase + ) + + +@cute.jit +def _seg(i0, i1, k): + """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: + the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" + uf = i0 // cutlass.Int32(S1) + ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) + unit = uf + k - cutlass.Int32(1) + unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) + unit = cutlass.select_(ul == uf, uf, unit) + j0 = i0 - unit * cutlass.Int32(S1) + j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) + j1 = i1 - unit * cutlass.Int32(S1) + j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) + return unit, j0, j1 + + +@cute.jit +def _nsegs(i0, i1): + n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) + return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) + + +@cute.jit +def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): + """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile + 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" + q = lane // cutlass.Int32(QS) + tl = lane - q * cutlass.Int32(QS) + sf_ok = (lane < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) + sfv = [] + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + j = a // NUM_KB + kb = a % NUM_KB + wsf = src.load( + idx=cutlass.select_( + sf_ok, + tl * cutlass.Int32(H // 64) + + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) + + cutlass.Int32(kb), + cutlass.Int32(0), + ) + ) + sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + prims.tcgen05_st( + "32x32b", + _tptr(base_row + warp * cutlass.Int32(32), base_col + SFB1_COL + a * 4), + sfv[a], + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(b1_ready) + + +@cute.jit +def _b1_addr(t, kb16): + """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" + kt = kb16 >> cutlass.Int32(4) + j = kt >> cutlass.Int32(2) + n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t + cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) + return ( + j * cutlass.Int32(NUM_BYTES_B1) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) + ) // cutlass.Int32(4) + + +@cute.jit +def _build_b1( + tidx, + warp, + lane, + base_row, + base_col, + x32: cutlass.Array, + xsf32: cutlass.Array, + sB1: cutlass.Array, + b1_ready: cutlass.Array, +): + """B (the M tokens' whole K) into the resident tiles from the NVFP4 activation, its scales to TMEM.""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + for r in cutlass.range_constexpr((M_MAX * (H // 16) + EPI_THREADS - 1) // EPI_THREADS): + b = tidx + cutlass.Int32(r * EPI_THREADS) + if b < cutlass.Int32(M_MAX * (H // 16)): + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + v2 = x32.load(idx=b * cutlass.Int32(2), vector_size=2, alignment=8) + sB1_32.store( + (cutlass.Int32(v2[0]), cutlass.Int32(v2[1])), idx=_b1_addr(t, kb), alignment=8 + ) + _sfb1_to_tmem(warp, lane, base_row, base_col, xsf32, b1_ready) + + +@cute.jit +def _build_b1_fused( + tidx, + warp, + lane, + base_row, + base_col, + xb32: cutlass.Array, + sfsc: cutlass.Array, + sB1: cutlass.Array, + s_xsf: cutlass.Array, + b1_ready: cutlass.Array, +): + """_build_b1 from the bf16 hidden states, quantized exactly as fp4_quantize(x, sf_scale, 16).""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + s_xsf8 = cutlass.Array( + s_xsf.data_ptr(0), shape=(M_MAX * (H // 16),), dtype=cutlass.Int8, alignment=16 + ) + sc = cutlass.Float32(sfsc.load(idx=0)) + rsc = _rcp(sc) + rsix = _rcp(cutlass.Float32(6.0)) + NBLK = M_MAX * (H // 16) // EPI_THREADS + QB = ( + 3 if NBLK % 3 == 0 else (2 if NBLK % 2 == 0 else 1) + ) # blocks whose loads are in flight together + # a rolled loop over groups of QB blocks (a fully unrolled one is ~50 KB of SASS at M = 8: instruction fetch bound) + for rg in cutlass.range(NBLK // QB, unroll=1): + wall = [] + for rr in cutlass.range_constexpr(QB): + bq = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + for h in cutlass.range_constexpr(2): + wall.append( + xb32.load( + idx=bq * cutlass.Int32(8) + cutlass.Int32(4 * h), + vector_size=4, + alignment=16, + ) + ) + for rr in cutlass.range_constexpr(QB): + b = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + w = [] + for h in cutlass.range_constexpr(2): + v4 = wall[rr * 2 + h] + for z in cutlass.range_constexpr(4): + w.append(cutlass.Int32(v4[z])) + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), + _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), + cutlass.Float32(0.0), + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + sB1_32.store((lo, hi), idx=_b1_addr(t, kb), alignment=8) + s_xsf8.store(cutlass.Int8(sfb), idx=b) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + + +@cute.jit +def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): + """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, + gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 + and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" + neg_max = cutlass.Float32(-3.40282346638528859812e38) + gc = cutlass.select_(gate_raw > lim, lim, gate_raw) + gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) + uc = cutlass.select_(xv > lim, lim, xv) + uc = cutlass.select_(xv < -lim, -lim, uc) + c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) + sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) + gs = _mul_rn(gc, sgm) + h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) + h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) + amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) + for off in (16, 8, 4, 2, 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) + sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) + sff = _e4m3_to_f32(sfb) + osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) + return _mul_rn(_mul_rn(h, osc), c), sfb + + +@cute.jit +def _plan_range(s_plan: cutlass.Array, p): + w = s_plan.load(idx=p) + return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) + + +@cute.jit +def _sh_issue( + lane, + sid, + g, + s_big: cutlass.Array, + sh_full: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, +): + """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both + slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its + raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile + (g - SH_GT) % 6.""" + st = g % cutlass.Int32(SH_NST) + bar = sh_full.subview(st) + if cutlass.const_expr(SH_PRETILE): + if lane == 0: + src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) + prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + gusf.subview(src + cutlass.Int32(16384)), + bar, + 512, + ) + cute.arch.sync_warp() + else: + if g < cutlass.Int32(SH_GT): + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) + cute.arch.sync_warp() + if lane < 4: + q = lane >> cutlass.Int32(1) + u = lane & cutlass.Int32(1) + row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) + kt = q * cutlass.Int32(SH_GT) + g + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), + tma_gu, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + cutlass.Int32(_SB_RAW) + + st * cutlass.Int32(2048) + + lane * cutlass.Int32(512) + ), + gusf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), + tma_dn, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), + dnsf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + cute.arch.sync_warp() + + +@dsl_user_op +def _red_release_add(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None): + _asm(None, [], "fence.proxy.async.global;", "", True, loc, ip) + + +@cute.jit +def _sfa_to_tmem(sSFA: cutlass.Array, stage, base_row, base_col): + """(SFA_CP, elected MMA thread) the stage's 4 scale atoms smem -> TMEM (what the splicer warp does otherwise).""" + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SFA1_COL + stage * cutlass.Int32(16) + kb * 4), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + + +@cute.jit +def _sf1_cp( + lane, stage, sf_line, c, sSFA: cutlass.Array, s1w: cutlass.Array, ab_full: cutlass.Array +): + """(SFA_CP) FC1 stage scale atoms: word [r][q] of atom kb <- raw atom 4 q + kb (raw line sf_line = 4 x atom) + word [r][c % 4] (what the splicer copies out of the raw box), 4-byte cp.async by the warp's 32 lanes; their + arrivals (noinc) complete the stage's full barrier together with the A boxes' bytes.""" + src0 = ( + (sf_line >> cutlass.Int32(2)) * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (c & cutlass.Int32(3)) + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for q in cutlass.range_constexpr(C1): + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + q)), + s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _sf2_cp( + lane, + stage, + g, + blk, + t2: cutlass.Constexpr, + n_here, + s_el: cutlass.Array, + sSFA: cutlass.Array, + s2w: cutlass.Array, + ab_full: cutlass.Array, +): + """(SFA_CP) FC2 k-tile t2 scale atoms: word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) of slot + jj's expert, word [r][blk % 4].""" + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + src0 = ( + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (blk & cutlass.Int32(3)) + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + jj)), + s2w.subview(src0 + cutlass.Int32(kb * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _shared_cta( + tidx, + warp, + lane, + sid, + s_big: cutlass.Array, + sh_bar: cutlass.Array, + epochs: cutlass.Array, + tdbg: cutlass.Array, + bx, + xb32: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + tmem_ptr_i32: cutlass.Array, + gu8: cutlass.Array, + dn8: cutlass.Array, +): + """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one + 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles + of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" + big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) + bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) + big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) + sh_full = sh_bar # [0, SH_NST): stage full + sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) + b_gu_rdy = sh_bar.subview(16) + b_dn_rdy = sh_bar.subview(17) + acc_gu = sh_bar.subview(18) + acc_dn = sh_bar.subview(19) + if warp == 0: + prims.prefetch_tensormap(tma_gu) + prims.prefetch_tensormap(tma_dn) + if lane < SH_NST: + prims.mbarrier_init(sh_full.subview(lane), 1) + prims.mbarrier_init( + sh_empty.subview(lane), 33 + ) # the MMAs' commit + each splicer lane's reads done + if lane == 0: + prims.mbarrier_init(b_gu_rdy, 1) + prims.mbarrier_init(b_dn_rdy, 1) + prims.mbarrier_init(acc_gu, 1) + prims.mbarrier_init(acc_dn, 1) + if warp == 1: + prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency + _sh_issue(lane, sid, cutlass.Int32(g), s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) + if cutlass.const_expr(SH_PRETILE): + # this CTA's stages beyond the ring, 4 per lane + for r in cutlass.range_constexpr(2): + pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(2) + cutlass.Int32(r) + if pg < cutlass.Int32(SH_STAGES): + _l2_prefetch( + gusf.subview( + (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) + ) + .data_ptr() + .toint(), + SH_STAGE_BYTES, + ) + else: + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + NQ = THREADS - 64 # activation quant threads (warps 2 ..) + if warp == 0: + for g in cutlass.range(SH_NST, SH_STAGES, unroll=1): + _wait( + sh_empty.subview(g % cutlass.Int32(SH_NST)), + ((g // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), + ) + _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + elif warp == 1: + # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down + idesc_gu = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=16, + m_dim=128, + ) + idesc_dn = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=8, + m_dim=128, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + desc_a0 = prims.Tcgen05SmemDesc.build( + s_big.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bx = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BX), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bh = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BH), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + sq = sid & cutlass.Int32(3) + _wait(b_gu_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for g in cutlass.range(SH_STAGES, unroll=1): + st = g % cutlass.Int32(SH_NST) + if g == cutlass.Int32(SH_GT): + _wait(b_dn_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _wait(sh_full.subview(st), (g // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) + raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) + if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): + # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 + w4 = [ + big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) + for qq in range(4) + ] + big32.store( + ( + cutlass.Int32(w4[0]), + cutlass.Int32(w4[1]), + cutlass.Int32(w4[2]), + cutlass.Int32(w4[3]), + ), + idx=cutlass.Int32(_SB_SPL // 4) + + st * cutlass.Int32(128) + + lane * cutlass.Int32(4), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + prims.mbarrier_arrive( + sh_empty.subview(st) + ) # this lane's generic reads of the slot are done + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + if g < cutlass.Int32(SH_GT): + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bx + cutlass.Int64(g * cutlass.Int32(2048 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), + da + 2 * kb, + db + 2 * kb, + idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_GT - 1): + prims.tcgen05_commit(acc_gu) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview( + cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) + + st * cutlass.Int32(512 if SH_PRETILE else 2048) + ), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bh + cutlass.Int64(kt * cutlass.Int32(1024 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr( + tcol + SH_ACC_DN + jt * cutlass.Int32(8), 6, cutlass.Float32 + ), + da + 2 * kb, + db + 2 * kb, + idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_STAGES - 1): + prims.tcgen05_commit(acc_dn) + cute.arch.sync_warp() + if warp >= 2: + # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = + # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes + qt = tidx - cutlass.Int32(64) + for r in cutlass.range_constexpr((M_MAX * SH_KB + NQ - 1) // NQ): + q_ = qt + cutlass.Int32(r * NQ) + if q_ < cutlass.Int32(M_MAX * SH_KB): + t = q_ // cutlass.Int32(SH_KB) + b = q_ - t * cutlass.Int32(SH_KB) + wv = [ + xb32.load( + idx=t * cutlass.Int32(H // 2) + + b * cutlass.Int32(16) + + cutlass.Int32(4 * z), + vector_size=4, + alignment=16, + ) + for z in range(4) + ] + v = [] + for z in cutlass.range_constexpr(4): + for zz in cutlass.range_constexpr(4): + w = cutlass.Int32(wv[z][zz]) + v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) + v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(v[z])) + code, osc = _mx_osc(amax) + words = _mx_words(v, osc) + kt = b >> cutlass.Int32(2) + sl = kt // cutlass.Int32(SH_GT) + it = kt - sl * cutlass.Int32(SH_GT) + n = sl * cutlass.Int32(8) + t + c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BX) + it * cutlass.Int32(2048) + n * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) + if warp >= 4: + if warp < 8: + qd = warp - cutlass.Int32(4) + # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of k-tile SH_GT q + i + n = lane + t = n & cutlass.Int32(7) + ok = (n < cutlass.Int32(16)) & (t < cutlass.Int32(M_MAX)) + for i in cutlass.range_constexpr(SH_GT): + kt = (n >> cutlass.Int32(3)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) + wsf = big32.load( + idx=( + cutlass.Int32(_SB_XC) + + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) + + kt * cutlass.Int32(4) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), + cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + et = tidx - cutlass.Int32(128) + if et == 0: + prims.mbarrier_arrive(b_gu_rdy) + _stamp(tdbg, bx, 54) + # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t at column 8 q + t + _wait(acc_gu, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=16, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(2): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[8 + tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd < cutlass.Int32(2): + # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag + for r in cutlass.range_constexpr((32 * M_MAX + 127) // 128): + pq = et + cutlass.Int32(r * 128) + if pq < cutlass.Int32(32 * M_MAX): + p = pq // cutlass.Int32(M_MAX) + tt = pq - p * cutlass.Int32(M_MAX) + ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(8) + tt) + up = bigf.load( + idx=cutlass.Int32(_SB_GU // 4) + + (cutlass.Int32(32) + p) * cutlass.Int32(8) + + tt + ) + q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) + q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) + q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) + q_e2 = _ex2_full( + _mul_rn( + _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), + cutlass.Float32(1.4426950216293334961), + ) + ) + q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) + q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) + _st_relaxed_u32( + hx.subview(tt * cutlass.Int32(SH_I) + sid * cutlass.Int32(32) + p) + .data_ptr() + .toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), + ) + if et == 0: + _stamp(tdbg, bx, 50) + # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes + for rr in cutlass.range_constexpr((M_MAX * SH_DKB + 127) // 128): + q_ = et + cutlass.Int32(rr * 128) + if q_ < cutlass.Int32(M_MAX * SH_DKB): + t = q_ // cutlass.Int32(SH_DKB) + kb = q_ - t * cutlass.Int32(SH_DKB) + base = t * cutlass.Int32(SH_I) + kb * cutlass.Int32(32) + hv = [cutlass.Int32(0) for _ in range(32)] + pend = cutlass.Int32(1) + while pend != cutlass.Int32(0): + for q4 in cutlass.range_constexpr(8): + v4 = prims.load_ext( + hx.subview(base + cutlass.Int32(4 * q4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for z in cutlass.range_constexpr(4): + hv[4 * q4 + z] = cutlass.Int32(v4[z]) + okw = cutlass.Boolean(True) + for z in cutlass.range_constexpr(32): + okw = okw & ( + ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 + ) + pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) + vals = [ + (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) + ] + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) + code, osc = _mx_osc(amax) + words = _mx_words(vals, osc) + kt = kb >> cutlass.Int32(2) + c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BH) + kt * cutlass.Int32(1024) + t * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + (((c0 + cutlass.Int32(1)) ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + okh = lane < cutlass.Int32(M_MAX) + for kt in cutlass.range_constexpr(6): + wsf = big32.load( + idx=( + cutlass.Int32(_SB_HC) + + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) + + cutlass.Int32(4 * kt) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), + cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if et == 0: + prims.mbarrier_arrive(b_dn_rdy) + _stamp(tdbg, bx, 51) + # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l + _wait(acc_dn, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jt in cutlass.range_constexpr(2): + dv = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_DN + jt * 8, + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + rowo = ( + sid * cutlass.Int32(256) + + cutlass.Int32(128 * jt) + + qd * cutlass.Int32(32) + + lane + ) + for tt in cutlass.range_constexpr(M_MAX): + _st_relaxed_u32( + dnx.subview(cutlass.Int32(tt * H) + rowo).data_ptr().toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), + ) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if et == 0: + _stamp(tdbg, bx, 52) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if warp == 1: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) + if tidx == 0: + epochs.store(ep + cutlass.Int32(1), idx=bx) + + +@cute.kernel +def m3_moe_kernel( + tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], # w3_w1_weight_scale as 128-byte lines + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], # w2_weight_scale as 128-byte lines + tma_gu_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: gate_up e4m3 [1536, 6144], box 128 B x 32 rows + tma_dn_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: down e4m3 [6144, 768], box 128 B x 128 rows + logits: cutlass.Array, # f32 [M, 128] (not FUSED_ROUTER) + bias: cutlass.Array, # f32 [128] + x32: cutlass.Array, # int32 words of x_fp4 [M, 3072 B] (not FUSED_ROUTER) + xsf32: cutlass.Array, # int32 words of x_sf [M, 384 B] (not FUSED_ROUTER) + escal: cutlass.Array, # f32 [6, 128]: fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha + plan: cutlass.Array, # int32 [4 M + 1, NCTA, NG_CAP + 1] by G: per group i0 | i1 << 16, then b0 | nb << 16 + partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums + hbuf32: cutlass.Array, # int32 [G_CAP * M, 48, 4]: tagged intermediate records + epochs: cutlass.Array, # int32 [NCTA] + out: cutlass.Array, # bf16 [M, 6144] + route_dbg: cutlass.Array, # int32 [M, 8]: ids, weight bits (CTA 0) + tdbg: cutlass.Array, # int64 [NCTA, 32] timing stamps (TIMING builds) + xb32: cutlass.Array, # FUSED_ROUTER: int32 words (bf16 pairs) of the hidden states [M, 6144] + rw: cutlass.Array, # FUSED_ROUTER: f32 router weight [128, 6144] + sfsc: cutlass.Array, # FUSED_ROUTER: f32 [1] fc31_input_scale + lgbuf: cutlass.Array, # FUSED_ROUTER: int64 [LG_REPL, LG_STRIDE] logits: tag << 32 | f32 bits + guw: cutlass.Array, # SHARED: uint8 gate_up e4m3 [1536, 6144] + gusf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + dnw: cutlass.Array, # SHARED: uint8 down e4m3 [6144, 768] + dnsf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + hx: cutlass.Array, # SHARED: int32 [M, 768]: h as (call tag << 16 | bf16), tag 0 = never written + dnx: cutlass.Array, # SHARED: int32 [M, 6144]: y_sh as (call tag << 16 | bf16), tag 0 = never written + arp: cutlass.Array, # int64 [8]: all-reduce push (enable, mc, flags, rank, world, maxt); enable 0 = off + s1w: cutlass.Array, # SFA_CP: int32 words of w3_w1_weight_scale + s2w: cutlass.Array, # SFA_CP: int32 words of w2_weight_scale + img32: cutlass.Array, # B2IMG: int32 [NG_CAP, IMG_BYTES / 4] per-group FC2 B image (tile + scale atoms) + b2cnt: cutlass.Array, # B2IMG: int32 [NG_CAP * 8 + 8] per-group unit counters (index 8 g) +): + tidx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + if tidx == 0: + _stamp(tdbg, bx, 0) + if cutlass.const_expr(TIMING): + if tidx == 0: + tdbg.store( + cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) + ) + + s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) + sA = cutlass.Array( + s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 + ) + sSFAraw = cutlass.Array( + s_big.data_ptr(_OFF_SFAR), + shape=(max(1, _RAW_STAGE // 4) * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sSFA = cutlass.Array( + s_big.data_ptr(_OFF_SFA), + shape=(NUM_WORDS_SFA * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sB1 = cutlass.Array( + s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 + ) + sB2 = cutlass.Array( + s_big.data_ptr(_OFF_B2), shape=(2 * KT2 * NUM_BYTES_B2,), dtype=cutlass.Int8, alignment=1024 + ) + sSFB2 = cutlass.Array( + s_big.data_ptr(_OFF_SFB2), + shape=(2 * KT2 * NUM_KB * 128,), + dtype=cutlass.Int32, + alignment=1024, + ) + ys = cutlass.Array( + cutlass.Float32, NB2_CAP * M_MAX * TOP_K * R2, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_part = cutlass.Array( + cutlass.Float32, 2 * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_red = cutlass.Array(cutlass.Float32, 9 * M_MAX, space=cutlass.AddressSpace.smem, alignment=16) + s_xsf = cutlass.Array( + cutlass.Int32, + M_MAX * (H // 64) if FUSED_ROUTER else 4, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_free = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + # this CTA's FC1 share of group g is out: one barrier per group parity (a single barrier could complete twice + # while warp 5 still works on an earlier group, and its parity wait would then miss them) + fc1_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) + b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + x_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # x landed, x consumed + img_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # (B2IMG) image landed + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + s_el = cutlass.Array( + cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16 + ) # slots, then G + s_kpos = cutlass.Array( + cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_w = cutlass.Array( + cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_ids = cutlass.Array( + cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + # this CTA's plan rows for every G (static data: loaded before the grid dependency) + s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_plan_all = cutlass.Array( + cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 + ) + s_esc = cutlass.Array( + cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32 + ) # slot s at 8 s + + is_routed = cutlass.Boolean(True) + if cutlass.const_expr(SHARED): + sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) + is_routed = bx < cutlass.Int32(NCTA) + if bx >= cutlass.Int32(NCTA): + _shared_cta( + tidx, + warp, + lane, + bx - cutlass.Int32(NCTA), + s_big, + sh_full, + epochs, + tdbg, + bx, + xb32, + tma_gu_desc.get_ptr(), + tma_dn_desc.get_ptr(), + gusf, + dnsf, + hx, + dnx, + arp, + tmem_ptr_i32, + guw, + dnw, + ) + if is_routed: + # ---- before the grid dependency + if warp == 0: + if tidx < STAGES: + prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) + prims.mbarrier_init(ab_empty.subview(tidx), 1) + prims.mbarrier_init(sfa_ready.subview(tidx), 1) + if tidx < 2: + prims.mbarrier_init(acc_full.subview(tidx), 1) + prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(acc2_full.subview(tidx), 1) + prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(b2_ready.subview(tidx), 1) + prims.mbarrier_init(b2_free.subview(tidx), EPI_THREADS) + if cutlass.const_expr(B2IMG): + prims.mbarrier_init(img_bar.subview(tidx), 1) + if tidx < 4: + prims.mbarrier_init(fc1_done.subview(tidx), EPI_THREADS) + if tidx == 0: + prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) + if cutlass.const_expr(XSTAGE): + prims.mbarrier_init(x_bar.subview(0), 1) + prims.mbarrier_init(x_bar.subview(1), EPI_THREADS) + if warp == MMA_WARP: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == W_WARP: + prims.prefetch_tensormap(tma_a1_desc.get_ptr()) + prims.prefetch_tensormap(tma_a2_desc.get_ptr()) + prims.prefetch_tensormap(tma_s1_desc.get_ptr()) + prims.prefetch_tensormap(tma_s2_desc.get_ptr()) + for i in cutlass.range_constexpr(E_LOCAL // 32): + s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) + for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): + q = lane + cutlass.Int32(32 * i) + if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): + gq = q // cutlass.Int32(NG_CAP + 1) + s_plan_all.store( + plan.load( + idx=(gq * cutlass.Int32(NCTA) + bx) * cutlass.Int32(NG_CAP + 1) + + q + - gq * cutlass.Int32(NG_CAP + 1) + ), + idx=q, + ) + rwv = [] + if cutlass.const_expr(FUSED_ROUTER): + tg = cutlass.select_(tidx < cutlass.Int32(GEMV_THREADS), tidx, cutlass.Int32(0)) + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + rbx = cutlass.select_(eg < cutlass.Int32(NROUTE), eg, cutlass.Int32(0)) + rwv.append( + [ + rw.load( + idx=rbx * cutlass.Int32(H) + + tg * cutlass.Int32(4) + + cutlass.Int32(1024 * j), + vector_size=4, + alignment=16, + ) + for j in range(H // 1024) + ] + ) + if cutlass.const_expr(TIMING): + if tidx == 0: + s_el.store(cutlass.Int32(-7), idx=G_CAP) + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(XSTAGE): + if warp == 0: + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_bar.subview(0), M_MAX * H * 2) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(lane * cutlass.Int32(H * 2)), + xb32.subview(lane * cutlass.Int32(H // 2)), + x_bar.subview(0), + H * 2, + ) + ar_on = arp.load(idx=0) != cutlass.Int64(0) + ar_off = arp.load(idx=0) == cutlass.Int64(0) + if cutlass.const_expr(not LATE_TRIGGER): + if ar_off: + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + + ep = epochs.load(idx=bx) + tagv = ((ep & cutlass.Int32(0x7F)) | cutlass.Int32(0x80)) << cutlass.Int32(24) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + if cutlass.const_expr(FUSED_ROUTER): + # ---- router GEMV (experts bx, bx + NCTA, .., threads 0..255), m3_router_gemv's order exactly + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + if (eg < cutlass.Int32(NROUTE)) & (tidx < cutlass.Int32(GEMV_THREADS)): + accs = [cutlass.Float32(0.0) for _ in range(M_MAX)] + xsrc = xb32 + if cutlass.const_expr(GEMV_SMEM and XSTAGE): + _wait(x_bar.subview(0), 0) + xsrc = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + for j in cutlass.range_constexpr(H // 1024): + for t in cutlass.range_constexpr(M_MAX): + xw = xsrc.load( + idx=cutlass.Int32(t * (H // 2) + 512 * j) + tidx * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + x0 = cutlass.Int32(xw[0]) + x1 = cutlass.Int32(xw[1]) + a = accs[t] + a = _fma_rn( + (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][0]), + a, + ) + a = _fma_rn( + (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][1]), + a, + ) + a = _fma_rn( + (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][2]), + a, + ) + a = _fma_rn( + (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][3]), + a, + ) + accs[t] = a + for t in cutlass.range_constexpr(M_MAX): + v = accs[t] + for off in (16, 8, 4, 2, 1): + v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) + if lane == 0: + s_red.store(v, idx=warp * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < M_MAX: + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX) + tidx)) + s_red.store(sacc, idx=cutlass.Int32(8 * M_MAX) + tidx) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < LG_REPL: + for t in cutlass.range_constexpr(M_MAX): + lw = ltag | ( + cutlass.Int64( + cutlass.Float32(s_red.load(idx=8 * M_MAX + t)).bitcast( + cutlass.Int32 + ) + ) + & cutlass.Int64(0xFFFFFFFF) + ) + _st_relaxed_u64( + lgbuf.subview( + tidx * cutlass.Int32(LG_STRIDE) + + cutlass.Int32(t * E_LOCAL) + + eg + ) + .data_ptr() + .toint(), + lw, + ) + if cutlass.const_expr(NGEMV > 1): + cute.arch.sync_warp() + cute.arch.barrier( + barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS + ) # s_red reused next + if tidx == 0: + _stamp(tdbg, bx, 21) + + # ---- routing (warps 4..7, token t on warp 4 + t % 4); the epilogue warps build B meanwhile + if ((warp >= W_WARP) & (warp < cutlass.Int32(W_WARP + NRW0))) | ( + (warp >= cutlass.Int32(8)) & (warp < cutlass.Int32(8 + RW_EXTRA)) + ): + rw_i = cutlass.select_( + warp < cutlass.Int32(8), warp - cutlass.Int32(W_WARP), warp - cutlass.Int32(4) + ) + # ILP: poll every token's logits first, then run the independent top-k chains back to back + lgv_all = [] + t_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = rw_i + cutlass.Int32(NRW * tr) + if cutlass.const_expr(M_MAX % NRW != 0): + t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) + t_all.append(t) + lgv = [] + if cutlass.const_expr(FUSED_ROUTER): + want = ltag >> cutlass.Int64(32) + rbase = ( + (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) + + t * cutlass.Int32(E_LOCAL) + + lane + ) + done = cutlass.Int32(0) + res = [cutlass.Int64(0) for _ in range(4)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(4): + res[z] = _ld_relaxed_u64( + lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() + ) + ok = ok & ((res[z] >> cutlass.Int64(32)) == want) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32( + -1 + ): + done = cutlass.Int32(1) + else: + _nanosleep(POLL_NS) + for z in cutlass.range_constexpr(4): + lgv.append(cutlass.Int32(res[z]).bitcast(cutlass.Float32)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 22) + else: + for i in cutlass.range_constexpr(4): + lgv.append( + cutlass.Float32( + logits.load( + idx=t * cutlass.Int32(E_LOCAL) + cutlass.Int32(i * 32) + lane + ) + ) + ) + lgv_all.append(lgv) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 61) + res_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = t_all[tr] + lgv = lgv_all[tr] + own = (rw_i + cutlass.Int32(NRW * tr)) < cutlass.Int32(M_MAX) + # top-4 (keys sigmoid + bias, ties to the lower id) + keys = [] + for i in cutlass.range_constexpr(4): + sg = _fma_rn( + cutlass.Float32(0.5), + cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), + cutlass.Float32(0.5), + ) + keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) + hi = [_twiddle(keys[i]) for i in range(4)] + taken = [cutlass.Boolean(False) for _ in range(4)] + ids_t = [] + kt_t = [] + for k in cutlass.range_constexpr(TOP_K): + bh = cutlass.Int32(0) + bl = cutlass.Int32(0) + for i in cutlass.range_constexpr(4): + lo_i = cutlass.Int32(65535 - i * 32) - lane + hv = cutlass.select_(taken[i], cutlass.Int32(0), hi[i]) + lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) + better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( + (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) + ) + bh = cutlass.select_(better, hv, bh) + bl = cutlass.select_(better, lv, bl) + mh = _redux_max_u32(bh) + ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) + e_k = cutlass.Int32(65535) - ml + for i in cutlass.range_constexpr(4): + taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) + ids_t.append(e_k) + kv = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(4): + kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) + kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) + sig = cutlass.Float32(0.0) + e_l = cutlass.Int32(0) + for k in cutlass.range_constexpr(TOP_K): + sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) + sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) + e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) + total = sig + for off in (16, 8, 4, 2, 1): + total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) + wv = _bf16_round( + _div_rn( + _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), + _add_rn(total, cutlass.Float32(SUM_EPS)), + ) + ) + res_all.append((t, own, wv, e_l)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 62) + for tr in cutlass.range_constexpr( + (M_MAX + NRW - 1) // NRW + ): # stores after every chain (one branch-free block) + t, own, wv, e_l = res_all[tr] + if (lane < TOP_K) & own: + s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) + s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) + if bx == 0: + route_dbg.store( + wv.bitcast(cutlass.Int32), + idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, + ) + cute.arch.sync_warp() + if cutlass.const_expr(EARLY0 and not E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + _bar_arrive( + E0_BAR, 64 + ) # token 0's experts are in s_ids: warp 7 issues the early stages + if cutlass.const_expr(B1_AFTER_TOPK): + _bar_arrive(TOPK2_BAR, 128 + 32 * NRW) + cute.arch.barrier(barrier_id=TOPK_BAR, number_of_threads=32 * NRW) + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(E0_WARP): + cute.arch.barrier(barrier_id=E0_BAR, number_of_threads=64) + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + _bar_arrive(E1_BAR, 64) + if warp == W_WARP: + if lane == 0: + _stamp(tdbg, bx, 20) + # slots: lane l < 4 M is (token l / 4, top-k l % 4); the distinct experts in lane order (match.any) + valid = lane < cutlass.Int32(TOP_K * M_MAX) + my_e = cutlass.select_( + valid, + s_ids.load(idx=cutlass.select_(valid, lane, cutlass.Int32(0))), + cutlass.Int32(-1) - lane, + ) + same = _match_any(my_e) + leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) + first = valid & (lane == leader) + bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) + rank = cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) + slot_l = cute.arch.shuffle_sync(rank, leader) + for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): + s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + cute.arch.sync_warp() + if first: + s_el.store(my_e, idx=rank) + if valid: + s_kpos.store( + lane % cutlass.Int32(TOP_K), + idx=slot_l * cutlass.Int32(M_MAX) + lane // cutlass.Int32(TOP_K), + ) + if bx == 0: + route_dbg.store( + my_e, + idx=(lane // cutlass.Int32(TOP_K)) * cutlass.Int32(8) + + lane % cutlass.Int32(TOP_K), + ) + gcnt = cute.arch.popc(bal) + if lane == 0: + s_el.store(gcnt, idx=G_CAP) + if lane < cutlass.Int32(NG_CAP + 1): # the plan row of this G (preloaded) + s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) + cute.arch.sync_warp() + if cutlass.const_expr(B2IMG): + # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) + if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): + nh = gcnt - lane * cutlass.Int32(4) + nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) + nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) + pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) + if pad > cutlass.Int32(0): + _red_release_add( + b2cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad + ) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 24) + cute.arch.sync_warp() + _bar_arrive(ROUTE_BAR, THREADS) + if warp < 4: + if cutlass.const_expr(B1_AFTER_TOPK): + cute.arch.barrier(barrier_id=TOPK2_BAR, number_of_threads=128 + 32 * NRW) + if cutlass.const_expr(FUSED_ROUTER): + if tidx == 0: + _stamp(tdbg, bx, 57) + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 58) + xs32 = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xs32, sfsc, sB1, s_xsf, b1_ready + ) + prims.mbarrier_arrive(x_bar.subview(1)) # the ring is free for the weights + else: + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xb32, sfsc, sB1, s_xsf, b1_ready + ) + if tidx == 0: + _stamp(tdbg, bx, 23) + else: + _build_b1(tidx, warp, lane, base_row, base_col, x32, xsf32, sB1, b1_ready) + if warp != W_WARP: # one sync site for every thread but warp 4's (which arrived) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ROUTE_BAR, number_of_threads=THREADS) + if tidx == 0: + _stamp(tdbg, bx, 2) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(s_el.load(idx=G_CAP)), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(28), + ) + + g_n = s_el.load(idx=G_CAP) + ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) + pr2 = s_plan.load(idx=NG_CAP) + blk0 = pr2 & cutlass.Int32(0xFFFF) + nb = pr2 >> cutlass.Int32(16) + + # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG + if warp == W_WARP: + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) # the staged hidden states are quantized + stage = cutlass.Int32(0) + phase = cutlass.Int32(1) + skip = cutlass.Int32(0) + if cutlass.const_expr(E0_OFFLOAD): + cute.arch.barrier(barrier_id=E1_BAR, number_of_threads=64) + if cutlass.const_expr(EARLY0): + skip = s_misc.load(idx=0) # group-0 stages already in the ring + stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) + phase = cutlass.select_( + skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) + ) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if skip > cutlass.Int32(0): + skip = skip - cutlass.Int32(1) + else: + _waitd(ab_empty.subview(stage), phase, tdbg, 1) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + ), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, stage, sf_line, c, sSFA, s1w, ab_full) + stage, phase = _adv(stage, phase) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + _waitd(ab_empty.subview(stage), phase, tdbg, 2) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + n_here + * cutlass.Int32( + NUM_BYTES_BOX1 + if SFA_CP + else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4 + ), + ) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(jj * NUM_BYTES_BOX1) + ), + tma_a2_desc.get_ptr(), + ( + cutlass.Int32(t2 * KT_BYTES), + blk * cutlass.Int32(R2), + e, + ), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) + ), + tma_s2_desc.get_ptr(), + ( + cutlass.Int32(0), + ( + ( + e * cutlass.Int32(SF2_RB) + + (blk >> cutlass.Int32(2)) + ) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(4), + ), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf2_cp(lane, stage, g, blk, t2, n_here, s_el, sSFA, s2w, ab_full) + stage, phase = _adv(stage, phase) + if lane == 0: + _stamp(tdbg, bx, 3) + + # ---- scale splicer (warp 6): every stage, atom kb word [r][q] = raw atom 4 q + kb word [r][m1] + if (warp == S_WARP) & cutlass.Boolean(not SFA_CP): # (SFA_CP: no splicer) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + i0 = cutlass.Int32(0) + i1 = cutlass.Int32(0) + for part in cutlass.range_constexpr(2): + cnt = cutlass.Int32(0) + m1 = cutlass.Int32(0) + nk = cutlass.Int32(0) + if cutlass.const_expr(part == 0): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + nk = _nsegs(i0, i1) + else: + if p >= cutlass.Int32(LAG): + nk = nb + for k in range(nk): + if cutlass.const_expr(part == 0): + ug, j0, j1 = _seg(i0, i1, k) + m1 = ( + ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + ) & cutlass.Int32(3) + cnt = j1 - j0 + else: + m1 = (blk0 + k) & cutlass.Int32(3) + cnt = cutlass.Int32(KT2) + for _j in cutlass.range(cnt, unroll=1): + _waitd(ab_full.subview(stage), phase, tdbg, 3) + raw0 = ( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + lane * cutlass.Int32(4) + + m1 + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for kb in cutlass.range_constexpr(NUM_KB): + w0 = sSFAraw.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) + w1 = sSFAraw.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) + w2 = sSFAraw.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) + w3 = sSFAraw.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) + sSFA.store( + ( + cutlass.Int32(w0), + cutlass.Int32(w1), + cutlass.Int32(w2), + cutlass.Int32(w3), + ), + idx=dst0 + cutlass.Int32(kb * 128), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + prims.tcgen05_commit(sfa_ready.subview(stage)) + stage, phase = _adv(stage, phase) + + if cutlass.const_expr(B2IMG): + # ---- (B2IMG) group releases (warp 6): this CTA's finished units of group p, one release add per group + if warp == cutlass.Int32(S_WARP): + for p in range(ng): + i0, i1 = _plan_range(s_plan, p) + n_done = cutlass.Int32(0) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + n_done = n_done + cutlass.select_( + (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), + cutlass.Int32(1), + cutlass.Int32(0), + ) + _wait( + fc1_done.subview(p & cutlass.Int32(3)), + (p >> cutlass.Int32(2)) & cutlass.Int32(1), + ) + if cutlass.const_expr(B2_DLY > 0): + t0d = _gtime() + dl = cutlass.Int64( + ( + (bx * cutlass.Int32(37) + p * cutlass.Int32(11) + ep) + & cutlass.Int32(63) + ) + * cutlass.Int32(B2_DLY // 64) + ) + while _gtime() - t0d < dl: + pass + if (lane == 0) & (n_done > cutlass.Int32(0)): + _fence_proxy_async_global() + _red_release_add( + b2cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done + ) + cute.arch.sync_warp() + # ---- (B2IMG) FC2 B loader (warp 5): group g's image once its counter is complete + if warp == cutlass.Int32(X_WARP): + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + if lane == 0: + _stamp(tdbg, bx, 29) + caddr = b2cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() + if lane == 0: + cv = _ld_acquire_u32(caddr) + while (cv - cnt_target) < cutlass.Int32(0): + _nanosleep(B2IMG_NS) + cv = _ld_acquire_u32(caddr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) + prims.cp_async_bulk_shared_cluster_global( + sB2.subview(buf * cutlass.Int32(IMG_B2)), + img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), + img_bar.subview(buf), + IMG_B2, + ) + prims.cp_async_bulk_shared_cluster_global( + sSFB2.subview(buf * cutlass.Int32(IMG_SF // 4)), + img32.subview( + g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) + ), + img_bar.subview(buf), + IMG_SF, + ) + cute.arch.sync_warp() + _waitd( + img_bar.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 108 + ) + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 25) + # ---- FC2 B builder (warps 5, 8 ..): group g's records (all 48 per (slot, token), the call's tag on + # every word) into B2 buffer g % 2 (row j M + t, k-tile t2 = record / 16) and its scale atoms, then + # tcgen05.cp to TMEM. + if ((warp == X_WARP) | (warp >= cutlass.Int32(8))) & cutlass.Boolean(not B2IMG): + xi = cutlass.select_( + warp == cutlass.Int32(X_WARP), cutlass.Int32(0), warp - cutlass.Int32(7) + ) + if cutlass.const_expr(G_CAP > TOP_K): + if (bx == 0) & (xi == cutlass.Int32(0)): + # rows of slots this call does not use: tag 0, never taken for a later call's + for i in range( + g_n * cutlass.Int32(M_MAX * UNITS_ROW) + lane, + cutlass.Int32(G_CAP * M_MAX * UNITS_ROW), + 32, + ): + hbuf32.store( + ( + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + ), + idx=i * cutlass.Int32(4), + alignment=16, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sSFB2_8 = cutlass.Array( + sSFB2.data_ptr(0), shape=(2 * KT2 * NUM_KB * 512,), dtype=cutlass.Int8, alignment=16 + ) + sB2_32 = cutlass.Array( + sB2.data_ptr(0), + shape=(2 * KT2 * NUM_BYTES_B2 // 4,), + dtype=cutlass.Int32, + alignment=16, + ) + nrec_cap = 4 * M_MAX * UNITS_ROW # records of one group + NZ = (nrec_cap + 32 * XW - 1) // (32 * XW) # records per lane: lane + 32 (XW z + xi) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + _waitd( + fc1_done.subview(g & cutlass.Int32(3)), + (g >> cutlass.Int32(2)) & cutlass.Int32(1), + tdbg, + 4, + ) # this CTA's own share of group g is out + if cutlass.const_expr(TIMING): + if (lane == 0) & (xi == cutlass.Int32(0)) & (g < cutlass.Int32(8)): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(32) + g + ) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + nrec = n_here * cutlass.Int32(M_MAX * UNITS_ROW) + rec0 = g * cutlass.Int32( + 4 * M_MAX * UNITS_ROW + ) # rows (4 g + j) M + t are contiguous + done = cutlass.Int32(0) + rspin = cutlass.Int32(0) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 29) + npoll = cutlass.Int32(0) + got = [[cutlass.Int32(0) for _ in range(3)] for _ in range(NZ)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + v4 = prims.load_ext( + hbuf32.subview((rec0 + qi) * cutlass.Int32(4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for wq in cutlass.range_constexpr(3): + got[z][wq] = cutlass.Int32(v4[wq]) + ok = ok & ( + ( + ((got[z][0] ^ tagv) | (got[z][1] ^ tagv) | (got[z][2] ^ tagv)) + & cutlass.Int32(-16777216) + ) + == cutlass.Int32(0) + ) + npoll = npoll + cutlass.Int32(1) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): + done = cutlass.Int32(1) + else: + _nanosleep(B2_POLL_NS) + if cutlass.const_expr(HANG_DBG): + rspin = rspin + cutlass.Int32(1) + if rspin > cutlass.Int32(1 << 14): + done = cutlass.Int32(1) + tdbg.store( + cutlass.Int64(200), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + warp, + ) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + w0, w1, w2 = got[z][0], got[z][1], got[z][2] + rr = qi // cutlass.Int32( + UNITS_ROW + ) # row within the group: j M + t = B row n + c = qi - rr * cutlass.Int32(UNITS_ROW) + t2 = c >> cutlass.Int32(4) + cc = (c >> cutlass.Int32(1)) & cutlass.Int32(7) + lo = (w0 & cutlass.Int32(0xFFFFFF)) | (w1 << cutlass.Int32(24)) + hi = ((w1 >> cutlass.Int32(8)) & cutlass.Int32(0xFFFF)) | ( + (w2 & cutlass.Int32(0xFFFF)) << cutlass.Int32(16) + ) + sB2_32.store( + (lo, hi), + idx=( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_BYTES_B2) + + rr * cutlass.Int32(KT_BYTES) + + (cc ^ (rr & cutlass.Int32(7))) * cutlass.Int32(16) + + (c & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + # scale of block c: atom (t2, kb = (c % 16) / 4), row rr, byte c % 4 + sSFB2_8.store( + cutlass.Int8((w2 >> cutlass.Int32(16)) & cutlass.Int32(0xFF)), + idx=( + ( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_KB) + + ((c & cutlass.Int32(15)) >> cutlass.Int32(2)) + ) + * cutlass.Int32(128) + + rr * cutlass.Int32(4) + ) + * cutlass.Int32(4) + + (c & cutlass.Int32(3)), + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + if cutlass.const_expr(XW > 1): + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=B2_BAR, number_of_threads=32 * XW) + cute.arch.sync_warp() + if (xi == cutlass.Int32(0)) & prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + SFB2_COL + buf * cutlass.Int32(KT2 * NUM_KB * 4) + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 25) + if cutlass.const_expr(TIMING): + if g < cutlass.Int32(8): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(11) + g + ) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(npoll), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(30), + ) + + # ---- MMA (warp 7) + if warp == MMA_WARP: + idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N1, + m_dim=MMA_M, + sparsity_version=0, + ) + idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N2, + m_dim=MMA_M, + sparsity_version=0, + ) + desc_a0 = prims.Tcgen05SmemDesc.build( + sA.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b10 = prims.Tcgen05SmemDesc.build( + sB1.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b20 = prims.Tcgen05SmemDesc.build( + sB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + _waitd(b1_ready, 0, tdbg, 5) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + acc_b2 = cutlass.Int64(0) + acc_sf = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + _waitd( + acc_empty.subview(ab), + ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 102, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_ptr = cutlass.inttoptr( + tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 + ) + for j in cutlass.range(j0, j1, unroll=1): + tw0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 6) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 6) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc1, + (j > j0) | cutlass.Boolean(kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(ab)) + k1 = k1 + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + buf = g & cutlass.Int32(1) + tw2 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + _waitd( + b2_ready.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 103 + ) + if cutlass.const_expr(TIMING): + acc_b2 = acc_b2 + (_gtime() - tw2) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 4) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_empty.subview(a2), + ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 104, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc2_ptr = cutlass.inttoptr( + tcol + ACC2_COL + a2 * cutlass.Int32(N2), 6, cutlass.Float32 + ) + for t2 in cutlass.range_constexpr(KT2): + tw1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 7) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 7) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw1) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b20 + cutlass.Int64( + (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) + * cutlass.Int32(NUM_BYTES_B2 >> 4) + ) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = ( + tcol + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + cutlass.Int32(t2 * NUM_KB * 4) + ) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc2_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc2, + cutlass.Boolean(t2 > 0 or kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc2_full.subview(a2)) + k2 = k2 + cutlass.Int32(1) + if cutlass.const_expr(TIMING): + if lane == 0: + tdbg.store(acc_b2, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(26)) + tdbg.store(acc_sf, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(27)) + + # ---- epilogue warps (0-3) + if warp < 4: + row_w = base_row + warp * 32 + for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): + qi = tidx + cutlass.Int32(r * EPI_THREADS) + sl_e = qi // cutlass.Int32(6) + if (qi < cutlass.Int32(6 * G_CAP)) & (sl_e < g_n): + fld = qi - sl_e * cutlass.Int32(6) + s_esc.store( + cutlass.Float32( + escal.load(idx=fld * cutlass.Int32(E_LOCAL) + s_el.load(idx=sl_e)) + ), + idx=sl_e * cutlass.Int32(8) + fld, + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + SH_ADD = SHARED + wys = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + pars = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + ard = (cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0)) + for p in range(ng + cutlass.Int32(LAG)): + if cutlass.const_expr(AR_PRE): + if (p == ng + cutlass.Int32(LAG - 1)) & ar_on: + # one 8-byte load each (the vectorizer merges adjacent arp.load()s into 16-byte reads + # at 8-byte offsets: misaligned) + ard = tuple( + cutlass.Int64(_ld_relaxed_u64(arp.subview(i_).data_ptr().toint())) + for i_ in (1, 3, 4, 5) + ) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_cl = cutlass.select_( + f_bb < nb, (blk0 + f_bb) * cutlass.Int32(R2), cutlass.Int32(0) + ) + pars[bt] = _ar_parity(arp, f_t, f_cl) + if cutlass.const_expr(SH_ADD): + if p == ng + cutlass.Int32(LAG - 1): + # this CTA's y_sh words, loaded while the last group's FC2 runs (re-polled in finalize) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_ch = cutlass.select_( + f_bb < nb, + f_t * cutlass.Int32(H) + + (blk0 + f_bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + cutlass.Int32(0), + ) + wys[bt] = _ld_relaxed_u32(dnx.subview(f_ch).data_ptr().toint()) + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + unit = slot * cutlass.Int32(U1) + cblk + _waitd( + acc_full.subview(ab), + (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 105, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 + ), + num=N1, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc_empty.subview(ab)) + pb = ab * cutlass.Int32(M_MAX * C1 * R1) + for t in cutlass.range_constexpr(M_MAX): + part = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + part = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part + ) + s_part.store(part, idx=pb + (t * C1 + warp) * R1 + lane) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + for tr in cutlass.range_constexpr((M_MAX + 3) // 4): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + xv = _add_rn( + _add_rn( + _add_rn( + s_part.load(idx=pb + (t * C1 + 0) * R1 + lane), + s_part.load(idx=pb + (t * C1 + 1) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 2) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 3) * R1 + lane), + ) + finish = cutlass.Int32(1) + if (j0 != cutlass.Int32(0)) | (j1 != cutlass.Int32(S1)): + paddr = ( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint() + ) + if j0 == cutlass.Int32(0): + _st_relaxed_u32(paddr, xv.bitcast(cutlass.Int32)) + finish = cutlass.Int32(0) + else: + pbits = _ld_relaxed_u32(paddr) + spins = cutlass.Int32(0) + while pbits == cutlass.Int32(PARTIAL_EMPTY): + pbits = _ld_relaxed_u32(paddr) + spins = spins + cutlass.Int32(1) + _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) + xv = _add_rn(pbits.bitcast(cutlass.Float32), xv) + cute.arch.sync_warp() + if cutlass.const_expr( + XVDUMP + ): # debug: FC1 sums into the (unused) partial slots + partials.store( + xv.bitcast(cutlass.Int32), + idx=(unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane, + ) + if finish == cutlass.Int32(1): + gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) + is_up = ( + (lane >> cutlass.Int32(3)) & cutlass.Int32(1) + ) == cutlass.Int32(0) + es = slot * cutlass.Int32(8) + qv, sfb = _swiglu_quant( + xv, + gate_raw, + is_up, + s_esc.load(idx=es), + s_esc.load(idx=es + cutlass.Int32(1)), + s_esc.load(idx=es + cutlass.Int32(2)), + s_esc.load(idx=es + cutlass.Int32(3)), + s_esc.load(idx=es + cutlass.Int32(4)), + ) + byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) + bt = [ + cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) + for i in range(8) + ] + rw0 = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | tagv + ) + rw1 = ( + bt[3] + | (bt[4] << cutlass.Int32(8)) + | (bt[5] << cutlass.Int32(16)) + | tagv + ) + rw2 = ( + bt[6] + | (bt[7] << cutlass.Int32(8)) + | ((sfb & cutlass.Int32(0xFF)) << cutlass.Int32(16)) + | tagv + ) + if cutlass.const_expr(B2IMG): + if lane == 0: + # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 + # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), + # row n at byte 16 n, byte c % 4 + n = (ug // cutlass.Int32(U1)) * cutlass.Int32(M_MAX) + t + t2i = cblk >> cutlass.Int32(4) + cci = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) + gbase = p * cutlass.Int32(IMG_BYTES) + lo = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | (bt[3] << cutlass.Int32(24)) + ) + hi = ( + bt[4] + | (bt[5] << cutlass.Int32(8)) + | (bt[6] << cutlass.Int32(16)) + | (bt[7] << cutlass.Int32(24)) + ) + img32.store( + (lo, hi), + idx=( + gbase + + t2i * cutlass.Int32(NUM_BYTES_B2) + + n * cutlass.Int32(KT_BYTES) + + (cci ^ (n & cutlass.Int32(7))) + * cutlass.Int32(16) + + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + img8 = cutlass.Array( + img32.data_ptr(0), + shape=(NG_CAP * IMG_BYTES,), + dtype=cutlass.Int8, + alignment=16, + ) + img8.store( + cutlass.Int8(sfb & cutlass.Int32(0xFF)), + idx=gbase + + cutlass.Int32(IMG_B2) + + ( + ( + t2i * cutlass.Int32(NUM_KB) + + ( + (cblk & cutlass.Int32(15)) + >> cutlass.Int32(2) + ) + ) + * cutlass.Int32(512) + ) + + n * cutlass.Int32(16) + + (cblk & cutlass.Int32(3)), + ) + else: + if lane == 0: + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4), + alignment=16, + ) + k1 = k1 + cutlass.Int32(1) + prims.mbarrier_arrive( + fc1_done.subview(p & cutlass.Int32(3)) + ) # warp 5 may poll group p + if tidx == 0: + _stamp(tdbg, bx, 10) + if cutlass.const_expr(TIMING): + if p < cutlass.Int32(8): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(40) + p, + ) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_full.subview(a2), + (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 106, + ) + if cutlass.const_expr(TIMING): + if (tidx == 0) & (g == ng - cutlass.Int32(1)): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(8) + bb, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, + base_col + + ACC2_COL + + a2 * cutlass.Int32(N2) + + warp * cutlass.Int32(M_MAX), + cutlass.Float32, + ), + num=LD_POW2, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc2_empty.subview(a2)) + sl = g * cutlass.Int32(4) + warp + if sl < g_n: + a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) + for t in cutlass.range_constexpr(M_MAX): + kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + if kp >= cutlass.Int32(0): + ys.store( + _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)), + idx=( + (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + * cutlass.Int32(TOP_K) + + kp + ) + * cutlass.Int32(R2) + + lane, + ) + k2 = k2 + cutlass.Int32(1) + prims.mbarrier_arrive(b2_free.subview(g & cutlass.Int32(1))) + if tidx == 0: + _stamp(tdbg, bx, 6) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + if tidx == 0: + _stamp(tdbg, bx, 19) + # finalize: (block bb, token t) per warp round: the fma chain over t's top-4 in order + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + pair = warp + cutlass.Int32(bt * 4) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + if bb < nb: + acc_o = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(TOP_K): + acc_o = _fma_rn( + s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ), + acc_o, + ) + if cutlass.const_expr( + YDUMP + ): # debug: y of (token, k) into the partials workspace + partials.store( + cutlass.Float32( + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ) + ).bitcast(cutlass.Int32), + idx=(t * cutlass.Int32(TOP_K) + cutlass.Int32(k)) * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + ) + ch = ( + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(SH_ADD): + # + y_sh (this row's shared output, from its shared CTA), as shared_output.add_(routed_output) + wy = wys[bt] + while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: + wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + acc_o = _add_rn( + (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) + ) + out.store(cutlass.BFloat16(acc_o), idx=ch) + if ar_on: + cl = ( + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(AR_PRE): + _ar_push_pre( + ard, + pars[bt], + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + else: + _ar_push( + arp, + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + + # ---- teardown + if tidx == 0: + _stamp(tdbg, bx, 7) + epochs.store(ep + cutlass.Int32(1), idx=bx) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if ar_on | cutlass.Boolean(LATE_TRIGGER): + # the all-reduce polls the pushed slots without a grid wait: release it after every push of this CTA + cute.arch.griddepcontrol_launch_dependents() + if warp == MMA_WARP: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS + ) + + +@cute.jit +def m3_moe( + a1_tensor: cute.Tensor, + sfa1_tensor: cute.Tensor, + a2_tensor: cute.Tensor, + sfa2_tensor: cute.Tensor, + logits: cute.Tensor, + bias: cute.Tensor, + x32: cute.Tensor, + xsf32: cute.Tensor, + escal: cute.Tensor, + plan: cute.Tensor, + partials: cute.Tensor, + hbuf32: cute.Tensor, + epochs: cute.Tensor, + out: cute.Tensor, + route_dbg: cute.Tensor, + tdbg: cute.Tensor, + xb32: cute.Tensor, + rw: cute.Tensor, + sfsc: cute.Tensor, + lgbuf: cute.Tensor, + guw: cute.Tensor, + gusf: cute.Tensor, + dnw: cute.Tensor, + dnsf: cute.Tensor, + hx: cute.Tensor, + dnx: cute.Tensor, + arp: cute.Tensor, + s1w: cute.Tensor, + s2w: cute.Tensor, + img32: cute.Tensor, + b2cnt: cute.Tensor, + stream: cuda_driver.CUstream, +): + tma_a1_desc = cuda.create_tensor_map_tiled( + global_address=a1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H // 2, TWO_I, E_LOCAL], + global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], + box_dims=(KT_BYTES, R1, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_a2_desc = cuda.create_tensor_map_tiled( + global_address=a2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[I_TP // 2, H, E_LOCAL], + global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], + box_dims=(KT_BYTES, R2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_s1_desc = cuda.create_tensor_map_tiled( + global_address=sfa1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_s2_desc = cuda.create_tensor_map_tiled( + global_address=sfa2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_gu_desc = tma_a1_desc + tma_dn_desc = tma_a2_desc + if cutlass.const_expr(SHARED and not SH_PRETILE): + tma_gu_desc = cuda.create_tensor_map_tiled( + global_address=guw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H, 2 * SH_I], + global_strides=[H // 16], + box_dims=(KT_BYTES, 32), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_dn_desc = cuda.create_tensor_map_tiled( + global_address=dnw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[SH_I, H], + global_strides=[SH_I // 16], + box_dims=(KT_BYTES, 128), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + m3_moe_kernel( + tma_a1_desc, + tma_a2_desc, + tma_s1_desc, + tma_s2_desc, + tma_gu_desc, + tma_dn_desc, + logits, + bias, + x32, + xsf32, + escal, + plan, + partials, + hbuf32, + epochs, + out, + route_dbg, + tdbg, + xb32, + rw, + sfsc, + lgbuf, + guw, + gusf, + dnw, + dnsf, + hx, + dnx, + arp, + s1w, + s2w, + img32, + b2cnt, + ).launch( + grid=(NCTA + NSH, 1, 1), + block=(THREADS, 1, 1), + cluster=(1, 1, 1), + stream=stream, + use_pdl=True, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/op.py new file mode 100644 index 000000000000..f521a23589ea --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/op.py @@ -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. +"""Host side of the MiniMax-M3 routed-expert kernel on one TP4 rank, over the TRTLLM-Gen NVFP4 buffers read in place. + +``m3_moe`` (m3_moe_kernel.py, :class:`M3MoeState`): 1..8 decode tokens (decode batches, spec-decode verify). The call +takes the bf16 hidden states and folds in the router GEMV (bit-exact with m3_router_gemv_fp4) and the activation's +NVFP4 quantization. With ``shared=True`` it also runs the shared expert on dedicated CTAs and returns +bf16(shared + routed); otherwise it returns the rank's routed partial, bit for bit what the stock fused MoE returns. + +Usage (build before CUDA-graph capture; every call is capture-safe afterwards):: + + state = M3MoeState( + device, num_tokens=6, shared=True + ) # workspace + compiled kernel, shared by all layers + layer = state.layer( + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared_weights=(gate_up_w, gate_up_sf, down_w, down_sf), + ) + out = layer(hidden) # bf16 [M, 6144] +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +from typing import Dict, Optional, Tuple + +import torch + +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size +INTER = 768 +NUM_EXPERTS = 128 +# The shared expert's intermediate per TP rank. +SHARED_INTER = 768 # capacity: fixed: the shared expert's width per tp4 rank +# m3_moe's shared-expert CTAs (beside num_ctas routed CTAs). +SHARED_CTAS = 24 # capacity: fixed: an SM split, not a size +TOP_K = 4 +U1 = 2 * INTER // 32 # FC1 units per expert +# The router GEMV takes at most 2 experts per CTA (router weights held in registers). +MIN_CTAS = 64 # capacity: fixed: an SM split, not a size + +_UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) # one FC1 unit's weight + scale bytes +_ROW2_BYTES = INTER // 2 + INTER // 16 # one FC2 weight row's bytes (values + scales) +_S1 = 6 # FC1 stages (items) per unit # capacity: fixed: pipeline stages +_ITEM_BYTES = _UNIT_BYTES // _S1 # capacity: fixed: a ring item's bytes +_modules: Dict[tuple, object] = {} +_compiled: Dict[tuple, object] = {} +# m3_moe trace-time options by token count 1..8 (index num_tokens - 1). +# ``sfa_cp``: 18 KB ring stages, the A scales copied by cp.async into the instruction layout (every token count). +SFA_CP_BY_M: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1, 1) +# ``b2img``: FC1 -> FC2 hand-off through per-group global B images + release counters (7 and 8 tokens). +B2IMG_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 1, 1) +# ``ar_pre``: the finalize's all-reduce push with its call parities / descriptor preloaded (all but 2 tokens). +AR_PRE_BY_M: Tuple[int, ...] = (1, 0, 1, 1, 1, 1, 1, 1) +# Ring ``stages`` (0 = the kernel's default, the most that fit). 6 at 8 tokens measured 0.4-0.5 us faster than the +# default 7 (same MMAs, bit-identical outputs). +STAGES_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 0, 6) +# m3_moe phase lag: phase p streams FC1 of group p and FC2 of group p - M3_MOE_LAG (the kernel's default lag). +M3_MOE_LAG = 2 + + +_ENGINE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe_kernel.py") + + +def _kernel_module(config: dict, path: str): + key = (os.path.basename(path),) + tuple(sorted(config.items())) + mod = _modules.get(key) + if mod is None: + name = ( + f"{__name__}_" + + os.path.splitext(os.path.basename(path))[0] + + "_" + + "_".join(f"{k}{v}" for k, v in key[1:]) + ) + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + mod.M3_CONFIG = dict(config) + sys.modules[name] = mod + spec.loader.exec_module(mod) + _modules[key] = mod + return mod + + +def _view(t: torch.Tensor, align: int): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align) + + +def is_supported(device: Optional[torch.device] = None) -> tuple: + if not torch.cuda.is_available(): + return False, "no CUDA" + major, _ = torch.cuda.get_device_capability(device) + if major != 10: + return False, "needs the sm_100 family (tcgen05 block-scaled NVFP4 MMA)" + try: + import cutlass # noqa: F401 + except ImportError: + return False, "CuTe DSL (nvidia-cutlass-dsl) is not installed" + return True, "" + + +def ar_push_desc( + mc_ptr: int, flags_ptr: int, rank: int, world: int, max_tokens: int, device +) -> torch.Tensor: + """The ``arp`` descriptor that makes a call also push its bf16 output into a multicast all-reduce exchange (set A + of m3_mc_allreduce's workspace at ``mc_ptr``, slot half = the all-reduce CTA's call parity in ``flags_ptr``). + With it the kernel releases its PDL dependents only after every push; the default (enable 0) keeps the old + behaviour.""" + return torch.tensor( + [1, mc_ptr, flags_ptr, rank, world, max_tokens, 0, 0], dtype=torch.int64, device=device + ) + + +def _shared_views(shared_weights, who): + """(gate_up, its scales, down, its scales) of the shared expert as flat uint8 views, checked.""" + if shared_weights is None: + raise ValueError(f"{who} (shared): shared_weights are required") + gu_w, gu_sf, dn_w, dn_sf = shared_weights + want = {"gate_up": (gu_w, (2 * SHARED_INTER, HIDDEN)), "down": (dn_w, (HIDDEN, SHARED_INTER))} + for name, (t, shape) in want.items(): + if t.dtype != torch.float8_e4m3fn or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError(f"{who} (shared): {name} weight must be contiguous e4m3 {shape}") + if ( + gu_sf.numel() != 2 * SHARED_INTER * HIDDEN // 32 + or dn_sf.numel() != HIDDEN * SHARED_INTER // 32 + ): + raise ValueError( + f"{who} (shared): the UE8M0 scales must be the swizzled [N, K / 32] buffers" + ) + return ( + gu_w.view(torch.uint8).view(-1), + gu_sf.view(torch.uint8).reshape(-1), + dn_w.view(torch.uint8).view(-1), + dn_sf.view(torch.uint8).reshape(-1), + ) + + +def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24) -> torch.Tensor: + """The shared expert's weights as m3_moe's shared CTAs stream them (SH_PRETILE): per CTA s, per stage g, the 16 KB + shared-memory image (128-byte swizzle: 16-byte chunk j of row r at j ^ (r % 8)) of the stage's 128-row MMA tile and + its 512-byte instruction scale atom (row R's 4 codes at word 4 (R % 32) + R / 32). Stages 0..23: gate_up k-tiles + g (slice 0) / 24 + g (slice 1) of rows [gate 32 s .., up 768 + 32 s .., gate (slice 1), up (slice 1)]; 24..35: + down tile (g - 24) / 6 (rows 256 s + 128 j ..), k-tile (g - 24) % 6. Built once per layer.""" + dev = gu8.device + H_, I_ = HIDDEN, SHARED_INTER + gu = gu8.view(2 * I_, H_) + dn = dn8.view(H_, I_) + gkb, dkb = H_ // 32, I_ // 32 + + def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast + return sf[ + ((n >> 7) * (kbn // 4) + (b >> 2)) * 512 + (n & 31) * 16 + ((n >> 5) & 3) * 4 + (b & 3) + ] + + r = torch.arange(128, device=dev) + q = r >> 5 + ri = r & 31 + chunk = torch.arange(8, device=dev) + swz = chunk[None, :] ^ (r[:, None] & 7) # source chunk of image chunk j: j ^ (r % 8) + out = torch.empty( + n_cta, 36, 16384 + 512, dtype=torch.uint8, device=dev + ) # capacity: fixed: a workspace of the kernel's fixed geometry + for s in range(n_cta): + # gate_up: MMA row R -> weight row, slice + nrow = torch.where((q & 1) == 0, 32 * s + ri, I_ + 32 * s + ri) + sl = q >> 1 + for g in range(24): + kt = sl * 24 + g # per row + tile = gu[ + nrow[:, None], (kt[:, None] * 128 + torch.arange(128, device=dev)[None, :]) + ] # [128, 128] + img = tile.view(128, 8, 16)[r[:, None], swz] # image chunk j = tile chunk j ^ (r % 8) + out[s, g, :16384] = img.reshape(-1) + kb = torch.arange(4, device=dev) + codes = sf_bytes(gusf8, nrow[:, None], kt[:, None] * 4 + kb[None, :], gkb) # [128, 4] + atom = torch.empty(512, dtype=torch.uint8, device=dev) + atom[(ri[:, None] * 16 + q[:, None] * 4 + kb[None, :]).reshape(-1)] = codes.reshape(-1) + out[s, g, 16384:] = atom + for gd in range(12): + jt, kt = gd // 6, gd % 6 + row0 = 256 * s + 128 * jt + tile = dn[row0 : row0 + 128, kt * 128 : (kt + 1) * 128] + img = tile.reshape(128, 8, 16)[r[:, None], swz] + out[s, 24 + gd, :16384] = img.reshape(-1) + out[s, 24 + gd, 16384:] = dnsf8[ + ((row0 >> 7) * (dkb // 4) + kt) * 512 : ((row0 >> 7) * (dkb // 4) + kt) * 512 + 512 + ] + return out.view(-1) + + +# --------------------------------------------------------------------------------------------------------------------- +# m3_moe (m3_moe_kernel.py): 1..8 tokens per call, group-phased weight stream. +# --------------------------------------------------------------------------------------------------------------------- +# Pre-tiled shared weights per layer, keyed by the source buffers (which the engines read in place): every token-count +# state of a layer streams the same copy. An entry holds its sources, so their addresses are not reused while cached. +_PRETILED: Dict[tuple, tuple] = {} + + +def _pretiled(gu8, gusf8, dn8, dnsf8) -> torch.Tensor: + key = (gu8.device, gu8.data_ptr(), gusf8.data_ptr(), dn8.data_ptr(), dnsf8.data_ptr()) + hit = _PRETILED.get(key) + if hit is None: + hit = _PRETILED[key] = ((gu8, gusf8, dn8, dnsf8), _pretile_shared(gu8, gusf8, dn8, dnsf8)) + return hit[1] + + +def make_phase_plan( + num_ctas: int, m_max: int, lag: int = 2, tail_w: float = 1.0, early0: bool = False +): + """The work split of ``m3_moe`` for each distinct-expert count G (4 .. 4 m_max): FC2 32-row output blocks in + contiguous per-CTA ranges (the same for every group), and per group p (slots 4 p ..) a range of its unit-major FC1 + stages per CTA, balancing the bytes each CTA streams in phase p (FC1 of group p plus FC2 of group p - lag, and the + trailing FC2s with the last group). Ranges are empty or >= 6 stages (a unit is split over at most two CTAs); a + group too small for every CTA gets whole units on the CTAs with the least FC2 work. Returns int32 + [4 m_max + 1, num_ctas, m_max + 1] (per group i0 | i1 << 16, then b0 | nb << 16) and the largest block count.""" + gcap = TOP_K * m_max + nb2 = HIDDEN // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + blk_bytes = 32 * _ROW2_BYTES # one block, one slot + for g in range(TOP_K, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) + for p in range(ng): + n_p = min(4, g - 4 * p) + total = n_p * U1 * _S1 + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * min(4, g - 4 * (p - lag)) * blk_bytes + if p == ng - 1 and not ( + early0 and p == 0 and m_max >= 2 + ): # early0: group 0's ranges the same for every G + for q in range(max(0, ng - lag), ng): + fc2[c] += int(tail_w * nbs[c] * min(4, g - 4 * q) * blk_bytes) + items = [0] * num_ctas + if total >= _S1 * num_ctas: + + def counts_for(limit): + return [max(_S1, (limit - f) // _ITEM_BYTES) for f in fc2] + + lo, hi = 0, max(fc2) + total * _ITEM_BYTES + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > _S1), + key=lambda c: (-(fc2[c] + items[c] * _ITEM_BYTES), c), + ) + items[c] -= 1 + extra -= 1 + else: + chosen = sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // _S1]) + for c in chosen: + items[c] = _S1 + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | ((i0 + items[c]) << 16) + i0 += items[c] + assert i0 == total + return plan, max(nbs) + + +class M3MoeState: + """Workspace + compiled kernel of ``m3_moe`` (1..8 tokens per call) on one device, shared by the layers that run + on it in one stream order. The call takes the bf16 hidden states and computes the router logits (bit-exact with + m3_router_gemv_fp4) and the NVFP4 activation itself. ``shared``: the shared expert runs on SHARED_CTAS dedicated + CTAs beside the ``num_ctas`` routed ones and the call returns bf16(shared + routed). Build before CUDA-graph + capture; the first call compiles.""" + + fused_router = True + + def __init__( + self, + device: torch.device, + num_tokens: int, + num_ctas: Optional[int] = None, + shared: bool = False, + ): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "M3MoeState allocates and compiles: build it before CUDA-graph capture" + ) + sms_all = torch.cuda.get_device_properties(device).multi_processor_count + if shared and (num_ctas or sms_all) + SHARED_CTAS > sms_all: + raise ValueError( + f"m3_moe (shared): num_ctas + {SHARED_CTAS} shared-expert CTAs must fit the {sms_all} SMs " + f"(all resident at once)" + ) + if not 1 <= num_tokens <= 8: + raise ValueError(f"m3_moe takes 1..8 tokens per call, not {num_tokens}") + self.device = torch.device(device) + ok, why = is_supported(self.device) + if not ok: + raise RuntimeError(f"m3_moe: {why}") + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.num_ctas = num_ctas or sms + if not MIN_CTAS <= self.num_ctas <= sms: + raise ValueError( + f"m3_moe runs {MIN_CTAS} .. {sms} CTAs (all resident at once), not {self.num_ctas}" + ) + self.m = num_tokens + # One token: the trailing FC2s weigh nothing in the last phase's balance. Group 0's FC1 ranges are the same + # for every distinct-expert count (early0). + plan, nb_cap = make_phase_plan( + self.num_ctas, num_tokens, M3_MOE_LAG, 0.0 if num_tokens == 1 else 1.0, True + ) + self.config = { + "m_max": num_tokens, + "num_ctas": self.num_ctas, + "nb2_cap": nb_cap, + "fused_router": 1, + "timing": 0, + } + if shared: + self.config["shared"] = 1 + self.shared = shared + # 18 KB ring stages (A scales by cp.async into the instruction layout) per token count, unless configured + self.config["sfa_cp"] = SFA_CP_BY_M[num_tokens - 1] + # FC1 -> FC2 hand-off through per-group global B images + release counters, per token count + self.config["b2img"] = B2IMG_BY_M[num_tokens - 1] + # the finalize's all-reduce push with its call parities / descriptor preloaded, per token count + self.config["ar_pre"] = AR_PRE_BY_M[num_tokens - 1] + if STAGES_BY_M[num_tokens - 1] > 0: + self.config["stages"] = STAGES_BY_M[num_tokens - 1] + try: + self.mod = mod = _kernel_module(self.config, _ENGINE_PATH) + except AssertionError as err: + raise ValueError( + f"m3_moe: num_tokens={num_tokens}, num_ctas={self.num_ctas} does not fit the shared " + f"memory / TMEM budget (fewer CTAs means more FC2 blocks per CTA): {err}" + ) from err + kw = dict(device=self.device) + self.plan = plan.to(self.device) + self.partials = torch.full((mod.G_CAP * U1 * num_tokens * 32,), -1, dtype=torch.int32, **kw) + self.hbuf32 = torch.zeros(mod.G_CAP * num_tokens * U1 * 4, dtype=torch.int32, **kw) + self.epochs = torch.zeros(self.num_ctas + mod.NSH, dtype=torch.int32, **kw) + self.route_dbg = torch.zeros(num_tokens, 8, dtype=torch.int32, **kw) + self.tdbg = torch.zeros(self.num_ctas + mod.NSH, 64, dtype=torch.int64, **kw) + self.lgbuf = torch.zeros(mod.LG_REPL * mod.LG_STRIDE, dtype=torch.int64, **kw) + self._dummy = torch.zeros(4, dtype=torch.int32, **kw) + self._ar_off = torch.zeros(8, dtype=torch.int64, **kw) + # B2IMG: per-group FC2 B images and their unit counters (int32; every group gains CNT_PER_CALL per call) + self.img32 = torch.zeros(max(1, mod.NG_CAP * mod.IMG_BYTES // 4), dtype=torch.int32, **kw) + self.cnt = torch.zeros(mod.NG_CAP * 8 + 8, dtype=torch.int32, **kw) + # shared expert: h and y_sh as (call tag << 16 | bf16), tag 0 = never written + self.hx = torch.zeros(num_tokens * SHARED_INTER, dtype=torch.int32, **kw) + self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) + self._fn = None + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> "M3MoeLayer": + """One MoE layer's operands (the TRTLLM-Gen NVFP4 expert buffers, read in place, never copied). + ``shared_weights`` (a ``shared`` state): the shared expert's (gate_up e4m3 [1536, 6144], its swizzled UE8M0 + scales, down e4m3 [6144, 768], its swizzled UE8M0 scales).""" + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + ) + if self.shared: + layer.shared_weights = _shared_views(shared_weights, "m3_moe") + layer.shared_mxfp8 = tuple(shared_weights) + if self.mod.SH_PRETILE: + # (gate_up, pre-tiled stages, down, dummy): the kernel streams the pre-tiled copy + gu8, _, dn8, _ = layer.shared_weights + layer.shared_weights = ( + gu8, + _pretiled(*layer.shared_weights), + dn8, + self._dummy.view(torch.uint8), + ) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError( + f"m3_moe (fused router): expects contiguous bf16 hidden states [{m}, {HIDDEN}]" + ) + fz = ( + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + ) + lgv, xq, xs = self._dummy.view(torch.float32), self._dummy, self._dummy + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + sh = layer.shared_weights if self.shared else (self._dummy.view(torch.uint8),) * 4 + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + lgv, + layer.bias, + xq, + xs, + layer.escal, + self.plan.view(-1), + self.partials, + self.hbuf32, + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + *fz, + *sh, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + layer.w3_w1_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + layer.w2_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + self.img32, + self.cnt, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_moe compiles on its first call: call it once before capture") + import cutlass.cute as cute + + aligns = [ + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + ] + sig = [_view(t, a) for t, a in zip(args, aligns)] + key = (_ENGINE_PATH,) + tuple(sorted(self.config.items())) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out + + +class M3MoeLayer: + """One MoE layer on a :class:`M3MoeState` (or op16's state): ``layer(hidden)``. The expert buffers are read in + place, never copied; the per-expert scalars are packed once (layer constants: rebuild the layer after reloading + weights).""" + + shared_weights = None # the kernel's view of the shared expert (a shared state) + shared_mxfp8 = None # the shared expert's (gate_up, its scales, down, its scales) as given + + def __init__( + self, + state, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + bias, + router_weight=None, + fc31_input_scale=None, + ): + e = NUM_EXPERTS + expected = { + "w3_w1_weight": (w3_w1_weight, (e, 2 * INTER, HIDDEN // 2)), + "w3_w1_weight_scale": (w3_w1_weight_scale, (e, 2 * INTER, HIDDEN // 16)), + "w2_weight": (w2_weight, (e, HIDDEN, INTER // 2)), + "w2_weight_scale": (w2_weight_scale, (e, HIDDEN, INTER // 16)), + } + for name, (t, shape) in expected.items(): + if t.element_size() != 1 or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError( + f"m3_moe: {name} is {t.dtype} {tuple(t.shape)}, expected contiguous 1-byte {shape}" + ) + self.state = state + self.w3_w1_weight = w3_w1_weight.view(torch.uint8) + self.w3_w1_weight_scale = w3_w1_weight_scale.view(torch.uint8) + self.w2_weight = w2_weight.view(torch.uint8) + self.w2_weight_scale = w2_weight_scale.view(torch.uint8) + self.escal = torch.stack( + [fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha] + ).float() + self.escal = self.escal.reshape(6, e).contiguous() + self.bias = bias.float().contiguous() + if ( + router_weight is None + or fc31_input_scale is None + or router_weight.dtype != torch.float32 + or tuple(router_weight.shape) != (e, HIDDEN) + or not router_weight.is_contiguous() + or fc31_input_scale.dtype != torch.float32 + or fc31_input_scale.numel() != 1 + ): + raise ValueError( + "m3_moe (fused router): needs router_weight f32 [128, 6144] and fc31_input_scale f32 [1]" + ) + self.router_weight = router_weight + self.fc31_input_scale = fc31_input_scale + + def __call__( + self, + hidden: torch.Tensor, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``hidden``: contiguous bf16 [M, 6144]. ``arp`` (int64 [8], see :func:`ar_push_desc`): also push the output + into an all-reduce exchange.""" + return self.state.run(self, hidden, out, arp=arp) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/op16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/op16.py new file mode 100644 index 000000000000..446776419315 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/op16.py @@ -0,0 +1,288 @@ +# 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. +"""Host side of ``m3_moe16`` (m3_moe16_kernel.py): the MiniMax-M3 MoE layer (fused router + routed experts + shared +expert) for 9..16 decode tokens on one TP4 rank. The 1..8-token engine (op.M3MoeState) is untouched; select this +state only for num_tokens > 8. + +Output: bf16 [M, 6144] = bf16(shared + routed), the routed part bit for bit TRTLLM-Gen's fused MoE on the fp32 router +GEMV logits (m3_router_gemv's math), the shared part as the in-model m3_mx_gemv shared expert. + + state = M3Moe16State(device, num_tokens=12) # before CUDA-graph capture; one per token count + layer = state.layer(w3_w1_weight, w3_w1_weight_scale, w2_weight, w2_weight_scale, fc31_alpha, fc31_scale_c, + act_alpha, act_beta, act_clamp, fc2_alpha, e_score_correction_bias, + router_weight=router_weight, fc31_input_scale=fc31_input_scale, + shared_weights=(gate_up_w, gate_up_sf, down_w, down_sf)) + out = layer(hidden) # first call compiles: call once before capture + +Protocol state (per state object, shared by the layers that run on it in stream order): per-group FC2-B images and +unit counters, the B1 image + counter, per-CTA epochs. Every call adds exactly 192 to each group counter and NCTA to +the B1 counter, so nothing is reset between calls (graph replay safe). +""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch + +from .op import ( + HIDDEN, + SHARED_CTAS, + SHARED_INTER, + U1, + M3MoeLayer, + _compiled, + _kernel_module, + _shared_views, + _view, + is_supported, +) + + +def _group_slots(g: int, p: int, perm: bool, min_ng: int = 5) -> int: + """Slots of group p when the call has g distinct experts (the kernel's slot order).""" + ng = (g + 3) // 4 + tail = g - 4 * (ng - 1) + if p >= ng: + return 0 + if perm and tail < 4 and ng >= max(3, min_ng): + return tail if p == 1 else 4 + return min(4, g - 4 * p) + + +def make_phase_plan16(num_ctas: int, m_max: int, lag: int = 2, perm: bool = True, min_ng: int = 5): + """op.make_phase_plan (early0, tail_w 1) with the kernel's group sizes (PERM: a partial last group is group 1).""" + S1 = 6 + UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) + ITEM = UNIT_BYTES // S1 + blk_bytes = 32 * (768 // 2 + 768 // 16) + gcap = 4 * m_max + nb2 = HIDDEN // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + for g in range(4, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) + for p in range(ng): + n_p = _group_slots(g, p, perm, min_ng) + total = n_p * U1 * S1 + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * _group_slots(g, p - lag, perm, min_ng) * blk_bytes + if p == ng - 1 and not (p == 0 and m_max >= 2): + for q in range(max(0, ng - lag), ng): + fc2[c] += nbs[c] * _group_slots(g, q, perm, min_ng) * blk_bytes + items = [0] * num_ctas + if total >= S1 * num_ctas: + + def counts_for(limit): + return [max(S1, (limit - f) // ITEM) for f in fc2] + + lo, hi = 0, max(fc2) + total * ITEM + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > S1), + key=lambda c: (-(fc2[c] + items[c] * ITEM), c), + ) + items[c] -= 1 + extra -= 1 + else: + for c in sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // S1]): + items[c] = S1 + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | ((i0 + items[c]) << 16) + i0 += items[c] + assert i0 == total + return plan, max(nbs) + + +_ENGINE16_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe16_kernel.py") +M16_MIN_TOKENS, M16_MAX_TOKENS = 9, 16 +# Default routed CTAs (beside the SHARED_CTAS shared-expert CTAs). +M16_NUM_CTAS = 128 +# Phase lag: phase p streams FC1 of group p and FC2 of group p - M16_LAG (the kernel's default lag). +M16_LAG = 3 + + +class M3Moe16State: + """Workspace + compiled kernel of ``m3_moe16`` for ``num_tokens`` (9..16) tokens per call.""" + + fused_router = True + shared = True + + def __init__(self, device: torch.device, num_tokens: int, num_ctas: Optional[int] = None): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "M3Moe16State allocates and compiles: build it before CUDA-graph capture" + ) + if not M16_MIN_TOKENS <= num_tokens <= M16_MAX_TOKENS: + raise ValueError( + f"m3_moe16 takes {M16_MIN_TOKENS}..{M16_MAX_TOKENS} tokens per call, not {num_tokens}" + ) + self.device = torch.device(device) + ok, why = is_supported(self.device) + if not ok: + raise RuntimeError(f"m3_moe16: {why}") + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.num_ctas = num_ctas or M16_NUM_CTAS + if not 128 <= self.num_ctas <= sms - SHARED_CTAS: + raise ValueError( + f"m3_moe16 runs 128 .. {sms - SHARED_CTAS} routed CTAs + {SHARED_CTAS} shared (all resident)" + ) + self.m = num_tokens + plan, nb_cap = make_phase_plan16(self.num_ctas, num_tokens, M16_LAG) + self.config = { + "m_max": num_tokens, + "num_ctas": self.num_ctas, + "nb2_cap": nb_cap, + "fused_router": 1, + "timing": 0, + "shared": 1, + } + try: + self.mod = mod = _kernel_module(self.config, _ENGINE16_PATH) + except AssertionError as err: + raise ValueError( + f"m3_moe16: num_tokens={num_tokens} does not fit the smem / TMEM budget: {err}" + ) from err + kw = dict(device=self.device) + self.plan = plan.to(self.device) + # split units' head sums: every slot empty (0xFFFFFFFF) between calls + self.partials = torch.full((mod.G_CAP * U1 * num_tokens * 32,), -1, dtype=torch.int32, **kw) + # FC2-B images per group + the B1 image; counters: group g at 8 g, B1 at 8 NG_CAP, CTA tickets at 8 NG_CAP + 8 + self.img = torch.zeros(mod.IMG_TOTAL // 4, dtype=torch.int32, **kw) + self.cnt = torch.zeros(mod.NG_CAP * 8 + 16, dtype=torch.int32, **kw) + self.epochs = torch.zeros(self.num_ctas + mod.NSH, dtype=torch.int32, **kw) + self.route_dbg = torch.zeros(num_tokens, 8, dtype=torch.int32, **kw) + self.tdbg = torch.zeros(self.num_ctas + mod.NSH, 64, dtype=torch.int64, **kw) + self.lgbuf = torch.zeros( + mod.LG_REPL * mod.LG_STRIDE, dtype=torch.int64, **kw + ) # logits: tag << 32 | f32 + self._ar_off = torch.zeros(8, dtype=torch.int64, **kw) + # shared expert: h and y_sh as (call tag << 16 | bf16), tag 0 = never written + self.hx = torch.zeros(num_tokens * SHARED_INTER, dtype=torch.int32, **kw) + self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) + self._fn = None + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> M3MoeLayer: + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + ) + layer.shared_weights = _shared_views(shared_weights, "m3_moe16") + layer.shared_mxfp8 = tuple(shared_weights) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError(f"m3_moe16: expects contiguous bf16 hidden states [{m}, {HIDDEN}]") + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + layer.bias, + layer.escal, + self.plan.view(-1), + self.partials, + self.img, + self.cnt, + layer.w3_w1_weight_scale.reshape(-1).view(torch.int32), + layer.w2_weight_scale.reshape(-1).view(torch.int32), + layer.w3_w1_weight.reshape(-1).view(torch.int8), + layer.w2_weight.reshape(-1).view(torch.int8), + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + *layer.shared_weights, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_moe16 compiles on its first call: call it once before capture" + ) + import cutlass.cute as cute + + sig = [_view(t, 16 if t.element_size() * t.numel() % 16 == 0 else 4) for t in args] + key = (_ENGINE16_PATH,) + tuple(sorted(self.config.items())) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_prequant.py new file mode 100644 index 000000000000..48a7827f59b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_prequant.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hand-off of an MXFP8-quantized activation from the kernel that produced it to the next MXFP8 GEMM. + +A producer that already holds a BF16 row in registers (the multicast all-reduce + RMSNorm, the split-K attention +merge) also writes its MXFP8 form (E4M3 rows + UE8M0 scales in the 128x4-swizzled layout of +``trtllm::mxfp8_quantize(x, True)``) and publishes it here keyed by the BF16 tensor. The MXFP8 linear whose input +is that tensor takes it and skips its own quantize kernel. Every take clears the slot, so a hand-off is only ever +consumed by the first MXFP8 linear that runs after the producer. + +Module state shared across files (one instance each per process): + +* ``SLOT`` (PrequantSlot): published by ``m3_mc_allreduce`` (``quant=True``), the sparse-decode merge and the fused + decode attention (MXFP8 output); cleared by the sparse decode before it runs; taken by the MXFP8 linear dispatch + and by the one-launch QKV + attention + o_proj path (which also peeks ``pending``). +* ``AR_PUSH`` (ArPushSlot): armed by the fast forward around the attention call; the o_proj GEMV takes the target + and marks ``pushed``; the fast forward checks ``take_pushed`` before the attention boundary all-reduce. +* ``MOE_PUSH`` (MoePushSlot): armed by the fast forward before the MoE; the routed-expert engine takes ``engine`` + and the shared expert's down projection takes ``shared``, each marks its pushed output; the fast forward + checks ``take_pushed`` before the MoE boundary all-reduce. +""" + + +class PrequantSlot: + def __init__(self): + self.pending = None # (data_ptr, rows, width, e4m3 [rows, width], scales) + + def publish(self, x, xq, xsf): + self.pending = (x.data_ptr(), x.shape[0], x.shape[-1], xq, xsf) + + def clear(self): + self.pending = None + + def take(self, x): + """(e4m3 [x.shape[0], width], scales) if x is the published tensor, else None.""" + pending, self.pending = self.pending, None + if ( + pending is None + or x.dim() != 2 + or x.data_ptr() != pending[0] + or x.shape[0] != pending[1] + or x.shape[1] != pending[2] + ): + return None + return pending[3], pending[4] + + +SLOT = PrequantSlot() + + +class ArPushSlot: + """The attention boundary's all-reduce exchange for the o_proj GEMV to push into. + + The fast forward arms ``target`` (m3_mx_gemv ``ar`` tuple) around the attention call; the o_proj marks the output + it pushed in ``pushed``; the boundary then runs the all-reduce with ``prepushed`` = the ``ar`` tuple's mode.""" + + def __init__(self): + self.target = None + self.pushed = None # data_ptr of the pushed output + + def arm(self, target): + self.target = target + self.pushed = None + + def take_target(self): + target, self.target = self.target, None + return target + + def take_pushed(self, x): + pushed, self.pushed = self.pushed, None + return pushed is not None and x.data_ptr() == pushed + + +AR_PUSH = ArPushSlot() + + +class MoePushSlot: + """The MoE boundary's all-reduce exchange for its two producers to push into: the routed-expert engine (set A, + ``engine`` = its ``arp`` descriptor) and the shared expert's down projection (set B, ``shared`` = m3_mx_gemv + ``ar`` tuple). Each marks the output it pushed; the boundary then runs the all-reduce with ``prepushed`` = 3.""" + + def __init__(self): + self.engine = None + self.shared = None + self.engine_pushed = None # data_ptr of the pushed routed output + self.shared_pushed = None # data_ptr of the pushed shared output + + def arm(self, engine, shared): + self.engine, self.shared = engine, shared + self.engine_pushed = self.shared_pushed = None + + def disarm(self): + self.engine = self.shared = None + + def take_engine(self): + engine, self.engine = self.engine, None + return engine + + def take_shared(self): + shared, self.shared = self.shared, None + return shared + + def take_pushed(self, shared, routed): + """(routed pushed, shared pushed) for these outputs; clears the marks.""" + e = routed is not None and self.engine_pushed == routed.data_ptr() + s = shared is not None and self.shared_pushed == shared.data_ptr() + self.engine_pushed = self.shared_pushed = None + return e, s + + +MOE_PUSH = MoePushSlot() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_router_gemv.py new file mode 100644 index 000000000000..30855eb7f05d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_router_gemv.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-M fp32 router projection: ``logits[M, E] = float(x[M, K]) @ W[E, K]^T`` in one kernel. + +Replaces the bf16->fp32 cast, the TF32 GEMM and its split-K reduce on decode +steps. One CTA per ``ROWS`` experts: the fp32 weight rows are bulk-copied into +shared memory before ``griddepcontrol.wait`` (overlapping the producer kernel +under PDL), the bf16 activations after it. Products and sums are exact fp32 FMA. + +The kernels are ``torch.ops.trtllm.minimax_m3_router_gemv`` / ``_router_gemv_fp4`` / ``_router_tc``, reached +through the ``ops/moe`` entries. +""" + +from typing import Optional + +import torch + +from ..ops.moe.m3_router_gemv import m3_router_gemv as _router_gemv_op +from ..ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 as _router_gemv_fp4_op +from ..ops.moe.m3_router_tc import m3_router_tc as _router_tc_op + +# Opt-in dynamic shared memory left after the kernel's static reduction buffer (SM100: 227 KiB). +_MAX_DYNAMIC_SMEM = 227 * 1024 - 1024 # capacity: fixed: the device's shared-memory limit + + +def m3_router_gemv( + x: torch.Tensor, weight: torch.Tensor, rows: int = 1, use_pdl: bool = False +) -> Optional[torch.Tensor]: + """fp32 ``x @ weight.T`` for bf16 ``x`` with at most 16 rows; None when unsupported (only ``rows=1`` is built).""" + if ( + rows != 1 + or not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 16 + or weight.dtype != torch.float32 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or weight.shape[0] % rows != 0 + or x.data_ptr() % 16 != 0 + or (x.stride(0) * 2) % 16 != 0 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or rows * x.shape[1] * 4 + x.shape[0] * x.shape[1] * 2 > _MAX_DYNAMIC_SMEM + ): + return None + return _router_gemv_op(x, weight, rows, use_pdl) + + +def m3_router_gemv_fp4( + x: torch.Tensor, weight: torch.Tensor, sf_scale: torch.Tensor, use_pdl: bool = False +) -> Optional[tuple]: + """Router logits plus the NVFP4 activation the routed experts consume (bit-identical to + ``fp4_quantize(x, sf_scale, 16, False, False)``): (logits [M, E] fp32, x_fp4 [M, K / 2] uint8, + x_sf [M * K / 16] uint8 linear). None when unsupported.""" + if ( + not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 16 + or weight.dtype != torch.float32 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or x.data_ptr() % 16 != 0 + or (x.stride(0) * 2) % 16 != 0 + or sf_scale is None + or sf_scale.dtype != torch.float32 + or sf_scale.numel() != 1 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or x.shape[1] * 4 + x.shape[0] * x.shape[1] * 2 > _MAX_DYNAMIC_SMEM + ): + return None + return tuple(_router_gemv_fp4_op(x, weight, sf_scale.reshape(1), use_pdl)) + + +def m3_router_tc( + x: torch.Tensor, + weight: torch.Tensor, + sf_scale: Optional[torch.Tensor] = None, + use_pdl: bool = False, +) -> Optional[tuple]: + """Router logits (TF32 tensor cores, like the native TF32 GEMM) for 1..48 tokens with 128 experts, plus the + NVFP4 activation when ``sf_scale`` is given. Returns (logits,) or (logits, x_fp4, x_sf); None when unsupported.""" + if ( + not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 48 + or weight.dtype != torch.float32 + or weight.shape[0] != 128 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or (x.stride(0) * 2) % 16 != 0 + or x.data_ptr() % 16 != 0 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or 16 * (x.shape[1] // 8 + 4) * 4 + x.shape[0] * (x.shape[1] // 8 + 8) * 2 + > _MAX_DYNAMIC_SMEM + or (sf_scale is not None and (sf_scale.dtype != torch.float32 or sf_scale.numel() != 1)) + ): + return None + res = _router_tc_op(x, weight, None if sf_scale is None else sf_scale.reshape(1), use_pdl) + return tuple(res) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_sampled_accept.py new file mode 100644 index 000000000000..86b08a0b24fc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_sampled_accept.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). + +Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the +strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p +active must stay on the generic path; the caller decides. + +Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL +vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z +(the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it +is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens heavier +than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the nucleus +with probability >= top_p, so most rows decide on their second exchange. + +Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, +staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local +candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the +candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based Philox +keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one 4-word +record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and reads +the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the same +order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing and +two call halves keep a rank one call ahead from overwriting a slot still being read. +""" + +from typing import Optional + +import torch + +from ..ops.spec.m3_sampled_accept import m3_sampled_accept as _sampled_accept_op +from ..ops.spec.m3_sampled_accept_ws_words import m3_sampled_accept_ws_words as _ws_words_op + +# The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.cu) and reached through +# the ops/spec entry m3_sampled_accept. +MAX_ROWS = 48 # capacity: fixed: the kernel's built rows (decode steps) + + +class SampledAcceptWorkspace: + """Exchange slots behind one multicast mapping (two call halves), the (finished rows, call) counters, the per-row + token scratch and an extra-attempt counter. Created collectively, outside CUDA-graph capture.""" + + def __init__(self, uc, mc_ptr, rank, world, handle=None, comm=None): + self.uc, self.mc_ptr, self.rank, self.world = uc, mc_ptr, rank, world + self.handle, self.comm = handle, comm + dev = uc.device + self.counters = torch.zeros(2, dtype=torch.int32, device=dev) + self.tokens = torch.zeros(MAX_ROWS, dtype=torch.int32, device=dev) + self.extra_attempts = torch.zeros(1, dtype=torch.int32, device=dev) + + @classmethod + def create(cls, mapping, device) -> "SampledAcceptWorkspace": + words = int(_ws_words_op(max(mapping.tp_size, 1))) + if mapping.tp_size == 1: + uc = torch.zeros(words, dtype=torch.int64, device=device) + return cls(uc, uc.data_ptr(), 0, 1) + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_disabled + + comm = _get_mnnvl_workspace_comm(mapping) + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + if not mpi_disabled(): + comm.Free() + raise RuntimeError( + "SampledAcceptWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 8, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("SampledAcceptWorkspace.create failed on another rank") + return cls(uc, mc.data_ptr(), mapping.tp_rank, mapping.tp_size, handle, comm) + + +_WS: Optional[SampledAcceptWorkspace] = None + + +def workspace(mapping, device) -> Optional[SampledAcceptWorkspace]: + global _WS + if _WS is None: + if torch.cuda.is_current_stream_capturing(): + # A capture that finds it missing is the caller's miss (it is built on the eager steps). + return None + _WS = SampledAcceptWorkspace.create(mapping, device) + return _WS + + +def sampled_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + temperatures: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: Optional[torch.Tensor], + seeds: torch.Tensor, + offsets: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, + ws: Optional[SampledAcceptWorkspace] = None, + use_pdl: bool = True, +): + """Sampled verification (temperature + top-p) from this rank's bf16 logits shard [nreq * (draft_len + 1), V / tp]: + outputs as shard_accept (accepted [nreq, D+1], num_accepted [nreq], gather_ids [nreq]); None + when unavailable. Per-row sampling parameters and Philox (seed, offset) are the spec metadata's, one per row. + top_ks / min_ps are not applied: rows with top-k or min-p active must be routed to the generic path by the + caller.""" + del top_ks, min_ps + nreq, draft_len = draft_tokens.shape + rows = nreq * (draft_len + 1) + if logits.dim() != 2 or logits.shape[0] != rows or rows > MAX_ROWS or logits.shape[1] > 51200: + return None + if logits.dtype != torch.bfloat16 or logits.stride(1) != 1: + return None + ws = ws if ws is not None else workspace(mapping, logits.device) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + _sampled_accept_op( + logits, + int(vocab_offset), + temperatures[:rows].float().contiguous(), + top_ps[:rows].float().contiguous(), + seeds[:rows].long().contiguous(), + offsets[:rows].long().contiguous(), + ws.uc, + ws.mc_ptr, + ws.counters, + ws.tokens, + ws.rank, + ws.world, + draft_tokens.int().contiguous(), + batch_idx.int().contiguous(), + accepted, + num_accepted, + gather_ids, + int(force_total), + next_new, + force_vec, + ws.extra_attempts, + bool(use_pdl), + ) + return accepted, num_accepted, gather_ids diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_shard_accept.py new file mode 100644 index 000000000000..21521591bfa9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_shard_accept.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused greedy Eagle3 verification for the MiniMax-M3 target head (vocab-sharded TP head, verify rows <= 48). + +Replaces the all-gather of the target logits, the argmax and the strict-acceptance ops of a greedy verify step with +one kernel over this rank's BF16 logits shard (the target lm_head's own local GEMM output): + +* Every logit becomes a 64-bit key ``ordered(value) << 32 | ~global_index`` (ties pick the lowest index, as + ``torch.argmax``), max-reduced per row: warp shuffles, then ``atomicMax``. +* The last CTA to finish pushes the rank's per-row key into every rank's exchange slot through a multicast mapping + (Lamport: 0 = empty, the key is never 0), polls its own slots for all ranks, takes the global argmax, re-arms the + slots and runs the strict acceptance. Call n uses slot half n % 2. + +The argmax is taken over the BF16 logits without arithmetic, so the tokens are those of the in-tree path's argmax +over the gathered logits. The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/ +shardAccept.cu) and reached through the ``ops/spec`` entry ``m3_shard_accept``. +""" + +from dataclasses import dataclass +from typing import Any, Optional + +import torch + +from ..ops.spec.m3_shard_accept import m3_shard_accept as _shard_accept_op + +#: Verify rows per call, and exchange rows per call half of the workspace. +SHARD_ROWS = 48 # capacity: fixed: DECODE_MAX_TOKENS verify rows + + +@dataclass(eq=False) +class AcceptWorkspace: + """Per-rank exchange slots ([2 halves][SHARD_ROWS rows][world] uint64) behind one multicast mapping, the local + key accumulators and the (finished CTAs, call) counters.""" + + uc: torch.Tensor + mc_ptr: int + keys: torch.Tensor + counters: torch.Tensor + rank: int + world_size: int + handle: Any = None + comm: Any = None + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def create_workspace(mapping, device, slot_rows: int) -> AcceptWorkspace: + """Exchange slots ([2 halves][slot_rows rows][world] uint64) behind one multicast mapping over the TP group, the + local key accumulators [slot_rows] and the (finished CTAs, call) counters. Collective over the TP group; call + outside CUDA-graph capture, in the same order on every rank.""" + words = 2 * slot_rows * max(mapping.tp_size, 1) + if mapping.tp_size > 1: + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_disabled + + comm = _get_mnnvl_workspace_comm(mapping) + if not _mnnvl_workspace_all_succeeded(comm, True): + if not mpi_disabled(): + comm.Free() + raise RuntimeError("exchange workspace creation failed on another rank") + handle = _make_mnnvl_mcast_buffer( + comm, max(words * 8, 4096), mapping, mapping.is_multi_node() + ) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("exchange workspace creation failed on another rank") + mc_ptr, rank = mc.data_ptr(), mapping.tp_rank + else: + uc = torch.zeros(words, dtype=torch.int64, device=device) + mc_ptr, rank, handle, comm = uc.data_ptr(), 0, None, None + return AcceptWorkspace( + uc=uc, + mc_ptr=mc_ptr, + keys=torch.zeros(slot_rows, dtype=torch.int64, device=uc.device), + counters=torch.zeros(2, dtype=torch.int32, device=uc.device), + rank=rank, + world_size=max(mapping.tp_size, 1), + handle=handle, + comm=comm, + ) + + +def _workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace, created collectively outside CUDA-graph capture on first use.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + # A capture that finds it missing is the caller's miss (it is built on the eager steps). + return None + _WS = create_workspace(mapping, device, SHARD_ROWS) + return _WS + + +def shard_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, +): + """Greedy verification from this rank's BF16 logits shard. + + Args: + logits: [nreq * (draft_len + 1) <= SHARD_ROWS, V / tp] BF16, unit column stride: this rank's columns + (global ids ``vocab_offset ..``) of the verify rows, request-major. + draft_tokens: [nreq, draft_len] int32, the drafts being verified. + batch_idx: [nreq] int32, each request's row in the verify layout. + force_total / force_vec: forced acceptance (see ``m3_shard_accept``). + next_new: optional int32 [nreq, >= 1]; column 0 receives the last accepted token. + + Returns: + (accepted_tokens [nreq, draft_len + 1] int32, num_accepted [nreq] int32, first-pass gather ids [nreq] + int64), or None when the shape is not served or the workspace is missing during a capture. Collective over + the TP group. + """ + nreq, draft_len = draft_tokens.shape + if ( + logits.dim() != 2 + or logits.shape[0] != nreq * (draft_len + 1) + or logits.shape[0] > SHARD_ROWS + or logits.dtype != torch.bfloat16 + or logits.stride(1) != 1 + ): + return None + ws = _workspace(mapping, logits.device) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + _shard_accept_op( + logits, + vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_tokens.contiguous(), + batch_idx, + accepted, + num_accepted, + gather_ids, + force_total, + next_new, + _num_sms(), + force_vec, + ) + return accepted, num_accepted, gather_ids diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_spec_tail.py new file mode 100644 index 000000000000..0d4dc9cd1b64 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_spec_tail.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Small fused kernels of the MiniMax-M3 Eagle3 step (GB300, decode). + +``chunk_norm`` / ``capture_norm`` apply the drafter's per-layer fc_norm RMSNorms to the captured target hidden states +(the in-tree ``apply_eagle3_fc`` norms, same fp32 arithmetic); ``draft_input`` is a fused draft pass's input (the +embedding lookup, the drafter layer's input_layernorm / hidden_norm and their concat, same arithmetic); +``step_gather`` / ``slot_scatter`` move the step's integer outputs between the spec sampler's store and the engine's +input buffers. +""" + +from typing import Optional, Tuple + +import torch + +# The kernels are compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu) and reached through the +# ops/spec entries. +from ..ops.spec.m3_capture_norm import m3_capture_norm as _capture_norm_op +from ..ops.spec.m3_chunk_norm import m3_chunk_norm as _chunk_norm_op +from ..ops.spec.m3_draft_input import m3_draft_input as _draft_input_op +from ..ops.spec.m3_slot_scatter import m3_slot_scatter as _slot_scatter_op +from ..ops.spec.m3_step_gather import m3_step_gather as _step_gather_op + + +def draft_input( + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + eps_embed: float, + w_hidden: torch.Tensor, + eps_hidden: float, + use_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """(cat([rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden], -1) [M, 2H], hidden[rows] [M, H]) for + int32 ids [M] and int64 rows [M] (rows None: hidden[:M]); ids outside the table embed to zero.""" + m, h = ids.shape[0], table.shape[1] + out = torch.empty((m, 2 * h), dtype=table.dtype, device=table.device) + residual = torch.empty((m, h), dtype=table.dtype, device=table.device) + return _draft_input_op( + ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, use_pdl + ) + + +def chunk_norm( + src: torch.Tensor, weights: torch.Tensor, eps: float, use_pdl: bool = True +) -> torch.Tensor: + """Per-chunk RMSNorm of src [M, >= C*H] (row stride may exceed C*H) with weights [C, H] -> [M, C*H].""" + c, h = weights.shape + out = torch.empty((src.shape[0], c * h), dtype=src.dtype, device=src.device) + return _chunk_norm_op(src, weights, eps, out, use_pdl) + + +# The overlap scheduler's per-step gathers and the spec sampler's slot stores as one kernel each (step_gather / +# slot_scatter) instead of ~10 and 4 small torch ops; called by the target's SpecStepFold. + + +def step_gather( + new_tok, next_draft, lens, slots, ppi, P, T, D, ids, ids0, draft, d0, pos_off, kv_off, b0 +) -> bool: + """See step_gather_kernel; False when the inputs are outside what the kernel covers (caller falls back).""" + ts = (new_tok, next_draft, lens, slots, ppi, ids, draft, pos_off, kv_off) + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in ts) + or next_draft.dim() != 2 + or new_tok.dim() < 2 + ): + return False + if (new_tok.dim() == 3 and new_tok.shape[2] != 1) or not ( + lens.is_contiguous() + and ids.is_contiguous() + and draft.is_contiguous() + and pos_off.is_contiguous() + and kv_off.is_contiguous() + ): + return False + _step_gather_op( + new_tok, next_draft, lens, slots, ppi, P, T, D, ids, ids0, draft, d0, pos_off, kv_off, b0 + ) + return True + + +def slot_scatter( + slots, src_new, dst_new, src_next_new, dst_next_new, src_lens, dst_lens, src_draft, dst_draft +) -> bool: + """See slot_scatter_kernel (2-D int32 views, zero padding); False to fall back.""" + ts = (src_new, dst_new, src_next_new, dst_next_new, src_lens, dst_lens, src_draft, dst_draft) + if any(t.dtype != torch.int32 or not t.is_cuda for t in ts) or slots.dtype != torch.long: + return False + two = (src_new, dst_new, src_next_new, dst_next_new, src_draft, dst_draft) + if any(t.dim() != 2 for t in two) or src_lens.dim() != 1 or dst_lens.dim() != 1: + return False + _slot_scatter_op( + slots, + src_new, + dst_new, + src_next_new, + dst_next_new, + src_lens, + dst_lens, + src_draft, + dst_draft, + ) + return True + + +def capture_norm( + h: torch.Tensor, + residual: Optional[torch.Tensor], + weight: torch.Tensor, + eps: float, + dst: torch.Tensor, +) -> None: + """dst [M, H] (row stride may exceed H) = rmsnorm(h + residual) * weight with chunk_norm's rounding.""" + _capture_norm_op(h, residual, weight, eps, dst) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_target_head.py new file mode 100644 index 000000000000..f5a1b229c777 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_target_head.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target lm_head shard of a verify step of at most 8 rows, streamed from a lossless 12-bit coding of its BF16 +weight (75% of the bytes; ``ops.gemm.m3_target_head`` has the format). + +``prepare_head`` codes this rank's shard once, before the executor's warmup, and keeps it only when it decodes to the +BF16 weight bit for bit (``round_trip_ok``); a shard that does not keeps the in-tree BF16 GEMM. ``logits`` runs the +kernel: BF16 x BF16 MMAs with fp32 accumulation over the full K and one BF16 rounding per logit, the in-tree head's +arithmetic in another summation order. The kernel is compiled into TensorRT-LLM +(cpp/tensorrt_llm/kernels/minimaxM3/targetHead.cu) and called through the catalog entry ``ops.gemm.m3_target_head``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from ..ops.gemm.m3_target_head import CTA_WARPS, MAX_ROWS, encode_chunked, round_trip_ok +from ..ops.gemm.m3_target_head import m3_target_head as _target_head_op + + +@dataclass(eq=False) +class TargetHead: + """One rank's target lm_head shard in the 12-bit coding, and the kernel's split-tile scratch. + + ``coded`` is False when the coding did not reproduce the weight bit for bit (the in-tree GEMM then runs).""" + + weight: torch.Tensor + coded: bool + w: Optional[torch.Tensor] = None + base: int = 0 + exc_off: Optional[torch.Tensor] = None + exc_pos: Optional[torch.Tensor] = None + exc_w: Optional[torch.Tensor] = None + part: Optional[torch.Tensor] = None + counter: Optional[torch.Tensor] = None + grid: int = 0 + + +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def head_supported(lm_head) -> bool: + """A bias-free 2-D BF16 head of whole 16-row tiles whose K the coding covers (a multiple of 64).""" + weight = getattr(lm_head, "weight", None) + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: + return False + return ( + getattr(lm_head, "bias", None) is None + and weight.shape[0] % 16 == 0 + and weight.shape[1] % 64 == 0 + ) + + +def prepare_head(lm_head, grid: Optional[int] = None) -> TargetHead: + """Code ``lm_head``'s local shard and check it decodes to the BF16 weight bit for bit (fails closed: a shard that + does not is returned with ``coded=False`` and nothing kept). Built once, eagerly, before the executor's warmup + (``M3Eagle3Worker.prebuild_target_head``).""" + weight = lm_head.weight + with torch.no_grad(): + w = weight.detach() + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + if not round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w): + return TargetHead(weight=weight, coded=False) + if grid is None: + grid = _num_sms() + n = weight.shape[0] + return TargetHead( + weight=weight, + coded=True, + w=packed, + base=base, + exc_off=exc_off, + exc_pos=exc_pos, + exc_w=exc_w.view(-1), + part=torch.zeros(grid * CTA_WARPS * 2 * 4 * 32, dtype=torch.float32, device=weight.device), + counter=torch.zeros((n + 15) // 16, dtype=torch.int32, device=weight.device), + grid=grid, + ) + + +def logits(head: TargetHead, x: torch.Tensor, use_pdl: bool = True) -> torch.Tensor: + """This rank's BF16 logits shard [rows, n] of the rows x [rows <= MAX_ROWS, K] (BF16) from the coded shard.""" + assert head.coded and 1 <= x.shape[0] <= MAX_ROWS, (head.coded, x.shape) + if x.stride(1) != 1 or x.stride(0) % 8 or x.data_ptr() % 16: + x = x.contiguous() + n = head.weight.shape[0] + out = torch.empty((x.shape[0], n), dtype=torch.bfloat16, device=x.device) + return _target_head_op( + head.w, + x, + out, + head.part, + head.counter, + head.exc_off, + head.exc_pos, + head.exc_w, + head.base, + n, + head.grid, + use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/__init__.py new file mode 100644 index 000000000000..77d7b84d9bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Mixture-of-experts entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe16_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe16_engine.py new file mode 100644 index 000000000000..950891839022 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe16_engine.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3's MoE layer for 9..13 decode rows in one CuTe DSL launch (``m3_moe16``): the router, the routed NVFP4 +experts and the shared MXFP8 expert, on caller-owned state (an ``M3Moe16State`` and one ``M3MoeLayer`` per MoE +layer).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..m3_moe.op import M3MoeLayer +from ..m3_moe.op16 import M3Moe16State +from .m3_moe_engine import moe_reference + +__all__ = ["M3Moe16State", "m3_moe16_engine"] + + +class M3Moe16Engine(OpWrapper): + """One MiniMax-M3 MoE layer for M (9..13 here; the state takes 9..16) decode rows: + ``layer.state.run(layer, hidden, out, arp)``, bf16(shared + routed) [M, 6144]. + + The same arguments and the same math as ``m3_moe_engine`` (the routed part operation for operation, the shared + expert always folded in) on an ``M3Moe16State``: a different schedule for more rows. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # One build per row count (m_max is a trace-time constant); the target routes 9..13 rows here (spec-decode verify + # of 2 x 6 rows lands at 12). + CELLS: tuple[Cell, ...] = ( + Cell(why="9 rows: the smallest build, a partial last slot group", spec=dict(tokens=9)), + Cell(why="10 rows: its own build (m_max is a trace-time constant)", spec=dict(tokens=10)), + Cell(why="11 rows: its own build, an odd row count", spec=dict(tokens=11)), + Cell(why="12 rows: two requests verifying 5-token drafts", spec=dict(tokens=12)), + Cell(why="13 rows: the largest the target routes here", spec=dict(tokens=13)), + ) + + note = """ + Bands as `m3_moe_engine` (same approximations, same quantization boundaries). Not covered: 14..16 rows, which the + state takes and this target never routes here; the all-reduce push (`arp`). No `is_valid`, as `m3_moe_engine`. + """ + + def __call__( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return layer.state.run(layer, hidden, out, arp=arp) + + def reference( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``m3_moe_engine``'s reference with the shared expert; `out` and `arp` are ignored.""" + return moe_reference(hidden, layer, shared=True) + + +m3_moe16_engine = M3Moe16Engine() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe_engine.py new file mode 100644 index 000000000000..b6b0111b949f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe_engine.py @@ -0,0 +1,326 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3's MoE layer for 1..8 decode tokens in one CuTe DSL launch (``m3_moe``): the router, the routed NVFP4 +experts and the shared MXFP8 expert, on caller-owned state (an ``M3MoeState`` and one ``M3MoeLayer`` per MoE layer).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.quantization.utils.fp4_utils import ( + get_reorder_rows_for_gated_act_gemm_row_indices, + get_shuffle_matrix_a_row_indices, +) + +from ..m3_moe.op import ( + HIDDEN, + INTER, + SHARED_INTER, + TOP_K, + M3MoeLayer, + M3MoeState, + ar_push_desc, + is_supported, +) + +__all__ = [ + "M3MoeLayer", + "M3MoeState", + "ar_push_desc", + "is_supported", + "m3_moe_engine", + "moe_reference", + "moe_route", + "mxfp8_dequant", + "mxfp8_quant_dequant", + "unswizzle_sf", +] + +# The model's constants the kernels are built with: routed_scaling_factor, the shared expert's SwiGLU-OAI +# (swiglu_alpha, swiglu_limit, the (up + 1) offset). +ROUTED_SCALING = 2.0 +SHARED_ALPHA = 1.702 +SHARED_LIMIT = 7.0 + +_LOG2E = 1.4426950216293334961 +_RCP6 = 0.16666667163372039795 # fp32(1 / 6), the kernels' reciprocal of the E2M1 maximum +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) +_F32_MAX = 3.40282346638528859812e38 + + +def _f32(x: float) -> torch.Tensor: + return torch.tensor(x, dtype=torch.float32) + + +def _e2m1_decode(packed: torch.Tensor) -> torch.Tensor: + """fp32 values of packed E2M1 bytes [..., K / 2] (low nibble first) -> [..., K].""" + lut = torch.tensor(_E2M1, dtype=torch.float32, device=packed.device) + lo = lut[(packed & 0xF).long()] + hi = lut[(packed >> 4).long()] + return torch.stack((lo, hi), dim=-1).flatten(-2) + + +def _e2m1_round(x: torch.Tensor) -> torch.Tensor: + """x (fp32) to the nearest E2M1 value, ties to even, saturating at 6 (cvt.rn.satfinite.e2m1x2).""" + a = x.abs() + q = torch.full_like(a, 6.0) + for hi, value, inclusive in ( + (5.0, 4.0, True), + (3.5, 3.0, False), + (2.5, 2.0, True), + (1.75, 1.5, False), + (1.25, 1.0, True), + (0.75, 0.5, False), + (0.25, 0.0, True), + ): + q = torch.where(a <= hi if inclusive else a < hi, torch.full_like(a, value), q) + return torch.where(x < 0, -q, q) + + +def _e4m3_round(x: torch.Tensor) -> torch.Tensor: + """x (fp32) to E4M3, nearest even, saturating (cvt.rn.satfinite.e4m3x2), as fp32.""" + return x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() + + +def _ue8m0_up(x: torch.Tensor) -> torch.Tensor: + """UE8M0 code of x >= 0 rounded toward +inf, saturating (cvt.rp.satfinite.ue8m0x2), int32.""" + bits = x.contiguous().view(torch.int32) + code = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) + return code.clamp(0, 254) + + +def unswizzle_sf(sf: torch.Tensor, rows: int, blocks: int) -> torch.Tensor: + """The [rows, blocks] scale bytes of a 128x4-swizzled scale buffer (int64 indices, any byte dtype).""" + flat = sf.reshape(-1).view(torch.uint8) + r = torch.arange(rows, device=flat.device).view(-1, 1) + b = torch.arange(blocks, device=flat.device).view(1, -1) + padded_blocks = (blocks + 3) // 4 * 4 + idx = ( + (r // 128) * (128 * padded_blocks) + + (b // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (b % 4) + ) + return flat[idx] + + +def mxfp8_dequant(q: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: + """fp32 values of e4m3 ``q`` [R, K] with its 128x4-swizzled UE8M0 scales (one per 32 elements).""" + rows, k = q.shape + codes = unswizzle_sf(sf, rows, k // 32).float() + return q.view(torch.float8_e4m3fn).float() * torch.exp2(codes - 127.0).repeat_interleave( + 32, dim=1 + ) + + +def mxfp8_quant_dequant(x: torch.Tensor) -> torch.Tensor: + """x [R, K] through MXFP8 (one UE8M0 scale per 32, the smallest power of two >= amax / 448; e4m3 data), back to + fp32 -- mxfp8_quantize's recipe, which every MXFP8 activation here reproduces.""" + rows, k = x.shape + blocks = x.float().view(rows, k // 32, 32) + code = _ue8m0_up(blocks.abs().amax(dim=-1) * _f32(1.0 / 448.0).to(x.device)) + scale = torch.exp2(code.float() - 127.0).unsqueeze(-1) + return (_e4m3_round(blocks / scale) * scale).view(rows, k) + + +def _bf16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.bfloat16).float() + + +def _gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """a [M, K] @ b [N, K]^T, accumulated in fp64 and rounded once to fp32.""" + return (a.double() @ b.double().t()).float() + + +def _row_perm(rows: int, gated: bool) -> torch.Tensor: + """The TRTLLM-Gen row order: stored row j holds logical row perm[j] (the gated interleave, then the 32-row shuffle + of epilogue tile 128).""" + probe = torch.empty(rows, 1) + perm = get_shuffle_matrix_a_row_indices(probe, epilogue_tile_m=128) + if gated: + perm = get_reorder_rows_for_gated_act_gemm_row_indices(probe)[perm] + return perm + + +def _nvfp4_expert(w: torch.Tensor, sf: torch.Tensor, perm_inv: torch.Tensor) -> torch.Tensor: + """One expert's stored NVFP4 matrix (values [R, K / 2], swizzled e4m3 scales) as raw fp32 [R, K] in logical row + order (the global scale is not applied: it is the per-expert alpha).""" + rows, half = w.shape + codes = unswizzle_sf(sf, rows, 2 * half // 16).view(torch.float8_e4m3fn).float() + vals = _e2m1_decode(w) * codes.repeat_interleave(16, dim=1) + return vals[perm_inv] + + +def moe_route(x: torch.Tensor, router_weight: torch.Tensor, bias: torch.Tensor): + """MiniMax2 routing on the fp32 router GEMV: (top-4 expert ids [M, 4] int64, bf16-rounded weights [M, 4] fp32).""" + logits = _gemm(x, router_weight) + sig = (0.5 * torch.tanh(0.5 * logits).double() + 0.5).float() + keys = sig + bias + ids = torch.sort(keys, dim=-1, descending=True, stable=True).indices[:, :TOP_K] + sk = keys.gather(1, ids) - bias[ids] + total = (sk[:, 0] + sk[:, 2]) + (sk[:, 1] + sk[:, 3]) + w = (sk * _f32(ROUTED_SCALING).to(x.device)) / (total + _f32(1e-20).to(x.device)).unsqueeze(1) + return ids, _bf16(w) + + +def _nvfp4_act(x: torch.Tensor, sc: torch.Tensor) -> torch.Tensor: + """The activation's NVFP4 quantization (fp4_quantize(x, sc, 16)), as raw fp32 values (q x block scale).""" + rows, k = x.shape + blocks = x.view(rows, k // 16, 16) + vmax = blocks.abs().amax(dim=-1, keepdim=True) + sff = _e4m3_round(sc * (vmax * _f32(_RCP6).to(x.device))) + # The kernels' rcp.approx(sf x rcp.approx(sc)) lands where the correctly rounded sc / sf does on every element + # measured (the 1 / (sf x (1 / sc)) spelling differs on ~0.06% of them). + osc = torch.where(vmax != 0, sc / sff, torch.zeros_like(vmax)) + return (_e2m1_round(blocks * osc) * sff).view(rows, k) + + +def _fc1_epilogue(up: torch.Tensor, gate: torch.Tensor, esc: torch.Tensor) -> torch.Tensor: + """TRTLLM-Gen's SwiGLU-OAI FC1 epilogue on the raw FC1 sums (fp32 [I]) and its NVFP4 requantization: the raw FC2 + input values. ``esc`` = (fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha) of the expert, with + act_beta and act_clamp already divided by fc31_alpha (the loader's convention).""" + a, c, alpha, beta, lim = (esc[i] for i in range(5)) + gc = torch.where(gate > lim, lim, gate) + gc = torch.where(gate < -_F32_MAX, torch.full_like(gc, -_F32_MAX), gc) + uc = torch.minimum(torch.maximum(up, -lim), lim) + c1 = a * (alpha * _f32(_LOG2E).to(up.device)) + sgm = 1.0 / (torch.exp2(-(gc * c1)) + 1.0) + h = (uc.double() * a.double() + (a * beta).double()).float() * (gc * sgm) + blocks = h.view(-1, 16) + amax = (blocks.abs() * c).amax(dim=-1, keepdim=True) + sff = _e4m3_round(amax * _f32(_RCP6).to(up.device)) + osc = torch.where(sff != 0, 1.0 / sff, torch.zeros_like(sff)) + return (_e2m1_round((blocks * osc) * c) * sff).view(-1) + + +def _shared_expert(x: torch.Tensor, shared_mxfp8) -> torch.Tensor: + """The shared expert as the target runs it outside the engine: MXFP8 gate_up (bf16 out), SwiGLU-OAI with the + model's alpha / limit / (up + 1) in fp32 (bf16 out), MXFP8 down (bf16 out).""" + gu_w, gu_sf, dn_w, dn_sf = shared_mxfp8 + gu = _bf16(_gemm(mxfp8_quant_dequant(x), mxfp8_dequant(gu_w, gu_sf))) + gate, up = gu[:, :SHARED_INTER], gu[:, SHARED_INTER:] + a = torch.minimum(gate, torch.full_like(gate, SHARED_LIMIT)) + ub = up.clamp(-SHARED_LIMIT, SHARED_LIMIT) + 1.0 + sg = 1.0 / ( + torch.exp2((a * _f32(-SHARED_ALPHA).to(x.device)) * _f32(_LOG2E).to(x.device)) + 1.0 + ) + h = _bf16(ub * (sg * a)) + return _bf16(_gemm(mxfp8_quant_dequant(h), mxfp8_dequant(dn_w, dn_sf))) + + +def moe_reference(hidden: torch.Tensor, layer: M3MoeLayer, shared: bool) -> torch.Tensor: + """The MoE layer of the decode engines (both token ranges compute this), in plain torch: bf16 [M, 6144].""" + x = hidden.float() + m = x.shape[0] + ids, weights = moe_route(x, layer.router_weight, layer.bias) + xr = _nvfp4_act(x, layer.fc31_input_scale.reshape(()).float()) + inv31 = torch.argsort(_row_perm(2 * INTER, gated=True)).to(x.device) + inv2 = torch.argsort(_row_perm(HIDDEN, gated=False)).to(x.device) + y = torch.zeros(m, TOP_K, HIDDEN, dtype=torch.float32, device=x.device) + for e in torch.unique(ids).tolist(): + w31 = _nvfp4_expert(layer.w3_w1_weight[e], layer.w3_w1_weight_scale[e], inv31) + w2 = _nvfp4_expert(layer.w2_weight[e], layer.w2_weight_scale[e], inv2) + esc = layer.escal[:, e] + for t, k in (ids == e).nonzero().tolist(): + acc1 = _gemm(xr[t : t + 1], w31)[0] + # The buffer holds [w3 (up); w1 (gate)]. + h = _fc1_epilogue(acc1[:INTER], acc1[INTER:], esc) + y[t, k] = _bf16(_gemm(h.unsqueeze(0), w2)[0] * esc[5]) + acc = torch.zeros(m, HIDDEN, dtype=torch.float64, device=x.device) + for k in range(TOP_K): + acc = (acc + weights[:, k : k + 1].double() * y[:, k].double()).float().double() + routed = acc.float() + if not shared: + return routed.to(torch.bfloat16) + return (_shared_expert(x, layer.shared_mxfp8) + _bf16(routed)).to(torch.bfloat16) + + +class M3MoeEngine(OpWrapper): + """One MiniMax-M3 MoE layer for M (1..8) decode tokens: ``layer.state.run(layer, hidden, out, arp)``. + + ``hidden`` bf16 [M, 6144] contiguous, M = the state's token count. ``layer`` (from ``M3MoeState.layer``) carries the + TRTLLM-Gen NVFP4 expert buffers read in place, the per-expert scalars, the router weight / bias, the activation's + NVFP4 global scale and, on a ``shared`` state, the shared expert's MXFP8 weights; its state carries the compiled + kernel and the workspace every layer on it shares. Returns bf16 [M, 6144]: bf16(shared + routed) on a ``shared`` + state, else the routed partial. ``out``: write there instead of a new tensor. ``arp`` (``ar_push_desc``): also + push the rows into a multicast all-reduce exchange. + + The routed part is TRTLLM-Gen's NVFP4 MoE on the fp32 router logits: sigmoid (tanh form) + correction bias top-4 + of 128 (ties to the lower id), weights sigmoid / sum * 2.0 rounded to bf16; the activation quantized to NVFP4 with + the layer's global scale; per expert the raw FC1 sums, SwiGLU-OAI (alpha / beta / clamp), the NVFP4 requant with + fc31_scale_c, FC2, bf16(fc2_alpha x sum); the fma chain over the top-4 in order. The shared expert is the target's + MXFP8 SwiGLU-OAI MLP. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # One cell per token count: m_max is a trace-time constant, so each is its own build, and the target builds all + # eight (decode batches of 1..8 requests, spec-decode verify rows up to 8). The whys name what the build differs in. + CELLS: tuple[Cell, ...] = ( + Cell( + why="1 token: pre-tiled shared expert, last phase balanced without trailing FC2", + spec=dict(tokens=1), + ), + Cell( + why="2 tokens: pre-tiled shared expert, the one build without AR_PRE", + spec=dict(tokens=2), + ), + Cell( + why="3 tokens: 4 FC2-B builder warps, the shared expert from swizzled scales", + spec=dict(tokens=3), + ), + Cell( + why="4 tokens: one token per routing warp, logits replicated 128x", spec=dict(tokens=4) + ), + Cell( + why="5 tokens: 5 builder warps, a routing warp takes two tokens, logits 32x", + spec=dict(tokens=5), + ), + Cell(why="6 tokens: one request verifying a 5-token draft", spec=dict(tokens=6)), + Cell(why="7 tokens: the B2IMG global FC1 -> FC2 hand-off", spec=dict(tokens=7)), + Cell( + why="8 tokens: B2IMG with a 6-stage ring, 32 distinct experts at most", + spec=dict(tokens=8), + ), + ) + + note = """ + The kernel's elementwise math uses the PTX approximations (rcp.approx, ex2.approx, tanh) where `reference` uses + torch's correctly rounded ones, and its GEMMs sum in tensor-core order where `reference` sums in fp64: a + quantization boundary (NVFP4 activation / intermediate code, routing near-tie) can fall on the other side, so the + two agree within the band the test derives, not bit for bit. The routed part is bit for bit the stock TRTLLM-Gen + NVFP4 MoE (fp4_quantize + fp4_block_scale_moe_runner) on the same routing; the test checks that too. + + Not covered: the all-reduce push (`arp`), which needs a multicast exchange (the decode forward's boundary all-reduce + exercises it); the non-shared state (`shared=False`), which this target never builds. Every layer on one state + must run in one stream order: the workspace is shared and tagged per call, never reset. + + No `is_valid`: the compiled call checks every buffer's dtype, shape and alignment against its signature and + raises on a mismatch (a misaligned `hidden`, an `out` that is not bf16 [M, 6144], a malformed `arp`). + """ + + def __call__( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return layer.state.run(layer, hidden, out, arp=arp) + + def reference( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Plain torch over the layer's buffers (see the class docstring); `out` and `arp` are ignored.""" + return moe_reference(hidden, layer, shared=layer.state.shared) + + +m3_moe_engine = M3MoeEngine() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_router_gemv.py new file mode 100644 index 000000000000..37870d0e6cd7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_router_gemv.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""fp32 router GEMV of the decode MoE: ``logits[M, E] = float(x) @ W^T`` at one token.""" + +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, +) + +_E4M3_MAX = 448.0 +_E2M1_MAX = 6.0 +#: Midpoints between consecutive e2m1 magnitudes; ``_TIE_UP[i]``: round-to-nearest-even takes the upper code. +_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) +_TIE_UP = (False, True, False, True, False, True, False) +#: The kernel builds each block's output scale through two ``rcp.approx.ftz.f32`` reciprocals (~2^-23 relative): +#: a value this close to an e2m1 midpoint may round to either neighbour, everything further out must match. +TIE_WINDOW = 2.0**-20 + + +def nvfp4_quantize( + x: torch.Tensor, sf_scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``fp4_quantize(x, sf_scale, 16, False, False)``: (packed e2m1 [M, K / 2], linear E4M3 scales [M * K / 16], + the pre-rounding scaled values [M, K] that ``near_tie`` reads). + + Per 16 columns: scale = e4m3(sf_scale * amax / 6), data = x * sf_scale / scale rounded to nearest-even e2m1. + """ + gs = sf_scale.float().reshape(()) + m, k = x.shape + blocks = x.float().view(m, k // 16, 16) + vecmax = blocks.abs().amax(-1) + sf = (gs * (vecmax / _E2M1_MAX)).clamp(max=_E4M3_MAX).to(torch.float8_e4m3fn) + out_scale = torch.where(vecmax != 0, gs / sf.float(), torch.zeros_like(vecmax)) + scaled = (blocks * out_scale.unsqueeze(-1)).view(m, k) + mid = torch.tensor(_E2M1_MIDPOINTS, device=x.device, dtype=torch.float32) + tie_up = torch.tensor(_TIE_UP, device=x.device, dtype=torch.bool) + a = scaled.abs().unsqueeze(-1) + code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8) + negative = (scaled < 0) | ((scaled == 0) & torch.signbit(scaled)) + code = code | (negative.to(torch.uint8) << 3) + packed = code[:, 0::2] | (code[:, 1::2] << 4) + return packed, sf.view(torch.uint8).reshape(-1), scaled + + +def near_tie(scaled: torch.Tensor) -> torch.Tensor: + """Elements whose e2m1 rounding the kernel's approximate reciprocals decide (see ``TIE_WINDOW``).""" + mid = torch.tensor(_E2M1_MIDPOINTS, device=scaled.device, dtype=torch.float32) + a = scaled.abs().unsqueeze(-1) + return ((a - mid).abs() <= TIE_WINDOW * a).any(-1) + + +def fp32_logits(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """``float(x) @ w^T`` as a true fp32 product (TF32 off).""" + with true_fp32_matmul(): + return x.float() @ w.t() + + +class M3RouterGemv(OpWrapper): + """fp32 router logits ``float(x[M, K]) @ w[E, K]^T`` with exact fp32 FMA, one CTA per expert. + + ``x`` is BF16 with 1..16 rows, unit column stride and 16-byte rows; ``w`` is the fp32 router weight; K a + multiple of 1024. ``rows`` experts per CTA; only 1 is built. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The target takes the GEMV at exactly one token (2..48 go to the tensor-core router): 128 experts, hidden 6144. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token, 128 experts x hidden 6144", + spec=dict(m=1, e=128, k=6144), + ), + ) + + note = """ + The op accepts up to 16 rows; the target sends one, so wider calls are not cells. The logits differ from a true + fp32 product only in accumulation order (per-thread FMA chains over K / 1024 strided column groups, a warp tree + and an 8-warp sum). + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_router_gemv(x, w, rows, pdl) + + def reference(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> torch.Tensor: + """True fp32 product; ``rows`` and ``pdl`` do not enter the arithmetic.""" + return fp32_logits(x, w) + + def is_valid(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> None: + # The activation rows are bulk-copied from x's base address; the op checks the row stride but not the base, + # and a misaligned bulk copy faults the context rather than raising. + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + + +m3_router_gemv = M3RouterGemv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_router_gemv_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_router_gemv_fp4.py new file mode 100644 index 000000000000..58f809e82cf6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_router_gemv_fp4.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""fp32 router GEMV plus the routed experts' NVFP4 activation, one token.""" + +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 + +from .m3_router_gemv import fp32_logits, nvfp4_quantize + + +class M3RouterGemvFp4(OpWrapper): + """``[logits fp32 [M, E], x_fp4 [M, K / 2] uint8, x_sf [M * K / 16] uint8]`` from one kernel. + + The logits are ``m3_router_gemv``'s; the NVFP4 activation is ``fp4_quantize(x, sf_scale, 16, False, False)`` + (16-element blocks, E4M3 scales in the linear layout, per-tensor ``sf_scale``) for the routed experts' FC1. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The fused-router decode step at one token: 128 experts, hidden 6144, the experts' input scale. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token: logits for 128 experts and the NVFP4 activation of hidden 6144", + spec=dict(m=1, e=128, k=6144, sf_scale=448.0 * 6.0 / 8.0), + ), + Cell( + why="a global scale that puts block scales near the E4M3 top, where the scale cast saturates", + spec=dict(m=1, e=128, k=6144, sf_scale=448.0 * 6.0 / 2.0), + ), + ) + + note = """ + The e2m1 codes are exact except within `TIE_WINDOW` of a rounding midpoint, where the kernel's approximate + reciprocals decide the direction. Up to 16 rows are accepted; the target sends one. + """ + + def __call__( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_router_gemv_fp4(x, w, sf_scale, pdl) + + def reference( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """True fp32 logits and the NVFP4 recipe; ``pdl`` does not enter the arithmetic.""" + packed, sf, _ = nvfp4_quantize(x, sf_scale) + return fp32_logits(x, w), packed, sf + + def is_valid(self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool) -> None: + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + # A zero or non-finite global scale makes every block scale 0 or NaN and the data codes garbage; the op + # reads the scalar on the device and cannot refuse it. + value = float(sf_scale.reshape(())) + assert value > 0.0 and value != float("inf"), ( + f"sf_scale must be positive and finite, got {value}" + ) + + +m3_router_gemv_fp4 = M3RouterGemvFp4() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_router_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_router_tc.py new file mode 100644 index 000000000000..e88af51f7fa2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_router_tc.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TF32 tensor-core router of the decode MoE for 2..48 tokens, optionally with the NVFP4 activation.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from .m3_router_gemv import nvfp4_quantize + + +def tf32_rna(w: torch.Tensor) -> torch.Tensor: + """fp32 -> TF32 as ``cvt.rna.tf32.f32``: keep 10 mantissa bits, round to nearest with ties away from zero.""" + bits = w.contiguous().view(torch.int32) + rounded = (bits + 0x1000) & ~0x1FFF + finite = (bits & 0x7F800000) != 0x7F800000 + return torch.where(finite, rounded, bits).view(torch.float32) + + +class M3RouterTc(OpWrapper): + """Router logits on TF32 tensor cores for 128 experts and 1..48 tokens: ``x @ tf32(w)^T``, fp32 accumulate. + + ``x`` is BF16 (exact in TF32), ``w`` the fp32 router weight, rounded to TF32 in the kernel. The 16 K slices' + partials are summed in slice order whatever cluster size the device co-schedules (16 or 8 CTAs), so the result + does not depend on it. With ``sf_scale`` also the NVFP4 activation as ``m3_router_gemv_fp4``: returns + ``[logits]`` or ``[logits, x_fp4, x_sf]``. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The target routes 2..48 tokens here: plain logits on the unfused path, logits + NVFP4 on the fused-router path. + CELLS: tuple[Cell, ...] = ( + Cell( + why="two tokens, logits only (the smallest call the target makes here)", + spec=dict(m=2, k=6144, fp4=False), + ), + Cell( + why="8 tokens with the NVFP4 activation: one 8-token MMA tile, every expert tile quantizes one row", + spec=dict(m=8, k=6144, fp4=True), + ), + Cell( + why="13 tokens: a partial second token tile", + spec=dict(m=13, k=6144, fp4=True), + ), + Cell( + why="the 48-token maximum, six token tiles, logits only", + spec=dict(m=48, k=6144, fp4=False), + ), + Cell( + why="the 48-token maximum with the NVFP4 activation", + spec=dict(m=48, k=6144, fp4=True), + ), + # Every other row count the router takes (dispatch.ROUTER_TC), with and without the NVFP4 activation. + *( + Cell( + why=f"{m} rows: a router row count between the ones above", + spec=dict(m=m, k=6144, fp4=m % 2 == 1), + ) + for m in range(2, 49) + if m not in (2, 8, 13, 48) + ), + ) + + note = """ + The kernel uses 16-CTA clusters when the device co-schedules eight of them and 8-CTA clusters otherwise; both + sum the same 16 slice partials in the same order. Only the one the test device picks is driven. + """ + + def __call__( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_router_tc(x, w, sf_scale, pdl) + + def reference( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> tuple[torch.Tensor, ...]: + """fp32 product of ``x`` and the TF32-rounded weight, plus the NVFP4 recipe with ``sf_scale``.""" + with true_fp32_matmul(): + logits = x.float() @ tf32_rna(w).t() + if sf_scale is None: + return (logits,) + packed, sf, _ = nvfp4_quantize(x, sf_scale) + return logits, packed, sf + + def is_valid( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> None: + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + if sf_scale is not None: + value = float(sf_scale.reshape(())) + assert value > 0.0 and value != float("inf"), ( + f"sf_scale must be positive and finite, got {value}" + ) + + +m3_router_tc = M3RouterTc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/msa_kv_update.py new file mode 100644 index 000000000000..8c2a67356d3f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/msa_kv_update.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One kernel for MSA's device-side KV-length correction (MsaAttentionMetadata.on_update_kv_lens). + +The torch formulation is ~23 one-CTA kernels (min, cumsum, gathers, div / clamp / remainder, the block-table lookup, +the valid-block counts) and runs at the head of every captured decode step, twice. This kernel computes the same +int32 results in one launch: + + kv_true[b] = min(kv_lens[b], staged[b]) -> seq_lens_out[b], cu_kv[b + 1] = prefix sum (cu_kv[0] = 0) + pos[t] = kv_true[r] - qo[r] + intra[t] (r = q_batch_row[t]) + out_loc[t] = table[r, min(max(pos, 0) // page, ncols - 1)] * page + max(pos, 0) % page + n_valid[t] = (max(pos + 1, 0) + page - 1) // page + +``msa_kv_update`` is the target's ``MsaKvUpdate`` (``tensorrt_llm._torch.pyexecutor.runtime_hooks``): the MSA +metadata calls it with its own buffers and falls back to the torch formulation when it returns False. The kernel is +``torch.ops.trtllm.minimax_m3_msa_kv_update``, called through this target's ``msa_kv_update`` entry. +""" + +from typing import Optional + +import torch + +from ..ops.attention.msa_kv_update import msa_kv_update as _kv_update_op + + +def msa_kv_update( + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], +) -> bool: + """The fused correction (see module doc); False, with nothing written, when an operand is outside what the + kernel covers. Arguments in ``MsaKvUpdate`` order; ``seq_lens`` are the per-request query lengths (``qo``).""" + ts = [ + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + msa_seq_lens, + out_cache_loc, + ] + ts += [t for t in (cu_kv_lens, n_valid_blocks) if t is not None] + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in ts) + or block_table.dim() != 2 + or block_table.stride(1) != 1 + ): + return False + if any(not t.is_contiguous() for t in ts if t is not block_table): + return False + if ( + kv_lens.numel() < batch + or kv_lens_staged.numel() < batch + or seq_lens.numel() < batch + or msa_seq_lens.numel() < batch + or q_batch_row.numel() < total_q + or q_intra.numel() < total_q + or out_cache_loc.numel() < total_q + or (cu_kv_lens is not None and cu_kv_lens.numel() < batch + 1) + or (n_valid_blocks is not None and n_valid_blocks.numel() < total_q) + ): + return False + _kv_update_op( + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + tokens_per_block, + batch, + total_q, + msa_seq_lens, + cu_kv_lens, + out_cache_loc, + n_valid_blocks, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mx_gemv/__init__.py new file mode 100644 index 000000000000..ef3a4c5847ea --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mx_gemv/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) as a CuTe DSL kernel (see op.py).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mx_gemv/m3_mx_gemv_kernel.py new file mode 100644 index 000000000000..e3d7751231ee --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mx_gemv/m3_mx_gemv_kernel.py @@ -0,0 +1,1238 @@ +# 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. +"""MXFP8 decode GEMV on tcgen05: ``y[M, N] = x[M, K] @ W[N, K]^T`` for M <= 8, block-scaled E4M3 x E4M3. + +The weight is the MMA's A operand (128 output features per CTA), the activation is B (8 token columns), +fp32 accumulation in TMEM. Both operands come in their MXFP8 storage: E4M3 values (K contiguous) and UE8M0 +scales in the 128x4-swizzled layout, whose 512-byte atom per (128 rows, 128 K) is exactly the source of a +``tcgen05.cp`` 32x128b.warpx4 scale copy, so both scale tensors are bulk-copied unchanged. For an activation of at +most 32 rows the atom of k-tile ``k`` sits at byte ``512 k``. + +A cluster of SPLIT CTAs shares one 128-row weight tile; rank r streams the interleaved k-tiles r, r + SPLIT, ... +through a RING-stage ring. Before the grid dependency each CTA issues its first RING weight tiles (and their scale +atoms) and prefetches the rest into L2; after it, the activation slice of the CTA (values and scale atoms) lands in +one buffer. Rank r owns rows [r * 128 / SPLIT, (r + 1) * 128 / SPLIT) of the tile: the other ranks store their fp32 +partials of those rows into its shared memory and arrive on its barrier (release, cluster scope); the owner adds the +SPLIT partials in rank order and stores bf16. + +Warps: 0 weight TMA, 1 activation TMA, 2 TMEM allocation + scale copies + MMA, 3 idle, 4-7 epilogue (warp 4 + i +reads TMEM lanes 32 i .. 32 i + 31). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass import dsl_user_op +from cutlass.experimental import primitives as prims + +CTA_M = 128 # weight rows per CTA = the MMA's M # capacity: fixed: a kernel tile constant +MMA_N = 8 # token columns +CTA_K = 128 # one k-tile: 128 E4M3 bytes per row = one 128-byte swizzle row # capacity: fixed: a kernel tile constant +MMA_INST_K = 32 +K_BLOCKS = ( + CTA_K // MMA_INST_K +) # 4 MMAs per k-tile, one per scale byte of the atom # capacity: fixed: a kernel tile constant +THREADS = 256 +SF_ATOM_BYTES = 512 +SF_ATOM_U16 = SF_ATOM_BYTES // 2 +SF_COLS = 4 # TMEM columns of one scale atom +W_TILE_BYTES = CTA_M * CTA_K +X_TILE_BYTES = MMA_N * CTA_K +EVICT_FIRST = 0x12F0000000000000 +A_COLS_PER_MMA = MMA_INST_K // 4 # TMEM columns of one 128 x K32 E4M3 granule of A + +a_dtype = cutlass.Float8E4M3FN +b_dtype = cutlass.Float8E4M3FN + + +def num_k_tiles(k_in: int) -> int: + return k_in // CTA_K + + +def acc_cols(nacc: int, mma_n: int = MMA_N) -> int: + # nacc independent accumulators of MMA_N columns: MMA j of the CTA accumulates into j % nacc, so the M128 x N8 + # MMAs do not form one serial dependency chain; the epilogue adds them. + return max(16, nacc * mma_n) + + +def a_col0(nacc: int, my_tiles: int, mma_n: int = MMA_N) -> int: + # a_tmem: [acc | SFA per k-tile | one SFB slot (reused in issue order) | A: 32 columns (128 x 128 E4M3) per k-tile] + return (acc_cols(nacc, mma_n) + (my_tiles + 1) * SF_COLS + 31) // 32 * 32 + + +def tmem_cols( + ring: int, nacc: int, a_tmem: bool = False, my_tiles: int = 0, mma_n: int = MMA_N +) -> int: + need = acc_cols(nacc, mma_n) + 2 * ring * SF_COLS + if a_tmem: + need = a_col0(nacc, my_tiles, mma_n) + my_tiles * K_BLOCKS * A_COLS_PER_MMA + cols = 32 + while cols < need: + cols *= 2 + return cols + + +@dsl_user_op +def read_globaltimer(*, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_async_v4(dst, w0, w1, w2, w3, mbar, *, loc=None, ip=None): + """st.async of four int32 words to a shared::cluster address, completing ``mbar`` there by 16 bytes.""" + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Int32(w0).ir_value(loc=loc, ip=ip), + cutlass.Int32(w1).ir_value(loc=loc, ip=ip), cutlass.Int32(w2).ir_value(loc=loc, ip=ip), + cutlass.Int32(w3).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "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, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _ld_shared_v4_f32(smem_ptr, *, loc=None, ip=None): + """ld.shared.v4.f32 of four fp32 at a 16-byte-aligned shared-memory pointer.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + out = _llvm.inline_asm( + _llvm.StructType.get_literal([_T.f32()] * 4), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", "=f,=f,=f,=f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + return [ + cutlass.Float32(_llvm.extractvalue(_T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4) + ] + + +@dsl_user_op +def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA completed by other CTAs' st.async.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@dsl_user_op +def _absmax(acc, v, *, loc=None, ip=None): + """max(acc, |v|).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(acc).ir_value(loc=loc, ip=ip), cutlass.Float32(v).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 a;\n\tabs.f32 a, $2;\n\tmax.f32 $0, $1, a;\n\t}", "=f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rcp_approx_ftz(x, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _e4m3x4(v0, v1, v2, v3, *, loc=None, ip=None): + """Four f32 to E4M3 (RN, satfinite), packed little-endian into one int32 word (v0 in the low byte).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n\t.reg .b16 lo, hi;\n\tcvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\tmov.b32 $0, {lo, hi};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).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, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16(x) (RN) as the fp32 bit pattern of the rounded value.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 h, z;\n\tcvt.rn.bf16.f32 h, $1;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _bf16_round(x): + return _bf16_bits(x).bitcast(cutlass.Float32) + + +@dsl_user_op +def _rope_bits(c, a, s, b, *, loc=None, ip=None): + """bf16(c * a + s * b) with separately rounded products (no FMA), as fp32 bits: the stock RoPE arithmetic.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (c, a, s, b)], + "{\n\t.reg .f32 p, q, r;\n\t.reg .b16 h, z;\n\tmul.rn.f32 p, $1, $2;\n\tmul.rn.f32 q, $3, $4;\n\t" + "add.rn.f32 r, p, q;\n\tcvt.rn.bf16.f32 h, r;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _din_load4(table, hidden, hidden_half, word_idx): + """Four int32 words at word_idx of the hidden rows (hidden_half == 1) or the embedding table: one CTA-uniform + base-address select instead of a branch.""" + addr = cutlass.select_( + hidden_half == cutlass.Int32(1), hidden.data_ptr(0).toint(), table.data_ptr(0).toint() + ) + ptr = cutlass.inttoptr( + addr + cutlass.Int64(word_idx) * cutlass.Int64(4), mem_space=1, dtype=cutlass.Int32 + ) + return ptr.load(count=4, alignment=16) + + +def _stamp(timers, timed: bool, bx, slot: int): + # Diagnostic: %globaltimer of milestone `slot` of CTA bx (timers [grid, 8] int64). + if timed: + timers.store(read_globaltimer(), idx=bx * cutlass.Int32(8) + cutlass.Int32(slot)) + + +def _k_tile(rank, i: int, split: int, my_tiles: int, contig: bool): + # contig: rank r streams k-tiles [r * my_tiles, (r + 1) * my_tiles); else the interleaved r, r + split, ... + if contig: + return rank * cutlass.Int32(my_tiles) + cutlass.Int32(i) + return rank + cutlass.Int32(i * split) + + +@cute.kernel +def m3_mx_gemv_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] E4M3 as bytes, box 128 x 128 + tma_desc_wsf: cutlass.GridConstant[ + cuda.TensorMap + ], # W scale atoms [N / 128 * K / 128, 256] u16 + tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] E4M3 as bytes, box 128 x 8 + tma_desc_xsf: cutlass.GridConstant[cuda.TensorMap], # x scale atoms [K / 128, 256] u16 + y: cutlass.Array, # [M * N] bf16, token-major + x_words: cutlass.Array, # x_bf16: the bf16 activation [M, K] as int32 words + timers: cutlass.Array, # [grid, 8] int64 milestones when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, # ar_push: the all-reduce workspace's multicast mapping (int32 words) + ar_flags: cutlass.Int64, # ar_push: its call flags (parity of the next call = flags[0] & 1) + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[ + int + ], # 1: push the output rows into the next all-reduce; 2: and trigger after it + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], # token columns of the MMA (8, 16 or 48): M <= mma_n + din_i: cutlass.Array, # din: int32 token ids [M] + din_table: cutlass.Array, # din: embedding [V, K / 2] bf16 as int32 words + din_hidden: cutlass.Array, # din: hidden rows [M, K / 2] bf16 as int32 words + din_w: cutlass.Array, # din: [embedding norm | hidden norm] weights [K] bf16 as int32 words + din_f: cutlass.Array, # din: fp32 [eps_e, eps_h] + din_vocab: cutlass.Int32, + din_g: cutlass.Array, # din, flag 1: int64 hidden row of each token, as int32 words + din_res: cutlass.Array, # din, flag 2: the gathered hidden rows out [M, K / 2] bf16 as int32 words + din_flags: cutlass.Int32, # din: 1 = gather the hidden rows, 2 = write them to din_res + rope_cs: cutlass.Array, # rope: fp32 (cos, sin) [positions][64][2] + rope_i: cutlass.Array, # rope: int32 cache lengths [M / T] + rope_bt: cutlass.Array, # rope: int32 K / V page tables [M / T][2][max_blocks] + rope_q: cutlass.Array, # rope: E4M3 Q out [M][hq][128] as uint8 + rope_pool: cutlass.Array, # rope: paged E4M3 cache [pages][hkv][tokens per page][128] as uint8 + rope_scale: cutlass.Array, # rope: fp32 [1] E4M3 quantization scale + rope_dims: cutlass.Array, # rope: int32 [max_blocks, tokens-per-page log2, hq, hkv, T] + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], +): + k_tiles = num_k_tiles(k_in) + x_tile_bytes = mma_n * CTA_K + my_tiles = k_tiles // split + rows_per_rank = CTA_M // split + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + rank = cutlass.Int32(0) + if cutlass.const_expr(split > 1): + rank = cute.arch.block_idx_in_cluster() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_wsf = tma_desc_wsf.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + tma_ptr_xsf = tma_desc_xsf.get_ptr() + n_tile = bx // cutlass.Int32(split) + m_offset = n_tile * cutlass.Int32(CTA_M) + + smem_a = cutlass.Array( + cutlass.Uint8, ring * W_TILE_BYTES, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + cutlass.Uint8, my_tiles * x_tile_bytes, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_sfa = cutlass.Array( + cutlass.Uint8, ring * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + smem_sfb = cutlass.Array( + cutlass.Uint8, my_tiles * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + w_full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + mma_done = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + act_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + reduce_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + # Partials of the rows this rank owns: [source rank][row][token] fp32 (its own slot stays unused). Rows are padded + # to mma_n + 4 words: the owner's 16-byte row reads (lane = row) then hit all 32 banks once per 8 lanes; at a + # stride of mma_n they were 8- to 16-way conflicted (4.3 us of the 48-token QKV tail). + mb_stride = mma_n + 4 + mailbox = cutlass.Array( + cutlass.Float32, + split * rows_per_rank * mb_stride, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + # din: per-token RMS factor of the half of x this CTA's k-tiles lie in. + din_rs = None + if cutlass.const_expr(din): + din_rs = cutlass.Array(cutlass.Float32, 8, space=cutlass.AddressSpace.smem, alignment=16) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_wsf) + prims.prefetch_tensormap(tma_ptr_x) + prims.prefetch_tensormap(tma_ptr_xsf) + if prims.elect_sync(): + for s in cutlass.range_constexpr(ring): + prims.mbarrier_init(w_full.subview(s), 1) + prims.mbarrier_init(mma_done.subview(s), 1) + prims.mbarrier_init(act_full, 1) + prims.mbarrier_init(acc_done, 1) + if cutlass.const_expr(split > 1 and st_async): + # Completed by the peers' st.async bytes; this CTA's own arrival carries the expected count. + prims.mbarrier_init(reduce_ready, 1) + prims.mbarrier_arrive_expect_tx( + reduce_ready, (split - 1) * rows_per_rank * mma_n * 4 + ) + elif cutlass.const_expr(split > 1): + prims.mbarrier_init(reduce_ready, (split - 1) * rows_per_rank) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n)) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + if cutlass.const_expr(split > 1): + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + if tx == 0: + _stamp(timers, timed, bx, 0) + + if warp_id == 0: + if prims.elect_sync(): + for i in cutlass.range_constexpr(ring, my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [k_global * cutlass.Int32(CTA_K), m_offset], [] + ) + if cutlass.const_expr((n_out // CTA_M) * split >= 80): + # Wide grids launch in part only after the predecessor (e.g. the fused attention) frees its SMs, which + # loses their pre-wait weight prefetch: the CTAs of the first half also stage their partner's + # (bx + grid / 2) weights into L2. + if bx < cutlass.Int32((n_out // CTA_M) * split // 2): + pb = bx + cutlass.Int32((n_out // CTA_M) * split // 2) + p_rank = pb % cutlass.Int32(split) + p_m_offset = (pb // cutlass.Int32(split)) * cutlass.Int32(CTA_M) + for i in cutlass.range_constexpr(my_tiles): + p_k = _k_tile(p_rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [p_k * cutlass.Int32(CTA_K), p_m_offset], [] + ) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + if cutlass.const_expr(i >= ring): + while not cute.arch.mbarrier_try_wait( + mma_done.subview(stage).data_ptr(), (i // ring - 1) % 2 + ): + pass + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.mbarrier_arrive_expect_tx(w_full.subview(stage), W_TILE_BYTES + SF_ATOM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * W_TILE_BYTES), + tma_ptr_w, + (k_global * cutlass.Int32(CTA_K), m_offset), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfa.subview(stage * SF_ATOM_BYTES), + tma_ptr_wsf, + (cutlass.Int32(0), n_tile * cutlass.Int32(k_tiles) + k_global), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + elif warp_id == 1: + prims.griddepcontrol(prims.GridDepAction.WAIT) + if cutlass.const_expr(trigger_early): + # Dependents launch once this grid's own dependency is satisfied (not at entry): consumers that read + # step metadata or caches before their own grid wait rely on every grid ahead of them having passed its + # wait, so the chain of releases never runs ahead of a producer still in flight. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if tx == 32: + _stamp(timers, timed, bx, 1) + if cutlass.const_expr(not x_bf16): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(act_full, my_tiles * (x_tile_bytes + SF_ATOM_BYTES)) + for i in cutlass.range_constexpr(my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(i * x_tile_bytes), + tma_ptr_x, + (k_global * cutlass.Int32(CTA_K), cutlass.Int32(0)), + act_full, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfb.subview(i * SF_ATOM_BYTES), + tma_ptr_xsf, + (cutlass.Int32(0), k_global), + act_full, + ) + elif warp_id == 2: + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + ring * SF_COLS + idesc = prims.Tcgen05MxInstrDesc.build( + a_dtype=a_dtype, b_dtype=b_dtype, scale_format=1, n_dim=mma_n, m_dim=CTA_M + ) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + if cutlass.const_expr(a_tmem): + # A staged in TMEM: before the grid wait every weight k-tile and its scale atom go on into TMEM as they + # land (the copies of a stage refilled later are committed to its mma_done); once the activation lands, + # its scale atoms and the MMAs, A read from TMEM. + base_row = tmem_raw >> 16 + base_col = tmem_raw & 0xFFFF + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + my_tiles * SF_COLS + a_base = tmem_raw + cutlass.Int32(a_col0(nacc, my_tiles, mma_n)) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + if prims.elect_sync(): + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_cp( + prims.Tcgen05CpShape.SHAPE_128X256B, + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), + 6, + cutlass.Int32, + ), + desc_a + ((MMA_INST_K >> 4) * kb), + group=prims.CTAGroup.CTA_1, + ) + prims.tcgen05_cp( + s2t_shape, + cutlass.inttoptr((base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32), + desc_sfa, multicast=s2t_multicast, + ) # fmt: skip + if cutlass.const_expr(i + ring < my_tiles): + prims.tcgen05_commit(mma_done.subview(stage)) + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + for i in cutlass.range_constexpr(my_tiles): + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + sfa_tmem = cutlass.inttoptr( + (base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32 + ) + sfb_tmem = cutlass.inttoptr((base_row << 16) | sfb_col0, 6, cutlass.Int32) + prims.tcgen05_cp(s2t_shape, sfb_tmem, desc_sfb, multicast=s2t_multicast) + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), 6, cutlass.Int32 + ), + desc_b + ((MMA_INST_K >> 4) * kb), idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + i > 0 or kb > 0, + sfa_tmem, sfb_tmem, + ) # fmt: skip + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + else: + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + sfa_tmem = (base_row << 16) | (sfa_col0 + stage * SF_COLS) + sfb_tmem = (base_row << 16) | (sfb_col0 + stage * SF_COLS) + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), desc_sfa, + multicast=s2t_multicast) # fmt: skip + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), desc_sfb, + multicast=s2t_multicast) # fmt: skip + for kb in cutlass.range_constexpr(K_BLOCKS): + inc = ((MMA_INST_K * 8 // 8) >> 4) * kb + j = i * K_BLOCKS + kb + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw + cutlass.Int32((j % nacc) * mma_n), 6, cutlass.Float32), + desc_a + inc, desc_b + inc, idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), j >= nacc, + cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), + ) # fmt: skip + prims.tcgen05_commit(mma_done.subview(stage)) + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + elif warp_id >= 4: + lane = tx % 32 + row = (warp_id % 4) * 32 + lane + if cutlass.const_expr(x_bf16): + # The bf16 activation of this CTA's k-tiles to MXFP8 in shared memory, as + # torch.ops.trtllm.mxfp8_quantize(x, True) does (UE8M0 = round-up(qz_amax * rcp(448)), + # E4M3 RN-satfinite): one 32-element block (token qz_t, k-tile qz_i, block qz_kb) per thread and round, + # the E4M3 bytes into the 128-byte-swizzled B tile, the qz_scale byte into the tile's qz_scale atom. + qz_tid = tx - cutlass.Int32(128) + prims.griddepcontrol(prims.GridDepAction.WAIT) + din_half_k = k_in // 2 + din_hidden_half = cutlass.Int32(0) + if cutlass.const_expr(din): + # Draft input x = [rmsnorm(table[ids]) * w_e, rmsnorm(hidden) * w_h] (an out-of-range id embeds to + # zero): this CTA's k-tiles lie in one half, so it needs that half's statistic of every token. 16 + # threads per token, 16-byte loads, all of a thread's loads in flight at once. + din_hidden_half = cutlass.select_( + rank * cutlass.Int32(my_tiles * CTA_K) >= cutlass.Int32(din_half_k), + cutlass.Int32(1), + cutlass.Int32(0), + ) + din_tok = qz_tid // cutlass.Int32(16) + din_part = qz_tid % cutlass.Int32(16) + din_tok_ld = cutlass.select_(din_tok < num_tokens, din_tok, cutlass.Int32(0)) + din_id = din_i.load(idx=din_tok_ld) + din_valid = (din_id >= cutlass.Int32(0)) & (din_id < din_vocab) + din_hrow = din_tok_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + din_hrow = din_g.load(idx=din_tok_ld * cutlass.Int32(2)) + din_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + din_hrow, + cutlass.select_(din_valid, din_id, cutlass.Int32(0)), + ) + # The first hidden-half CTA of tile 0 reads whole hidden rows here: it also writes them out. + din_copy = ( + ((din_flags & cutlass.Int32(2)) != cutlass.Int32(0)) + & (bx == cutlass.Int32(split // 2)) + & (din_hidden_half == cutlass.Int32(1)) + & (din_tok < num_tokens) + ) + din_base = din_row * cutlass.Int32(din_half_k // 2) + din_ss = cutlass.Float32(0.0) + din_vecs_per_part = din_half_k // 8 // 16 + for din_round in cutlass.range_constexpr(0, din_vecs_per_part, 16): + din_n = min(16, din_vecs_per_part - din_round) + din_loaded = [] + for din_j in cutlass.range_constexpr(din_n): + din_widx = din_base + ( + din_part + cutlass.Int32((din_round + din_j) * 16) + ) * cutlass.Int32(4) + din_loaded.append( + _din_load4(din_table, din_hidden, din_hidden_half, din_widx) + ) + if din_copy: + for din_j in cutlass.range_constexpr(din_n): + din_v = din_loaded[din_j] + din_res.store( + (din_v[0], din_v[1], din_v[2], din_v[3]), + idx=din_tok * cutlass.Int32(din_half_k // 2) + + (din_part + cutlass.Int32((din_round + din_j) * 16)) + * cutlass.Int32(4), + alignment=16, + ) + for din_j in cutlass.range_constexpr(din_n): + for din_e in cutlass.range_constexpr(4): + din_word = cutlass.Int32(din_loaded[din_j][din_e]) + din_lo = (din_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + din_hi = (din_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + din_ss = din_ss + din_lo * din_lo + din_hi * din_hi + for din_o in cutlass.range_constexpr(4): + din_ss = din_ss + cute.arch.shuffle_sync_bfly(din_ss, offset=1 << din_o) + if din_part == cutlass.Int32(0): + din_eps = din_f.load(idx=din_hidden_half) + din_r = cute.math.rsqrt( + din_ss / cutlass.Float32(din_half_k) + din_eps, fastmath=True + ) + din_keep = din_valid | (din_hidden_half == cutlass.Int32(1)) + din_rs.store( + cutlass.select_(din_keep, din_r, cutlass.Float32(0.0)), idx=din_tok + ) + prims.barrier_cta_sync(3, thread_count=128) + qz_b_words = cutlass.Array(smem_b.data_ptr(0), shape=(my_tiles * x_tile_bytes // 4,), dtype=cutlass.Int32, + alignment=16) # fmt: skip + qz_items = my_tiles * K_BLOCKS * mma_n + qz_rounds = (qz_items + 127) // 128 + # A group's loads first (clamped to a valid row), then its math: one L2 round trip per group of + # QZ_GROUP rounds instead of one per round, without holding every round's data in registers. + QZ_GROUP = qz_rounds if qz_rounds <= 8 else 4 + qz_loaded = [] + for qz_r in cutlass.range_constexpr(qz_rounds): + qz_item = qz_tid + cutlass.Int32(qz_r * 128) + qz_t = qz_item % cutlass.Int32(mma_n) + qz_kb = (qz_item // cutlass.Int32(mma_n)) % cutlass.Int32(K_BLOCKS) + qz_i = qz_item // cutlass.Int32(mma_n * K_BLOCKS) + if cutlass.const_expr(qz_r * 128 + 127 >= qz_items): + qz_i = cutlass.select_( + qz_item < cutlass.Int32(qz_items), qz_i, cutlass.Int32(0) + ) + qz_k_global = rank + qz_i * cutlass.Int32(split) + if cutlass.const_expr(contig): + qz_k_global = rank * cutlass.Int32(my_tiles) + qz_i + qz_t_ld = cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + qz_x_base = ( + qz_t_ld * cutlass.Int32(k_in // 2) + + qz_k_global * cutlass.Int32(CTA_K // 2) + + qz_kb * cutlass.Int32(MMA_INST_K // 2) + ) + qz_wvecs = None + if cutlass.const_expr(din): + # The draft input's raw values (table row of the id / hidden row) and its norm weights. + qz_col = qz_k_global * cutlass.Int32(CTA_K) + qz_kb * cutlass.Int32(MMA_INST_K) + qz_hcol = qz_col - din_hidden_half * cutlass.Int32(din_half_k) + qz_id = din_i.load(idx=qz_t_ld) + qz_hrow = qz_t_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + qz_hrow = din_g.load(idx=qz_t_ld * cutlass.Int32(2)) + qz_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + qz_hrow, + cutlass.select_( + (qz_id >= cutlass.Int32(0)) & (qz_id < din_vocab), + qz_id, + cutlass.Int32(0), + ), + ) + qz_src_w = qz_row * cutlass.Int32(din_half_k // 2) + qz_hcol // cutlass.Int32(2) + qz_vecs = [ + _din_load4( + din_table, + din_hidden, + din_hidden_half, + qz_src_w + cutlass.Int32(qz_v * 4), + ) + for qz_v in range(4) + ] + qz_wvecs = [ + din_w.load( + idx=qz_col // cutlass.Int32(2) + cutlass.Int32(qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + else: + qz_vecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(qz_v * 4), vector_size=4, alignment=16 + ) + for qz_v in range(4) + ] + qz_loaded.append((qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs)) + if cutlass.const_expr((qz_r + 1) % QZ_GROUP == 0 or qz_r + 1 == qz_rounds): + for qz_g in cutlass.range_constexpr(len(qz_loaded)): + qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs = qz_loaded[qz_g] + qz_vals = [] + qz_amax = cutlass.Float32(0.0) + if cutlass.const_expr(din): + qz_rs = din_rs.load( + idx=cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + ) + for qz_v in cutlass.range_constexpr(4): + for qz_e4 in cutlass.range_constexpr(4): + qz_word = cutlass.Int32(qz_vecs[qz_v][qz_e4]) + qz_lo = (qz_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + qz_hi = (qz_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + if cutlass.const_expr(din): + # bf16(x * rs * w), the draft-input kernel's rounding. + qz_wword = cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + qz_lo = _bf16_round( + qz_lo + * qz_rs + * (qz_wword << cutlass.Int32(16)).bitcast(cutlass.Float32) + ) + qz_hi = _bf16_round( + qz_hi + * qz_rs + * (qz_wword & cutlass.Int32(-65536)).bitcast( + cutlass.Float32 + ) + ) + qz_vals.append(qz_lo) + qz_vals.append(qz_hi) + qz_amax = _absmax(_absmax(qz_amax, qz_lo), qz_hi) + qz_bits = (qz_amax * _rcp_approx_ftz(cutlass.Float32(448.0))).bitcast( + cutlass.Int32 + ) + qz_e = ( + (qz_bits >> cutlass.Int32(23)) & cutlass.Int32(0xFF) + ) + cutlass.select_( + (qz_bits & cutlass.Int32(0x7FFFFF)) != cutlass.Int32(0), + cutlass.Int32(1), + cutlass.Int32(0), + ) + qz_e = cutlass.select_(qz_e > cutlass.Int32(254), cutlass.Int32(254), qz_e) + qz_scale = ((cutlass.Int32(254) - qz_e) << cutlass.Int32(23)).bitcast( + cutlass.Float32 + ) + qz_scale = cutlass.select_( + qz_amax != cutlass.Float32(0.0), qz_scale, cutlass.Float32(0.0) + ) + qz_q = [_e4m3x4(qz_vals[4 * qz_j] * qz_scale, qz_vals[4 * qz_j + 1] * qz_scale, + qz_vals[4 * qz_j + 2] * qz_scale, qz_vals[4 * qz_j + 3] * qz_scale) + for qz_j in range(8)] # fmt: skip + if qz_item < cutlass.Int32(qz_items): + if qz_t < num_tokens: + qz_row_w = qz_i * cutlass.Int32( + x_tile_bytes // 4 + ) + qz_t * cutlass.Int32(CTA_K // 4) + qz_c0 = qz_kb * cutlass.Int32(2) + qz_b_words.store( + (qz_q[0], qz_q[1], qz_q[2], qz_q[3]), + idx=qz_row_w + + (qz_c0 ^ (qz_t & cutlass.Int32(7))) * cutlass.Int32(4), + alignment=16, + ) + qz_b_words.store( + (qz_q[4], qz_q[5], qz_q[6], qz_q[7]), + idx=qz_row_w + + ((qz_c0 + cutlass.Int32(1)) ^ (qz_t & cutlass.Int32(7))) + * cutlass.Int32(4), + alignment=16, + ) + smem_sfb.store( + cutlass.Uint8(qz_e), + idx=qz_i * cutlass.Int32(SF_ATOM_BYTES) + + (qz_t % cutlass.Int32(32)) * cutlass.Int32(16) + + (qz_t // cutlass.Int32(32)) * cutlass.Int32(4) + + qz_kb, + ) + qz_loaded = [] + # Generic-proxy stores made visible to the MMA's async-proxy reads, then one arrival for all. + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.barrier_cta_sync(2, thread_count=128) + if qz_tid == cutlass.Int32(0): + prims.mbarrier_arrive(act_full) + if cutlass.const_expr(st_async): + while not cute.arch.mbarrier_test_wait(acc_done.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + # tcgen05.ld widths are powers of two: the accumulator columns in descending power-of-two pieces. + acc_all = [] + ld_col = 0 + for ld_w in (128, 64, 32, 16, 8): + if cutlass.const_expr((mma_n * nacc - ld_col) >= ld_w): + piece = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr( + tmem_ptr_i32.load() + cutlass.Int32(ld_col), 6, cutlass.Float32 + ), + num=ld_w, + ) + acc_all = acc_all + [piece[c] for c in range(ld_w)] + ld_col += ld_w + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + acc = [] + for t in cutlass.range_constexpr(mma_n): + v = cutlass.Float32(acc_all[t]) + for a in cutlass.range_constexpr(1, nacc): + v = v + cutlass.Float32(acc_all[a * mma_n + t]) + acc.append(v) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if tx == 128: + _stamp(timers, timed, bx, 4) + n = m_offset + row + if cutlass.const_expr(split == 1): + for t in cutlass.range_constexpr(mma_n): + if cutlass.Int32(t) < num_tokens: + y.store( + cutlass.BFloat16(cutlass.Float32(acc[t])), idx=cutlass.Int32(t * n_out) + n + ) + if cutlass.const_expr(ar_push > 0): + prims.griddepcontrol(prims.GridDepAction.WAIT) + # The parity of each (token, column slice) slot: the all-reduce's CTA t * 8 + slice flag. All loads are + # issued before the first store so they overlap (one L2 round trip, not one per token). + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_v = cutlass.Float32(acc[t]) + ar_hi = cute.arch.shuffle_sync_bfly(ar_v, offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), _bf16x2_word(ar_v, ar_hi) + ) + else: + owner = row // cutlass.Int32(rows_per_rank) + local = row % cutlass.Int32(rows_per_rank) + if cutlass.const_expr(rope): + # Rank r owns head dims [r * P, (r + 1) * P) and their rotation partners 64 + ..., P = rows / 2. + pairs = rows_per_rank // 2 + owner = (row % cutlass.Int32(64)) // cutlass.Int32(pairs) + local = (row // cutlass.Int32(64)) * cutlass.Int32(pairs) + ( + row % cutlass.Int32(64) + ) % cutlass.Int32(pairs) + if owner != rank: + base = (rank * cutlass.Int32(rows_per_rank) + local) * cutlass.Int32(mb_stride) + if cutlass.const_expr(st_async): + dst = _mapa_u32(mailbox.subview(base).data_ptr(), owner) + mb = _mapa_u32(reduce_ready.data_ptr(), owner) + words = [cutlass.Float32(acc[t]).bitcast(cutlass.Int32) for t in range(mma_n)] + for g in cutlass.range_constexpr(mma_n // 4): + _st_async_v4( + dst + cutlass.Int32(16 * g), + words[4 * g], + words[4 * g + 1], + words[4 * g + 2], + words[4 * g + 3], + mb, + ) + else: + for t in cutlass.range_constexpr(mma_n): + prims.mapa(mailbox.subview(base + cutlass.Int32(t)), owner).store( + cutlass.Float32(acc[t]) + ) + prims.mbarrier_arrive( + prims.mapa(reduce_ready, owner), scope=prims.MemScope.CLUSTER + ) + else: + if cutlass.const_expr(st_async): + while not _test_wait_cluster(reduce_ready.data_ptr(), 0): + pass + else: + while not prims.mbarrier_try_wait_parity( + reduce_ready, 0, scope=prims.MBarrierScope.CLUSTER + ): + pass + if lane == 0: + if row // cutlass.Int32(rows_per_rank) == rank: + if (row % cutlass.Int32(rows_per_rank)) == 0: + _stamp(timers, timed, bx, 5) + totals = [] + for g in cutlass.range_constexpr(mma_n // 4): + sums = [cutlass.Float32(0.0)] * 4 + for q in cutlass.range_constexpr(split): + parts = _ld_shared_v4_f32( + mailbox.subview( + (cutlass.Int32(q * rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g) + ).data_ptr() + ) + sums = [sums[j] + cutlass.select_(rank == cutlass.Int32(q), cutlass.Float32(acc[4 * g + j]), + parts[j]) for j in range(4)] # fmt: skip + for j in cutlass.range_constexpr(4): + totals.append(sums[j]) + if cutlass.const_expr(rope): + # The own-rank mailbox slot is unused: it carries the totals to the RoPE stage. + mailbox.store( + sums[j], + idx=(rank * cutlass.Int32(rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g + j), + ) + elif cutlass.Int32(4 * g + j) < num_tokens: + y.store( + cutlass.BFloat16(sums[j]), + idx=cutlass.Int32((4 * g + j) * n_out) + n, + ) + if cutlass.const_expr(ar_push > 0): + # Every lane of an owner warp is an owner (rows_per_rank >= 32), so lane pairs hold row pairs. + prims.griddepcontrol(prims.GridDepAction.WAIT) + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_hi = cute.arch.shuffle_sync_bfly(totals[t], offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), + _bf16x2_word(totals[t], ar_hi), + ) + prims.barrier_cta_sync(1, thread_count=128) + if cutlass.const_expr(rope): + # Token t of request t / T sits at cache position kv_len - T + t % T. Q / K heads: GPT-NeoX RoPE over + # the 128-wide head, bf16(cos * x -/+ sin * pair) on the bf16 projection; then E4M3 of x * scale, Q to + # rope_q, K / V into the page of that position. + pairs = rows_per_rank // 2 + rp_e = tx - cutlass.Int32(128) + rp_p = rp_e % cutlass.Int32(pairs) + rp_t = rp_e // cutlass.Int32(pairs) + if rp_t < num_tokens: + rp_max_blocks = rope_dims.load(idx=0) + rp_tpb_log2 = rope_dims.load(idx=1) + rp_hq = rope_dims.load(idx=2) + rp_hkv = rope_dims.load(idx=3) + rp_T = rope_dims.load(idx=4) + rp_d = rank * cutlass.Int32(pairs) + rp_p + rp_a = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + rp_p) * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_b = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + cutlass.Int32(pairs) + rp_p) + * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_req = rp_t // rp_T + rp_tok = rope_i.load(idx=rp_req) - rp_T + (rp_t - rp_req * rp_T) + rp_head = n_tile + if rp_head < rp_hq + rp_hkv: + rp_c = rope_cs.load(idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2)) + rp_s = rope_cs.load( + idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2) + + cutlass.Int32(1) + ) + rp_ra = _rope_bits(rp_c, rp_a, -rp_s, rp_b).bitcast(cutlass.Float32) + rp_rb = _rope_bits(rp_c, rp_b, rp_s, rp_a).bitcast(cutlass.Float32) + rp_a = rp_ra + rp_b = rp_rb + rp_qs = rope_scale.load(idx=0) + rp_w = _e4m3x4( + rp_a * rp_qs, rp_b * rp_qs, cutlass.Float32(0.0), cutlass.Float32(0.0) + ) + rp_lo = cutlass.Uint8(rp_w & cutlass.Int32(0xFF)) + rp_hi = cutlass.Uint8((rp_w >> cutlass.Int32(8)) & cutlass.Int32(0xFF)) + if rp_head < rp_hq: + rp_q = (rp_t * rp_hq + rp_head) * cutlass.Int32(128) + rp_d + rope_q.store(rp_lo, idx=rp_q) + rope_q.store(rp_hi, idx=rp_q + cutlass.Int32(64)) + else: + rp_kv = cutlass.select_( + rp_head < rp_hq + rp_hkv, cutlass.Int32(0), cutlass.Int32(1) + ) + rp_kh = rp_head - rp_hq - rp_kv * rp_hkv + rp_page = rope_bt.load( + idx=(rp_req * cutlass.Int32(2) + rp_kv) * rp_max_blocks + + (rp_tok >> rp_tpb_log2) + ) + # 64-bit byte offset: the pool view can exceed 2^31 bytes. + rp_dst = ( + (cutlass.Int64(rp_page) * cutlass.Int64(rp_hkv) + cutlass.Int64(rp_kh)) + << cutlass.Int64(rp_tpb_log2 + cutlass.Int32(7)) + ) + cutlass.Int64( + (rp_tok & ((cutlass.Int32(1) << rp_tpb_log2) - cutlass.Int32(1))) + * cutlass.Int32(128) + + rp_d + ) + rp_addr = rope_pool.data_ptr(0).toint() + rp_dst + cutlass.inttoptr(rp_addr, mem_space=1, dtype=cutlass.Uint8).store(rp_lo) + cutlass.inttoptr( + rp_addr + cutlass.Int64(64), mem_space=1, dtype=cutlass.Uint8 + ).store(rp_hi) + if tx == 128: + _stamp(timers, timed, bx, 6) + if cutlass.const_expr(ar_push == 2): + # The all-reduce polls the pushed words; it launches after this CTA's rows are out, not at grid end. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if warp_id == 4: + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32), + tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n), + ) + + +def _bytes_map(t, cols, rows, box_rows, swizzle): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[cols, rows], + global_strides=[cols // 16], + box_dims=[CTA_K, box_rows], + swizzle=swizzle, + ) + + +def _atom_map(t, atoms): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint16, + global_dims=[SF_ATOM_U16, atoms], + global_strides=[SF_ATOM_BYTES // 16], + box_dims=[SF_ATOM_U16, 1], + swizzle=cuda.TensorMapSwizzle.none, + ) + + +@cute.jit +def m3_mx_gemv( + w: cute.Tensor, # [N, K] E4M3 as uint8, K contiguous + wsf: cute.Tensor, # weight scales, 128x4-swizzled, as uint8 + x: cute.Tensor, # [M, K] E4M3 as uint8, M <= 8 + xsf: cute.Tensor, # activation scales, 128x4-swizzled (atom k at byte 512 k), as uint8 + y: cute.Tensor, # [M * N] bf16 + x_words: cute.Tensor, # x_bf16: bf16 x [M, K] as int32 words (x / xsf are then unused placeholders) + timers: cute.Tensor, # [grid * 8] int64, written when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[int], + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, + din_i: cute.Tensor, # din / rope inputs: see m3_mx_gemv_kernel (placeholders when off) + din_table: cute.Tensor, + din_hidden: cute.Tensor, + din_w: cute.Tensor, + din_f: cute.Tensor, + din_vocab: cutlass.Int32, + din_g: cute.Tensor, + din_res: cute.Tensor, + din_flags: cutlass.Int32, + rope_cs: cute.Tensor, + rope_i: cute.Tensor, + rope_bt: cute.Tensor, + rope_q: cute.Tensor, + rope_pool: cute.Tensor, + rope_scale: cute.Tensor, + rope_dims: cute.Tensor, + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], +) -> None: + k_tiles = num_k_tiles(k_in) + tma_desc_w = _bytes_map(w, k_in, n_out, CTA_M, cuda.TensorMapSwizzle.s128b) + tma_desc_wsf = _atom_map(wsf, (n_out // CTA_M) * k_tiles) + tma_desc_x = _bytes_map(x, k_in, num_tokens, mma_n, cuda.TensorMapSwizzle.s128b) + tma_desc_xsf = _atom_map(xsf, k_tiles) + m3_mx_gemv_kernel( + tma_desc_w, tma_desc_wsf, tma_desc_x, tma_desc_xsf, y, x_words, timers, num_tokens, n_out, k_in, split, ring, + trigger_early, contig, timed, nacc, a_tmem, st_async, x_bf16, ar_ws, ar_flags, ar_rank, ar_push, ar_world, + ar_maxt, mma_n, + din_i, din_table, din_hidden, din_w, din_f, din_vocab, din_g, din_res, din_flags, rope_cs, + rope_i, rope_bt, rope_q, rope_pool, rope_scale, rope_dims, din, rope, + ).launch( + grid=((n_out // CTA_M) * split, 1, 1), + block=(THREADS, 1, 1), + cluster=(split, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mx_gemv/op.py new file mode 100644 index 000000000000..d5440088cbc8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mx_gemv/op.py @@ -0,0 +1,474 @@ +# 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. +"""MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) for already-quantized activations of at most 8 tokens. + +Compiled on the first call per (N, K, split, ring, early trigger, PDL, push, MMA width), which must happen outside +CUDA-graph capture, or ahead of any call by ``precompile``; once ``freeze`` closes compiles, a call whose kernel is not +compiled declines instead. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +MMA_WIDTHS = ( + 8, + 16, + 32, + 48, +) # token columns of the compiled MMA; a call uses the smallest that holds M +_SMEM_LIMIT = 220 * 1024 # capacity: fixed: the device's shared-memory limit + + +def _mma_n(m: int) -> int: + return next(w for w in MMA_WIDTHS if m <= w) + + +def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: + return ( + ring * (128 * 128 + 512) + + tiles * (mma_n * 128 + 512) + + split * (128 // split) * (mma_n + 4) * 4 + + 1024 + ) + + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + +# (N, K) -> (split, ring): GB300 exposed time behind a 4.25 us PDL predecessor at M 1 / 8, vs the fused low-M GEMM: +# QKV 2.78 / 3.16 vs 4.58 / 6.21 us, o_proj 2.23 / 2.41 vs 4.40 / 5.09, shared gate_up 2.43 / 2.74 vs 3.56 / 4.30, +# shared down 1.74 / 1.66 vs 2.92 / 3.38. A config is (split, ring[, contig, nacc, a_tmem, st_async]); the optional +# fields default to _CONTIG, _NACC, _A_TMEM, _ST_ASYNC. +_CONFIGS: Dict[Tuple[int, int], tuple] = { + (2560, 6144): (4, 12), + (6144, 2048): (2, 8), + # Shared-expert gate_up: split 2, ring 12, single accumulator, A operand from shared memory. + (1536, 6144): (2, 12, 1, 1, 0, 1), + # Shared-expert down: unsplit, ring 3, single accumulator, A operand from shared memory. + (6144, 768): (1, 3, 1, 1, 0, 1), + (2304, 6144): (4, 12), + (6144, 3072): (2, 12), + # Dense-layer gate_up: 6.4 / 6.3 us at M 1 / 8 vs FlashInfer 8.8. + (6144, 6144): (2, 12), +} +# Per-token-count configs (N, K, M): the grid that suits the concurrent routed-expert engine depends on how long +# that engine runs. At 2 tokens the shared-expert gate_up splits K over 4 CTAs. +_CONFIGS_BY_M: Dict[Tuple[int, int, int], tuple] = { + (1536, 6144, 2): (4, 12, 1, 1, 0, 1), +} +# Above 8 tokens (spec-decode verify rows): the shared-expert gate_up splits K over 4 CTAs. +_CONFIGS_WIDE: Dict[Tuple[int, int], tuple] = { + (1536, 6144): (4, 12, 1, 1, 0, 1), +} + +# Defaults of the optional config fields: contiguous k-tile ranges per split rank, one TMEM accumulator, the A +# operand staged in TMEM (when its columns fit), st.async split-K partial exchange. +_CONTIG = True +_NACC = 1 +_A_TMEM = True +_ST_ASYNC = True +_ensured = set() + + +_DUMMY: Dict[torch.device, torch.Tensor] = {} + + +def _dummy_timers(device) -> torch.Tensor: + """Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on.""" + t = _DUMMY.get(device) + if t is None: + t = _DUMMY[device] = torch.zeros(8, dtype=torch.int64, device=device) + return t + + +def _arg(t: torch.Tensor, align: int = 16): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def get_config(n: int, k: int, m: Optional[int] = None) -> Optional[tuple]: + cfg = _CONFIGS_BY_M.get((n, k, m)) if m is not None else None + if cfg is None and m is not None and m > 8: + cfg = _CONFIGS_WIDE.get((n, k)) + cfg = cfg or _CONFIGS.get((n, k)) + # The split-K epilogue gives each of the split CTAs 128 // split rows of a 128-row output tile. + if cfg is None or n % 128 != 0 or k % 128 != 0 or (k // 128) % cfg[0] != 0 or 128 % cfg[0] != 0: + return None + return cfg + + +def _declines(n: int, k: int, m: int, x_bf16: bool, cfg: tuple) -> bool: + """Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds.""" + split = cfg[0] + if not 0 < m <= MAX_TOKENS or 128 % split != 0 or (k // 128) % split != 0: + return True + tiles = k // 128 // split + mma_n = _mma_n(m) + # GB300 exposed time vs the FlashInfer prequantized GEMM (behind a PDL predecessor), at 24 / 48 tokens: QKV 3.76 / + # 4.63 vs 7.46 / 11.44 us, o_proj 2.74 / 3.33 vs 4.46 / 5.84, shared gate_up 4.57 / 5.71 vs 7.16 / 11.25. Behind + # only for K = 768 (unsplit shared down) at 17-32 tokens: 3.1 vs 2.9. + # The in-kernel bf16 -> MXFP8 quantization (x_bf16) costs one 32-element block per thread and round on 128 threads; + # beyond 8 rounds per CTA, mxfp8_quantize + this GEMV (or the fallback GEMM) is faster. + if (16 < m <= 32 and k < 1024) or (x_bf16 and tiles * 4 * mma_n > 8 * 128): + return True + return _smem_bytes(min(cfg[1], tiles, 2), tiles, split, mma_n) > _SMEM_LIMIT + + +def serves(n: int, k: int, m: int, x_bf16: bool) -> bool: + """Whether the GEMV with its default config takes an [M, K] activation (bf16 or MXFP8) against an [N, K] weight.""" + cfg = get_config(n, k, m) + return cfg is not None and not _declines(n, k, m, x_bf16, cfg) + + +def _resolve( + m: int, + n: int, + k: int, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple], +) -> Optional[tuple]: + """The compile key of an [M, K] x [N, K]^T call and its runtime push words, or None when the call is declined. + + ``ar`` is the push tuple of ``_run``; only its (world, bucket tokens, mode) enter the key, the addresses and the + rank are runtime arguments.""" + cfg = cfg or get_config(n, k, m) + if cfg is None or _declines(n, k, m, x_bf16, cfg): + return None + split, ring = cfg[0], cfg[1] + contig = bool(cfg[2]) if len(cfg) > 2 else _CONTIG + nacc = int(cfg[3]) if len(cfg) > 3 else _NACC + a_tmem = bool(cfg[4]) if len(cfg) > 4 else _A_TMEM + st_async = bool(cfg[5]) if len(cfg) > 5 else _ST_ASYNC + ring = min(ring, k // 128 // split) + tiles = k // 128 // split + mma_n = _mma_n(m) + while ring > 2 and _smem_bytes(ring, tiles, split, mma_n) > _SMEM_LIMIT: + ring -= 1 + if a_tmem: + from .m3_mx_gemv_kernel import tmem_cols + + if tmem_cols(ring, nacc, True, tiles, mma_n) > 512: + a_tmem = False + # ar = (multicast workspace address, flags address, rank, world, bucket tokens, mode): push the output rows into + # the next m3_mc_allreduce call's exchange slots (mode 2 also defers the dependents' launch until after the push). + ar_ws, ar_flags, ar_rank, ar_mode, ar_world, ar_maxt = 0, 0, 0, 0, 1, 8 + if ar is not None and split <= 4 and n % 2 == 0: + ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_mode = ar + if ar_mode == 2: + trigger_early = False + if use_pdl is None: + use_pdl = get_env_enable_pdl() + # The kernel's ``timed`` flag (per-CTA milestone timers) is always off. + key = ( + n, + k, + split, + ring, + trigger_early, + contig, + False, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + ) + return key, (ar_ws, ar_flags, ar_rank) + + +def _compile(key: tuple, args: tuple, m: int, runtime_ar: tuple, stream, extra: tuple): + """cute.compile the kernel of ``key`` (see ``_resolve``) on these argument views; cached.""" + import cutlass.cute as cute + + from . import m3_mx_gemv_kernel as kernel + + ( + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + ) = key + ar_ws, ar_flags, ar_rank = runtime_ar + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.m3_mx_gemv, + *args, + m, + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_ws, + ar_flags, + ar_rank, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + stream, + *extra, + False, + False, + ) + return fn + + +def _views( + x: torch.Tensor, xsf: Optional[torch.Tensor], weight, weight_scale, y, x_bf16: bool +) -> tuple: + k = x.shape[1] + if x_bf16: + x_bytes = x.view(torch.uint8) + xsf = _placeholder_scales(k, x.device) + x_words = x.view(torch.int32) + else: + x_bytes = x.view(torch.uint8) + x_words = _placeholder_words(x.device) + return ( + _arg(weight.view(torch.uint8)), + _arg(weight_scale.view(torch.uint8).reshape(-1)), + _arg(x_bytes), + _arg(xsf.view(torch.uint8).reshape(-1)), + _arg(y.view(-1) if y.numel() else _placeholder_words(x.device).view(torch.bfloat16)), + _arg(x_words.reshape(-1)), + _arg(_dummy_timers(x.device)), + ) + + +def _run( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight: torch.Tensor, + weight_scale: torch.Tensor, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + global last_pushed, last_uncompiled + last_pushed = False + last_uncompiled = None + import cuda.bindings.driver as cuda_driver + + m, k = x.shape + n = weight.shape[0] + if weight.shape[1] != k: + return None + resolved = _resolve(m, n, k, x_bf16, trigger_early, cfg, use_pdl, ar) + if resolved is None: + return None + key, runtime_ar = resolved + fn = _compiled.get(key) + if fn is None: + if _frozen[0]: + # Compiles are closed (the model built every kernel it dispatches to before warmup): decline, the caller + # reports the miss. + last_uncompiled = key + return None + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_mx_gemv must run once per shape outside CUDA-graph capture first" + ) + y = torch.empty(m, n, dtype=torch.bfloat16, device=x.device) + args = _views(x, xsf, weight, weight_scale, y, x_bf16) + extra = _extra_args(x.device) + stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + if fn is None: + fn = _compile(key, args, m, runtime_ar, stream, extra) + ar_ws, ar_flags, ar_rank = runtime_ar + fn(*args, m, ar_ws, ar_flags, ar_rank, stream, *extra) + last_pushed = key[11] > 0 + return y + + +def precompile( + weight: torch.Tensor, + weight_scale: torch.Tensor, + m: int, + x_bf16: bool, + use_pdl: Optional[bool] = None, + ar: Optional[Tuple[int, int, int]] = None, +) -> Optional[tuple]: + """Compile, without launching, the kernel an M-row call against ``weight`` takes; returns its key, or None when + such a call is declined. ``ar`` = (world, bucket tokens, mode) of a pushing call (its addresses and rank are + runtime arguments, so none is needed here). Runs outside capture, even while compiles are frozen.""" + import cuda.bindings.driver as cuda_driver + + n, k = weight.shape + push = None if ar is None else (0, 0, 0, *ar) + resolved = _resolve(m, n, k, x_bf16, True, None, use_pdl, push) + if resolved is None: + return None + key, runtime_ar = resolved + if key in _compiled: + return key + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_mx_gemv.precompile allocates and compiles: call it outside capture") + dev = weight.device + x = torch.zeros(m, k, dtype=torch.bfloat16 if x_bf16 else torch.float8_e4m3fn, device=dev) + xsf = None if x_bf16 else torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) + y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) + args = _views(x, xsf, weight, weight_scale, y, x_bf16) + stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) + _compile(key, args, m, runtime_ar, stream, _extra_args(dev)) + return key + + +def freeze(frozen: bool = True) -> None: + """Close (or reopen) compiles: while frozen, a call whose kernel is not compiled yet declines (returns None and + sets ``last_uncompiled``) instead of compiling. ``precompile`` and ``ensure_compiled`` still compile.""" + _frozen[0] = bool(frozen) + + +def frozen() -> bool: + return _frozen[0] + + +def is_compiled(key: tuple) -> bool: + return key in _compiled + + +def _extra_args(device) -> tuple: + """Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off).""" + i32 = _placeholder_words(device) + f32 = i32.view(torch.float32) + u8 = i32.view(torch.uint8) + d = (_arg(i32), _arg(i32), _arg(i32), _arg(i32), _arg(f32), 0, _arg(i32), _arg(i32), 0) + r = (_arg(f32), _arg(i32), _arg(i32), _arg(u8), _arg(u8), _arg(f32), _arg(i32)) + return d + r + + +def m3_mx_gemv( + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``xq`` E4M3 [M <= 8, K] with ``xsf`` its 128x4-swizzled UE8M0 scales; ``weight`` E4M3 [N, K] with + ``weight_scale`` swizzled. Returns bf16 [M, N], or None when the shape is not served.""" + return _run(xq, xsf, weight, weight_scale, False, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_bf16( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """The same GEMV for a bf16 ``x`` [M <= 8, K] (contiguous), quantized to MXFP8 inside the kernel exactly as + torch.ops.trtllm.mxfp8_quantize(x, True) does.""" + if x.dtype != torch.bfloat16 or x.dim() != 2 or not x.is_contiguous() or x.data_ptr() % 16 != 0: + return None + return _run(x, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar) + + +_PLACEHOLDERS: Dict[tuple, torch.Tensor] = {} +# Whether the last call pushed its output into an all-reduce (see _run's ``ar``). +last_pushed = False +# The key of the last call declined because compiles are frozen and its kernel was not compiled (see ``freeze``). +last_uncompiled: Optional[tuple] = None +_frozen = [False] + + +def _placeholder_scales(k: int, device) -> torch.Tensor: + key = ("sf", k, device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(k // 128 * 512, dtype=torch.uint8, device=device) + return t + + +def _placeholder_words(device) -> torch.Tensor: + key = ("w", device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(16, dtype=torch.int32, device=device) + return t + + +def ensure_compiled(weight: torch.Tensor, weight_scale: torch.Tensor, use_pdl: bool) -> None: + """Compiles the kernel of this weight's shape outside capture (a one-token dummy call).""" + if ( + weight.dim() != 2 + or weight.dtype != torch.float8_e4m3fn + or get_config(*weight.shape) is None + ): + return + n, k = weight.shape + if (n, k, use_pdl) in _ensured: + return + _ensured.add((n, k, use_pdl)) + xq = torch.zeros(1, k, dtype=torch.float8_e4m3fn, device=weight.device) + xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=weight.device) + m3_mx_gemv(xq, xsf, weight, weight_scale, use_pdl=use_pdl) + for width in MMA_WIDTHS[1:]: + m3_mx_gemv( + torch.zeros(width, k, dtype=torch.float8_e4m3fn, device=weight.device), + xsf, + weight, + weight_scale, + use_pdl=use_pdl, + ) + m3_mx_gemv_bf16( + torch.zeros(1, k, dtype=torch.bfloat16, device=weight.device), + weight, + weight_scale, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mxfp8_low_m_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mxfp8_low_m_gemm.py new file mode 100644 index 000000000000..cc635f93a836 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mxfp8_low_m_gemm.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused MXFP8 activation quantization + low-M MXFP8 GEMM for decode (SM100+). + +``out[M, N] = mxfp8(x)[M, K] @ mxfp8(W)[N, K]^T`` for ``1 <= M <= 8`` in one kernel +(``torch.ops.trtllm.minimax_m3_mxfp8_low_m_gemm``, reached through the ``ops/gemm`` entry). + +* Each CTA owns ``16 * ROW_TILES`` weight rows and ``K / CLUSTER_K`` columns. + Before ``griddepcontrol.wait`` it issues ``cp.async.bulk`` copies of its weight + tile and UE8M0 scale chunks (existing swizzled 128x4 layout) into shared memory, + so the weight stream overlaps the producer kernel under PDL. Configurations + keep the whole grid resident in one wave so every CTA prefetches. +* Activation quantization reproduces ``trtllm::mxfp8_quantize`` with identical + intrinsics (block 32, UE8M0 rounded toward +inf, saturating RN E4M3); each + thread issues all of its loads before converting, one L2 round trip per CTA. +* One ``mma.sync.m16n8k32`` E4M3 per 32-wide scale block and row tile with exact + power-of-two rescaling and FP32 accumulation. Warps reduce through shared + memory (reusing the consumed weight tile); split-K peers push their partial + tiles into the cluster root with ``st.async`` and exit, and the root sums the + ranks in order. +""" + +from typing import Optional + +import torch + +from tensorrt_llm.logger import logger + +from ..ops.gemm.mxfp8_low_m_gemm import mxfp8_low_m_gemm as _low_m_gemm_op + +# (N, K) -> (warps, cluster_k, max_m, row_tiles); measured on GB300 TP4 M3 shapes. The fused kernel beats native +# quantize + GEMM only up to max_m rows; larger M, and every other shape, keeps the native path. These four +# configurations are the ones built into ``torch.ops.trtllm.minimax_m3_mxfp8_low_m_gemm``. +_CONFIGS = { + (2560, 6144): (4, 4, 8, 2), # QKV (q, k, v, index q / k) + (6144, 2048): (8, 2, 8, 4), # o_proj + (6144, 6144): (4, 4, 1, 1), # dense MLP gate_up + (6144, 3072): (4, 2, 8, 2), # dense MLP down +} +# Opt-in shared memory per CTA on SM100 (227 KiB) less a margin for the static mbarriers. +_MAX_DYNAMIC_SMEM = 227 * 1024 - 1024 # capacity: fixed: the device's shared-memory limit + +_logged_shapes = set() + + +def _smem_bytes(m: int, k: int, warps: int, cluster_k: int, row_tiles: int) -> int: + """Dynamic shared memory of one CTA, mirroring the host-side size in ``launch``.""" + kc = k // cluster_k + kc_blocks = kc // 32 + tile = 16 * row_tiles * 8 * (6 if m > 32 else 4 if m > 16 else 2 if m > 8 else 1) + w_bytes = 16 * row_tiles * (kc + 32) + red_bytes = warps * tile * 4 + return ( + w_bytes + + row_tiles * (kc_blocks // 4) * 256 + + m * kc + + (m * kc_blocks + 15) // 16 * 16 + + (cluster_k - 1) * tile * 4 + + (0 if red_bytes <= w_bytes else red_bytes) + ) + + +def get_config(n: int, k: int) -> Optional[tuple]: + """Return (warps, cluster_k, max_m, row_tiles) for an (N, K) weight, or None when no configuration is built.""" + return _CONFIGS.get((n, k)) + + +def mxfp8_low_m_gemm( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + use_pdl: bool = False, + config: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """Fused MXFP8 quantize + GEMM; ``weight_scale`` is the swizzled 128x4 buffer. + + Returns None for unsupported operands so callers keep the native path. + """ + if ( + output_dtype != torch.bfloat16 + or x.dtype != torch.bfloat16 + or torch.cuda.get_device_capability(x.device)[0] < 10 + ): + return None + n, k = weight.shape + config = config or get_config(n, k) + if config is None or config not in _CONFIGS.values(): + return None + warps, cluster_k, max_m, row_tiles = config + if ( + x.shape[0] > max_m + or n % (16 * row_tiles) != 0 + or k % (cluster_k * 128) != 0 + or (k // cluster_k // 32) % warps != 0 + or x.data_ptr() % 16 != 0 + or x.stride(0) % 8 != 0 + or weight.data_ptr() % 16 != 0 + or weight_scale.data_ptr() % 16 != 0 + or _smem_bytes(x.shape[0], k, warps, cluster_k, row_tiles) > _MAX_DYNAMIC_SMEM + ): + return None + if (n, k) not in _logged_shapes: + _logged_shapes.add((n, k)) + logger.info( + f"[MXFP8 low-M] fused decode GEMM engaged: N={n} K={k} " + f"warps={warps} cluster_k={cluster_k} row_tiles={row_tiles} pdl={use_pdl}" + ) + return _low_m_gemm_op(x, weight, weight_scale, warps, cluster_k, row_tiles, use_pdl) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/__init__.py new file mode 100644 index 000000000000..acec8fdbb26f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Normalization entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/m3_embedding_firstnorm.py new file mode 100644 index 000000000000..1842171bde1a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/m3_embedding_firstnorm.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-step embedding gather and the first layer's input RMSNorm in one CuTe DSL launch, into caller buffers.""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm + +__all__ = ["DecodeEmbedding", "m3_embedding_firstnorm"] + +VOCAB = 200064 # MiniMax-M3's vocabulary; the decode table is the full (replicated) embedding + + +class M3EmbeddingFirstnorm(OpWrapper): + """Write ``table[ids]`` into ``raw`` (bf16 [M, 6144], a zero row for an id outside ``[0, V)``) and + ``raw * rsqrt(mean(raw^2) + eps) * weight`` into ``out``; returns ``(out, raw)`` -- the first layer's normed input + and its residual. ``ids`` int32 or int64 [M], M = 1..8; ``table`` bf16 [V, 6144] contiguous; ``weight`` bf16 + [6144]. The first call per (M, id dtype, table shape, PDL) compiles (``DecodeEmbedding`` does it at build).""" + + ARCHS = frozenset({Arch.SM_103}) + + # The target builds every (rows 1..8, id dtype) owner; the kernel is one CTA per row, so the rows differ only in + # grid size: the ends of the range and both id dtypes. + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, int32 ids", spec=dict(rows=1, id_dtype=torch.int32)), + Cell(why="eight requests, int32 ids", spec=dict(rows=8, id_dtype=torch.int32)), + Cell(why="one request, int64 ids", spec=dict(rows=1, id_dtype=torch.int64)), + Cell( + why="six rows (one request verifying a draft), int64 ids", + spec=dict(rows=6, id_dtype=torch.int64), + ), + # Every other row count the decode dispatch takes (dispatch.EMBED_FIRSTNORM), both id dtypes. + *( + Cell( + why=f"{rows} rows: a decode row count between the ones above", + spec=dict(rows=rows, id_dtype=torch.int32 if rows % 2 else torch.int64), + ) + for rows in (2, 3, 4, 5, 7) + ), + ) + + note = """ + The gathered rows are exact. The norm is flashinfer's RMSNorm arithmetic in its reduction order with a fast + rsqrt, where `reference` sums in fp64 and takes torch's rsqrt: within the band the test derives, not bit for bit. + An id outside [0, V) gives a zero row and a zero norm (torch's embedding would raise). No accepted input is known + to compute wrongly: the host side rejects any other dtype, shape or a strided buffer. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return m3_embed_norm(ids, table, weight, eps, raw, out, use_pdl) + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """New tensors (``raw`` / ``out`` are not written); `use_pdl` is ignored.""" + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] * valid.unsqueeze(1).to( + table.dtype + ) + x = rows.double() + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + return (x * rstd * weight.double()).to(torch.bfloat16), rows + + +m3_embedding_firstnorm = M3EmbeddingFirstnorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/__init__.py new file mode 100644 index 000000000000..d25c41dc9557 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Speculative-decoding entries this target alone calls: the Eagle3 verify ops, the capture / fc_norm RMSNorms and +the step-input staging.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_capture_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_capture_norm.py new file mode 100644 index 000000000000..131269ae7d07 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_capture_norm.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Eagle3 hidden-state capture with the drafter's fc_norm folded in (``m3_spec_tail.capture_norm``).""" + +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 + +from .m3_chunk_norm import rmsnorm_rows + + +class M3CaptureNorm(OpWrapper): + """``dst [M, H] = rmsnorm(bf16(h + r)) * w`` (``r`` optional), the rounding chunk_norm applies, into a capture + buffer slice whose row stride may exceed H. Returns ``dst``.""" + + ARCHS = frozenset({Arch.SM_103}) + + # The target's captured layers (hidden 6144) into the 3-layer capture buffer (row stride 3 x 6144), over the + # verify rows of draft_len 3 (4 .. 32 rows); the residual is present for every captured decoder layer. + CELLS: tuple[Cell, ...] = ( + Cell( + why="4 verify rows of one request, layer slot 0, residual added", + spec=dict(m=4, slot=0, residual=True), + ), + Cell( + why="cap 8 (32 rows) into the last layer slot", spec=dict(m=32, slot=2, residual=True) + ), + Cell( + why="no residual (a capture point without one): plain RMSNorm", + spec=dict(m=8, slot=1, residual=False), + ), + ) + + note = """ + fp32 statistic (block sum, rsqrtf), one bf16 rounding after the bf16 sum h + r; the band states the occasional + one-ulp move against torch's statistic. + """ + + def __call__( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_capture_norm(h, r, w, eps, dst) + return dst + + def reference( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + """bf16 sum, then the normalized row; ``dst`` only says where the op writes.""" + x = (h.float() + r.float()).to(torch.bfloat16) if r is not None else h + return rmsnorm_rows(x, w, eps) + + def is_valid( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> None: + # The norm weight is read as H dense bf16 values; a strided view (a column of a stacked weight) is read + # past its elements without an error. + assert w.is_contiguous(), "w must be dense [H]" + + +m3_capture_norm = M3CaptureNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_chunk_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_chunk_norm.py new file mode 100644 index 000000000000..f6adb2207a20 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_chunk_norm.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 fc_norm over the captured layers in one launch: per-chunk RMSNorm of the concatenated hidden states +(``m3_spec_tail.chunk_norm``).""" + +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 + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3ChunkNorm(OpWrapper): + """``dst[r, c*H:(c+1)*H] = rmsnorm(src[r, c*H:(c+1)*H]) * w[c]`` for the C chunks of ``w [C, H]``; ``src`` may + have a row stride wider than C * H (the capture buffer). Writes ``dst [M, C*H]`` and returns it.""" + + ARCHS = frozenset({Arch.SM_103}) + + # The Eagle3 capture buffer of the target (3 captured layers x hidden 6144, buffer rows wider than the layers + # when the final hidden is also captured), over the verify rows of draft_len 3: 4 / 8 / 32 rows. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request's 4 verify rows, dense [4, 3 x 6144]", + spec=dict(m=4, c=3, h=6144, extra_cols=0), + ), + Cell( + why="cap 8 (32 rows) read from the capture buffer with a fourth (final-hidden) slot in the row stride", + spec=dict(m=32, c=3, h=6144, extra_cols=6144), + ), + ) + + note = """ + fp32 statistic per chunk (block sum, rsqrtf), one bf16 rounding: an output moves by one bf16 ulp now and then + against torch's statistic; the cells' band states it. + """ + + def __call__( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + torch.ops.trtllm.minimax_m3_chunk_norm(src, w, eps, dst, pdl) + return dst + + def reference( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + """Each chunk normalized on its own; ``dst`` only says where the op writes, ``pdl`` shapes the launch.""" + c, h = w.shape + x = src[:, : c * h].reshape(src.shape[0], c, h) + return rmsnorm_rows(x, w, eps).reshape(src.shape[0], c * h) + + def is_valid( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ) -> None: + c, h = w.shape + assert w.is_contiguous(), "w must be dense [C, H] (chunk c read at c * H)" + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + + +m3_chunk_norm = M3ChunkNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_draft_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_draft_input.py new file mode 100644 index 000000000000..75f7a44dfcfb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_draft_input.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 draft pass's input in one launch: embedding lookup, both RMSNorms and the concat +(``m3_spec_tail.draft_input``).""" + +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 + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3DraftInput(OpWrapper): + """``out [M, 2H] = [rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden]`` and ``residual [M, H] = + hidden[rows]`` (rows default to 0..M-1) in one launch; ids outside the table embed to zero (the masked lookup of + a vocab-local embedding). Returns ``(out, residual)``.""" + + ARCHS = frozenset({Arch.SM_103}) + + # Every pass of the fused draft loop (dispatch.DRAFT_INPUT) at draft_len 3, drafter hidden 6144, eps 1e-6: the + # first pass of 1..8 requests (4..32 rows, the carried hidden read in place) and the later passes of 1..8 requests + # (one row each; the second pass gathers the last accepted rows from the first pass's 4 rows per request). + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"first pass of {n} request(s): {4 * n} rows, hidden read in place", + spec=dict(m=4 * n, h=6144, gather=False, invalid_id=False), + ) + for n in range(1, 9) + ), + *( + Cell( + why=f"second pass of {n} request(s): the last accepted rows gathered from {4 * n} rows", + spec=dict(m=n, h=6144, gather=True, invalid_id=False), + ) + for n in range(1, 9) + ), + Cell( + why="third pass of 3 requests: one row each, hidden read in place", + spec=dict(m=3, h=6144, gather=False, invalid_id=False), + ), + Cell( + why="one row with an out-of-table id: the embedding half is zero", + spec=dict(m=1, h=6144, gather=False, invalid_id=True), + ), + ) + + note = """ + The arithmetic is the in-tree RMSNorm's (fp32 statistic, (x * rsqrt) * w, one bf16 rounding); the statistic's + summation order differs from FlashInfer's, which moves an output by one bf16 ulp now and then (the test's band, + also held against the in-tree modules). The out-of-table id is the masked lookup of a vocab-local embedding; the + drafter's embedding is the whole table, so its ids are always in range. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + torch.ops.trtllm.minimax_m3_draft_input( + ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, pdl + ) + return out, residual + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """The two normalized halves (fp32 statistic, one bf16 rounding) and the gathered hidden rows. ``out`` / + ``residual`` only say where the op writes; ``pdl`` only shapes the launch.""" + valid = (ids >= 0) & (ids < table.shape[0]) + emb = rmsnorm_rows(table[torch.where(valid, ids, 0).long()], w_embed, eps_embed) + emb = torch.where(valid[:, None], emb, torch.zeros_like(emb)) + hid = hidden[rows.long()] if rows is not None else hidden[: ids.shape[0]] + ref = torch.cat([emb, rmsnorm_rows(hid, w_hidden, eps_hidden)], dim=-1) + return ref, (hid.clone() if residual is not None else None) + + def is_valid( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> None: + h = table.shape[1] + # Rows are read as hidden + rows[r] * H and table + id * H: both must be dense H-wide rows. + assert table.is_contiguous() and hidden.is_contiguous(), ( + "table / hidden must be dense [*, H]" + ) + if rows is not None: + assert bool(((rows >= 0) & (rows < hidden.shape[0])).all()), "rows must index hidden" + else: + assert hidden.shape[0] >= ids.shape[0], "hidden must hold a row per id" + # Vector loads of 16 bytes: H must be a multiple of 8 within the register tile, and the weights aligned. + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + assert residual is None or ( + residual.is_contiguous() and residual.shape == (ids.shape[0], h) + ), "residual must be contiguous [M, H]" + + +m3_draft_input = M3DraftInput() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_draft_step.py new file mode 100644 index 000000000000..a4750f4c1238 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_draft_step.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One pass of the fused greedy Eagle3 draft loop: the drafter's vocab-sharded BF16 head (streamed from its lossless +12-bit coding, ``ops.gemm.m3_target_head``, or as BF16 fragments), the cross-rank argmax and the loop's per-pass +bookkeeping in one kernel (``m3_draft_step.draft_step``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from ..gemm.m3_target_head import decode as decode_coded +from .m3_shard_accept import exchange_tokens, local_argmax, make_keys + +#: Requests (head rows) per pass, and exchange rows per call half of the workspace. +MAX_TOKENS = 8 +#: 16-row weight tiles one CTA can own (the kernel's shared-memory accumulator). +MAX_TILES_PER_CTA = 32 +#: Draft length of the fused loop, and the drafter's TP4 head shard (200064-id vocabulary) and hidden size. +DRAFT_LEN = 3 +DRAFT_SHARD = 50016 +HIDDEN = 6144 + + +def bf16_fragments(weight: torch.Tensor) -> torch.Tensor: + """BF16 [N, K] -> the same values in A-fragment order [N / 16, K / 16, 256]. + + One lane's 16 bytes of a 16-row tile's k16 step are the m16n8k16 A fragment: element (row h * 8 + g, k ks * 16 + + hk * 8 + jj * 2 + e) goes to [tile][ks][g][jj][hk][h][e] (lane = g * 4 + jj). + """ + n, k = weight.shape + assert n % 16 == 0 and k % 16 == 0 and weight.dtype == torch.bfloat16, (n, k, weight.dtype) + q = weight.view(n // 16, 2, 8, k // 16, 2, 4, 2).permute(0, 3, 2, 5, 4, 1, 6).contiguous() + return q.view(n // 16, k // 16, 256) + + +def from_fragments(wf: torch.Tensor) -> torch.Tensor: + """The inverse of ``bf16_fragments``: [N / 16, K / 16, 256] -> [N, K].""" + tiles, ks, _ = wf.shape + return ( + wf.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, ks * 16) + ) + + +def head_weight( + wf: torch.Tensor, + k: int, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The BF16 head shard [N, K] the kernel multiplies by: the fragment-ordered copy, or (``exc_off`` given) the + 12-bit coded stream with its exception lists decoded.""" + if exc_off is None: + return from_fragments(wf) + n = (exc_off.numel() - 1) * 16 + return decode_coded(wf, exp_base, exc_off, exc_pos, exc_w, n, k) + + +def shard_logits( + x: torch.Tensor, + wf: torch.Tensor, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The in-tree head's BF16 logits of the rows x [M, K] against the shard (``head_weight``): the fp32 product, + rounded once to BF16 (the BF16 GEMM's output), as fp32.""" + w = head_weight(wf, x.shape[1], exc_off, exc_pos, exc_w, exp_base) + with true_fp32_matmul(): + logits = x.float() @ w.float().t() + return logits.to(torch.bfloat16).float() + + +class M3DraftStep(OpWrapper): + """One fused greedy draft pass over ``rows <= 8`` requests (one head row each). + + Head input ``x[x_rows[b]]`` (pass 0: the last accepted verify row of each request) or ``x[b]``. The head is this + rank's BF16 lm_head shard in fragment order (``m3_draft_step.bf16_fragments``, global ids ``vocab_offset ..``), + or with ``exc_off`` / ``exc_pos`` / ``exc_w`` / ``exp_base`` its lossless 12-bit coding (``wf`` = the uint8 stream + of ``m3_target_head.encode_chunked``, which must round-trip to the BF16 shard); each logit is the fp32 dot product + rounded to BF16 and the row's token is the global argmax over all ranks' + shards (ties, -0 / +0 included, to the lowest global id), exchanged through the multicast slots + ``[2][8][world]``. The last CTA then writes ``draft_out[b, step]`` and ``ids_next[b]`` (the token), ``pos_next[b] + = pos_src[x_rows[b] or b] + 1``, and advances the draft KV lengths as the in-tree draft loop does between passes + (pass 0 with ``num_acc``: ``kv_lens -= draft_len - num_acc`` and ``seq_lens = 1``; later passes: ``kv_lens += + 1``). ``out_keys`` (optional) receives each row's winning key. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # Every pass of the fused draft loop (dispatch.DRAFT_STEP): draft_len 3 at 1..8 requests, the drafter's + # full-vocabulary BF16 head (50016 rows per rank of 4, hidden 6144). Pass 0 reads the verify rows' hidden (4 per + # request) through x_rows and rolls the KV back; the later passes read one row per request. + _PASSES: tuple[Cell, ...] = ( + *( + Cell( + why=f"pass 0 at {n} request(s): head rows gathered from {4 * n} verify rows, KV rollback, seq_lens = 1", + spec=dict( + nreq=n, + draft_len=DRAFT_LEN, + step=0, + rank=n % 4, + world=4, + vs=DRAFT_SHARD, + k=HIDDEN, + ), + ) + for n in range(1, MAX_TOKENS + 1) + ), + Cell( + why="pass 1 at two requests: one row each, kv_lens + 1, rank 3's ids", + spec=dict( + nreq=2, draft_len=DRAFT_LEN, step=1, rank=3, world=4, vs=DRAFT_SHARD, k=HIDDEN + ), + ), + Cell( + why="last pass (step 2) at 8 requests, rank 0", + spec=dict( + nreq=8, draft_len=DRAFT_LEN, step=2, rank=0, world=4, vs=DRAFT_SHARD, k=HIDDEN + ), + ), + Cell( + why="pass 2 at 5 requests with exact cross-rank ties (the lower global id wins)", + spec=dict( + nreq=5, + draft_len=DRAFT_LEN, + step=2, + rank=1, + world=4, + vs=DRAFT_SHARD, + k=HIDDEN, + ties=True, + ), + ), + ) + # The coded head (the fused draft loop's default) at every pass shape, and the BF16 fragments (the coding's + # fail-closed path) at the first and last pass shapes and with ties. + CELLS: tuple[Cell, ...] = ( + *(Cell(why=f"coded: {c.why}", spec=dict(c.spec, coded=True)) for c in _PASSES), + *( + Cell(why=f"bf16: {c.why}", spec=dict(c.spec, coded=False)) + for c in (_PASSES[0], _PASSES[MAX_TOKENS - 1], _PASSES[-1]) + ), + ) + + note = """ + The head is the in-tree BF16 head's arithmetic (BF16 weights and activations, fp32 accumulation, a BF16 logit; + the coded head's weight is the BF16 shard bit for bit; its exceptions decode to +0 in the MMAs and their exact + products are added at the tile's flush) + and the argmax is the in-tree one over BF16 logits, so a token differs from the in-tree path's only where the fp32 + sum, accumulated in another order, rounds to the other side of a BF16 boundary and that decides the row's maximum. + The cells' rows are decidable and the peers' keys sit 0.25 away (or tie exactly); the edge tests (forced BF16 + ties, -0 / +0, top-2 one fp32 step apart that round to one BF16 value, adjacent BF16 values) use integer data + whose fp32 sums are exact in any order and hold the kernel to the in-tree path bit for bit. The cross-rank + exchange is certified on one device with the peers' keys written into the slots (see m3_shard_accept). + """ + + def __call__( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + ) -> None: + torch.ops.trtllm.minimax_m3_draft_step( + x, + x_rows, + rows, + wf, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + num_ctas, + pdl, + out_keys, + exc_off, + exc_pos, + exc_w, + exp_base, + ) + + def reference( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + ) -> dict[str, torch.Tensor]: + """Every buffer the pass writes, computed on copies: "logits" (this rank's BF16 logits as fp32), "local_keys" + (this rank's key per row), "keys" (the winning key), "tokens", "draft_out", "ids_next", "pos_next", "kv_lens" + and "seq_lens". + + Reads ``ws_uc``, ``counters`` and ``kv_lens`` before the call. ``keys``, ``ws_mc_ptr``, ``num_ctas``, ``pdl`` + and ``out_keys`` do not change the result. + """ + h = x[x_rows.long()] if x_rows is not None else x[:rows] + logits = shard_logits(h, wf, exc_off, exc_pos, exc_w, exp_base) + best, idx = local_argmax(logits, vocab_offset) + local_keys = make_keys(best, idx) + tokens = exchange_tokens(local_keys, ws_uc, counters, rank, world, MAX_TOKENS) + half = int(counters[1].item()) & 1 + slots = ws_uc.view(torch.int64)[ + half * MAX_TOKENS * world : (half * MAX_TOKENS + rows) * world + ] + slots = slots.view(rows, world).clone() + slots[:, rank] = local_keys + sign = torch.iinfo(torch.int64).min + win = (slots ^ sign).max(dim=1).values ^ sign + d_out = draft_out.clone() + d_out[:rows, step] = tokens + src = x_rows.long() if x_rows is not None else torch.arange(rows, device=x.device) + lens = kv_lens.clone() + sl = seq_lens.clone() if seq_lens is not None else None + if step == 0 and num_acc is not None: + lens[:rows] -= draft_len - num_acc[:rows] + sl[:rows] = 1 + elif step > 0: + lens[:rows] += 1 + out = { + "logits": logits, + "local_keys": local_keys, + "keys": win, + "tokens": tokens, + "draft_out": d_out, + "ids_next": tokens.clone(), + "pos_next": pos_src[src] + 1, + "kv_lens": lens, + } + if sl is not None: + out["seq_lens"] = sl + return out + + def is_valid( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + ) -> None: + assert not bool(keys[:rows].any()), "keys must start zeroed" + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + # The kernel reads the head row x + x_rows[b] * K and the position pos_src[x_rows[b]] with no bound. + if x_rows is not None: + assert bool(((x_rows[:rows] >= 0) & (x_rows[:rows] < x.shape[0])).all()), ( + "x_rows must index x" + ) + assert bool((x_rows[:rows] < pos_src.shape[0]).all()), ( + "pos_src must cover the head rows" + ) + else: + assert pos_src.shape[0] >= rows, "pos_src must cover the head rows" + # Each CTA accumulates at most MAX_TILES_PER_CTA tiles in shared memory. + tiles = wf.shape[0] if exc_off is None else exc_off.numel() - 1 + assert -(-tiles // num_ctas) <= MAX_TILES_PER_CTA, "too many weight tiles per CTA" + if exc_off is not None: + # The exceptions are read per tile in list order with no bound against the coded stream. + assert int(exc_off[-1].item()) == exc_pos.numel(), ( + "exc_off must end at the exception count" + ) + assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" + # A NaN logit maps to a key above every finite one and wins the row silently. + h = x[x_rows.long()] if x_rows is not None else x[:rows] + assert not bool(torch.isnan(h).any()), "NaN head inputs win the argmax" + if num_acc is not None: + na = num_acc[:rows] + assert bool(((na >= 1) & (na <= draft_len + 1)).all()), ( + "num_acc counts must lie in [1, draft_len + 1]" + ) + + +m3_draft_step = M3DraftStep() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_force_vec.py new file mode 100644 index 000000000000..af4b06e3999f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_force_vec.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance, drawn as the in-tree worker draws them.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..m3_force_vec import force_vec_kernel + + +class M3ForceVec(OpWrapper): + """Bump ``counter`` (int64 [1]) by one, then for each of the ``batch_size`` requests ``b`` write + ``min(base_total + (pool[(c * counter_stride + b * slot_stride) & pool_mask] < frac), max_total)`` into ``out`` + (int32), ``c`` the bumped counter: the in-tree ``_apply_force_accepted_tokens`` draw over its fixed-seed fp32 + ``pool``. ``block`` is the launch width (a power of two >= batch_size; 16, 32 or 64 are compiled at warmup). + + Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) + runs it. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), + Cell(why="eight requests, fraction 0.3", spec=dict(batch=8, block=16, frac=0.3)), + ) + + note = """ + Integer arithmetic and one fp32 comparison: `reference` must match exactly. + """ + + def __call__( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + force_vec_kernel[(1,)]( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK=block, + ) + return out + + def reference( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + """The counts (a new tensor; ``counter`` and ``out`` are not written).""" + c = int(counter.item()) + 1 + slots = torch.arange(batch_size, dtype=torch.int64, device=pool.device) + idx = (c * counter_stride + slots * slot_stride) & pool_mask + r = pool[idx] + return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) + + +m3_force_vec = M3ForceVec() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_sampled_accept.py new file mode 100644 index 000000000000..e904608fec13 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_sampled_accept.py @@ -0,0 +1,269 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sampled (temperature + top-p) Eagle3 verification over the TP-sharded vocabulary in one kernel +(``m3_sampled_accept.sampled_accept``).""" + +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 + +from .m3_shard_accept import strict_accept + +#: Verify rows per call. +MAX_ROWS = 48 +#: The fused sampling kernel's 1 / T guard. +TEMP_EPSILON = 1e-6 + + +def nucleus_probs( + logits: torch.Tensor, temperatures: torch.Tensor, top_ps: torch.Tensor +) -> torch.Tensor: + """Per-row sampling distribution [rows, Vs] (fp32): w = exp(x / (T + 1e-6) - max), kept iff the mass strictly + above its weight is at most top_p * sum(w) (ties included), renormalized over the kept tokens.""" + inv = 1.0 / (temperatures.float() + TEMP_EPSILON) + x = logits.float() * inv[:, None] + w = torch.exp(x - x.max(dim=1, keepdim=True).values).double() + z = w.sum(dim=1, keepdim=True) + order = torch.sort(w, dim=1, descending=True) + sw = order.values + # mass strictly above each sorted weight: the cumulative sum before the first entry of its tie group + csum = torch.cumsum(sw, dim=1) - sw + first = torch.ones_like(sw, dtype=torch.bool) + first[:, 1:] = sw[:, 1:] != sw[:, :-1] + group = torch.cumsum(first.long(), dim=1) - 1 + starts = torch.zeros_like(csum) + starts.scatter_reduce_(1, group, csum, reduce="amin", include_self=False) + above = torch.gather(starts, 1, group) + keep_sorted = above <= top_ps.double()[:, None] * z + keep = torch.zeros_like(keep_sorted) + keep.scatter_(1, order.indices, keep_sorted) + p = torch.where(keep, w, torch.zeros_like(w)) + return (p / p.sum(dim=1, keepdim=True)).float() + + +class M3SampledAccept(OpWrapper): + """Sampled verification of ``rows = nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard. + + Each row draws one token from its top-p nucleus of softmax(logits / T) over the full (all-rank) vocabulary by + distributed rejection sampling (Philox keyed by the row's seed / offset, the row, the attempt and the CTA), then + the strict acceptance runs against ``draft`` as in ``m3_shard_accept``: ``acc_tokens``, ``num_acc``, + ``gather_ids``, optional ``next_new[:, 0]``. ``tokens`` receives the sampled token per row; ``extra_attempts`` + accumulates exchange rounds beyond the second. The exchange goes through the multicast workspace (two call + halves of ``minimax_m3_sampled_accept_ws_words(world)`` words) with the (finished rows, call) ``counters``. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # draft_len 3 (T = 4 rows per request) at the TP4 caps: 1 / 2 / 8 requests (4 / 8 / 32 rows, 8 CTAs per row) + # and the 48-row maximum (4 CTAs per row); the 50016-column target shard. Single rank (see note). + CELLS: tuple[Cell, ...] = ( + Cell( + why="cap 1 (4 rows), top_p 0.8 over a 12-token head: the draws follow the nucleus distribution", + spec=dict( + nreq=1, vs=50016, vocab_offset=0, check="distribution", top_p=0.8, temperature=0.7 + ), + ), + Cell( + why="cap 8 (32 rows, 256 co-resident CTAs): top_p ~ 0 keeps only the argmax, token exact", + spec=dict( + nreq=8, + vs=50016, + vocab_offset=2 * 50016, + check="argmax", + top_p=1e-7, + temperature=1.0, + ), + ), + Cell( + why="cap 2 (8 rows) with top_p = 1 (no cut: the first draw stands), next_new and forced counts", + spec=dict( + nreq=2, + vs=50016, + vocab_offset=50016, + check="distribution", + top_p=1.0, + temperature=1.0, + force="vec", + ), + ), + Cell( + why="48 rows (4 CTAs per row): the second launch geometry, nucleus distribution at top_p 0.5", + spec=dict( + nreq=12, + vs=50016, + vocab_offset=3 * 50016, + check="distribution", + top_p=0.5, + temperature=1.3, + force="total", + ), + ), + ) + + note = """ + What is sampled is gated statistically, the rest exactly. Every drawn token must lie in the reference nucleus; + the empirical distribution over many calls (distinct Philox offsets) must sit within a stated total-variation + band of the reference's; the acceptance outputs must equal the strict acceptance of the drawn tokens exactly. + The Philox stream itself is not reproduced in torch: the kernel's draw uses the hardware exp / log + approximations (__expf, __logf) and a CTA-order inverse CDF, which torch does not reproduce bit for bit. + + Certified at world 1 only. The decision is replicated from every rank's per-round records, which another rank + computes from its own shard on the fly, so peers cannot be stood in for on one device; the cross-rank path + needs the multi-device job and is not covered here. Rows with top-k or min-p must not reach this op (it applies + neither); is_valid states it. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> None: + torch.ops.trtllm.minimax_m3_sampled_accept( + logits, + vocab_offset, + temperatures, + top_ps, + seeds, + offsets, + ws_uc, + ws_mc_ptr, + counters, + tokens, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + force_vec, + extra_attempts, + use_pdl, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> dict: + """{"probs": the per-row nucleus distribution over this shard [rows, Vs] (global id = vocab_offset + + column), "accept": tokens -> the strict acceptance outputs}. The draw itself is random; ``seeds`` / + ``offsets`` select it and do not change the distribution. World 1: this shard is the whole vocabulary.""" + rows = logits.shape[0] + probs = nucleus_probs(logits, temperatures[:rows], top_ps[:rows]) + + def accept(drawn: torch.Tensor) -> dict[str, torch.Tensor]: + return strict_accept(drawn, draft, batch_idx, force_total, next_new, force_vec) + + return {"probs": probs, "accept": accept} + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> None: + rows = logits.shape[0] + # top-k and min-p are not applied: a row that has them gets a token outside its real sampling set. + if top_ks is not None: + k = top_ks[:rows] + assert bool(((k <= 0) | (k >= logits.shape[1] * world)).all()), ( + "top-k rows must take the generic path" + ) + if min_ps is not None: + assert not bool((min_ps[:rows] > 0).any()), "min-p rows must take the generic path" + # Rows are request-major verify rows: tokens[b * (D + 1) + k]; a mismatched layout pairs drafts with the + # wrong rows. + assert rows == draft.shape[0] * (draft.shape[1] + 1), ( + "logits rows must be nreq * (draft_len + 1)" + ) + # Per-row parameters are read at the row index, unchecked. + for t in (temperatures, top_ps, seeds, offsets): + assert t.numel() >= rows and t.is_contiguous(), ( + "per-row parameters must cover every row" + ) + assert temperatures.dtype == torch.float32 and top_ps.dtype == torch.float32, ( + "fp32 sampling parameters" + ) + assert seeds.dtype == torch.int64 and offsets.dtype == torch.int64, ( + "int64 Philox seeds / offsets" + ) + # The finished-row counter elects the CTA that runs the acceptance. + assert int(counters[0].item()) == 0, ( + "counters[0] (finished rows) must be zero between calls" + ) + + +m3_sampled_accept = M3SampledAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_sampled_accept_ws_words.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_sampled_accept_ws_words.py new file mode 100644 index 000000000000..3e42fd5c50ad --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_sampled_accept_ws_words.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Size of the sampled-verify exchange workspace (``m3_sampled_accept.SampledAcceptWorkspace.create``).""" + +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 + +#: Exchange rounds per row, verify rows per call, CTAs per row (at most), uint64 words per record. +MAX_ATTEMPTS = 16 +MAX_ROWS = 48 +MAX_SPLITS = 8 +RECORD_WORDS = 4 + + +class M3SampledAcceptWsWords(OpWrapper): + """uint64 words of the multicast workspace ``minimax_m3_sampled_accept`` exchanges through for ``world`` ranks: + two call halves of ``[attempts][rows][world * splits][record words]``.""" + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="the tp4 target's workspace", spec=dict(world=4)), + Cell(why="a single rank (local workspace, no multicast)", spec=dict(world=1)), + ) + + note = """ + A host computation (no device work); it is an op so the workspace size and the kernel's indexing come from one + definition. + """ + + def __call__(self, world: int) -> int: + return torch.ops.trtllm.minimax_m3_sampled_accept_ws_words(world) + + def reference(self, world: int) -> int: + """Two halves of attempts x rows x (world x splits) records x record words.""" + return 2 * MAX_ATTEMPTS * MAX_ROWS * world * MAX_SPLITS * RECORD_WORDS + + +m3_sampled_accept_ws_words = M3SampledAcceptWsWords() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_shard_accept.py new file mode 100644 index 000000000000..90bb11ef84b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_shard_accept.py @@ -0,0 +1,284 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Greedy Eagle3 verification from this rank's logits shard: per-row argmax, the cross-rank (value, index) exchange +and the strict acceptance in one kernel (``m3_shard_accept.shard_accept``).""" + +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 + +#: Verify rows per call, and exchange rows per call half of its workspace. +SHARD_ROWS = 48 + +_LOW32 = 0xFFFFFFFF + + +def ordered_bits(v: torch.Tensor) -> torch.Tensor: + """The kernel's order-preserving uint32 image of fp32 values, as int64 (-0.0 maps to +0.0's image: they tie).""" + b = v.float().contiguous().view(torch.int32).to(torch.int64) & _LOW32 + b = torch.where(b == 0x80000000, torch.zeros_like(b), b) + return torch.where(b >= 0x80000000, _LOW32 - b, b | 0x80000000) + + +def make_keys(values: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """Exchange keys of (fp32 value, global token index) pairs: larger value wins, ties take the lower index.""" + return (ordered_bits(values) << 32) | (_LOW32 - index.to(torch.int64)) + + +def unsigned_max(keys: torch.Tensor, dim: int) -> torch.Tensor: + """Max of int64-stored uint64 keys (the kernel compares them unsigned).""" + sign = torch.iinfo(torch.int64).min + return (keys ^ sign).max(dim=dim).values ^ sign + + +def key_token(key: torch.Tensor) -> torch.Tensor: + return (_LOW32 - (key & _LOW32)).to(torch.int32) + + +def local_argmax(logits: torch.Tensor, vocab_offset: int) -> tuple[torch.Tensor, torch.Tensor]: + """(max value, lowest global index reaching it) per row of a [M, Vs] fp32 shard.""" + best = logits.max(dim=1).values + idx = (logits == best[:, None]).float().argmax(dim=1) # first (lowest) index at the max + return best, idx.to(torch.int64) + vocab_offset + + +def exchange_tokens( + own_keys: torch.Tensor, + ws_uc: torch.Tensor, + counters: torch.Tensor, + rank: int, + world: int, + slot_rows: int, +) -> torch.Tensor: + """Global tokens of the cross-rank pick: this rank's keys against the other ranks' keys already in its slots. + + Slot (half * slot_rows + t) * world + r of ``ws_uc`` holds rank r's key of row t; half = call counter & 1. A + single-device run stands in for the other ranks by writing their keys there before the call. + """ + m = own_keys.shape[0] + half = int(counters[1].item()) & 1 + slots = ( + ws_uc.view(torch.int64)[half * slot_rows * world : (half * slot_rows + m) * world] + .view(m, world) + .clone() + ) + slots[:, rank] = own_keys + return key_token(unsigned_max(slots, dim=1)) + + +def strict_accept( + tokens: torch.Tensor, + draft: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], +) -> dict[str, torch.Tensor]: + """The strict acceptance epilogue from the verify-row tokens [nreq * (D + 1)] (request-major). + + accepted[b] = tokens of request b's rows; num_accepted[b] = 1 + the matching prefix of draft[b] (overridden by + force_vec[b], else by force_total > 0); gather_ids[b] = batch_idx[b] * (D + 1) + num_accepted[b] - 1, and + next_new[b, 0] (when given) = the token at that row. Returns the outputs plus ``next_new`` as written. + """ + nreq, d = draft.shape + k1 = d + 1 + tokens = tokens.to(torch.int32) + accepted = tokens.view(nreq, k1).clone() + match = (draft.to(torch.int32) == accepted[:, :d]).to(torch.int32) + num = 1 + match.cumprod(dim=1).sum(dim=1) + if force_vec is not None: + num = force_vec[:nreq].to(num.dtype) + elif force_total > 0: + num = torch.full_like(num, force_total) + gather = batch_idx.to(torch.int64) * k1 + num.to(torch.int64) - 1 + out = {"accepted": accepted, "num_accepted": num.to(torch.int32), "gather_ids": gather} + if next_new is not None: + nn = next_new.clone() + nn[:, 0] = tokens[gather] + out["next_new"] = nn + return out + + +class M3ShardAccept(OpWrapper): + """Greedy verification of ``nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard ``[rows, Vs]`` + (global ids ``vocab_offset ..``). + + Each row's token is its global argmax over all ranks' shards: keys ``ordered(value) << 32 | ~index`` (ties take + the lowest index, as ``torch.argmax``) are max-reduced over the shard, pushed into every rank's exchange slot + through the multicast slots ``[2][48][world]`` (Lamport: 0 = empty) and max-reduced over the ranks by the last + CTA, which then runs the strict acceptance: ``acc_tokens [nreq, D + 1]``, ``num_acc [nreq]`` (``force_vec`` / + ``force_total`` override the counts), ``gather_ids [nreq]`` (the first draft pass's row per request) and, when + given, ``next_new[:, 0]`` (the last accepted token). + """ + + ARCHS = frozenset({Arch.SM_103}) + + # Every greedy verify step of the target (dispatch.SHARD_ACCEPT): draft_len 3 at 1..12 requests (4..48 rows), the + # target lm_head's BF16 shard logits (50016 columns per rank at TP4). The logits come from the head GEMM, so the + # op only reduces. + CELLS: tuple[Cell, ...] = ( + Cell( + why="cap 1: 4 verify rows (1 request x T = 4), rank 0 of 4, next_new staged", + spec=dict(nreq=1, draft_len=3, vs=50016, world=4, rank=0, force="none", strided=False), + ), + Cell( + why="cap 2: 8 verify rows, rank 2, global ids offset by 2 x 50016", + spec=dict(nreq=2, draft_len=3, vs=50016, world=4, rank=2, force="none", strided=False), + ), + Cell( + why="cap 4: 16 verify rows (4 requests x T = 4), rank 1 of 4, next_new staged", + spec=dict(nreq=4, draft_len=3, vs=50016, world=4, rank=1, force="none", strided=False), + ), + Cell( + why="cap 8: 32 verify rows, rank 3, global ids offset by 3 x 50016", + spec=dict(nreq=8, draft_len=3, vs=50016, world=4, rank=3, force="none", strided=False), + ), + Cell( + why="12 rows (3 requests) read from a row-strided view of a wider logits buffer, forced counts", + spec=dict(nreq=3, draft_len=3, vs=50016, world=4, rank=0, force="vec", strided=True), + ), + Cell( + why="the 48-row maximum (12 requests at draft 3), force_total, rows with exact cross-rank ties", + spec=dict( + nreq=12, + draft_len=3, + vs=50016, + world=4, + rank=2, + force="total", + strided=False, + ties=True, + ), + ), + # Every other request count of the verify range, on the 4-rank 50016-column shards. + *( + Cell( + why=f"{n} requests ({4 * n} verify rows): a request count between the ones above", + spec=dict( + nreq=n, draft_len=3, vs=50016, world=4, rank=n % 4, force="none", strided=False + ), + ) + for n in (5, 6, 7, 9, 10, 11) + ), + ) + + note = """ + The op reduces BF16 logits exactly, so the reference is bit-exact, including cross-rank ties (lowest global + index wins). The cross-rank exchange is certified on one device: the other ranks' keys are written into this + rank's slots before the call, which is what their multicast pushes would have put there; the multicast mapping + itself (ws_mc_ptr = this rank's slots, as at world 1) is not exercised across devices here. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + torch.ops.trtllm.minimax_m3_shard_accept( + logits, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + num_ctas, + force_vec, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> dict[str, torch.Tensor]: + """Lowest-index maximum of each row's BF16 logits, the largest key among it and the peers' keys in + ``ws_uc``, then the strict acceptance. Reads ``ws_uc`` / ``counters`` before the call; the scratch and + output buffers and ``num_ctas`` do not change the result.""" + best, idx = local_argmax(logits.float(), vocab_offset) + tokens = exchange_tokens(make_keys(best, idx), ws_uc, counters, rank, world, SHARD_ROWS) + out = strict_accept(tokens, draft, batch_idx, force_total, next_new, force_vec) + out["tokens"] = tokens + return out + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + assert not bool(keys[: logits.shape[0]].any()), "keys must start zeroed" + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + assert ws_uc.numel() >= 2 * SHARD_ROWS * world, "ws_uc must hold [2][48][world] slots" + # A NaN logit maps to a key above every finite one and wins the row silently. + assert not bool(torch.isnan(logits).any()), "NaN logits win the argmax" + nreq = draft.shape[0] + assert bool(((batch_idx[:nreq] >= 0) & (batch_idx[:nreq] < nreq)).all()), ( + "batch_idx must index the nreq requests" + ) + if force_vec is not None: + fv = force_vec[:nreq] + assert bool(((fv >= 1) & (fv <= draft.shape[1] + 1)).all()), ( + "force_vec counts must lie in [1, D + 1]" + ) + + +m3_shard_accept = M3ShardAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_slot_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_slot_scatter.py new file mode 100644 index 000000000000..470a6372f2b6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_slot_scatter.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The spec sampler's per-slot stores of one step's outputs, in one launch (``M3SpecStepFold.scatter`` -> +``m3_spec_tail.slot_scatter``).""" + +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 padded_rows(src: torch.Tensor, width: int) -> torch.Tensor: + """src rows zero-padded (or cut) to ``width`` columns.""" + out = torch.zeros(src.shape[0], width, dtype=src.dtype, device=src.device) + w = min(width, src.shape[1]) + out[:, :w] = src[:, :w] + return out + + +class M3SlotScatter(OpWrapper): + """For i < len(slots), all int32 2-D strided views: ``da[slots[i]] = sa[i]``, ``db[slots[i]] = sb[i]``, + ``dd[slots[i]] = sd[i]`` (each zero-padded / cut to the destination width) and ``dl[slots[i]] = sl[i]``. The + destinations are the store's transposed [slots, width] views.""" + + ARCHS = frozenset({Arch.SM_103}) + + # The sampler store at draft_len 3 (8 slots): new tokens / next new tokens [slots, 4] (transposed views of + # [4, slots, 1]), next draft tokens [slots, 3]; a worker output narrower than the store (dynamic draft length) + # is zero-padded. + CELLS: tuple[Cell, ...] = ( + Cell( + why="8 rows into shuffled slots, widths equal to the store's", + spec=dict(n=8, slots=8, src_w=4, src_d=3, store_w=4, store_d=3), + ), + Cell( + why="3 rows of a narrower step (runtime draft 2): zero padding to the store width", + spec=dict(n=3, slots=8, src_w=3, src_d=2, store_w=4, store_d=3), + ), + ) + + note = """ + Pure index copies: the reference is exact. Slots are not bounds-checked; is_valid states it. + """ + + def __call__(self, slots, sa, da, sb, db, sl, dl, sd, dd): + torch.ops.trtllm.minimax_m3_slot_scatter(slots, sa, da, sb, db, sl, dl, sd, dd) + + def reference(self, slots, sa, da, sb, db, sl, dl, sd, dd): + """The four destinations after the call, on copies.""" + s = slots.long() + out = {} + for name, src, dst in (("da", sa, da), ("db", sb, db), ("dd", sd, dd)): + r = dst.clone() + r[s] = padded_rows(src[: s.numel()], dst.shape[1]) + out[name] = r + lr = dl.clone() + lr[s] = sl[: s.numel()] + out["dl"] = lr + return out + + def is_valid(self, slots, sa, da, sb, db, sl, dl, sd, dd): + assert bool(((slots >= 0) & (slots < da.shape[0]) & (slots < dl.numel())).all()), ( + "slots must index the store" + ) + # Two rows into one slot race (which one lands is unordered). + assert slots.unique().numel() == slots.numel(), "slots must be distinct" + assert sa.shape[0] >= slots.numel() and sl.numel() >= slots.numel(), "a source row per slot" + + +m3_slot_scatter = M3SlotScatter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_step_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_step_gather.py new file mode 100644 index 000000000000..abea2cf620e9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_step_gather.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The overlap scheduler's per-step input gathers from the previous step's device outputs, in one launch +(``M3SpecStepFold.gather`` -> ``m3_spec_tail.step_gather``).""" + +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 + + +class M3StepGather(OpWrapper): + """For the ``P`` previous requests (sequence slots ``slots``, ``T`` tokens and ``D`` draft tokens each), all + int32:: + + ids[ids0 + p*T + t] = new_tok[t, slots[p]] (new_tok [T', slots, (1)]) + draft[d0 + p*D + d] = next_draft[slots[p], d] + pos_off[i] = lens[ppi[i - b0*T]] for i in [b0*T, (b0+P)*T), else 0 (every i < pos_off.numel()) + kv_off[i] = lens[slots[i - b0]] - T for i in [b0, b0+P), else 0 (every i < kv_off.numel()) + + Every other element of ``ids`` / ``draft`` is left as it was. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # draft_len 3: T = 4, D = 3; the store's [T, slots, 1] new tokens and [slots, D] next drafts (8 slots, the + # max batch); the engine's flat buffers sized for the max batch. + CELLS: tuple[Cell, ...] = ( + Cell( + why="an unchanged batch of 8 previous requests in shuffled slots, nothing before them (b0 = 0)", + spec=dict(p=8, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + Cell( + why="2 previous requests after 3 new ones: offsets ids0 / d0 / b0 nonzero, zero-filled offsets outside", + spec=dict(p=2, t=4, d=3, slots=8, b0=3, ids0=12, d0=9, extra_rows=3), + ), + Cell( + why="one request, the store's 2-D view of new tokens (beam dim squeezed)", + spec=dict(p=1, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0, two_d=True), + ), + ) + + note = """ + Pure index copies: the reference is exact. The op does not bound slots, ppi or the destination offsets; is_valid + states them. + """ + + def __call__( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + torch.ops.trtllm.minimax_m3_step_gather( + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ) + + def reference( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + """The four destinations after the call, on copies (``ids`` / ``draft`` keep their other elements).""" + nt = new_tok.reshape(new_tok.shape[0], new_tok.shape[1]) + s = slots[:P].long() + ids_r, draft_r = ids.clone(), draft.clone() + ids_r[ids0 : ids0 + P * T] = nt[:T, s].t().reshape(-1) + draft_r[d0 : d0 + P * D] = next_draft[s, :D].reshape(-1) + pos_r = torch.zeros_like(pos_off) + pos_r[b0 * T : (b0 + P) * T] = lens[ppi[: P * T].long()] + kv_r = torch.zeros_like(kv_off) + kv_r[b0 : b0 + P] = lens[s] - T + return {"ids": ids_r, "draft": draft_r, "pos_off": pos_r, "kv_off": kv_r} + + def is_valid( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + # Reads through slots / ppi and writes at the offsets are unchecked: any of them out of range reads or + # writes another buffer. + s = slots[:P] + assert bool( + ( + (s >= 0) & (s < next_draft.shape[0]) & (s < new_tok.shape[1]) & (s < lens.numel()) + ).all() + ), "slots must index the store" + assert bool(((ppi[: P * T] >= 0) & (ppi[: P * T] < lens.numel())).all()), ( + "ppi must index lens" + ) + assert new_tok.shape[0] >= T and next_draft.shape[1] >= D, "store narrower than T / D" + assert ids.numel() >= ids0 + P * T and draft.numel() >= d0 + P * D, "destinations too short" + assert pos_off.numel() >= (b0 + P) * T and kv_off.numel() >= b0 + P, ( + "offset buffers too short" + ) + + +m3_step_gather = M3StepGather() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/triton_sparse_decode.py new file mode 100644 index 000000000000..064a3c0b2b10 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/triton_sparse_decode.py @@ -0,0 +1,1651 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/common/ops/sparse_attn.py +"""Triton block-sparse GQA decode attention for MiniMax-M3. + +Flash-decoding over the blocks the indexer selected: one CTA per (query token, +top-k chunk, KV head) accumulates a partial output plus its log-sum-exp, and a +second kernel merges the chunks by LSE weight. That beats running the +context-schedule FMHA kernel at decode, where a single query token leaves most +of a 128-row Q tile idle. + +GB300/SM103 tuning targets GQA group 16, top-k 16 and four-token Eagle3 +verification. Split-K, query pairing and launch tables require four KV heads; +the linear-softmax policy also covers one and two KV heads. Other hardware +and shapes retain the default policies, including supported SM100 NVFP4. + +Vendored from the vLLM source linked in the file header (v0.26.1rc0-77-g6f91edf96). +Differences from upstream: + +* K and V are separate HND paged views ([num_pages, num_kv_heads, page_size, + head_dim]) with independent strides rather than one cache fused along the + last dim, because that is how the M3 KV pool is laid out here. +* Only the scalar KV-scale mode is kept; M3 stores unscaled E4M3 K/V. +* A packed NVFP4 cache is read in place, dequantized in-register from its E2M1 + nibbles and E4M3 block scales. +* Partial-output and LSE scratch come from the persistent buffer arena so their + addresses survive CUDA graph replay. +* A single top-k chunk writes the final output directly, avoiding scratch and + the otherwise redundant LSE merge. +* The merge kernel zeroes rows whose chunks are all empty instead of letting + them produce NaN, since CUDA-graph padding rows flow on into the rest of the + network here. + +This is the MiniMax-M3 NVFP4 / sm_103 / tp4 target's copy. On top of the +in-tree decode it routes long-context TP4 steps to the one-launch +``m3_attn_lc`` kernel, short-sequence verify steps of at least +``MQ_MIN_TOKENS`` rows to the multi-query producer, releases PDL dependents +early, and has the merge emit the MXFP8 o_proj hand-off (``m3_prequant.SLOT``). +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + check_decode_span_shape, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers + +from . import m3_attn_lc, m3_prequant + +# One sparse block is exactly one KV page. +SPARSE_BLOCK_SIZE = 128 # capacity: fixed: the MSA block size of the checkpoint + +_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) + +_TL_DTYPES = { + torch.bfloat16: tl.bfloat16, + torch.float16: tl.float16, + torch.float32: tl.float32, +} + +# Total CTAs the split-K partitioning aims for before it stops splitting. +_TARGET_GRID = 256 # capacity: fixed: a launch-grid target, not a size + +# NVFP4 quantizes in groups of 16 cache elements along head_dim, one E4M3 scale +# byte each; see the fp4_quantize call in msa_scatter.fused_write_layer_caches_nvfp4. +NVFP4_SF_VEC_SIZE = 16 + + +# Fewest query rows for which the merge writes the MXFP8 o_proj hand-off: every decode step, so the MXFP8 o_proj +# never runs its own quantize kernel after attention. +MERGE_MXFP8_MIN_TOKENS = 1 # capacity: fixed: a row lower bound + + +class MergeMxfp8: + """MXFP8 copy of the split-K merge output for the o_proj GEMM. + + From ``MERGE_MXFP8_MIN_TOKENS`` to ``MAX_ROWS`` rows the merge kernel writes the E4M3 rows and swizzled UE8M0 + scales next to the BF16 output and publishes them to ``m3_prequant.SLOT``, so the MXFP8 o_proj skips its own + quantization. One 128-row buffer per width serves every layer: each o_proj consumes it before the next layer's + merge rewrites it, in stream order. + """ + + MAX_ROWS = 128 # capacity: fixed: the merge hand-off's rows (decode steps) + + def __init__(self): + self._buffers = {} + + def buffers(self, width: int, device: torch.device): + key = (width, device) + if key not in self._buffers: + self._buffers[key] = ( + torch.empty((self.MAX_ROWS, width), dtype=torch.float8_e4m3fn, device=device), + torch.zeros((self.MAX_ROWS * (width // 32),), dtype=torch.uint8, device=device), + ) + return self._buffers[key] + + +MERGE_MXFP8 = MergeMxfp8() + + +@triton.jit +def _e2m1x2_to_f16x2(packed): + """Convert one packed E2M1 byte into its low/high FP16 values.""" + return tl.inline_asm_elementwise( + asm="""{ + .reg .b8 src; + .reg .b32 pair; + mov.b32 {src, _, _, _}, $2; + cvt.rn.f16x2.e2m1x2 pair, src; + mov.b32 {$0, $1}, pair; + }""", + constraints="=h,=h,r", + args=(packed,), + dtype=(tl.float16, tl.float16), + is_pure=True, + pack=1, + ) + + +@triton.jit +def _dequant_nvfp4_rows( + packed, # [BLOCK_N, D // 2] uint8, two E2M1 nibbles per byte + scale_bytes, # [BLOCK_N, SCALE_COLS] uint8, one E4M3 scale per 16 elements + BLOCK_N: tl.constexpr, + D: tl.constexpr, + SCALE_COLS: tl.constexpr, + SF_VEC: tl.constexpr, +): + """Dequantize one packed NVFP4 [BLOCK_N, D] tile to fp16. + + The caller gathers the block scales in whichever layout its role stores + them, so this handles only the value math and the scale broadcast. The + per-tensor scale is applied by the caller. + """ + # Blackwell converts both E2M1 nibbles with one instruction. Joining the + # low/high results along a new innermost axis and flattening it rebuilds + # the original element order without scalar exponent/sign reconstruction. + value_lo, value_hi = _e2m1x2_to_f16x2(packed) + values = tl.reshape(tl.join(value_lo, value_hi), (BLOCK_N, D)) + + scales = scale_bytes.to(tl.float8e4nv, bitcast=True).to(tl.float16) + scales = tl.reshape( + tl.broadcast_to(scales[:, :, None], (BLOCK_N, SCALE_COLS, SF_VEC)), (BLOCK_N, D) + ) + return values * scales + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max(16, triton.next_power_of_2(args["gqa_group_size"])), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +# Mixed ctx+gen steps launch on [num_ctx:] slices (unaligned seq_lens / block_table) with arbitrary +# total_q; specializing on those compiled a new variant mid-serving (~650 ms TTFT stall at c8). +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + out_ptr, # final [total_q, num_heads, head_dim] for FUSED_MERGE, else unused + counter_ptr, # int32 [total_q, num_kv_heads] arrival counters for FUSED_MERGE, zero at rest + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + stride_out_n, + stride_out_h, + stride_out_d, + num_kv_heads, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + FUSED_MERGE: tl.constexpr, # the last chunk to finish merges the partials in place + EARLY_TRIGGER: tl.constexpr, # release PDL dependents (weight prefetch) at entry +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query token, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + pid_b = pid_bc % total_q + pid_c = pid_bc // total_q + req_id = pid_b // decode_query_len + q_offset = pid_b - req_id * decode_query_len + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + if EARLY_TRIGGER: + # Dependents still wait for this grid to complete; triggering here only + # lets them start their input-independent prologue (e.g. weight loads). + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + req_id) + query_pos = seq_len - decode_query_len + q_offset + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + + # Bound the walk by the token's own valid block count rather than a + # sentinel: the selector emits ascending ids and only pads with -1 past + # min(topk, cdiv(kv_len, blk)) entries, so those are never dereferenced. + num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk = tl.minimum(max_topk, num_blocks) + chunk_end = tl.minimum(chunk_start + chunk_size_topk, real_topk) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + off_h = tl.arange(0, BLOCK_SIZE_H) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + h_mask = off_h < gqa_group_size + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + pid_b * stride_qn + + (pid_h + off_h[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so the result + # matches a caller that widened q itself before the call. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + cur_idx_ptr = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + for it in tl.range(chunk_start, chunk_end): + blk = tl.load(cur_idx_ptr) + cur_idx_ptr = cur_idx_ptr + stride_tk + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + page = tl.load(bt_row + blk).to(tl.int64) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len + if KV_NVFP4: + # K block scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + qk = tl.where(pos_mask[None, :], 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.exp2(qk - m_ij[:, None]) + l_ij = tl.sum(p, axis=1) + if USE_LINEAR_SOFTMAX: + alpha = tl.exp2(m_i - m_ij) + acc_o = acc_o * alpha[:, None] + else: + acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] + if KV_NVFP4: + # V block scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = l_i * alpha + l_ij + else: + lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + o_base = o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h + tl.store( + o_base + off_h[:, None] * stride_o_h + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + lse_base = lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h + tl.store(lse_base + off_h * stride_l_h, lse_out, mask=h_mask) + + if FUSED_MERGE: + # Split-K serial reduction: every chunk publishes its partial, and the + # chunk that arrives last for this (token, KV head) merges all of them + # exactly as _merge_topk_attn_out_kernel does, then re-arms the counter. + tl.debug_barrier() + counter = counter_ptr + pid_b * num_kv_heads + pid_kh + arrived = tl.atomic_add(counter, 1, sem="acq_rel", scope="gpu") + if arrived == NUM_TOPK_CHUNKS - 1: + tl.debug_barrier() + o_rows = ( + o_ptr + + pid_b * stride_o_b + + (pid_h + off_h[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d + ) + l_rows = lse_ptr + pid_b * stride_l_b + (pid_h + off_h) * stride_l_h + lse_max = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + lse_max = tl.maximum(lse_max, lse_c) + # All chunks empty (a CUDA-graph padding row): zero, never NaN. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + denom = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + denom += tl.exp2(lse_c - lse_max) + denom = tl.where(denom > 0, denom, 1.0) + merged = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + w_c = tl.exp2(lse_c - lse_max) / denom + o_c = tl.load( + o_rows + c * stride_o_c, mask=hd_mask, other=0.0, cache_modifier=".cg" + ).to(tl.float32) + merged += o_c * w_c[:, None] + tl.store( + out_ptr + + pid_b * stride_out_n + + (pid_h + off_h[:, None]) * stride_out_h + + off_d[None, :] * stride_out_d, + merged.to(out_ptr.dtype.element_ty), + mask=hd_mask, + ) + tl.atomic_xchg(counter, 0, sem="relaxed", scope="gpu") + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max( + 16, + triton.next_power_of_2(args["gqa_group_size"] * args["QUERY_GROUP_SIZE"]), + ), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_pair_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + QUERY_GROUP_SIZE: tl.constexpr, # consecutive Eagle queries sharing one block set + MAX_TOPK: tl.constexpr, + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query group, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + num_query_groups = total_q // QUERY_GROUP_SIZE + pid_g = pid_bc % num_query_groups + pid_c = pid_bc // num_query_groups + groups_per_req = decode_query_len // QUERY_GROUP_SIZE + req_id = pid_g // groups_per_req + q_group_offset = (pid_g - req_id * groups_per_req) * QUERY_GROUP_SIZE + pid_b = req_id * decode_query_len + q_group_offset + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + + seq_len = tl.load(seq_lens + req_id) + # Query rows are flattened as [query-in-group, GQA head]. + off_h = tl.arange(0, BLOCK_SIZE_H) + q_lane = off_h // gqa_group_size + h_lane = off_h - q_lane * gqa_group_size + h_mask = off_h < gqa_group_size * QUERY_GROUP_SIZE + q_token = pid_b + q_lane + query_pos = seq_len - decode_query_len + q_group_offset + q_lane + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + kv_len_max = tl.max(tl.where(h_mask, kv_len, 0), axis=0) + + # Each query can have a different valid-block count at a page boundary. + # The selector puts those valid entries first and pads only the tail. + kv_len_0 = tl.maximum(seq_len - decode_query_len + q_group_offset + 1, 0) + num_blocks_0 = (kv_len_0 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_0 = tl.minimum(max_topk, num_blocks_0) + chunk_end_0 = tl.minimum(chunk_start + chunk_size_topk, real_topk_0) + if QUERY_GROUP_SIZE == 2: + kv_len_1 = tl.maximum(seq_len - decode_query_len + q_group_offset + 2, 0) + num_blocks_1 = (kv_len_1 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_1 = tl.minimum(max_topk, num_blocks_1) + chunk_end_1 = tl.minimum(chunk_start + chunk_size_topk, real_topk_1) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + q_token[:, None] * stride_qn + + (pid_h + h_lane[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so this + # reproduces the values the caller's standalone widening kernel produced. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + idx_ptr_0 = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + if QUERY_GROUP_SIZE == 2: + idx_ptr_1 = idx_ptr_0 + stride_tn + idx_0 = chunk_start + idx_1 = chunk_start + count_0 = chunk_end_0 - chunk_start + count_1 = chunk_end_1 - chunk_start + chunk_capacity: tl.constexpr = (MAX_TOPK + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_offsets = tl.arange(0, chunk_capacity) + sentinel = 2147483647 + blocks_0 = tl.load( + idx_ptr_0 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_0, + other=sentinel, + ) + blocks_1 = tl.load( + idx_ptr_1 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_1, + other=sentinel, + ) + chunk_equal = tl.sum(blocks_0 != blocks_1, axis=0) == 0 + # Equal chunks take the measured shared-page loop. A divergent chunk + # walks the sorted union, bounded by the sum of both row lengths. + num_candidates = tl.where(chunk_equal, count_0, count_0 + count_1) + else: + num_candidates = chunk_end_0 - chunk_start + + for _ in tl.range(0, num_candidates): + if QUERY_GROUP_SIZE == 2: + if chunk_equal: + has_block = True + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + else: + valid_0 = idx_0 < chunk_end_0 + valid_1 = idx_1 < chunk_end_1 + blk_0 = tl.load(idx_ptr_0, mask=valid_0, other=sentinel) + blk_1 = tl.load(idx_ptr_1, mask=valid_1, other=sentinel) + use_0 = valid_0 & (blk_0 <= blk_1) + use_1 = valid_1 & (blk_1 <= blk_0) + has_block = use_0 | use_1 + blk = tl.minimum(blk_0, blk_1) + idx_0 += use_0.to(tl.int32) + idx_1 += use_1.to(tl.int32) + idx_ptr_0 += use_0.to(tl.int64) * stride_tk + idx_ptr_1 += use_1.to(tl.int64) * stride_tk + query_active = h_mask & (((q_lane == 0) & use_0) | ((q_lane == 1) & use_1)) + else: + has_block = True + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + + if has_block: + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + page = tl.load(bt_row + blk).to(tl.int64) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len_max + q_pos_mask = blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None] + if KV_NVFP4: + # K scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + active_pos = query_active[:, None] & q_pos_mask + qk = tl.where(active_pos, 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.where(query_active, tl.maximum(m_i, tl.max(qk, axis=1)), m_i) + p = tl.where(active_pos, tl.exp2(qk - m_ij[:, None]), 0.0) + l_ij = tl.sum(p, axis=1) + alpha = tl.where(query_active, tl.exp2(m_i - m_ij), 1.0) + acc_o = acc_o * alpha[:, None] + if KV_NVFP4: + # V scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = tl.where(query_active, l_i * alpha + l_ij, l_i) + else: + lse_ij = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + lse_i = tl.where(query_active, lse_ij, lse_i) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + tl.store( + o_ptr + + pid_c * stride_o_c + + q_token[:, None] * stride_o_b + + (pid_h + h_lane[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + tl.store( + lse_ptr + pid_c * stride_l_c + q_token * stride_l_b + (pid_h + h_lane) * stride_l_h, + lse_out, + mask=h_mask, + ) + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _rcp_approx_ftz(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 _store_mxfp8_head(x, q8_ptr, sf_ptr, row, head, num_sf_k_tiles, HEAD_DIM: tl.constexpr): + """Quantize one row's [head * HEAD_DIM, (head + 1) * HEAD_DIM) slice to MXFP8 as + torch.ops.trtllm.mxfp8_quantize(x, True) does: UE8M0 scale = round-up(amax * rcp(448)), E4M3 RN-satfinite, + scales in the 128x4-swizzled layout.""" + NUM_SF: tl.constexpr = HEAD_DIM // 32 + xr = tl.reshape(x.to(tl.float32), (NUM_SF, 32)) + amax = tl.max(tl.abs(xr), axis=1) + sfv = amax * _rcp_approx_ftz(tl.full((NUM_SF,), 448.0, tl.float32)) + bits = sfv.to(tl.int32, bitcast=True) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32) + e = tl.minimum(e, 254) + scale = ((254 - e) << 23).to(tl.float32, bitcast=True) + scale = tl.where(amax != 0, scale, 0.0) + q = (xr * scale[:, None]).to(tl.float8e4nv) + off_q = tl.reshape(tl.arange(0, HEAD_DIM), (NUM_SF, 32)) + tl.store(q8_ptr + row * (num_sf_k_tiles * 128) + head * HEAD_DIM + off_q, q) + k_sf = head * NUM_SF + tl.arange(0, NUM_SF) + off_sf = ( + (row // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + + k_sf % 4 + ) + tl.store(sf_ptr + off_sf, e.to(tl.uint8)) + + +@triton.heuristics({"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}) +@triton.jit +def _merge_topk_attn_out_kernel( + o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim] + lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads] + out_ptr, # merged out: [total_q, num_heads, head_dim] + head_dim, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_out_n, + stride_out_h, + stride_out_d, + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr = False, + q8_ptr=None, # MXFP8 copy of out: [total_q, num_heads * head_dim] E4M3 + sf_ptr=None, # its UE8M0 block scales, 128x4-swizzled + num_sf_k_tiles=0, # ceil(num_heads * head_dim / 128) + MXFP8_OUT: tl.constexpr = False, +): + pid_b, pid_h = tl.program_id(0), tl.program_id(1) + + if USE_PDL: + if EARLY_TRIGGER: + # Release the dependent (the o_proj GEMM) before waiting on the attention producer, so its weight + # prefetch overlaps the attention kernel as well as this merge; its griddepcontrol.wait still waits + # for this grid (and therefore the producer) to complete. + tl.extra.cuda.gdc_launch_dependents() + tl.extra.cuda.gdc_wait() + + off_c = tl.arange(0, NUM_TOPK_CHUNKS) + off_d = tl.arange(0, BLOCK_SIZE_D) + o = tl.load( + o_ptr + + pid_b * stride_o_b + + pid_h * stride_o_h + + off_c[:, None] * stride_o_c + + off_d[None, :] * stride_o_d, + mask=off_d[None, :] < head_dim, + other=0.0, + ).to(tl.float32) + # Empty chunks contribute -inf, hence weight 0. + lse = tl.load(lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c) + lse_max = tl.max(lse, axis=0) + # A row whose every chunk is empty (a CUDA-graph padding row) would give + # -inf - -inf here. Zero it instead: the padded output is discarded, but a + # NaN would survive into the residual stream and the all-reduce. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + weights = tl.exp2(lse - lse_max) + denom = tl.sum(weights, axis=0) + weights = weights / tl.where(denom > 0, denom, 1.0) + o_merged = tl.sum(o * weights[:, None], axis=0) + out_ptrs = out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d + tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim) + if MXFP8_OUT: + _store_mxfp8_head( + o_merged.to(out_ptr.dtype.element_ty), + q8_ptr, + sf_ptr, + pid_b, + pid_h, + num_sf_k_tiles, + BLOCK_SIZE_D, + ) + + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> int: + """Split-K factor over the top-k blocks, as a power of two. + + Depends only on shapes that are fixed for a captured batch size, so the + launch geometry is frozen inside a CUDA graph. + """ + target = max(1, min(max_topk, _TARGET_GRID // max(1, total_q * num_kv_heads))) + return 1 << (target.bit_length() - 1) + + +# NVFP4 instructions are supported on both datacenter Blackwell targets. +_SM100F_CAPABILITIES = ((10, 0), (10, 3)) +# Shape-specific tuning was measured on GB300 only. SM100 uses the defaults. +_NVFP4_TUNED_CAPABILITIES = ((10, 3),) + + +def _sm100f_nvfp4_num_topk_chunks( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> Optional[int]: + """Select GB300-measured split-K factors for M3's SM100-family decode shape.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return None + local_batch = total_q // decode_query_len + return { + 2: 8, + 3: 8, + 4: 8, + 5: 16, + 6: 4, + 7: 4, + 8: 2, + 9: 2, + 10: 8, + 11: 4, + 12: 2, + 13: 2, + 14: 2, + }.get(local_batch) + + +def _sm100f_nvfp4_use_linear_softmax( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> bool: + """Use the GB300-measured lower-register recurrence on validated M3 shapes.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads not in (1, 2, 4) + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return False + local_batch = total_q // decode_query_len + # The TP4/DEP32 single-request graph is 8.9% faster with the original + # log-domain recurrence; every other measured DEP8/16/32 bucket is neutral + # or faster with the linear denominator. + return not (num_kv_heads == 1 and local_batch == 1) + + +def _sm100f_nvfp4_query_group_size( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> int: + """Group Eagle queries using the GB300-measured shared-page policy.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return 1 + # The sorted-union kernel is 6.5%, 1.6%, and 24.2% faster at local batches + # 11, 12, and 14 with split-K 8. Other measured buckets are neutral or + # slower and retain the single-query producer. + return 2 if total_q // decode_query_len in (11, 12, 14) else 1 + + +def _sm100f_nvfp4_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select native-NVFP4 launch geometry for the SM100 family. + + The FP16 QK/PV path is register-sensitive. GB300 dense replay shows four + warps/two stages wins through local batch 28, while two warps/one stage + avoids the next residency cliff at local batches 29--32. Leave every + unmeasured architecture, model geometry, and larger batch on Triton's + established defaults. + """ + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + ): + return {} + local_batch = total_q // decode_query_len + if local_batch == 1 or local_batch in (2, 3, 5, 6, 7, 10, 11, 12): + return {"num_warps": 4, "num_stages": 1} + if local_batch == 4: + return {"num_warps": 2, "num_stages": 1} + if 8 <= local_batch <= 28: + return {"num_warps": 4, "num_stages": 2} + if 29 <= local_batch <= 32: + return {"num_warps": 2, "num_stages": 1} + return {} + + +def _sm100f_nvfp4_merge_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select GB300-measured one-warp merges for M3's SM100-family graphs.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability in _NVFP4_TUNED_CAPABILITIES + and num_kv_heads == 4 + and gqa_group_size == 16 + and max_topk == 16 + and decode_query_len == 4 + and total_q % decode_query_len == 0 + and total_q // decode_query_len in (2, 4, 5, 6, 7, 10, 11) + ): + return {"num_warps": 1} + return {} + + +def _check_nvfp4_inputs( + k_paged: torch.Tensor, + v_paged: torch.Tensor, + nvfp4_args: tuple[Optional[torch.Tensor], ...], + *, + head_dim: int, + scale_cols: int, +) -> None: + """Validate the packed-NVFP4 cache against what the kernel assumes. + + The kernel indexes block scales by a flat within-page offset and reshapes a + [., SCALE_COLS] scale tile up to head_dim, so it needs contiguous scale rows + and head_dim exactly BLOCK_SIZE_D. Both hold for MiniMax-M3; check them + rather than let a mismatch read neighboring pages. + """ + k_block_scale, v_block_scale, k_global_scale, v_global_scale = nvfp4_args + if any(arg is None for arg in nvfp4_args): + raise ValueError( + "MiniMax-M3 NVFP4 sparse decode needs all four of k_block_scale, " + "v_block_scale, k_global_scale and v_global_scale." + ) + if head_dim != triton.next_power_of_2(head_dim): + raise ValueError(f"NVFP4 sparse decode requires a power-of-two head_dim; got {head_dim}.") + for name, data in (("k_paged", k_paged), ("v_paged", v_paged)): + if int(data.shape[3]) * 2 != head_dim: + raise ValueError( + f"NVFP4 {name} must pack two elements per byte: expected " + f"head_dim/2 ({head_dim // 2}) columns, got {int(data.shape[3])}." + ) + for name, scale in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + if tuple(scale.shape) != tuple(k_paged.shape[:3]) + (scale_cols,): + raise ValueError( + f"NVFP4 {name} must be [num_pages, num_kv_heads, page_size, " + f"{scale_cols}] to match the cache; got {tuple(scale.shape)}." + ) + if scale.stride(3) != 1 or scale.stride(2) != scale_cols: + raise ValueError( + f"NVFP4 {name} needs contiguous per-page scale bytes; got strides " + f"{tuple(scale.stride())}." + ) + for name, gs in (("k_global_scale", k_global_scale), ("v_global_scale", v_global_scale)): + if gs.dtype != torch.float32 or gs.numel() != 1: + raise ValueError( + f"NVFP4 {name} must be a single-element FP32 tensor; got " + f"{gs.numel()} element(s) of {gs.dtype}." + ) + + +# Multi-query producer (triton_sparse_decode_mq) for identity-selection verify steps with at least this many query +# tokens; below it the per-token producer is as fast. +MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query producer's crossover +# Warps per multi-query producer program. +MQ_NUM_WARPS = 8 + + +def _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, +) -> None: + """minimax_m3_sparse_attn_decode through m3_attn_lc: same table, output and MXFP8 hand-off contract.""" + use_pdl = get_env_enable_pdl() + total_q, num_heads, head_dim = q.shape + linear = _sm100f_nvfp4_use_linear_softmax( + total_q=total_q, + num_kv_heads=1, + gqa_group_size=num_heads, + max_topk=int(topk_idx.shape[-1]), + decode_query_len=decode_query_len, + ) + mxfp8 = None + m3_prequant.SLOT.clear() + if MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS: + mxfp8 = MERGE_MXFP8.buffers(num_heads * head_dim, output.device) + # The early trigger lets the o_proj launch and stream its weights while attention runs. + m3_attn_lc.m3_lc_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + linear_softmax=linear, + mxfp8=mxfp8, + use_pdl=use_pdl, + early_trigger=True, + ) + if mxfp8 is not None: + m3_prequant.SLOT.publish( + output.view(total_q, num_heads * head_dim), mxfp8[0][:total_q], mxfp8[1] + ) + + +@torch.no_grad() +def minimax_m3_sparse_attn_decode( + q: torch.Tensor, # [total_q, num_heads, head_dim] + k_paged: torch.Tensor, # [num_pages, num_kv_heads, page_size, head_dim] + v_paged: torch.Tensor, # same layout as k_paged + topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk] + block_table: torch.Tensor, # [num_reqs, max_blocks] + seq_lens: torch.Tensor, # [num_reqs] int32 + *, + sm_scale: float, + output: torch.Tensor, # [total_q, num_heads, head_dim] + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, +) -> None: + """Block-sparse GQA decode attention, written into output in place. + + q may be FP8, as a fused QK-norm/RoPE producer emits it. The kernel widens + it to output.dtype in-register rather than making the caller materialize a + widened copy, which would cost a standalone kernel on every sparse layer. + The attention math is unchanged either way, since E4M3 widens exactly. + + kv_scale is an optional scalar dequantization factor for an FP8 cache; + MiniMax-M3 stores unscaled E4M3, so it is normally None. + + Passing the four NVFP4 scale tensors selects the packed-NVFP4 cache instead: + k_paged/v_paged then hold two E2M1 elements per byte, the block scales are + paged E4M3 bytes in the layouts msa_scatter writes, and the two per-tensor + dequant scales are single-element FP32 device tensors. + + num_topk_chunks overrides the split-K factor and exists for tests: the + merged result must not depend on it. + """ + total_q, num_heads, head_dim = q.shape + num_kv_heads = int(k_paged.shape[1]) + check_decode_span_shape( + "MiniMax-M3 Triton sparse decode", total_q, int(seq_lens.shape[0]), decode_query_len + ) + if int(k_paged.shape[2]) != SPARSE_BLOCK_SIZE: + raise ValueError( + f"MiniMax-M3 sparse decode requires page_size={SPARSE_BLOCK_SIZE}; " + f"got {int(k_paged.shape[2])}." + ) + if num_heads % num_kv_heads: + # The kernel covers heads pid_kh * gqa_group_size + [0, group), so a + # head count that is not a whole number of groups leaves the tail heads + # unwritten and the merge kernel reads uninitialized partials. + raise ValueError( + f"MiniMax-M3 sparse decode requires num_heads ({num_heads}) to be a " + f"multiple of num_kv_heads ({num_kv_heads})." + ) + nvfp4_args = (k_block_scale, v_block_scale, k_global_scale, v_global_scale) + kv_nvfp4 = any(arg is not None for arg in nvfp4_args) + scale_cols = head_dim // NVFP4_SF_VEC_SIZE + if kv_nvfp4: + capability = torch.cuda.get_device_capability(q.device) + if capability not in _SM100F_CAPABILITIES: + raise NotImplementedError( + "MiniMax-M3 NVFP4 sparse decode requires SM100/SM103; " + f"got SM{capability[0] * 10 + capability[1]}." + ) + _check_nvfp4_inputs(k_paged, v_paged, nvfp4_args, head_dim=head_dim, scale_cols=scale_cols) + # Two E2M1 elements per byte; Triton must see them unsigned so that the + # high-nibble shift does not sign-extend. + k_paged, v_paged = k_paged.view(torch.uint8), v_paged.view(torch.uint8) + k_block_scale, v_block_scale = ( + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + ) + max_topk = int(topk_idx.shape[-1]) + gqa_group_size = num_heads // num_kv_heads + use_scale = k_paged.dtype in _FP8_DTYPES and kv_scale is not None + # Triton needs a real pointer even for the unused arguments. + scale_arg = kv_scale if use_scale else output + k_sc_arg, v_sc_arg = (k_block_scale, v_block_scale) if kv_nvfp4 else (output, output) + k_gs_arg, v_gs_arg = (k_global_scale, v_global_scale) if kv_nvfp4 else (output, output) + + widen_q = q.dtype in _FP8_DTYPES + if widen_q and output.dtype not in _TL_DTYPES: + raise ValueError( + f"MiniMax-M3 sparse decode cannot widen FP8 q into {output.dtype}; " + f"supported compute dtypes are {sorted(d.__str__() for d in _TL_DTYPES)}." + ) + # Packed NVFP4 dequantizes to FP16. The production producer supplies FP8 + # q, which widens exactly to FP16, so keep both dot operands there instead + # of converting every K/V tile to BF16. Accumulation remains FP32. + if widen_q: + compute_dtype = tl.float16 if kv_nvfp4 else _TL_DTYPES[output.dtype] + else: + compute_dtype = tl.float32 + + # Long-context TP4 steps: attention, split-K merge and MXFP8 hand-off in one launch. + if ( + kv_nvfp4 + and widen_q + and not identity_blocks + and num_topk_chunks is None + and m3_attn_lc.lc_applies( + q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output + ) + ): + _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + ) + return + + query_group_size = ( + _sm100f_nvfp4_query_group_size( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else 1 + ) + if num_topk_chunks is None: + if query_group_size == 2: + measured_chunks = 8 + else: + measured_chunks = ( + _sm100f_nvfp4_num_topk_chunks( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else None + ) + num_topk_chunks = ( + measured_chunks + if measured_chunks is not None + else resolve_num_topk_chunks(total_q, num_kv_heads, max_topk) + ) + if kv_nvfp4 and total_q * num_kv_heads <= 64: + # GB300 TP4 Eagle verify (1 KV head, group 16, 6 query tokens): one block per program is fastest up to + # batch 8 (lowm5/bench_decode_chunks.py: b8 11.7 vs 14.8 us with the default split of 4). + num_topk_chunks = max_topk + elif num_topk_chunks & (num_topk_chunks - 1): + raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") + + # Every query token of a request reads the same blocks under identity selection (short-sequence mode): one + # program per (request, block) covers all of them and dequantizes each block once. + use_mq = ( + identity_blocks + and total_q >= MQ_MIN_TOKENS + and kv_nvfp4 + and num_kv_heads == 1 + and decode_query_len > 1 + and head_dim == 128 + and gqa_group_size * decode_query_len <= 128 + ) + if use_mq: + num_topk_chunks = max_topk + direct_output = num_topk_chunks == 1 + if direct_output: + # With one chunk, acc_o is already normalized and final. Point the + # producer at the caller's output and compile away the LSE store. + # Keeping a real dummy pointer preserves Triton's pointer contract. + o_arg = output + lse_arg = output + o_strides = (0, output.stride(0), output.stride(1), output.stride(2)) + lse_strides = (0, 0, 0) + else: + # Persistent arena rather than torch.empty, so the partials keep one + # address across CUDA graph replays. FP32 keeps the split-K factor from + # perturbing the merged result. + reserve = torch.cuda.is_current_stream_capturing() + o_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads, head_dim], + torch.float32, + buffer_name="m3_sparse_decode_o_partial", + reserve_buffer=reserve, + ) + lse_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads], + torch.float32, + buffer_name="m3_sparse_decode_lse_partial", + reserve_buffer=reserve, + ) + o_strides = tuple(o_arg.stride()) + lse_strides = tuple(lse_arg.stride()) + + use_pdl = get_env_enable_pdl() + pdl_launch = {"launch_pdl": True} if use_pdl else {} + use_linear_softmax = kv_nvfp4 and _sm100f_nvfp4_use_linear_softmax( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if query_group_size == 2: + launch_options = {"num_warps": 4, "num_stages": 1} + else: + launch_options = ( + _sm100f_nvfp4_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + + # The pair path merges the two sorted selector rows on device. Equal blocks + # share one K/V dequantization; divergent blocks update only their owning + # query, so the optimization does not rely on runtime selection equality. + if total_q % query_group_size or decode_query_len % query_group_size: + raise ValueError( + f"query group {query_group_size} must divide total_q={total_q} and " + f"decode_query_len={decode_query_len}." + ) + + producer_kernel = ( + _gqa_sparse_decode_pair_kernel if query_group_size == 2 else _gqa_sparse_decode_kernel + ) + single_kernel_args = ( + { + # No split-K arrival counters: the separate merge kernel combines the chunks. + "merge": (output, None), + "tail": (output.stride(0), output.stride(1), output.stride(2), num_kv_heads), + # An early release pays off when the dependent is the consumer of the + # final output (it can prefetch), or a merge grid that itself releases + # its dependent at entry: the o_proj then launches during this kernel and + # streams its weights while attention runs. Above 8 query rows the + # concurrent weight stream slows the larger attention grid more than the + # prefetch saves (GB300 TP4: -1.1% at 48 rows). + "options": {"FUSED_MERGE": False, "EARLY_TRIGGER": direct_output or total_q <= 8}, + } + if query_group_size == 1 + else {"merge": (), "tail": (), "options": {}} + ) + pair_kernel_options = ( + {"QUERY_GROUP_SIZE": query_group_size, "MAX_TOPK": max_topk} + if query_group_size == 2 + else {} + ) + if use_mq: + from .triton_sparse_decode_mq import launch_mq_producer + + launch_mq_producer( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + block_table, + seq_lens, + o_arg, + lse_arg, + sm_scale=sm_scale, + decode_query_len=decode_query_len, + use_pdl=use_pdl, + early_trigger=False, + num_warps=MQ_NUM_WARPS, + ) + else: + producer_kernel[ + (total_q // query_group_size) * num_topk_chunks, + num_kv_heads, + ]( + q, + k_paged, + v_paged, + scale_arg, + k_sc_arg, + v_sc_arg, + k_gs_arg, + v_gs_arg, + topk_idx, + o_arg, + lse_arg, + *single_kernel_args["merge"], + block_table, + seq_lens, + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + q.stride(0), + q.stride(1), + q.stride(2), + k_paged.stride(0), + k_paged.stride(1), + k_paged.stride(2), + k_paged.stride(3), + v_paged.stride(0), + v_paged.stride(1), + v_paged.stride(2), + v_paged.stride(3), + k_sc_arg.stride(0), + k_sc_arg.stride(1), + v_sc_arg.stride(0), + v_sc_arg.stride(1), + topk_idx.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + *o_strides, + *lse_strides, + block_table.stride(0), + *single_kernel_args["tail"], + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_SCALE=use_scale, + WIDEN_Q=widen_q, + COMPUTE_DTYPE=compute_dtype, + USE_PDL=use_pdl, + KV_NVFP4=kv_nvfp4, + USE_LINEAR_SOFTMAX=use_linear_softmax, + SCALE_COLS=scale_cols, + DIRECT_OUTPUT=direct_output, + **single_kernel_args["options"], + **pair_kernel_options, + **pdl_launch, + **launch_options, + ) + m3_prequant.SLOT.clear() + if not direct_output: + width = num_heads * head_dim + mxfp8 = ( + MERGE_MXFP8.buffers(width, output.device) + if MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS + and head_dim == 128 + and output.dim() == 3 + and output.is_contiguous() + and output.dtype == torch.bfloat16 + else None + ) + mxfp8_args = ( + { + "q8_ptr": mxfp8[0], + "sf_ptr": mxfp8[1], + "num_sf_k_tiles": width // 128, + "MXFP8_OUT": True, + } + if mxfp8 is not None + else {} + ) + merge_launch_options = ( + _sm100f_nvfp4_merge_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + _merge_topk_attn_out_kernel[(total_q, num_heads)]( + o_arg, + lse_arg, + output, + head_dim, + *o_strides, + *lse_strides, + output.stride(0), + output.stride(1), + output.stride(2), + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_PDL=use_pdl, + # Release the o_proj at the merge's entry so it streams its weights during the merge. + EARLY_TRIGGER=True, + **mxfp8_args, + **pdl_launch, + **merge_launch_options, + ) + if mxfp8 is not None: + m3_prequant.SLOT.publish(output.view(total_q, width), mxfp8[0][:total_q], mxfp8[1]) + + +__all__ = [ + "MERGE_MXFP8", + "MERGE_MXFP8_MIN_TOKENS", + "MergeMxfp8", + "NVFP4_SF_VEC_SIZE", + "SPARSE_BLOCK_SIZE", + "_sm100f_nvfp4_merge_launch_options", + "_sm100f_nvfp4_num_topk_chunks", + "_sm100f_nvfp4_query_group_size", + "_sm100f_nvfp4_use_linear_softmax", + "minimax_m3_sparse_attn_decode", + "resolve_num_topk_chunks", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/triton_sparse_decode_mq.py new file mode 100644 index 000000000000..048d31b82d68 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/triton_sparse_decode_mq.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multi-query block-sparse decode producer for MiniMax-M3 short-sequence verify steps. + +In short-sequence mode every request attends at most ``topk`` KV blocks and block selection is the identity, so +all ``decode_query_len`` query tokens of a request read the same blocks. The per-token producer +(``_gqa_sparse_decode_kernel``) then dequantizes every block once per query token. This producer runs one program +per (request, block) for all of the request's query tokens: rows are (token, head) pairs, each block is +dequantized once, and the causal limit is a per-row mask. + +It writes the same per-block partials (output scaled by 1/l and the V global scale, log2 LSE) as the per-token +producer with one block per chunk, so ``_merge_topk_attn_out_kernel`` merges them unchanged. +""" + +import triton +import triton.language as tl + +from .triton_sparse_decode import _dequant_nvfp4_rows + + +@triton.jit( + do_not_specialize=["num_reqs"], do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"] +) +def _gqa_sparse_decode_mq_kernel( + q_ptr, # [num_reqs * QLEN, num_heads, head_dim], FP8 or FP16/BF16 + k_ptr, # packed NVFP4 K: [num_pages, 1, page_size, head_dim / 2] uint8 + v_ptr, # packed NVFP4 V: same layout + k_sc_ptr, # K block scales [num_pages, 1, page_size, head_dim / 16], token-linear + v_sc_ptr, # V block scales, 4x4 token-quad order + k_gs_ptr, + v_gs_ptr, + t_ptr, # topk table [1, total_q, topk]; the request's last token's row is read + o_ptr, # partials [topk, total_q, num_heads, head_dim] fp32 + lse_ptr, # partials [topk, total_q, num_heads] fp32 (log2) + block_table_ptr, + seq_lens, + num_reqs, + sm_scale, + stride_qn, + stride_qh, + stride_k_blk, + stride_k_pos, + stride_v_blk, + stride_v_pos, + stride_ksc_blk, + stride_vsc_blk, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + QLEN: tl.constexpr, + GROUP: tl.constexpr, # query heads per KV head (16) + BLOCK_R: tl.constexpr, # next power of two of QLEN * GROUP + HEAD_DIM: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + SCALE_COLS: tl.constexpr, + USE_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr, +): + pid_r = tl.program_id(0) + pid_c = tl.program_id(1) # top-k position == block index under identity selection + sm_scale_log2e = sm_scale * 1.4426950409 * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + + off_r = tl.arange(0, BLOCK_R) + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, HEAD_DIM) + off_dp = tl.arange(0, HEAD_DIM // 2) + off_s = tl.arange(0, SCALE_COLS) + row_t = off_r // GROUP + row_h = off_r % GROUP + r_mask = off_r < QLEN * GROUP + tok = pid_r * QLEN + row_t + + if USE_PDL: + tl.extra.cuda.gdc_wait() + if EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + pid_r) + # Token t of the request sits at seq_len - QLEN + t and attends positions [0, kv_len_t). + kv_len = tl.maximum(seq_len - QLEN + row_t + 1, 0) + blk = tl.load(t_ptr + (pid_r * QLEN + QLEN - 1) * stride_tn + pid_c * stride_tk) + blk_ok = blk >= 0 + blk = tl.where(blk_ok, blk, 0) + page = tl.load(block_table_ptr + pid_r * stride_bt_b + blk, mask=blk_ok, other=0).to(tl.int64) + # Rows past the last token's position are never read: their cache bytes may be stale (0 * NaN in the PV dot). + n_valid = tl.where( + blk_ok, tl.minimum(tl.maximum(seq_len - blk * BLOCK_SIZE_K, 0), BLOCK_SIZE_K), 0 + ) + n_mask = off_n < n_valid + + q = tl.load( + q_ptr + tok[:, None] * stride_qn + row_h[:, None] * stride_qh + off_d[None, :], + mask=r_mask[:, None], + other=0.0, + ).to(tl.float16) + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + page * stride_k_blk + off_n[:, None] * stride_k_pos + off_dp[None, :], + mask=n_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + page * stride_ksc_blk + off_n[:, None] * SCALE_COLS + off_s[None, :], + mask=n_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=HEAD_DIM, + SCALE_COLS=SCALE_COLS, + SF_VEC=HEAD_DIM // SCALE_COLS, + ) + valid = ( + (blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None]) & n_mask[None, :] & r_mask[:, None] + ) + qk = tl.dot(q, tl.trans(k_deq).to(tl.float16)) * sm_scale_log2e + qk = tl.where(valid, qk, float("-inf")) + m = tl.max(qk, axis=1) + m_safe = tl.where(m == float("-inf"), 0.0, m) + p = tl.exp2(qk - m_safe[:, None]) + p_sum = tl.sum(p, axis=1) + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + page * stride_v_blk + off_n[:, None] * stride_v_pos + off_dp[None, :], + mask=n_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=n_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=HEAD_DIM, + SCALE_COLS=SCALE_COLS, + SF_VEC=HEAD_DIM // SCALE_COLS, + ).to(tl.float16) + acc = tl.dot(p.to(tl.float16), v) + # An empty partial stores zero with LSE -inf, so the merge gives it weight 0. + scale = tl.where(p_sum > 0.0, 1.0 / p_sum, 0.0) * v_global_scale + lse = tl.where(p_sum > 0.0, m + tl.log2(p_sum), float("-inf")) + acc = acc * scale[:, None] + tl.store( + o_ptr + + pid_c * stride_o_c + + tok[:, None] * stride_o_b + + row_h[:, None] * stride_o_h + + off_d[None, :], + acc, + mask=r_mask[:, None], + ) + tl.store(lse_ptr + pid_c * stride_l_c + tok * stride_l_b + row_h * stride_l_h, lse, mask=r_mask) + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +def launch_mq_producer( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + block_table, + seq_lens, + o_arg, + lse_arg, + *, + sm_scale, + decode_query_len, + use_pdl, + early_trigger, + num_warps=8, + num_stages=1, +): + """Launch the multi-query producer; arguments as minimax_m3_sparse_attn_decode after its NVFP4 views.""" + total_q, num_heads, head_dim = q.shape + num_reqs = total_q // decode_query_len + max_topk = int(topk_idx.shape[-1]) + rows = decode_query_len * num_heads + pdl_launch = {"launch_pdl": True} if use_pdl else {} + _gqa_sparse_decode_mq_kernel[(num_reqs, max_topk)]( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + o_arg, + lse_arg, + block_table, + seq_lens, + num_reqs, + sm_scale, + q.stride(0), + q.stride(1), + k_paged.stride(0), + k_paged.stride(2), + v_paged.stride(0), + v_paged.stride(2), + k_block_scale.stride(0), + v_block_scale.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + o_arg.stride(0), + o_arg.stride(1), + o_arg.stride(2), + lse_arg.stride(0), + lse_arg.stride(1), + lse_arg.stride(2), + block_table.stride(0), + QLEN=decode_query_len, + GROUP=num_heads, + BLOCK_R=triton.next_power_of_2(rows), + HEAD_DIM=head_dim, + BLOCK_SIZE_K=int(k_paged.shape[2]), + SCALE_COLS=head_dim // 16, + USE_PDL=use_pdl, + EARLY_TRIGGER=early_trigger, + num_warps=num_warps, + num_stages=num_stages, + **pdl_launch, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/weights.py new file mode 100644 index 000000000000..da27f2fba67f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/weights.py @@ -0,0 +1,865 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight tables and loaders: MiniMax-M3 NVFP4 / sm_103 / tp4. + +Two tables, one per checkpoint the target reads: + +* ``MODEL_WEIGHTS`` (``MiniMaxM3Tp4Weights``): the vision-language checkpoint's language model, + ``language_model.*``. Every rank is handed the whole checkpoint, so every tensor-parallel split is this + table's: the Q heads, the K/V/index-Q heads (one each per rank), the dense MLPs and the shared expert's + intermediate, the o_proj / down-projection input columns, every routed expert's intermediate (all 128 + experts live on every rank), and the vocabulary of the shell's ``lm_head``. The index-K head is replicated. +* ``DRAFT_WEIGHTS`` (``MiniMaxM3Tp4DraftWeights``): the one-layer BF16 Eagle3 drafter's own checkpoint + (``speculative_model``), loaded through the shell's ``load_draft_weights`` when speculation is on. + +Storage is HF ``[out, in]`` row-major. Five families need a transform, each a layout a kernel demands: + +* **MXFP8 block scales** (every attention projection, the dense MLPs and the shared expert): the checkpoint + stores UE8M0 bytes ``[N, K / 32]``; the GEMMs read the 128x4-swizzled order + (``trtllm::block_scale_interleave``). A fused projection concatenates its shards' scales first and swizzles + once. +* **The five-way sparse-layer projection** ``[Q | K | V | index-Q | index-K]`` and the three-way dense one + ``[Q | K | V]``: the per-rank shards concatenated along N. +* **Gated MLPs** (dense and shared expert): ``[gate ; up]`` rows. +* **NVFP4 routed experts** (layers 3..59): ``[up (w3) ; gate (w1)]`` rows, then the trtllm-gen row shuffle of + the weight and of its E4M3 scales (gated-activation reorder + 128-row epilogue tile shuffle, + ``trtllmgen_maybe_get_cached_*_permute_indices``) and the 128x4 swizzle of the scales. +* **Gemma boundary norms**: ``input_layernorm``, ``post_attention_layernorm`` and ``model.norm`` are stored as + Gemma norms, ``(1 + w) * x``; the fused all-reduce + RMSNorm kernels apply ``w * x``, so ``1`` is added + (in BF16) at load. The per-head Q/K/index norms stay Gemma: their kernels apply the offset themselves. + +The routed experts' per-tensor scalars (input scales and second-level weight scales, every expert) load raw; +``modeling`` derives the kernels' global scales, alphas and activation constants from them once the weights +are real. + +Loading contract: ``load(model, weights)`` consumes the engine-provided checkpoint dict, fills every declared +parameter exactly once and asserts coverage in both directions: every declared parameter written, and every +checkpoint key either consumed or in ``expected_unconsumed`` (the vision tower and its projector, which this +text-only target does not build, and the language model's multi-token-prediction head). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, Callable + +import torch +import torch.nn as nn + +from tensorrt_llm._torch.moe.fused_moe.quantization import ( + trtllmgen_maybe_get_cached_w2_permute_indices, + trtllmgen_maybe_get_cached_w3_w1_permute_indices, +) + +#: ``layers=ALL_LAYERS`` is the common case: a weight every decoder layer carries. +ALL_LAYERS = object() + +_LM = "language_model." +#: Checkpoint key families this target does not build: the vision tower and its two projectors (text only), and the +#: language model's multi-token-prediction head (speculation, when on, is the separate Eagle3 drafter). +_UNBUILT_PREFIXES = ( + "vision_tower.", + "multi_modal_projector.", + "patch_merge_mlp.", + _LM + "model.mtp.", +) + +MX_BLOCK = 32 # MXFP8: one UE8M0 scale per 32 elements along K +NV_BLOCK = 16 # NVFP4: one E4M3 scale per 16 elements along K +EPILOGUE_TILE_M = ( + 128 # the trtllm-gen MoE runner's row tile: fixes its weight and scale row shuffles +) + + +@dataclass(frozen=True) +class W: + """One weight role: what it is, where its bytes come from, what happens on the way. + + ``shape`` takes the dimension bundle and returns the declared shape. ``src`` is a callable + ``(d, i) -> [(key or key tuple, index into the param | None), ...]`` with keys relative to the checkpoint + root; a key tuple hands the transform those tensors together. ``layers`` is ``ALL_LAYERS``, ``None`` (not + per-layer) or ``d -> layer indices``. ``transform(src, d)`` runs on the destination's device before the copy. + """ + + name: str + shape: Callable[[Any], tuple[int, ...]] + src: Callable + dtype: torch.dtype | None = None + transform: Callable | None = None + layers: object | None = ALL_LAYERS + + def layer_indices(self, d: SimpleNamespace) -> list[int] | None: + if self.layers is None: + return None + if self.layers is ALL_LAYERS: + return list(range(d.num_layers)) + return list(self.layers(d)) + + +# ------------------------------------------------------------------------------------------------ shards + + +def _col( + t: torch.Tensor, d: SimpleNamespace, tp: int | None = None, rank: int | None = None +) -> torch.Tensor: + """This rank's rows (column-parallel output split), even split, sliced on the host and moved to the + destination device (only this rank's bytes are read).""" + tp = d.tp if tp is None else tp + rank = d.rank if rank is None else rank + n = t.shape[0] // tp + return _dev(t[rank * n : (rank + 1) * n], d) + + +def _row(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """This rank's input columns (row-parallel input split), even split, on the destination device.""" + n = t.shape[1] // d.tp + return _dev(t[:, d.rank * n : (d.rank + 1) * n], d) + + +def _dev(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """``t`` on the destination device (``d.dev``, set per fill).""" + return t.to(d.dev, non_blocking=True) + + +def _swizzle(scale: torch.Tensor) -> torch.Tensor: + """128x4 swizzle of a row-major ``[N, K / 32]`` UE8M0 scale matrix, flat.""" + return torch.ops.trtllm.block_scale_interleave(scale.contiguous()) + + +# --------------------------------------------------------------------------------------------- transforms + + +def _gemma_fold(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """A Gemma boundary norm as the plain norm the fused all-reduce epilogue applies: ``w + 1`` in BF16.""" + return t + 1.0 + + +def _f32(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return t.to(torch.float32) + + +def _qkv_sparse_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[Q | K | V | index-Q | index-K]``: Q by attention heads, K/V/index-Q by KV heads, index-K replicated.""" + q, k, v, iq, ik = t + return torch.cat( + [_col(q, d), _col(k, d), _col(v, d), _col(iq, d), _dev(ik, d)], dim=0 + ).contiguous() + + +def _qkv_sparse_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v, iq, ik = t + return _swizzle( + torch.cat([_col(q, d), _col(k, d), _col(v, d), _col(iq, d), _dev(ik, d)], dim=0) + ) + + +def _qkv_dense_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + return torch.cat([_col(q, d), _col(k, d), _col(v, d)], dim=0).contiguous() + + +def _qkv_dense_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + return _swizzle(torch.cat([_col(q, d), _col(k, d), _col(v, d)], dim=0)) + + +def _index_qk_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[index-Q | index-K]`` in BF16: this rank's index-Q heads and the replicated index-K, the E4M3 values cast + (the scale tensors, the last two sources, are consumed only).""" + iq, ik = t[0], t[1] + return torch.cat([_col(iq, d), _dev(ik, d)], dim=0).to(d.dtype).contiguous() + + +def _gate_up_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return torch.cat([_col(gate, d), _col(up, d)], dim=0).contiguous() + + +def _gate_up_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return _swizzle(torch.cat([_col(gate, d), _col(up, d)], dim=0)) + + +def _row_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _row(t, d).contiguous() + + +def _row_s(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _swizzle(_row(t, d)) + + +def _col_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _col(t, d).contiguous() + + +#: Permute indices of the trtllm-gen expert relayouts, keyed by (shape, family, scale group): one expert's are +#: every expert's, so each is computed once per load. +_PERMUTE_CACHE: dict = {} + + +def _expert_w31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 operand: ``[up ; gate]`` packed-FP4 rows of this rank, then the gated-activation row + reorder and the epilogue-tile shuffle.""" + up, gate = t + w = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + w, _PERMUTE_CACHE, EPILOGUE_TILE_M, is_gated_act_gemm=True + ) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 E4M3 scales: ``[up ; gate]`` rows of this rank, the same row reorder at scale + granularity, then the 128x4 swizzle.""" + up, gate = t + s = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK, is_gated_act_gemm=True + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled.view(torch.uint8).reshape(shape)) + .view(torch.float8_e4m3fn) + .reshape(shape) + ) + + +def _expert_w2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 operand: this rank's input columns, then the epilogue-tile row shuffle.""" + w = _row(t, d).contiguous() + perm = trtllmgen_maybe_get_cached_w2_permute_indices(w, _PERMUTE_CACHE, EPILOGUE_TILE_M) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 E4M3 scales: this rank's columns, the row shuffle at scale granularity, the swizzle.""" + s = _row(t, d).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w2_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled).view(torch.float8_e4m3fn).reshape(shape) + ) + + +# ------------------------------------------------------------------------------------------------ sources + + +def _one(template: str): + return lambda d, i, t=template: [(_LM + t.format(i=i), None)] + + +def _many(*templates: str): + return lambda d, i, ts=templates: [(tuple(_LM + t.format(i=i) for t in ts), None)] + + +def _per_expert(*suffixes: str): + """One source per routed expert, written at its expert index (all experts are local at tp4).""" + + def build(d, i, sx=suffixes): + out = [] + for e in range(d.num_experts): + keys = tuple(_LM + f"model.layers.{i}.block_sparse_moe.experts.{e}.{s}" for s in sx) + out.append((keys if len(keys) > 1 else keys[0], (e,))) + return out + + return build + + +def _dense_layers(d): + return range(d.first_sparse) + + +def _sparse_layers(d): + return range(d.first_sparse, d.num_layers) + + +def _three_way_qkv_layers(d): + """Layers whose QKV is the three-way [Q | K | V]: the dense layers, and every layer without the five-way packed + projection (``fuse_qkv_index_projection=False``).""" + return range(d.first_sparse) if d.fused else range(d.num_layers) + + +def _five_way_qkv_layers(d): + return range(d.first_sparse, d.num_layers) if d.fused else range(0) + + +def _unfused_sparse_layers(d): + return range(0) if d.fused else range(d.first_sparse, d.num_layers) + + +_ATTN = "model.layers.{i}.self_attn" +_MOE = "model.layers.{i}.block_sparse_moe" + + +class MiniMaxM3Tp4Weights: + """MiniMax-M3 NVFP4 / sm_103 / tp4: the language model's table and the passes over it. + + ``declare`` allocates every weight it names, ``manifest`` maps checkpoint keys onto them, ``load`` fills + them; ``_targets`` is the one place that decides which parameters exist. + """ + + #: Free the per-expert relayout scratch after each expert stack. + RELEASE_AFTER = "_e_s2" + + WEIGHTS: tuple[W, ...] = ( + # --- the two boundary norms of every layer (Gemma, folded) --- + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.input_layernorm.weight"), + transform=_gemma_fold, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.post_attention_layernorm.weight"), + transform=_gemma_fold, + ), + # --- attention: per-head Q/K norms (Gemma, unfolded), the packed QKV and the output projection --- + W("q_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".q_norm.weight")), + W("k_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".k_norm.weight")), + W( + "qkv_w", + shape=lambda d: (d.qkv_dense_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in "qkv")), + transform=_qkv_dense_w, + layers=_three_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_dense_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many(*(f"{_ATTN}.{p}_proj.weight_scale_inv" for p in "qkv")), + transform=_qkv_dense_s, + layers=_three_way_qkv_layers, + ), + W( + "qkv_w", + shape=lambda d: (d.qkv_sparse_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in ("q", "k", "v", "index_q", "index_k"))), + transform=_qkv_sparse_w, + layers=_five_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_sparse_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + *( + f"{_ATTN}.{p}_proj.weight_scale_inv" + for p in ("q", "k", "v", "index_q", "index_k") + ) + ), + transform=_qkv_sparse_s, + layers=_five_way_qkv_layers, + ), + # The separate BF16 index projection [index-Q | index-K] of the unfused sparse layers. The built-in model + # builds it as an unquantized Linear and fills it through the fused gate/up loader, which casts the + # checkpoint's E4M3 values to BF16 and does not apply their block scales: so does this table (the scale + # tensors are consumed, not applied). + W( + "idx_w", + shape=lambda d: (d.index_heads * d.index_dim + d.index_dim, d.hidden), + src=_many( + _ATTN + ".index_q_proj.weight", + _ATTN + ".index_k_proj.weight", + _ATTN + ".index_q_proj.weight_scale_inv", + _ATTN + ".index_k_proj.weight_scale_inv", + ), + transform=_index_qk_w, + layers=_unfused_sparse_layers, + ), + W( + "idx_q_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_q_norm.weight"), + layers=_sparse_layers, + ), + W( + "idx_k_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_k_norm.weight"), + layers=_sparse_layers, + ), + W( + "o_w", + shape=lambda d: (d.hidden, d.q_width), + dtype=torch.float8_e4m3fn, + src=_one(_ATTN + ".o_proj.weight"), + transform=_row_w, + ), + W( + "o_s", + shape=lambda d: (d.hidden * (d.q_width // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_ATTN + ".o_proj.weight_scale_inv"), + transform=_row_s, + ), + # --- dense SwiGLU-OAI MLP (layers 0..2) --- + W( + "mlp_gu_w", + shape=lambda d: (2 * d.dense_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight", "model.layers.{i}.mlp.up_proj.weight" + ), + transform=_gate_up_w, + layers=_dense_layers, + ), + W( + "mlp_gu_s", + shape=lambda d: (2 * d.dense_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight_scale_inv", + "model.layers.{i}.mlp.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_dense_layers, + ), + W( + "mlp_dn_w", + shape=lambda d: (d.hidden, d.dense_inter), + dtype=torch.float8_e4m3fn, + src=_one("model.layers.{i}.mlp.down_proj.weight"), + transform=_row_w, + layers=_dense_layers, + ), + W( + "mlp_dn_s", + shape=lambda d: (d.hidden * (d.dense_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one("model.layers.{i}.mlp.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_dense_layers, + ), + # --- MoE (layers 3..59): the fp32 router and its bias, the shared expert --- + W( + "router", + shape=lambda d: (d.num_experts, d.hidden), + dtype=torch.float32, + src=_one(_MOE + ".gate.weight"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "router_bias", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_one(_MOE + ".e_score_correction_bias"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "sh_gu_w", + shape=lambda d: (2 * d.shared_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + _MOE + ".shared_experts.gate_proj.weight", _MOE + ".shared_experts.up_proj.weight" + ), + transform=_gate_up_w, + layers=_sparse_layers, + ), + W( + "sh_gu_s", + shape=lambda d: (2 * d.shared_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + _MOE + ".shared_experts.gate_proj.weight_scale_inv", + _MOE + ".shared_experts.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_sparse_layers, + ), + W( + "sh_dn_w", + shape=lambda d: (d.hidden, d.shared_inter), + dtype=torch.float8_e4m3fn, + src=_one(_MOE + ".shared_experts.down_proj.weight"), + transform=_row_w, + layers=_sparse_layers, + ), + W( + "sh_dn_s", + shape=lambda d: (d.hidden * (d.shared_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_MOE + ".shared_experts.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_sparse_layers, + ), + # --- the routed experts, NVFP4 (trtllm-gen layouts), all 128 on every rank --- + W( + "e_w31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // 2), + dtype=torch.uint8, + src=_per_expert("w3.weight", "w1.weight"), + transform=_expert_w31, + layers=_sparse_layers, + ), + W( + "e_w2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2), + dtype=torch.uint8, + src=_per_expert("w2.weight"), + transform=_expert_w2, + layers=_sparse_layers, + ), + W( + "e_s31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w3.weight_scale", "w1.weight_scale"), + transform=_expert_s31, + layers=_sparse_layers, + ), + W( + "e_s2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w2.weight_scale"), + transform=_expert_s2, + layers=_sparse_layers, + ), + # The per-tensor scalars of every expert, raw; modeling derives the kernels' scalars from them. + *( + W( + f"e_{name}", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_per_expert(suffix), + layers=_sparse_layers, + ) + for name, suffix in ( + ("isc1", "w1.input_scale"), + ("isc3", "w3.input_scale"), + ("isc2", "w2.input_scale"), + ("ws2_1", "w1.weight_scale_2"), + ("ws2_3", "w3.weight_scale_2"), + ("ws2_2", "w2.weight_scale_2"), + ) + ), + # --- not per-layer --- + W( + "embed", + shape=lambda d: (d.vocab, d.hidden), + src=_one("model.embed_tokens.weight"), + layers=None, + ), + W( + "final_norm", + shape=lambda d: (d.hidden,), + src=_one("model.norm.weight"), + transform=_gemma_fold, + layers=None, + ), + ) + + def dims(self, core) -> SimpleNamespace: + cfg = core.model_config.pretrained_config + mapping = core.model_config.mapping + sparse = cfg.sparse_attention_config + freq = [int(f) for f in sparse["sparse_attention_freq"]] + tp = mapping.tp_size + head_dim = int(cfg.head_dim) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + index_heads = max(int(sparse["sparse_num_index_heads"]) // tp, 1) + index_dim = int(sparse["sparse_index_dim"]) + return SimpleNamespace( + fused=bool( + getattr( + core.model_config.sparse_attention_config, "fuse_qkv_index_projection", False + ) + ), + num_layers=cfg.num_hidden_layers, + first_sparse=freq.index(1), + hidden=cfg.hidden_size, + head_dim=head_dim, + heads=heads, + kv_heads=kv_heads, + index_heads=index_heads, + index_dim=index_dim, + q_width=heads * head_dim, + qkv_dense_rows=(heads + 2 * kv_heads) * head_dim, + qkv_sparse_rows=(heads + 2 * kv_heads) * head_dim + index_heads * index_dim + index_dim, + dense_inter=int(cfg.dense_intermediate_size) // tp, + shared_inter=int(cfg.shared_intermediate_size) * int(cfg.n_shared_experts) // tp, + moe_inter=int(cfg.intermediate_size) // tp, + num_experts=int(cfg.num_local_experts), + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + layers = entry.layer_indices(d) + if layers is None: + return [(entry.name, None)] + return [(f"l{i}_{entry.name}", i) for i in layers] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + """Allocate every weight the table declares (meta under the engine's meta init).""" + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, core) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + """param key -> [(checkpoint key or key tuple, index into the param | None, transform | None)].""" + d = self.dims(core) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, core, weights) -> set: + """The vision tower and its projectors, and the multi-token-prediction head: in the checkpoint, never built + by this target.""" + return {k for k in weights.keys() if k.startswith(_UNBUILT_PREFIXES)} + + def load(self, model, weights) -> None: + """Fill every declared weight and the shell's ``lm_head`` shard; assert coverage both ways.""" + core = model.model + d = self.dims(core) + manifest = self.manifest(core) + assert set(manifest) == set(core.w.keys()), ( + "manifest/parameter drift", + set(manifest) ^ set(core.w.keys()), + ) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(core, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(core.w[param_key], ckpt, index, transform, weights, d, consumed) + if param_key.endswith(self.RELEASE_AFTER): + torch.cuda.empty_cache() + # The shell owns the vocab-parallel head: this rank's rows of the full table. + _fill(model.lm_head.weight, _LM + "lm_head.weight", None, _col_w, weights, d, consumed) + shell = { + n for n, _ in model.named_parameters() if not n.startswith(("model.w.", "draft_model.")) + } + assert shell == {"lm_head.weight"}, f"unfed shell parameters: {sorted(shell)}" + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + "checkpoint coverage: leftover keys are not exactly the unbuilt families (vision, MTP head); " + f"unexpected {sorted(leftover - expected)[:8]}, missed {sorted(expected - leftover)[:8]}" + ) + + +def _materialize(entry) -> torch.Tensor: + """``[:]`` on a lazy safetensors slice; a 0-dim tensor (the per-tensor scalars) rejects that index.""" + return entry if getattr(entry, "ndim", 1) == 0 else entry[:] + + +def _fill(param, ckpt, index, transform, weights, d, consumed) -> None: + keys = ckpt if isinstance(ckpt, tuple) else (ckpt,) + for key in keys: + assert key in weights, f"checkpoint key missing: {key}" + dst = param.data if index is None else param.data[index] + d.dev = dst.device + # Host tensors: a transform slices this rank's part on the host and moves only that to the device. + src_all = tuple(_materialize(weights[key]) for key in keys) + if transform is not None: + src = transform(src_all if isinstance(ckpt, tuple) else src_all[0], d) + else: + assert len(src_all) == 1, "a multi-key row needs a transform" + src = src_all[0] + if src.dim() == 0 and dst.dim() == 0: + pass + elif src.numel() == 1 and dst.numel() == 1: + src = src.reshape(dst.shape) + src = src.to(dst.device, non_blocking=True) + assert tuple(dst.shape) == tuple(src.shape), (ckpt, tuple(dst.shape), tuple(src.shape)) + assert src.dtype == dst.dtype, (ckpt, src.dtype, dst.dtype) + dst.copy_(src, non_blocking=True) + consumed.update(keys) + _release(weights, keys) + + +def _release(weights, keys) -> None: + """Drop consumed checkpoint tensors from the engine's dict as the load goes (bounded host memory).""" + if hasattr(weights, "mark_consumed_keys"): + weights.mark_consumed_keys(keys) + else: + for key in keys: + weights.pop(key, None) + + +MODEL_WEIGHTS = MiniMaxM3Tp4Weights() + + +# -------------------------------------------------------------------------------------------- the drafter + + +def _draft_one(template: str): + return lambda d, i, t=template: [(t, None)] + + +def _draft_many(*templates: str): + return lambda d, i, ts=templates: [(tuple(ts), None)] + + +def _draft_fc_norms(d): + return range(d.capture_layers) + + +def _draft_qkv(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + return torch.cat([_col(q, d), _col(k, d), _col(v, d)], dim=0).contiguous() + + +class MiniMaxM3Tp4DraftWeights: + """The one-layer BF16 Eagle3 drafter (``speculative_model``), tp4. + + The fc over the three captured target layers and its per-layer fc_norms, the layer's three norms, its QKV + (input ``[embeds | hidden]``, so twice the hidden width; Q by heads, K/V by KV heads) and o_proj, its gated + MLP, the final norm and the full-vocabulary head (vocab-parallel). The checkpoint's ``embed_tokens`` is the + target's table (the drafter reads the target's); it is a predicted non-load. + """ + + WEIGHTS: tuple[W, ...] = ( + W( + "fc", + shape=lambda d: (d.hidden, d.capture_layers * d.hidden), + src=_draft_one("fc.weight"), + layers=None, + ), + W( + "fc_norm", + shape=lambda d: (d.hidden,), + src=lambda d, i: [(f"fc_norm.{i}.weight", None)], + layers=_draft_fc_norms, + ), + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.input_layernorm.weight"), + layers=None, + ), + W( + "hidden_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.hidden_norm.weight"), + layers=None, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.post_attention_layernorm.weight"), + layers=None, + ), + W( + "qkv", + shape=lambda d: (d.qkv_rows, 2 * d.hidden), + src=_draft_many(*(f"layers.0.self_attn.{p}_proj.weight" for p in "qkv")), + transform=_draft_qkv, + layers=None, + ), + W( + "o", + shape=lambda d: (d.hidden, d.q_width), + src=_draft_one("layers.0.self_attn.o_proj.weight"), + transform=_row_w, + layers=None, + ), + W( + "gu", + shape=lambda d: (2 * d.inter, d.hidden), + src=_draft_many("layers.0.mlp.gate_proj.weight", "layers.0.mlp.up_proj.weight"), + transform=_gate_up_w, + layers=None, + ), + W( + "dn", + shape=lambda d: (d.hidden, d.inter), + src=_draft_one("layers.0.mlp.down_proj.weight"), + transform=_row_w, + layers=None, + ), + W("final_norm", shape=lambda d: (d.hidden,), src=_draft_one("norm.weight"), layers=None), + ) + + def dims(self, draft) -> SimpleNamespace: + cfg = draft.draft_config.pretrained_config + mapping = draft.draft_config.mapping + tp = mapping.tp_size + head_dim = int(getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads)) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + return SimpleNamespace( + hidden=cfg.hidden_size, + capture_layers=draft.capture_layers, + head_dim=head_dim, + q_width=heads * head_dim, + qkv_rows=(heads + 2 * kv_heads) * head_dim, + inter=cfg.intermediate_size // tp, + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + if entry.layers is None: + return [(entry.name, None)] + return [(f"{entry.name}{i}", i) for i in entry.layers(d)] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, draft) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + d = self.dims(draft) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, draft, weights) -> set: + """The drafter checkpoint's own embedding: the drafter reads the target's table.""" + return {"embed_tokens.weight"} & set(weights.keys()) + + def load(self, draft, weights) -> None: + d = self.dims(draft) + manifest = self.manifest(draft) + assert set(manifest) == set(draft.w.keys()), set(manifest) ^ set(draft.w.keys()) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(draft, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(draft.w[param_key], ckpt, index, transform, weights, d, consumed) + _fill(draft.lm_head.weight, "lm_head.weight", None, _col_w, weights, d, consumed) + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + f"drafter checkpoint coverage: unexpected {sorted(leftover - expected)[:8]}, " + f"missed {sorted(expected - leftover)[:8]}" + ) + + +DRAFT_WEIGHTS = MiniMaxM3Tp4DraftWeights() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/__init__.py new file mode 100644 index 000000000000..91238370999c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/__init__.py @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 NVFP4 / sm_103 / tp8.""" + +from ..kernels import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py new file mode 100644 index 000000000000..1e8c520ca9d6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py @@ -0,0 +1,4312 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ModelingV2 target: MiniMax-M3 NVFP4 / sm_103 / tp8 -- self-contained modeling code. + +MiniMax-M3's language model: 60 layers, hidden 6144, 64 query heads and 4 KV heads of dim 128. Layers 0-2 are dense +attention with a dense SwiGLU-OAI MLP; layers 3-59 are MSA block-sparse attention (an index branch scores 128-token +blocks and the top 16 are attended) with a MoE of 128 routed experts (top 4, sigmoid scores with a correction bias, +renormalized and scaled by 2) plus one shared expert. The checkpoint is a modelopt mixed-precision export: NVFP4 routed +experts, MXFP8 projections, fp32 router, bf16 embedding and head; the KV cache is NVFP4 for the sparse layers and FP8 +for the dense ones. + +``tp8`` is eight GB300s over two trays (multicast all-reduces over MNNVL): head-split attention (8 query heads, one +KV head and one index head per rank; the 4 KV and index heads each duplicated on two ranks) and tensor-parallel +experts (every expert on every rank, intermediate 384 per rank), no attention DP. Speculative decoding, when +configured, is one-model Eagle3 (draft length 3): the BF16 drafter (``M3Eagle3Draft``) and the BF16 target head, both +at checkpoint precision; decode steps of up to 8 requests verify and draft through fused kernels at that precision +(``M3Eagle3Worker``). The decode kernels are specialized for one request verifying a draft (4 query rows), the +deployment this target was measured on. + +**The core** (``MiniMaxM3Tp8ModelingV2``) holds the weights (``weights.MODEL_WEIGHTS``) and the runtime resources the +steps use -- one MSA attention backend per layer, the TP all-reduce communicator, the RoPE table, the MXFP8 GEMM +dispatch state the engine tunes -- and dispatches each step, on the host, from its context count alone +(``_select_target``): + +* ``DecodeTarget`` takes every pure generation step. Up to ``DECODE_MAX_TOKENS`` tokens (up to 8 requests, each one + token or a draft verify of 1 + 3 rows) with 1-D input ids it runs the decode-step kernels and picks each layer + boundary's collective by the step's token bucket (``DECODE_PLANS``); it runs its own generic layer for the first + layer when the fused embedding does not take the step, for every layer whose output Eagle3 captures, and for every + layer of a wider step. +* ``PrefillTarget`` takes every step with context rows: prefill chunks and mixed steps, every layer in its generic + form over the MSA backend. + +The two targets share no body: each owns its layer, attention and MoE compositions and its ``step_args``, the one +projection of the step metadata its forward reads. What both call are the module-level op adapters, each picking the +kernel for one call from shapes and layouts (``_linear``, ``_qk_norm_rope``, ``_sparse_qkv_producer``, the router and +MoE-runner calls, the all-reduce + norm calls), which read no step metadata themselves. + +Every fast path the decode step can take, the rows each covers (``FAST_PATHS``, each paired with the entry whose +``CELLS`` certify it) and what a step that cannot take one does are below: **uncovered** (outside every covered +range: log once, run the generic path) or **missed** (inside a covered range but a precondition is missing: raise +``FastPathUnavailable`` under ``modeling_v2='require'``, else log once and fall back). + +**Hand-offs between kernels.** A producer that already holds an activation in registers writes its MXFP8 form for the +next MXFP8 projection: the decode boundary all-reduce for the next layer's QKV, the sparse-decode merge for the +o_proj. It publishes it in ``ops/m3_prequant.SLOT`` keyed by the BF16 tensor; the next +MXFP8 projection on a decode graph takes it (``_linear``), and every take clears the slot. The decode boundaries' +push targets (the o_proj / shared-expert GEMVs and the MoE engine writing straight into the boundary all-reduce's +exchange) are passed explicitly from the layer to the producer and back. + +**What this target asserts rather than adapts**: SM 10.3, the topology and checkpoint (``routing.py``); the +MiniMax-M3 MSA sparse attention runtime with the fused five-way QKV/index projection, an FP8 indexer and an FP8 or +NVFP4 KV cache, no torch.compile and no piecewise prefill graphs (``_validate_runtime``); text-only input (the vision +tower is not built and a step carrying image input raises). +""" + +from __future__ import annotations + +import atexit +import copy +import dataclasses +import os +import threading +import time +from dataclasses import dataclass +from typing import ( + Any, + Callable, + Dict, + FrozenSet, + Hashable, + Iterable, + List, + Mapping, + Optional, + Tuple, +) + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + ModelingV2Mode, + NoTarget, + resolved_mode, + step_contract_enabled, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( + fp4_block_scale_moe_runner, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 + flashinfer_fused_add_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import _aligned_nvfp4_dequant_scales +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + MSA_REQUIRED_TOPK, + msa_paged_kv, +) +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.distributed import AllReduce, 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_speculative import ( + _ensure_draft_vocab_size, + _get_requested_draft_moe_backend, + _set_draft_kv_cache_quant_algo, + _slice_spec_position_ids, +) +from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import ( + MXFP8LinearMethod, + TensorParallelMode, + flashinfer_mxfp8_autotune_active, + flashinfer_mxfp8_decode_graph_capture_active, +) +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.moe.fused_moe.routing import RoutingMethodType +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy, RuntimeHooks +from tensorrt_llm._torch.speculative import SpecMetadata +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + should_use_separate_draft_kv_cache, +) +from tensorrt_llm._torch.utils import ActType_TrtllmGen +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantAlgo + +from . import weights as _weights +from .ops import m3_draft_step as m3_draft_step_kernel +from .ops import m3_fused_score_select as _fused_score_select_kernel +from .ops import m3_index_score as _index_score_kernel +from .ops import m3_mc_allreduce, m3_prequant, m3_spec_tail +from .ops import m3_sampled_accept as m3_sampled_accept_kernel +from .ops import m3_shard_accept as m3_shard_accept_kernel +from .ops import msa_kv_update as _msa_kv_update_kernel +from .ops.activation.m3_silu_and_mul import m3_silu_and_mul +from .ops.attention.m3_draft_attention import m3_draft_attention +from .ops.attention.m3_fp8_indexer_qk_norm_rope import m3_fp8_indexer_qk_norm_rope +from .ops.attention.m3_fused_score_select import ( + MAX_REQUESTS as FUSED_SCORE_SELECT_KERNEL_MAX_REQUESTS, +) +from .ops.attention.m3_fused_score_select import m3_fused_score_select +from .ops.attention.m3_index_score import SPLIT_K as INDEX_SCORE_SPLIT_K +from .ops.attention.m3_index_score import m3_index_score +from .ops.attention.m3_msa_attention import m3_msa_attention +from .ops.attention.m3_producer import m3_producer +from .ops.attention.m3_qk_norm_rope_fp8 import m3_qk_norm_rope_fp8 +from .ops.attention.m3_qkv_insert import m3_qkv_insert +from .ops.attention.m3_select_blocks import m3_select_blocks +from .ops.attention.m3_select_long import MAX_BLOCKS as SELECT_LONG_MAX_BLOCKS +from .ops.attention.m3_select_long import m3_select_long +from .ops.attention.m3_sparse_decode import m3_sparse_decode +from .ops.attention.m3_sparse_decode_lc8 import m3_sparse_decode_lc8 +from .ops.comm.m3_allreduce import m3_allreduce +from .ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm +from .ops.comm.m3_mc_finalize_allreduce_norm import m3_mc_finalize_allreduce_norm +from .ops.embedding.m3_embedding import m3_embedding +from .ops.gemm.m3_bf16_linear import m3_bf16_linear +from .ops.gemm.m3_draft_lowm import TACTICS as DRAFT_LOWM_TACTICS +from .ops.gemm.m3_draft_lowm import m3_draft_lowm, tactic_for +from .ops.gemm.m3_fp32_linear import m3_fp32_linear +from .ops.gemm.m3_mx_gemv import DENSE_DOWN, DENSE_GATE_UP, O_PROJ, QKV, QKV_SPARSE, m3_mx_gemv +from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 +from .ops.gemm.m3_mxfp8_gemm import m3_mxfp8_gemm +from .ops.gemm.m3_mxfp8_linear import m3_mxfp8_linear +from .ops.long_select import long_select_blocks +from .ops.m3_router_gemv import m3_router_gemv as _router_gemv +from .ops.m3_router_gemv import m3_router_gemv_fp4 as _router_gemv_fp4 +from .ops.m3_router_gemv import m3_router_tc as _router_tc +from .ops.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine +from .ops.moe.m3_router_gemv import m3_router_gemv +from .ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 +from .ops.moe.m3_router_tc import m3_router_tc +from .ops.mx_gemv import op as mx_gemv +from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm +from .ops.spec.m3_draft_input import m3_draft_input +from .ops.spec.m3_draft_step import MAX_TOKENS as DRAFT_STEP_KERNEL_MAX_TOKENS +from .ops.spec.m3_draft_step import m3_draft_step +from .ops.spec.m3_force_vec import m3_force_vec +from .ops.spec.m3_shard_accept import m3_shard_accept + +_SM = (10, 3) +# -------------------------------------------------------------------------------------------------------------------- +# Row ranges +# -------------------------------------------------------------------------------------------------------------------- + +# The largest generation step the decode kernels take: 8 requests verifying 1 + 3 drafts and their padding rows. A +# wider pure generation step stays on the decode target and runs its generic layers. +DECODE_MAX_TOKENS = 48 # capacity: fixed: the decode kernels' step bound + + +@dataclass(frozen=True) +class DecodePlan: + """The layer-boundary collectives of one decode token bucket (GB300 measurements in the comments).""" + + name: str + max_tokens: int + # The attention boundary's multicast all-reduce + residual + RMSNorm (5.05 vs 5.27 us native at 1 token). + attn_mc_ar: bool + # The same kernel at the MoE boundary with the shared + routed add folded in: wins up to 48 tokens. + moe_mc_ar: bool + # The MoE-boundary all-reduce also emits the next QKV activation in MXFP8 for the decode GEMV / GEMM. + qkv_prequant: bool + + +# tiny: every step up to 8 tokens (c1..c8 without speculation, the c1 verify, every draft pass); small: the c2 +# verify; verify: the c4 / c8 verifies. Each is one multicast all-reduce bucket (ops/m3_mc_allreduce.BUCKETS). +DECODE_PLANS = ( + DecodePlan("tiny", 8, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), + DecodePlan("small", 16, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), + DecodePlan("verify", DECODE_MAX_TOKENS, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), +) + +# The sparse layers' NVFP4 QKV + indexer producer (ops/attention/m3_producer, every step its geometry fits) triggers its +# dependent kernel on entry, so the kernel after it launches and stages its pre-wait loads while the producer runs. +PRODUCER_EARLY_TRIGGER = True + +# The fused embedding + first RMSNorm kernel: one owner per row count and index dtype. +EMBED_FIRSTNORM_MAX_TOKENS = 8 # capacity: fixed: built rows; wider steps embed generically +# The o_proj pushes its rows into the attention boundary's multicast all-reduce (prepushed mode 2) up to this many +# tokens; wider steps run the GEMV and a plain multicast all-reduce. +ATTN_PUSH_MAX_TOKENS = 16 # capacity: fixed: measured push crossover +AR_PUSH_MODE = 2 +# The 1..8-token MoE engine (router, routed experts at the TP8 width 384 and the shared expert in one launch) pushes +# its sum into the MoE boundary's all-reduce (prepushed mode 2). +MOE_ENGINE_MAX_TOKENS = 8 # capacity: fixed: the engine's built rows +MOE_PUSH_MAX_TOKENS = MOE_ENGINE_MAX_TOKENS # capacity: fixed: the engine's rows +# From this many tokens the MoE runs unfinalized through the TRTLLM-Gen runner and its finalize (top-k gather, +# routing weights, shared expert add) runs inside the boundary all-reduce. +MC_FINALIZE_MIN_TOKENS = MOE_ENGINE_MAX_TOKENS + 1 # capacity: fixed: the rows above the engine +# The router: the fp32 GEMV at one row, the tensor-core router (fused with the routed experts' NVFP4 activation +# quantization) at 2..48 rows. +ROUTER_GEMV_MAX_TOKENS = 1 # capacity: fixed: measured router crossover +ROUTER_TC_MAX_TOKENS = 48 # capacity: fixed: the tensor-core router's built rows +# A greedy Eagle3 verify step of up to this many requests verifies with the fused shard argmax + acceptance kernel +# over the target lm_head's BF16 shard logits (ops/m3_shard_accept), at the configured draft length. The lm_head +# is the in-tree BF16 GEMM: at the TP8 shard (25008 rows) a coded copy of the head saves no time. +SHARD_ACCEPT_DRAFT_LEN = 3 # capacity: fixed: the built Eagle3 draft length +SHARD_ACCEPT_MAX_REQUESTS = DECODE_MAX_TOKENS // ( + SHARD_ACCEPT_DRAFT_LEN + 1 +) # capacity: fixed: decode rows / (draft length + 1) +# On a fused-verify step of up to this many verify rows the drafter's first fc is forked onto the aux stream right +# after the last captured target layer (M3Eagle3Worker). At TP8 it forks beside the verify head instead: forked early +# its GEMM lands beside the target tail it was meant to hide under. +FC_EARLY_MAX_ROWS = 0 # capacity: fixed: measured fork crossover (never at TP8) +# A decode-only Eagle3 step of up to this many requests drafts with the fused draft loop (M3Eagle3Worker): per pass +# the draft input in one launch (ops/m3_spec_tail.draft_input), the drafter layer's attention, MLP and +# all-reduces, and one kernel for the BF16 draft head, the cross-rank argmax and the pass's bookkeeping +# (ops/m3_draft_step), at the configured draft length. The head kernel takes one row per request, at most 8, and +# streams the head's BF16 weight fragments. +DRAFT_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN # capacity: fixed: the built draft length +DRAFT_STEP_MAX_REQUESTS = DRAFT_STEP_KERNEL_MAX_TOKENS +# The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows +# (ops/triton_sparse_decode.MQ_MIN_TOKENS). +SPARSE_DECODE_MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query crossover +# Decode spans of at most this many requests score and select their top-16 blocks in one launch +# (ops/attention/m3_fused_score_select), which loads the older index-K pages before its grid wait; larger spans run +# the index scorer and the long-row selector. +FUSED_SCORE_SELECT_MAX_REQUESTS = FUSED_SCORE_SELECT_KERNEL_MAX_REQUESTS +# The long-row top-16 selector takes engines of 129..LONG_SELECT_MAX_BLOCKS KV blocks of 128 tokens; it holds 2, 4 or +# 9 keys per thread by the row's valid-block count. At most 128 blocks the built-in selector runs. +LONG_SELECT_MIN_BLOCKS = 129 # capacity: fixed: past the built-in selector's 128 +LONG_SELECT_MAX_BLOCKS = SELECT_LONG_MAX_BLOCKS +LONG_SELECT_KEY_WIDTHS = (2, 4, 9) +# The long-row select releases the sparse decode after it right after its own grid wait up to this many query rows; +# wider steps keep the trigger at the select's exit (the early-resident attention CTAs crowd the select grid). +SELECT_EARLY_MAX_ROWS = 8 # capacity: fixed: measured early-release crossover +# The 8-head long-context sparse decode (ops/attention/m3_sparse_decode_lc8) stages the request length, its query rows +# and the request's block-table row before its grid wait up to this many query rows; above, the staged row's shared +# memory costs occupancy and the kernel stages only after the wait (same results bit for bit). +LC8_PREWAIT_MAX_ROWS = 8 # capacity: fixed: measured occupancy crossover +# The MSA index scorer (ops/attention/m3_index_score) prefetches the older index-K pages into L2 before its grid wait +# at every batch it runs (more than FUSED_SCORE_SELECT_MAX_REQUESTS requests); it never runs its one-request +# early-issue schedule here. +INDEX_SCORE_EARLY_MAX_REQUESTS = 0 # capacity: fixed: the early-issue schedule is off at TP8 + +# The decode MXFP8 GEMVs, per projection (N, K): the rows each serves on an activation its producer quantized (the +# MoE boundary for the QKV, the attention merge for the o_proj). +_ALL_ROWS = frozenset(range(1, DECODE_MAX_TOKENS + 1)) +MX_GEMV_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { + QKV_SPARSE: _ALL_ROWS, + QKV: _ALL_ROWS, + O_PROJ: _ALL_ROWS, + DENSE_GATE_UP: _ALL_ROWS, + DENSE_DOWN: _ALL_ROWS, +} +# On a bf16 activation the GEMV quantizes it in its prologue, which pays only up to 8 32-element blocks per thread: +# the most rows each projection takes. +MX_GEMV_BF16_MAX_TOKENS: Dict[Tuple[int, int], int] = { + QKV_SPARSE: 16, + QKV: 16, + O_PROJ: 48, + DENSE_GATE_UP: 16, + DENSE_DOWN: 32, +} +# The Eagle3 drafter's BF16 linears (qkv, o_proj, gate_up, down, fc) through FlashInfer's fixed-tactic low-M kernels +# (ops/gemm/m3_draft_lowm): per (N, K), the rows its tactic table covers; more rows keep the in-tree GEMM. +DRAFT_LOWM_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { + nk: frozenset(range(1, table[-1][0] + 1)) for nk, table in DRAFT_LOWM_TACTICS.items() +} + +Key = Hashable + + +@dataclass(frozen=True) +class RowRange: + """The shapes one fast path takes, and the catalog entry whose cells certify each of them. + + ``keys`` are the covered shapes: row counts for most paths, (query rows per request, requests) for the attention + kernels, a variant name for thresholds. ``certifies`` maps a cell's spec of ``entry`` to the keys that cell + drives (empty for a cell of another variant). + """ + + path: str + keys: FrozenSet[Key] + entry: OpWrapper + certifies: Callable[[Mapping[str, Any]], Iterable[Key]] + + def covers(self, key: Key) -> bool: + return key in self.keys + + def certified(self) -> FrozenSet[Key]: + return frozenset(k for cell in self.entry.CELLS for k in self.certifies(cell.spec)) + + +def _rows(lo: int, hi: int) -> FrozenSet[int]: + return frozenset(range(lo, hi + 1)) + + +def _spec_rows(field: str, **match: Any) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + """A cell's ``spec[field]`` when the cell's other fields match ``match``.""" + + def certifies(spec: Mapping[str, Any]) -> Iterable[Key]: + if any(spec.get(name) != value for name, value in match.items()): + return () + return (spec[field],) + + return certifies + + +def _const(*keys: Key) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + return lambda spec: keys + + +def _sparse_decode_producer(spec: Mapping[str, Any]) -> Iterable[Key]: + if not spec["identity"]: + return () + rows = spec["decode_query_len"] * len(spec["seq_lens"]) + return ("multi_query" if rows >= SPARSE_DECODE_MQ_MIN_TOKENS else "per_row",) + + +def _long_select_widths(spec: Mapping[str, Any]) -> Iterable[Key]: + """The key widths a cell's rows reach: every width up to its blocks' when rows of mixed lengths are drawn.""" + blocks = spec["num_blocks"] + widest = 2 if blocks <= 2048 else 4 if blocks <= 4096 else 9 + if spec["valid"] == "mixed": + return tuple(w for w in LONG_SELECT_KEY_WIDTHS if w <= widest) + return (widest,) + + +EMBED_FIRSTNORM = RowRange( + "embed_firstnorm", + _rows(1, EMBED_FIRSTNORM_MAX_TOKENS), + m3_embedding_firstnorm, + _spec_rows("rows"), +) +SPARSE_DECODE_PRODUCER = RowRange( + "sparse_decode_producer", + frozenset({"per_row", "multi_query"}), + m3_sparse_decode, + _sparse_decode_producer, +) +# The long-context sparse decode: "prewait" up to LC8_PREWAIT_MAX_ROWS query rows, "plain" above. +SPARSE_DECODE_LC8 = RowRange( + "m3_sparse_decode_lc8", + frozenset({"prewait", "plain"}), + m3_sparse_decode_lc8, + lambda spec: ("prewait" if spec["prewait"] else "plain",), +) +# Every token count of a sparse layer on an NVFP4 KV cache at the TP8 heads (prefill and decode): the producer takes the +# geometry, not a row range. +PRODUCER = RowRange("m3_producer", frozenset({"nvfp4_tp8"}), m3_producer, _const("nvfp4_tp8")) +# Requests of a pure decode span whose blocks are scored and selected in one launch. +FUSED_SCORE_SELECT = RowRange( + "m3_fused_score_select", + _rows(1, FUSED_SCORE_SELECT_MAX_REQUESTS), + m3_fused_score_select, + _spec_rows("batch"), +) +INDEX_SCORE = RowRange( + "m3_index_score", + frozenset({"prefetch"}), + m3_index_score, + lambda spec: () if spec["early"] else ("prefetch",), +) +LONG_SELECT = RowRange( + "long_select", frozenset(LONG_SELECT_KEY_WIDTHS), m3_select_long, _long_select_widths +) +MC_ALLREDUCE = RowRange( + "mc_allreduce_norm", + _rows(1, DECODE_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=0), +) +ATTN_PUSH = RowRange( + "attn_push", + _rows(1, ATTN_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +MOE_PUSH = RowRange( + "moe_push", + _rows(1, MOE_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +MC_FINALIZE = RowRange( + "mc_finalize_allreduce_norm", + _rows(MC_FINALIZE_MIN_TOKENS, DECODE_MAX_TOKENS), + m3_mc_finalize_allreduce_norm, + _spec_rows("m"), +) +MOE_ENGINE = RowRange( + "m3_moe_engine", _rows(1, MOE_ENGINE_MAX_TOKENS), m3_moe_engine, _spec_rows("tokens") +) +ROUTER_GEMV = RowRange( + "m3_router_gemv", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv, _spec_rows("m") +) +ROUTER_GEMV_FP4 = RowRange( + "m3_router_gemv_fp4", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv_fp4, _spec_rows("m") +) +ROUTER_TC = RowRange( + "m3_router_tc", + _rows(ROUTER_GEMV_MAX_TOKENS + 1, ROUTER_TC_MAX_TOKENS), + m3_router_tc, + _spec_rows("m"), +) +# The GEMM paths are keyed (projection (N, K), rows): each projection runs its own split-K / ring / MMA +# configuration, so a row count is certified per projection. +MX_GEMV = RowRange( + "m3_mx_gemv", + frozenset((nk, m) for nk, rows in MX_GEMV_ROWS.items() for m in rows), + m3_mx_gemv, + lambda spec: ((spec["nk"], spec["m"]),), +) +MX_GEMV_BF16 = RowRange( + "m3_mx_gemv_bf16", + frozenset((nk, m) for nk, hi in MX_GEMV_BF16_MAX_TOKENS.items() for m in _rows(1, hi)), + m3_mx_gemv_bf16, + lambda spec: ((spec["nk"], spec["m"]),), +) +DRAFT_LOWM = RowRange( + "m3_draft_lowm", + frozenset((nk, m) for nk, rows in DRAFT_LOWM_ROWS.items() for m in rows), + m3_draft_lowm, + lambda spec: ((spec["nk"], spec["m"]),), +) +# (draft length, requests) of a greedy verify step. +SHARD_ACCEPT = RowRange( + "m3_shard_accept", + frozenset((SHARD_ACCEPT_DRAFT_LEN, n) for n in range(1, SHARD_ACCEPT_MAX_REQUESTS + 1)), + m3_shard_accept, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) + +# (draft length, requests) of a fused draft loop, and the rows of its passes' draft input: the first pass's +# requests x (draft length + 1), one per request on the later passes. +DRAFT_STEP = RowRange( + "m3_draft_step", + frozenset((DRAFT_STEP_DRAFT_LEN, n) for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1)), + m3_draft_step, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) +DRAFT_INPUT = RowRange( + "m3_draft_input", + frozenset( + rows + for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1) + for rows in (n, n * (DRAFT_STEP_DRAFT_LEN + 1)) + ), + m3_draft_input, + _spec_rows("m"), +) + +FAST_PATHS: Tuple[RowRange, ...] = ( + EMBED_FIRSTNORM, + SPARSE_DECODE_PRODUCER, + SPARSE_DECODE_LC8, + PRODUCER, + FUSED_SCORE_SELECT, + INDEX_SCORE, + LONG_SELECT, + MC_ALLREDUCE, + ATTN_PUSH, + MOE_PUSH, + MC_FINALIZE, + MOE_ENGINE, + ROUTER_GEMV, + ROUTER_GEMV_FP4, + ROUTER_TC, + MX_GEMV, + MX_GEMV_BF16, + DRAFT_LOWM, + SHARD_ACCEPT, + DRAFT_STEP, + DRAFT_INPUT, +) + +# -------------------------------------------------------------------------------------------------------------------- +# Selection: pure functions of the step shape +# -------------------------------------------------------------------------------------------------------------------- + + +def plan_for(num_tokens: int) -> DecodePlan: + for plan in DECODE_PLANS: + if num_tokens <= plan.max_tokens: + return plan + raise AssertionError(f"no decode plan covers {num_tokens} tokens") + + +def is_decode_step(num_tokens: int) -> bool: + """Whether a pure generation step of ``num_tokens`` rows runs the decode kernels (else generic layers).""" + return num_tokens <= DECODE_MAX_TOKENS + + +def moe_path(num_tokens: int) -> Optional[str]: + """The MoE kernel of a decode step: the engine, the unfinalized MoE + finalize all-reduce, or None (generic).""" + if MOE_ENGINE.covers(num_tokens): + return MOE_ENGINE.path + if MC_FINALIZE.covers(num_tokens): + return MC_FINALIZE.path + return None + + +def router_path(num_tokens: int) -> Optional[str]: + """The router kernel when the MoE runs unfused: the fp32 GEMV, the tensor-core router, or None (built-in gate).""" + if ROUTER_GEMV.covers(num_tokens): + return ROUTER_GEMV.path + if ROUTER_TC.covers(num_tokens): + return ROUTER_TC.path + return None + + +MX_GEMV_PATH = MX_GEMV.path +MX_GEMV_BF16_PATH = MX_GEMV_BF16.path +# FlashInfer's CuTe DSL MXFP8 GEMM on the producer's MXFP8 activation (rows the GEMV does not serve). +PREQUANT_GEMM_PATH = "flashinfer_mxfp8_gemm" +# The MXFP8 quantize (tuned backend) and FlashInfer's CuTe DSL GEMM: every row the decode kernels do not serve. +GENERIC_GEMM_PATH = "mxfp8_gemm" + + +def linear_path(n: int, k: int, num_tokens: int, prequant: bool) -> str: + """The decode-graph MXFP8 GEMM of an [N, K] projection on ``num_tokens`` rows. + + ``prequant``: the producer published the activation in MXFP8. + """ + if prequant: + if num_tokens in MX_GEMV_ROWS.get((n, k), ()): + return MX_GEMV_PATH + return PREQUANT_GEMM_PATH + if num_tokens <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0): + return MX_GEMV_BF16_PATH + return GENERIC_GEMM_PATH + + +# -------------------------------------------------------------------------------------------------------------------- +# Fallback policy +# -------------------------------------------------------------------------------------------------------------------- + + +class FastPathUnavailable(RuntimeError): + """A covered decode shape could not take its fast path under ``modeling_v2='require'``.""" + + +_lock = threading.Lock() +_require = [False] +_logged: set = set() + + +def set_require(require: bool) -> None: + """Whether a miss raises (``modeling_v2='require'``); set by the target at construction.""" + _require[0] = bool(require) + + +def require_mode() -> bool: + return _require[0] + + +def _log_once(level: Callable[[str], None], path: str, reason: str, message: str) -> None: + key = (path, reason) + with _lock: + if key in _logged: + return + _logged.add(key) + level(message) + + +def missed(path: str, rows: Optional[int], reason: str) -> None: + """``path`` covers this step but cannot run: raise under require, else log once; the caller falls back.""" + at = "" if rows is None else f" at {rows} rows" + if _require[0]: + raise FastPathUnavailable( + f"modeling_v2='require': the MiniMax-M3 decode path {path!r} covers this step{at} but cannot run: " + f"{reason}" + ) + _log_once( + logger.warning, + path, + reason, + f"[M3 dispatch] {path}{at}: {reason}; falling back (logged once per reason)", + ) + + +def capture_missed(path: str, rows: Optional[int], what: str) -> None: + """``what`` is built on eager steps but a CUDA-graph capture found it missing (see ``missed``).""" + missed(path, rows, f"{what} is not built and a CUDA graph is being captured") + + +def uncovered(path: str, reason: str) -> None: + """The step is outside ``path``'s covered shapes: log once per (path, reason); the generic path runs.""" + _log_once(logger.info, path, reason, f"[M3 dispatch] {path} not taken: {reason} (logged once)") + + +def reset_logged() -> None: + """Forget which (path, reason) pairs were logged.""" + with _lock: + _logged.clear() + + +# -------------------------------------------------------------------------------------------------------------------- +# Target constants +# -------------------------------------------------------------------------------------------------------------------- + +# The generic boundary takes the low-latency multicast all-reduce + residual + RMSNorm up to this many tokens; with +# the shared + routed add folded in, up to its own maximum (m3_mc_allreduce.MAX_TOKENS). +MC_AR_MAX_TOKENS = 8 # capacity: fixed: measured crossover vs the one-shot all-reduce +# The shared expert folds into the routed engine from one token (dedicated shared-expert CTAs). +MOE_FOLD_MIN_TOKENS = 1 # capacity: fixed: a row lower bound +# SMs left to the concurrent kernels while a persistent routed engine runs. +MOE_ENGINE_RESERVED_SMS = 24 # capacity: fixed: an SM split, not a size +# Id of the target's final hidden state in Eagle3's ``layers_to_capture``. +FINAL_HIDDEN_LAYER_ID = -1 # capacity: fixed: an Eagle3 layer id, not a size +# The trtllm-gen MoE runner's routing kind for MiniMax-M3 (sigmoid + correction bias, renormalized, scaled) and its +# gated activation (SwiGLU with the per-expert alpha / beta / clamp operands). +ROUTING_MINIMAX2 = int(RoutingMethodType.MiniMax2) +ACT_SWIGLU = int(ActType_TrtllmGen.SwiGlu) + +# -------------------------------------------------------------------------------------------------------------------- +# The runtime-value dump (TRTLLM_MODELING_V2_VALUE_DUMP=1): the metadata of every fast-path and entry decision a step +# takes, recorded on eager steps and captures (replay runs no Python), printed once at exit. +# -------------------------------------------------------------------------------------------------------------------- + +_VALUE_DUMP = os.environ.get("TRTLLM_MODELING_V2_VALUE_DUMP") == "1" +_VALUES: set = set() + + +def _record(*meta: Hashable) -> None: + if _VALUE_DUMP: + _VALUES.add(meta) + + +def _print_values() -> None: + if _VALUES: + rank = os.environ.get("OMPI_COMM_WORLD_RANK", os.environ.get("RANK", "0")) + for meta in sorted(_VALUES, key=repr): + print(f"[modeling_v2 value dump] rank {rank}: {meta}", flush=True) + + +if _VALUE_DUMP: + atexit.register(_print_values) + +# -------------------------------------------------------------------------------------------------------------------- +# Configuration: the text tower and the runtime the target assumes +# -------------------------------------------------------------------------------------------------------------------- + + +def _is_vl_config(pretrained_config: PretrainedConfig) -> bool: + model_type = getattr(pretrained_config, "model_type", None) + if model_type == "minimax_m3_vl": + return True + if model_type == "minimax_m3": + return False + return any( + "ForConditionalGeneration" in a + for a in (getattr(pretrained_config, "architectures", None) or []) + ) + + +def _text_model_config(model_config: ModelConfig) -> ModelConfig: + """The model config with the language model's ``text_config`` as its pretrained config. + + The checkpoint declares the vision-language wrapper; its ``text_config`` arrives as a plain dict (no typed + sub-config is registered), so it is materialized into a ``PretrainedConfig`` with the wrapper's dtype and + embedding tie carried over. A text config is returned unchanged. + """ + if not _is_vl_config(model_config.pretrained_config): + return model_config + cfg = copy.deepcopy(model_config) + text = getattr(cfg.pretrained_config, "text_config", None) + if text is None: + raise ValueError("MiniMax-M3 VL config has no `text_config`") + if isinstance(text, dict): + out = PretrainedConfig() + for k, v in text.items(): + setattr(out, k, v) + text = out + if ( + getattr(text, "torch_dtype", None) is None + and getattr(cfg.pretrained_config, "torch_dtype", None) is not None + ): + text.torch_dtype = cfg.pretrained_config.torch_dtype + if ( + getattr(text, "tie_word_embeddings", None) is None + and getattr(cfg.pretrained_config, "tie_word_embeddings", None) is not None + ): + text.tie_word_embeddings = cfg.pretrained_config.tie_word_embeddings + return dataclasses.replace(cfg, pretrained_config=text) + + +def _validate_runtime(model_config: ModelConfig) -> None: + """The runtime this target is written for. It is set by engine arguments rather than the checkpoint, so routing + sees only its KV cache algorithm; this check at construction refuses every other runtime. + + The MiniMax-M3 MSA backend (it owns the step metadata and the index-K cache), the fused five-way QKV/index + projection or the separate QKV and index projections, an FP8 indexer and an FP8 or NVFP4 main KV cache. + """ + sparse = model_config.sparse_attention_config + if ( + sparse is None + or getattr(sparse, "algorithm", None) != "minimax_m3" + or sparse.implementation != "msa" + ): + raise NoTarget( + "this MiniMax-M3 target requires sparse_attention_config.algorithm='minimax_m3' with implementation='msa'" + ) + if getattr(sparse, "indexer_kv_dtype", None) != "fp8": + raise NoTarget( + "this MiniMax-M3 target requires sparse_attention_config.indexer_kv_dtype='fp8'" + ) + mode = model_config.quant_config.quant_mode + if mode is None or not (mode.has_fp8_kv_cache() or mode.has_fp4_kv_cache()): + raise NoTarget("this MiniMax-M3 target requires an FP8 or NVFP4 KV cache") + + +def _sparse_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = cfg.sparse_attention_config["sparse_attention_freq"] + return [i for i, f in enumerate(freq) if int(f) != 0] + + +def _moe_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = getattr(cfg, "moe_layer_freq", None) + if freq is None: + return list(range(cfg.num_hidden_layers)) + return [i for i, f in enumerate(freq) if int(f) != 0] + + +class _Mxfp8GemmState(nn.Module): + """The MXFP8 GEMM backend state the engine's warmup drives. + + The engine finds MXFP8 projections by a module's ``quant_method`` and, for every one, opts it into native GEMM + tuning and (on decode CUDA graphs) the auto-dispatched FlashInfer GEMM, then marks each tuned. Every projection of + this target dispatches identically, so one method object carries that state for all of them; the forward reads it + at each call (``_mx_mode``) and passes the resolved backend to the GEMM entries. + """ + + def __init__(self) -> None: + super().__init__() + self.quant_method = MXFP8LinearMethod() + + +@dataclass(frozen=True) +class MxMode: + """The MXFP8 GEMM dispatch of one call, resolved from ``_Mxfp8GemmState`` and the engine's contexts.""" + + decode_graph: bool # inside a decode CUDA-graph warmup or capture + autotune: bool # inside the engine's FlashInfer MXFP8 tuning pass + tuned_quantize_gemm: ( + bool # per-bucket tuned quantize + GEMM (decode graphs, tune_decode_graph_backends) + ) + flashinfer: bool # FlashInfer's MXFP8 GEMM on the quantized activation + native_tuning: ( + bool # the native GEMM through the autotuner (warmup, before its tactics are registered) + ) + + +def _mx_mode(state: _Mxfp8GemmState) -> MxMode: + method = state.quant_method + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + autotune = flashinfer_mxfp8_autotune_active() + return MxMode( + decode_graph=decode_graph, + autotune=autotune, + tuned_quantize_gemm=method.tune_decode_graph_backends and decode_graph, + flashinfer=method.backend == "flashinfer" + or ( + method.backend == "auto" + and (autotune or (method._flashinfer_autotuned and decode_graph)) + ), + native_tuning=method.needs_native_autotune, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# Projections: the MXFP8 GEMM dispatch (``linear_path``) and its hand-offs +# -------------------------------------------------------------------------------------------------------------------- + + +def _mxfp8_linear( + core, x: torch.Tensor, w: torch.Tensor, s: torch.Tensor, mode: MxMode +) -> torch.Tensor: + """One MXFP8 projection on the engine-tuned path: MXFP8 quantize + GEMM, backend per ``mode``.""" + _record( + "m3_mxfp8_linear", + tuple(w.shape), + mode.tuned_quantize_gemm, + mode.flashinfer, + mode.native_tuning, + ) + return m3_mxfp8_linear( + x.contiguous(), + w, + s, + core.mx_alpha, + torch.bfloat16, + mode.tuned_quantize_gemm, + mode.autotune, + mode.flashinfer, + mode.native_tuning, + ) + + +def _gemv_variants(n: int, k: int, push: Optional[str], world: int): + """Every (rows, bf16 activation, push (world, bucket tokens, mode)) call the decode dispatch routes to an [N, K] + projection's GEMV: an MXFP8 activation on the rows ``MX_GEMV_ROWS`` gives it, a bf16 one up to + ``MX_GEMV_BF16_MAX_TOKENS``; each pushing when the projection's boundary push covers the rows.""" + push_rows = {"attn": ATTN_PUSH.keys, "moe": MOE_PUSH.keys}.get(push, frozenset()) + variants = [] + for rows in range(1, DECODE_MAX_TOKENS + 1): + for x_bf16, served in ( + (False, rows in MX_GEMV_ROWS.get((n, k), ())), + (True, rows <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0)), + ): + if not served: + continue + variants.append((rows, x_bf16, None)) + if rows in push_rows: + bucket = next(b for b in m3_mc_allreduce.BUCKETS if rows <= b) + variants.append((rows, x_bf16, (world, bucket, AR_PUSH_MODE))) + return variants + + +def _uncompiled_note() -> str: + key = mx_gemv.last_uncompiled + return "" if key is None else f" (its kernel {key} was not compiled before warmup)" + + +def _linear( + core, x: torch.Tensor, wk: str, sk: str, push: Optional[tuple] = None +) -> Tuple[torch.Tensor, bool]: + """One MXFP8 projection ``x @ w[wk]^T``; returns (output, whether its GEMV pushed into ``push``). + + Outside decode CUDA graphs: the engine-tuned MXFP8 quantize + GEMM. Inside them the projection's (N, K) and the + step's rows pick one path (``linear_path``): + + 1. an activation its producer already quantized to MXFP8 (``m3_prequant.SLOT``, taken here): the tcgen05 decode + GEMV on the rows it serves, else the decode-graph MXFP8 GEMM on that activation (AutoTuner-selected backend + per token bucket); + 2. a bf16 activation on the rows the GEMV serves quantizing in its prologue; + 3. otherwise the engine-tuned MXFP8 quantize + GEMM. + + A kernel declining rows its path covers is a miss; the chain then continues down the list. ``push`` (an + ``m3_mx_gemv`` ``ar`` tuple) is the boundary all-reduce exchange the GEMV writes its rows into. + """ + w, s = core.w[wk], core.w[sk] + x = x.contiguous() + mode = _mx_mode(core.mxfp8) + if not mode.decode_graph: + return _mxfp8_linear(core, x, w, s, mode), False + use_pdl = get_env_enable_pdl() + rows = x.shape[0] + n, k = w.shape + prequant = m3_prequant.SLOT.take(x) + if prequant is None and x.dtype != torch.bfloat16: + path = GENERIC_GEMM_PATH + else: + path = linear_path(n, k, rows, prequant is not None) + _record("linear", (n, k), rows, path, push is not None) + output = None + if prequant is not None: + act, act_scale = prequant + if path == MX_GEMV_PATH: + output = m3_mx_gemv(act, act_scale, w, s, True, None, use_pdl, push) + if output is not None: + return output, push is not None and mx_gemv.last_pushed + missed( + path, + rows, + f"the GEMV declined the [{n}, {k}] projection's MXFP8 activation" + + _uncompiled_note(), + ) + if path == PREQUANT_GEMM_PATH: + uncovered(MX_GEMV_PATH, f"the [{n}, {k}] projection at rows the GEMV does not serve") + return m3_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16, mode.autotune), False + if path == MX_GEMV_BF16_PATH: + output = m3_mx_gemv_bf16(x, w, s, True, None, use_pdl, push) + if output is not None: + return output, push is not None and mx_gemv.last_pushed + missed( + path, rows, f"the GEMV declined the [{n}, {k}] projection's input" + _uncompiled_note() + ) + if path == GENERIC_GEMM_PATH: + uncovered( + MX_GEMV_BF16_PATH, + f"the [{n}, {k}] projection's bf16 input beyond the decode GEMV's rows", + ) + return _mxfp8_linear(core, x, w, s, mode), False + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer boundaries: all-reduce + residual + RMSNorm +# -------------------------------------------------------------------------------------------------------------------- + + +def _can_mc(x, folded: bool = False) -> bool: + """Whether the generic boundary takes the multicast all-reduce + residual + RMSNorm.""" + return ( + isinstance(x, torch.Tensor) + and x.dim() == 2 + and x.shape[0] <= (m3_mc_allreduce.MAX_TOKENS if folded else MC_AR_MAX_TOKENS) + and x.dtype == torch.bfloat16 + and x.is_contiguous() + ) + + +def _mc_allreduce_norm(core, x, x2, residual, norm_w, *, quant=False, prepushed=0): + """The multicast all-reduce + optional second addend + residual + RMSNorm (``m3_mc_allreduce_norm``), or None + without a workspace or for operands the kernel does not take (the same on every rank). With ``quant`` the normed + rows are also published, MXFP8-quantized, to ``m3_prequant.SLOT`` for the next MXFP8 projection.""" + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + if workspace is None: + return None + _record("m3_mc_allreduce_norm", x.shape[0], x2 is not None, quant, prepushed) + return m3_mc_allreduce.m3_mc_allreduce_norm( + x, + residual.contiguous(), + norm_w, + core.rms_eps, + workspace, + use_pdl=get_env_enable_pdl(), + x2=x2, + quant=quant, + prepushed=prepushed, + ) + + +def _native_allreduce_norm(core, x, residual, norm_w): + """The built-in fused all-reduce + residual + RMSNorm.""" + _record("m3_allreduce", x.shape[0]) + return m3_allreduce( + x, + core.allreduce, + AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=norm_w, + eps=core.rms_eps, + trigger_completion_at_end=False, + ), + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The MoE: router, routed experts (trtllm-gen NVFP4), shared expert +# -------------------------------------------------------------------------------------------------------------------- + + +@dataclass +class _MoeScalars: + """The routed experts' kernel scalars of one layer, derived from the checkpoint's per-tensor scales once the + weights are real (``MiniMaxM3Tp8ModelingV2._derive_moe``).""" + + fc31_input_scale: torch.Tensor + fc2_input_scale: torch.Tensor + fc31_alpha: torch.Tensor + fc2_alpha: torch.Tensor + fc31_scale_c: torch.Tensor + act_alpha: torch.Tensor + act_beta: torch.Tensor + act_clamp: torch.Tensor + + +def _router_logits(core, i: int, h: torch.Tensor) -> torch.Tensor: + """The fp32 router logits of MoE layer ``i``: the fp32 GEMV at one row, the tensor-core router up to 48, else the + fp32 linear.""" + rows = h.shape[0] if h.dim() == 2 else 0 + path = router_path(rows) + w = core.w[f"l{i}_router"] + if path == ROUTER_GEMV.path: + logits = _router_gemv(h, w, use_pdl=get_env_enable_pdl()) + if logits is not None: + return logits + missed(path, rows, "the router GEMV declined the operands") + elif path == ROUTER_TC.path: + fused = _router_tc(h, w, use_pdl=get_env_enable_pdl()) + if fused is not None: + return fused[0] + missed(path, rows, "the tensor-core router declined the operands") + else: + uncovered(ROUTER_TC.path, f"more than {ROUTER_TC_MAX_TOKENS} rows") + return m3_fp32_linear(h, w) + + +def _router_fp4(core, i: int, h: torch.Tensor): + """Router logits and the routed experts' NVFP4 activation from one kernel, or None (then ``_router_logits``).""" + if h.dim() != 2 or h.shape[0] > ROUTER_TC_MAX_TOKENS: + return None + rows = h.shape[0] + path = router_path(rows) + if path == ROUTER_GEMV.path: + path = ROUTER_GEMV_FP4.path + scale = core.moe[i].fc31_input_scale + if path == ROUTER_TC.path: + fused = _router_tc(h, core.w[f"l{i}_router"], scale, use_pdl=get_env_enable_pdl()) + else: + fused = _router_gemv_fp4(h, core.w[f"l{i}_router"], scale, use_pdl=get_env_enable_pdl()) + if fused is None: + missed(path, rows, "the router declined the operands") + return fused + + +def _routed_experts(core, i: int, x, x_sf, router_logits, do_finalize: bool): + """The trtllm-gen NVFP4 MoE runner over this rank's tensor-parallel slice of every expert, routing in-kernel. + + ``x_sf`` None means ``x`` is the bf16 activation, quantized here with the experts' global input scale. Returns + the finalized rows, or the unfinalized ``[gemm2 out, expert scales, expanded->permuted index]``. + """ + moe = core.moe[i] + if x_sf is None: + x, x_sf = fp4_quantize(x, moe.fc31_input_scale, 16, False, False) + rows = x.shape[0] + x_sf = x_sf.view(rows, -1).flatten() + _record("fp4_block_scale_moe_runner", rows, do_finalize) + outs = fp4_block_scale_moe_runner( + router_logits, + core.w[f"l{i}_router_bias"], + x, + x_sf.view(torch.float8_e4m3fn), + core.w[f"l{i}_e_w31"], + core.w[f"l{i}_e_s31"].view(torch.float8_e4m3fn), + None, + moe.act_alpha, + moe.act_beta, + moe.act_clamp, + core.w[f"l{i}_e_w2"], + core.w[f"l{i}_e_s2"].view(torch.float8_e4m3fn), + None, + moe.fc31_scale_c, + moe.fc31_alpha, + moe.fc2_alpha, + core.num_experts, + core.top_k, + None, + None, + core.moe_inter, + 0, + core.num_experts, + core.routed_scaling_factor, + ROUTING_MINIMAX2, + do_finalize, + act_type=ACT_SWIGLU, + topk_weights=None, + topk_ids=None, + output=None, + tune_max_num_tokens=core.max_num_tokens, + use_dp=False, + ) + return outs[0] if do_finalize else outs + + +def _split_chunk(rows: int, num_chunks: int) -> List[int]: + """The MoE scheduler's even split: the first ``rows % num_chunks`` chunks take one more row.""" + base, extra = divmod(rows, num_chunks) + return [base + 1] * extra + [base] * (num_chunks - extra) + + +# -------------------------------------------------------------------------------------------------------------------- +# Attention +# -------------------------------------------------------------------------------------------------------------------- + + +def _nvfp4_layer(manager, i: int) -> bool: + return bool(getattr(manager, "is_nvfp4_layer", lambda _layer_idx: False)(i)) + + +def _producer_covered(core) -> bool: + """Whether ``PRODUCER`` covers this target's sparse layers: an NVFP4 KV cache at the TP8 heads (8 query, one KV + and one index head). A covered layer that cannot take ``m3_producer`` is a miss.""" + return core.main_kv_is_nvfp4 and (core.heads, core.kv_heads, core.index_heads) == (8, 1, 1) + + +def _producer_layout(kv_data, kv_scale) -> bool: + """Whether a sparse layer's NVFP4 pools are the layout ``m3_producer`` is built for: [pages, 2, 1, 128, 64] with + its scale pool [pages, 2, 1, 128, 8].""" + return ( + kv_data is not None + and kv_scale is not None + and kv_data.dim() == 5 + and tuple(kv_data.shape[1:]) == (2, 1, 128, 64) + and kv_data.stride()[2:] == (128 * 64, 64, 1) + and kv_scale.dim() == 5 + and tuple(kv_scale.shape[1:]) == (2, 1, 128, 8) + and kv_scale.stride()[2:] == (128 * 8, 8, 1) + and kv_scale.shape[0] == kv_data.shape[0] + ) + + +def _nvfp4_fused_write_layout(data: torch.Tensor, scales: torch.Tensor, num_heads: int) -> bool: + """Reject invalid HND pools; select the fused producer only for its packed layout.""" + if ( + data.dim() != 5 + or scales.dim() != 5 + or data.shape[0] <= 0 + or data.shape[0] != scales.shape[0] + or tuple(data.shape[1:]) != (2, num_heads, 128, 64) + or tuple(scales.shape[1:]) != (2, num_heads, 128, 8) + or data.dtype not in (torch.int8, torch.uint8) + or scales.dtype != torch.uint8 + ): + raise ValueError("MiniMax-M3 NVFP4 data and scale caches must have matching HND geometry") + for cache in (data, scales): + if ( + cache.stride(4) != 1 + or cache.stride(3) != cache.shape[4] + or cache.stride(2) < 128 * cache.stride(3) + or cache.stride(1) < num_heads * cache.stride(2) + or cache.stride(0) < 2 * cache.stride(1) + ): + raise ValueError( + "MiniMax-M3 NVFP4 cache rows must be packed and heads/planes/pages must not overlap" + ) + if scales.stride(2) != 128 * 8: + raise ValueError("MiniMax-M3 NVFP4 MSA scale cache must have packed head blocks") + return data.stride(2) == 128 * 64 and data.stride(0) % 2 == 0 and data.stride(1) % 2 == 0 + + +def _producer_miss(covered: bool, rows: Optional[int], reason: str) -> None: + if covered: + missed(PRODUCER.path, rows, reason) + return None + + +def _sparse_qkv_producer( + core, i: int, packed, position_ids, kv_cache_manager, msa_idx_k_cache, msa_out_cache_loc +): + """The horizontal producer of a sparse layer: Q / index-Q (FP8) out, K / V / index-K written into the paged + caches (NVFP4 or FP8 main cache), for every step type (token-major). An NVFP4 cache at the TP8 heads takes + ``m3_producer`` (``PRODUCER``); a covered layer that cannot is a miss. Returns (q, index_q), or None for caches + neither kernel takes. + + ``kv_cache_manager``, ``msa_idx_k_cache`` (the per-layer index-K cache getter) and ``msa_out_cache_loc`` are the + calling target's ``step_args`` entries.""" + covered = _producer_covered(core) + rows = int(packed.shape[0]) + if packed.dtype != torch.bfloat16 or position_ids is None: + return _producer_miss( + covered, + rows, + "the packed QKV dtype, position ids, head dims, norm or RoPE configuration", + ) + manager = kv_cache_manager + if manager is None: + return _producer_miss(covered, rows, "the step has no KV cache manager") + buffers = manager.get_buffers(i, kv_layout="HND") + index_k_cache = msa_idx_k_cache(i) + out_cache_loc = msa_out_cache_loc + layer_uses_nvfp4 = _nvfp4_layer(manager, i) + scale_buffers = ( + manager.get_block_scale_buffers(i, kv_layout="HND") if layer_uses_nvfp4 else None + ) + kv_quant_scale = core.inv_kv_scales[i] if layer_uses_nvfp4 else None + if layer_uses_nvfp4: + supported_main = ( + buffers is not None + and scale_buffers is not None + and _nvfp4_fused_write_layout(buffers, scale_buffers, core.kv_heads) + and buffers.is_cuda + and scale_buffers.is_cuda + and scale_buffers.device == buffers.device + and kv_quant_scale.is_cuda + and kv_quant_scale.dtype == torch.float32 + and kv_quant_scale.numel() >= 3 + and kv_quant_scale.is_contiguous() + ) + else: + supported_main = ( + buffers is not None + and buffers.is_cuda + and buffers.dtype == torch.float8_e4m3fn + and buffers.dim() == 5 + and tuple(buffers.shape[1:]) == (2, core.kv_heads, 128, 128) + and buffers.stride(4) == 1 + and buffers.stride(3) == 128 + and buffers.stride(2) == 128 * 128 + and buffers.stride(1) >= core.kv_heads * buffers.stride(2) + and buffers.stride(0) >= 2 * buffers.stride(1) + and buffers.stride(0) % 4 == 0 + and buffers.stride(1) % 4 == 0 + ) + supported_index = ( + index_k_cache.is_cuda + and index_k_cache.dtype == torch.float8_e4m3fn + and index_k_cache.dim() == 4 + and tuple(index_k_cache.shape[1:]) == (1, 128, 128) + and index_k_cache.stride(3) == 1 + and index_k_cache.stride(2) == 128 + and index_k_cache.stride(1) >= 128 * 128 + and index_k_cache.stride(0) >= index_k_cache.stride(1) + and index_k_cache.stride(0) % 4 == 0 + and index_k_cache.stride(1) % 4 == 0 + and buffers is not None + and index_k_cache.shape[0] == buffers.shape[0] + ) + cos_sin = core.rotary_cos_sin + supported_rope = ( + cos_sin.is_cuda + and cos_sin.dtype == torch.float32 + and cos_sin.is_contiguous() + and cos_sin.dim() == 3 + and tuple(cos_sin.shape[1:]) == (2, 32) + ) + if ( + not supported_main + or not supported_index + or not supported_rope + or out_cache_loc is None + or not out_cache_loc.is_cuda + or out_cache_loc.dtype != torch.int32 + or not out_cache_loc.is_contiguous() + or out_cache_loc.numel() < rows + ): + return _producer_miss( + covered and layer_uses_nvfp4, + rows, + "the layer's NVFP4 / index caches, RoPE table, KV quant scale or cache slots are not the layout the " + "producer is built for", + ) + w = core.w + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + if ( + layer_uses_nvfp4 + and (core.heads, core.kv_heads, core.index_heads) == (8, 1, 1) + and _producer_layout(buffers, scale_buffers) + ): + _record("m3_producer", rows) + q, index_q = m3_producer( + packed.contiguous(), + buffers.view(torch.uint8), + scale_buffers.view(torch.uint8), + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + core.heads, + get_env_enable_pdl(), + PRODUCER_EARLY_TRIGGER, + ) + else: + if layer_uses_nvfp4 and covered: + missed( + PRODUCER.path, + rows, + f"the NVFP4 pool {tuple(buffers.shape)} / scale pool {tuple(scale_buffers.shape)} is not the " + "producer's layout", + ) + _record("m3_qkv_insert", rows, layer_uses_nvfp4) + q, index_q = m3_qkv_insert( + packed.contiguous(), + buffers.view(torch.uint8) if layer_uses_nvfp4 else buffers, + scale_buffers, + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.heads, + core.kv_heads, + core.index_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + ) + return q.flatten(1), index_q.flatten(1) + + +def _qk_norm_rope(core, i: int, qkv: torch.Tensor, position_ids, out_fp8: bool) -> torch.Tensor: + """The per-head Gemma RMSNorm of Q and K and the partial NeoX RoPE of a dense layer's packed [Q | K | V], in one + kernel: in place, or an FP8 copy (Q / K normed and roped, V cast) when ``out_fp8``.""" + qkv = qkv.contiguous() + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + args = ( + qkv, + core.heads, + core.kv_heads, + core.kv_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_q_norm"], + core.w[f"l{i}_k_norm"], + core.rope_theta, + True, + pos, + 1.0, # factor: no YARN + 0.0, # low + 0.0, # high + 1.0, # attention_factor + True, # is_qk_norm + True, # use_gemma + False, # use_mrope + 0, + 0, + ) + _record("qk_norm_rope", qkv.shape[0], out_fp8) + if out_fp8: + return m3_qk_norm_rope_fp8(*args) + return fused_qk_norm_rope(*args) + + +def _fp8_index_qk_norm_rope( + core, i: int, idx_qk: torch.Tensor, position_ids, msa_idx_k_cache, msa_out_cache_loc +) -> torch.Tensor: + """The index branch of an unfused sparse layer: index-Q in FP8 out, index-K into the FP8 index-K cache, normed and + roped in one kernel (the built-in ``_fused_fp8_index_qk_norm_rope``). ``msa_idx_k_cache`` / ``msa_out_cache_loc`` + are the calling target's ``step_args`` entries.""" + if position_ids is None or idx_qk.dtype != torch.bfloat16: + raise NotImplementedError( + "MiniMax-M3 fused FP8 indexer requires BF16 activations and position_ids." + ) + index_k_cache = msa_idx_k_cache(i) + if index_k_cache.dtype != torch.float8_e4m3fn: + raise ValueError( + "MiniMax-M3 fused FP8 indexer requires an E4M3 index-K cache, " + f"got {index_k_cache.dtype}." + ) + rows = int(idx_qk.shape[0]) + _record("m3_fp8_indexer_qk_norm_rope", rows) + return m3_fp8_indexer_qk_norm_rope( + idx_qk.contiguous(), + index_k_cache, + msa_out_cache_loc[:rows], + core.index_heads, + core.index_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_idx_q_norm"], + core.w[f"l{i}_idx_k_norm"], + core.rope_theta, + position_ids.reshape(-1).contiguous().to(torch.int32), + ).flatten(1) + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer helpers both targets' bodies call: the Eagle3 capture hook and the norm after a layer +# -------------------------------------------------------------------------------------------------------------------- + + +def _plain_capture(spec_metadata, layer_id: int, hidden_states, residual) -> None: + """The in-tree Eagle3 hidden-state capture.""" + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + + +def _next_norm(core, i: int) -> torch.Tensor: + """The norm the boundary after layer ``i`` applies: the next layer's input norm, the final norm after the last.""" + return core.w[f"l{i + 1}_in_norm"] if i + 1 < core.num_layers else core.w["final_norm"] + + +# -------------------------------------------------------------------------------------------------------------------- +# The core +# -------------------------------------------------------------------------------------------------------------------- + +# Engine states (compiled kernel + workspace) per (device, token count[, shared fold]), shared by every layer: a step +# runs the layers one after another on one stream. +_ENGINE_STATES: Dict[tuple, object] = {} + + +class MiniMaxM3Tp8ModelingV2(DecoderModel): + """The MiniMax-M3 text model of this target: its weights, its runtime resources, and the per-step dispatch.""" + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + _validate_runtime(model_config) + super().__init__(model_config) + cfg = model_config.pretrained_config + mapping = model_config.mapping + self.mapping = mapping + self.vocab_size = cfg.vocab_size + d = _weights.MODEL_WEIGHTS.dims(self) + self.num_layers = d.num_layers + self.hidden = d.hidden + self.heads = d.heads + self.kv_heads = d.kv_heads + self.head_dim = d.head_dim + self.index_heads = d.index_heads + self.index_dim = d.index_dim + self.q_size = d.heads * d.head_dim + self.kv_size = d.kv_heads * d.head_dim + self.moe_inter = d.moe_inter + self.num_experts = d.num_experts + self.top_k = int(cfg.num_experts_per_tok) + self.routed_scaling_factor = float(getattr(cfg, "routed_scaling_factor", 1.0)) + self.rms_eps = float(cfg.rms_norm_eps) + self.swiglu_alpha = float(getattr(cfg, "swiglu_alpha", 1.702)) + self.swiglu_beta = 1.0 # the (up + 1) offset of SwiGLU-OAI + self.swiglu_limit = float(getattr(cfg, "swiglu_limit", 7.0)) + self.max_num_tokens = model_config.max_num_tokens + self.moe_max_num_tokens = model_config.moe_max_num_tokens + self.sparse_layers = frozenset(_sparse_layer_ids(cfg)) + self.moe_layers = frozenset(_moe_layer_ids(cfg)) + mode = model_config.quant_config.quant_mode + self.main_kv_is_fp8 = bool(mode is not None and mode.has_fp8_kv_cache()) + self.main_kv_is_nvfp4 = bool(mode is not None and mode.has_fp4_kv_cache()) + # The sparse layers' QKV and index projections: one five-way packed projection, or the QKV projection and a + # separate BF16 index projection (fuse_qkv_index_projection=False) run on two streams. + self.fused_projection = bool( + getattr(model_config.sparse_attention_config, "fuse_qkv_index_projection", False) + ) + + self.w = _weights.MODEL_WEIGHTS.declare(d) + + # The engine reads these two off the module tree: the MXFP8 GEMM state it tunes, and the opt-in to the + # automatic FlashInfer MXFP8 path on decode graphs. + self.mxfp8 = _Mxfp8GemmState() + self._use_flashinfer_mxfp8_decode_graph_default = True + + # The built-in fused all-reduce + residual + RMSNorm of the layer boundaries the multicast kernel does not take. + self.allreduce = AllReduce( + mapping=mapping, strategy=model_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + + # Streams: the shared / routed expert overlap and, when the MoE chunks its input, the chunking overlap. + self.aux_stream_moe = torch.cuda.Stream() + chunked = self.moe_max_num_tokens < model_config.max_num_tokens * mapping.dp_size + self.aux_stream_moe_chunk = torch.cuda.Stream() if chunked else None + self.moe_events = {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers} + self.moe_chunk_events = { + i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers + } + # The unfused sparse layers' index branch overlaps the QKV branch on its own stream. + self.aux_stream_attn = None if self.fused_projection else torch.cuda.Stream() + self.attn_events = ( + {} + if self.fused_projection + else {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.sparse_layers} + ) + + # Partial NeoX RoPE (rotary_dim of head_dim): the cos/sin table every attention kernel reads. + rope = RopeParams.from_config(cfg) + self.rope_theta = rope.theta + self.rotary_dim = int(rope.dim) + self.rotary_cos_sin = rope.create_rope_const_params(interleave=False)[1].reshape( + rope.max_positions, 2, -1 + ) + + # One MSA attention backend per layer (dense and sparse; the step metadata is the backend's). + quant_config = model_config.get_quant_config() + self.attn: List[MiniMaxM3MsaSparseAttention] = [] + for i in range(self.num_layers): + backend = create_attention( + model_config.attn_backend, + i, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=None, + quant_config=quant_config, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + q_scaling=1.0, + sparse_params=model_config.sparse_attention_config.to_sparse_params( + pretrained_config=cfg, layer_idx=i + ), + ) + backend.update_quant_config(quant_config) + if not isinstance(backend, MiniMaxM3MsaSparseAttention): + raise ValueError( + f"layer {i}: expected the MiniMax-M3 MSA backend, got {type(backend).__name__}" + ) + self.attn.append(backend) + + # Built once the weights are real (``post_load_weights``): the KV (Q, K, V) scales, the MoE scalars, the + # decode engines and workspaces, the two targets. + self.kv_scales: Dict[int, torch.Tensor] = {} + self.inv_kv_scales: Dict[int, torch.Tensor] = {} + self.moe: Dict[int, _MoeScalars] = {} + self.mx_alpha: Optional[torch.Tensor] = None + self.moe_engines: Dict[int, Dict[Tuple[int, bool, bool], Optional[object]]] = {} + self.embed_owners: Dict[tuple, DecodeEmbedding] = {} + self._built_for: Optional[tuple] = None + self._prefill: Optional[PrefillTarget] = None + self._decode: Optional[DecodeTarget] = None + + def sm_scale(self, i: int) -> float: + return (self.head_dim**-0.5) / float(self.attn[i].q_scaling) + + # ---------------------------------------------------------------------------------------------- after the load + + def post_load_weights(self) -> None: + """Derive the kernels' state from the loaded weights, then build both targets and every decode kernel. + + Idempotent per set of weight storages: the loader's post-load walk reaches this module directly and through + the shell, and a second call on the same storages does nothing. + """ + identity = ( + self.w["embed"].data_ptr(), + self.w["l0_in_norm"].data_ptr(), + self.w["final_norm"].data_ptr(), + ) + if self._built_for == identity: + return + device = self.w["embed"].device + derived = nn.Module() + derived.register_buffer( + "mx_alpha", torch.ones(1, dtype=torch.float32, device=device), persistent=False + ) + for i in range(self.num_layers): + # No KV calibration in the checkpoint: unit (Q, K, V) scales and their reciprocals. + derived.register_buffer( + f"l{i}_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + derived.register_buffer( + f"l{i}_inv_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + for i in sorted(self.moe_layers): + for name, value in dataclasses.asdict(self._derive_moe(i)).items(): + derived.register_buffer(f"l{i}_{name}", value, persistent=False) + self.derived = derived + self.mx_alpha = derived.mx_alpha + self.kv_scales = {i: getattr(derived, f"l{i}_kv_scales") for i in range(self.num_layers)} + self.inv_kv_scales = { + i: getattr(derived, f"l{i}_inv_kv_scales") for i in range(self.num_layers) + } + self.moe = { + i: _MoeScalars( + **{ + f.name: getattr(derived, f"l{i}_{f.name}") + for f in dataclasses.fields(_MoeScalars) + } + ) + for i in sorted(self.moe_layers) + } + self._build_embed_owners() + self.prebuild_kernels() + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + self._built_for = identity + + def cache_derived_state(self) -> None: + """The staged loaders' entry to the same build (``post_load_weights``).""" + self.post_load_weights() + + def _derive_moe(self, i: int) -> _MoeScalars: + """The routed experts' kernel scalars from the checkpoint's per-tensor scales (trtllm-gen NVFP4 layout). + + Global input scales: the reciprocal of the largest per-expert input scale (gate and up must agree). Per expert: + ``alpha = 1 / (input_scale * (1 / weight_scale_2))`` for the gate (fc31_alpha), the up half (its GEMM2-input + rescale ``fc31_scale_c = fc2_input_scale * up_alpha``) and the down projection; SwiGLU-OAI's activation + constants per expert, with beta and the clamp in the pre-dequant GEMM output domain (divided by fc31_alpha). + """ + w = self.w + isc1, isc3, isc2 = w[f"l{i}_e_isc1"].data, w[f"l{i}_e_isc3"].data, w[f"l{i}_e_isc2"].data + if not torch.allclose(isc1, isc3): + raise ValueError(f"layer {i}: the routed experts' w1 and w3 input scales differ") + fc31_input_scale = isc1.max().reciprocal() + fc2_input_scale = isc2.max().reciprocal() + ws2_1, ws2_3, ws2_2 = ( + w[f"l{i}_e_ws2_1"].data, + w[f"l{i}_e_ws2_3"].data, + w[f"l{i}_e_ws2_2"].data, + ) + fc31_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_1)) + up_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_3)) + fc2_alpha = 1.0 / (fc2_input_scale * (1.0 / ws2_2)) + fc31_scale_c = fc2_input_scale * up_alpha + n = self.num_experts + device = fc31_alpha.device + act_alpha = torch.full((n,), self.swiglu_alpha, dtype=torch.float32, device=device) + act_beta = torch.full((n,), self.swiglu_beta, dtype=torch.float32, device=device) + act_clamp = torch.full((n,), self.swiglu_limit, dtype=torch.float32, device=device) + act_beta.div_(fc31_alpha) + act_clamp.div_(fc31_alpha) + return _MoeScalars( + fc31_input_scale=fc31_input_scale, + fc2_input_scale=fc2_input_scale, + fc31_alpha=fc31_alpha, + fc2_alpha=fc2_alpha, + fc31_scale_c=fc31_scale_c, + act_alpha=act_alpha, + act_beta=act_beta, + act_clamp=act_clamp, + ) + + def _build_embed_owners(self) -> None: + """The fused embedding + first RMSNorm's caller-owned buffers, one per (rows, id dtype) of a decode step.""" + table, weight = self.w["embed"], self.w["l0_in_norm"] + self.embed_owners = { + (rows, dtype): DecodeEmbedding(table, weight, self.rms_eps, rows, dtype) + for rows in range(1, EMBED_FIRSTNORM_MAX_TOKENS + 1) + for dtype in (torch.int32, torch.int64) + } + + def projections(self) -> List[Tuple[str, str, Optional[str]]]: + """Every MXFP8 projection in layer order: (weight key, scale key, the boundary its GEMV pushes into).""" + out = [] + for i in range(self.num_layers): + out.append((f"l{i}_qkv_w", f"l{i}_qkv_s", None)) + out.append((f"l{i}_o_w", f"l{i}_o_s", "attn")) + if i in self.moe_layers: + out.append((f"l{i}_sh_gu_w", f"l{i}_sh_gu_s", None)) + out.append((f"l{i}_sh_dn_w", f"l{i}_sh_dn_s", "moe")) + else: + out.append((f"l{i}_mlp_gu_w", f"l{i}_mlp_gu_s", None)) + out.append((f"l{i}_mlp_dn_w", f"l{i}_mlp_dn_s", None)) + return out + + def decode_query_lens(self) -> List[int]: + """The query rows per request a decode step can have: 1, and 1 + each draft length speculation runs.""" + lens = {1} + spec = getattr(self.model_config, "spec_config", None) + if spec is not None: + for drafts in ( + getattr(spec, "max_draft_len", None), + getattr(spec, "max_total_draft_tokens", None), + *(getattr(spec, "draft_len_schedule", None) or {}).values(), + ): + if drafts: + lens.add(1 + int(drafts)) + return sorted(lens) + + def prebuild_kernels(self) -> Dict[str, object]: + """Compile (and load) every kernel a decode forward dispatches to that would otherwise compile on its first + call: each projection's decode GEMVs (every row count, activation and boundary push), the MoE engines of + ``MOE_ENGINE``'s token counts, and the index scorer at the decode query lengths. Runs on the + current (main) stream before the executor's warmup, in the same order on every rank, then synchronizes; after + it, GEMV compiles are frozen and the forward only looks kernels up (a missing one is a miss). Returns the + seconds each part took and the scorer keys.""" + times: Dict[str, object] = {} + if not self.mxfp8.quant_method.use_cutlass: + raise ValueError("the MiniMax-M3 decode dispatch needs the CUTLASS MXFP8 GEMM") + mx_gemv.freeze(False) + use_pdl = get_env_enable_pdl() + world = int(self.mapping.tp_size) + start = time.perf_counter() + for wk, sk, push in self.projections(): + weight, scale = self.w[wk], self.w[sk] + mx_gemv.ensure_compiled(weight, scale, use_pdl) + n, k = weight.shape + for rows, x_bf16, ar in _gemv_variants(n, k, push, world): + mx_gemv.precompile(weight, scale, rows, x_bf16, use_pdl=use_pdl, ar=ar) + torch.cuda.synchronize() + times["gemv_s"] = time.perf_counter() - start + start = time.perf_counter() + for i in sorted(self.moe_layers): + self._prebuild_engines(i) + torch.cuda.synchronize() + times["moe_s"] = time.perf_counter() - start + start = time.perf_counter() + times["scorer_keys"] = self._prebuild_index_scorer( + self.decode_query_lens(), self.w["embed"].device + ) + torch.cuda.synchronize() + times["scorer_s"] = time.perf_counter() - start + mx_gemv.freeze(True) + logger.info( + "[M3 prebuild] decode kernels built before warmup: GEMVs %.1f s, MoE engines %.1f s, index scorer " + "%.1f s (%d variants)" + % (times["gemv_s"], times["moe_s"], times["scorer_s"], len(times["scorer_keys"])) + ) + return times + + # ------------------------------------------------------------------------------------------ the MoE engines + + def _expert_operands(self, i: int): + """The routed experts' operands the engines read (trtllm-gen NVFP4 SwiGLU-OAI layout), and the input scale.""" + w, m = self.w, self.moe[i] + tensors = [ + w[f"l{i}_e_w31"].data, + w[f"l{i}_e_s31"].data, + w[f"l{i}_e_w2"].data, + w[f"l{i}_e_s2"].data, + m.fc31_alpha, + m.fc31_scale_c, + m.act_alpha, + m.act_beta, + m.act_clamp, + m.fc2_alpha, + ] + return tensors, m.fc31_input_scale + + def _shared_weights(self, i: int): + w = self.w + return ( + w[f"l{i}_sh_gu_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_gu_s"].data, + w[f"l{i}_sh_dn_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_dn_s"].data, + ) + + def engine(self, i: int, num_tokens: int, shared: bool) -> Optional[object]: + """Layer ``i``'s prebuilt routed-expert engine for ``num_tokens`` (1..8), router folded in (and the shared + expert with ``shared``). None after a miss: the build failed or was never made (a forward never builds one).""" + return self._lookup( + i, self.moe_engines[i], (num_tokens, True, shared), MOE_ENGINE.path, num_tokens + ) + + def _lookup(self, i: int, engines: dict, key, path: str, num_tokens: int) -> Optional[object]: + if key not in engines: + what = f"layer {i}'s engine" + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, what) + else: + missed(path, num_tokens, f"{what} was not built before warmup") + return None + engine = engines[key] + if engine is None: + missed(path, num_tokens, "its build failed (logged)") + return engine + + def _prebuild_engines(self, i: int) -> None: + """Build (compile, on first use of each shared state, and warm) every engine a decode step of layer ``i`` can + take: ``MOE_ENGINE``'s token counts with the shared expert folded in. Whether an + engine is kept is agreed across the TP group (one host collective per engine): a build that fails on one rank + drops that engine on every rank, so the ranks never mix an engine (and its boundary push) with the generic + MoE.""" + engines = self.moe_engines.setdefault(i, {}) + for num_tokens in sorted(MOE_ENGINE.keys): + key = (num_tokens, True, True) + if key not in engines: + engines[key] = m3_mc_allreduce.tp_agree_built( + self.mapping, + self._try_build( + i, MOE_ENGINE.path, num_tokens, self._build_engine, shared=True + ), + f"[M3 dispatch] {MOE_ENGINE.path} at {num_tokens} rows: layer {i}'s engine", + ) + + def _try_build(self, i: int, path: str, num_tokens: int, build, **kwargs) -> Optional[object]: + try: + return build(i, num_tokens, **kwargs) + except (ImportError, RuntimeError, ValueError) as err: + logger.warning( + f"[M3 dispatch] {path} at {num_tokens} rows: layer {i}'s engine failed to build: {err}" + ) + return None + + def _build_engine(self, i: int, num_tokens: int, shared: bool) -> object: + tensors, scale = self._expert_operands(i) + device = tensors[0].device + sms = torch.cuda.get_device_properties(device).multi_processor_count + kwargs = {} + if shared: + kwargs["shared_weights"] = self._shared_weights(i) + state = _ENGINE_STATES.get((str(device), num_tokens, shared)) + if state is None: + state = _ENGINE_STATES[(str(device), num_tokens, shared)] = M3MoeState( + device, num_tokens=num_tokens, num_ctas=sms - MOE_ENGINE_RESERVED_SMS, shared=shared + ) + engine = state.layer( + *tensors, + self.w[f"l{i}_router_bias"].data, + router_weight=self.w[f"l{i}_router"].data, + fc31_input_scale=scale.reshape(1).float(), + **kwargs, + ) + # Compile (first use of the state) and warm. + m3_moe_engine( + torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device), + engine, + ) + return engine + + # ------------------------------------------------------------------------------------------ the index scorer + + def _scorer_variants(self, query_lens) -> list: + """The scorer keys a decode step can take: every sparse layer's (index dtype, index heads per rank, index dim) + at each decode query length, with the early-issue schedule (up to ``INDEX_SCORE_EARLY_MAX_REQUESTS`` requests, + none here) and without it, where the kernel serves it.""" + runner = _index_score_kernel.IndexScoreRunner + keys = [] + for i in sorted(self.sparse_layers): + backend = self.attn[i] + config = backend.m3_config + dtype = torch.float8_e4m3fn if backend.indexer_kv_dtype == "fp8" else torch.bfloat16 + for qlen in sorted(set(query_lens)): + if not runner.is_supported( + q_dtype=dtype, + num_heads=config.num_index_heads, + head_dim=config.sparse_index_dim, + page_size=runner.SUPPORTED_PAGE_SIZE, + max_decode_query_len=qlen, + ): + continue + for early in (True, False) if INDEX_SCORE_EARLY_MAX_REQUESTS else (False,): + key = ( + dtype, + int(config.num_index_heads), + int(qlen), + int(config.sparse_index_dim), + int(INDEX_SCORE_SPLIT_K), + bool(early), + ) + if key not in keys: + keys.append(key) + return keys + + def _prebuild_index_scorer(self, query_lens, device) -> list: + """Compile every scorer key and launch each once on real-shaped dummy operands (loading its module), on the + current stream. The kernel's batch, token, page and block counts are symbolic: its compile needs no step + metadata. Returns the keys.""" + runner = _index_score_kernel.IndexScoreRunner + keys = self._scorer_variants(query_lens) + page = runner.SUPPORTED_PAGE_SIZE + for key in keys: + dtype, heads, qlen, dim, split_k, early = key + runner._compile(dtype, heads, qlen, dim, split_k, early) + rows = 1 if early else INDEX_SCORE_EARLY_MAX_REQUESTS + 1 + blocks = 2 + tokens = rows * qlen + idx_q = torch.zeros(tokens, heads, dim, dtype=torch.bfloat16, device=device).to(dtype) + idx_k = torch.zeros( + rows * blocks, 1, page, dim, dtype=torch.bfloat16, device=device + ).to(dtype) + block_table = torch.arange(rows * blocks, dtype=torch.int32, device=device).view( + rows, blocks + ) + seq_lens = torch.full((rows,), blocks * page, dtype=torch.int32, device=device) + score = torch.empty(heads, blocks, tokens, dtype=torch.float32, device=device) + runner.kernel_cache[key]( + idx_q, idx_k.squeeze(1), block_table, score.transpose(1, 2), seq_lens + ) + return keys + + # --------------------------------------------------------------------------------------------- the dispatch + + def _select_target(self, attn_metadata: AttentionMetadata): + """Which target runs this step: one host-side integer, fixed per CUDA-graph capture. A step with context rows + (a mixed step included) is prefill's; every pure generation step is decode's, whatever its width.""" + return self._decode if attn_metadata.num_contexts == 0 else self._prefill + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + capture: Callable[..., None] = _plain_capture, + **kwargs, + ) -> torch.Tensor: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + if kwargs.get("lora_params") is not None: + raise ValueError("this MiniMax-M3 target serves no LoRA adapters") + return self._select_target(attn_metadata).forward( + attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, capture + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The targets +# -------------------------------------------------------------------------------------------------------------------- + + +class PrefillTarget: + """A step with context rows: prefill chunks, and the mixed steps that carry generation rows beside them. Every + layer runs in its generic form; the attention core is the MSA backend's for every row.""" + + def __init__(self, core: MiniMaxM3Tp8ModelingV2) -> None: + self.core = core + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + Counts and per-step host scalars are read here; device tensors are the metadata's persistent buffers, held + by reference; ``msa_idx_k_cache`` is the metadata's per-layer index-K cache getter. ``metadata`` is the + prepared object itself: the MSA backend's step (``m3_msa_attention``) takes it as its argument. + """ + return dict( + num_contexts=md.num_contexts, + num_ctx_tokens=md.num_ctx_tokens, + num_tokens=md.num_tokens, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + metadata=md, + ) + + def forward( + self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, capture + ): + core = self.core + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target does not run under piecewise (breakable) prefill CUDA graphs" + ) + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + _record("prefill", int(step["num_tokens"]), int(step["num_contexts"])) + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + hidden, residual = self._layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + return hidden + + def _swiglu_mlp(self, h: torch.Tensor, prefix: str) -> torch.Tensor: + """A gated MXFP8 MLP (the dense MLP, the shared expert): gate_up, SwiGLU-OAI, down, no reduction.""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s") + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + out, _ = _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s") + return out + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: + the MoE scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor): + """Router + routed experts, finalized: the fused router with the NVFP4 activation up to 48 rows, else the + router logits and the experts quantizing their bf16 input (chunked).""" + core = self.core + fused = _router_fp4(core, i, h) + if fused is not None: + router_logits, x_fp4, x_sf = fused + return _routed_experts(core, i, x_fp4, x_sf, router_logits, True) + router_logits = _router_logits(core, i, h) + return self._routed_chunks(i, h, router_logits) + + def _moe(self, i: int, h: torch.Tensor, *, return_parts: bool): + """The MoE block of layer ``i`` without its reduction. Returns ``(out, routed)``: with ``return_parts`` the + shared and routed parts (routed None when the engine folded the shared expert in), else their sum in + ``out``.""" + core = self.core + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, False) + folded = m3_moe_engine(h.contiguous(), engine, None, None) + return folded, None + routed_out, shared_out = maybe_execute_in_parallel( + lambda: self._routed(i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh"), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the + # same collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if return_parts: + return shared_out, routed_out + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + if _can_mc(hidden): + fused = _mc_allreduce_norm(core, hidden, None, residual, core.w[f"l{i}_post_norm"]) + hidden, residual = ( + fused + if fused is not None + else _native_allreduce_norm(core, hidden, residual, core.w[f"l{i}_post_norm"]) + ) + else: + hidden, residual = _native_allreduce_norm( + core, hidden, residual, core.w[f"l{i}_post_norm"] + ) + if i in core.moe_layers: + out, routed = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out = self._swiglu_mlp(hidden, f"l{i}_mlp") + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step): + """Attention of layer ``i`` up to its o_proj (no reduction). Returns (o_proj output, pushed): a prefill step + arms no boundary push, so pushed is always False.""" + core = self.core + if i not in core.sparse_layers: + # Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v for an FP8 cache, else in place in bf16), the + # MSA backend (it writes the layer's cache), o_proj. + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s") + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s") + if not core.fused_projection: + return self._unfused_sparse_attention(i, position_ids, hidden, step) + packed, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s") + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s") + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection + and its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and + the FP8 indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the + o_proj. A layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s") + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s") + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): the MSA + backend (``m3_msa_attention``: the KV write unless a producer wrote the caches, the sparse layers' top-k block + selection, the FMHA over the context rows and any generation rows beside them).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + live = int(step["num_tokens"]) + q_live, out_live = q[:live], output[:live] + k_live = k[:live] if k is not None else None + v_live = v[:live] if v is not None else None + idx_q_live = idx_q[:live] if idx_q is not None else None + sparse = i in core.sparse_layers + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k_live is not None) + m3_msa_attention( + core.attn[i], + q_live, + k_live, + v_live, + idx_q_live, + step["metadata"], + out_live, + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + return output + + +class DecodeTarget: + """A pure generation step (no context rows), of any width. + + Up to ``DECODE_MAX_TOKENS`` tokens with 1-D input ids, per layer: attention (o_proj all-reduce deferred, its rows + pushed into the boundary all-reduce up to 16 tokens) -> attention boundary (all-reduce + residual + post-attention + RMSNorm) -> MoE (the 1..8-row engine, pushing its sum into the boundary all-reduce; the unfinalized MoE with the + finalize in the all-reduce from 9 rows) or dense MLP -> MoE boundary (all-reduce + shared/routed add + residual + + the next layer's RMSNorm, emitting the next QKV's MXFP8 activation). A layer whose output Eagle3 captures runs the + generic layer, as does the first layer when the fused embedding does not take the step. A wider step (or one + carrying embeddings) runs every layer in its generic form, with this target's attention (the decode selection and + sparse decode where the producer wrote the caches). + """ + + def __init__(self, core: MiniMaxM3Tp8ModelingV2) -> None: + self.core = core + self._moe_push_descs: Dict[tuple, Any] = {} + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + States the two context counts instead of reading them back: this target runs only when ``num_contexts`` is + zero, so reading them would ask the metadata to confirm what the routing already decided. + + Graph classes: counts, the generation span, the identity-selection flag and the worst-case score extent are + host values fixed per capture (prepare() stages them before the forward); the block table, lengths, cache + slots, identity table and valid-block counts are the metadata's persistent buffers, held by reference and + refreshed in place every step; ``msa_idx_k_cache`` / ``msa_proxy_max_score_view`` are the metadata's getters + of the per-layer index-K cache and of the score scratch view. ``metadata`` is the prepared object itself: the + MSA backend's step (``m3_msa_attention``) takes it as its argument. + """ + return dict( + num_contexts=0, + num_ctx_tokens=0, + num_tokens=md.num_tokens, + num_seqs=md.num_seqs, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + msa_decode_span=md.msa_decode_span, + msa_block_table=md.msa_block_table, + msa_seq_lens_cuda=md.msa_seq_lens_cuda, + msa_short_ok=md._msa_short_ok, + msa_short_table=md.msa_short_table, + msa_proxy_max_score_view=md.msa_proxy_max_score_view, + msa_worst_case_max_k_tiles=md.msa_worst_case_max_k_tiles, + msa_n_valid_blocks=md.msa_n_valid_blocks, + metadata=md, + ) + + def forward( + self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, capture + ): + core = self.core + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + # The decode kernels' rows (a per-capture constant): up to DECODE_MAX_TOKENS tokens of 1-D input ids. + kernels = ( + inputs_embeds is None and input_ids.ndim == 1 and is_decode_step(input_ids.numel()) + ) + if kernels: + num_tokens = int(input_ids.numel()) + plan = plan_for(num_tokens) + _record("decode", num_tokens, plan.name) + hidden, residual = self._embed(input_ids, num_tokens) + else: + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target does not run under piecewise (breakable) CUDA graphs" + ) + plan = None + _record("decode", int(step["num_tokens"]), "generic") + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + if ( + not kernels + or residual is None + or (spec_metadata is not None and spec_metadata.is_layer_capture(i)) + ): + hidden, residual = self._generic_layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + else: + hidden, residual = self._layer(i, plan, position_ids, hidden, residual, step) + return hidden + + def _embed(self, input_ids: torch.Tensor, num_tokens: int): + """(hidden, residual) of the fused embedding + first RMSNorm; (embedding, None) when it does not run, so the + first layer runs its own input norm.""" + core = self.core + if not EMBED_FIRSTNORM.covers(num_tokens): + uncovered(EMBED_FIRSTNORM.path, f"more than {EMBED_FIRSTNORM_MAX_TOKENS} tokens") + return m3_embedding(input_ids, core.w["embed"]), None + owner = core.embed_owners.get((num_tokens, input_ids.dtype)) + if owner is None or not input_ids.is_contiguous(): + missed( + EMBED_FIRSTNORM.path, + num_tokens, + f"no embedding owner for {input_ids.dtype} ids (built for int32 / int64 after the load)" + if owner is None + else "the input ids are not contiguous", + ) + return m3_embedding(input_ids, core.w["embed"]), None + _record("m3_embedding_firstnorm", num_tokens, input_ids.dtype) + return m3_embedding_firstnorm( + input_ids, owner.table, owner.weight, owner.eps, owner.raw, owner.output, owner.use_pdl + ) + + def _bucket_flags(self, num_tokens: int): + """(bucket workspace, its 8-CTA call counters) of the boundary all-reduce for ``num_tokens``, or None.""" + workspace = m3_mc_allreduce.get_workspace(self.core.mapping, self.core.hidden) + workspace = None if workspace is None else m3_mc_allreduce._bucket(workspace, num_tokens) + flags = None if workspace is None else workspace.flags_for(8) + return None if flags is None else (workspace, flags) + + @staticmethod + def _no_bucket(path: str, num_tokens: int) -> None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, "the multicast all-reduce workspace") + else: + missed(path, num_tokens, "no multicast all-reduce bucket holds the step") + + def _attn_push_target(self, num_tokens: int): + """The o_proj's push tuple into the attention boundary's all-reduce, or None (after a miss). + + None comes only from the token count or from workspace state created collectively in the same eager call on + every rank, so all ranks agree. (A lone non-pushing rank would still be safe: its all-reduce pushes the same + bf16 rows into the same set A.)""" + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(ATTN_PUSH.path, num_tokens) + return None + workspace, flags = found + return ( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + AR_PUSH_MODE, + ) + + def _moe_push_targets(self, num_tokens: int): + """(engine push descriptor, shared down-projection push tuple) for the MoE boundary's exchange sets, or None + (after a miss). + + This decision must match across ranks: prepushed 3 reads set B of every rank, which a non-pushing rank never + fills. None comes only from the token count or from state created in the same eager call on every rank + (workspace, flags, the descriptor cache), and the engine it is armed for is agreed at build time + (``_prebuild_engines``), so it does.""" + hidden = self.core.hidden + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(MOE_PUSH.path, num_tokens) + return None + workspace, flags = found + key = (workspace.mc.data_ptr(), flags.data_ptr()) + desc = self._moe_push_descs.get(key) + if desc is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(MOE_PUSH.path, num_tokens, "the MoE engine's push descriptor") + return None + desc = self._moe_push_descs[key] = ar_push_desc( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + flags.device, + ) + set_b = workspace.mc.data_ptr() + 2 * 4 * workspace.max_tokens * workspace.world_size * ( + hidden // 2 + ) + return desc, ( + set_b, + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + 2, + ) + + def _boundary(self, use_mc: bool, x, residual, norm_w, x2=None, quant=False): + """All-reduce + (second addend) + residual + RMSNorm: the multicast kernel, else the built-in all-reduce.""" + core = self.core + if use_mc: + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w, quant=quant) + if fused is not None: + return fused + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + missed( + MC_ALLREDUCE.path, + x.shape[0], + m3_mc_allreduce.decline_reason(x, residual, workspace, x2=x2), + ) + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _prepushed(self, i: int, x, residual, norm_w, quant, prepushed, x2=None): + fused = _mc_allreduce_norm( + self.core, x, x2, residual, norm_w, quant=quant, prepushed=prepushed + ) + if fused is None: + raise RuntimeError(f"layer {i}: rows were pushed into an all-reduce that did not run") + return fused + + def _layer(self, i: int, plan: DecodePlan, position_ids, hidden, residual, step): + core = self.core + num_tokens = hidden.shape[0] + # Padded rows: the boundary all-reduce polls and re-arms rows t < M = residual.shape[0], CUDA-graph padding + # rows included. Every producer pushes all rows of the same tensor: the o_proj / shared down GEMV t < m = + # x.shape[0] (its row count comes from the host shape, not a live count), the MoE engine every row of the + # num_tokens it was built for (= num_tokens here). The all-reduce checks x.sizes() == residual.sizes(), so + # no padded row is left with a stale flag. + # The GEMV dispatch (and with it every push into a boundary all-reduce) is a decode-graph one. + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + if not decode_graph: + uncovered( + "decode_graph", + "an eager decode step: the decode-graph GEMV dispatch and the boundary pushes are off", + ) + push = plan.attn_mc_ar and decode_graph and ATTN_PUSH.covers(num_tokens) + armed = self._attn_push_target(num_tokens) if push else None + hidden, attn_pushed = self._attention(i, position_ids, hidden, step, armed) + post_norm = core.w[f"l{i}_post_norm"] + if armed is not None and attn_pushed: + hidden, residual = self._prepushed(i, hidden, residual, post_norm, False, AR_PUSH_MODE) + else: + if armed is not None: + missed( + ATTN_PUSH.path, + num_tokens, + "the o_proj did not push its rows into the armed boundary all-reduce", + ) + hidden, residual = self._boundary(plan.attn_mc_ar, hidden, residual, post_norm) + + norm_w = _next_norm(core, i) + if i not in core.moe_layers: + out, _ = self._swiglu_mlp(hidden, f"l{i}_mlp") + return self._boundary(plan.attn_mc_ar, out, residual, norm_w) + + quant = plan.qkv_prequant and i + 1 < core.num_layers + path = moe_path(num_tokens) + moe_push = None + if path == MOE_ENGINE.path: + # The engine runs inside the MoE; looked up (or found missing) here first so a push is armed only for an + # engine that will take it. + engine = core.engine(i, num_tokens, shared=True) + if engine is not None and plan.moe_mc_ar and decode_graph: + moe_push = self._moe_push_targets(num_tokens) + if path == MC_FINALIZE.path and plan.moe_mc_ar: + if self._bucket_flags(num_tokens) is None: + self._no_bucket(MC_FINALIZE.path, num_tokens) + else: + (shared, fc2, scales, perm), _, _, _ = self._moe( + i, hidden, return_parts=False, do_finalize=False + ) + _record("m3_mc_finalize_allreduce_norm", num_tokens, quant) + fused = m3_mc_allreduce.m3_mc_finalize_allreduce_norm( + fc2, + perm.view(num_tokens, -1).to(torch.int32), + scales.reshape(num_tokens, -1), + shared, + residual, + norm_w, + core.rms_eps, + m3_mc_allreduce.get_workspace(core.mapping, core.hidden), + use_pdl=get_env_enable_pdl(), + quant=quant, + ) + if fused is None: + raise RuntimeError( + "the MC finalize all-reduce rejected the unfinalized MoE outputs" + ) + return fused + shared, routed, engine_pushed, shared_pushed = self._moe( + i, hidden, return_parts=True, push=moe_push + ) + if moe_push is not None: + if routed is None: + # The engine folded the shared expert in and pushed routed + shared into set A. + if engine_pushed: + return self._prepushed(i, shared, residual, norm_w, quant, 2) + else: + if engine_pushed and shared_pushed: + return self._prepushed(i, shared, residual, norm_w, quant, 3, x2=routed) + if engine_pushed or shared_pushed: + raise RuntimeError("only one MoE part pushed into the boundary all-reduce") + missed( + MOE_PUSH.path, + num_tokens, + "the MoE did not push its output into the armed boundary all-reduce", + ) + return self._boundary(plan.moe_mc_ar, shared, residual, norm_w, x2=routed, quant=quant) + + def _swiglu_mlp(self, h: torch.Tensor, prefix: str, push: Optional[tuple] = None): + """A gated MXFP8 MLP (the dense MLP, the shared expert): gate_up, SwiGLU-OAI, down, no reduction. Returns + (output, whether the down projection's GEMV pushed into ``push``).""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s") + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + return _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s", push) + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: the MoE + scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor, do_finalize: bool): + """Router + routed experts: the fused router with the NVFP4 activation up to 48 rows, else the router logits and + the experts quantizing their bf16 input.""" + core = self.core + fused = _router_fp4(core, i, h) + if fused is not None: + router_logits, x_fp4, x_sf = fused + return _routed_experts(core, i, x_fp4, x_sf, router_logits, do_finalize) + router_logits = _router_logits(core, i, h) + if not do_finalize: + return _routed_experts(core, i, h, None, router_logits, False) + return self._routed_chunks(i, h, router_logits) + + def _moe( + self, i: int, h: torch.Tensor, *, return_parts: bool, do_finalize: bool = True, push=None + ): + """The MoE block of layer ``i`` without its reduction. + + ``push`` = (the engine's ``arp`` descriptor, the shared down projection's ``ar`` tuple) the decode layer armed + for the MoE boundary's exchange sets, or None. Returns ``(out, routed, engine_pushed, shared_pushed)``: with + ``return_parts`` the shared and routed parts (routed None when the engine folded the shared expert in), else + their sum in ``out``; with ``do_finalize=False`` ``out`` is ``[shared, fc2 out, expert scales, permuted + index]``. + """ + core = self.core + engine_arp, shared_ar = push if push is not None else (None, None) + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + do_finalize + and h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, engine_arp is not None) + folded = m3_moe_engine(h.contiguous(), engine, None, engine_arp) + return folded, None, engine_arp is not None, False + routed_out, (shared_out, shared_pushed) = maybe_execute_in_parallel( + lambda: self._routed(i, h, do_finalize), + lambda: self._swiglu_mlp(h, f"l{i}_sh", shared_ar), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the same + # collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if not do_finalize: + return [shared_out, *routed_out], None, False, shared_pushed + if return_parts: + return shared_out, routed_out, False, shared_pushed + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None, False, shared_pushed + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _generic_layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + if _can_mc(hidden): + fused = _mc_allreduce_norm(core, hidden, None, residual, core.w[f"l{i}_post_norm"]) + hidden, residual = ( + fused + if fused is not None + else _native_allreduce_norm(core, hidden, residual, core.w[f"l{i}_post_norm"]) + ) + else: + hidden, residual = _native_allreduce_norm( + core, hidden, residual, core.w[f"l{i}_post_norm"] + ) + if i in core.moe_layers: + out, routed, _, _ = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out, _ = self._swiglu_mlp(hidden, f"l{i}_mlp") + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step, push: Optional[tuple] = None): + """Attention of layer ``i`` up to its o_proj (no reduction). ``push`` is the attention boundary's exchange the + o_proj may write into. Returns (o_proj output, pushed).""" + core = self.core + if i not in core.sparse_layers: + # Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v: an FP8 cache, or the FP8 pool the dense + # layers sit in within a hybrid NVFP4 cache on a decode step), the MSA backend, o_proj. + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s") + out_fp8 = core.main_kv_is_fp8 or core.main_kv_is_nvfp4 + fused = _qk_norm_rope(core, i, qkv, position_ids, out_fp8) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", push) + if not core.fused_projection: + return self._unfused_sparse_attention(i, position_ids, hidden, step, push) + packed, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s") + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", push) + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step, push): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection and + its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and the FP8 + indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the o_proj. A + layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s") + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", push) + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): a sparse + layer whose caches the producer wrote runs this target's selection and sparse decode when the generation span + starts at row 0; everything else the MSA backend (``m3_msa_attention``: the KV write unless a producer wrote the + caches, the sparse layers' top-k block selection, the FMHA).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + live = int(step["num_tokens"]) + q_live, out_live = q[:live], output[:live] + k_live = k[:live] if k is not None else None + v_live = v[:live] if v is not None else None + idx_q_live = idx_q[:live] if idx_q is not None else None + span = step["msa_decode_span"] + sparse = i in core.sparse_layers + if sparse and k_live is None and span is not None and span.row_first == 0: + self._decode_sparse_core(i, q_live, idx_q_live, step, out_live) + return output + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k_live is not None) + m3_msa_attention( + core.attn[i], + q_live, + k_live, + v_live, + idx_q_live, + step["metadata"], + out_live, + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + return output + + def _decode_sparse_core(self, i: int, q, idx_q, step, output) -> None: + """Top-16 block selection and the block-sparse decode of one pure generation step of a sparse layer whose caches + the producer already wrote. + + The selection is the identity when the metadata staged it (every request within 16 blocks). Otherwise a span of + at most ``FUSED_SCORE_SELECT_MAX_REQUESTS`` requests scores and selects in one launch; a wider span runs the + CuTe DSL scorer (the per-index-head block maxima, reduced to the KV head) and the long-row selector (the + built-in selector beyond its shapes). + """ + core = self.core + backend = core.attn[i] + config = backend.m3_config + num_tokens = int(q.shape[0]) + rows = int(step["num_seqs"]) + qlen = int(step["msa_decode_span"].query_len) + block_table = step["msa_block_table"][:rows] + seq_lens = step["msa_seq_lens_cuda"][:rows] + num_kv_heads = int(backend.indexer.config.num_kv_heads) + identity = step["msa_short_ok"] and step["msa_short_table"] is not None + fused = ( + None if identity else self._fused_score_select(i, idx_q, step, rows, qlen, num_tokens) + ) + if identity: + topk = step["msa_short_table"][: num_kv_heads * num_tokens * MSA_REQUIRED_TOPK].view( + num_kv_heads, num_tokens, MSA_REQUIRED_TOPK + ) + elif fused is not None: + topk = fused.permute(1, 0, 2) + else: + idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) + idx_k_cache = step["msa_idx_k_cache"](i) + max_score = step["msa_proxy_max_score_view"]( + config.num_index_heads, step["msa_worst_case_max_k_tiles"], num_tokens + ) + early = rows <= INDEX_SCORE_EARLY_MAX_REQUESTS + key = ( + idx_q_view.dtype, + int(config.num_index_heads), + int(qlen), + int(idx_q_view.shape[-1]), + int(INDEX_SCORE_SPLIT_K), + bool(early), + ) + if key not in _index_score_kernel.IndexScoreRunner.kernel_cache: + # Outside require mode the op below compiles it (the only scorer this step has). + missed( + INDEX_SCORE.path, + num_tokens, + f"the index scorer {key} was not compiled before warmup", + ) + _record("m3_index_score", key) + m3_index_score( + idx_q_view, + idx_k_cache.squeeze(1), + block_table, + seq_lens, + max_score.transpose(1, 2), + qlen, + INDEX_SCORE_SPLIT_K, + early, + ) + n_valid_blocks = step["msa_n_valid_blocks"][:num_tokens] + group = config.num_index_heads // config.num_kv_heads + max_score_kv = ( + max_score.view( + config.num_kv_heads, group, max_score.shape[1], max_score.shape[2] + ).amax(dim=1) + if group > 1 + else max_score + ) + table = long_select_blocks( + max_score_kv, + n_valid_blocks=n_valid_blocks, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + use_pdl=get_env_enable_pdl(), + early_trigger=num_tokens <= SELECT_EARLY_MAX_ROWS, + ) + if table is None: + uncovered( + LONG_SELECT.path, + f"{max_score_kv.shape[1]} KV blocks, outside {LONG_SELECT_MIN_BLOCKS}..{LONG_SELECT_MAX_BLOCKS}: " + "the built-in selector", + ) + table = m3_select_blocks( + max_score_kv, + MSA_REQUIRED_TOPK, + n_valid_blocks, + config.init_blocks, + config.local_blocks, + True, + ) + topk = table.permute(1, 0, 2) + manager = step["kv_cache_manager"] + k_paged, v_paged = msa_paged_kv(manager, i) + nvfp4 = {} + if manager.is_nvfp4_layer(i): + k_scale, v_scale = _aligned_nvfp4_dequant_scales(backend, core.kv_scales[i]) + scales = manager.get_block_scale_buffers(i, "HND") + nvfp4 = dict( + k_block_scale=scales[:, 0], + v_block_scale=scales[:, 1], + k_global_scale=k_scale, + v_global_scale=v_scale, + ) + _record("m3_sparse_decode", num_tokens, qlen, identity) + m3_sparse_decode( + q.view(num_tokens, core.heads, core.head_dim), + k_paged, + v_paged, + topk, + block_table, + seq_lens, + sm_scale=core.sm_scale(i), + output=output.view(num_tokens, core.heads, core.head_dim), + decode_query_len=qlen, + identity_blocks=identity, + **nvfp4, + ) + + def _fused_score_select(self, i: int, idx_q, step, rows: int, qlen: int, num_tokens: int): + """The top-16 table [total_q, num_kv_heads, 16] of a covered decode span from the fused score + select, or None + (after a miss when the span is covered) for the scorer + selector.""" + core = self.core + path = FUSED_SCORE_SELECT.path + if not FUSED_SCORE_SELECT.covers(rows): + uncovered( + path, f"more than {FUSED_SCORE_SELECT_MAX_REQUESTS} requests: scorer + selector" + ) + return None + config = core.attn[i].m3_config + idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) + idx_k_pages = step["msa_idx_k_cache"](i).squeeze(1) + max_tiles = int(step["msa_worst_case_max_k_tiles"]) + n_valid = step["msa_n_valid_blocks"] + reason = _fused_score_select_kernel.supported( + idx_q_view, idx_k_pages, qlen, config.num_kv_heads, max_tiles + ) + if reason is None and ( + num_tokens != rows * qlen + or n_valid is None + or n_valid.dtype != torch.int32 + or int(n_valid.shape[0]) < num_tokens + ): + reason = "the step's query rows or valid-block counts are not one decode span" + if reason is not None: + missed(path, rows, reason) + return None + bufs = _fused_score_select_kernel.buffers(idx_q.device, max_tiles) + if bufs is None: + capture_missed(path, rows, "the fused score + select scratch") + return None + _record("m3_fused_score_select", rows, qlen) + return _fused_score_select_kernel.fused_score_select( + idx_q_view, + idx_k_pages, + step["msa_block_table"][:rows], + step["msa_seq_lens_cuda"][:rows], + n_valid[:num_tokens].contiguous(), + decode_query_len=qlen, + num_kv_heads=config.num_kv_heads, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + max_tiles=max_tiles, + scratch=bufs[0], + counters=bufs[1], + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# Speculative decoding: one-model Eagle3 at checkpoint precision +# -------------------------------------------------------------------------------------------------------------------- + + +class M3SpecStepFold: + """The target's speculative step-input staging (the in-tree ``SpecStepFold`` protocol). + + Between two decode graphs of a one-model Eagle3 step, the spec sampler stores the step's outputs into its per-slot + store and the engine gathers them back into its flat input buffers for the next step. This object does both in + one launch each (``m3_spec_tail.slot_scatter`` / ``step_gather``): integer copies only, the same values the + engine's and the sampler's own copies write. The forward never stores the sampler's outputs itself, so + ``begin_step`` and ``sampler_rows_stored`` always report that nothing was staged. + """ + + def __init__(self) -> None: + self.store: Optional[object] = None + self.step_inputs: Optional[tuple] = None + + def bind_engine_buffers( + self, input_ids, draft_tokens, position_id_offsets, kv_lens_offsets, max_batch_size + ): + self.step_inputs = (input_ids, draft_tokens, position_id_offsets, kv_lens_offsets) + + def bind_sampler_store(self, store: object) -> None: + self.store = store + + def begin_step(self, seq_slots: Optional[list], num_rows: int) -> bool: + return False + + def gather( + self, + new_tokens, + next_draft_tokens, + new_tokens_lens, + previous_slots, + previous_pos_indices, + previous_batch_len: int, + tokens_per_gen_step: int, + draft_token_buffer_width: int, + num_tokens: int, + num_draft_tokens: int, + first_previous_row: int, + ) -> bool: + if self.step_inputs is None: + return False + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets = self.step_inputs + return m3_spec_tail.step_gather( + new_tokens, + next_draft_tokens, + new_tokens_lens, + previous_slots, + previous_pos_indices, + previous_batch_len, + tokens_per_gen_step, + draft_token_buffer_width, + input_ids, + num_tokens, + draft_tokens, + num_draft_tokens, + position_id_offsets, + kv_lens_offsets, + first_previous_row, + ) + + def sampler_rows_stored(self, outputs: dict, num_skip: int, seq_slots: list) -> bool: + return False + + def scatter( + self, seq_slots, new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens + ) -> bool: + store = self.store + if store is None: + return False + return m3_spec_tail.slot_scatter( + seq_slots, + new_tokens, + store.new_tokens.squeeze(-1).T, + next_new_tokens, + store.next_new_tokens.squeeze(-1).T, + new_tokens_lens, + store.new_tokens_lens, + next_draft_tokens, + store.next_draft_tokens, + ) + + +# (N, K, M, PDL) of every drafter low-M call compiled before the executor's warmup (``_prebuild_draft_lowm``). +_DRAFT_LOWM_BUILT: set = set() + + +def _draft_linear(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """A drafter BF16 linear ``x @ weight^T``: FlashInfer's fixed-tactic low-M kernel on the rows ``DRAFT_LOWM`` + covers (compiled before warmup; a covered call found uncompiled is a miss), else the in-tree BF16 GEMM.""" + n, k = weight.shape + m = int(x.shape[0]) + path = DRAFT_LOWM.path + if x.dim() != 2 or not DRAFT_LOWM.covers(((n, k), m)): + uncovered(path, f"the drafter's {n}x{k} linear past its tactic table's rows") + elif not x.is_contiguous(): + # The first pass's fc may read a column window of the captured hidden states; the in-tree GEMM takes strides. + uncovered(path, "a strided activation (the in-tree GEMM reads it in place)") + elif x.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: + missed(path, m, "the activation or weight is not BF16") + elif (n, k, m, get_env_enable_pdl()) not in _DRAFT_LOWM_BUILT: + missed(path, m, f"the {n}x{k} kernel at {m} rows was not compiled before warmup") + else: + kind, tactic = tactic_for(n, k, m) + out = torch.empty((m, n), dtype=x.dtype, device=x.device) + _record("m3_draft_lowm", (n, k), m) + return m3_draft_lowm(x, weight, out, kind, tactic, get_env_enable_pdl()) + return m3_bf16_linear(x, weight) + + +def _prebuild_draft_lowm(draft: "M3Eagle3Draft") -> int: + """Compile (and run once) every covered row count of the drafter's linears, outside capture, on the current + stream. Returns the number of calls compiled.""" + use_pdl = get_env_enable_pdl() + built = 0 + shapes = {tuple(draft.w[key].shape): draft.w[key] for key in ("qkv", "o", "gu", "dn", "fc")} + for (n, k), weight in sorted(shapes.items()): + for _, m in sorted(key for key in DRAFT_LOWM.keys if key[0] == (n, k)): + if (n, k, m, use_pdl) in _DRAFT_LOWM_BUILT: + continue + kind, tactic = tactic_for(n, k, m) + x = torch.zeros(m, k, dtype=torch.bfloat16, device=weight.device) + out = torch.empty(m, n, dtype=torch.bfloat16, device=weight.device) + m3_draft_lowm(x, weight, out, kind, tactic, use_pdl) + _DRAFT_LOWM_BUILT.add((n, k, m, use_pdl)) + built += 1 + return built + + +class _DraftForward: + """``draft_model.model`` for the in-tree worker: the drafter layer's forward, called as + ``model(input_ids=, position_ids=, hidden_states=, attn_metadata=, spec_metadata=, all_rank_num_tokens=)`` and + returning (hidden, hidden to carry to the next pass). ``d2t`` None: the draft vocabulary is the target's.""" + + d2t = None + + def __init__(self, draft: "M3Eagle3Draft") -> None: + self._draft = draft + + def __call__(self, **kwargs): + return self._draft.forward_layer(**kwargs) + + +class M3Eagle3Draft(nn.Module): + """The one-layer BF16 Eagle3 drafter (``speculative_model``): its weights (``weights.DRAFT_WEIGHTS``), its + vocab-parallel BF16 head and logits processor, the drafter layer's TRTLLM attention backend (layer index = the + target's layer count, FP8 KV) and its TP all-reduce. + + What the worker reads: ``model`` (the layer forward), ``apply_eagle3_fc`` (the per-layer fc_norms + fc over the + captured target hidden states), ``lm_head`` and ``logits_processor``; the fused draft loop reads the layer's parts + (``draft_input_weights``, ``attention``, ``mlp``, ``ar_norm``). The drafter embeds with the target's table. + """ + + def __init__( + self, + core: MiniMaxM3Tp8ModelingV2, + draft_config: ModelConfig, + start_layer_idx: int, + capture_layers: int, + ) -> None: + super().__init__() + cfg = draft_config.pretrained_config + _ensure_draft_vocab_size(cfg) + eagle = getattr(cfg, "eagle_config", None) or {} + mapping = draft_config.mapping + self.draft_config = draft_config + # The engine's draft weight mapper binds to these two. + self.model_config = draft_config + self.config = cfg + self.capture_layers = int(capture_layers) + self.layer_idx = int(start_layer_idx) + self.dtype = cfg.torch_dtype + self.hidden = int(cfg.hidden_size) + self.hidden_size_in = int(getattr(cfg, "target_hidden_size", cfg.hidden_size)) + self.rms_eps = float(cfg.rms_norm_eps) + self.return_hidden_post_norm = bool( + eagle.get("return_hidden_post_norm", False) or getattr(cfg, "norm_output", False) + ) + if ( + int(cfg.num_hidden_layers) != 1 + or eagle.get("eh_proj_before_attn", False) + or eagle.get("norm_before_fc", False) + or not getattr(cfg, "fc_norm", False) + or getattr(cfg, "kv_lora_rank", None) + or self.capture_layers <= 1 + or self.hidden_size_in != self.hidden + or cfg.draft_vocab_size != cfg.vocab_size + or getattr(cfg, "attention_bias", False) + or getattr(cfg, "mlp_bias", False) + or getattr(cfg, "bias", False) + or getattr(cfg, "tie_word_embeddings", False) + ): + raise NoTarget( + "this MiniMax-M3 target drafts with a one-layer Llama-style Eagle3 drafter with per-layer " + "fc_norms over " + "several captured layers, the target's hidden size and vocabulary and no biases" + ) + tp = mapping.tp_size + self.heads = cfg.num_attention_heads // tp + self.kv_heads = (cfg.num_key_value_heads + tp - 1) // tp + self.head_dim = int( + getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads + ) + self.q_size = self.heads * self.head_dim + self.kv_size = self.kv_heads * self.head_dim + self.w = _weights.DRAFT_WEIGHTS.declare(_weights.DRAFT_WEIGHTS.dims(self)) + self.lm_head = LMHead( + cfg.draft_vocab_size, + cfg.hidden_size, + dtype=cfg.torch_dtype, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=draft_config.lm_head_gather_output, + reduce_output=False, + use_custom_cublas_mm=False, + quant_config=None, + ) + self.logits_processor = LogitsProcessor() + self.allreduce = AllReduce( + mapping=mapping, strategy=draft_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + quant_config = draft_config.get_quant_config() + self.attn = create_attention( + draft_config.attn_backend, + self.layer_idx, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, rope=RopeParams.from_config(cfg) + ), + quant_config=quant_config, + skip_create_weights_in_init=draft_config.skip_create_weights_in_init, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + if not self.attn.support_fused_rope() or not self.attn.support_fused_qkv(): + raise NoTarget( + f"the drafter's attention backend {type(self.attn).__name__} must fuse RoPE and QKV" + ) + self.attn.update_quant_config(quant_config) + # The drafter embeds with the target's table, read off the core on every call (the engine materializes the + # core's weights after construction). + self._core = [core] + self.model = _DraftForward(self) + + @property + def embed_table(self) -> torch.Tensor: + return self._core[0].w["embed"] + + def load_weights_from_target_model(self, target_model: nn.Module) -> None: + """Nothing to alias: the drafter reads the target's embedding table through the core and owns its head.""" + + # ------------------------------------------------------------------------------------------------- the layer + + def fc(self, hs: torch.Tensor) -> torch.Tensor: + return _draft_linear(hs, self.w["fc"]) + + def apply_eagle3_fc(self, hidden_states: torch.Tensor) -> torch.Tensor: + """The captured target hidden states [T, layers x H] to the drafter's input width: each layer's fc_norm, + then the fc (the in-tree ``Eagle3ForCausalLM.apply_eagle3_fc``).""" + hidden_states = hidden_states.to(self.dtype) + if hidden_states.shape[-1] != self.hidden: + chunks = hidden_states.chunk(self.capture_layers, dim=-1) + hidden_states = torch.cat( + [ + flashinfer_rmsnorm(c, self.w[f"fc_norm{i}"], self.rms_eps) + for i, c in enumerate(chunks) + ], + dim=-1, + ) + hidden_states = self.fc(hidden_states) + return hidden_states + + def attention(self, x: torch.Tensor, md, reduce: bool) -> torch.Tensor: + """QKV, the backend's fused RoPE + KV append + causal attention, o_proj (all-reduced unless ``reduce`` is + False, where the caller fuses the reduction).""" + qkv = _draft_linear(x, self.w["qkv"]) + _record("m3_draft_attention", int(md.num_tokens)) + o = m3_draft_attention(self.attn, qkv, md) + out = _draft_linear(o, self.w["o"]) + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def mlp(self, h: torch.Tensor, reduce: bool) -> torch.Tensor: + gu = _draft_linear(h, self.w["gu"]) + act = m3_silu_and_mul(gu, None, None, None) + out = _draft_linear(act, self.w["dn"]) + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def ar_norm(self, x: torch.Tensor, residual: torch.Tensor, norm_key: str): + """One draft-layer boundary in the in-tree structure: the projection's plain all-reduce (BF16 out), then the + fused residual add + RMSNorm (FP32 add, norm of the FP32 sum). Returns (normed, residual sum).""" + return flashinfer_fused_add_rmsnorm( + m3_allreduce(x, self.allreduce, None), residual, self.w[norm_key], self.rms_eps + ) + + def forward_layer( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata, + spec_metadata, + all_rank_num_tokens: Optional[List[int]] = None, + ): + """The drafter layer (the in-tree ``Eagle3DraftModel.forward`` for one layer): the embedded tokens and the + fc'd hidden states normed and concatenated, attention, the post-attention norm with the residual, the MLP and + the final norm. Returns (normed output, the hidden state the next pass carries).""" + previous = attn_metadata.all_rank_num_tokens + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = all_rank_num_tokens + try: + embeds = m3_embedding(input_ids, self.embed_table).to(self.dtype) + residual = hidden_states + h = flashinfer_rmsnorm(hidden_states, self.w["hidden_norm"], self.rms_eps) + e = flashinfer_rmsnorm(embeds, self.w["in_norm"], self.rms_eps) + h = torch.cat([e, h], dim=-1) + h = self.attention(h, attn_metadata, reduce=True) + h, residual = flashinfer_fused_add_rmsnorm( + h, residual, self.w["post_norm"], self.rms_eps + ) + h = self.mlp(h, reduce=True) + spec_metadata.maybe_capture_hidden_states(self.layer_idx, h, residual) + h, to_save = flashinfer_fused_add_rmsnorm( + h, residual, self.w["final_norm"], self.rms_eps + ) + return (h, h) if self.return_hidden_post_norm else (h, to_save) + finally: + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = previous + + +def _eagle3_capture(spec_metadata, layer_id: int, hidden_states, residual, capture_norm) -> None: + """The Eagle3 hidden-state capture of one target layer. + + Mirrors ``Eagle3OneModelSpecMetadata.maybe_capture_hidden_states``. With ``capture_norm`` = (the drafter's stacked + fc_norm weights [layers, H], eps), the capture writes rmsnorm(bf16(h + residual)) * fc_norm[i] in one kernel (the + in-tree capture's bf16 add, then the in-tree fc_norm arithmetic), which the fused first drafter inputs read as is; + otherwise it is the in-tree capture. + """ + if capture_norm is None: + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + return + for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): + if captured_layer_id == layer_id: + assert spec_metadata.hidden_states is not None + num_tokens = spec_metadata.num_capture_tokens + if num_tokens > spec_metadata.hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds buffer capacity: " + f"num_tokens={num_tokens}, capacity={spec_metadata.hidden_states.shape[0]}" + ) + if num_tokens > hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds available hidden states: " + f"num_tokens={num_tokens}, hidden_states={hidden_states.shape[0]}" + ) + hidden_size = spec_metadata.hidden_size + m3_spec_tail.capture_norm( + hidden_states[:num_tokens], + residual[:num_tokens] if residual is not None else None, + capture_norm[0][i], + capture_norm[1], + spec_metadata.hidden_states[:num_tokens, i * hidden_size : (i + 1) * hidden_size], + ) + break + + +def _sharded_head(lm_head: nn.Module) -> bool: + """A vocab-sharded TP head without padding (the fused kernels read the local shard).""" + return ( + getattr(lm_head, "tp_mode", None) == TensorParallelMode.COLUMN + and getattr(lm_head, "padding_size", 0) == 0 + and getattr(lm_head, "tp_size", 1) > 1 + ) + + +class M3Eagle3Worker(Eagle3OneModelWorker): + """``Eagle3OneModelWorker`` with the MiniMax-M3 fused verify and fused draft loop, at the in-tree numerics. + + * Verify: on decode-only batches whose verify rows the fused kernels cover (``tail_eligible``), the target's BF16 + lm_head computes this rank's vocabulary shard of the logits (the in-tree head's own local GEMM) and one kernel + replaces the logits all-gather, the argmax (or the temperature + top-p draw) and the strict acceptance. + * Draft: the drafter's first pass may take the per-layer fc_norms in one launch (or folded into the hidden-state + capture) and run its fc beside the target head GEMM. Its BF16 linears run FlashInfer's fixed-tactic low-M + kernels at the rows ``DRAFT_LOWM`` covers (``_draft_linear``). On a decode-only step ``DRAFT_STEP`` covers, + the draft loop is fused: each pass's draft input in one launch, the drafter layer's attention and MLP with + their all-reduces fused with the residual add and the next RMSNorm, and one kernel for the BF16 head, the + cross-rank argmax and the pass's bookkeeping. + + Every other batch takes the in-tree paths. Methods whose docstring says "Mirrors" copy the named in-tree method + with the fused branches inserted. + """ + + def __init__(self, spec_config, mapping, use_separate_draft_kv_cache: bool = False) -> None: + super().__init__( + spec_config, mapping=mapping, use_separate_draft_kv_cache=use_separate_draft_kv_cache + ) + # Per-step state, reset at the top of every _forward_impl: the first draft pass's gather rows when this step + # takes the fused first drafter inputs, and the fc output a fork computed beside the target head. + self._m3_first_gather_ids: Optional[torch.Tensor] = None + self._m3_fc_hidden: Optional[torch.Tensor] = None + self._m3_target_lm_head: Optional[nn.Module] = None + # (fc_norm0 weight, stacked fc_norm weights [layers, H]), built eagerly on the first fused first pass. + self._m3_fc_norm_w: Optional[tuple] = None + # Armed per step by arm_capture_norm: (stacked fc_norm weights, eps) or None. + self._m3_capture_norm: Optional[tuple] = None + self._m3_aux_stream: Optional[torch.cuda.Stream] = None + # Armed per step by arm_capture_norm: the drafter whose first fc is forked at the last captured target layer, + # and the fork that did it, joined by _m3_target_accept. + self._m3_fc_early_drafter: Optional[M3Eagle3Draft] = None + self._m3_fc_early_stream: Optional[torch.cuda.Stream] = None + self._m3_force_vec_warm = False + # The streamed copy of the drafter's BF16 head shard the fused draft loop reads, built before the warmup. + self._m3_draft_head = None + # Per step: the next step's input tokens [nreq, D + 1], column 0 written by the fused verify (or the seed) and + # columns 1.. by the fused draft passes; set as _m3_next_new_tokens once the fused draft loop filled it. + self._m3_next_new: Optional[torch.Tensor] = None + self._m3_next_new_tokens: Optional[torch.Tensor] = None + # Per step: the first draft inputs are the fused ones (accepted tokens as they lie, fc output of every row). + self._m3_fused_first_inputs = False + + def capture_hidden_states( + self, spec_metadata, layer_id: int, hidden_states, residual=None + ) -> None: + """The Eagle3 capture with this step's capture norm; on a step armed for the early fc fork + (``arm_capture_norm``), the last captured layer then forks the drafter's first fc onto the aux stream, where + it overlaps the remaining target layers and the head (joined by ``_m3_target_accept``).""" + _eagle3_capture(spec_metadata, layer_id, hidden_states, residual, self._m3_capture_norm) + drafter = self._m3_fc_early_drafter + layers = spec_metadata.layers_to_capture + if drafter is not None and layers and layer_id == layers[-1]: + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = self._m3_fc_fork( + spec_metadata, spec_metadata.num_capture_tokens, drafter + ) + + def prebuild(self, lm_head: nn.Module, draft: Optional[M3Eagle3Draft]) -> None: + """The fused draft loop's head copy, the drafter linears' low-M kernels and the drafter attention's KV-split + cap, from the loaded weights, before the executor's warmup, on every rank.""" + start = time.perf_counter() + if self.prebuild_draft_head(draft): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] fused draft head built before warmup: %.1f s" + % (time.perf_counter() - start) + ) + start = time.perf_counter() + built = self.prebuild_drafter_linears(draft) + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] drafter low-M GEMMs built before warmup: %d calls, %.1f s" + % (built, time.perf_counter() - start) + ) + + def _forward_impl( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + """Mirrors Eagle3OneModelWorker._forward_impl, with the fused verify branch.""" + runtime_draft_len = spec_metadata.runtime_draft_len + # skip the draft forward if the runtime draft length is 0 + if runtime_draft_len == 0: + return self.skip_drafting( + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + ) + + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + + raw_logits = logits + + self._m3_first_gather_ids = None + self._m3_next_new = None + self._m3_next_new_tokens = None + self._m3_fused_first_inputs = False + if logits is None: + # tail_eligible() skipped the gathered target logits: verify from this rank's logits shard. + accepted_tokens, num_accepted_tokens, self._m3_first_gather_ids = ( + self._m3_target_accept( + hidden_states, spec_metadata, batch_size, runtime_draft_len, draft_model + ) + ) + else: + self._execute_guided_decoder_if_present(logits) + + # Sample and accept tokens. ``input_ids`` is required by the relaxed-acceptance path (scans for + # thinking-phase tokens); ignored otherwise. + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + input_ids, logits, attn_metadata, spec_metadata + ) + if self.draft_only_eligible(attn_metadata, spec_metadata, draft_model): + self._m3_seed_first_pass(num_accepted_tokens, spec_metadata, batch_size) + + # Mamba hybrid models need state updates after token acceptance (a no-op on this target's cache manager). + if self._is_mamba_hybrid_cache is None: + self._is_mamba_hybrid_cache = isinstance( + attn_metadata.kv_cache_manager, MambaHybridCacheManager + ) + if num_gens > 0 and self._is_mamba_hybrid_cache: + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices, + ) + + sa_manager = getattr(spec_metadata.spec_resource_manager, "sa_manager", None) + if self.sa_enhancer is not None and sa_manager is not None: + self.sa_enhancer.extend_and_prepare( + sa_manager=sa_manager, + request_ids=spec_metadata.request_ids, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_gens=num_gens, + num_contexts=num_contexts, + max_draft_len=runtime_draft_len, + ) + + if self._uses_external_shared_target_kv: + next_draft_tokens = self._forward_external_shared_target_kv_draft_loop( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + ) + else: + # Save the old attn_metadata and spec_metadata + self._prepare_attn_metadata_for_spec_dec(attn_metadata) + + # Prepare inputs for the 1st draft model forward + position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + # Predict draft tokens. ``original_all_rank_num_tokens`` is saved here so the post-loop restore (below) + # can put attn_metadata back into a state the target model expects. + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + + # Get the draft KV cache manager if using separate layouts + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + + next_draft_tokens = self._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + # restore attn_metadata to support cuda graph + self._restore_attn_metadata_from_spec_dec(attn_metadata) + # restore all_rank_num_tokens for attention DP + if original_all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + + # prepare next new tokens to support overlap scheduler + if self._m3_next_new_tokens is not None: + # Filled in place by the fused verify or the seed (column 0) and the fused draft passes (columns 1..). + next_new_tokens = self._m3_next_new_tokens + else: + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + attn_metadata.use_spec_decoding = True + + # Commit model-owned target state only after draft preparation has succeeded; SpecWorkerBase aborts the + # still-pending snapshots on any exception above this point. + if self._auxiliary_state_handlers: + self.commit_auxiliary_speculative_states( + num_accepted_tokens, + attn_metadata.mamba_metadata.state_indices[:batch_size], + num_contexts, + ) + + return { + "logits": raw_logits, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + } + + def prepare_1st_drafter_inputs( + self, + input_ids, + position_ids, + hidden_states, + accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + ): + """The fused first draft inputs on a step that verified with the fused kernels (or was seeded for them), + else the in-tree ones.""" + fused = self._m3_first_drafter_inputs( + position_ids, + accepted_tokens, + input_ids.shape[0], + attn_metadata.num_contexts, + spec_metadata, + draft_model, + ) + self._m3_fused_first_inputs = fused is not None + if fused is not None: + fused.update(attn_metadata=attn_metadata, spec_metadata=spec_metadata) + return fused + return super().prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + def _m3_structural_ok(self, attn_metadata, spec_metadata) -> bool: + """The batch conditions shared by the fused verify and the fused first drafter inputs: a decode-only plain + Eagle3 step without attention DP, guided decoding, relaxed acceptance or a suffix-automaton enhancer, of at + most DECODE_MAX_TOKENS verify rows. Host state only, so the TP ranks agree.""" + if self.is_mtp_eagle or self.use_dynamic_tree: + return False + if self._uses_external_shared_target_kv or self.sa_enhancer is not None: + return False + if ( + self.guided_decoder is not None + or self.mapping is None + or self.mapping.enable_attention_dp + ): + return False + if getattr(self.spec_config, "use_relaxed_acceptance_for_thinking", False): + return False + if attn_metadata.num_contexts != 0 or attn_metadata.num_seqs == 0: + return False + draft_len = spec_metadata.runtime_draft_len + if draft_len == 0 or attn_metadata.num_seqs * (draft_len + 1) > DECODE_MAX_TOKENS: + return False + return not spec_metadata.wants_advanced_draft_sampling + + def tail_eligible(self, attn_metadata, spec_metadata, lm_head: nn.Module) -> bool: + """Whether this step verifies with the fused kernels instead of the gathered target logits. + + Decode-only Eagle3 batches (``_m3_structural_ok``) over an unpadded vocab-sharded bias-free target head, all + greedy at a request count ``SHARD_ACCEPT`` covers, or sampled under the temperature + top-p contract (see + _m3_sampled_verify_ok), with nothing else reading the target logits. Every input is host state shared by all + TP ranks, so the ranks agree. + """ + if not self._m3_structural_ok(attn_metadata, spec_metadata): + return False + if getattr(spec_metadata, "enable_penalty", False): + return False + if not _sharded_head(lm_head) or getattr(lm_head, "bias", None) is not None: + return False + draft_len = spec_metadata.runtime_draft_len + if spec_metadata.is_all_greedy_sample: + key = (draft_len, attn_metadata.num_seqs) + if not SHARD_ACCEPT.covers(key): + uncovered( + SHARD_ACCEPT.path, + f"draft length {draft_len} at {attn_metadata.num_seqs} requests", + ) + return False + elif not self._m3_sampled_verify_ok(spec_metadata): + return False + self._m3_target_lm_head = lm_head + return True + + def draft_only_eligible(self, attn_metadata, spec_metadata, draft_model) -> bool: + """Whether a step verified by the in-tree path can still take the fused first drafter inputs: the structural + conditions of tail_eligible, without the target-sampling ones.""" + return self._m3_structural_ok(attn_metadata, spec_metadata) and draft_model is not None + + def _m3_seed_first_pass(self, num_accepted_tokens, spec_metadata, batch_size: int) -> None: + """The first draft pass's gather rows (the last accepted row of each request) from the in-tree acceptance, as + _m3_target_accept leaves them.""" + k1 = spec_metadata.runtime_draft_len + 1 + n = num_accepted_tokens[:batch_size].to(torch.int64) + rows = spec_metadata.batch_indices_cuda[:batch_size].to(torch.int64) + self._m3_first_gather_ids = rows * k1 + n - 1 + + def _m3_first_drafter_inputs( + self, + position_ids, + accepted_tokens, + num_tokens: int, + num_contexts: int, + spec_metadata, + draft_model: Optional[M3Eagle3Draft], + ): + """The in-tree first draft inputs of a decode-only step, with the per-layer fc_norm RMSNorms of + ``apply_eagle3_fc`` in one kernel (unless the capture already applied them) before the fc; the input ids are + the accepted tokens as they lie. None to use the in-tree path.""" + if ( + self._m3_first_gather_ids is None + or self.is_mtp_eagle + or num_contexts != 0 + or draft_model is None + ): + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + ): + return None + cache = self._m3_fc_norm_w + if cache is None or cache[0] is not draft_model.w["fc_norm0"]: + if torch.cuda.is_current_stream_capturing(): + capture_missed("spec_tail", None, "the stacked Eagle3 fc_norm weights") + return None + cache = ( + draft_model.w["fc_norm0"], + torch.stack( + [draft_model.w[f"fc_norm{i}"].detach() for i in range(layers)] + ).contiguous(), + ) + self._m3_fc_norm_w = cache + hidden = self._m3_fc_hidden + self._m3_fc_hidden = None + if hidden is None or hidden.shape[0] != num_tokens: + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + hs = m3_spec_tail.chunk_norm(hs, cache[1], draft_model.rms_eps) + hidden = draft_model.fc(hs) + input_ids = accepted_tokens[:, : spec_metadata.runtime_draft_len + 1].reshape(-1) + return {"input_ids": input_ids, "position_ids": position_ids, "hidden_states": hidden} + + def _m3_fc_fork( + self, + spec_metadata, + num_tokens: int, + draft_model: Optional[M3Eagle3Draft], + ) -> Optional[torch.cuda.Stream]: + """Launch the first draft pass's fc_norm + fc on a forked stream beside the target head GEMM (after everything + queued on the current stream); the result is picked up by _m3_first_drafter_inputs. Returns the stream to + join, or None when the fc path does not apply (the first drafter inputs then compute it).""" + self._m3_fc_hidden = None + if draft_model is None: + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + or self._m3_fc_norm_w is None + ): + return None + aux = self._m3_aux_stream + if aux is None: + aux = self._m3_aux_stream = torch.cuda.Stream() + aux.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(aux): + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + hs = m3_spec_tail.chunk_norm( + hs, self._m3_fc_norm_w[1], draft_model.rms_eps, use_pdl=False + ) + self._m3_fc_hidden = draft_model.fc(hs) + return aux + + def _m3_fc_join(self, fork: torch.cuda.Stream) -> None: + """Join the fc fork onto the current stream and keep its output's memory alive for that stream. + + _m3_fc_hidden is allocated in fork (aux) stream order and read on the current stream. Without the record, + dropping it returns its block to the aux stream's pool, where later aux work can take it while the current + stream's reads are still queued (eager warm-up / mixed steps). Nothing is recorded under graph capture, as in + maybe_execute_in_parallel: a captured graph relies on the next fork's aux.wait_stream(main) to order the + reuse.""" + main = torch.cuda.current_stream() + main.wait_stream(fork) + hidden = self._m3_fc_hidden + if ( + hidden is not None + and hidden.is_cuda + and not (torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing()) + ): + hidden.record_stream(main) + + def arm_capture_norm( + self, attn_metadata, spec_metadata, lm_head: nn.Module, draft_model: Optional[M3Eagle3Draft] + ) -> None: + """Before every target forward: whether this step's Eagle3 captures apply the drafter's fc_norm themselves. + + On a fused-verify step whose fc_norm weights are already stacked (by an earlier fused first pass), the + capture writes rmsnorm(h + r) * fc_norm[i] (one kernel per captured layer instead of add + copy) and the fused + first drafter inputs read it without their own chunk_norm. A fused-verify step of at most + ``FC_EARLY_MAX_ROWS`` verify rows is also armed to fork the drafter's first fc at the last captured layer + (``capture_hidden_states``). + """ + self._m3_capture_norm = None + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = None + if draft_model is None: + return + if not self.tail_eligible(attn_metadata, spec_metadata, lm_head): + return + # The fork reads every captured column at the last captured layer, so a final-state capture, which is written + # after the model returns, keeps the fc at its stock place. + if ( + attn_metadata.num_seqs * (spec_metadata.runtime_draft_len + 1) <= FC_EARLY_MAX_ROWS + and FINAL_HIDDEN_LAYER_ID not in spec_metadata.layers_to_capture + ): + self._m3_fc_early_drafter = draft_model + fc_norm_w = self._m3_fc_norm_w + if ( + fc_norm_w is not None + and fc_norm_w[0] is draft_model.w["fc_norm0"] + and len(spec_metadata.layers_to_capture) == fc_norm_w[1].shape[0] + ): + self._m3_capture_norm = (fc_norm_w[1], draft_model.rms_eps) + + @staticmethod + def _m3_sampled_verify_ok(spec_metadata) -> bool: + """Whether sampled verify rows can use the fused sampled op: only under advanced_sampling_mode=no_topk, the + deploy-time guarantee that no row uses top-k or min-p (that op applies neither).""" + mode = getattr(spec_metadata, "advanced_sampling_mode", None) + return mode is not None and mode.skips_top_k and not mode.skips_top_p + + def _m3_target_accept( + self, + hidden_states, + spec_metadata, + batch_size: int, + draft_len: int, + draft_model: Optional[M3Eagle3Draft] = None, + ): + """Fused verification (see tail_eligible); same outputs as the strict acceptance path.""" + force_total = 0 + force_vec = None + if self.force_num_accepted_tokens != 0.0: + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if not torch.cuda.is_current_stream_capturing(): + # The draw pool is built by an H2D copy and the draw kernel JIT-compiles per block width, so both + # happen here, before the first captured step (a scratch counter keeps the real draw sequence). + self._ensure_force_accept_rng_state(hidden_states.device) + if not self._m3_force_vec_warm: + scratch = torch.zeros(1, dtype=torch.int64, device=hidden_states.device) + dummy = torch.empty(64, dtype=torch.int32, device=hidden_states.device) + for block in (16, 32, 64): + m3_force_vec( + self._force_accept_rng_pool, + scratch, + dummy, + block, + 1, + 2, + 0.5, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + self._m3_force_vec_warm = True + if not is_warmup: + force_total = min(int(self.force_num_accepted_tokens) + 1, draft_len + 1) + force_vec = self._m3_force_vec( + batch_size, force_total, draft_len + 1, hidden_states.device + ) + draft_tokens = spec_metadata.draft_tokens.view(batch_size, draft_len) + if draft_tokens.dtype != torch.int32: + draft_tokens = draft_tokens.to(torch.int32) + # Column 0 receives the last accepted token, the next step's first input (the fused draft loop adds the rest). + self._m3_next_new = torch.empty( + (batch_size, draft_len + 1), dtype=torch.int32, device=hidden_states.device + ) + rows = hidden_states[: batch_size * (draft_len + 1)] + lm_head = self._m3_target_lm_head + batch_idx = spec_metadata.batch_indices_cuda[:batch_size] + # The Eagle3 fc on the aux stream: the fork made at the last captured layer (none at FC_EARLY_MAX_ROWS = 0), + # else one beside the head. + fork = self._m3_fc_early_stream + self._m3_fc_early_stream = None + if fork is None: + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model) + # This rank's vocabulary shard of the target logits: the head's own local BF16 GEMM, whose all-gather the + # fused kernels replace (they exchange per-row records across ranks, not logits). + logits = lm_head.apply_linear(rows, None) + vocab_offset = self.mapping.tp_rank * lm_head.weight.shape[0] + if spec_metadata.is_all_greedy_sample: + result = m3_shard_accept_kernel.shard_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + path = SHARD_ACCEPT.path + else: + n_rows = rows.shape[0] + seeds, offsets = self._rng_state_per_token(spec_metadata, n_rows) + result = m3_sampled_accept_kernel.sampled_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + spec_metadata.temperatures[:n_rows], + spec_metadata.top_ks[:n_rows], + spec_metadata.top_ps[:n_rows], + None, + seeds, + offsets, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + path = "spec_tail" + if fork is not None: + self._m3_fc_join(fork) + if result is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, None, "the fused verify's exchange workspace") + raise RuntimeError( + "M3 fused verify unavailable after tail_eligible() accepted the step" + ) + return result + + def _m3_force_vec( + self, batch_size: int, base_total: int, max_total: int, device + ) -> Optional[torch.Tensor]: + """Per-request forced accepted counts for a fractional forced acceptance, else None. + + Draws exactly as _apply_force_accepted_tokens does for a decode-only batch (same fixed-seed pool, one counter + bump per step, the same (counter, slot) hash), so every TP rank and the generic path agree on the counts. + """ + frac = self.force_num_accepted_tokens - int(self.force_num_accepted_tokens) + if frac <= 0.0 or base_total >= max_total: + return None + self._ensure_force_accept_rng_state(device) + out = torch.empty(batch_size, dtype=torch.int32, device=device) + return m3_force_vec( + self._force_accept_rng_pool, + self._force_accept_rng_counter, + out, + batch_size, + base_total, + max_total, + float(frac), + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + max(16, 1 << (batch_size - 1).bit_length()), + ) + + # ---------------------------------------------------------------------------------------------- fused draft loop + + def prebuild_draft_head(self, draft_model: Optional[M3Eagle3Draft]) -> bool: + """Build the streamed copy of the drafter's BF16 head shard that the fused draft loop reads (rebuilt from the + current weight values on every call), on every rank. False when the drafter is not one the loop serves.""" + self._m3_draft_head = None + if draft_model is None or not self._m3_drafter_supported(draft_model): + return False + self._m3_draft_head = m3_draft_step_kernel.prepare_head( + draft_model.lm_head, self.mapping.tp_rank + ) + return True + + def prebuild_drafter_linears(self, draft_model: Optional[M3Eagle3Draft]) -> int: + """Compile every drafter linear call ``DRAFT_LOWM`` covers before the executor's warmup, on every + rank. Returns the number of low-M calls compiled.""" + if draft_model is None: + return 0 + return _prebuild_draft_lowm(draft_model) + + def _m3_drafter_supported(self, draft_model: M3Eagle3Draft) -> bool: + """The drafter the fused draft loop runs (``M3Eagle3Draft`` validates its shape at construction) over this + worker's tensor-parallel group with a bias-free vocab-sharded BF16 head without padding.""" + if self.mapping is None or self.mapping.tp_size <= 1 or self.mapping.enable_attention_dp: + return False + if self._d2t is not None: + return False + lm_head = draft_model.lm_head + return _sharded_head(lm_head) and m3_draft_step_kernel.head_supported(lm_head) + + def _m3_fused_loop_ok( + self, attn_metadata, spec_metadata, batch_size: int, draft_model: Optional[M3Eagle3Draft] + ) -> bool: + """Whether this step drafts with the fused loop: a decode-only step the fused verify or the seed prepared + (``_m3_first_gather_ids``) with the fused first draft inputs, a supported drafter, a (draft length, requests) + ``DRAFT_STEP`` covers, the prebuilt head and the exchange workspace. Host state only, so the TP ranks + agree.""" + if ( + self._m3_first_gather_ids is None + or not self._m3_fused_first_inputs + or draft_model is None + ): + return False + path = DRAFT_STEP.path + if not self._m3_drafter_supported(draft_model) or not hasattr( + attn_metadata, "kv_lens_cuda" + ): + uncovered( + path, "the drafter or its attention metadata is not one the fused loop serves" + ) + return False + draft_len = spec_metadata.runtime_draft_len + if not DRAFT_STEP.covers((draft_len, batch_size)): + uncovered(path, f"draft length {draft_len} at {batch_size} requests") + return False + head = self._m3_draft_head + if head is None or head.weight is not draft_model.lm_head.weight: + missed( + path, + batch_size, + "the drafter head's streamed copy is not built for its current weight (prebuild)", + ) + return False + if m3_draft_step_kernel.workspace(self.mapping, head.wf.device) is None: + capture_missed(path, None, "the draft-step exchange workspace") + return False + return True + + def _forward_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ): + """The fused draft loop on a step it covers (``_m3_fused_loop_ok``), else the in-tree loop.""" + if self._m3_fused_loop_ok(attn_metadata, spec_metadata, batch_size, draft_model): + return self._m3_fused_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ) + return super()._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + + def _m3_fused_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model: M3Eagle3Draft, + draft_kv_cache_manager, + batch_size: int, + num_accepted_tokens, + ): + """Mirrors Eagle3OneModelWorker._forward_linear_draft_loop for a decode-only Eagle3 step, with each pass's + glue fused (see ``_m3_fused_loop_ok``): + + * the draft input -- the embedding lookup, the layer's input norm and hidden norm and their concat -- in one + launch (``m3_spec_tail.draft_input``; the second pass gathers the last accepted rows there); + * the drafter layer's attention and MLP, each followed by the all-reduce with the residual add and the next + RMSNorm (the post-attention norm, then the drafter's final norm); + * one kernel for the BF16 head over the head rows (pass 0: the last accepted row of each request), the + cross-rank argmax and the bookkeeping the in-tree loop does between passes: the draft token, the next pass's + input id and position, the draft KV lengths (and pass 0's sequence lengths) (``m3_draft_step``). + + The host metadata transitions are the in-tree loop's. The drafts land in columns 1.. of the step's next-input + tokens (column 0 is the last accepted token), which ``_forward_impl`` returns as they lie. + """ + draft_len = spec_metadata.runtime_draft_len + table = draft_model.embed_table + dev = table.device + head = self._m3_draft_head + gather_ids = self._m3_first_gather_ids + ids = inputs["input_ids"] + hidden = inputs["hidden_states"] + pos = inputs["position_ids"] + if ids.dtype != torch.int32: + ids = ids.to(torch.int32) + if pos.dtype != torch.int32: + pos = pos.to(torch.int32) + num_acc = num_accepted_tokens[:batch_size] + if num_acc.dtype != torch.int32: + num_acc = num_acc.to(torch.int32) + next_new = self._m3_next_new + if next_new is None: + # Seeded from the in-tree acceptance: the last accepted token is the first-pass input at the gather row. + next_new = torch.empty((batch_size, draft_len + 1), dtype=torch.int32, device=dev) + next_new[:, 0] = ids[gather_ids] + draft_tokens = next_new[:, 1:] + w = draft_model.w + eps = draft_model.rms_eps + + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager) as md: + for i in range(draft_len): + x, residual = m3_spec_tail.draft_input( + ids, + table, + hidden, + gather_ids if i == 1 else None, + w["in_norm"], + eps, + w["hidden_norm"], + eps, + ) + h = draft_model.attention(x, md, reduce=False) + h, residual = draft_model.ar_norm(h, residual, "post_norm") + h = draft_model.mlp(h, reduce=False) + h, to_save = draft_model.ar_norm(h, residual, "final_norm") + if draft_model.return_hidden_post_norm: + to_save = h + ids_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + pos_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + if not m3_draft_step_kernel.draft_step( + head, + h, + batch_size, + gather_ids if i == 0 else None, + self.mapping, + draft_tokens, + i, + ids_next, + pos, + pos_next, + md.kv_lens_cuda, + num_acc if i == 0 else None, + md._seq_lens_cuda if i == 0 else None, + draft_len, + ): + raise RuntimeError( + "M3 fused draft step unavailable after _m3_fused_loop_ok() accepted the step" + ) + ids, pos, hidden = ids_next, pos_next, to_save + if i == 0: + # The draft step set the device sequence lengths to 1 and rolled the KV lengths back. + md._seq_lens[:batch_size].fill_(1) + md.on_update() + has_kv_cache = md.kv_cache_manager is not None + if has_kv_cache and hasattr(md, "host_request_types"): + md.host_request_types[: md.num_contexts].fill_(1) + if has_kv_cache: + md.num_contexts = 0 + self._prepare_flash_mla_generation_layout(md, 0, batch_size) + md.update_for_spec_dec() + md.use_spec_decoding = False + else: + md.update_for_spec_dec() + self._m3_next_new_tokens = next_new + return draft_tokens + + +def _build_speculation(shell, model_config: ModelConfig): + """The drafter config, the drafter and the worker of one-model Eagle3 (the in-tree one-engine construction for + the llama3 drafter architecture).""" + spec = model_config.spec_config + if ( + not spec.spec_dec_mode.use_one_engine() + or not spec.spec_dec_mode.is_eagle3_one_model() + or spec.eagle3_model_arch != "llama3" + or getattr(spec, "use_dynamic_tree", False) + ): + raise NoTarget( + "this MiniMax-M3 target speculates with linear one-model Eagle3 (llama3 drafter) only" + ) + draft_config = ModelConfig.from_pretrained( + spec.speculative_model, + trust_remote_code=True, + attn_backend=model_config.attn_backend, + moe_backend=_get_requested_draft_moe_backend(model_config, spec), + mapping=model_config.mapping, + spec_config=model_config.spec_config, + max_num_tokens=model_config.max_num_tokens, + moe_max_num_tokens=model_config.moe_max_num_tokens, + ) + _set_draft_kv_cache_quant_algo(draft_config, model_config) + draft_config.extra_attrs = model_config.extra_attrs + use_separate = should_use_separate_draft_kv_cache(spec) + draft = M3Eagle3Draft( + shell.model, + draft_config, + model_config.pretrained_config.num_hidden_layers, + spec.num_capture_layers, + ) + worker = M3Eagle3Worker(spec, model_config.mapping, use_separate_draft_kv_cache=use_separate) + worker.set_draft_model(draft) + shell.epilogue.append(draft) + shell.epilogue.append(worker) + return draft_config, draft, worker, use_separate + + +# -------------------------------------------------------------------------------------------------------------------- +# Runtime hooks and the shell +# -------------------------------------------------------------------------------------------------------------------- + + +def _runtime_hooks(spec_fold: Optional[M3SpecStepFold]) -> RuntimeHooks: + """The step-loop hooks this target opts into, measured on GB300 AgentX decode (tp4 and tp8). + + * ``pack_step_inputs``: a decode step's small input copies go out as one upload and one scatter at the graph + replay instead of a chain of copies the next graph waits on (~30 us of idle GPU at one request). + * ``skip_unchanged_h2d``: steady decode re-uploads slot indices, query lengths and page tables that rarely change. + * ``loop_on_execution_stream``: the overlap loop's own stream work and the forward on one stream, so step N+1's + inputs wait on graph N in program order rather than through two cross-stream event hops. + * ``merge_tuned_tactics_across_ranks``: tensor-parallel peers capture the same GEMM tactics (rank-local profiling + noise otherwise splits them, and the slowest rank sets every step). + * ``skip_greedy_rng_upload``: all-greedy speculative batches never read the Philox seeds. + * ``spec_step_fold``: the spec sampler's slot stores and the engine's next-step input gathers in one launch each. + * MSA: identity selections for short sequences (and short/long graph families), the one-launch KV-length update, + unchanged block-table copies skipped, pool views built once. + """ + return RuntimeHooks( + pack_step_inputs=True, + skip_unchanged_h2d=True, + loop_on_execution_stream=True, + merge_tuned_tactics_across_ranks=True, + skip_greedy_rng_upload=True, + spec_step_fold=spec_fold, + msa=MsaRuntimePolicy( + short_sequence_identity=True, + kv_update=_msa_kv_update_kernel.msa_kv_update, + skip_unchanged_block_offsets=True, + reuse_pool_views=True, + ), + ) + + +class ModelingV2MiniMaxM3Nvfp4Sm103Tp8( + DecoderModelForCausalLM[MiniMaxM3Tp8ModelingV2, PretrainedConfig] +): + """MiniMax-M3 NVFP4 / sm_103 / tp8, text only, with one-model Eagle3 or no speculation. + + The shell owns the vocab-parallel BF16 ``lm_head`` and the logits processor (``DecoderModelForCausalLM``); with + speculation, the drafter (``M3Eagle3Draft``, its weights from the speculative checkpoint through + ``load_draft_weights``) and the worker (``M3Eagle3Worker``) that verifies and drafts each step. + """ + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + assert torch.cuda.get_device_capability() == _SM, ( + f"this target is built for sm_{_SM[0]}{_SM[1]}, the device is sm_" + f"{''.join(map(str, torch.cuda.get_device_capability()))}" + ) + # Under modeling_v2='require' a covered decode shape that cannot take its fast path raises. + set_require(resolved_mode(type(self)) is ModelingV2Mode.REQUIRE) + model_config = _text_model_config(model_config) + cfg = model_config.pretrained_config + spec_config = getattr(model_config, "spec_config", None) + if model_config.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4: + # The one-model Eagle3 layer has no NVFP4 trtllm-gen decode kernel: keep its modules and the shared + # draft cache on FP8 while the target stays NVFP4. + model_config.extra_attrs["draft_kv_cache_quant_algo_override"] = QuantAlgo.FP8 + super().__init__( + MiniMaxM3Tp8ModelingV2(model_config), + config=model_config, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + self.spec_config = spec_config + self.draft_model: Optional[M3Eagle3Draft] = None + self.draft_config: Optional[ModelConfig] = None + self.spec_worker: Optional[M3Eagle3Worker] = None + self.use_separate_draft_kv_cache = False + self.layer_idx = FINAL_HIDDEN_LAYER_ID + spec_fold = None + if spec_config is not None: + ( + self.draft_config, + self.draft_model, + self.spec_worker, + self.use_separate_draft_kv_cache, + ) = _build_speculation(self, model_config) + spec_fold = M3SpecStepFold() + self.runtime_hooks = _runtime_hooks(spec_fold) + + def load_weights(self, weights: Dict, *args, **kwargs) -> None: + _weights.MODEL_WEIGHTS.load(self, weights) + + def load_draft_weights(self, weights: Dict, weight_mapper=None) -> None: + """The drafter's checkpoint (``speculative_model``), loaded by the engine after the target's.""" + _weights.DRAFT_WEIGHTS.load(self.draft_model, weights) + + def set_guided_decoder(self, guided_decoder) -> bool: + """Mirrors ``SpecDecOneEngineForCausalLM.set_guided_decoder``: guided decoding with speculation is installed + on the worker.""" + if hasattr(self.spec_worker, "set_guided_decoder"): + return self.spec_worker.set_guided_decoder(guided_decoder) + return False + + def cache_derived_state(self) -> None: + """Build the targets and every decode kernel once the weights are real, before the executor's warmup; with + speculation also the fused draft loop's head copy and the drafter linears' low-M kernels.""" + super().cache_derived_state() + self.model.post_load_weights() + if self.spec_worker is not None: + self.spec_worker.prebuild(self.lm_head, self.draft_model) + + def post_load_weights(self) -> None: + super().post_load_weights() + self.model.post_load_weights() + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + return_context_logits: bool = False, + spec_metadata: Optional[SpecMetadata] = None, + resource_manager: Optional[Any] = None, + **kwargs, + ) -> Any: + """One engine step: the target forward, then the logits (or, with speculation, the worker's verify and + draft). Text only.""" + multimodal_params = kwargs.pop("multimodal_params", None) or [] + if any(p is not None and getattr(p, "multimodal_data", None) for p in multimodal_params): + raise ValueError( + "this MiniMax-M3 target is text only; the request carries multimodal input" + ) + worker = self.spec_worker + capture = _plain_capture + if worker is not None: + capture = worker.capture_hidden_states + if spec_metadata is not None: + worker.arm_capture_norm( + attn_metadata, spec_metadata, self.lm_head, self.draft_model + ) + hidden_states = self.model( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + capture=capture, + **kwargs, + ) + if getattr(attn_metadata, "csa2_remote_tail_mode", None) == "source": + # A context worker stopped at the encoder/decoder boundary: a disposable token for the context-only + # completion, and no draft worker. + return self.logits_processor.forward(hidden_states, self.lm_head, attn_metadata, False) + if spec_metadata is not None and spec_metadata.is_layer_capture(FINAL_HIDDEN_LAYER_ID): + capture(spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None) + if attn_metadata.padded_num_tokens is not None: + hidden_states = hidden_states[: attn_metadata.num_tokens] + if worker is None: + return self.logits_processor.forward( + hidden_states, self.lm_head, attn_metadata, return_context_logits + ) + # The fused verify computes this rank's logits shard itself and needs no gathered logits. + if worker.tail_eligible(attn_metadata, spec_metadata, self.lm_head): + logits = None + else: + logits = self.logits_processor.forward( + hidden_states[spec_metadata.gather_ids], self.lm_head, attn_metadata, True + ) + spec_input_ids = input_ids if input_ids is not None else kwargs.get("orig_input_ids") + spec_position_ids = position_ids + if attn_metadata.padded_num_tokens is not None: + if spec_input_ids is not None: + spec_input_ids = spec_input_ids[: attn_metadata.num_tokens] + if position_ids is not None: + spec_position_ids = _slice_spec_position_ids(position_ids, attn_metadata.num_tokens) + return worker( + input_ids=spec_input_ids, + position_ids=spec_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, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/__init__.py new file mode 100644 index 000000000000..1059445198db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this target. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/activation/__init__.py new file mode 100644 index 000000000000..b8c7426e3933 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/activation/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Activation ops of this target: the gated MLPs' fused SwiGLU.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/activation/m3_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/activation/m3_silu_and_mul.py new file mode 100644 index 000000000000..7a494d97d788 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/activation/m3_silu_and_mul.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The gated MLPs' fused SwiGLU: ``silu_and_mul`` with MiniMax-M3's SwiGLU-OAI gain, offset and clamp.""" + +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 M3SiluAndMul(OpWrapper): + """``out = g * sigmoid(alpha * g) * (u + beta)`` over ``x = [g | u]`` (bf16 [M, 2I] -> [M, I]), with + ``g = min(gate, limit)`` and ``u = clamp(up, -limit, limit)``; ``limit``, ``alpha`` and ``beta`` None give plain + SwiGLU ``silu(gate) * up``. + + The MXFP8 dense MLPs (layers 0-2, intermediate 3072 per rank) and the MoE's shared expert (intermediate 768 per + rank) run it at MiniMax-M3's (limit 7, alpha 1.702, beta 1); the BF16 Eagle3 drafter's MLP (intermediate 4608 per + rank) runs it plain. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert, one decode token", spec=dict(m=1, inter=384, oai=True)), + Cell(why="shared expert, a 48-row decode verify", spec=dict(m=48, inter=384, oai=True)), + Cell(why="dense MLP, a prefill chunk", spec=dict(m=2048, inter=1536, oai=True)), + Cell(why="dense MLP, one decode token", spec=dict(m=1, inter=1536, oai=True)), + Cell( + why="drafter MLP (plain SwiGLU), eight decode rows", + spec=dict(m=8, inter=2304, oai=False), + ), + Cell( + why="drafter MLP (plain SwiGLU), a prefill chunk", + spec=dict(m=1024, inter=2304, oai=False), + ), + ) + + note = """ + The kernel computes in fp32 and rounds once to bf16; `reference` does the same arithmetic in fp32, so the two + differ only by the fast sigmoid's last bits (the test's band). + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + + def reference( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + gate, up = x.float().chunk(2, dim=-1) + if limit is not None: + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + a = 1.0 if alpha is None else alpha + b = 0.0 if beta is None else beta + return (gate * torch.sigmoid(a * gate) * (up + b)).to(x.dtype) + + def is_valid( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> None: + assert x.stride(-1) == 1, "the gate/up halves are read from contiguous rows" + + +m3_silu_and_mul = M3SiluAndMul() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/__init__.py new file mode 100644 index 000000000000..a7f7fee7a234 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Attention ops of this target: MSA block selection, the sparse decode and the KV-length update.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_dense_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_dense_decode.py new file mode 100644 index 000000000000..31cbd5462d43 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_dense_decode.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' decode attention over the full page table with valid sequence-length bounds.""" + +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.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( # noqa: E501 + minimax_m3_trtllm_gen_dense_decode, +) + +__all__ = ["m3_dense_decode"] + +#: The TP8 dense layers' geometry: 8 query heads and one KV head of dim 128 per rank, an FP8 KV pool. +HEADS = 8 +KV_HEADS = 1 +HEAD_DIM = 128 +#: Valid upper bounds used to compare short-context and model-maximum dispatch. +SHORT_BOUND = 16384 +MODEL_MAX = 1 << 20 + + +def dense_attention_reference( + q: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, +) -> torch.Tensor: + """fp32 causal decode attention: query row ``b * qlen + i`` of request b attends its first + ``seq_lens[b] - qlen + i + 1`` keys. ``keys`` / ``values``: [batch, kv_heads, max_len, head_dim] (the request's + tokens in order). Returns fp32 [total_q, heads, head_dim].""" + batch, kv_heads, max_len, _ = keys.shape + heads = q.shape[1] + group = heads // kv_heads + out = torch.zeros(q.shape, dtype=torch.float32, device=q.device) + pos = torch.arange(max_len, device=q.device) + with true_fp32_matmul(): + for b in range(batch): + full = int(seq_lens[b]) + rows = q[b * decode_query_len : (b + 1) * decode_query_len].float() # [qlen, heads, d] + limit = full - decode_query_len + 1 + torch.arange(decode_query_len, device=q.device) + valid = pos[None, :] < limit[:, None] # [qlen, max_len] + for h in range(kv_heads): + qh = rows[:, h * group : (h + 1) * group] # [qlen, group, d] + logits = torch.einsum("qgd,kd->qgk", qh, keys[b, h].float()) * sm_scale + logits = logits.masked_fill(~valid[:, None, :], float("-inf")) + probs = torch.softmax(logits, dim=-1) + out[ + b * decode_query_len : (b + 1) * decode_query_len, h * group : (h + 1) * group + ] = torch.einsum("qgk,kd->qgd", probs, values[b, h].float()) + return out + + +def gather_paged_kv( + kv_cache_manager, layer_idx: int, block_table: torch.Tensor, seq_lens: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """The K and V of each request in token order, [batch, kv_heads, max_len, head_dim], from the flat sub-page pool + (``get_kv_subpage_pool``: K of slot s at sub-page s * subpages_per_slot, V one sub-page later).""" + pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + page = pool.shape[2] + pages = (int(seq_lens.max()) + page - 1) // page + slots = block_table[:, :pages].long() + k = pool[slots * subpages] # [batch, pages, kv_heads, page, d] + v = pool[slots * subpages + 1] + batch, _, kv_heads, _, d = k.shape + + def order(t): + return t.permute(0, 2, 1, 3, 4).reshape(batch, kv_heads, pages * page, d) + + return order(k), order(v) + + +class M3DenseDecode(OpWrapper): + """Decode over each request's full causal context in the paged FP8 pool. + + The sequence-length bound covers every request and the BF16 output is written in place. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over a 20000-token context, its full sequence-length bound", + spec=dict(kv_lens=(20000,), qlen=4, max_seq_len=20000, tight_bound=True), + ), + Cell( + why="one request, one row, over a 65536-token context with its full sequence-length bound", + spec=dict(kv_lens=(65536,), qlen=1, max_seq_len=65536, tight_bound=True), + ), + Cell( + why="one request verifying 1 + 3 rows over a 3000-token context: an upper bound above the length", + spec=dict(kv_lens=(3000,), qlen=4, max_seq_len=SHORT_BOUND, tight_bound=True), + ), + Cell( + why="two requests with the model maximum as the upper bound", + spec=dict(kv_lens=(20000, 5000), qlen=4, max_seq_len=MODEL_MAX, tight_bound=False), + ), + ) + + note = """ + Outputs at two valid sequence-length bounds agree to a few BF16 ulps and both match the fp32 + reference within the FP8-KV tolerance of the in-tree dense decode test. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + minimax_m3_trtllm_gen_dense_decode( + q, + kv_cache_manager, + layer_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + max_seq_len=max_seq_len, + max_num_requests=max_num_requests, + ) + return output + + def reference( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + """fp32 attention over the pool's K / V (q rounded to the pool's dtype first, as the kernel runs it). + + ``output``, ``max_seq_len`` and ``max_num_requests`` do not change the result.""" + keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) + if keys.dtype == torch.float8_e4m3fn: + q = q.to(torch.float8_e4m3fn) + return dense_attention_reference(q, keys, values, seq_lens, decode_query_len, sm_scale) + + def is_valid( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> None: + # The multi-CTA KV counters are sized for max_num_requests: more requests overrun them. + assert seq_lens.shape[0] <= max_num_requests, "more requests than the counters cover" + assert bool((seq_lens <= max_seq_len).all()), "sequence length exceeds its bound" + # Each verify row attends a causal prefix of at least one key. + assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" + + +m3_dense_decode = M3DenseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_draft_attention.py new file mode 100644 index 000000000000..6347348d1681 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_draft_attention.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's attention on the TRTLLM backend: RoPE, KV-cache append and causal attention.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + PredefinedAttentionMask, +) +from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType + + +class M3DraftAttention(OpWrapper): + """The drafter layer's attention over its fused QKV rows ``qkv`` [T, (q + 2 kv) heads x 128] (bf16): the + backend applies the drafter's RoPE (fused), appends K / V to the drafter layer's paged cache (FP8, the target's + NVFP4 cache keeps its draft layer in FP8) and runs causal attention over the step's live rows (``metadata`` + ``num_tokens``); returns [num_tokens, q heads x 128] bf16. + + ``backend`` (``TrtllmAttention``) and ``metadata`` are the runtime's objects for the drafter layer and step. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="first draft pass of a 4-request verify (4 x 4 rows)", spec=dict(reqs=4, qlen=4)), + Cell(why="a later draft pass, one row per request", spec=dict(reqs=8, qlen=1)), + Cell(why="a context chunk's draft rows", spec=dict(reqs=1, context=512)), + ) + + note = """ + `reference` is exact (float64) causal softmax attention: Q rotated by the backend's NeoX RoPE at each token's + position, over the drafter layer's K / V rows read back from its FP8 cache (dequantized by the backend's + ``kv_scale_quant_orig``). What remains between the two is the FMHA's own arithmetic (Q's rounding to the + kernel's input precision, its P rounding and accumulation order), which the test bounds. + """ + + def __call__(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + num_tokens = metadata.num_tokens + out = backend.forward( + qkv[:num_tokens, :], + None, + None, + metadata, + forward_args=AttentionForwardArgs(attention_mask=PredefinedAttentionMask.CAUSAL), + ) + return out[0] if isinstance(out, tuple) else out + + def reference(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + """The call's [num_tokens, q heads x 128] output, in ``qkv``'s dtype. + + Query token ``j`` of request ``r`` sits at position ``cached_r + j`` (the metadata's cached lengths) and + attends positions up to it. Its Q is rotated here; every K / V row, the step's own included, is read from + the drafter layer's paged cache at the slots the metadata's KV cache manager holds, so this is evaluated + after the call (whose roped, quantized append it then reads back; the test checks that append itself). + Mirrors one KV head per rank and un-scaled NeoX RoPE, the drafter's configuration. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + assert backend.num_kv_heads == 1, "mirrors the one-KV-head-per-rank layout only" + rope = backend.rope_params + assert backend.position_embedding_type == int(PositionEmbeddingType.rope_gpt_neox) + assert rope.scale_type == RotaryScalingType.none and rope.alpha == 1.0 + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + kv_scale = float(backend.kv_scale_quant_orig.reshape(-1)[0]) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = int(metadata.num_tokens) + assert sum(qlens) == total + device = qkv.device + + rot = int(rope.dim) + inv_freq = float(rope.theta) ** ( + -torch.arange(0, rot, 2, dtype=torch.float64, device=device) / rot + ) + pool = manager.get_buffers(layer, "HND") + q = qkv[:total, : heads * dim].to(torch.float64).reshape(total, heads, dim) + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][ + : (length + page - 1) // page + ], + dtype=torch.int64, + device=device, + ) + keys = pool[pages, 0, 0].to(torch.float64).reshape(-1, dim)[:length] * kv_scale + values = pool[pages, 1, 0].to(torch.float64).reshape(-1, dim)[:length] * kv_scale + qpos = cached[row] + torch.arange(n, device=device) + angles = qpos.to(torch.float64)[:, None] * inv_freq[None, :] + cos, sin = angles.cos()[:, None, :], angles.sin()[:, None, :] + rows = q[first : first + n] + x1, x2 = rows[..., : rot // 2], rows[..., rot // 2 : rot] + rows = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, rows[..., rot:]], dim=-1) + allowed = torch.arange(length, device=device)[None, :] <= qpos[:, None] + logits = torch.einsum("thd,kd->thk", rows, keys) * sm_scale + logits = logits.masked_fill(~allowed[:, None, :], float("-inf")) + out[first : first + n] = torch.softmax(logits, dim=-1) @ values + first += n + return out.reshape(total, heads * dim).to(qkv.dtype) + + +m3_draft_attention = M3DraftAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..bcd928fccf5b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The unfused sparse layers' index branch: per-head Gemma RMSNorm + partial NeoX RoPE of [index-Q | index-K], index-Q +out in FP8 E4M3 and index-K inserted into the FP8 index-K cache, in one kernel.""" + +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 +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3Fp8IndexerQkNormRope(OpWrapper): + """``qk`` [T, (heads + 1) * 128] BF16 (the index projection's output): returns index-Q [T, heads, 128] E4M3 and + writes index-K (E4M3) at ``out_cache_loc`` of ``index_k_cache`` [pages, 1, 128, 128]. Gemma RMSNorm of each head, + then RoPE on its leading ``rotary_dim`` channels (NeoX). + + The unfused sparse layers (``fuse_qkv_index_projection=False``) run it on the index branch with an FP8 indexer; + the fused layers' producer writes the index cache itself. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, one index head per rank", spec=dict(m=1, heads=1)), + Cell(why="a 48-row decode verify", spec=dict(m=48, heads=1)), + Cell(why="a prefill chunk", spec=dict(m=2048, heads=1)), + ) + + note = """ + `reference` is the in-place fused norm + RoPE in BF16 followed by the E4M3 cast (the kernel rounds through BF16 + before the cast); its index-K is what the kernel writes to the cache. The two kernels may round a handful of + values differently: the test allows fewer than 0.1% of the bytes to differ, as the in-tree test does. + """ + + def __call__( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_fp8_indexer_qk_norm_rope( + qk, + index_k_cache, + out_cache_loc, + num_heads_q, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + position_ids, + ) + + def reference( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """(index-Q [T, heads, head_dim] E4M3, index-K [T, head_dim] E4M3): the catalog's in-place kernel on a copy, + then the E4M3 cast.""" + out = fused_qk_norm_rope( + qk.clone(), + num_heads_q, + 1, + 0, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + True, + position_ids, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + q, k = out.split([num_heads_q * head_dim, head_dim], dim=-1) + return ( + q.reshape(q.shape[0], num_heads_q, head_dim).to(torch.float8_e4m3fn), + k.to(torch.float8_e4m3fn), + ) + + +m3_fp8_indexer_qk_norm_rope = M3Fp8IndexerQkNormRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_fused_score_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_fused_score_select.py new file mode 100644 index 000000000000..0457f456024e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_fused_score_select.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MSA index scoring and top-16 block selection of a pure decode span of at most 2 requests, in one launch.""" + +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 + +from .m3_index_score import m3_index_score +from .m3_select_long import m3_select_long + +#: Requests per launch the kernel takes. +MAX_REQUESTS = 2 +#: Query rows per request it takes (decode, or a verify of up to 3 drafts). +MAX_DECODE_QUERY_LEN = 4 +#: (KV head x query row) score rows of one request its stage ring holds. +MAX_ROWS = 8 +#: Bytes of the stage ring the select scratch shares. +RING_BYTES = 12 * 16384 + +# M3's checkpoint: no forced initial blocks, the one block holding the query forced in. +_M3 = dict(num_kv_heads=1, init_blocks=0, local_blocks=1) + + +def ring_fits(rows: int, max_tiles: int) -> bool: + """Whether a request's select scratch of ``rows`` score rows over ``max_tiles`` blocks fits the stage ring.""" + return rows * ((max_tiles + 3) // 4 * 4) * 4 + rows * 1024 <= RING_BYTES + + +def scratch_shape(max_tiles: int) -> tuple[int, int, int]: + """The scratch [requests, rows, L] that serves every request count and geometry the op takes.""" + return (MAX_REQUESTS, MAX_ROWS, (max_tiles + 3) // 4 * 4) + + +class M3FusedScoreSelect(OpWrapper): + """Fill ``out`` [total_q, num_kv_heads, 16] with each query row's top-16 KV blocks, straight from the FP8 index Q + and the paged FP8 index-K cache. + + The table is ``m3_select_long`` (ascending ids, -1 tail, ties to the smaller block, the last + ``local_blocks`` valid blocks and the first ``init_blocks`` forced in) applied to the per-KV-head maxima over its + index heads of the ``m3_index_score`` scores, with each row's valid-block count ``n_valid`` clamped to + ``max_tiles``. Grid (splits, requests): each CTA scores a share of its request's blocks through a 12-stage TMA + ring with ``mma.sync`` (the products and the 128-term sums of the CuTe DSL scorer, so the same fp32 scores); the + last CTA of a request (an arrival counter it returns to zero) selects. The request's block-table row and length + are read and its older index-K pages are issued before the grid wait; only the last CTA releases the PDL + dependents. ``scratch`` holds the per-request scores between the two phases and needs no initialization; + ``counters`` must be zero before the first call and stays zero between calls. Returns ``out``. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="Eagle3 verify at one request (the TP8 operating point): 4 rows over 20000 tokens", + spec=dict(batch=1, qlen=4, seq_lens=(20000,), max_tiles=8192, **_M3), + ), + Cell( + why="one request decoding one row over 65536 tokens: 512 blocks across every split", + spec=dict(batch=1, qlen=1, seq_lens=(65536,), max_tiles=8192, **_M3), + ), + Cell( + why="one request of 129 tokens: two blocks, most splits empty", + spec=dict(batch=1, qlen=4, seq_lens=(129,), max_tiles=8192, **_M3), + ), + Cell( + why="one request of 200000 tokens: more blocks than one split holds in its 128 local slots", + spec=dict(batch=1, qlen=4, seq_lens=(200000,), max_tiles=8192, **_M3), + ), + Cell( + why="two requests verifying 4 rows each, of different lengths", + spec=dict(batch=2, qlen=4, seq_lens=(20000, 3000), max_tiles=8192, **_M3), + ), + Cell( + why="two requests, one of them shorter than 16 blocks: its whole row is selected and -1 padded", + spec=dict(batch=2, qlen=4, seq_lens=(1500, 70000), max_tiles=8192, **_M3), + ), + ) + + note = """ + Certified at one index head and one KV head per rank (TP8) with M3's forcing (`init_blocks=0`, + `local_blocks=1`), FP8 index Q and index-K, a 1M-token engine (`max_tiles=8192`). The selection is checked bit + for bit against `reference`; scores are compared through the table only. Rows with more than 1024 blocks tied at + the 16th score are not driven. + """ + + def __call__( + self, + idx_q: torch.Tensor, + idx_k_pages: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + n_valid: torch.Tensor, + out: torch.Tensor, + scratch: torch.Tensor, + counters: torch.Tensor, + decode_query_len: int, + num_kv_heads: int, + init_blocks: int, + local_blocks: int, + max_tiles: int, + use_pdl: bool, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_tp8_fused_score_select( + idx_q, + idx_k_pages, + block_table, + seq_lens, + n_valid, + out, + scratch, + counters, + decode_query_len, + num_kv_heads, + init_blocks, + local_blocks, + max_tiles, + use_pdl, + ) + return out + + def reference( + self, + idx_q: torch.Tensor, + idx_k_pages: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + n_valid: torch.Tensor, + out: torch.Tensor, + scratch: torch.Tensor, + counters: torch.Tensor, + decode_query_len: int, + num_kv_heads: int, + init_blocks: int, + local_blocks: int, + max_tiles: int, + use_pdl: bool, + ) -> torch.Tensor: + """The int32 [total_q, num_kv_heads, 16] table: ``m3_select_long.reference`` over the per-KV-head maxima of + ``m3_index_score.reference``. ``out``, ``scratch``, ``counters`` and ``use_pdl`` only say where and how the op + works.""" + total_q, heads, _ = idx_q.shape + placeholder = torch.empty(heads, total_q, max_tiles, device=idx_q.device) + scores = m3_index_score.reference( + idx_q, idx_k_pages, block_table, seq_lens, placeholder, decode_query_len, 1, False + ) # [heads, total_q, max_tiles] + per_kv = scores.view(num_kv_heads, heads // num_kv_heads, total_q, max_tiles).amax(dim=1) + table = torch.empty(total_q, num_kv_heads, 16, dtype=torch.int32, device=idx_q.device) + return m3_select_long.reference( + per_kv.transpose(1, 2), n_valid, table, init_blocks, local_blocks + ) + + def is_valid( + self, + idx_q: torch.Tensor, + idx_k_pages: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + n_valid: torch.Tensor, + out: torch.Tensor, + scratch: torch.Tensor, + counters: torch.Tensor, + decode_query_len: int, + num_kv_heads: int, + init_blocks: int, + local_blocks: int, + max_tiles: int, + use_pdl: bool, + ) -> None: + """The inputs this op takes and answers wrongly. + + A non-zero counter makes a request's selection run early, on partial scores. The block-table row is read up + to ``ceil(seq_len / 128)`` entries unchecked. Reads ``counters`` and ``seq_lens`` back to the host: a test-only + cost. + """ + assert bool((counters[: seq_lens.shape[0]] == 0).all()), ( + "counters must be zero on entry: the kernel selects once a request's arrivals reach its split count" + ) + nblocks = (seq_lens + 127) // 128 + assert bool((nblocks <= block_table.shape[1]).all()), ( + "a request runs past its block-table row" + ) + assert bool((nblocks <= max_tiles).all()), "a request has more blocks than max_tiles" + assert init_blocks >= 0 and local_blocks >= 0, "forcing counts must be non-negative" + + +m3_fused_score_select = M3FusedScoreSelect() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_index_score.py new file mode 100644 index 000000000000..f5c385977d4c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_index_score.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# PyTorch oracle ported from vLLM (Apache-2.0), _reference_decode_index_score: +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/tests/kernels/attention/test_minimax_m3.py#L188 +"""The MSA index scorer of a sparse decode step: per (index head, query token, KV block) the causal max of Q . K over +the block's 128 index-K positions, with the pre-wait schedule of ``ops/m3_index_score``.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +__all__ = ["m3_index_score"] + +PAGE = 128 +HEAD_DIM = 128 +#: CTAs per request (each scores blocks split, split + SPLIT_K, ...; the scores do not depend on it). +SPLIT_K = 512 + + +def compiled_scorer(key: tuple): + """The compiled kernel of ``key`` (dtype, heads, query length, head dim, split_k, early), compiling it on a miss + (the target compiles every key it can take before warmup: ``layers.prebuild_index_scorer``).""" + from ..m3_index_score import IndexScoreRunner + + if key not in IndexScoreRunner.kernel_cache: + IndexScoreRunner._compile(*key) + return IndexScoreRunner.kernel_cache[key] + + +class M3IndexScore(OpWrapper): + """``score[h, t, j]`` (fp32, a transposed view of the [heads, blocks, tokens] selector buffer) = the max over the + positions ``p`` of block ``j`` of request ``b = t // qlen`` with ``p <= seq_lens[b] - qlen + t % qlen`` of + ``idx_q[t, h] . index_k[block_table[b, j], p % 128]``, for ``j < ceil(seq_lens[b] / 128)`` (no other entry is + written). FP8 or BF16 q and index-K, fp32 products and sums. + + ``early`` (the target's one-request kernel) issues the Q and first K tiles right after the grid wait from a + pre-wait length hint and prefetches one older block per CTA into L2 before the wait; otherwise every older block of + the CTA is prefetched. Neither changes a score. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="three requests verifying 1 + 3 rows (the narrowest span the fused score + select leaves to it)", + spec=dict(seq_lens=(20000, 3000, 129), qlen=4, early=False), + ), + Cell( + why="three requests, one longer than 512 blocks (the CTAs loop)", + spec=dict(seq_lens=(70000, 2049, 4096), qlen=4, early=False), + ), + Cell( + why="eight requests of mixed lengths: the prefetching kernel", + spec=dict( + seq_lens=(300, 1500, 2049, 4096, 5000, 7777, 8192, 33001), qlen=4, early=False + ), + ), + Cell( + why="eight requests decoding one row each", + spec=dict( + seq_lens=(300, 1500, 2049, 4096, 5000, 7777, 8192, 33001), qlen=1, early=False + ), + ), + ) + + note = """ + The products of FP8 (or BF16) values are exact in fp32 and the kernel sums 128 of them in its k-split order, so it + differs from the reference matmul by a few fp32 ulps; it is held to the in-tree scorer bit for bit. + """ + + def __call__( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + _, num_heads, head_dim = idx_q.shape + key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + compiled_scorer(key)(idx_q, index_k, block_table, score, seq_lens) + return score + + def reference( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + """The fp32 [heads, total_q, blocks] scores (``-inf`` where nothing is written). ``score``, ``split_k`` and + ``early`` do not change them.""" + total_q, heads, _ = idx_q.shape + qlen = max_decode_query_len + out = torch.full( + (heads, total_q, score.shape[2]), + -float("inf"), + dtype=torch.float32, + device=idx_q.device, + ) + with true_fp32_matmul(): + for b, seq in enumerate(seq_lens.tolist()): + nblocks = (seq + PAGE - 1) // PAGE + if nblocks == 0: + continue + t0 = b * qlen + q = idx_q[t0 : t0 + qlen].float() + k = index_k[block_table[b, :nblocks].long()].reshape(nblocks * PAGE, -1).float() + s = torch.einsum("qhd,kd->hqk", q, k) + q_pos = seq - qlen + torch.arange(qlen, device=idx_q.device) + k_pos = torch.arange(k.shape[0], device=idx_q.device) + s.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) + out[:, t0 : t0 + qlen, :nblocks] = s.reshape(heads, qlen, nblocks, PAGE).amax(dim=3) + return out + + def is_valid( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> None: + # The block-table row and the score row are read and written up to ceil(seq_len / 128) with no bound. + nblocks = (seq_lens + PAGE - 1) // PAGE + assert bool((nblocks <= block_table.shape[1]).all()), ( + "a request runs past its block-table row" + ) + assert bool((nblocks <= score.shape[2]).all()), "a request runs past the score buffer" + assert idx_q.shape[0] == seq_lens.shape[0] * max_decode_query_len, ( + "idx_q must hold max_decode_query_len rows per request" + ) + + +m3_index_score = M3IndexScore() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_msa_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_msa_attention.py new file mode 100644 index 000000000000..48b0ef2c1c35 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_msa_attention.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MiniMax-M3 MSA attention backend's step: KV write, sparse top-k block selection, FMHA.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + SparseBackendForwardArgs, +) + +from .m3_sparse_decode_lc8 import dequant_nvfp4_pages + + +class M3MsaAttention(OpWrapper): + """One layer's attention core on the MSA backend (``MiniMaxM3MsaSparseAttention``), written into ``output`` + [T, heads * 128] (bf16, the step's live rows). + + * ``k`` / ``v`` given (the dense layers): written into the layer's paged main cache first (NVFP4 with + ``kv_scale_orig_quant``, else FP8 / BF16 as the pool is). None: a producer already wrote the caches. + * ``sparse``: the indexer scores the layer's index-K cache (already written) with ``idx_q`` and publishes the + top-k blocks; the FMHA attends those blocks. Otherwise the FMHA attends the whole page table. + + The FMHA dispatches on the step's metadata (context rows: the MSA prefill FMHA; generation rows: the MSA decode + FMHA, the dense layers' trtllm-gen decode over their full context). ``backend`` and ``metadata`` are the runtime's + objects for this layer and step. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="dense layer, a prefill chunk, FP8 pool", + spec=dict(layer="dense", phase="context", tokens=1024), + ), + Cell( + why="dense layer, a decode step of 4 requests", + spec=dict(layer="dense", phase="generation", reqs=4), + ), + Cell( + why="sparse layer, a prefill chunk beyond 16 blocks, NVFP4 pool", + spec=dict(layer="sparse", phase="context", tokens=4096), + ), + Cell( + why="sparse layer, a mixed step (context + generation rows)", + spec=dict(layer="sparse", phase="mixed", tokens=512, reqs=2), + ), + ) + + note = """ + `reference` is exact (float64) softmax attention over the cache rows the step attends, read back from the paged + cache the call wrote (or a producer wrote before it) and dequantized as the FMHA dequantizes them: FP8 rows as + E4M3, NVFP4 rows as E2M1 x E4M3 block scale x the per-tensor ``kv_scale_quant_orig``. A sparse layer's top-k blocks + are re-selected from ``idx_q`` and the index-K cache with the backend's rule (per-block max over the causal + keys, the query's own block forced in, ties to the lower block). On an FP8 pool the FMHA takes Q as a plain E4M3 + cast, which the reference applies too. What remains between the two is the FMHA's own arithmetic (its P + rounding and accumulation order), which the test bounds. + """ + + def __call__( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + if k is not None: + backend.write_layer_caches( + k, v, None, metadata, kv_scale_orig_quant=kv_scale_orig_quant + ) + if sparse: + blocks = backend.run_indexer(idx_q, None, metadata, idx_k_prewritten=True) + args = AttentionForwardArgs( + output=output, sparse_backend_args=SparseBackendForwardArgs(topk_indices=blocks) + ) + else: + args = AttentionForwardArgs(output=output) + args.kv_scale_orig_quant = kv_scale_orig_quant + args.kv_scale_quant_orig = kv_scale_quant_orig + backend.forward(q, None, None, metadata, forward_args=args) + return output + + def reference( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + """``output`` as the call is supposed to write it, [T, heads * 128] in ``output``'s dtype. + + Every attended row -- the step's own included -- is read from the layer's paged cache, at the slots the + metadata's KV cache manager holds for each request, so this is evaluated after the call (whose write of + ``k`` / ``v`` it then reads back). ``k``, ``v`` and ``kv_scale_orig_quant`` are therefore not read; the + test checks the write itself. Query token ``j`` of request ``r`` sits at position ``cached_r + j`` and + attends positions up to it; on a sparse layer only those in its top-k blocks. Mirrors one KV head per rank, + the target's layout. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + assert backend.num_kv_heads == 1, "mirrors the one-KV-head-per-rank layout only" + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = sum(qlens) + assert q.shape[0] == total and output.shape[0] == total + + pool = manager.get_buffers(layer, "HND") + nvfp4 = manager.is_nvfp4_layer(layer) + if nvfp4: + scale_pool = manager.get_block_scale_buffers(layer, "HND") + k_global = float(kv_scale_quant_orig.reshape(-1)[1]) + v_global = float(kv_scale_quant_orig.reshape(-1)[2]) + q_rows = q.reshape(total, heads, dim) + if q_rows.dtype != torch.float8_e4m3fn and pool.dtype == torch.float8_e4m3fn: + # The FP8 FMHA shares one dtype across Q / K / V: Q is cast to E4M3, unscaled. + q_rows = q_rows.to(torch.float8_e4m3fn) + q64 = q_rows.to(torch.float64) + if sparse: + config = backend.m3_config + idx_cache = metadata.msa_idx_k_cache(layer) + idx64 = idx_q.reshape(total, config.num_index_heads, config.sparse_index_dim).to( + torch.float64 + ) + + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=q.device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + blocks = (length + page - 1) // page + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][:blocks], + dtype=torch.int64, + device=q.device, + ) + if nvfp4: + keys = dequant_nvfp4_pages( + pool[:, 0].view(torch.uint8), scale_pool[:, 0], pages, quad_scales=False + ) + values = dequant_nvfp4_pages( + pool[:, 1].view(torch.uint8), scale_pool[:, 1], pages, quad_scales=True + ) + keys, values = keys * k_global, values * v_global + else: + keys = pool[pages, 0, 0].to(torch.float64) + values = pool[pages, 1, 0].to(torch.float64) + keys = keys.reshape(-1, dim)[:length] + values = values.reshape(-1, dim)[:length] + qpos = cached[row] + torch.arange(n, device=q.device) + allowed = torch.arange(length, device=q.device)[None, :] <= qpos[:, None] + if sparse: + idx_keys = idx_cache[pages, 0].to(torch.float64).reshape(-1, dim)[:length] + chosen = _selected_blocks( + idx64[first : first + n], idx_keys, allowed, qpos, page, config + ) + allowed &= chosen.repeat_interleave(page, dim=1)[:, :length] + for lo in range(0, n, 256): + hi = min(n, lo + 256) + logits = torch.einsum("thd,kd->thk", q64[first + lo : first + hi], keys) * sm_scale + logits = logits.masked_fill(~allowed[lo:hi, None, :], float("-inf")) + out[first + lo : first + hi] = torch.softmax(logits, dim=-1) @ values + first += n + return out.reshape(total, heads * dim).to(output.dtype) + + +def _selected_blocks( + idx_q: torch.Tensor, + idx_keys: torch.Tensor, + allowed: torch.Tensor, + qpos: torch.Tensor, + page: int, + config, +) -> torch.Tensor: + """The blocks each query token of one request attends, [n, blocks] bool. + + A block's score is the maximum, over the KV head's index heads and the block's causally visible keys, of + ``idx_q . idx_k``. Of the token's ``ceil((pos + 1) / page)`` valid blocks the first ``init_blocks`` and the last + ``local_blocks`` are forced in, and the ``topk`` highest-scoring are kept, ties to the lower block (the + backend's ``select_blocks_from_maxscore`` rule). The positive score scale the scorers apply does not change the + ranking, so none is applied. + """ + n, length = allowed.shape + blocks = (length + page - 1) // page + scores = torch.einsum("thd,kd->thk", idx_q, idx_keys).amax(dim=1) + scores = scores.masked_fill(~allowed, float("-inf")) + padded = torch.full((n, blocks * page), float("-inf"), dtype=scores.dtype, device=scores.device) + padded[:, :length] = scores + block_scores = padded.view(n, blocks, page).amax(dim=-1) + ids = torch.arange(blocks, device=scores.device)[None, :] + n_valid = torch.div(qpos + page, page, rounding_mode="floor")[:, None] + block_scores = block_scores.masked_fill(ids >= n_valid, float("-inf")) + forced = (ids < torch.clamp(n_valid, max=config.init_blocks)) | ( + (ids >= n_valid - config.local_blocks) & (ids < n_valid) + ) + block_scores = block_scores.masked_fill(forced, float("inf")) + order = torch.sort(block_scores, dim=1, descending=True, stable=True).indices[:, : config.topk] + keep = torch.arange(order.shape[1], device=scores.device)[None, :] < n_valid + chosen = torch.zeros(n, blocks, dtype=torch.bool, device=scores.device) + rows = torch.arange(n, device=scores.device)[:, None].expand_as(order) + chosen[rows[keep], order[keep]] = True + return chosen + + +m3_msa_attention = M3MsaAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_producer.py new file mode 100644 index 000000000000..7dcfbd4624b9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_producer.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer, with its constant loads staged +before the PDL grid-dependency wait and an early dependents trigger (``kernels/minimaxM3/producerPw.cu``).""" + +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 + +__all__ = ["m3_producer"] + +HEAD_DIM = 128 +ROTARY_DIM = 64 +PAGE = 128 +#: Query heads per rank at TP8 (MiniMax-M3: 64 query heads; 1 KV head and 1 index head per rank). +HEADS_Q = 8 +#: NVFP4 block: 16 values share one E4M3 scale. +FP4_BLOCK = 16 +E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def _norm_rope( + x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, eps: float +): + """Gemma RMSNorm (``x * rsqrt(mean(x^2) + eps) * (1 + w)``) then NeoX RoPE on the first 64 dims, in fp32. + x [T, 128] fp32; cos / sin [T, 32].""" + y = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) * (1.0 + weight.float()) + lo, hi, rest = y[:, :32], y[:, 32:64], y[:, 64:] + return torch.cat([lo * cos - hi * sin, hi * cos + lo * sin, rest], dim=-1) + + +def decode_nvfp4_token( + data: torch.Tensor, scales: torch.Tensor, quant_scale: float +) -> torch.Tensor: + """fp32 [128] of one token's packed E2M1 row (64 bytes, low nibble first) and its 8 E4M3 block scales.""" + table = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=data.device) + nib = torch.stack([data & 0xF, data >> 4], dim=-1).reshape(-1).long() + vals = table[nib & 7] * torch.where(nib >= 8, -1.0, 1.0) + s = scales.view(torch.float8_e4m3fn).float() / quant_scale + return vals * s.repeat_interleave(FP4_BLOCK) + + +def scale_index(within_page: int, plane: int) -> torch.Tensor: + """Positions of a token's 8 block scales in its page's [128 * 8] scale plane: row-major for K, the 4-token + swizzle for V.""" + cols = torch.arange(8) + if plane == 0: + return within_page * 8 + cols + return (within_page // 4) * 32 + cols * 4 + within_page % 4 + + +class M3Producer(OpWrapper): + """The NVFP4 QKV + indexer producer of a sparse layer: ``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert`` for one + KV head, one index head, head dim 128 and rotary dim 64, with the same arguments and outputs bit for bit. + + ``packed`` [T, (heads_q + 4) * 128] BF16 holds each token's q heads, k, v, index-q and index-k. Every head but v + gets the Gemma RMSNorm with its weight and NeoX RoPE on its first 64 dims (``rotary_cos_sin[position]``). Returns + (q FP8 [T, heads_q, 128], index-q FP8 [T, 1, 128]); writes the token's slot ``out_cache_loc[t]`` (skipped when + negative or past the pool): k and v as NVFP4 (16-value blocks, E4M3 scales against ``kv_quant_scale[1]`` / + ``[2]``, the V scales in the 4-token swizzle) and index-k as FP8. ``early`` triggers the dependent kernel on + entry (its outputs are read only after the dependent's own grid wait). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token (c1 without speculation)", spec=dict(tokens=1, scale=1.0)), + Cell(why="a c1 verify: 4 tokens", spec=dict(tokens=4, scale=3.0)), + Cell(why="a c4 verify: 16 tokens, small activations", spec=dict(tokens=16, scale=1e-3)), + Cell( + why="a c8 verify: 32 tokens, large activations, a skipped and an out-of-pool slot", + spec=dict(tokens=32, scale=40.0, bad_slots=True), + ), + Cell(why="the widest decode step: 48 tokens", spec=dict(tokens=48, scale=1.0)), + Cell( + why="a prefill chunk of 1000 tokens, early trigger off", + spec=dict(tokens=1000, scale=1.0, early=False), + ), + ) + + note = """ + The kernel is the native producer's arithmetic, compiled with the library's --use_fast_math, so it is held to that + kernel bit for bit on every output and cache byte. The torch reference states the math in fp32: the FP8 outputs are + held to it within one E4M3 step, and the NVFP4 K/V, decoded, within one E2M1 step of their block scale. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + q, index_q = torch.ops.trtllm.minimax_m3_producer_pw( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + heads_q, + pdl, + early, + ) + return q, index_q + + def reference( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + ) -> dict[str, torch.Tensor]: + """fp32 values before their quantization: "q" [T, heads_q, 128], "index_q" [T, 128], and per token "k", "v", + "index_k" [T, 128] (the values written to the token's slot). The caches, ``pdl`` and ``early`` do not change + them.""" + t = packed.shape[0] + x = packed.float().view(t, heads_q + 4, HEAD_DIM) + rope = rotary_cos_sin[position_ids[:t].long()] + cos, sin = rope[:, 0], rope[:, 1] + q = torch.stack( + [_norm_rope(x[:, h], q_weight, cos, sin, eps) for h in range(heads_q)], dim=1 + ) + return { + "q": q, + "k": _norm_rope(x[:, heads_q], k_weight, cos, sin, eps), + "v": x[:, heads_q + 1], + "index_q": _norm_rope(x[:, heads_q + 2], index_q_weight, cos, sin, eps), + "index_k": _norm_rope(x[:, heads_q + 3], index_k_weight, cos, sin, eps), + } + + def is_valid( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + ) -> None: + t = packed.shape[0] + # The RoPE row is read at the position with no bound. + pos = position_ids[:t] + assert bool(((pos >= 0) & (pos < rotary_cos_sin.shape[0])).all()), ( + "position_ids must index rotary_cos_sin" + ) + # Two tokens writing one slot race. + slots = out_cache_loc[:t] + live = slots[(slots >= 0) & (slots < kv_data.shape[0] * PAGE)] + assert live.unique().numel() == live.numel(), "out_cache_loc must not repeat a slot" + + +m3_producer = M3Producer() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..681aec29472e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_qk_norm_rope_fp8.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' per-head Gemma RMSNorm + partial NeoX RoPE of the packed [Q | K | V], into a fresh FP8 copy.""" + +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 +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3QkNormRopeFp8(OpWrapper): + """``fused_qk_norm_rope`` written out of place in FP8 E4M3: Q and K heads normed and roped in fp32 and rounded + once to E4M3 (saturating at +/-448), V cast, ``qkv`` left untouched. Same arguments as the catalog's in-place + entry. + + The dense layers (0-2) feed the FP8 main KV cache with it: on every step of an FP8 cache, and on the decode steps + of a hybrid NVFP4 cache (whose dense layers sit in its FP8 pool). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, tp8 heads (8 q, 1 kv)", spec=dict(m=1, q=8, kv=1)), + Cell(why="a 48-row decode verify", spec=dict(m=48, q=8, kv=1)), + Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=8, kv=1)), + ) + + note = """ + `reference` is the in-place kernel followed by a cast to E4M3. V matches it bit for bit; Q and K round once from + fp32 where the composition rounds through bf16 first, so they are within one E4M3 step of it. + """ + + def __call__( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + return torch.ops.trtllm.fused_qk_norm_rope_to_fp8( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + + def reference( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + """The catalog's in-place kernel on a copy, then the E4M3 cast (V exact; Q and K within one E4M3 step).""" + out = fused_qk_norm_rope( + qkv.clone(), + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + return out.to(torch.float8_e4m3fn) + + +m3_qk_norm_rope_fp8 = M3QkNormRopeFp8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_qkv_insert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_qkv_insert.py new file mode 100644 index 000000000000..fc46936cc839 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_qkv_insert.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' generic horizontal producer: packed QKV/index rows -> FP8 Q / index-Q, K / V / index-K into +the paged caches.""" + +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 M3QkvInsert(OpWrapper): + """The packed ``[Q | K | V | index-Q | index-K]`` bf16 rows of a sparse layer: per-head Gemma RMSNorm and the + partial NeoX RoPE of Q, K, index-Q and index-K, K / V written to the main paged cache at ``out_cache_loc`` and + index-K to the FP8 index cache; returns (Q, index-Q) in FP8 E4M3 [M, heads, 128]. + + ``kv_scale`` given: the main cache is NVFP4 (packed E2M1 data [pages, 2, kv, 128, 64] uint8 + E4M3 scale pool + [pages, 2, kv, 128, 8], ``kv_quant_scale`` the fp32 (Q, K, V) quantization scales); None: FP8 E4M3 + [pages, 2, kv, 128, 128]. Token-major and step-type agnostic. ``m3_producer`` takes the NVFP4 cache at the tp8 + heads; this entry takes the FP8 cache and NVFP4 pools outside that layout. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="FP8 main cache, a prefill chunk", spec=dict(m=1024, cache="fp8")), + Cell(why="FP8 main cache, one decode token", spec=dict(m=1, cache="fp8")), + ) + + note = """ + The arithmetic is the CUDA producer's (fp32 norm and RoPE with bf16 roundings, then the E4M3 casts); `reference` + returns the fp32 values, and the test holds the FP8 outputs and cache rows within one E4M3 step of them. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + common = ( + num_heads, + num_kv_heads, + num_index_heads, + head_dim, + rotary_dim, + eps, + q_norm, + k_norm, + index_q_norm, + index_k_norm, + rotary_cos_sin, + position_ids, + ) + if kv_scale is not None: + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, kv_scale, index_k_cache, out_cache_loc, kv_quant_scale, *common + ) + return torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, index_k_cache, out_cache_loc, *common + ) + + def reference( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + """The fp32 values the producer rounds (a new dict, nothing written): Gemma RMSNorm + NeoX partial RoPE of + Q [M, heads, D], index-Q [M, index heads, D], K [M, kv heads, D] and index-K [M, D]; V [M, kv heads, D] as + packed. The cache rows at ``out_cache_loc`` and the returned FP8 Q / index-Q are these values cast.""" + m = packed.shape[0] + sizes = [ + num_heads * head_dim, + num_kv_heads * head_dim, + num_kv_heads * head_dim, + num_index_heads * head_dim, + head_dim, + ] + q, k, v, iq, ik = packed.float().split(sizes, dim=-1) + cos, sin = rotary_cos_sin[position_ids.long()].unbind(1) + half = rotary_dim // 2 + + def norm_rope(x, w, heads): + x = x.reshape(m, heads, head_dim) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w.float()) + x1, x2, rest = x[..., :half], x[..., half:rotary_dim], x[..., rotary_dim:] + c, s = cos.unsqueeze(1), sin.unsqueeze(1) + return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s, rest], dim=-1) + + return { + "q": norm_rope(q, q_norm, num_heads), + "index_q": norm_rope(iq, index_q_norm, num_index_heads), + "k": norm_rope(k, k_norm, num_kv_heads), + "v": v.reshape(m, num_kv_heads, head_dim), + "index_k": norm_rope(ik, index_k_norm, 1)[:, 0], + } + + +m3_qkv_insert = M3QkvInsert() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_select_blocks.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_select_blocks.py new file mode 100644 index 000000000000..6bf1764bfc0d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_select_blocks.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MSA top-k block selection from per-KV-head block scores (the built-in selector).""" + +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 + + +class M3SelectBlocks(OpWrapper): + """Per query row the ``topk`` highest-scoring KV blocks of ``max_score_kv`` [kv_heads, n_blocks, total_q] (fp32), + with the first ``init_blocks`` and the last ``local_blocks`` valid blocks forced in and blocks at or beyond the + row's ``n_valid_blocks`` excluded; [total_q, kv_heads, topk] int32 ascending ids, -1 padded. ``head_major_output`` + backs the result head-major so ``permute(1, 0, 2)`` is contiguous. + + The decode dispatch runs it where the long-row selector does not apply: requests of 17..128 KV blocks. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="eight decode rows of 17..128 blocks", + spec=dict(total_q=8, n_blocks=128, mixed=True), + ), + Cell( + why="one decode row at the 128-block edge", + spec=dict(total_q=1, n_blocks=128, mixed=False), + ), + Cell(why="a 4-row verify at 40 blocks", spec=dict(total_q=4, n_blocks=40, mixed=False)), + ) + + note = """ + Exact: ties between equal scores break toward the lower block id; `reference` sorts with the same rule. + """ + + def __call__( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous() + return torch.ops.trtllm.minimax_m3_select_blocks( + max_score_kv, nvb, topk, init_blocks, local_blocks, head_major_output + ) + + def reference( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + """Plain torch selection; ``head_major_output`` changes the backing only.""" + heads, n_blocks, total_q = max_score_kv.shape + out = torch.full((total_q, heads, topk), -1, dtype=torch.int32, device=max_score_kv.device) + blocks = torch.arange(n_blocks, device=max_score_kv.device) + for t in range(total_q): + nv = int(n_valid_blocks[t]) + for h in range(heads): + s = max_score_kv[h, :, t].clone().double() + s[blocks >= nv] = float("-inf") + forced = (blocks < min(init_blocks, nv)) | ( + (blocks >= nv - local_blocks) & (blocks < nv) + ) + s[forced] = float("inf") + order = sorted(range(nv), key=lambda b: (-float(s[b]), b))[: min(topk, nv)] + ids = torch.tensor(sorted(order), dtype=torch.int32, device=out.device) + out[t, h, : ids.numel()] = ids + return out + + +m3_select_blocks = M3SelectBlocks() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_select_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_select_long.py new file mode 100644 index 000000000000..867822e7fe41 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_select_long.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 MSA block selection for decode rows of up to 9216 blocks.""" + +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 + +#: Blocks per decode row the op covers (9 keys per thread of one 1024-thread CTA). +MAX_BLOCKS = 9216 +#: Blocks the selection keeps per (query, KV head). +TOPK = 16 +#: Scores the forcing substitutes: initial blocks rank above local ones, and both above any finite score. +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 + +# M3's checkpoint: no forced initial blocks, the one block holding the query forced in. +_M3 = dict(num_kv_heads=1, init_blocks=0, local_blocks=1) + + +class M3SelectLong(OpWrapper): + """Fill ``out`` [total_q, num_kv_heads, 16] with each query's top-16 blocks. + + ``scores`` is the per-KV-head block maxima [num_kv_heads, num_blocks, + total_q] the CuTe DSL index scorer wrote (amax over the KV head's index + heads). Per (query, KV head) row the op keeps the ``n_valid_blocks[query]`` + leading blocks (clamped to ``[0, num_blocks]``), forces the last + ``local_blocks`` of them (counted from the *unclamped* valid count) to + ``LOCAL_SCORE`` and the first ``init_blocks`` to ``INIT_SCORE`` -- local wins + where both apply -- and selects the 16 largest, ties to the smaller block. + A pick whose forced score is -inf is emitted as -1. The ids are written + ascending with a -1 tail. + + Same contract as ``torch.ops.trtllm.minimax_m3_select_blocks`` (the in-tree + ``select_blocks_from_maxscore``) and bit-identical to it, except on a row + where more than 1024 valid blocks tie at the 16th-largest score: this op + breaks that tie to the smallest blocks as the contract says, the in-tree op + does not. The test checks both against ``reference``. Returns ``out``. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="32K-context engine (257 worst-case k-tiles), Eagle3 verify at one request: 4 query rows", + spec=dict(num_blocks=257, total_q=4, valid="near_full", scores="random", **_M3), + ), + Cell( + why="decode without speculation at batch 8 on a 2048-block engine: the 2-key-per-thread path at its edge", + spec=dict(num_blocks=2048, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="1M-context engine at batch 8 x qlen 4: rows on all three key widths, short rows and padding rows", + spec=dict(num_blocks=8193, total_q=32, valid="mixed", scores="random", **_M3), + ), + Cell( + why="duplicated scores: ties broken to the smaller block, and > 1024 keys at the bound (re-read path)", + spec=dict(num_blocks=4500, total_q=4, valid="near_full", scores="ties", **_M3), + ), + Cell( + why="non-finite scores: +inf, NaN ranked above +inf, and rows whose picks run into -inf (emitted -1)", + spec=dict(num_blocks=600, total_q=8, valid="mixed", scores="nonfinite", **_M3), + ), + ) + + note = """ + Certified at one KV head per rank (TP8) with M3's forcing (`init_blocks=0`, + `local_blocks=1`). The kernel takes any head count and forcing; neither is + driven here. + + `scores` may be strided (the op reads it by stride); the target hands over + the scorer's contiguous view, which is what the cells drive. Rows of at + most 128 blocks are within the kernel's range, but the target keeps those on + the in-tree op, so they are not certified here. + + Ordering is by the descending-score key `orderedKey`: a positive NaN ranks + above +inf, +0 above -0, and a negative NaN below -inf -- which, unlike + -inf, is emitted as a block id rather than -1. The cells drive positive + NaN only. + """ + + def __call__( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_select_long( + scores, n_valid_blocks, out, init_blocks, local_blocks, use_pdl, early_trigger + ) + return out + + def reference( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + ) -> torch.Tensor: + """The [total_q, num_kv_heads, 16] int32 table the op is supposed to write. + + Exact integer ranking, no floating point beyond the forcing: each + score's bit pattern maps to an unsigned key that is smaller for a larger + score (NaN above +inf), packed with the block id so one sort ranks by + (score descending, block ascending). ``out``, ``use_pdl`` and + ``early_trigger`` only say where and how the op writes; they do not enter the result. + """ + num_kv_heads, num_blocks, total_q = scores.shape + device = scores.device + blocks = torch.arange(num_blocks, device=device, dtype=torch.int64) + raw = n_valid_blocks[:total_q].to(torch.int64) + valid = raw.clamp(0, num_blocks) + local_start = (raw - local_blocks).clamp_min(0) + result = torch.full((total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device=device) + for head in range(num_kv_heads): + forced = scores[head].float().t().clone() # [total_q, num_blocks] + forced = torch.where(blocks[None, :] < init_blocks, INIT_SCORE, forced) + forced = torch.where(blocks[None, :] >= local_start[:, None], LOCAL_SCORE, forced) + bits = forced.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + key = torch.where(bits >= 0x80000000, bits, (~bits) & 0x7FFFFFFF) + in_row = blocks[None, :] < valid[:, None] + packed = torch.where( + in_row, (key << 16) | blocks[None, :], torch.iinfo(torch.int64).max + ) + k = min(TOPK, num_blocks) + picked = torch.topk(packed, k, dim=1, largest=False).values + picked_key = picked >> 16 + keep = (picked != torch.iinfo(torch.int64).max) & (picked_key != 0xFF800000) + ids = torch.where(keep, picked & 0xFFFF, torch.iinfo(torch.int64).max) + ids = ids.sort(dim=1).values + ids = torch.where(ids == torch.iinfo(torch.int64).max, -1, ids) + result[:, head, :k] = ids.to(torch.int32) + return result + + def is_valid( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + ) -> None: + """The inputs this op takes and answers wrongly. + + Negative forcing counts are accepted and push the forced ranges past the + row instead of failing; the in-tree op this one stands in for refuses + them, so a caller relying on this one would silently diverge from it. + An ``out`` overlapping ``scores`` is overwritten while it is still read. + """ + assert init_blocks >= 0 and local_blocks >= 0, ( + f"init_blocks ({init_blocks}) and local_blocks ({local_blocks}) must be non-negative; the " + "op accepts negative counts and selects as if the forced ranges were elsewhere" + ) + storage = scores.untyped_storage() + assert not (storage.data_ptr() <= out.data_ptr() < storage.data_ptr() + storage.nbytes()), ( + "out must not alias scores: the op writes it while scores are still read" + ) + + +m3_select_long = M3SelectLong() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode.py new file mode 100644 index 000000000000..0438d20109b5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This target's block-sparse decode attention (Triton, with the one-launch long-context route).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..triton_sparse_decode import minimax_m3_sparse_attn_decode +from .m3_sparse_decode_lc8 import PAGE, sparse_attention_reference + + +class M3SparseDecode(OpWrapper): + """``minimax_m3_sparse_attn_decode`` as ``M3Attention._decode_sparse_core`` calls it. + + Attention of each query token over the blocks ``topk_idx`` selected for it + (``[num_kv_heads, T, topk]``, ascending ids with a -1 tail), each block + truncated at the token's causal extent, written into ``output`` in place. + The target passes FP8 q, the packed NVFP4 sparse-layer cache with its block + and per-tensor scales, and BF16 ``output`` [T, 8, 128]. Returns ``output``. + + Which kernel runs is decided inside, per step: a selected (non-identity) + table on the TP8 geometry goes to the one-launch ``m3_sparse_decode_lc8`` + kernel; an identity table (every request within 16 blocks, staged by the + metadata) runs the Triton producer, or from 17 query rows the multi-query + producer, then the LSE merge. Every route also publishes the MXFP8 o_proj + hand-off to ``m3_prequant.SLOT``; this entry certifies the attention output. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="identity selection, decode without speculation at batch 8: the Triton producer + merge", + spec=dict( + seq_lens=(1, 128, 129, 700, 2048, 1500, 64, 0), decode_query_len=1, identity=True + ), + ), + Cell( + why="identity selection, Eagle3 verify at one request (T=4): one block per producer program", + spec=dict(seq_lens=(1900,), decode_query_len=4, identity=True), + ), + Cell( + why="identity selection, Eagle3 verify at batch 4 (T=16): the last row count below the MQ producer", + spec=dict(seq_lens=(300, 2048, 4, 1000), decode_query_len=4, identity=True), + ), + Cell( + why="identity selection, Eagle3 verify at batch 8 (T=32): the multi-query producer, padding rows", + spec=dict( + seq_lens=(2048, 130, 5, 0, 1777, 900, 2000, 0), decode_query_len=4, identity=True + ), + ), + Cell( + why="selected blocks, Eagle3 verify at batch 2: routed to the one-launch long-context kernel", + spec=dict(seq_lens=(9000, 40000), decode_query_len=4, identity=False), + ), + Cell( + why="selected blocks, decode without speculation at batch 8: the long-context route at T=8", + spec=dict( + seq_lens=(3000, 2100, 9999, 2048, 2049, 4500, 16000, 2200), + decode_query_len=1, + identity=False, + ), + ), + ) + + note = """ + Certified only as the target calls it: one KV head of 8 query heads, + head_dim 128, FP8 q, packed NVFP4 cache, BF16 output, `num_topk_chunks` + left to the dispatcher. The function also takes BF16 / FP8 caches, other + head geometries and an explicit split-K factor (the in-tree tests drive + those); none is certified here. + + The MXFP8 hand-off each route publishes is certified on the long-context + kernel's own entry (`m3_sparse_decode_lc8`); the Triton merge writes it with + the same recipe (`_store_mxfp8_head`). + """ + + def __call__( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + minimax_m3_sparse_attn_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + kv_scale=kv_scale, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + num_topk_chunks=num_topk_chunks, + identity_blocks=identity_blocks, + ) + return output + + def reference( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + """``output`` as the call is supposed to write it: float64 block-sparse attention over the + dequantized NVFP4 blocks (``sparse_attention_reference``), rounded once to ``output``'s dtype. + + Only the packed NVFP4 cache with one KV head is mirrored -- the target's + only configuration. ``kv_scale``, ``num_topk_chunks`` and + ``identity_blocks`` do not change what is computed. + """ + assert k_block_scale is not None and k_paged.shape[1] == 1, ( + "mirrors the one-KV-head NVFP4 cache only" + ) + attn = sparse_attention_reference( + q, + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx[0], + block_table, + seq_lens, + decode_query_len, + sm_scale, + ) + return attn.to(output.dtype) + + def is_valid( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> None: + """The inputs this call takes and answers wrongly. + + Every producer walks a token's first ``min(topk, ceil(kv_len / 128))`` + table entries and indexes the block table with them unchecked, so they + must be real blocks of the request. ``identity_blocks`` is a promise the + multi-query producer relies on: every token of a request reads the same + blocks, all of them. Reads ``topk_idx`` and ``seq_lens`` back to the + host: a test-only cost. + """ + total_q = q.shape[0] + seq = seq_lens.tolist() + max_topk = topk_idx.shape[-1] + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk_idx[:, tok, : min(max_topk, nblocks)] + assert bool(((ids >= 0) & (ids < min(nblocks, block_table.shape[1]))).all()), ( + f"token {tok}: the first min(topk, block count) selected ids must be blocks of the request; " + "the block table is indexed with them unchecked" + ) + if identity_blocks: + assert nblocks <= max_topk and bool( + (ids == torch.arange(nblocks, device=ids.device, dtype=ids.dtype)).all() + ), ( + f"token {tok}: identity_blocks promises the table holds every block of the request in order" + ) + + +m3_sparse_decode = M3SparseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode_lc8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode_lc8.py new file mode 100644 index 000000000000..d0676c3495a2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode_lc8.py @@ -0,0 +1,433 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-launch long-context sparse decode over the selected NVFP4 blocks at 8 query heads (16-CTA clusters). + +Also the plain-torch block-sparse attention and MXFP8 hand-off mirrors the +Triton sparse decode entry (``m3_sparse_decode``) reuses: both compute the same +thing over the same cache. +""" + +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 + +#: The TP8 geometry the kernel is compiled for: 8 query heads on one KV head. +NUM_HEADS = 8 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +#: NVFP4 elements per E4M3 block scale. +SF_VEC = 16 +#: MXFP8 elements per UE8M0 scale, and the o_proj width the hand-off is laid out for. +MX_BLOCK = 32 +WIDTH = NUM_HEADS * HEAD_DIM + +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + + +def dequant_nvfp4_pages( + packed: torch.Tensor, block_scale: torch.Tensor, pages: torch.Tensor, quad_scales: bool +) -> torch.Tensor: + """Dequantize KV head 0 of the given pages of a packed NVFP4 cache to float64 [n, PAGE, HEAD_DIM]. + + ``packed`` is the [pages, kv_heads, PAGE, HEAD_DIM // 2] byte view (low + nibble = even element), ``block_scale`` the [pages, kv_heads, PAGE, + HEAD_DIM // 16] E4M3 bytes ``fused_write_layer_caches_nvfp4`` writes: token + ``t``'s column ``c`` at byte ``t * 8 + c`` of the page's region for K, and at + ``(t // 4) * 32 + c * 4 + t % 4`` for V (``quad_scales``). The per-tensor + global scale is not applied. Every product is exact in float64. + """ + table = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) + rows = packed[pages, 0].to(torch.int64) # [n, PAGE, HEAD_DIM // 2] + nibbles = torch.stack([rows & 0x0F, rows >> 4], dim=-1).reshape(rows.shape[0], PAGE, HEAD_DIM) + cols = HEAD_DIM // SF_VEC + region = block_scale[pages, 0].reshape(rows.shape[0], PAGE * cols) + t = torch.arange(PAGE, device=packed.device)[:, None] + c = torch.arange(cols, device=packed.device)[None, :] + index = (t // 4) * 32 + c * 4 + t % 4 if quad_scales else t * cols + c + scales = region[:, index.reshape(-1)].reshape(rows.shape[0], PAGE, cols) + scales = scales.view(torch.float8_e4m3fn).to(torch.float64) + return table[nibbles] * scales.repeat_interleave(SF_VEC, dim=-1) + + +def sparse_attention_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, +) -> torch.Tensor: + """Block-sparse decode attention over one KV head of a packed NVFP4 cache, float64 [T, heads, HEAD_DIM]. + + Query token ``tok`` belongs to request ``tok // decode_query_len`` at offset + ``tok % decode_query_len`` and attends the first ``min(16, ceil(kv_len / + PAGE))`` entries of ``topk[tok]`` that are not -1, each block truncated at + ``kv_len = max(seq_len - decode_query_len + offset + 1, 0)``. A token that + attends nothing (a CUDA-graph padding row) is all zeros. ``q`` is E4M3 (or + its byte view); softmax and both products are float64. + """ + q64 = ( + q.view(torch.float8_e4m3fn).to(torch.float64) + if q.element_size() == 1 + else q.to(torch.float64) + ) + total_q, num_heads, _ = q64.shape + kgs = float(k_global_scale.reshape(-1)[0]) + vgs = float(v_global_scale.reshape(-1)[0]) + out = torch.zeros(total_q, num_heads, HEAD_DIM, dtype=torch.float64, device=q.device) + seq = seq_lens.tolist() + positions = torch.arange(PAGE, device=q.device) + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + real_topk = min(TOPK, (kv_len + PAGE - 1) // PAGE) + blocks = topk[tok, :real_topk] + blocks = blocks[blocks >= 0].to(torch.int64) + if blocks.numel() == 0: + continue + pages = block_table[req, blocks].to(torch.int64) + keys = ( + dequant_nvfp4_pages(k, k_block_scale, pages, quad_scales=False).reshape(-1, HEAD_DIM) + * kgs + ) + values = ( + dequant_nvfp4_pages(v, v_block_scale, pages, quad_scales=True).reshape(-1, HEAD_DIM) + * vgs + ) + valid = ((blocks[:, None] * PAGE + positions[None, :]) < kv_len).reshape(-1) + logits = (q64[tok] @ keys.T) * sm_scale + logits = logits.masked_fill(~valid[None, :], float("-inf")) + out[tok] = torch.softmax(logits, dim=-1) @ values + return out + + +def mxfp8_handoff_reference( + out: torch.Tensor, num_rows_alloc: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """The MXFP8 o_proj hand-off the merge writes for BF16 rows ``out`` [T, WIDTH]: ``(q8 [T, WIDTH] uint8, + sf [num_rows_alloc * WIDTH // 32] uint8)``. + + Per 32-element block: ``e`` is the biased exponent of ``amax * rcp(448)`` + rounded up when its mantissa is non-zero (capped at 254), the elements are + scaled by ``2^(127 - e)`` and rounded to E4M3 (saturating), and ``e`` lands + in the 128x4-swizzled scale layout the MXFP8 GEMM reads. Entries of ``sf`` + for rows past ``T`` are zero. Computed from the op's own BF16 rows, so a + comparison against it is exact. + """ + total_q = out.shape[0] + x = out.float().reshape(total_q, WIDTH // MX_BLOCK, MX_BLOCK) + amax = x.abs().amax(dim=-1) + sfv = amax * torch.tensor(1.0 / 448.0, dtype=torch.float32, device=out.device) + bits = sfv.view(torch.int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) + e = e.clamp(max=254) + qs = torch.where(amax != 0, ((254 - e) << 23).view(torch.float32), torch.zeros_like(amax)) + xq = (x * qs[..., None]).clamp(-448.0, 448.0) + q8 = xq.to(torch.float8_e4m3fn).view(torch.uint8).reshape(total_q, WIDTH) + num_sf_k_tiles = WIDTH // 128 + sf = torch.zeros(num_rows_alloc * (WIDTH // MX_BLOCK), dtype=torch.uint8, device=out.device) + tok = torch.arange(total_q, device=out.device)[:, None] + k_sf = torch.arange(WIDTH // MX_BLOCK, device=out.device)[None, :] + offset = ( + (tok // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (tok % 32) * 16 + + ((tok % 128) // 32) * 4 + + k_sf % 4 + ) + sf[offset.reshape(-1)] = e.to(torch.uint8).reshape(-1) + return q8, sf + + +class M3SparseDecodeLc8(OpWrapper): + """Sparse decode attention, split-K merge and MXFP8 hand-off of one TP8 decode step in one launch. + + One cluster of 16 CTAs per query token: CTA ``r`` attends the token's + ``r``-th selected block (``topk[tok, r]``) for all 8 heads (transposed MMAs, + the heads on the n dimension), and the cluster merges the 16 partials over + distributed shared memory into the BF16 row ``out[tok]`` [1024]. With ``q8`` / + ``sf`` it also writes the row's MXFP8 copy and swizzled UE8M0 scales for the + o_proj GEMM. All operands are byte views (E4M3 ``q``, packed E2M1 pages, E4M3 + block scales). The kernel releases its PDL dependents right after its grid + wait. Returns ``out``. + + ``prewait`` stages the request length, the query rows (widened to FP16) and + the request's block-table row (up to 8192 entries, longer rows are read from + global memory) before the grid wait; the arithmetic is the same, so both + variants write the same bits. + + Per block: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16 + (exact), QK and PV on FP16 tensor cores with FP32 accumulation, the softmax + probabilities rounded to FP16 for PV, log2-domain softmax with ``ex2.approx``. + ``linear_softmax`` picks the denominator recurrence (the two agree to + rounding); the target picks it per batch shape. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode without speculation, one request: T=1, log-domain recurrence, MXFP8 hand-off, prewait", + spec=dict( + seq_lens=(9000,), decode_query_len=1, linear_softmax=False, mxfp8=True, prewait=True + ), + ), + Cell( + why="Eagle3 verify at one request (the TP8 operating point): T=4, log-domain recurrence, prewait", + spec=dict( + seq_lens=(20000,), + decode_query_len=4, + linear_softmax=False, + mxfp8=True, + prewait=True, + ), + ), + Cell( + why="Eagle3 verify at two requests: T=8 on the linear denominator, prewait at its row limit", + spec=dict( + seq_lens=(5000, 12345), + decode_query_len=4, + linear_softmax=True, + mxfp8=True, + prewait=True, + ), + ), + Cell( + why="a 1M-token request: a block-table row longer than the staged 8192 entries, read from global", + spec=dict( + seq_lens=(1048576,), + decode_query_len=4, + linear_softmax=False, + mxfp8=True, + prewait=True, + ), + ), + Cell( + why="Eagle3 verify at batch 8 (T=32): short requests with -1 tails and zero-length padding rows", + spec=dict( + seq_lens=(16384, 2049, 130, 0, 7777, 4, 30000, 0), + decode_query_len=4, + linear_softmax=True, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="decode without speculation at batch 8: T=8 rows of distinct requests, staged after the wait", + spec=dict( + seq_lens=(3000, 2100, 9999, 2048, 2049, 4500, 16000, 2200), + decode_query_len=1, + linear_softmax=False, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="no hand-off buffers (the variant the call site takes past 128 rows), driven at T=8", + spec=dict( + seq_lens=(5000, 12345), + decode_query_len=4, + linear_softmax=True, + mxfp8=False, + prewait=False, + ), + ), + ) + + note = """ + Certified: `linear_softmax` x hand-off on / off x `prewait`, the variants + the cells drive. The hand-off cells are certified bit-exactly against + `mxfp8_handoff_reference` applied to the op's own BF16 rows; the BF16 rows + themselves within the band the test derives, and the two `prewait` variants + against each other bit for bit. + + `topk` is read for the first `min(16, ceil(kv_len / 128))` entries of each + token; a -1 among them skips that CTA's block, entries past them are never + read. `use_pdl` only changes the launch attribute. The hand-off assumes + `num_sf_k_tiles = 8` (a 1024-wide o_proj input) and buffers with at least + as many rows as the step; `MERGE_MXFP8` sizes them at 128 rows. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + prewait: bool, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_tp8_sparse_decode_lc8( + q, + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + out, + q8, + sf, + num_sf_k_tiles, + decode_query_len, + sm_scale, + linear_softmax, + use_pdl, + prewait, + ) + return out + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + prewait: bool, + ) -> torch.Tensor: + """The BF16 ``out`` [T, 1024] the op is supposed to write: ``sparse_attention_reference`` in float64, + rounded once. ``linear_softmax``, ``use_pdl`` and ``prewait`` do not change what is computed; the + hand-off buffers are checked separately with ``mxfp8_handoff_reference``. + """ + attn = sparse_attention_reference( + q.view(-1, NUM_HEADS, HEAD_DIM), + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + decode_query_len, + sm_scale, + ) + return attn.reshape(out.shape).to(out.dtype) + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + prewait: bool, + ) -> None: + """The inputs this op takes and answers wrongly (``m3_attn_lc8.lc8_applies`` is the target's form of the + layout half). + + The kernel addresses a page's token rows at a fixed 64-byte (data) and + 8-byte (scale) pitch and loads them 16 / 8 bytes at a time, so only the + page stride is read from the views. It reads the selected block ids, + the block table and the hand-off buffers unchecked. Reads ``topk`` and + ``seq_lens`` back to the host: a test-only cost. + """ + for name, t in (("k", k), ("v", v)): + assert t.dim() == 4 and t.shape[1] == 1 and t.shape[2:] == (PAGE, HEAD_DIM // 2), ( + f"{name} must be the [pages, 1, 128, 64] byte view of one KV head" + ) + assert t.stride(3) == 1 and t.stride(2) == HEAD_DIM // 2, ( + f"{name} rows must be dense 64-byte rows: the kernel ignores the row stride" + ) + assert t.data_ptr() % 16 == 0 and t.stride(0) % 16 == 0, ( + f"{name} pages must be 16-byte aligned" + ) + for name, t in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + assert t.dim() == 4 and t.stride(3) == 1 and t.stride(2) == HEAD_DIM // SF_VEC, ( + f"{name} rows must be dense 8-byte rows: the kernel ignores the row stride" + ) + assert t.data_ptr() % 8 == 0 and t.stride(0) % 8 == 0, ( + f"{name} pages must be 8-byte aligned" + ) + assert q.data_ptr() % 16 == 0 and q.stride(0) % 16 == 0 and q.stride(1) % 16 == 0, ( + "q rows are loaded 16 bytes at a time and must be 16-byte aligned" + ) + assert k_global_scale.numel() == 1 and v_global_scale.numel() == 1, ( + "global scales are single values" + ) + total_q = q.shape[0] + assert ( + total_q % decode_query_len == 0 and seq_lens.numel() >= total_q // decode_query_len + ), ( + "every query token must belong to a request: seq_lens is read at tok // decode_query_len unchecked" + ) + seq = seq_lens.tolist() + num_cols = block_table.shape[1] + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk[tok, : min(TOPK, nblocks)] + assert bool(((ids >= -1) & (ids < min(nblocks, num_cols))).all()), ( + f"token {tok}: selected block ids must be -1 or below the request's block count; the block " + "table is indexed with them unchecked" + ) + if q8 is not None: + assert num_sf_k_tiles == WIDTH // 128, ( + f"num_sf_k_tiles must be {WIDTH // 128}: the hand-off rows are {WIDTH} wide and a smaller " + "value overlaps them silently" + ) + assert ( + q8.numel() >= total_q * WIDTH + and sf is not None + and sf.numel() >= ((total_q + 127) // 128 * 128 * (WIDTH // MX_BLOCK)) + ), "the hand-off buffers must hold every row of the step: they are written unchecked" + + +m3_sparse_decode_lc8 = M3SparseDecodeLc8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/msa_kv_update.py new file mode 100644 index 000000000000..ad05b934d4a9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/msa_kv_update.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MSA KV-length correction in one launch (this target's ``MsaKvUpdate`` hook).""" + +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 + +#: The MSA pool's page: one sparse block is one 128-token KV page. +TOKENS_PER_BLOCK = 128 + + +class MsaKvUpdate(OpWrapper): + """Re-derive the MSA decode state after the overlap scheduler moved ``kv_lens``. + + ``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` calls this with its + own buffers at the head of every captured speculative decode step. One + single-CTA launch writes, for ``batch`` requests and ``total_q`` query + tokens (all int32): + + * ``msa_seq_lens[b] = min(kv_lens[b], kv_lens_staged[b])`` and, when given, + ``cu_kv_lens`` as its exclusive prefix sum (``batch + 1`` entries); + * ``out_cache_loc[t]``, the KV write slot of the position token ``t`` + attends up to, through ``block_table``; + * ``n_valid_blocks[t]``, when given, the block count top-k selection is + bounded by. + + It returns None; the cells read the four buffers. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode without speculation at one request: qo = 1, every buffer written", + spec=dict(batch=1, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="decode without speculation at the captured batch 8", + spec=dict(batch=8, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="Eagle3 verify (draft 3): qo = 4 per request, partial acceptance shortens kv_lens below staged", + spec=dict(batch=8, query_len=4, accept="partial", cu_kv=True, n_valid=True), + ), + Cell( + why="the metadata without cu_kv_lens (no extend_ctx promotion): the optional prefix sum is skipped", + spec=dict(batch=4, query_len=4, accept="partial", cu_kv=False, n_valid=True), + ), + Cell( + why="CUDA-graph padding rows (kv_lens 0, pos -1) and no n_valid_blocks buffer", + spec=dict(batch=8, query_len=4, accept="padded", cu_kv=True, n_valid=False), + ), + Cell( + why="positions crossing a page boundary: the slot moves to the previous page of the request", + spec=dict(batch=3, query_len=4, accept="page_edge", cu_kv=True, n_valid=True), + ), + ) + + note = """ + The kernel covers any batch up to the block size's loop (it strides over + 256 threads), and any page size; only `tokens_per_block = 128`, the MSA + pool's page, is certified. + + The op does not bounds-check the request index a token reads + (`q_batch_row`) or the block-table column it derives: both come from the + metadata's own staging, which keeps them in range by construction, and an + out-of-range value reads whatever memory sits there. `is_valid` states both. + """ + + def __call__( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + torch.ops.trtllm.minimax_m3_msa_kv_update( + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + tokens_per_block, + batch, + total_q, + msa_seq_lens, + cu_kv_lens, + out_cache_loc, + n_valid_blocks, + ) + + def reference( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]: + """``(msa_seq_lens[:batch], cu_kv_lens[:batch + 1], out_cache_loc[:total_q], + n_valid_blocks[:total_q])`` as the op is supposed to write them. + + The in-tree metadata's own torch re-derivation + (``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` past its hook + call), statement for statement, returning new tensors rather than + writing the buffers. The buffer arguments only say which outputs exist: + a ``None`` buffer gives a ``None`` result. + """ + page = int(tokens_per_block) + kv_true = torch.minimum(kv_lens[:batch], kv_lens_staged[:batch]) + cu_kv = None + if cu_kv_lens is not None: + cu_kv = torch.zeros(batch + 1, dtype=torch.int32, device=kv_lens.device) + cu_kv[1:] = torch.cumsum(kv_true, 0) + qbr = q_batch_row[:total_q].to(torch.long) + qo_dev = seq_lens[:batch] + pos = kv_true[qbr] - qo_dev[qbr] + q_intra[:total_q] + pos_slot = pos.clamp_min(0) + block_col = torch.div(pos_slot, page, rounding_mode="floor").clamp( + max=int(block_table.shape[1]) - 1 + ) + slots = block_table[qbr, block_col.to(torch.long)] * page + torch.remainder(pos_slot, page) + n_valid = None + if n_valid_blocks is not None: + n_valid = torch.div( + (pos + 1).clamp_min(0) + (page - 1), page, rounding_mode="floor" + ).to(torch.int32) + return kv_true.to(torch.int32), cu_kv, slots.to(torch.int32), n_valid + + def is_valid( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + """The inputs this op takes and answers wrongly. + + Every operand but the block table is read as a dense run from its first + element, so a strided view reads the wrong elements. The request index of + each token and the table row it selects are not bounds-checked. Reads + ``q_batch_row`` back to the host: a test-only cost. + """ + for name, t in ( + ("kv_lens", kv_lens), + ("kv_lens_staged", kv_lens_staged), + ("seq_lens", seq_lens), + ("q_batch_row", q_batch_row), + ("q_intra", q_intra), + ("msa_seq_lens", msa_seq_lens), + ("out_cache_loc", out_cache_loc), + ("cu_kv_lens", cu_kv_lens), + ("n_valid_blocks", n_valid_blocks), + ): + assert t is None or t.is_contiguous(), ( + f"{name} must be contiguous: the kernel indexes it densely from data_ptr and ignores strides" + ) + rows = q_batch_row[:total_q] + assert total_q == 0 or ( + int(rows.min()) >= 0 and int(rows.max()) < min(batch, block_table.shape[0]) + ), ( + "q_batch_row must index a request below batch (and a block-table row): the op reads the " + "request's lengths and table row unchecked" + ) + + +msa_kv_update = MsaKvUpdate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/__init__.py new file mode 100644 index 000000000000..5da86fe24bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collective-communication entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_allreduce.py new file mode 100644 index 000000000000..7e03223a5dba --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_allreduce.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The built-in TP all-reduce, fused with the residual add and the next RMSNorm.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.distributed import AllReduce, AllReduceParams + + +class M3AllReduce(OpWrapper): + """``comm(x, all_reduce_params=params)``: the runtime's ``AllReduce`` over the TP group. ``params`` None: the + plain sum over the ranks (the drafter's o_proj and down projection). With ``RESIDUAL_RMS_NORM`` params: + ``updated = residual + sum_ranks(x)``, ``normed = rmsnorm(updated) * norm_weight``; + returns (normed, updated). The strategy (``AllReduceStrategy.AUTO``: one-shot / two-shot Lamport, NCCL) is the + communicator's per call. + + The layer boundaries the multicast all-reduce does not take run it: prefill and mixed steps, and generation steps + beyond the multicast kernel's rows. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # No cell names a world size: the result does not depend on it, and the certifying matrix runs at every world + # size one node can form (the target's two-tray group of 8 cannot be formed on one node). + CELLS: tuple[Cell, ...] = ( + Cell(why="a prefill chunk", spec=dict(m=2048, hidden=6144)), + Cell(why="a 64-row generation step", spec=dict(m=64, hidden=6144)), + Cell(why="plain sum of eight drafter rows", spec=dict(m=8, hidden=6144, plain=True)), + ) + + note = """ + The reduction order across ranks is the strategy's; `reference` sums the ranks' rows in fp32 in rank order, then + adds the residual and normalizes (one bf16 rounding of the sum and one of the norm). A strategy may round the sum + to bf16 before the residual add, so the residual sum agrees within one bf16 ulp of its operands' magnitude. + """ + + def __call__(self, x: torch.Tensor, comm: AllReduce, params: Optional[AllReduceParams]): + return comm(x, all_reduce_params=params) + + def reference(self, xs, comm: AllReduce, params: Optional[AllReduceParams]): + """``xs``: every rank's ``x`` in rank order (a collective's reference takes the group's inputs).""" + total = torch.stack([x.float() for x in xs]).sum(0) + if params is None: + return total.to(xs[0].dtype) + updated = (params.residual.float() + total).to(xs[0].dtype) + h = updated.float() + normed = ( + h + * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + params.eps) + * params.norm_weight.float() + ) + return normed.to(xs[0].dtype), updated + + +m3_allreduce = M3AllReduce() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_mc_allreduce_norm.py new file mode 100644 index 000000000000..f7054de397dc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_mc_allreduce_norm.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode layer boundary, 8 ranks.""" + +from __future__ import annotations + +from typing import 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 + +_E4M3_MAX = 448.0 + + +def mxfp8_rows(normed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """MXFP8 of BF16 rows as ``trtllm::mxfp8_quantize(normed, True)``: E4M3 [M, H] and 128x4-swizzled UE8M0 bytes. + + One scale per 32 columns, the smallest power of two at or above ``amax / 448`` (biased exponent, round up), + data = row / scale rounded to nearest E4M3. The scale buffer is padded to 128 rows, padding bytes left as the + op leaves them (the comparison masks them out). + """ + m, h = normed.shape + blocks = normed.float().view(m, h // 32, 32) + amax = blocks.abs().amax(-1) + bits = (amax / _E4M3_MAX).view(torch.int32) + byte = (((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) > 0).to(torch.int32)).clamp(0, 254) + scale = torch.exp2(byte.float() - 127.0).unsqueeze(-1) + data = torch.where(amax.unsqueeze(-1) != 0, blocks / scale, torch.zeros_like(blocks)) + data = data.view(m, h).to(torch.float8_e4m3fn) + return data, swizzle_sf(byte.to(torch.uint8)) + + +def swizzle_sf(byte: torch.Tensor) -> torch.Tensor: + """Scatter [rows, H / 32] scale bytes into the 128x4-swizzled buffer the op writes (zero padding).""" + rows, cols = byte.shape + padded_rows = (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (r // 128) * (128 * cols) + ) + out = torch.zeros(padded_rows * cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def sf_written(rows: int, cols: int, device) -> torch.Tensor: + """Mask of the swizzled scale bytes the op writes for ``rows`` rows (the padding rows are left unwritten).""" + return swizzle_sf(torch.ones(rows, cols, dtype=torch.uint8, device=device)).bool() + + +def boundary( + partials: Sequence[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, +) -> tuple[torch.Tensor, ...]: + """The boundary arithmetic once each rank's BF16 partial is known (shared with the finalize entry). + + Ranks are summed in FP32 in rank order, the residual added and rounded to BF16 once; the RMS statistic and the + normalization are done in float64 and rounded to BF16 once. + """ + total = torch.zeros(partials[0].shape, dtype=torch.float32, device=residual.device) + for p in partials: + total = total + p.float() + updated = (total + residual.float()).to(torch.bfloat16) + u = updated.double() + inv = torch.rsqrt(u.pow(2).mean(-1, keepdim=True) + eps) + normed = (u * inv * norm_w.double()).to(torch.bfloat16) + if not quant: + return normed, updated + return (normed, updated, *mxfp8_rows(normed)) + + +class M3McAllreduceNorm(OpWrapper): + """One TP group's boundary: ``updated = bf16(residual + sum_r bf16(x_r + x2_r))``, ``normed = RMSNorm(updated)``. + + Every rank pushes its partial through the multicast mapping ``ws_mc_ptr`` into every rank's exchange, then polls + its own view ``ws_uc`` (re-arming the words it read) and rotates the per-CTA call parities in ``flags``; the + workspace is the ``McAllReduceWorkspace`` bucket of ``maxt`` tokens. All ranks must make the same calls in the + same stream order. ``prepushed`` 2 means a producer already pushed ``x`` (and launched this grid after its + pushes); 3 that two producers pushed ``x`` into exchange set A and ``x2`` into set B. Returns + ``[normed, updated]`` or, with ``quant``, also the MXFP8 copy of ``normed`` (E4M3 rows, swizzled UE8M0). + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The tp8 decode boundaries (hidden 6144, 8 ranks over two trays, eps 1e-6). Buckets 8 / 16 / 48 are separate + # workspaces; the 8-token bucket runs epilogue option 7 at 1, 2, 7, 8 tokens and 15 at 3..6, the others 15. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token at the attention boundary, x only, bucket 8 option 7", + spec=dict(m=1, maxt=8, opts=7, x2=False, quant=False, prepushed=0), + ), + Cell( + why="MoE boundary at 4 tokens: shared + routed summed in-kernel, MXFP8 copy for the next QKV, option 15", + spec=dict(m=4, maxt=8, opts=15, x2=True, quant=True, prepushed=0), + ), + Cell( + why="a full 8-token bucket with the second addend", + spec=dict(m=8, maxt=8, opts=7, x2=True, quant=False, prepushed=0), + ), + Cell( + why="attention boundary whose rows the o_proj GEMV already pushed (prepushed 2)", + spec=dict(m=8, maxt=8, opts=7, x2=False, quant=False, prepushed=2), + ), + Cell( + why="routed engine pushed set A, shared down projection set B (prepushed 3), with the MXFP8 copy", + spec=dict(m=6, maxt=8, opts=15, x2=True, quant=True, prepushed=3), + ), + Cell( + why="12 tokens pushed by the o_proj GEMV into the 16-token bucket", + spec=dict(m=12, maxt=16, opts=15, x2=False, quant=False, prepushed=2), + ), + Cell( + why="the 9..13-token MoE kernel's sum into the 16-token bucket, MXFP8 copy", + spec=dict(m=16, maxt=16, opts=15, x2=False, quant=True, prepushed=0), + ), + Cell( + why="the 48-token bucket at a partial fill", + spec=dict(m=30, maxt=48, opts=15, x2=False, quant=False, prepushed=0), + ), + Cell( + why="the 48-token bucket full, second addend and MXFP8 copy", + spec=dict(m=48, maxt=48, opts=15, x2=True, quant=True, prepushed=0), + ), + # Every other row count of the decode boundaries (dispatch.MC_ALLREDUCE): the bucket and epilogue options the + # target picks for it (ops/m3_mc_allreduce), the second addend and the MXFP8 copy alternating. + *( + Cell( + why=f"{m} tokens: a boundary row count between the ones above", + spec=dict( + m=m, + maxt=8 if m <= 8 else 16 if m <= 16 else 48, + opts=(7, 7, 15, 15, 15, 15, 7, 7)[m - 1] if m <= 8 else 15, + x2=m % 2 == 0, + quant=m % 3 == 0, + prepushed=0, + ), + ) + for m in range(1, 49) + if m not in (1, 4, 8, 16, 30, 48) + ), + # Every other row count pushed by the o_proj (1..16) or the 1..8-token MoE engine (with the MXFP8 copy): + # dispatch.ATTN_PUSH / dispatch.MOE_PUSH. + *( + Cell( + why=f"{m} tokens prepushed: a pushed row count between the ones above", + spec=dict( + m=m, + maxt=8 if m <= 8 else 16, + opts=(7, 7, 15, 15, 15, 15, 7, 7)[m - 1] if m <= 8 else 15, + x2=False, + quant=m <= 8, + prepushed=2, + ), + ) + for m in range(1, 17) + if m not in (8, 12) + ), + ) + + note = """ + The epilogue option bits (1 early residual / weight load, 2 st.async cluster exchange, 4 early parity rotation, + 8 max carveout) are bit-identical by construction; only the four (bucket, option) pairs the target selects are + built. `prepushed` 1 (pushed, grid wait kept) is accepted by the op but no caller passes it. `pdl` changes only + launch overlap and is driven off in the cells. + + `normed` is gated to one BF16 rounding step: the kernel's RMS statistic is a warp / cluster tree of FP32 sums + and its `rsqrtf` is approximate, so a row whose scaled value lands next to a BF16 midpoint can round either + way. `updated` and the MXFP8 copy of the op's own `normed` are bit-exact. + """ + + def __call__( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_mc_allreduce_norm( + x, + x2, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + None, + None, + maxt, + quant, + prepushed, + opts, + ) + + def reference( + self, + per_rank_x: Sequence[torch.Tensor], + per_rank_x2: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute. + + The documented collective exception: it takes every rank's ``x`` (and ``x2``) in rank order, not the calling + rank's. Each rank's partial is ``bf16(x + x2)`` (``x`` alone without ``x2``), whatever pushed it. + """ + partials = [] + for r, x in enumerate(per_rank_x): + if per_rank_x2 is None: + partials.append(x) + else: + partials.append((x.float() + per_rank_x2[r].float()).to(torch.bfloat16)) + return boundary(partials, residual, norm_w, eps, quant) + + def is_valid( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + ) -> None: + h = residual.shape[-1] + # The op checks norm_w's element count, not its layout: a strided view is read as packed memory. + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + # The workspace is addressed as [set][parity][maxt][world][H / 2] words and never bounds-checked: a smaller + # buffer (another bucket's, or one built for fewer ranks) is polled and re-armed out of range. + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + # Flags are per (token, cluster CTA) of the bucket; the op checks the count, not that they belong to the + # workspace. Flags of another bucket hold another call history, and the parity they select is not the one the + # peers push into: the poll then waits on the other half forever or reads last call's rows. + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + + +m3_mc_allreduce_norm = M3McAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..64feae8122d5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MoE finalize + multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode MoE boundary, 8 ranks.""" + +from __future__ import annotations + +from typing import 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 .m3_mc_allreduce_norm import boundary + + +def finalize_partial( + fc2: torch.Tensor, perm: torch.Tensor, scales: torch.Tensor, shared: Optional[torch.Tensor] +) -> torch.Tensor: + """One rank's BF16 partial: ``bf16(bf16(sum_k scale_k * fc2[perm_k]) + shared)``, the routed sum in FP32 in k + order and rounded to BF16, then the shared expert added as stock adds it (one more BF16 rounding). + + The kernel's per-k step is a fused multiply-add; it is reproduced exactly by the float64 product (exact for two + FP32 operands) plus the FP32 accumulator, rounded back to FP32. Dropped entries (perm -1) are skipped. + """ + m, topk = perm.shape + acc = torch.zeros(m, fc2.shape[1], dtype=torch.float32, device=fc2.device) + sc = scales.reshape(m, topk).float() + for k in range(topk): + rows = perm[:, k].long() + keep = (rows >= 0).unsqueeze(-1) + gathered = fc2[rows.clamp_min(0)].double() + step = (gathered * sc[:, k : k + 1].double() + acc.double()).float() + acc = torch.where(keep, step, acc) + if shared is not None: + acc = acc.to(torch.bfloat16).float() + shared.float() + return acc.to(torch.bfloat16) + + +class M3McFinalizeAllreduceNorm(OpWrapper): + """The MoE boundary with the routed experts' finalize folded in. + + Per rank the partial is ``bf16(bf16(sum_k scales[t, k] * fc2[perm[t, k]]) + shared[t])`` (``perm`` -1 = dropped, + ``scales`` fp32 or bf16), then the group's all-reduce, residual add and RMSNorm exactly as + ``m3_mc_allreduce_norm``: same workspace, flags and call-order contract. Returns ``[normed, updated]`` or, with + ``quant``, also the MXFP8 copy of ``normed``. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The tp8 target takes this path from 9 tokens (top-4 of 128 experts, fp32 routing weights, the shared expert + # always present): the 16- and 48-token buckets, epilogue option 15. + CELLS: tuple[Cell, ...] = ( + Cell( + why="the smallest finalize call (9 tokens, 16-token bucket) with the next QKV's MXFP8 copy", + spec=dict(m=9, maxt=16, opts=15, topk=4, quant=True, dropped=0), + ), + Cell( + why="a full 16-token bucket with dropped expert slots", + spec=dict(m=16, maxt=16, opts=15, topk=4, quant=False, dropped=5), + ), + Cell( + why="the 48-token bucket at a partial fill", + spec=dict(m=24, maxt=48, opts=15, topk=4, quant=False, dropped=0), + ), + Cell( + why="the 48-token bucket full, MXFP8 copy", + spec=dict(m=48, maxt=48, opts=15, topk=4, quant=True, dropped=3), + ), + # Every other row count of the unfinalized MoE (dispatch.MC_FINALIZE): the bucket and options the target picks. + *( + Cell( + why=f"{m} tokens: a finalize row count between the ones above", + spec=dict( + m=m, + maxt=16 if m <= 16 else 48, + opts=15, + topk=4, + quant=m % 2 == 0, + dropped=m % 4, + ), + ) + for m in range(9, 49) + if m not in (9, 16, 24, 48) + ), + ) + + note = """ + BF16 routing weights and the 8-token bucket are accepted by the op (option 7 there) but the target only enters + this path at 9 tokens and up, with the built-in MoE's fp32 weights; neither is a cell. The rest is as in + `m3_mc_allreduce_norm`'s note. + """ + + def __call__( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_mc_allreduce_norm( + fc2, + shared, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + perm, + scales, + maxt, + quant, + 0, + opts, + ) + + def reference( + self, + per_rank_fc2: Sequence[torch.Tensor], + per_rank_perm: Sequence[torch.Tensor], + per_rank_scales: Sequence[torch.Tensor], + per_rank_shared: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute, from every rank's MoE outputs in rank order (the + collective exception, as in ``m3_mc_allreduce_norm.reference``).""" + partials = [ + finalize_partial( + fc2, + per_rank_perm[r], + per_rank_scales[r], + None if per_rank_shared is None else per_rank_shared[r], + ) + for r, fc2 in enumerate(per_rank_fc2) + ] + return boundary(partials, residual, norm_w, eps, quant) + + def is_valid( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + ) -> None: + h = residual.shape[-1] + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + # The kernel gathers fc2[perm] without a bound: an index past the FC2 rows reads whatever follows them. + lo, hi = int(perm.min()), int(perm.max()) + assert lo >= -1 and hi < fc2.shape[0], ( + f"perm spans [{lo}, {hi}]; entries must be -1 (dropped) or a row of fc2 (< {fc2.shape[0]})" + ) + + +m3_mc_finalize_allreduce_norm = M3McFinalizeAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding/__init__.py new file mode 100644 index 000000000000..e9f4d8e20b94 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Embedding ops of this target: the generic token-embedding gather.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding/m3_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding/m3_embedding.py new file mode 100644 index 000000000000..f979fdecf376 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding/m3_embedding.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The generic token-embedding gather over the replicated table (an id outside the table gives a zero row).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.embedding import pre_comm_embedding_ops + + +class M3Embedding(OpWrapper): + """``table[ids]`` with a zero row for an id outside ``[0, V)`` (the speculative sampler pads with -1). + + The table is the full (replicated) bf16 embedding [V, 6144]; ids int32 / int64 of any count. The fused decode + embedding (``m3_embedding_firstnorm``) takes decode steps of up to 8 rows; this entry takes every other step. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="a prefill chunk, int32 ids", spec=dict(rows=2048, id_dtype=torch.int32, pad=False) + ), + Cell( + why="a 48-row decode verify, int64 ids", + spec=dict(rows=48, id_dtype=torch.int64, pad=False), + ), + Cell( + why="draft rows padded with -1 ids", spec=dict(rows=32, id_dtype=torch.int32, pad=True) + ), + ) + + note = """ + An exact gather: `reference` must match bit for bit. + """ + + def __call__(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + return pre_comm_embedding_ops(ids, table, 1, 0, None, 0, table.shape[0], False, 0) + + def reference(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] + return rows * valid.unsqueeze(-1).to(table.dtype) + + +m3_embedding = M3Embedding() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding_firstnorm/__init__.py new file mode 100644 index 000000000000..4c6fc92cd49c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding_firstnorm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 decode embedding gather + first-layer RMSNorm as one CuTe DSL kernel (see op.py).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py new file mode 100644 index 000000000000..78563b4db40d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fixed M3 H6144 gather/raw residual/RMS, reusing pinned K3 norm kernel. + +128 threads x 8 BF16 x 6 iterations covers the row. Unlike the original +K3 wrapper this target does not promise FlashInfer's bit-identical reduction. +""" + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +from cutlass.experimental import primitives as prims + +NORM_THREADS = 128 +NORM_VEC = 8 +NORM_WARPS = 4 # capacity: fixed: the CTA's warps + + +@cute.kernel +def m3_embed_norm_kernel( + ids: cute.Tensor, # int32 or int64 [N] + table: cute.Tensor, # bf16 [V, H] + weight: cute.Tensor, # bf16 [H] + raw: cute.Tensor, # bf16 [N, H], out: the rows + out: cute.Tensor, # bf16 [N, H], out: the normed rows + vocab: cutlass.Int64, + eps: cutlass.Float32, + tv_layout: cute.Layout, + tiler_mn: cute.Shape, + hidden: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + prims.griddepcontrol(prims.GridDepAction.WAIT) + + smem = cutlass.utils.SmemAllocator() + reduction_buffer = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((1, NORM_WARPS)), byte_alignment=4 + ) + + tok = cutlass.Int64(ids[bidx]) + valid = (tok >= cutlass.Int64(0)) & (tok < vocab) + row = cutlass.Int64(cutlass.select_(valid, tok, cutlass.Int64(0))) + + gX = cute.local_tile(table, tiler_mn, (row, 0)) + gR = cute.local_tile(raw, tiler_mn, (bidx, 0)) + gY = cute.local_tile(out, tiler_mn, (bidx, 0)) + w_layout = cute.prepend(weight.layout, cute.make_layout((tiler_mn[0],), stride=(0,))) + gW = cute.local_tile(cute.make_tensor(weight.iterator, w_layout), tiler_mn, (0, 0)) + + copy_atom_load = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), table.element_type, num_bits_per_copy=128 + ) + copy_atom_store = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), out.element_type, num_bits_per_copy=128 + ) + thr_copy_X = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_W = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_O = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_mn).get_slice(tidx) + + tXgX = thr_copy_X.partition_S(gX) + tXrX = cute.make_fragment_like(tXgX) + tWgW = thr_copy_W.partition_S(gW) + tWrW = cute.make_fragment_like(tWgW) + tXrW = thr_copy_X.retile(tWrW) + tXgR = thr_copy_O.partition_D(gR) + tXgO = thr_copy_O.partition_D(gY) + tXrO = cute.make_fragment_like(tXgO) + + # The row (zero for an id outside [0, V)), straight to the snapshot slot; the weight. + tXrX.store(cute.zeros_like(tXrX, dtype=table.element_type)) + if valid: + cute.copy(copy_atom_load, tXgX, tXrX) + cute.copy(copy_atom_load, tWgW, tWrW) + cute.copy(copy_atom_store, tXrX, tXgR) + + # flashinfer's RMSNormKernel arithmetic, in its order (see the header). + x = tXrX.load().to(cutlass.Float32) + x_sq = x * x + sum_sq = x_sq.reduce(cute.ReductionOp.ADD, init_val=cutlass.Float32(0.0), reduction_profile=0) + for i in cutlass.range_constexpr(5): + sum_sq = sum_sq + cute.arch.shuffle_sync_bfly(sum_sq, offset=1 << i) + lane = cute.arch.lane_idx() + warp = cute.arch.warp_idx() + if lane == 0: + reduction_buffer[0, warp] = sum_sq + cute.arch.barrier() + total = cutlass.Float32(0.0) + if lane < NORM_WARPS: + total = reduction_buffer[0, lane] + for i in cutlass.range_constexpr(5): + total = total + cute.arch.shuffle_sync_bfly(total, offset=1 << i) + mean_sq = total / cutlass.Float32(hidden) + rstd = cute.math.rsqrt(mean_sq + eps, fastmath=True) + w = tXrW.load().to(cutlass.Float32) + y = x * rstd * (w + cutlass.Float32(0.0)) + tXrO.store(y.to(out.element_type)) + cute.copy(copy_atom_store, tXrO, tXgO) + + +@cute.jit +def m3_embed_norm( + ids: cute.Tensor, + table: cute.Tensor, + weight: cute.Tensor, + raw: cute.Tensor, + out: cute.Tensor, + vocab: cutlass.Int64, + eps: cutlass.Float32, + n_tokens: cutlass.Constexpr[int], + hidden: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + # flashinfer's thread-value layout for one row: thread t, value (v, k) -> column 8 t + v + 1024 k. + blocks = hidden // (NORM_VEC * NORM_THREADS) + tv_layout = cute.make_layout(((NORM_THREADS, 1), (NORM_VEC, blocks)), + stride=((NORM_VEC, 1), (1, NORM_VEC * NORM_THREADS))) # fmt: skip + m3_embed_norm_kernel( + ids, table, weight, raw, out, vocab, eps, tv_layout, (1, hidden), hidden + ).launch( + grid=[n_tokens, 1, 1], + block=[NORM_THREADS, 1, 1], + smem=NORM_WARPS * 4, + stream=stream, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding_firstnorm/op.py new file mode 100644 index 000000000000..276f8c32023b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding_firstnorm/op.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host side of the MiniMax-M3 decode embedding + first-layer RMSNorm (m3_embedding_firstnorm_kernel.py): H 6144, +1..8 rows, into caller-owned buffers. + +Compiled on the first call per (rows, id dtype, table shape, PDL), which must happen outside CUDA-graph capture +(:class:`DecodeEmbedding` compiles at construction). +""" + +import threading +from typing import Dict, Tuple + +import cuda.bindings.driver as cuda_driver +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +from .m3_embedding_firstnorm_kernel import m3_embed_norm as _m3_embed_norm_jit + +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size +MAX_ROWS = 8 # capacity: fixed: the kernel's built rows (decode steps) + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def arg(tensor): + return from_dlpack( + tensor.detach(), + assumed_align=tensor.element_size() + if tensor.ndim == 1 and tensor.dtype in (torch.int32, torch.int64) + else 16, + ) + + +def _check(ids, table, weight, raw, out) -> None: + rows = ids.shape[0] if ids.dim() == 1 else 0 + if not ( + 1 <= rows <= MAX_ROWS + and ids.dtype in (torch.int32, torch.int64) + and ids.is_contiguous() + and table.dtype == torch.bfloat16 + and table.dim() == 2 + and table.shape[1] == HIDDEN + and table.is_contiguous() + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (HIDDEN,) + and all( + t.dtype == torch.bfloat16 and tuple(t.shape) == (rows, HIDDEN) and t.is_contiguous() + for t in (raw, out) + ) + ): + raise ValueError( + "m3_embed_norm: expects ids int32/int64 [1..8], a contiguous bf16 table [V, 6144], bf16 weight [6144] " + "and contiguous bf16 raw / out [rows, 6144]" + ) + + +def _compile(ids, table, weight, eps, raw, out, use_pdl: bool): + key = (ids.shape[0], ids.dtype, tuple(table.shape), table.device, bool(use_pdl)) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_embed_norm compiles on its first call per shape: call it once before capture" + ) + with _lock: + fn = _compiled.get(key) + if fn is None: + stream = cuda_driver.CUstream(torch.cuda.current_stream(table.device).cuda_stream) + fn = _compiled[key] = cute.compile( + _m3_embed_norm_jit, + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + ids.shape[0], + HIDDEN, + bool(use_pdl), + stream, + ) + return fn + + +def m3_embed_norm( + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Write ``table[ids]`` into ``raw`` (a zero row for an id outside ``[0, V)``) and its RMSNorm with ``weight`` / + ``eps`` into ``out``; returns ``(out, raw)``.""" + _check(ids, table, weight, raw, out) + fn = _compile(ids, table, weight, eps, raw, out, use_pdl) + stream = cuda_driver.CUstream(torch.cuda.current_stream(ids.device).cuda_stream) + fn( + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + stream, + ) + return out, raw + + +class DecodeEmbedding: + """Caller-owned buffers of one (rows, id dtype) decode embedding: ``raw`` (the gathered rows, the first residual) + and ``output`` (their RMSNorm), over ``table`` and the first layer's plain (non-Gemma) input norm ``weight`` / + ``eps``. Built (and compiled) before CUDA-graph capture.""" + + def __init__(self, table, weight, eps, rows, id_dtype, use_pdl=True): + assert not torch.cuda.is_current_stream_capturing() + assert table.dtype == torch.bfloat16 and table.is_contiguous() and table.shape[1] == HIDDEN + assert weight.dtype == torch.bfloat16 and weight.shape == (HIDDEN,) + assert 1 <= rows <= MAX_ROWS and id_dtype in (torch.int32, torch.int64) + self.table, self.weight, self.rows, self.id_dtype, self.use_pdl = ( + table, + weight, + rows, + id_dtype, + use_pdl, + ) + self.eps = float(eps) + self.raw = torch.empty(rows, HIDDEN, device=table.device, dtype=torch.bfloat16) + self.output = torch.empty_like(self.raw) + ids = torch.empty(rows, device=table.device, dtype=id_dtype) + _compile(ids, table, weight, self.eps, self.raw, self.output, use_pdl) + + def __call__(self, ids): + assert ids.shape == (self.rows,) and ids.dtype == self.id_dtype and ids.is_contiguous() + return m3_embed_norm( + ids, self.table, self.weight, self.eps, self.raw, self.output, self.use_pdl + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/__init__.py new file mode 100644 index 000000000000..9cc374bdf1cc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GEMM entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_bf16_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_bf16_linear.py new file mode 100644 index 000000000000..5007e1da7386 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_bf16_linear.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The BF16 Eagle3 drafter's projections: ``x @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3Bf16Linear(OpWrapper): + """``F.linear(x, w)`` for bf16 ``x`` [M, K] and a bf16 weight [N, K] (cuBLAS, fp32 accumulation), bf16 out. + + The drafter's fc over the captured target layers (K = 3 x 6144), its fused QKV (K = 2 x 6144, this rank's heads), + o_proj (this rank's input columns), gate_up and down projections run it. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="drafter fc, 32 verify rows", spec=dict(m=32, n=6144, k=18432)), + Cell(why="drafter QKV, nine draft rows", spec=dict(m=9, n=1280, k=12288)), + Cell(why="drafter o_proj, a prefill chunk", spec=dict(m=1024, n=6144, k=1024)), + Cell(why="drafter gate_up, 16 rows", spec=dict(m=16, n=4608, k=6144)), + Cell(why="drafter down, eight rows", spec=dict(m=8, n=6144, k=2304)), + ) + + note = """ + cuBLAS's split and accumulation order: `reference` multiplies in fp32 (one bf16 rounding), within the test's band. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x, w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return (x.float() @ w.float().t()).to(x.dtype) + + +m3_bf16_linear = M3Bf16Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_draft_lowm.py new file mode 100644 index 000000000000..1caeda1aa98d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_draft_lowm.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""BF16 x BF16 -> BF16 low-M GEMM of the Eagle3 drafter's linears at fixed FlashInfer tactics (SM10x).""" + +from __future__ import annotations + +from typing import Mapping, Optional, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: (N, K) of the TP8 drafter shards -> ((max rows, kind, tactic), ...), rows ascending. Kind "d" is FlashInfer's direct +#: SIMT kernel with tactic (block, outputs per block, rows per block); kind "s" is its split-K tensor-core kernel with +#: tactic (mma_m, mma_n, split_k, ab_stages). A row count is served by the first entry whose max rows holds it; rows +#: past the last entry are not served. GB300 cold L2, CUDA-graph timing, us cuBLAS -> low-M at 1 / 4 rows: qkv 11.9 -> +#: 6.4 / 12.0 -> 7.8, o 3.6 -> 2.3 / -, gate_up 11.9 -> 9.4 / 12.1 -> 11.0, down 6.8 -> 5.6 / -, fc 37.9 -> 31.9 / +#: 39.2 -> 35.5. +TACTICS: Mapping[Tuple[int, int], tuple] = { + # qkv + (1280, 12288): ( + (1, "d", (384, 2, 1)), + (2, "d", (128, 4, 2)), + (3, "d", (128, 4, 3)), + (8, "s", (64, 8, 4, 12)), + ), + # o_proj + (6144, 1024): ((1, "d", (32, 4, 1)), (2, "d", (32, 4, 2))), + # gate_up + (4608, 6144): ( + (1, "d", (128, 1, 1)), + (2, "d", (256, 2, 2)), + (3, "s", (128, 8, 3, 6)), + (4, "s", (128, 8, 3, 5)), + (8, "s", (64, 8, 2, 10)), + ), + # down + (6144, 2304): ((1, "d", (96, 1, 1)), (2, "d", (96, 2, 2))), + # the Eagle3 fc over the three captured hidden states + (6144, 18432): ( + (1, "d", (384, 2, 1)), + (2, "d", (64, 4, 2)), + (3, "d", (32, 2, 3)), + (8, "s", (64, 8, 1, 12)), + ), +} + + +def tactic_for(n: int, k: int, m: int) -> Optional[Tuple[str, tuple]]: + """(kind, tactic) of an [m, k] x [n, k]^T call, or None when the table does not serve it.""" + for rows, kind, tactic in TACTICS.get((n, k), ()): + if m <= rows: + return kind, tactic + return None + + +#: The drafter linear each (N, K) of the table is. +_NAMES: Mapping[Tuple[int, int], str] = { + (1280, 12288): "qkv", + (6144, 1024): "o_proj", + (4608, 6144): "gate_up", + (6144, 2304): "down", + (6144, 18432): "fc", +} + + +class M3DraftLowM(OpWrapper): + """``out`` [M, N] BF16 = ``x`` [M, K] BF16 @ ``weight`` [N, K]^T BF16, FP32 accumulation. + + ``kind`` / ``tactic`` pick FlashInfer's kernel: "d" runs ``run_direct_dense`` (SIMT, one dot product per output in + K order), "s" runs ``run_splitk_dense`` (tcgen05 MMA, K split over ``split_k`` CTAs and summed in FP32). Both + round once to BF16. The kernel is compiled on its first call per (M, N, K, tactic, PDL); the target compiles every + covered call before warmup (modeling's ``_prebuild_draft_lowm``). + """ + + ARCHS = frozenset({Arch.SM_103}) + # Every row count of every table entry. + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"drafter {_NAMES[nk]}, {m} rows: {tactic_for(nk[0], nk[1], m)[0]} tactic", + spec=dict(nk=nk, m=m), + ) + for nk, table in TACTICS.items() + for m in range(1, table[-1][0] + 1) + ) + + note = """ + The table is the TP8 drafter's: its shapes are the per-rank shards of the Eagle3 layer at 8-way tensor + parallelism (64 query and 4 KV heads, intermediate 18432) and the replicated fc. Certified at those shapes and + rows only; FlashInfer's kernels take other shapes, which are not measured here. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + ) -> torch.Tensor: + if kind == "d": + from flashinfer.gemm.kernels.dense_bf16_gemm_direct import ( + DirectTactic, + run_direct_dense, + ) + + run_direct_dense(x, weight.t(), out, use_pdl, DirectTactic(*tactic)) + else: + from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import ( + SplitKTactic, + run_splitk_dense, + ) + + run_splitk_dense(x, weight.t(), None, out, use_pdl, SplitKTactic(*tactic)) + return out + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + ) -> torch.Tensor: + """bf16(x @ weight^T) summed in fp64 and rounded once; ``out``, ``kind``, ``tactic`` and ``use_pdl`` only + choose where and how the op writes.""" + return (x.double() @ weight.double().t()).to(torch.bfloat16) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + ) -> None: + """The kernels index ``x`` and ``out`` densely: a strided view is read as if it were packed.""" + assert x.dim() == 2 and x.is_contiguous(), "x must be a contiguous [M, K] matrix" + assert weight.dim() == 2 and weight.is_contiguous(), ( + "weight must be a contiguous [N, K] matrix" + ) + assert out.is_contiguous() and tuple(out.shape) == (x.shape[0], weight.shape[0]), ( + "out must be a contiguous [M, N] matrix" + ) + + +m3_draft_lowm = M3DraftLowM() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_fp32_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_fp32_linear.py new file mode 100644 index 000000000000..8d7374810675 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_fp32_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The router's fp32 projection beyond the decode routers' rows: ``float(x) @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + + +class M3Fp32Linear(OpWrapper): + """fp32 router logits ``F.linear(float(x[M, 6144]), w[128, 6144])`` for prefill and steps above 48 rows (the + fp32 GEMV and the tensor-core router take 1..48). Whatever fp32 matmul precision the process runs with applies + (the engine leaves torch's default).""" + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="a prefill chunk", spec=dict(m=2048, e=128, k=6144)), + Cell( + why="a generation step above the tensor-core router's 48 rows", + spec=dict(m=64, e=128, k=6144), + ), + ) + + note = """ + Under torch's default fp32 matmul precision the product may be TF32: the test's band covers that rounding against + the true fp32 reference. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x.to(torch.float32), w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + with true_fp32_matmul(): + return x.float() @ w.t() + + +m3_fp32_linear = M3Fp32Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv.py new file mode 100644 index 000000000000..17c48fe4f7c6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 x MXFP8 GEMV (tcgen05, block-scaled, CuTe DSL) on an activation its producer already quantized.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..moe.m3_moe_engine import mxfp8_dequant +from ..mx_gemv import op as _gemv + +__all__ = ["PROJECTIONS", "gemv_is_valid", "m3_mx_gemv", "served_cells"] + +# (N, K) of this target's MXFP8 projections at tp8, the GEMV's tuned shapes. The shared expert's projections +# (768 x 6144, 6144 x 384) run inside the MoE engine at decode rows and on the generic GEMM otherwise. +QKV_SPARSE = (1536, 6144) +QKV = (1280, 6144) +O_PROJ = (6144, 1024) +DENSE_GATE_UP = (3072, 6144) +DENSE_DOWN = (6144, 1536) +PROJECTIONS = { + "sparse-layer QKV": QKV_SPARSE, + "dense-attention QKV": QKV, + "o_proj": O_PROJ, + "dense gate_up": DENSE_GATE_UP, + "dense down": DENSE_DOWN, +} + + +def served_cells(named: tuple[Cell, ...], x_bf16: bool, max_m: int) -> tuple[Cell, ...]: + """A cell for every (projection, rows) the GEMV serves up to ``max_m`` rows that ``named`` does not list. + + Certification is per projection: a projection's split-K, ring and MMA width differ, so each (N, K, M) the + decode dispatch can send here is its own cell. + """ + listed = {(cell.spec["nk"], cell.spec["m"]) for cell in named} + return tuple( + Cell(why=f"{name}, {m} rows", spec=dict(nk=nk, m=m)) + for name, nk in PROJECTIONS.items() + for m in range(1, max_m + 1) + if (nk, m) not in listed and _gemv.serves(nk[0], nk[1], m, x_bf16) + ) + + +def gemv_reference( + x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor +) -> torch.Tensor: + """bf16(x @ dequant(weight)^T) for fp32 activation values x [M, K], summed in fp64 and rounded once.""" + return (x.double() @ mxfp8_dequant(weight, weight_scale).double().t()).to(torch.bfloat16) + + +def gemv_is_valid(weight: torch.Tensor, weight_scale: torch.Tensor) -> None: + n, k = weight.shape + # The kernel reads both from data_ptr() as dense buffers at the layout it was compiled for. + assert weight.dtype == torch.float8_e4m3fn and weight.is_contiguous(), ( + "weight must be contiguous e4m3 [N, K]" + ) + assert weight_scale.numel() == n * k // 32, ( + "weight_scale must be the 128x4-swizzled UE8M0 buffer [N * K / 32]" + ) + + +class M3MxGemv(OpWrapper): + """bf16 [M, N] = xq [M, K] (e4m3, 128x4-swizzled UE8M0 scales ``xsf``) x weight [N, K]^T (e4m3, swizzled + ``weight_scale``), M <= 48; None when the shape is not served (the caller falls back). + + ``trigger_early``: release the PDL dependents before the epilogue. ``cfg``: a (split, ring[, contig, nacc, a_tmem, + st_async]) tuning override; default the table's. ``use_pdl``: default the environment's. ``ar`` = (multicast + workspace address, flags address, rank, world, bucket tokens, mode): also push the rows into the next boundary + all-reduce's exchange (then ``op.last_pushed`` is True). + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The projections' (N, K) at the token counts each build differs in: the MMA width (8 / 16 / 32 / 48 columns) a + # call uses is the smallest that holds M, and the sparse-layer QKV takes another split at 2 tokens and above 8. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token (split 8)", spec=dict(nk=QKV_SPARSE, m=1)), + Cell( + why="sparse-layer QKV, 2 tokens: its split-4 config, A from shared memory", + spec=dict(nk=QKV_SPARSE, m=2), + ), + Cell( + why="sparse-layer QKV, 4 rows: one request verifying 3 drafts", + spec=dict(nk=QKV_SPARSE, m=4), + ), + Cell( + why="sparse-layer QKV, 12 rows: the wide split-4 config", + spec=dict(nk=QKV_SPARSE, m=12), + ), + Cell(why="dense-attention QKV, one token", spec=dict(nk=QKV, m=1)), + Cell(why="dense-attention QKV, 13 rows: the 16-column MMA", spec=dict(nk=QKV, m=13)), + Cell(why="o_proj, one token (split 4)", spec=dict(nk=O_PROJ, m=1)), + Cell(why="o_proj, 24 rows: the 32-column MMA", spec=dict(nk=O_PROJ, m=24)), + Cell(why="dense gate_up, one token", spec=dict(nk=DENSE_GATE_UP, m=1)), + Cell(why="dense down, 8 tokens", spec=dict(nk=DENSE_DOWN, m=8)), + ) + # Every other (projection, rows) the decode dispatch can send here (dispatch.MX_GEMV_ROWS). + CELLS = CELLS + served_cells(CELLS, x_bf16=False, max_m=48) + + note = """ + The kernel sums in fp32 in split-K order where `reference` sums in fp64: the bf16 outputs agree within the band the + test derives. Not covered: the all-reduce push (`ar`), which needs a multicast exchange; `cfg` overrides beyond the + table (the target passes none). Shapes outside the table and M > 48 return None. + """ + + def __call__( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv(xq, xsf, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The dequantized product; the launch options are ignored.""" + return gemv_reference(mxfp8_dequant(xq, xsf), weight, weight_scale) + + def is_valid( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + m, k = xq.shape + assert xq.dtype == torch.float8_e4m3fn and xq.is_contiguous(), ( + "xq must be contiguous e4m3 [M, K]" + ) + # A linear-layout [M, K / 32] scale buffer is read as the swizzled one: wrong scales, no error. + assert xsf.numel() >= (m + 127) // 128 * 128 * (k // 32), ( + "xsf must be the 128x4-swizzled scale buffer (rows padded to 128), not the linear one" + ) + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv = M3MxGemv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..08abc2249da3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv_bf16.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 GEMV (CuTe DSL) on a bf16 activation, quantized to MXFP8 in the kernel's prologue.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..moe.m3_moe_engine import mxfp8_quant_dequant +from ..mx_gemv import op as _gemv +from .m3_mx_gemv import ( + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + gemv_is_valid, + gemv_reference, + served_cells, +) + +__all__ = ["m3_mx_gemv_bf16"] + + +class M3MxGemvBf16(OpWrapper): + """bf16 [M, N] = mxfp8_quantize(x) x weight [N, K]^T for a contiguous, 16-byte aligned bf16 ``x`` [M, K]; the + quantization is mxfp8_quantize's (UE8M0 scale per 32 = the smallest power of two >= amax / 448). None when the + shape is not served. The other arguments as ``m3_mx_gemv``'s. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The prologue quantizes one 32-element block per thread and round: the projections at the token counts the + # target sends here (no producer-quantized activation), up to the widest each serves. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token", spec=dict(nk=QKV_SPARSE, m=1)), + Cell(why="sparse-layer QKV, 2 tokens: its split-4 config", spec=dict(nk=QKV_SPARSE, m=2)), + Cell( + why="sparse-layer QKV, 16 rows: the most its split-4 prologue serves", + spec=dict(nk=QKV_SPARSE, m=16), + ), + Cell(why="dense-attention QKV, 8 tokens", spec=dict(nk=QKV, m=8)), + Cell(why="o_proj, one token", spec=dict(nk=O_PROJ, m=1)), + Cell(why="o_proj, 48 rows: the 48-column MMA", spec=dict(nk=O_PROJ, m=48)), + Cell( + why="dense gate_up, 16 rows: the most split 4 at K 6144 serves", + spec=dict(nk=DENSE_GATE_UP, m=16), + ), + Cell(why="dense down, one token", spec=dict(nk=DENSE_DOWN, m=1)), + Cell( + why="dense down, 32 rows: the most split 2 at K 1536 serves", + spec=dict(nk=DENSE_DOWN, m=32), + ), + ) + # Every other (projection, rows) the decode dispatch can send here (dispatch.MX_GEMV_BF16_MAX_TOKENS). + CELLS = CELLS + served_cells(CELLS, x_bf16=True, max_m=48) + + note = """ + As `m3_mx_gemv`. The in-kernel quantization is mxfp8_quantize's bit for bit except where the kernel's + rcp.approx(448) and an exact 1/448 put a block's amax on different sides of a power of two (not observed on the + cells). Not covered: `ar`, `cfg` overrides. Returns None for rows beyond 8 rounds of the prologue per CTA. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv_bf16(x, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The MXFP8 round trip of x times the dequantized weight; the launch options are ignored.""" + return gemv_reference(mxfp8_quant_dequant(x), weight, weight_scale) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv_bf16 = M3MxGemvBf16() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mxfp8_gemm.py new file mode 100644 index 000000000000..8fc131a7bb61 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mxfp8_gemm.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-graph MXFP8 GEMM on an activation its producer already quantized (rows the GEMV does not serve).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_gemm_autotuned + +from .m3_mxfp8_linear import mxfp8_reference + + +class M3Mxfp8Gemm(OpWrapper): + """``act @ weight^T`` for an MXFP8 activation (e4m3 [M, K] + 128x4-swizzled UE8M0 scales) and an MXFP8 weight + [N, K] (swizzled scales), bf16 out. The GEMM backend per token bucket is the AutoTuner's (the same runner, op key + and tactic cache as the engine-tuned quantize + GEMM); ``tune`` profiles in the engine's tuning pass.""" + + ARCHS = frozenset({Arch.SM_103}) + + # The decode dispatch hands it the boundary / merge MXFP8 activation when the GEMV declines rows it serves: the + # sparse-layer QKV and the o_proj at the tp8 shards. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, 4 rows", spec=dict(m=4, n=1536, k=6144)), + Cell(why="o_proj, 32 rows", spec=dict(m=32, n=6144, k=1024)), + ) + + note = """ + Exactly the MXFP8 product (fp32 accumulation, one bf16 rounding): `reference` dequantizes both operands and + multiplies in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + return mxfp8_gemm_autotuned(act, act_scale, weight, weight_scale, output_dtype, tune=tune) + + def reference( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + """The dequantized product; ``tune`` is ignored.""" + return mxfp8_reference(act, act_scale, weight, weight_scale).to(output_dtype) + + +m3_mxfp8_gemm = M3Mxfp8Gemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mxfp8_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mxfp8_linear.py new file mode 100644 index 000000000000..1bdc1829508d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mxfp8_linear.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One MXFP8 projection on the engine-tuned path: dynamic MXFP8 quantize of a bf16 activation + block-scaled GEMM.""" + +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, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_quantize_gemm_autotuned + +_BLOCK = 32 + + +def _unswizzle(scale: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """[rows, cols] UE8M0 bytes from the 128x4-swizzled layout (rows padded to 128, columns to 4).""" + pad_r = (rows + 127) // 128 * 128 + pad_c = (cols + 3) // 4 * 4 + s = scale.view(torch.uint8).reshape(pad_r // 128, pad_c // 4, 32, 4, 4) + return s.permute(0, 3, 2, 1, 4).reshape(pad_r, pad_c)[:rows, :cols] + + +def _dequant(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + rows, k = x.shape + s = _unswizzle(scale, rows, k // _BLOCK).to(torch.int32) + factor = torch.pow(2.0, (s - 127).to(torch.float32)) + return ( + x.view(torch.float8_e4m3fn).float().view(rows, k // _BLOCK, _BLOCK) * factor.unsqueeze(-1) + ).view(rows, k) + + +def mxfp8_reference(act, act_scale, weight, weight_scale) -> torch.Tensor: + """fp32 product of two dequantized MXFP8 operands (swizzled UE8M0 scales).""" + with true_fp32_matmul(): + return _dequant(act, act_scale) @ _dequant(weight, weight_scale).t() + + +class M3Mxfp8Linear(OpWrapper): + """``x @ weight^T`` for bf16 ``x`` [M, K] and an MXFP8 weight [N, K] (e4m3 + 128x4-swizzled UE8M0 scales), bf16 + out, on the backend the engine selected: + + * ``tuned_quantize_gemm`` (decode CUDA graphs once the engine tuned per-bucket backends): the AutoTuner's quantize + and GEMM per token bucket (``tune`` profiles them in the engine's tuning pass); + * otherwise the native ``mxfp8_quantize`` (swizzled scales) and then FlashInfer's MXFP8 GEMM (``flashinfer``) or + the CUTLASS block-scaled GEMM, through the autotuner while the engine tunes it (``native_tuning``), scaled by + ``alpha`` (fp32 [1], 1.0 for these checkpoints). + + Every MXFP8 projection of the target (the attention QKV and o_proj, the dense MLPs and the shared expert) runs it + outside the decode-step kernels. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="sparse-layer QKV, a prefill chunk, native GEMM", + spec=dict(m=2048, n=1536, k=6144, path="native"), + ), + Cell( + why="o_proj, a prefill chunk, native GEMM", + spec=dict(m=2048, n=6144, k=1024, path="native"), + ), + Cell( + why="dense down projection, eager decode, native GEMM", + spec=dict(m=4, n=6144, k=1536, path="native"), + ), + Cell( + why="shared gate_up, FlashInfer GEMM", + spec=dict(m=64, n=768, k=6144, path="flashinfer"), + ), + Cell(why="o_proj, tuned quantize + GEMM", spec=dict(m=64, n=6144, k=1024, path="tuned")), + ) + + note = """ + Exactly the MXFP8 product of the MXFP8 round trip of x (fp32 accumulation, one bf16 rounding): `reference` + quantizes x with the same block rule (UE8M0 scale = the block's power-of-two ceiling over e4m3's 448) and + multiplies dequantized operands in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + if tuned_quantize_gemm: + return mxfp8_quantize_gemm_autotuned(x, weight, weight_scale, output_dtype, tune=tune) + act, act_scale = torch.ops.trtllm.mxfp8_quantize(x, True) + if flashinfer: + return torch.ops.trtllm.flashinfer_mm_mxfp8( + act, act_scale, weight, weight_scale, output_dtype + ) + gemm = ( + torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned + if native_tuning + else torch.ops.trtllm.mxfp8_mxfp8_gemm + ) + return gemm(act, act_scale, weight, weight_scale, alpha, output_dtype) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + """The MXFP8 round trip of ``x`` times the dequantized weight, times ``alpha``; the backend flags are + ignored.""" + m, k = x.shape + blocks = x.float().view(m, k // _BLOCK, _BLOCK) + amax = blocks.abs().amax(-1, keepdim=True) + exp = torch.ceil(torch.log2(torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)))) + scale = torch.pow(2.0, exp) + q = (blocks / scale).to(torch.float8_e4m3fn).float() * scale + with true_fp32_matmul(): + out = q.view(m, k) @ _dequant(weight, weight_scale).t() + return (out * alpha.float()).to(output_dtype) + + +m3_mxfp8_linear = M3Mxfp8Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/long_select.py new file mode 100644 index 000000000000..c01a82f0da41 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/long_select.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 block select for MiniMax-M3 decode rows of 129 to 9216 blocks. + +Same contract and bit-identical output as ``torch.ops.trtllm.minimax_m3_select_blocks`` on per-KV-head scores +[num_kv_heads, n_blocks, total_q]: init / local forcing, per-query valid-block count, ties to the smaller block index, +forced -inf picks emitted as -1, ascending block ids with -1 tail padding. + +One 1024-thread CTA per (query, KV head) row holds the row's ordered score keys in registers (2, 4 or 9 per +thread, by the row's valid count). The weakest of the 16 maxima of the interleaved groups ``block % 16`` is a lower +bound on the 16th-largest key, so the top 16 is among the keys at or above it; those are compacted to shared memory +and ranked exactly by (key, block). More than 1024 keys tied at the bound (degenerate rows) rank every candidate +against the whole row (re-read). + +The kernel is ``torch.ops.trtllm.minimax_m3_select_long``, called through this target's ``m3_select_long`` entry. +""" + +from typing import Optional + +import torch + +from ..ops.attention.m3_select_long import MAX_BLOCKS, m3_select_long + + +def long_select_blocks( + max_score_kv: torch.Tensor, + *, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, +) -> Optional[torch.Tensor]: + """Top-16 select over per-KV-head scores [num_kv_heads, n_blocks, total_q] with 129..9216 blocks. + + Returns the head-major backed [total_q, num_kv_heads, 16] int32 table, or None outside this kernel's shapes + (callers keep the reference op). ``early_trigger`` releases the dependent sparse decode right after the select's + grid wait (``SELECT_EARLY_MAX_ROWS`` in modeling.py). + """ + if ( + max_score_kv.dtype != torch.float32 + or not max_score_kv.is_cuda + or max_score_kv.dim() != 3 + or not 128 < max_score_kv.shape[1] <= MAX_BLOCKS + or n_valid_blocks.dtype != torch.int32 + or not n_valid_blocks.is_cuda + or n_valid_blocks.shape[0] < max_score_kv.shape[2] + ): + return None + num_kv_heads, _, total_q = max_score_kv.shape + out = torch.empty( + (num_kv_heads, total_q, 16), dtype=torch.int32, device=max_score_kv.device + ).permute(1, 0, 2) + return m3_select_long( + max_score_kv, + n_valid_blocks.contiguous(), + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_attn_lc8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_attn_lc8.py new file mode 100644 index 000000000000..2c4c55cf089f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_attn_lc8.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 long-context sparse decode attention at 8 query heads in one launch. + +Replaces the generic chain's split-K producer (``_gqa_sparse_decode_kernel``, one block per chunk) and the LSE +merge (``_merge_topk_attn_out_kernel``, with its MXFP8 copy for the o_proj) once the top-k block table exists. + +Grid: one cluster of 16 CTAs per query token; CTA ``r`` attends the token's ``r``-th selected block for all 8 heads, +keeps the partial in shared memory, and the cluster merges the 16 partials over DSMEM into the BF16 row and, +optionally, its MXFP8 copy and swizzled UE8M0 scales exactly as ``_store_mxfp8_head`` lays them out. + +Per-block math follows the Triton producer: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16, QK and PV +on FP16 tensor cores with FP32 accumulation, log2-domain softmax with either the linear denominator or the +log-domain recurrence. Accumulation order inside the MMAs and the reductions differs, so results match to rounding, +not bit for bit. + +Schedule: the kernel releases its PDL dependents right after its grid wait, so the o_proj launches and streams its +weights while attention runs. Up to ``LC8_PREWAIT_MAX_ROWS`` (modeling.py) query rows it also stages the request +length, the query rows and the request's block-table row before that wait, overlapping them with the select ahead +when it releases this grid early (``SELECT_EARLY_MAX_ROWS`` in modeling.py). Neither changes a value. + +Geometry: 8 query heads per KV head (TP8), one KV head, head_dim 128, page 128, top-k 16. The kernel is +``torch.ops.trtllm.minimax_m3_tp8_sparse_decode_lc8``, called through this target's ``m3_sparse_decode_lc8`` entry. +""" + +import torch + +from ..ops.attention.m3_sparse_decode_lc8 import NUM_HEADS, WIDTH, m3_sparse_decode_lc8 + + +def _prewait_max_rows() -> int: + """``LC8_PREWAIT_MAX_ROWS`` of the target's dispatch (imported at call time: modeling imports this module).""" + from .. import modeling + + return modeling.LC8_PREWAIT_MAX_ROWS + + +def lc8_applies(q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output) -> bool: + """Whether the step's operands fit the kernel: TP8 geometry, packed NVFP4 pages, FP8 q, BF16 [T, 8, 128] out.""" + if q.dim() != 3 or tuple(q.shape[1:]) != (NUM_HEADS, 128) or q.dtype != torch.float8_e4m3fn: + return False + if q.stride(2) != 1 or q.stride(1) % 16 or q.stride(0) % 16 or q.data_ptr() % 16: + return False + if int(k_paged.shape[1]) != 1 or int(topk_idx.shape[0]) != 1 or int(topk_idx.shape[-1]) != 16: + return False + for t in (k_paged, v_paged): + if ( + t.stride(3) != 1 + or t.stride(2) != 64 + or int(t.shape[3]) != 64 + or t.data_ptr() % 16 + or t.stride(0) % 16 + ): + return False + for t in (k_block_scale, v_block_scale): + if t is None or t.stride(3) != 1 or t.stride(2) != 8 or t.data_ptr() % 8 or t.stride(0) % 8: + return False + return ( + output.dtype == torch.bfloat16 + and output.dim() == 3 + and output.is_contiguous() + and tuple(output.shape) == tuple(q.shape) + ) + + +def m3_lc8_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + linear_softmax=True, + mxfp8=None, + use_pdl=True, +): + """Drop-in for ``minimax_m3_sparse_attn_decode`` on the TP8 NVFP4 shape; callers check ``lc8_applies`` first. + + ``topk_idx`` is the head-major [1, total_q, 16] view the Triton producer takes. ``mxfp8`` is ``(q8, sf)`` from + ``MERGE_MXFP8.buffers(1024, dev)`` to also emit the o_proj hand-off. + """ + total_q = int(q.shape[0]) + q8, sf = ( + (mxfp8[0].view(torch.uint8), mxfp8[1].view(torch.uint8)) + if mxfp8 is not None + else (None, None) + ) + m3_sparse_decode_lc8( + q.view(torch.uint8), + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx[0], + block_table, + seq_lens, + output.view(total_q, WIDTH), + q8, + sf, + WIDTH // 128, + int(decode_query_len), + float(sm_scale), + bool(linear_softmax), + bool(use_pdl), + total_q <= _prewait_max_rows(), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_draft_step.py new file mode 100644 index 000000000000..917f04005cc9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_draft_step.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused greedy Eagle3 draft pass for the MiniMax-M3 drafter's vocab-sharded BF16 head (TP decode, <= 8 requests). + +One kernel replaces a draft pass's head and the draft loop's per-pass bookkeeping, at the in-tree precision: + +* The drafter's BF16 lm_head shard is kept a second time, built by ``prepare_head`` before the executor's warmup, in + ``mma.m16n8k16`` A-fragment order (``bf16_fragments``, the same values). The kernel streams it with BF16 MMAs and + FP32 accumulation over the BF16 head input rows. +* Each logit is rounded to BF16 (the in-tree logits are the BF16 head GEMM's output) and keyed + ``ordered(value) << 32 | ~global_index`` with -0 as +0, so the larger value wins and ties take the lowest index, as + the in-tree ``torch.max`` / ``torch.argmax`` over the gathered logits (``greedy_sample_draft_with_tp_gather``). +* The last CTA exchanges each row's key across the TP ranks through a multicast mapping (Lamport slots, see + ``m3_shard_accept``), then stores the draft token, the next pass's input id and position, and advances the draft KV + lengths as the in-tree draft loop does between passes. + +The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/draftStep.cu) and called through the +catalog entry ``ops.spec.m3_draft_step``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from ..ops.spec.m3_draft_step import MAX_TILES_PER_CTA, MAX_TOKENS, bf16_fragments +from ..ops.spec.m3_draft_step import m3_draft_step as _draft_step_op +from .m3_shard_accept import AcceptWorkspace, create_workspace + + +@dataclass(eq=False) +class DraftHead: + """The streamed copy of one rank's draft lm_head shard and its first global id: ``wf`` is the fragment-ordered BF16 + copy (``bf16_fragments``).""" + + weight: torch.Tensor + wf: torch.Tensor + vocab_offset: int + + +def head_supported(lm_head) -> bool: + """A bias-free BF16 head whose local shard the kernel can own: whole 16-row tiles, K a multiple of 16, at most + MAX_TILES_PER_CTA tiles per CTA of a one-CTA-per-SM grid.""" + weight = getattr(lm_head, "weight", None) + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: + return False + if getattr(lm_head, "bias", None) is not None: + return False + n, k = weight.shape + if n % 16 or k % 16: + return False + return (n // 16 + _num_sms() - 1) // _num_sms() <= MAX_TILES_PER_CTA + + +def prepare_head(lm_head, tp_rank: int) -> DraftHead: + """The fragment-ordered BF16 copy of ``lm_head``'s local shard (global ids from tp_rank * rows). + Built once, eagerly, before the executor's warmup (``M3Eagle3Worker.prebuild_draft_head``).""" + weight = lm_head.weight + offset = tp_rank * weight.shape[0] + with torch.no_grad(): + wf = bf16_fragments(weight.detach()) + return DraftHead(weight=weight, wf=wf, vocab_offset=offset) + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace ([2][MAX_TOKENS][world] slots), created collectively on the first eager call; None + during a CUDA-graph capture that finds it missing.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + # A capture that finds it missing is the caller's miss (it is built on the eager steps). + return None + _WS = create_workspace(mapping, device, MAX_TOKENS) + return _WS + + +def draft_step( + head: DraftHead, + x: torch.Tensor, + rows: int, + x_rows: Optional[torch.Tensor], + mapping, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, +) -> bool: + """One fused greedy draft pass over ``rows`` requests (see the module docstring); False when the exchange + workspace is missing during a capture. Collective over the TP group. + + Args: + head: the prepared head of this rank (``prepare_head``), BF16 fragments. + x: [*, K] BF16 head input rows (the drafter's final-norm output); row ``x_rows[b]`` (int64) or ``b`` is + request b's. + draft_out: int32 [rows, >= step + 1], unit column stride: column ``step`` receives the tokens. + ids_next / pos_next: int32 [rows]: the tokens, and ``pos_src[x_rows[b] or b] + 1`` (pos_src int32). + kv_lens: int32 draft KV lengths: ``-= draft_len - num_acc`` with ``seq_lens = 1`` on pass 0 when ``num_acc`` + is given, ``+= 1`` on later passes. + """ + ws = workspace(mapping, x.device) + if ws is None: + return False + _draft_step_op( + x, + x_rows, + rows, + head.wf, + head.vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + _num_sms(), + True, + None, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_force_vec.py new file mode 100644 index 000000000000..b4de37794ca9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_force_vec.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance (the in-tree speculative worker's +``force_num_accepted_tokens``), drawn exactly as its ``_apply_force_accepted_tokens`` does, in one launch. + +Triton: one CTA bumps the shared step counter and draws each request's count from the fixed-seed pool with the +in-tree (counter, slot) hash. Reached through the ``ops/spec`` entry ``m3_force_vec``. +""" + +import triton +import triton.language as tl + + +@triton.jit +def force_vec_kernel( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK: tl.constexpr, +): + c = tl.load(counter) + 1 + slots = tl.arange(0, BLOCK) + tl.store(counter, c) + idx = (c * counter_stride + slots.to(tl.int64) * slot_stride) & pool_mask + live = slots < batch_size + r = tl.load(pool + idx, mask=live, other=1.0) + v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) + tl.store(out + slots, v, mask=live) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_fused_score_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_fused_score_select.py new file mode 100644 index 000000000000..3ef07ce6dbff --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_fused_score_select.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fused index score + top-16 select of a decode span of at most ``dispatch.FUSED_SCORE_SELECT_MAX_REQUESTS`` +requests (``ops/attention/m3_fused_score_select``), with its scratch and arrival counters. + +One scratch [2, 8, ceil4(max_tiles)] and one zeroed counter pair per (device, engine block count) serve every sparse +layer: the layers run one after another on one stream, and the kernel returns each counter to zero. They are +created on an eager decode step; a CUDA-graph capture only reads them. +""" + +from __future__ import annotations + +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +from ..ops.attention.m3_fused_score_select import ( + MAX_DECODE_QUERY_LEN, + MAX_ROWS, + m3_fused_score_select, + ring_fits, + scratch_shape, +) + +_BUFFERS: Dict[Tuple[str, int], Tuple[torch.Tensor, torch.Tensor]] = {} + + +def supported( + idx_q: torch.Tensor, + idx_k_pages: torch.Tensor, + decode_query_len: int, + num_kv_heads: int, + max_tiles: int, +): + """None when the kernel takes this span's geometry, else the reason it does not.""" + if ( + idx_q.dim() != 3 + or idx_q.dtype != torch.float8_e4m3fn + or idx_k_pages.dtype != torch.float8_e4m3fn + ): + return "the index Q / index-K cache is not FP8" + heads = int(idx_q.shape[1]) + rows = decode_query_len * num_kv_heads + if ( + idx_q.shape[2] != 128 + or idx_k_pages.dim() != 3 + or tuple(idx_k_pages.shape[1:]) != (128, 128) + or not 1 <= decode_query_len <= MAX_DECODE_QUERY_LEN + or heads % num_kv_heads + or heads * MAX_DECODE_QUERY_LEN > MAX_ROWS + or rows > MAX_ROWS + ): + return f"{heads} index heads x {decode_query_len} query rows is not the kernel's geometry" + if not ring_fits(rows, max_tiles): + return f"{max_tiles} blocks of select scratch do not fit the stage ring" + return None + + +def buffers(device: torch.device, max_tiles: int) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """(scratch, counters) for this engine's block count; None during a capture that finds them missing.""" + key = (str(device), int(max_tiles)) + found = _BUFFERS.get(key) + if found is None: + if torch.cuda.is_current_stream_capturing(): + return None + found = _BUFFERS[key] = ( + torch.empty(scratch_shape(max_tiles), dtype=torch.float32, device=device), + torch.zeros(scratch_shape(max_tiles)[0], dtype=torch.int32, device=device), + ) + return found + + +def fused_score_select( + idx_q, + idx_k_pages, + block_table, + seq_lens, + n_valid, + *, + decode_query_len, + num_kv_heads, + init_blocks, + local_blocks, + max_tiles, + scratch, + counters, +): + """The [total_q, num_kv_heads, 16] table (a view of a head-major [num_kv_heads, total_q, 16] backing).""" + total_q = int(idx_q.shape[0]) + out = torch.empty((num_kv_heads, total_q, 16), dtype=torch.int32, device=idx_q.device).permute( + 1, 0, 2 + ) + return m3_fused_score_select( + idx_q, + idx_k_pages, + block_table, + seq_lens, + n_valid, + out, + scratch, + counters, + int(decode_query_len), + int(num_kv_heads), + int(init_blocks), + int(local_blocks), + int(max_tiles), + get_env_enable_pdl(), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_index_score.py new file mode 100644 index 000000000000..fad75d44694a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_index_score.py @@ -0,0 +1,638 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py +"""MiniMax-M3 index decode block scorer (CuTe DSL, Blackwell) with its pre-wait schedule. + +The kernel of ``tensorrt_llm/_torch/attention/kernels/blackwell/minimax_m3_index_decode_score.py`` (the same scores, bit +for bit) with two schedule options this target compiles in: + +* ``prefetch_old``: before the grid-dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that hold only + tokens older than this step (the producer ahead writes only this step's tokens), up to ``prefetch_depth`` blocks per + CTA (0 = all of them). +* ``early_issue``: the TMA warp reads a ``seq_lens`` hint and its first page id before the wait and issues the Q tile + and the first K tile right after it. The tiles are read after the wait, so their contents are final; the block table + is only written between graph replays, so the early page id is the right one whenever the block is valid. A CTA + whose early tiles lie past the request's true length drains them before it exits. + +Every kernel is compiled before the executor's warmup (``layers.prebuild_index_scorer``) through ``IndexScoreRunner`` +and called through this target's ``m3_index_score`` entry. +""" + +import cutlass +import torch +from cuda.bindings.driver import CUstream +from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import ( + EVICT_FIRST, + fp8x4_to_fp16x4, + mma_sync, + simple_tma_copy, +) +from tensorrt_llm._torch.kernels.blackwell.utils import TRTLLM_ENABLE_PDL, griddepcontrol_wait + +__all__ = ["IndexDecodeScoreKernel", "IndexScoreRunner"] + + +@dsl_user_op +def _l2_prefetch_bulk(addr: Int64, size: Int32, *, loc=None, ip=None) -> None: + """Prefetch ``size`` contiguous bytes (multiple of 16) at global ``addr`` into L2.""" + llvm.inline_asm( + res=None, + operands_=[Int64(addr).ir_value(loc=loc, ip=ip), Int32(size).ir_value(loc=loc, ip=ip)], + asm_string="cp.async.bulk.prefetch.L2.global [$0], $1;", + constraints="l,r", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _fp8_to_f16_mma_fragments(src: cute.Tensor): + """Split an FP8 ldmatrix fragment into the two FP16 k-fragments MMA wants.""" + src_elems = cute.size(src) + src_u32 = cute.recast_tensor(src, Uint32) + src_f16 = cute.make_rmem_tensor(src_elems, Float16) + src_f16_u32 = cute.recast_tensor(src_f16, Uint32) + # Packed conversion; faster and fewer SASS instructions than + # src.load().to(Float16). + for i in cutlass.range_constexpr(src_elems // 4): + converted = fp8x4_to_fp16x4(src_u32[i]) + src_f16_u32[i * 2] = converted[0] + src_f16_u32[i * 2 + 1] = converted[1] + lower = cute.make_rmem_tensor(src_elems // 2, Float16) + upper = cute.make_rmem_tensor(src_elems // 2, Float16) + + # FP8 ldmatrix yields four consecutive values along K. Split each group + # into the lower two and upper two for the two FP16 MMA k-fragments. + for i in cutlass.range_constexpr(src_elems // 2): + lower[i] = src_f16[(i // 2) * 4 + i % 2] + upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] + return lower, upper + + +class IndexDecodeScoreKernel: + """Per-block max index score for one decode step. + + Grid is (batch, split_k): CTA (b, s) walks KV blocks + s, s + split_k, s + 2 * split_k, ... of request b and writes each block's + score directly, so no cross-CTA reduction is needed. CTAs whose split_id + exceeds the request's block count exit immediately. + """ + + BLOCK_K = 128 + BAR_MMA = 1 + num_stages = 2 + + def __init__( + self, + dtype: type[cutlass.Numeric], + num_heads: int, + max_decode_query_len: int, + split_k: int, + head_dim: int = 128, + prefetch_old: bool = False, + early_issue: bool = False, + prefetch_depth: int = 0, + ): + # prefetch_old: before the grid dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that + # hold only tokens older than this step (the predecessor writes only this step's tokens). + self.prefetch_old = prefetch_old + # early_issue: the TMA warp reads a seq_lens hint and this CTA's first page id before the grid dependency + # wait, then issues the Q tile and the first K tile right after it, overlapping the post-wait seq_lens load. + # The tiles are read after the wait, so their contents are final; the block table is only written between + # graph replays, so the early page id is the right one whenever the block is valid. A CTA whose early tiles + # turn out to lie past the request's true length drains them before exiting. + self.early_issue = early_issue + # prefetch_depth: blocks per CTA that prefetch_old covers (0 = all of them). Prefetching every older page at + # large batches floods L2; the first few per CTA are the ones whose latency the pipeline cannot hide. + self.prefetch_depth = prefetch_depth + self.dtype = dtype + self.num_heads = num_heads + self.max_decode_query_len = max_decode_query_len + self.split_k = split_k + self.head_dim = head_dim + + def _swizzle_elems(self) -> int: + """Elements spanned by one 128-byte swizzle atom.""" + return 128 * 8 // self.dtype.width + + def _sq_layout(self): + """Composed SMEM layout for the Q tile, shared by the descriptor and the kernel.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_q = self.num_heads * self.max_decode_query_len + layout = cute.make_layout( + (self.max_decode_query_len, self.num_heads, (elems, head_dim // elems)), + stride=(elems, self.max_decode_query_len * elems, (1, block_q * elems)), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + def _sk_layout(self): + """Composed SMEM layout for the pipelined K tiles.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_k = self.BLOCK_K + layout = cute.make_layout( + (1, block_k, (elems, head_dim // elems), self.num_stages), + stride=(0, elems, (1, block_k * elems), block_k * head_dim), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + @cute.jit + def __call__( + self, + gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] + gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] + block_table: cute.Tensor, # [bs, max_pages] + score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] + seq_lens: cute.Tensor, # [bs] + stream: CUstream, + ): + num_heads = self.num_heads + head_dim = self.head_dim + MAX_DQL = self.max_decode_query_len + assert num_heads * MAX_DQL <= 32 + + batch = seq_lens.shape[0] + decode_query_len = gQ.shape[0] // batch + grid = (batch, self.split_k, 1) + block = (32 * 5, 1, 1) + + tma_g2s = cpasync.CopyBulkTensorTileG2SOp() + elems = self._swizzle_elems() + + q_tma_atom, q_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gQ, (None, None, elems)), + self._sq_layout(), + (MAX_DQL, num_heads, head_dim), + ) + k_tma_atom, k_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gK_cache, (None, None, elems)), + self._sk_layout(), + (1, self.BLOCK_K, head_dim), + ) + + self.kernel( + q_tma_atom, + q_tma_tensor, + k_tma_atom, + k_tma_tensor, + gK_cache, + block_table, + score, + seq_lens, + decode_query_len, + ).launch(grid=grid, block=block, stream=stream, use_pdl=TRTLLM_ENABLE_PDL) + + @cute.kernel + def kernel( + self, + q_tma_atom: cute.CopyAtom, + q_tma_tensor: cute.Tensor, + k_tma_atom: cute.CopyAtom, + k_tma_tensor: cute.Tensor, + gK_cache: cute.Tensor, + block_table: cute.Tensor, + score: cute.Tensor, + seq_lens: cute.Tensor, + decode_query_len, + ): + batch_id, split_id, _ = cute.arch.block_idx() + _, split_k, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = cute.arch.lane_idx() + + NUM_HEADS = self.num_heads + MAX_DQL = self.max_decode_query_len + BLOCK_Q = NUM_HEADS * MAX_DQL + BLOCK_K = self.BLOCK_K + head_dim = self.head_dim + dtype = self.dtype + MMA_N = 8 + num_stages = self.num_stages + Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) + EPI_Q = Q_TILES * MMA_N + + sq_layout = self._sq_layout() + sk_layout = self._sk_layout() + + smem = SmemAllocator() + sK = smem.allocate_tensor( + dtype, + sk_layout.outer, + byte_alignment=128, + swizzle=sk_layout.inner, + )[0, None, None, None] + # sQ aliases the first K stage: Q is consumed into registers before the + # first K tile is needed, so the two never overlap in time. + sQ_tma = cute.make_tensor(sK[None, None, 0].iterator, layout=sq_layout.outer) + # TMA sees Q as (query, head, dim) while ldmatrix consumes a flattened Q + # column mode. The target profile keeps the rank-2 view even for + # degenerate shapes such as MAX_DQL == 1. + q_tma_elems = self._swizzle_elems() + sQ = cute.coalesce( + cute.group_modes(sQ_tma, 0, 2), + target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), + ) + epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + + # Barrier init and descriptor prefetch do not depend on the predecessor + # grid, so they are the work PDL can overlap. Every CTA runs them, even + # the ones that own no blocks, since ownership is not known until + # seq_lens is read below. + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, 1) + cute.arch.mbarrier_init(tma_empty_mbar + i, 128) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(q_tma_atom) + cpasync.prefetch_descriptor(k_tma_atom) + cute.arch.sync_threads() + + if cutlass.const_expr(self.prefetch_old): + if warp_id == 4: + with cute.arch.elect_one(): + seqlen0 = seq_lens[batch_id] + num_blocks0 = cute.ceil_div(seqlen0, BLOCK_K) + stable = (seqlen0 - decode_query_len) // BLOCK_K + if stable > num_blocks0: + stable = num_blocks0 + if cutlass.const_expr(self.prefetch_depth > 0): + if stable > split_id + self.prefetch_depth * split_k: + stable = split_id + self.prefetch_depth * split_k + for block_id in range(split_id, stable, split_k): + page = block_table[batch_id, block_id] + _l2_prefetch_bulk( + gK_cache[page, None, None].iterator.toint().to(Int64), + Int32(BLOCK_K * head_dim * (dtype.width // 8)), + ) + + # No explicit launch_dependents: the only consumer is the block selector, + # which reduces over every block this grid writes, so triggering it early + # would leave it resident and spinning in its own wait. The implicit + # trigger as the CTAs exit still gives it the launch overlap. + early_q = False + early_page = Int32(0) + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + hint_blocks = cute.ceil_div(seq_lens[batch_id], BLOCK_K) + if hint_blocks > block_table.shape[1]: + hint_blocks = block_table.shape[1] + if split_id < hint_blocks: + early_q = True + early_page = block_table[batch_id, split_id] + + griddepcontrol_wait() + + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + gQ_early = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar, BLOCK_Q * head_dim * (dtype.width // 8) + ) + simple_tma_copy(q_tma_atom, gQ_early, sQ_tma, tma_full_mbar) + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, BLOCK_K * head_dim * (dtype.width // 8) + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + # seq_lens must be read after the wait. It is a graph-stable buffer that + # on_update_kv_lens patches every step, and num_blocks bounds an unmasked + # block_table load and an unmasked score store, so a stale length is an + # out-of-bounds write rather than just a wrong answer. + seqlen = seq_lens[batch_id] + num_blocks = cute.ceil_div(seqlen, BLOCK_K) + + if split_id < num_blocks: + if warp_id == 4: + # TMA warp + tma_stage = 0 + tma_parity = 1 + first_block = split_id + + if not early_q: + gQ_tile = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) + with cute.arch.elect_one(): + Q_size = BLOCK_Q * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) + # TMA bounds-checks rows when the runtime decode_query_len is + # below MAX_DQL; padded Q columns are masked before the stores. + simple_tma_copy(q_tma_atom, gQ_tile, sQ_tma, tma_full_mbar) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + if early_q: + # The first K tile was issued right after the wait. + first_block = split_id + split_k + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + for block_id in range(first_block, num_blocks, split_k): + page_id = block_table[batch_id, block_id] + gK_tile = k_tma_tensor[page_id, None, None] + k_mbar = tma_full_mbar + tma_stage + + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) + with cute.arch.elect_one(): + K_size = BLOCK_K * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) + simple_tma_copy( + k_tma_atom, + gK_tile, + sK[None, None, tma_stage], + k_mbar, + cache_policy=EVICT_FIRST, + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + else: + # MMA warps; each handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T + sK_warp = cute.local_tile(sK, (32, head_dim, num_stages), (warp_id, 0, 0)) + q_start = seqlen - decode_query_len + + elems = 128 // dtype.width # 16B + MMA_K = 32 * 8 // dtype.width # 32B + + # Pre-compute ldmatrix addresses. + # sK loads a [16 x 16B] tile: + # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) + # sQ loads an [8 x 32B] tile: + # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = cute.zipped_divide(sK_warp, (16, cute.make_layout((elems, 2)), 1)) + sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) + + # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) + # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] + sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] + + ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) + ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) + + rQ = cute.make_rmem_tensor( + ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype + ) + rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) + rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + for q in cutlass.range_constexpr(Q_TILES): + cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) + cute.arch.mbarrier_arrive(tma_empty_mbar) + + tma_stage = 1 % self.num_stages + tma_parity = 0 + if tma_stage == 0: + tma_parity ^= 1 + + if cutlass.const_expr(dtype is Float8E4M3FN): + rQ_f16 = cute.make_rmem_tensor((4, head_dim // MMA_K, Q_TILES, 2), Float16) + q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) + rQ_f16[None, None, None, 0].store(q_lower.load()) + rQ_f16[None, None, None, 1].store(q_upper.load()) + + for block_id in range(split_id, num_blocks, split_k): + rC.fill(0.0) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + for k in cutlass.range_constexpr(head_dim // MMA_K): + cute.copy( + ldsm_atom, + sK_ldsm[None, (None, k, tma_stage)], + rK[None, None, k], + ) + for m in cutlass.range_constexpr(2): + if cutlass.const_expr(dtype is Float8E4M3FN): + rK_lower, rK_upper = _fp8_to_f16_mma_fragments(rK[None, m, k]) + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK_lower, + rQ_f16[None, k, n, 0], + rC[None, m, n], + ) + rC[None, m, n] = mma_sync( + rK_upper, + rQ_f16[None, k, n, 1], + rC[None, m, n], + ) + else: + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK[None, m, k], + rQ[(None, k % 2), k // 2, n], + rC[None, m, n], + ) + + cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) + + k_start = block_id * BLOCK_K + warp_id * 32 + + for q in cutlass.range_constexpr(Q_TILES): + for i in cutlass.range_constexpr(4): + for j in cutlass.range_constexpr(2): + col = q * 8 + (lane_id % 4) * 2 + j + q_local_pos = col % MAX_DQL + q_pos = q_start + q_local_pos + k_pos = k_start + i * 8 + lane_id // 4 + rC[q * 8 + i * 2 + j] = ( + rC[q * 8 + i * 2 + j] if q_pos >= k_pos else float("-inf") + ) + + for q in cutlass.range_constexpr(Q_TILES): + # Thread-local reduction along the BLOCK_K dim. + rScore = cute.make_rmem_tensor(2, Float32) + rScore.fill(float("-inf")) + for i in cutlass.range_constexpr(4): + rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) + rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) + + # Warp reduction among lanes 0, 4, 8, 12, ... + for i in cutlass.range_constexpr(3): + offset = 4 << i + other0 = cute.arch.shuffle_sync_bfly( + rScore[0], offset=offset, mask=-1, mask_and_clamp=31 + ) + other1 = cute.arch.shuffle_sync_bfly( + rScore[1], offset=offset, mask=-1, mask_and_clamp=31 + ) + rScore[0] = cute.arch.fmax(rScore[0], other0) + rScore[1] = cute.arch.fmax(rScore[1], other1) + + if lane_id * 2 < MMA_N: + epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] + epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + head_id = lane_id // MAX_DQL + q_local_pos = lane_id - head_id * MAX_DQL + valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len + if lane_id < BLOCK_Q and valid_q: + final_score = epi_buffer[lane_id, 0] + for i in cutlass.range_constexpr(1, 4): + final_score = cute.arch.fmax(final_score, epi_buffer[lane_id, i]) + + t = batch_id * decode_query_len + q_local_pos + score[head_id, t, block_id] = final_score + + tma_stage = (tma_stage + 1) % self.num_stages + if tma_stage == 0: + tma_parity ^= 1 + else: + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + # Early tiles past the request's true length: let them land before the CTA exits. + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.mbarrier_wait(tma_full_mbar + 1 % num_stages, 0) + + +class IndexScoreRunner: + """Compiled ``IndexDecodeScoreKernel`` variants, keyed by (index dtype, index heads, decode query length, head dim, + split_k, early issue). Batch, query-token count, page count and block-table width stay symbolic, so one compile + serves every decode step of its key. ``prefetch_old`` is always on; ``early`` also caps the prefetch at one block + per CTA.""" + + kernel_cache: dict = {} + SUPPORTED_HEAD_DIM = 128 # capacity: fixed: the checkpoint's index head dim + SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K + # BLOCK_Q must fit one warp's worth of epilogue lanes. + MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile + EARLY_PREFETCH_DEPTH = 1 + + @classmethod + def is_supported( + cls, + *, + q_dtype: torch.dtype, + num_heads: int, + head_dim: int, + page_size: int, + max_decode_query_len: int, + ) -> bool: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + from tensorrt_llm._utils import is_sm_100f + + return ( + is_sm_100f() + and q_dtype in TORCH_TO_CUTE_DTYPE + and head_dim == cls.SUPPORTED_HEAD_DIM + and page_size == cls.SUPPORTED_PAGE_SIZE + and num_heads * max_decode_query_len <= cls.MAX_BLOCK_Q + and max_decode_query_len >= 1 + ) + + @classmethod + def _compile( + cls, + q_dtype: torch.dtype, + num_heads: int, + max_decode_query_len: int, + head_dim: int, + split_k: int, + early: bool, + ) -> None: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + + key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + if key in cls.kernel_cache: + return + cute_dtype = TORCH_TO_CUTE_DTYPE[q_dtype] + page_size = cls.SUPPORTED_PAGE_SIZE + sym_total_tokens = cute.sym_int() + sym_batch = cute.sym_int() + + # 16-element divisibility on every non-innermost stride is what the TMA descriptors for Q and K assume. + def _sym_stride(): + return cute.sym_int64(divisibility=16) + + q_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (sym_total_tokens, num_heads, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + # The index-K pool may be coalesced with the main K/V cache (per-page stride above page_size * head_dim). + k_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (cute.sym_int(), page_size, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + bt_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch, cute.sym_int()), stride_order=(1, 0) + ) + # The target passes a transposed view of the [heads, blocks, tokens] selector buffer. + score_fake = cute.runtime.make_fake_tensor( + cutlass.Float32, + (num_heads, sym_total_tokens, cute.sym_int()), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + sl_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch,), stride_order=(0,) + ) + fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = IndexDecodeScoreKernel( + cute_dtype, + num_heads, + max_decode_query_len, + split_k, + head_dim, + prefetch_old=True, + early_issue=early, + prefetch_depth=cls.EARLY_PREFETCH_DEPTH if early else 0, + ) + cls.kernel_cache[key] = cute.compile( + kernel, + q_fake, + k_fake, + bt_fake, + score_fake, + sl_fake, + fake_stream, + options="--enable-tvm-ffi", + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_mc_allreduce.py new file mode 100644 index 000000000000..84a9d1e7b0dc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_mc_allreduce.py @@ -0,0 +1,365 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-latency TP all-reduce + residual + RMSNorm for decode (M <= 8 tokens) over a multicast mapping. + + partial = bf16(x + x2) (optional second input, e.g. shared + routed experts) + updated = bf16(residual + sum_r partial_r) (ranks summed in FP32, rank order) + normed = bf16(updated * rsqrt(mean(updated^2) + eps) * w) + +One cluster of CLUSTER CTAs per token; CTA c owns columns [c * H / CLUSTER, ...), one 16-byte chunk (8 BF16) +per thread. Each CTA pushes its slice of this rank's partial through the multicast mapping (one store lands in +every rank's buffer), releases dependents, polls its own buffer until all ranks' slices arrived, sums them, adds +the residual and reduces the RMS statistic across the cluster through DSMEM. + +Exchange protocol as in the K3 sandwich: int32 words, 0x80000000 = empty (pushed -0.0 halves become +0.0), each +CTA keeps a call counter whose parity selects the buffer half, readers re-arm the words they read. All ranks of +the TP group must issue the same calls in the same stream order. + +The workspace holds two exchange sets. With prepushed 3 two producers pushed this rank's x (set A) and x2 (set B), +e.g. the routed experts' finalize and the shared expert's down projection, and the kernel forms bf16(x + x2) per +rank from the two slots, exactly as the unpushed path does. + +Producers that push (the o_proj / shared-expert GEMV ``ar`` tuple, the routed-expert engine's ``arp`` descriptor) are +pointed at a bucket's exchange by its ``mc`` address, ``flags_for(8)`` counters, ``rank``, ``world_size`` and +``max_tokens``; ``get_workspace`` and ``_bucket`` select it. + +The kernel is ``torch.ops.trtllm.minimax_m3_mc_allreduce_norm``, reached through the ``ops/comm`` entries +``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``; this module owns the workspaces, the bucket and +epilogue-option choice and the MXFP8 publication. +""" + +from dataclasses import dataclass +from typing import Any + +import torch + +from ..ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm as _allreduce_norm_op +from ..ops.comm.m3_mc_finalize_allreduce_norm import ( + m3_mc_finalize_allreduce_norm as _finalize_allreduce_norm_op, +) + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +BUCKETS = (8, 16, 48) # token capacity of each workspace (grid = bucket * cluster CTAs) +EMPTY_WORD = -2147483648 + +# Attention / MoE boundary epilogue options (each bit-identical): 1 residual + norm weight loaded at entry (before +# the poll); 2 cluster RMS exchange by st.async + mbarrier complete_tx instead of cluster.sync; 4 call parity rotated +# right after the poll (no trailing __syncthreads); 8 max-smem carveout so the next kernel's CTAs co-reside early. +# The 8-token bucket, by token count 1..8 (GB300 tuning, the same at TP4 and TP8). +AR_OPTS_BY_M = (7, 7, 15, 15, 15, 15, 7, 7) +# The 16- and 48-token buckets. +AR_OPTS_BY_BUCKET = {16: 15, 48: 15} +# The MoE-finalize all-reduce, per bucket. +AR_FIN_OPTS_BY_BUCKET = {8: 7, 16: 15, 48: 15} + + +# Option 2 is also the safe ordering of the cluster RMS exchange: its stores into peer CTAs' shared memory follow a +# cluster barrier arrive / wait pair. Without it they run before any cluster barrier, when a peer need not have +# started, so every call keeps it. +_OPT_CLUSTER_BARRIER = 2 + + +def _ar_opts(m: int, bucket: int, finalize: bool = False) -> int: + """Epilogue options for an m-row call on the given bucket; every option is bit-identical. A bucket missing from + the tables gets 0, and every result keeps _OPT_CLUSTER_BARRIER.""" + if finalize: + opts = AR_FIN_OPTS_BY_BUCKET.get(bucket, 0) + elif bucket == 8: + opts = AR_OPTS_BY_M[m - 1] + else: + opts = AR_OPTS_BY_BUCKET.get(bucket, 0) + return opts | _OPT_CLUSTER_BARRIER + + +@dataclass(eq=False) +class McAllReduceWorkspace: + """Two halves of [max_tokens][world][H] BF16 (int32 words) per rank behind one multicast mapping, plus one call + counter per CTA per cluster size. Create collectively and eagerly on every TP rank.""" + + uc: torch.Tensor + mc: torch.Tensor + flags: dict + rank: int + world_size: int + hidden: int + handle: Any + comm: Any + max_tokens: int = 8 + + @classmethod + def create(cls, mapping, hidden: int, max_tokens: int = 8) -> "McAllReduceWorkspace": + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_disabled + + comm = _get_mnnvl_workspace_comm(mapping) + words = 4 * max_tokens * mapping.tp_size * (hidden // 2) # two exchange sets of two halves + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + if not mpi_disabled(): + comm.Free() + raise RuntimeError( + "McAllReduceWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 4, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("McAllReduceWorkspace.create failed on another rank") + return cls( + uc=uc, + mc=mc, + flags={}, + rank=mapping.tp_rank, + world_size=mapping.tp_size, + hidden=hidden, + handle=handle, + comm=comm, + max_tokens=max_tokens, + ) + + def flags_for(self, cluster: int): + """This cluster size's per-CTA call counters, or None when they do not exist yet during a capture. + get_workspace creates the cluster-8 counters with the workspace, so for cluster 8 (the only size built) + that None never happens and cannot differ between ranks.""" + flags = self.flags.get(cluster) + if flags is None: + if torch.cuda.is_current_stream_capturing(): + return None + flags = torch.zeros(self.max_tokens * cluster, dtype=torch.int32, device=self.uc.device) + self.flags[cluster] = flags + return flags + + +class McBuckets: + """One workspace per token bucket (8 / 16 / 48): calls of different sizes never share exchange slots.""" + + def __init__(self, workspaces): + self.workspaces = workspaces + self.hidden = workspaces[0].hidden + + def for_tokens(self, m: int): + for ws in self.workspaces: + if m <= ws.max_tokens: + return ws + return None + + +def _bucket(workspace, m): + return workspace.for_tokens(m) if isinstance(workspace, McBuckets) else workspace + + +def decline_reason(x, residual, workspace, cluster=8, x2=None) -> str: + """Why ``m3_mc_allreduce_norm`` returned None for these operands (for the dispatch's miss message).""" + if workspace is None: + return "no multicast workspace: it is created collectively on the first eager decode step" + bucket = _bucket(workspace, x.shape[0]) + if bucket is None: + return f"no workspace bucket holds {x.shape[0]} tokens (buckets {BUCKETS})" + if bucket.flags.get(cluster) is None: + return f"no {cluster}-CTA call counters for the {bucket.max_tokens}-token bucket" + return ( + f"operands outside the kernel: x {tuple(x.shape)} {x.dtype} contiguous={x.is_contiguous()}, " + f"residual contiguous={residual.is_contiguous()}, hidden {bucket.hidden}" + + ("" if x2 is None else f", x2 {tuple(x2.shape)} contiguous={x2.is_contiguous()}") + ) + + +def m3_mc_allreduce_norm( + x, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + x2=None, + quant=False, + prepushed=0, +): + """Return (normed, updated_residual) or None when unsupported. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM. + + A None makes the caller take another collective, so it must be the same on every TP rank. It is: it depends only + on the shapes, dtypes and layouts of the operands (the same code builds them on every rank) and on whether the + workspace's flags exist, which ``get_workspace`` creates collectively, before any capture, for cluster 8.""" + workspace = _bucket(workspace, x.shape[0]) + if workspace is None: + return None + if ( + x.dim() != 2 + or x.shape[0] > MAX_TOKENS + or x.dtype != torch.bfloat16 + or x.shape[1] != workspace.hidden + or not x.is_contiguous() + or not residual.is_contiguous() + or (x2 is not None and (not x2.is_contiguous() or x2.shape != x.shape)) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + outs = _allreduce_norm_op( + x, + x2, + residual, + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and x.shape[1] % 128 == 0), + int(prepushed), + _ar_opts(x.shape[0], workspace.max_tokens), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +def m3_mc_finalize_allreduce_norm( + fc2, + perm, + scales, + shared, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + quant=False, +): + """MoE finalize (sum_k scale * fc2[perm] + shared) + all-reduce + residual + RMSNorm in one kernel. + Returns (normed, updated_residual) or None when unsupported. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM.""" + m = residual.shape[0] + workspace = _bucket(workspace, m) + if workspace is None: + return None + if ( + residual.dim() != 2 + or m > MAX_TOKENS + or residual.shape[1] != workspace.hidden + or fc2.dtype != torch.bfloat16 + or not fc2.is_contiguous() + or fc2.dim() != 2 + or perm.dtype != torch.int32 + or perm.dim() != 2 + or perm.shape[0] != m + or not perm.is_contiguous() + or scales.dtype not in (torch.float32, torch.bfloat16) + or not scales.is_contiguous() + or (shared is not None and (shared.shape != residual.shape or not shared.is_contiguous())) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + outs = _finalize_allreduce_norm_op( + fc2, + perm, + scales, + shared, + residual.contiguous(), + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and residual.shape[1] % 128 == 0), + _ar_opts(m, workspace.max_tokens, finalize=True), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +_WORKSPACE: dict = {} +_AGREE_COMMS: dict = {} + + +def _tp_comm(mapping): + """A communicator over mapping's TP group, created once (collective over that group only).""" + key = tuple(mapping.tp_group) + comm = _AGREE_COMMS.get(key) + if comm is None: + from tensorrt_llm._utils import mpi_comm, mpi_disabled + + if mpi_disabled(): + comm = mapping.tp_group_pg + else: + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + _AGREE_COMMS[key] = comm + return comm + + +def tp_agree(mapping, ok: bool) -> bool: + """Whether ``ok`` holds on every rank of mapping's TP group (one blocking host collective). + + For a fast path whose availability is decided from rank-local state (a kernel build that can fail on one rank): + every rank must take the same path, so that they issue the same collectives with the same exchange protocol and + the same kernels. Call it once per decision, at build time, at the same point on every rank, outside CUDA-graph + capture; never per step. With one rank, or attention DP (whose ranks run different batches), it returns ``ok``. + """ + if mapping.tp_size <= 1 or mapping.enable_attention_dp: + return ok + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("tp_agree must run outside CUDA-graph capture") + from tensorrt_llm._torch.distributed.ops import _mnnvl_workspace_all_succeeded + + return _mnnvl_workspace_all_succeeded(_tp_comm(mapping), bool(ok)) + + +def tp_agree_built(mapping, built, what: str): + """``built`` when it is not None on every TP rank, else None on every rank (see ``tp_agree``).""" + if tp_agree(mapping, built is not None) or built is None: + return built + from tensorrt_llm.logger import logger + + logger.warning(f"{what}: built here but not on another TP rank; every rank falls back") + return None + + +def get_workspace(mapping, hidden: int): + """Process-wide workspace (created collectively on first eager use); None while capturing before creation. + + The first eager call for a hidden size (the warmup forward, before any CUDA-graph capture) creates one + McAllReduceWorkspace per bucket in BUCKETS on every TP rank (MPI collective over the TP group) plus each + bucket's 8-CTA-cluster call counters; later calls, captured or not, return the same McBuckets. The workspaces + live until process exit. + """ + ws = _WORKSPACE.get(hidden) + if ws is None: + if torch.cuda.is_current_stream_capturing(): + return None + wss = [McAllReduceWorkspace.create(mapping, hidden, b) for b in BUCKETS] + for w in wss: + w.flags_for(8) + ws = McBuckets(wss) + _WORKSPACE[hidden] = ws + return ws diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_moe/__init__.py new file mode 100644 index 000000000000..e361dc614231 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_moe/__init__.py @@ -0,0 +1,16 @@ +# 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. +"""MiniMax-M3 MoE layer of the TP8 target for low-batch decode as a CuTe DSL kernel: ``m3_moe`` for 1..8 tokens +(op.py) at the TP8 routed and shared intermediate of 384 per rank.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_moe/m3_moe_kernel.py new file mode 100644 index 000000000000..34592d700b99 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_moe/m3_moe_kernel.py @@ -0,0 +1,3820 @@ +# 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. +"""MiniMax-M3 routed experts for up to 8 decode tokens on one TP rank: ``m3_moe`` (group-phased weight stream). + +One persistent kernel replaces TRTLLM-Gen's MiniMax2 routing, FC1 (NVFP4 x NVFP4 + SwiGLU-OAI + NVFP4 requantization), +FC2 and finalize, optionally with the router GEMV and the activation's NVFP4 quantization folded in (FUSED_ROUTER). It +writes the rank's routed partial [M, 6144] bf16, bit for bit what the TRTLLM-Gen op returns. + +- Routing (warps 4..7, token t on warp 4 + t % 4): TRTLLM-Gen's MiniMax2 top-4, operation for operation; warp 4 then + makes the distinct experts slots in (token, top-k) order (match.any), slot s in group s / 4. +- Work: the groups' FC1 stages and FC2 tiles stream through one ring per CTA in phases: phase p = FC1 of group p (a + stream-K range of the group's unit-major stages), then FC2 of group p - LAG (this CTA's 32-row output blocks, 4 + slots' slices stacked in one 128-row MMA, B row j M + t). LAG phases of streaming hide the hand-off of group + p - LAG's intermediate between CTAs. +- FC1 stage: a 32-row unit's rows at 4 consecutive 256-K tiles (rows 32 q + r) against resident B rows QS q + t. + FC2 stage: 4 slots' 32 rows at one 256-K tile. Both carry 8 KB of raw scale atoms, spliced into the 4 instruction + atoms by warp 6 with one formula (word [r][q] <- raw atom 4 q + kb, word [r][m1]). +- FC1 epilogue: k-quarters summed in order, SwiGLU-OAI, NVFP4 requant into tagged 16-byte records (3 words of 3 data + bytes + the call's tag); a unit split between two CTAs is summed by its tail's holder (Lamport slot). +- Warps 5 and 8.. (XW warps) build each group's FC2 B (tiles + scale atoms) from the records once they carry the + call's tag; warp 5 copies the scales to TMEM. +- Finalize: bf16(fc2_alpha * acc) per (slot, token) into its top-k position, then per token the fma chain over its + top-4 in order, bf16 out. With an all-reduce descriptor, every output word is also pushed to the exchange + (AR_PRE: call parities and descriptor loaded during the last group's FC2). + +Configuration (trace-time constants) via M3_CONFIG = {"m_max", "num_ctas", "stages", "nb2_cap", "fused_router", "lag", +"xw", ...}; the plan (op.make_phase_plan) must be built for the same lag. +Launched with PDL: barrier setup, TMEM allocation, tensormap prefetch, bias / plan / router weight loads before +griddepcontrol.wait; every producer-output read and every global write after it. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +_CFG = globals().get("M3_CONFIG") or {} + + +def _cfg(key: str, default): + return type(default)(_CFG.get(key, default)) + + +# ============================================================================= +# Problem shape (trace-time constants). +# ============================================================================= +H = 6144 # capacity: fixed: the checkpoint's hidden size +# routed intermediate per rank: 768 at TP4, 384 at TP8 (FC2 K = 1.5 256-K tiles: the last tile's A box is half out +# of bounds, zero-filled by TMA, and only its first KB2_LAST K-64 MMAs / scale atoms are issued; no padding, no extra +# bytes) +I_TP = _cfg("i_tp", 768) +TWO_I = 2 * I_TP +assert I_TP % 128 == 0 and 128 <= I_TP <= 768, I_TP +E_LOCAL = 128 +TOP_K = 4 # capacity: fixed: the checkpoint's top-k +ROUTE_SCALE = 2.0 +SUM_EPS = 1e-20 +M_MAX = _cfg("m_max", 1) +NCTA = _cfg("num_ctas", 152) +FUSED_ROUTER = bool(_cfg("fused_router", 0)) +NROUTE = E_LOCAL if FUSED_ROUTER else 0 # experts whose logits the router GEMV computes +NGEMV = (NROUTE + NCTA - 1) // NCTA # experts per CTA: bx, bx + NCTA, .. +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +XVDUMP = bool( + _cfg("xvdump", 0) +) # debug only (needs a plan without split units): FC1 sums into `partials` +YDUMP = bool( + _cfg("ydump", 0) +) # debug only: FC2 outputs per (token, k) into `partials` (breaks its protocol) +LAG = _cfg("lag", 2) # phase p runs FC1 of group p and FC2 of group p - LAG +G_CAP = TOP_K * M_MAX # distinct experts per call +NG_CAP = M_MAX # groups of 4 slots +XW = _cfg("xw", 5 if M_MAX > 4 else (4 if M_MAX > 2 else 1)) # FC2 B builder warps: 5, then 8 .. +THREADS = 256 + 32 * (XW - 1) +GEMV_THREADS = 256 # capacity: fixed: the CTA size +# SHARED: the MoE's shared expert (MXFP8 gate_up [1536, 6144] -> SwiGLU-OAI -> MXFP8 down [6144, 768]) on NSH more +# CTAs (warps 0-3), its output added in finalize: out = bf16(shared + routed) (shared_output.add_(routed_output)). +# The math is mxfp8_low_m_gemm's (native MXFP8 activation quant, mma.sync m16n8k32 e4m3 per 32-K block, block scales +# as one fma, its partial-sum order, bf16) and the Triton silu_and_mul's, operation for operation. Shared CTA s owns +# intermediate positions 32 s .. 32 s + 31 (gate and up rows: MMA tiles = warps 0-3) and down rows 256 s .. + 255; +# 16-row x 768-byte weight chunks (+ their scale atoms) stream through NBW buffers per warp, the first NBW before +# the grid dependency. h crosses the shared CTAs, y_sh reaches the routed CTAs' finalize (both tagged words). +SHARED = bool(_cfg("shared", 0)) +# LATE_TRIGGER: release the dependents at teardown, not after the grid dependency. An early release lets the next +# kernels' CTAs (the all-reduce, then the attention it releases at entry) become resident during this kernel, packed +# onto the SMs it leaves free, where the attention then runs on a fraction of the GPU. +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +# AR_PRE: the finalize's all-reduce push uses the call parities and descriptor fields loaded while the last group's +# FC2 runs (not two dependent global loads per push on the kernel's tail). +AR_PRE = bool(_cfg("ar_pre", 1)) +# XSTAGE (fused router): the hidden states land in the (still idle) ring by one bulk copy per token right after the +# grid dependency; the NVFP4 quantization of B1 reads them from shared memory (one HBM / L2 round trip, not ~8). +XSTAGE = bool(_cfg("xstage", 1)) and bool(_cfg("fused_router", 0)) +NSH = 24 if SHARED else 0 +# the shared intermediate per rank: 768 (TP4), 384 (TP8). A shared CTA owns SH_ROWS intermediate positions: its +# gate_up MMA tile is SH_SPLIT K slices x (SH_ROWS gate + SH_ROWS up rows) = 128 rows (TP4: 2 x (32 + 32), the order +# of m3_mx_gemv split 2; TP8: 4 x (16 + 16), m3_mx_gemv split 4), and 6144 / NSH = 256 down rows over the full K. +SH_I = _cfg("sh_i", 768) +SH_ROWS = SH_I // 24 +# gate_up K slices (= m3_mx_gemv's cluster ranks, contiguous K ranges) +SH_SPLIT = 128 // (2 * SH_ROWS) +SH_NB = 8 * SH_SPLIT # gate_up B rows (8 q + t: token t's slice-q k-tile) = its MMA N +SH_KB = H // 32 # gate_up K blocks # capacity: fixed: a kernel tile constant +SH_DKB = SH_I // 32 # down K blocks # capacity: fixed: a kernel tile constant +SH_DKT = SH_I // 128 # down k-tiles +# Gate_up k-tiles (128 K) per slice. +SH_GT = SH_KB // 4 // SH_SPLIT # capacity: fixed: a kernel tile constant +SH_STAGES_GU = SH_GT +# + 2 down tiles x SH_DKT k-tiles. +SH_STAGES = SH_GT + 2 * SH_DKT # capacity: fixed: pipeline stages +SHX_BAR = 6 +# SH_PRETILE: the shared weights come pre-tiled (op._pretile_shared): per shared CTA its SH_STAGES stages, each the +# 16 KB shared-memory image of its 128-row MMA tile (128-byte swizzle) followed by the stage's 512-byte instruction +# scale atom, so a stage is two contiguous bulk copies (not 128 row segments of 128 bytes) and needs no splicing. +SH_PRETILE = bool( + _cfg("sh_pretile", 1 if M_MAX <= 2 or SH_SPLIT != 2 else 0) +) # (in-chain: a win at M <= 2, not 4..8) +if SHARED: + assert SH_I in (384, 768), SH_I + assert SH_PRETILE or SH_SPLIT == 2, ( + "the un-pretiled shared stream is the 32-row (TP4) layout only" + ) +SH_STAGE_BYTES = 16384 + 512 +MMA_M = 128 +KT_BYTES = 128 +INC_K = 2 +NUM_KB = 4 +# FC1 +N1 = 16 if M_MAX <= 4 else 32 # B rows: QS q + t (token t at the stage's k-tile 4 j + q) +QS = N1 // 4 +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 # 48 units per slot +K1_TILES = H // 256 +S1 = K1_TILES // C1 # 6 stages per unit # capacity: fixed: pipeline stages +NUM_BYTES_A = MMA_M * KT_BYTES # capacity: fixed: a stage's operand bytes +NUM_BYTES_BOX1 = R1 * KT_BYTES +SF1_KA = H // 64 +SF1_RB = TWO_I // 128 +NUM_BYTES_SFA_RAW = 16 * 512 +NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 +NUM_BYTES_SFA = NUM_KB * 512 # capacity: fixed: a stage's scale bytes +NUM_WORDS_SFA = NUM_BYTES_SFA // 4 +NUM_BYTES_B1 = N1 * KT_BYTES +NUM_SFB1_ATOMS = S1 * NUM_KB +# FC2 +R2 = 32 +NB2 = H // R2 # 192 +KT2 = (I_TP + 255) // 256 # 3 (TP4), 2 (TP8) +# K-64 instructions of the last FC2 k-tile: 4 (TP4), 2 (TP8) +KB2_LAST = (I_TP - 256 * (KT2 - 1)) // 64 + + +def _nkb2(t2: int) -> int: + """K-64 MMA instructions (= scale atoms) of FC2 k-tile t2.""" + return NUM_KB if t2 < KT2 - 1 else KB2_LAST + + +SF2_KA = I_TP // 64 +SF2_RB = H // 128 +N2 = max(8, (4 * M_MAX + 7) // 8 * 8) # B rows j M + t +NUM_BYTES_B2 = N2 * KT_BYTES +NB2_CAP = _cfg("nb2_cap", 2) +UNITS_ROW = TWO_I // R1 # 48 records per (slot, token) row +H_ROW = UNITS_ROW * 16 +# HB_REPL: the tagged FC1 -> FC2 records written to HB_REPL copies (HB_STRIDE int32 words apart); CTA bx polls copy +# bx % HB_REPL. Every routed CTA polls every record, so with one copy the hand-off latency depends on where the few +# hot lines land in L2 (TP8 M=1: 1.5 KB of records, +-1 us across allocations); copies spread the polls. +HB_REPL = _cfg("hb_repl", 1) +assert 1 <= HB_REPL <= 32, HB_REPL # (copy i is written by lane i) +HB_STRIDE = (G_CAP * M_MAX * UNITS_ROW * 4 + 255) // 256 * 256 + 32 +PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +B2_BAR = 5 +TD_STRIDE = 64 # tdbg columns per CTA +TOPK_BAR = 4 +B1_AFTER_TOPK = bool( + _cfg("b1_after_topk", 0) +) # (experiment) the B1 build waits for the routing warps' top-k +TOPK2_BAR = 7 +EARLY0 = bool( + _cfg("early0", 1) +) # warp 4 streams group 0 (token 0's experts) right after its own top-k +EARLY0_N = _cfg("early0_n", 99) # stages issued early (capped at STAGES) +# routing warps: 4..7, then the FC2 B builder warps 8.. (idle until the routing is done) up to one warp per token +RW_EXTRA = max(0, min(M_MAX - 4, (THREADS // 32) - 8)) if bool(_cfg("rw_extra", 1)) else 0 +NRW0 = ( + min(4, M_MAX) if bool(_cfg("rw_trim", 1)) else 4 +) # warps 4.. with a token (no idle pollers in the barrier) +NRW = NRW0 + RW_EXTRA +LG_REPL = _cfg("lg_repl", 32 if M_MAX > 4 else 128) +GEMV_SMEM = bool( + _cfg("gemv_smem", 1) +) # (XSTAGE) the router GEMV reads the staged hidden states from smem +LG_STRIDE = M_MAX * E_LOCAL + 32 +POLL_NS = _cfg("poll_ns", 150) +B2_POLL_NS = _cfg("b2_poll_ns", 200) +# B2IMG (needs SFA_CP: warp 6 is free): the FC1 epilogue writes each finished unit's FP4 bytes and scale byte straight +# into a per-group global image laid out as the FC2 B tile + its scale atoms (exactly what the builder warps assembled +# from the tagged records); per group, warp 6 joins the epilogue (fc1_done), fences the async proxy and does one +# release add of this CTA's finished units to the group counter; warp 5 acquires the counter (every group gains +# CNT_PER_CALL per call, CTA 0 pads absent slots), fences, and bulk-copies the image into the B2 buffer. +B2IMG = bool(_cfg("b2img", 0)) and bool(_cfg("sfa_cp", 0)) +IMG_B2 = KT2 * NUM_BYTES_B2 +IMG_SF = KT2 * NUM_KB * 512 +IMG_BYTES = IMG_B2 + IMG_SF +CNT_PER_CALL = 4 * (TWO_I // R1) +B2IMG_NS = _cfg("b2img_ns", 16) # (B2IMG) counter poll back-off +B2_DLY = _cfg( + "b2_dly", 0 +) # (stress) the release warp spins up to this many ns (CTA- and group-dependent) first +W_WARP = 4 # capacity: fixed: a warp role index +X_WARP = 5 # capacity: fixed: a warp role index +S_WARP = 6 # capacity: fixed: a warp role index +MMA_WARP = 7 # capacity: fixed: a warp role index +# (M <= 3) the early stages are issued by the idle MMA warp, not by warp 4 (routing's critical path). +E0_WARP = MMA_WARP # capacity: fixed: a warp role index +E0_BAR = 8 +E1_BAR = 9 +E0_OFFLOAD = EARLY0 and NRW0 <= 3 and bool(_cfg("e0_off", 1)) +EVICT_FIRST = 0x12F0000000000000 +_LOG2E = 1.4426950408889634 +LD_POW2 = ( + 1 if M_MAX == 1 else 2 if M_MAX == 2 else 4 if M_MAX <= 4 else 8 +) # tcgen05.ld .num covering M columns + +_FIXED_SMEM = ( + S1 * NUM_BYTES_B1 + + 2 * KT2 * NUM_BYTES_B2 + + 2 * KT2 * NUM_KB * 512 + + NB2_CAP * M_MAX * TOP_K * R2 * 4 + + 2 * M_MAX * C1 * R1 * 4 + + (M_MAX * (H // 16) if FUSED_ROUTER else 0) + + 6 * 1024 +) +# SFA_CP: the A scale words go by 4-byte cp.async straight into the instruction atom layout (no raw scale box, no +# splicer warp; the MMA warp tcgen05.cp's them), completing on the stage's full barrier: 18 KB stages instead of 26 KB +SFA_CP = bool(_cfg("sfa_cp", 0)) +_RAW_STAGE = 0 if SFA_CP else NUM_BYTES_SFA_RAW +_STAGE_SMEM = NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA +_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // _STAGE_SMEM +ACC1_COL = 0 +ACC2_COL = 2 * N1 +SFB1_COL = ACC2_COL + 2 * N2 +SFB2_COL = SFB1_COL + NUM_SFB1_ATOMS * 4 +SFA1_COL = SFB2_COL + 2 * KT2 * NUM_KB * 4 +STAGES = min(_cfg("stages", 12 if SFA_CP else 8), _STAGES_FIT, (512 - SFA1_COL) // 16) +TMEM_COLS = 512 +assert 1 <= M_MAX <= 8, M_MAX +assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) +assert M_MAX * H * 2 <= STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA), ( + "x staging must fit the ring" +) +# one shared-memory region for the ring and the FC1 / FC2 B operands (a shared CTA reuses all of it) +_OFF_SFAR = STAGES * NUM_BYTES_A +_OFF_SFA = _OFF_SFAR + STAGES * _RAW_STAGE +_OFF_B1 = _OFF_SFA + STAGES * NUM_BYTES_SFA +_OFF_B2 = _OFF_B1 + S1 * NUM_BYTES_B1 +_OFF_SFB2 = _OFF_B2 + 2 * KT2 * NUM_BYTES_B2 +BIG = _OFF_SFB2 + 2 * KT2 * NUM_KB * 512 +# shared CTA layout in BIG: weight buffers, x quarters (2), x scales [M, 192], gate / up [4, 16, M], h, h scales +# shared CTA layout in BIG: A ring (16 KB stages), raw scale atoms (4 per stage), spliced atoms, B (x: [stage][16 rows] +# [128 B]; h: [6][8][128 B]), x / h scale codes, gate_up partial exchange +_SB_A = 0 +_SB_RAW = None +# slices 1.. gate / up partials [slice - 1][u][p][8 tokens] f32 +_SH_XCH_BYTES = (SH_SPLIT - 1) * 2 * SH_ROWS * 8 * 4 +_SH_FIX = ( + SH_GT * SH_NB * 128 + + SH_DKT * 8 * 128 + + M_MAX * SH_KB + + M_MAX * SH_DKB + + _SH_XCH_BYTES + + 2 * SH_ROWS * 8 * 4 + + 1024 +) +SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 +_SB_RAW = SH_NST * NUM_BYTES_A +_SB_SPL = _SB_RAW + SH_NST * 2048 +_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 +_SB_BH = _SB_BX + SH_GT * SH_NB * 128 +_SB_XC = _SB_BH + SH_DKT * 8 * 128 +_SB_HC = _SB_XC + M_MAX * SH_KB +_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 +_SB_GU = _SB_XCH + _SH_XCH_BYTES +_SB_END = _SB_GU + 2 * SH_ROWS * 8 * 4 +# shared CTA TMEM columns: gate_up accumulator (N SH_NB), down accumulators (2 x N 8), A scales per stage, B scales +SH_ACC_GU = 0 +SH_ACC_DN = SH_NB +SH_SFA = SH_ACC_DN + 16 +SH_SFB_GU = SH_SFA + 4 * 8 +SH_SFB_DN = SH_SFB_GU + 4 * SH_GT # capacity: fixed: a kernel tile constant +SH_TMEM = 256 # capacity: fixed: TMEM columns +if SHARED: + assert FUSED_ROUTER, "shared expert: fused router only" + assert SH_STAGES - SH_NST <= 64, SH_STAGES + assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 24 <= SH_TMEM, (SH_NST, _SB_END, BIG) + + +# ============================================================================= +# Inline-PTX helpers. +# ============================================================================= +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- all-reduce push (the MoE boundary's multicast exchange, set A): +# arp = int64 [enable, mc, flags, rank, world, maxt] +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): + """st.global.b32 when pred != 0.""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _ar_push(arp, v, low, t, cl): + """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on + lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call + parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + world = arp.load(idx=4) + parity = cutlass.Int64( + _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) + ) + & cutlass.Int32(1) + ) + off = ( + parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +def _ar_parity(arp, t, cl): + """The call parity _ar_push reads for token t's channel cl (fixed between this kernel's grid dependency and the + next all-reduce, which this kernel releases at teardown): loaded ahead of the finalize (AR_PRE).""" + return _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cl // cutlass.Int32(H // 8)) + ) + + +def _ar_push_pre(ard, parity32, v, low, t, cl): + """_ar_push with the parity (_ar_parity) and the descriptor fields ard = (mc, rank, world, maxt) preloaded.""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + mc, rk, world, maxt = ard + parity = cutlass.Int64(parity32 & cutlass.Int32(1)) + off = ( + parity * maxt * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + rk) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(mc + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" + _llvm.inline_asm( + None, + [], + f"bar.arrive {bar_id}, {threads};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "st.relaxed.gpu.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int64(val)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _nanosleep(ns: int, *, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _gtime(*, loc=None, ip=None): + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): + if cutlass.const_expr(TIMING): + tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _sub_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_bits(x, *, loc=None, ip=None): + """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3_to_f32(b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(b)], + "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" + "cvt.f32.f16 $0, l;\n}", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x2(hi, lo, *, loc=None, ip=None): + """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """x rounded to bf16 (nearest even), as f32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(v, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(v)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16 (nearest even) bits of x in the low 16 bits.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ue8m0_rp(x, *, loc=None, ip=None): + """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3x2(hi, lo, *, loc=None, ip=None): + """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _min_xorsign_abs(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_full(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2_full(x, *, loc=None, ip=None): + """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): + """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" + _llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _mx_osc(amax): + """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" + code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) + cf = cutlass.Int32( + cutlass.select_( + code == cutlass.Int32(255), + cutlass.Int32(0x7FFFFFFF), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) + ), + ) + ).bitcast(cutlass.Float32) + return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) + + +def _mx_words(v, osc): + """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" + words = [] + for t in range(len(v) // 4): + lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) + hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) + words.append(lo | (hi << cutlass.Int32(16))) + return words + + +def _twiddle(score): + """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" + bits = ( + score.bitcast(cutlass.Int32) + if hasattr(score, "bitcast") + else cutlass.Float32(score).bitcast(cutlass.Int32) + ) + return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) + + +@cute.jit +def _wait(bar: cutlass.Array, parity): + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +def _tptr(base_row, col, dtype=cutlass.Int32): + return cutlass.inttoptr((base_row << 16) | col, 6, dtype) + + +@dsl_user_op +def _match_any(value, *, loc=None, ip=None): + """The mask of the warp's lanes holding the same 32-bit value.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(value)], + "match.any.sync.b32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@cute.jit +def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): + """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" + if cutlass.const_expr(HANG_DBG): + n = cutlass.Int32(0) + while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( + n < cutlass.Int32(1 << 22) + ): + n = n + cutlass.Int32(1) + if n >= cutlass.Int32(1 << 22): + bx, _, _ = cute.arch.block_idx() + w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tdbg.store( + cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w + ) + else: + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +@cute.jit +def _adv(stage, phase): + """The next ring position's (stage, phase).""" + stage = stage + cutlass.Int32(1) + wrap = stage == cutlass.Int32(STAGES) + return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( + wrap, phase ^ cutlass.Int32(1), phase + ) + + +@cute.jit +def _seg(i0, i1, k): + """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: + the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" + uf = i0 // cutlass.Int32(S1) + ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) + unit = uf + k - cutlass.Int32(1) + unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) + unit = cutlass.select_(ul == uf, uf, unit) + j0 = i0 - unit * cutlass.Int32(S1) + j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) + j1 = i1 - unit * cutlass.Int32(S1) + j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) + return unit, j0, j1 + + +@cute.jit +def _nsegs(i0, i1): + n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) + return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) + + +@cute.jit +def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): + """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile + 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" + q = lane // cutlass.Int32(QS) + tl = lane - q * cutlass.Int32(QS) + sf_ok = (lane < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) + sfv = [] + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + j = a // NUM_KB + kb = a % NUM_KB + wsf = src.load( + idx=cutlass.select_( + sf_ok, + tl * cutlass.Int32(H // 64) + + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) + + cutlass.Int32(kb), + cutlass.Int32(0), + ) + ) + sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + prims.tcgen05_st( + "32x32b", + _tptr(base_row + warp * cutlass.Int32(32), base_col + SFB1_COL + a * 4), + sfv[a], + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(b1_ready) + + +@cute.jit +def _b1_addr(t, kb16): + """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" + kt = kb16 >> cutlass.Int32(4) + j = kt >> cutlass.Int32(2) + n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t + cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) + return ( + j * cutlass.Int32(NUM_BYTES_B1) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) + ) // cutlass.Int32(4) + + +@cute.jit +def _build_b1( + tidx, + warp, + lane, + base_row, + base_col, + x32: cutlass.Array, + xsf32: cutlass.Array, + sB1: cutlass.Array, + b1_ready: cutlass.Array, +): + """B (the M tokens' whole K) into the resident tiles from the NVFP4 activation, its scales to TMEM.""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + for r in cutlass.range_constexpr((M_MAX * (H // 16) + EPI_THREADS - 1) // EPI_THREADS): + b = tidx + cutlass.Int32(r * EPI_THREADS) + if b < cutlass.Int32(M_MAX * (H // 16)): + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + v2 = x32.load(idx=b * cutlass.Int32(2), vector_size=2, alignment=8) + sB1_32.store( + (cutlass.Int32(v2[0]), cutlass.Int32(v2[1])), idx=_b1_addr(t, kb), alignment=8 + ) + _sfb1_to_tmem(warp, lane, base_row, base_col, xsf32, b1_ready) + + +@cute.jit +def _build_b1_fused( + tidx, + warp, + lane, + base_row, + base_col, + xb32: cutlass.Array, + sfsc: cutlass.Array, + sB1: cutlass.Array, + s_xsf: cutlass.Array, + b1_ready: cutlass.Array, +): + """_build_b1 from the bf16 hidden states, quantized exactly as fp4_quantize(x, sf_scale, 16).""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + s_xsf8 = cutlass.Array( + s_xsf.data_ptr(0), shape=(M_MAX * (H // 16),), dtype=cutlass.Int8, alignment=16 + ) + sc = cutlass.Float32(sfsc.load(idx=0)) + rsc = _rcp(sc) + rsix = _rcp(cutlass.Float32(6.0)) + NBLK = M_MAX * (H // 16) // EPI_THREADS + QB = ( + 3 if NBLK % 3 == 0 else (2 if NBLK % 2 == 0 else 1) + ) # blocks whose loads are in flight together + # a rolled loop over groups of QB blocks (a fully unrolled one is ~50 KB of SASS at M = 8: instruction fetch bound) + for rg in cutlass.range(NBLK // QB, unroll=1): + wall = [] + for rr in cutlass.range_constexpr(QB): + bq = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + for h in cutlass.range_constexpr(2): + wall.append( + xb32.load( + idx=bq * cutlass.Int32(8) + cutlass.Int32(4 * h), + vector_size=4, + alignment=16, + ) + ) + for rr in cutlass.range_constexpr(QB): + b = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + w = [] + for h in cutlass.range_constexpr(2): + v4 = wall[rr * 2 + h] + for z in cutlass.range_constexpr(4): + w.append(cutlass.Int32(v4[z])) + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), + _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), + cutlass.Float32(0.0), + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + sB1_32.store((lo, hi), idx=_b1_addr(t, kb), alignment=8) + s_xsf8.store(cutlass.Int8(sfb), idx=b) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + + +@cute.jit +def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): + """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, + gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 + and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" + neg_max = cutlass.Float32(-3.40282346638528859812e38) + gc = cutlass.select_(gate_raw > lim, lim, gate_raw) + gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) + uc = cutlass.select_(xv > lim, lim, xv) + uc = cutlass.select_(xv < -lim, -lim, uc) + c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) + sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) + gs = _mul_rn(gc, sgm) + h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) + h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) + amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) + for off in (16, 8, 4, 2, 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) + sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) + sff = _e4m3_to_f32(sfb) + osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) + return _mul_rn(_mul_rn(h, osc), c), sfb + + +@cute.jit +def _plan_range(s_plan: cutlass.Array, p): + w = s_plan.load(idx=p) + return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) + + +@cute.jit +def _sh_issue( + lane, + sid, + g, + s_big: cutlass.Array, + sh_full: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, +): + """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both + slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its + raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile + (g - SH_GT) % 6.""" + st = g % cutlass.Int32(SH_NST) + bar = sh_full.subview(st) + if cutlass.const_expr(SH_PRETILE): + if lane == 0: + src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) + prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + gusf.subview(src + cutlass.Int32(16384)), + bar, + 512, + ) + cute.arch.sync_warp() + else: + if g < cutlass.Int32(SH_GT): + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) + cute.arch.sync_warp() + if lane < 4: + q = lane >> cutlass.Int32(1) + u = lane & cutlass.Int32(1) + row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) + kt = q * cutlass.Int32(SH_GT) + g + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), + tma_gu, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + cutlass.Int32(_SB_RAW) + + st * cutlass.Int32(2048) + + lane * cutlass.Int32(512) + ), + gusf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(SH_DKT) + kt = gd - jt * cutlass.Int32(SH_DKT) + row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), + tma_dn, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), + dnsf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + cute.arch.sync_warp() + + +@dsl_user_op +def _red_release_add(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None): + _asm(None, [], "fence.proxy.async.global;", "", True, loc, ip) + + +@cute.jit +def _sfa_to_tmem(sSFA: cutlass.Array, stage, base_row, base_col): + """(SFA_CP, elected MMA thread) the stage's 4 scale atoms smem -> TMEM (what the splicer warp does otherwise).""" + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SFA1_COL + stage * cutlass.Int32(16) + kb * 4), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + + +@cute.jit +def _sf1_cp( + lane, stage, sf_line, c, sSFA: cutlass.Array, s1w: cutlass.Array, ab_full: cutlass.Array +): + """(SFA_CP) FC1 stage scale atoms: word [r][q] of atom kb <- raw atom 4 q + kb (raw line sf_line = 4 x atom) + word [r][c % 4] (what the splicer copies out of the raw box), 4-byte cp.async by the warp's 32 lanes; their + arrivals (noinc) complete the stage's full barrier together with the A boxes' bytes.""" + src0 = ( + (sf_line >> cutlass.Int32(2)) * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (c & cutlass.Int32(3)) + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for q in cutlass.range_constexpr(C1): + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + q)), + s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _sf2_cp( + lane, + stage, + g, + blk, + t2: cutlass.Constexpr, + n_here, + s_el: cutlass.Array, + sSFA: cutlass.Array, + s2w: cutlass.Array, + ab_full: cutlass.Array, +): + """(SFA_CP) FC2 k-tile t2 scale atoms: word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) of slot + jj's expert, word [r][blk % 4].""" + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + src0 = ( + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (blk & cutlass.Int32(3)) + ) + # (a partial last k-tile has no atoms past I_TP) + for kb in cutlass.range_constexpr(_nkb2(t2)): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + jj)), + s2w.subview(src0 + cutlass.Int32(kb * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _shared_cta( + tidx, + warp, + lane, + sid, + s_big: cutlass.Array, + sh_bar: cutlass.Array, + epochs: cutlass.Array, + tdbg: cutlass.Array, + bx, + xb32: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + tmem_ptr_i32: cutlass.Array, + gu8: cutlass.Array, + dn8: cutlass.Array, +): + """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one + 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles + of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" + big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) + bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) + big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) + sh_full = sh_bar # [0, SH_NST): stage full + sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) + b_gu_rdy = sh_bar.subview(16) + b_dn_rdy = sh_bar.subview(17) + acc_gu = sh_bar.subview(18) + acc_dn = sh_bar.subview(19) + if warp == 0: + prims.prefetch_tensormap(tma_gu) + prims.prefetch_tensormap(tma_dn) + if lane < SH_NST: + prims.mbarrier_init(sh_full.subview(lane), 1) + prims.mbarrier_init( + sh_empty.subview(lane), 33 + ) # the MMAs' commit + each splicer lane's reads done + if lane == 0: + prims.mbarrier_init(b_gu_rdy, 1) + prims.mbarrier_init(b_dn_rdy, 1) + prims.mbarrier_init(acc_gu, 1) + prims.mbarrier_init(acc_dn, 1) + if warp == 1: + prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency + _sh_issue(lane, sid, cutlass.Int32(g), s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) + if cutlass.const_expr(SH_PRETILE): + # this CTA's stages beyond the ring, 4 per lane + for r in cutlass.range_constexpr(2): + pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(2) + cutlass.Int32(r) + if pg < cutlass.Int32(SH_STAGES): + _l2_prefetch( + gusf.subview( + (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) + ) + .data_ptr() + .toint(), + SH_STAGE_BYTES, + ) + else: + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + NQ = THREADS - 64 # activation quant threads (warps 2 ..) + if warp == 0: + for g in cutlass.range(SH_NST, SH_STAGES, unroll=1): + _wait( + sh_empty.subview(g % cutlass.Int32(SH_NST)), + ((g // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), + ) + _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + elif warp == 1: + # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down + idesc_gu = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=SH_NB, + m_dim=128, + ) + idesc_dn = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=8, + m_dim=128, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + desc_a0 = prims.Tcgen05SmemDesc.build( + s_big.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bx = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BX), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bh = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BH), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + sq = sid & cutlass.Int32(3) + _wait(b_gu_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for g in cutlass.range(SH_STAGES, unroll=1): + st = g % cutlass.Int32(SH_NST) + if g == cutlass.Int32(SH_GT): + _wait(b_dn_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _wait(sh_full.subview(st), (g // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) + raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) + if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): + # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 + w4 = [ + big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) + for qq in range(4) + ] + big32.store( + ( + cutlass.Int32(w4[0]), + cutlass.Int32(w4[1]), + cutlass.Int32(w4[2]), + cutlass.Int32(w4[3]), + ), + idx=cutlass.Int32(_SB_SPL // 4) + + st * cutlass.Int32(128) + + lane * cutlass.Int32(4), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + prims.mbarrier_arrive( + sh_empty.subview(st) + ) # this lane's generic reads of the slot are done + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + if g < cutlass.Int32(SH_GT): + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bx + cutlass.Int64(g * cutlass.Int32(SH_NB * 128 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), + da + 2 * kb, + db + 2 * kb, + idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_GT - 1): + prims.tcgen05_commit(acc_gu) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(SH_DKT) + kt = gd - jt * cutlass.Int32(SH_DKT) + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview( + cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) + + st * cutlass.Int32(512 if SH_PRETILE else 2048) + ), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bh + cutlass.Int64(kt * cutlass.Int32(1024 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr( + tcol + SH_ACC_DN + jt * cutlass.Int32(8), 6, cutlass.Float32 + ), + da + 2 * kb, + db + 2 * kb, + idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_STAGES - 1): + prims.tcgen05_commit(acc_dn) + cute.arch.sync_warp() + if warp >= 2: + # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = + # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes + qt = tidx - cutlass.Int32(64) + for r in cutlass.range_constexpr((M_MAX * SH_KB + NQ - 1) // NQ): + q_ = qt + cutlass.Int32(r * NQ) + if q_ < cutlass.Int32(M_MAX * SH_KB): + t = q_ // cutlass.Int32(SH_KB) + b = q_ - t * cutlass.Int32(SH_KB) + wv = [ + xb32.load( + idx=t * cutlass.Int32(H // 2) + + b * cutlass.Int32(16) + + cutlass.Int32(4 * z), + vector_size=4, + alignment=16, + ) + for z in range(4) + ] + v = [] + for z in cutlass.range_constexpr(4): + for zz in cutlass.range_constexpr(4): + w = cutlass.Int32(wv[z][zz]) + v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) + v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(v[z])) + code, osc = _mx_osc(amax) + words = _mx_words(v, osc) + kt = b >> cutlass.Int32(2) + sl = kt // cutlass.Int32(SH_GT) + it = kt - sl * cutlass.Int32(SH_GT) + n = sl * cutlass.Int32(8) + t + c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BX) + it * cutlass.Int32(SH_NB * 128) + n * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) + if warp >= 4: + if warp < 8: + qd = warp - cutlass.Int32(4) + # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of k-tile SH_GT q + i + n = lane + t = n & cutlass.Int32(7) + ok = (n < cutlass.Int32(SH_NB)) & (t < cutlass.Int32(M_MAX)) + for i in cutlass.range_constexpr(SH_GT): + kt = (n >> cutlass.Int32(3)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) + wsf = big32.load( + idx=( + cutlass.Int32(_SB_XC) + + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) + + kt * cutlass.Int32(4) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), + cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + et = tidx - cutlass.Int32(128) + if et == 0: + prims.mbarrier_arrive(b_gu_rdy) + _stamp(tdbg, bx, 54) + # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t at column 8 q + t + _wait(acc_gu, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if cutlass.const_expr(SH_SPLIT == 2): + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=16, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(2): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[8 + tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) + * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd < cutlass.Int32(2): + # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + else: + # 16-row octants: warp qd = slice qd, lane l = (u = l / 16 gate / up, position l % 16); its tokens are + # accumulator columns 8 qd + t + av = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_GU + qd * cutlass.Int32(8), + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(1): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(1)) * cutlass.Int32(32) + lane) + * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd == cutlass.Int32(0): + # m3_mx_gemv's owner sum: (((0 + rank 0) + rank 1) + rank 2) + rank 3, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + for s_ in cutlass.range_constexpr(1, SH_SPLIT): + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (cutlass.Int32((s_ - 1) * 32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + lane * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag + for r in cutlass.range_constexpr((SH_ROWS * M_MAX + 127) // 128): + pq = et + cutlass.Int32(r * 128) + if pq < cutlass.Int32(SH_ROWS * M_MAX): + p = pq // cutlass.Int32(M_MAX) + tt = pq - p * cutlass.Int32(M_MAX) + ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(8) + tt) + up = bigf.load( + idx=cutlass.Int32(_SB_GU // 4) + + (cutlass.Int32(SH_ROWS) + p) * cutlass.Int32(8) + + tt + ) + q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) + q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) + q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) + q_e2 = _ex2_full( + _mul_rn( + _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), + cutlass.Float32(1.4426950216293334961), + ) + ) + q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) + q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) + _st_relaxed_u32( + hx.subview(tt * cutlass.Int32(SH_I) + sid * cutlass.Int32(SH_ROWS) + p) + .data_ptr() + .toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), + ) + if et == 0: + _stamp(tdbg, bx, 50) + # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes + for rr in cutlass.range_constexpr((M_MAX * SH_DKB + 127) // 128): + q_ = et + cutlass.Int32(rr * 128) + if q_ < cutlass.Int32(M_MAX * SH_DKB): + t = q_ // cutlass.Int32(SH_DKB) + kb = q_ - t * cutlass.Int32(SH_DKB) + base = t * cutlass.Int32(SH_I) + kb * cutlass.Int32(32) + hv = [cutlass.Int32(0) for _ in range(32)] + pend = cutlass.Int32(1) + while pend != cutlass.Int32(0): + for q4 in cutlass.range_constexpr(8): + v4 = prims.load_ext( + hx.subview(base + cutlass.Int32(4 * q4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for z in cutlass.range_constexpr(4): + hv[4 * q4 + z] = cutlass.Int32(v4[z]) + okw = cutlass.Boolean(True) + for z in cutlass.range_constexpr(32): + okw = okw & ( + ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 + ) + pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) + vals = [ + (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) + ] + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) + code, osc = _mx_osc(amax) + words = _mx_words(vals, osc) + kt = kb >> cutlass.Int32(2) + c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BH) + kt * cutlass.Int32(1024) + t * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + (((c0 + cutlass.Int32(1)) ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + okh = lane < cutlass.Int32(M_MAX) + for kt in cutlass.range_constexpr(SH_DKT): + wsf = big32.load( + idx=( + cutlass.Int32(_SB_HC) + + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) + + cutlass.Int32(4 * kt) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), + cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if et == 0: + prims.mbarrier_arrive(b_dn_rdy) + _stamp(tdbg, bx, 51) + # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l + _wait(acc_dn, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jt in cutlass.range_constexpr(2): + dv = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_DN + jt * 8, + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + rowo = ( + sid * cutlass.Int32(256) + + cutlass.Int32(128 * jt) + + qd * cutlass.Int32(32) + + lane + ) + for tt in cutlass.range_constexpr(M_MAX): + _st_relaxed_u32( + dnx.subview(cutlass.Int32(tt * H) + rowo).data_ptr().toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), + ) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if et == 0: + _stamp(tdbg, bx, 52) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if warp == 1: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) + if tidx == 0: + epochs.store(ep + cutlass.Int32(1), idx=bx) + + +@cute.kernel +def m3_moe_kernel( + tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], # w3_w1_weight_scale as 128-byte lines + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], # w2_weight_scale as 128-byte lines + tma_gu_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: gate_up e4m3 [1536, 6144], box 128 B x 32 rows + tma_dn_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: down e4m3 [6144, 768], box 128 B x 128 rows + logits: cutlass.Array, # f32 [M, 128] (not FUSED_ROUTER) + bias: cutlass.Array, # f32 [128] + x32: cutlass.Array, # int32 words of x_fp4 [M, 3072 B] (not FUSED_ROUTER) + xsf32: cutlass.Array, # int32 words of x_sf [M, 384 B] (not FUSED_ROUTER) + escal: cutlass.Array, # f32 [6, 128]: fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha + plan: cutlass.Array, # int32 [4 M + 1, NCTA, NG_CAP + 1] by G: per group i0 | i1 << 16, then b0 | nb << 16 + partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums + hbuf32: cutlass.Array, # int32 [G_CAP * M, 48, 4]: tagged intermediate records + epochs: cutlass.Array, # int32 [NCTA] + out: cutlass.Array, # bf16 [M, 6144] + route_dbg: cutlass.Array, # int32 [M, 8]: ids, weight bits (CTA 0) + tdbg: cutlass.Array, # int64 [NCTA, 32] timing stamps (TIMING builds) + xb32: cutlass.Array, # FUSED_ROUTER: int32 words (bf16 pairs) of the hidden states [M, 6144] + rw: cutlass.Array, # FUSED_ROUTER: f32 router weight [128, 6144] + sfsc: cutlass.Array, # FUSED_ROUTER: f32 [1] fc31_input_scale + lgbuf: cutlass.Array, # FUSED_ROUTER: int64 [LG_REPL, LG_STRIDE] logits: tag << 32 | f32 bits + guw: cutlass.Array, # SHARED: uint8 gate_up e4m3 [1536, 6144] + gusf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + dnw: cutlass.Array, # SHARED: uint8 down e4m3 [6144, 768] + dnsf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + hx: cutlass.Array, # SHARED: int32 [M, 768]: h as (call tag << 16 | bf16), tag 0 = never written + dnx: cutlass.Array, # SHARED: int32 [M, 6144]: y_sh as (call tag << 16 | bf16), tag 0 = never written + arp: cutlass.Array, # int64 [8]: all-reduce push (enable, mc, flags, rank, world, maxt); enable 0 = off + s1w: cutlass.Array, # SFA_CP: int32 words of w3_w1_weight_scale + s2w: cutlass.Array, # SFA_CP: int32 words of w2_weight_scale + img32: cutlass.Array, # B2IMG: int32 [NG_CAP, IMG_BYTES / 4] per-group FC2 B image (tile + scale atoms) + b2cnt: cutlass.Array, # B2IMG: int32 [NG_CAP * 8 + 8] per-group unit counters (index 8 g) +): + tidx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + if tidx == 0: + _stamp(tdbg, bx, 0) + if cutlass.const_expr(TIMING): + if tidx == 0: + tdbg.store( + cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) + ) + + s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) + sA = cutlass.Array( + s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 + ) + sSFAraw = cutlass.Array( + s_big.data_ptr(_OFF_SFAR), + shape=(max(1, _RAW_STAGE // 4) * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sSFA = cutlass.Array( + s_big.data_ptr(_OFF_SFA), + shape=(NUM_WORDS_SFA * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sB1 = cutlass.Array( + s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 + ) + sB2 = cutlass.Array( + s_big.data_ptr(_OFF_B2), shape=(2 * KT2 * NUM_BYTES_B2,), dtype=cutlass.Int8, alignment=1024 + ) + sSFB2 = cutlass.Array( + s_big.data_ptr(_OFF_SFB2), + shape=(2 * KT2 * NUM_KB * 128,), + dtype=cutlass.Int32, + alignment=1024, + ) + ys = cutlass.Array( + cutlass.Float32, NB2_CAP * M_MAX * TOP_K * R2, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_part = cutlass.Array( + cutlass.Float32, 2 * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_red = cutlass.Array(cutlass.Float32, 9 * M_MAX, space=cutlass.AddressSpace.smem, alignment=16) + s_xsf = cutlass.Array( + cutlass.Int32, + M_MAX * (H // 64) if FUSED_ROUTER else 4, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_free = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + # this CTA's FC1 share of group g is out: one barrier per group parity (a single barrier could complete twice + # while warp 5 still works on an earlier group, and its parity wait would then miss them) + fc1_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) + b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + x_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # x landed, x consumed + img_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # (B2IMG) image landed + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + s_el = cutlass.Array( + cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16 + ) # slots, then G + s_kpos = cutlass.Array( + cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_w = cutlass.Array( + cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_ids = cutlass.Array( + cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + # this CTA's plan rows for every G (static data: loaded before the grid dependency) + s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_plan_all = cutlass.Array( + cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 + ) + s_esc = cutlass.Array( + cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32 + ) # slot s at 8 s + + is_routed = cutlass.Boolean(True) + if cutlass.const_expr(SHARED): + sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) + is_routed = bx < cutlass.Int32(NCTA) + if bx >= cutlass.Int32(NCTA): + _shared_cta( + tidx, + warp, + lane, + bx - cutlass.Int32(NCTA), + s_big, + sh_full, + epochs, + tdbg, + bx, + xb32, + tma_gu_desc.get_ptr(), + tma_dn_desc.get_ptr(), + gusf, + dnsf, + hx, + dnx, + arp, + tmem_ptr_i32, + guw, + dnw, + ) + if is_routed: + # ---- before the grid dependency + if warp == 0: + if tidx < STAGES: + prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) + prims.mbarrier_init(ab_empty.subview(tidx), 1) + prims.mbarrier_init(sfa_ready.subview(tidx), 1) + if tidx < 2: + prims.mbarrier_init(acc_full.subview(tidx), 1) + prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(acc2_full.subview(tidx), 1) + prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(b2_ready.subview(tidx), 1) + prims.mbarrier_init(b2_free.subview(tidx), EPI_THREADS) + if cutlass.const_expr(B2IMG): + prims.mbarrier_init(img_bar.subview(tidx), 1) + if tidx < 4: + prims.mbarrier_init(fc1_done.subview(tidx), EPI_THREADS) + if tidx == 0: + prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) + if cutlass.const_expr(XSTAGE): + prims.mbarrier_init(x_bar.subview(0), 1) + prims.mbarrier_init(x_bar.subview(1), EPI_THREADS) + if warp == MMA_WARP: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == W_WARP: + prims.prefetch_tensormap(tma_a1_desc.get_ptr()) + prims.prefetch_tensormap(tma_a2_desc.get_ptr()) + prims.prefetch_tensormap(tma_s1_desc.get_ptr()) + prims.prefetch_tensormap(tma_s2_desc.get_ptr()) + for i in cutlass.range_constexpr(E_LOCAL // 32): + s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) + for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): + q = lane + cutlass.Int32(32 * i) + if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): + gq = q // cutlass.Int32(NG_CAP + 1) + s_plan_all.store( + plan.load( + idx=(gq * cutlass.Int32(NCTA) + bx) * cutlass.Int32(NG_CAP + 1) + + q + - gq * cutlass.Int32(NG_CAP + 1) + ), + idx=q, + ) + rwv = [] + if cutlass.const_expr(FUSED_ROUTER): + tg = cutlass.select_(tidx < cutlass.Int32(GEMV_THREADS), tidx, cutlass.Int32(0)) + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + rbx = cutlass.select_(eg < cutlass.Int32(NROUTE), eg, cutlass.Int32(0)) + rwv.append( + [ + rw.load( + idx=rbx * cutlass.Int32(H) + + tg * cutlass.Int32(4) + + cutlass.Int32(1024 * j), + vector_size=4, + alignment=16, + ) + for j in range(H // 1024) + ] + ) + if cutlass.const_expr(TIMING): + if tidx == 0: + s_el.store(cutlass.Int32(-7), idx=G_CAP) + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(XSTAGE): + if warp == 0: + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_bar.subview(0), M_MAX * H * 2) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(lane * cutlass.Int32(H * 2)), + xb32.subview(lane * cutlass.Int32(H // 2)), + x_bar.subview(0), + H * 2, + ) + ar_on = arp.load(idx=0) != cutlass.Int64(0) + ar_off = arp.load(idx=0) == cutlass.Int64(0) + if cutlass.const_expr(not LATE_TRIGGER): + if ar_off: + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + + ep = epochs.load(idx=bx) + tagv = ((ep & cutlass.Int32(0x7F)) | cutlass.Int32(0x80)) << cutlass.Int32(24) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + if cutlass.const_expr(FUSED_ROUTER): + # ---- router GEMV (experts bx, bx + NCTA, .., threads 0..255), m3_router_gemv's order exactly + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + if (eg < cutlass.Int32(NROUTE)) & (tidx < cutlass.Int32(GEMV_THREADS)): + accs = [cutlass.Float32(0.0) for _ in range(M_MAX)] + xsrc = xb32 + if cutlass.const_expr(GEMV_SMEM and XSTAGE): + _wait(x_bar.subview(0), 0) + xsrc = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + for j in cutlass.range_constexpr(H // 1024): + for t in cutlass.range_constexpr(M_MAX): + xw = xsrc.load( + idx=cutlass.Int32(t * (H // 2) + 512 * j) + tidx * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + x0 = cutlass.Int32(xw[0]) + x1 = cutlass.Int32(xw[1]) + a = accs[t] + a = _fma_rn( + (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][0]), + a, + ) + a = _fma_rn( + (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][1]), + a, + ) + a = _fma_rn( + (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][2]), + a, + ) + a = _fma_rn( + (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][3]), + a, + ) + accs[t] = a + for t in cutlass.range_constexpr(M_MAX): + v = accs[t] + for off in (16, 8, 4, 2, 1): + v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) + if lane == 0: + s_red.store(v, idx=warp * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < M_MAX: + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX) + tidx)) + s_red.store(sacc, idx=cutlass.Int32(8 * M_MAX) + tidx) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < LG_REPL: + for t in cutlass.range_constexpr(M_MAX): + lw = ltag | ( + cutlass.Int64( + cutlass.Float32(s_red.load(idx=8 * M_MAX + t)).bitcast( + cutlass.Int32 + ) + ) + & cutlass.Int64(0xFFFFFFFF) + ) + _st_relaxed_u64( + lgbuf.subview( + tidx * cutlass.Int32(LG_STRIDE) + + cutlass.Int32(t * E_LOCAL) + + eg + ) + .data_ptr() + .toint(), + lw, + ) + if cutlass.const_expr(NGEMV > 1): + cute.arch.sync_warp() + cute.arch.barrier( + barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS + ) # s_red reused next + if tidx == 0: + _stamp(tdbg, bx, 21) + + # ---- routing (warps 4..7, token t on warp 4 + t % 4); the epilogue warps build B meanwhile + if ((warp >= W_WARP) & (warp < cutlass.Int32(W_WARP + NRW0))) | ( + (warp >= cutlass.Int32(8)) & (warp < cutlass.Int32(8 + RW_EXTRA)) + ): + rw_i = cutlass.select_( + warp < cutlass.Int32(8), warp - cutlass.Int32(W_WARP), warp - cutlass.Int32(4) + ) + # ILP: poll every token's logits first, then run the independent top-k chains back to back + lgv_all = [] + t_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = rw_i + cutlass.Int32(NRW * tr) + if cutlass.const_expr(M_MAX % NRW != 0): + t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) + t_all.append(t) + lgv = [] + if cutlass.const_expr(FUSED_ROUTER): + want = ltag >> cutlass.Int64(32) + rbase = ( + (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) + + t * cutlass.Int32(E_LOCAL) + + lane + ) + done = cutlass.Int32(0) + res = [cutlass.Int64(0) for _ in range(4)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(4): + res[z] = _ld_relaxed_u64( + lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() + ) + ok = ok & ((res[z] >> cutlass.Int64(32)) == want) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32( + -1 + ): + done = cutlass.Int32(1) + else: + _nanosleep(POLL_NS) + for z in cutlass.range_constexpr(4): + lgv.append(cutlass.Int32(res[z]).bitcast(cutlass.Float32)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 22) + else: + for i in cutlass.range_constexpr(4): + lgv.append( + cutlass.Float32( + logits.load( + idx=t * cutlass.Int32(E_LOCAL) + cutlass.Int32(i * 32) + lane + ) + ) + ) + lgv_all.append(lgv) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 61) + res_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = t_all[tr] + lgv = lgv_all[tr] + own = (rw_i + cutlass.Int32(NRW * tr)) < cutlass.Int32(M_MAX) + # top-4 (keys sigmoid + bias, ties to the lower id) + keys = [] + for i in cutlass.range_constexpr(4): + sg = _fma_rn( + cutlass.Float32(0.5), + cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), + cutlass.Float32(0.5), + ) + keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) + hi = [_twiddle(keys[i]) for i in range(4)] + taken = [cutlass.Boolean(False) for _ in range(4)] + ids_t = [] + kt_t = [] + for k in cutlass.range_constexpr(TOP_K): + bh = cutlass.Int32(0) + bl = cutlass.Int32(0) + for i in cutlass.range_constexpr(4): + lo_i = cutlass.Int32(65535 - i * 32) - lane + hv = cutlass.select_(taken[i], cutlass.Int32(0), hi[i]) + lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) + better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( + (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) + ) + bh = cutlass.select_(better, hv, bh) + bl = cutlass.select_(better, lv, bl) + mh = _redux_max_u32(bh) + ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) + e_k = cutlass.Int32(65535) - ml + for i in cutlass.range_constexpr(4): + taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) + ids_t.append(e_k) + kv = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(4): + kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) + kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) + sig = cutlass.Float32(0.0) + e_l = cutlass.Int32(0) + for k in cutlass.range_constexpr(TOP_K): + sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) + sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) + e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) + total = sig + for off in (16, 8, 4, 2, 1): + total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) + wv = _bf16_round( + _div_rn( + _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), + _add_rn(total, cutlass.Float32(SUM_EPS)), + ) + ) + res_all.append((t, own, wv, e_l)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 62) + for tr in cutlass.range_constexpr( + (M_MAX + NRW - 1) // NRW + ): # stores after every chain (one branch-free block) + t, own, wv, e_l = res_all[tr] + if (lane < TOP_K) & own: + s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) + s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) + if bx == 0: + route_dbg.store( + wv.bitcast(cutlass.Int32), + idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, + ) + cute.arch.sync_warp() + if cutlass.const_expr(EARLY0 and not E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + _bar_arrive( + E0_BAR, 64 + ) # token 0's experts are in s_ids: warp 7 issues the early stages + if cutlass.const_expr(B1_AFTER_TOPK): + _bar_arrive(TOPK2_BAR, 128 + 32 * NRW) + cute.arch.barrier(barrier_id=TOPK_BAR, number_of_threads=32 * NRW) + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(E0_WARP): + cute.arch.barrier(barrier_id=E0_BAR, number_of_threads=64) + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + _bar_arrive(E1_BAR, 64) + if warp == W_WARP: + if lane == 0: + _stamp(tdbg, bx, 20) + # slots: lane l < 4 M is (token l / 4, top-k l % 4); the distinct experts in lane order (match.any) + valid = lane < cutlass.Int32(TOP_K * M_MAX) + my_e = cutlass.select_( + valid, + s_ids.load(idx=cutlass.select_(valid, lane, cutlass.Int32(0))), + cutlass.Int32(-1) - lane, + ) + same = _match_any(my_e) + leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) + first = valid & (lane == leader) + bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) + rank = cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) + slot_l = cute.arch.shuffle_sync(rank, leader) + for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): + s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + cute.arch.sync_warp() + if first: + s_el.store(my_e, idx=rank) + if valid: + s_kpos.store( + lane % cutlass.Int32(TOP_K), + idx=slot_l * cutlass.Int32(M_MAX) + lane // cutlass.Int32(TOP_K), + ) + if bx == 0: + route_dbg.store( + my_e, + idx=(lane // cutlass.Int32(TOP_K)) * cutlass.Int32(8) + + lane % cutlass.Int32(TOP_K), + ) + gcnt = cute.arch.popc(bal) + if lane == 0: + s_el.store(gcnt, idx=G_CAP) + if lane < cutlass.Int32(NG_CAP + 1): # the plan row of this G (preloaded) + s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) + cute.arch.sync_warp() + if cutlass.const_expr(B2IMG): + # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) + if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): + nh = gcnt - lane * cutlass.Int32(4) + nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) + nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) + pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) + if pad > cutlass.Int32(0): + _red_release_add( + b2cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad + ) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 24) + cute.arch.sync_warp() + _bar_arrive(ROUTE_BAR, THREADS) + if warp < 4: + if cutlass.const_expr(B1_AFTER_TOPK): + cute.arch.barrier(barrier_id=TOPK2_BAR, number_of_threads=128 + 32 * NRW) + if cutlass.const_expr(FUSED_ROUTER): + if tidx == 0: + _stamp(tdbg, bx, 57) + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 58) + xs32 = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xs32, sfsc, sB1, s_xsf, b1_ready + ) + prims.mbarrier_arrive(x_bar.subview(1)) # the ring is free for the weights + else: + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xb32, sfsc, sB1, s_xsf, b1_ready + ) + if tidx == 0: + _stamp(tdbg, bx, 23) + else: + _build_b1(tidx, warp, lane, base_row, base_col, x32, xsf32, sB1, b1_ready) + if warp != W_WARP: # one sync site for every thread but warp 4's (which arrived) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ROUTE_BAR, number_of_threads=THREADS) + if tidx == 0: + _stamp(tdbg, bx, 2) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(s_el.load(idx=G_CAP)), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(28), + ) + + g_n = s_el.load(idx=G_CAP) + ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) + pr2 = s_plan.load(idx=NG_CAP) + blk0 = pr2 & cutlass.Int32(0xFFFF) + nb = pr2 >> cutlass.Int32(16) + + # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG + if warp == W_WARP: + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) # the staged hidden states are quantized + stage = cutlass.Int32(0) + phase = cutlass.Int32(1) + skip = cutlass.Int32(0) + if cutlass.const_expr(E0_OFFLOAD): + cute.arch.barrier(barrier_id=E1_BAR, number_of_threads=64) + if cutlass.const_expr(EARLY0): + skip = s_misc.load(idx=0) # group-0 stages already in the ring + stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) + phase = cutlass.select_( + skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) + ) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if skip > cutlass.Int32(0): + skip = skip - cutlass.Int32(1) + else: + _waitd(ab_empty.subview(stage), phase, tdbg, 1) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + ), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, stage, sf_line, c, sSFA, s1w, ab_full) + stage, phase = _adv(stage, phase) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + _waitd(ab_empty.subview(stage), phase, tdbg, 2) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + n_here + * cutlass.Int32( + NUM_BYTES_BOX1 + if SFA_CP + else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4 + ), + ) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(jj * NUM_BYTES_BOX1) + ), + tma_a2_desc.get_ptr(), + ( + cutlass.Int32(t2 * KT_BYTES), + blk * cutlass.Int32(R2), + e, + ), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) + ), + tma_s2_desc.get_ptr(), + ( + cutlass.Int32(0), + ( + ( + e * cutlass.Int32(SF2_RB) + + (blk >> cutlass.Int32(2)) + ) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(4), + ), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf2_cp(lane, stage, g, blk, t2, n_here, s_el, sSFA, s2w, ab_full) + stage, phase = _adv(stage, phase) + if lane == 0: + _stamp(tdbg, bx, 3) + + # ---- scale splicer (warp 6): every stage, atom kb word [r][q] = raw atom 4 q + kb word [r][m1] + if (warp == S_WARP) & cutlass.Boolean(not SFA_CP): # (SFA_CP: no splicer) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + i0 = cutlass.Int32(0) + i1 = cutlass.Int32(0) + for part in cutlass.range_constexpr(2): + cnt = cutlass.Int32(0) + m1 = cutlass.Int32(0) + nk = cutlass.Int32(0) + if cutlass.const_expr(part == 0): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + nk = _nsegs(i0, i1) + else: + if p >= cutlass.Int32(LAG): + nk = nb + for k in range(nk): + if cutlass.const_expr(part == 0): + ug, j0, j1 = _seg(i0, i1, k) + m1 = ( + ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + ) & cutlass.Int32(3) + cnt = j1 - j0 + else: + m1 = (blk0 + k) & cutlass.Int32(3) + cnt = cutlass.Int32(KT2) + for _j in cutlass.range(cnt, unroll=1): + _waitd(ab_full.subview(stage), phase, tdbg, 3) + raw0 = ( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + lane * cutlass.Int32(4) + + m1 + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for kb in cutlass.range_constexpr(NUM_KB): + w0 = sSFAraw.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) + w1 = sSFAraw.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) + w2 = sSFAraw.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) + w3 = sSFAraw.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) + sSFA.store( + ( + cutlass.Int32(w0), + cutlass.Int32(w1), + cutlass.Int32(w2), + cutlass.Int32(w3), + ), + idx=dst0 + cutlass.Int32(kb * 128), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + prims.tcgen05_commit(sfa_ready.subview(stage)) + stage, phase = _adv(stage, phase) + + if cutlass.const_expr(B2IMG): + # ---- (B2IMG) group releases (warp 6): this CTA's finished units of group p, one release add per group + if warp == cutlass.Int32(S_WARP): + for p in range(ng): + i0, i1 = _plan_range(s_plan, p) + n_done = cutlass.Int32(0) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + n_done = n_done + cutlass.select_( + (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), + cutlass.Int32(1), + cutlass.Int32(0), + ) + _wait( + fc1_done.subview(p & cutlass.Int32(3)), + (p >> cutlass.Int32(2)) & cutlass.Int32(1), + ) + if cutlass.const_expr(B2_DLY > 0): + t0d = _gtime() + dl = cutlass.Int64( + ( + (bx * cutlass.Int32(37) + p * cutlass.Int32(11) + ep) + & cutlass.Int32(63) + ) + * cutlass.Int32(B2_DLY // 64) + ) + while _gtime() - t0d < dl: + pass + if (lane == 0) & (n_done > cutlass.Int32(0)): + _fence_proxy_async_global() + _red_release_add( + b2cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done + ) + cute.arch.sync_warp() + # ---- (B2IMG) FC2 B loader (warp 5): group g's image once its counter is complete + if warp == cutlass.Int32(X_WARP): + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + if lane == 0: + _stamp(tdbg, bx, 29) + caddr = b2cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() + if lane == 0: + cv = _ld_acquire_u32(caddr) + while (cv - cnt_target) < cutlass.Int32(0): + _nanosleep(B2IMG_NS) + cv = _ld_acquire_u32(caddr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) + prims.cp_async_bulk_shared_cluster_global( + sB2.subview(buf * cutlass.Int32(IMG_B2)), + img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), + img_bar.subview(buf), + IMG_B2, + ) + prims.cp_async_bulk_shared_cluster_global( + sSFB2.subview(buf * cutlass.Int32(IMG_SF // 4)), + img32.subview( + g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) + ), + img_bar.subview(buf), + IMG_SF, + ) + cute.arch.sync_warp() + _waitd( + img_bar.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 108 + ) + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 25) + # ---- FC2 B builder (warps 5, 8 ..): group g's records (all 48 per (slot, token), the call's tag on + # every word) into B2 buffer g % 2 (row j M + t, k-tile t2 = record / 16) and its scale atoms, then + # tcgen05.cp to TMEM. + if ((warp == X_WARP) | (warp >= cutlass.Int32(8))) & cutlass.Boolean(not B2IMG): + xi = cutlass.select_( + warp == cutlass.Int32(X_WARP), cutlass.Int32(0), warp - cutlass.Int32(7) + ) + if cutlass.const_expr(G_CAP > TOP_K): + if (bx == 0) & (xi == cutlass.Int32(0)): + # rows of slots this call does not use: tag 0, never taken for a later call's + for i in range( + g_n * cutlass.Int32(M_MAX * UNITS_ROW) + lane, + cutlass.Int32(G_CAP * M_MAX * UNITS_ROW), + 32, + ): + for rr in cutlass.range_constexpr(HB_REPL): + hbuf32.store( + ( + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + ), + idx=i * cutlass.Int32(4) + + (cutlass.Int32(rr * HB_STRIDE) if rr else 0), + alignment=16, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sSFB2_8 = cutlass.Array( + sSFB2.data_ptr(0), shape=(2 * KT2 * NUM_KB * 512,), dtype=cutlass.Int8, alignment=16 + ) + sB2_32 = cutlass.Array( + sB2.data_ptr(0), + shape=(2 * KT2 * NUM_BYTES_B2 // 4,), + dtype=cutlass.Int32, + alignment=16, + ) + nrec_cap = 4 * M_MAX * UNITS_ROW # records of one group + hb_off = cutlass.Int32(0) + if cutlass.const_expr(HB_REPL > 1): + # this CTA's copy of the records + hb_off = (bx % cutlass.Int32(HB_REPL)) * cutlass.Int32(HB_STRIDE) + NZ = (nrec_cap + 32 * XW - 1) // (32 * XW) # records per lane: lane + 32 (XW z + xi) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + _waitd( + fc1_done.subview(g & cutlass.Int32(3)), + (g >> cutlass.Int32(2)) & cutlass.Int32(1), + tdbg, + 4, + ) # this CTA's own share of group g is out + if cutlass.const_expr(TIMING): + if (lane == 0) & (xi == cutlass.Int32(0)) & (g < cutlass.Int32(8)): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(32) + g + ) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + nrec = n_here * cutlass.Int32(M_MAX * UNITS_ROW) + rec0 = g * cutlass.Int32( + 4 * M_MAX * UNITS_ROW + ) # rows (4 g + j) M + t are contiguous + done = cutlass.Int32(0) + rspin = cutlass.Int32(0) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 29) + npoll = cutlass.Int32(0) + got = [[cutlass.Int32(0) for _ in range(3)] for _ in range(NZ)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + v4 = prims.load_ext( + hbuf32.subview( + (rec0 + qi) * cutlass.Int32(4) + (hb_off if HB_REPL > 1 else 0) + ), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for wq in cutlass.range_constexpr(3): + got[z][wq] = cutlass.Int32(v4[wq]) + ok = ok & ( + ( + ((got[z][0] ^ tagv) | (got[z][1] ^ tagv) | (got[z][2] ^ tagv)) + & cutlass.Int32(-16777216) + ) + == cutlass.Int32(0) + ) + npoll = npoll + cutlass.Int32(1) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): + done = cutlass.Int32(1) + else: + _nanosleep(B2_POLL_NS) + if cutlass.const_expr(HANG_DBG): + rspin = rspin + cutlass.Int32(1) + if rspin > cutlass.Int32(1 << 14): + done = cutlass.Int32(1) + tdbg.store( + cutlass.Int64(200), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + warp, + ) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + w0, w1, w2 = got[z][0], got[z][1], got[z][2] + rr = qi // cutlass.Int32( + UNITS_ROW + ) # row within the group: j M + t = B row n + c = qi - rr * cutlass.Int32(UNITS_ROW) + t2 = c >> cutlass.Int32(4) + cc = (c >> cutlass.Int32(1)) & cutlass.Int32(7) + lo = (w0 & cutlass.Int32(0xFFFFFF)) | (w1 << cutlass.Int32(24)) + hi = ((w1 >> cutlass.Int32(8)) & cutlass.Int32(0xFFFF)) | ( + (w2 & cutlass.Int32(0xFFFF)) << cutlass.Int32(16) + ) + sB2_32.store( + (lo, hi), + idx=( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_BYTES_B2) + + rr * cutlass.Int32(KT_BYTES) + + (cc ^ (rr & cutlass.Int32(7))) * cutlass.Int32(16) + + (c & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + # scale of block c: atom (t2, kb = (c % 16) / 4), row rr, byte c % 4 + sSFB2_8.store( + cutlass.Int8((w2 >> cutlass.Int32(16)) & cutlass.Int32(0xFF)), + idx=( + ( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_KB) + + ((c & cutlass.Int32(15)) >> cutlass.Int32(2)) + ) + * cutlass.Int32(128) + + rr * cutlass.Int32(4) + ) + * cutlass.Int32(4) + + (c & cutlass.Int32(3)), + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + if cutlass.const_expr(XW > 1): + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=B2_BAR, number_of_threads=32 * XW) + cute.arch.sync_warp() + if (xi == cutlass.Int32(0)) & prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + SFB2_COL + buf * cutlass.Int32(KT2 * NUM_KB * 4) + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 25) + if cutlass.const_expr(TIMING): + if g < cutlass.Int32(8): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(11) + g + ) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(npoll), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(30), + ) + + # ---- MMA (warp 7) + if warp == MMA_WARP: + idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N1, + m_dim=MMA_M, + sparsity_version=0, + ) + idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N2, + m_dim=MMA_M, + sparsity_version=0, + ) + desc_a0 = prims.Tcgen05SmemDesc.build( + sA.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b10 = prims.Tcgen05SmemDesc.build( + sB1.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b20 = prims.Tcgen05SmemDesc.build( + sB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + _waitd(b1_ready, 0, tdbg, 5) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + acc_b2 = cutlass.Int64(0) + acc_sf = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + _waitd( + acc_empty.subview(ab), + ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 102, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_ptr = cutlass.inttoptr( + tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 + ) + for j in cutlass.range(j0, j1, unroll=1): + tw0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 6) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 6) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc1, + (j > j0) | cutlass.Boolean(kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(ab)) + k1 = k1 + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + buf = g & cutlass.Int32(1) + tw2 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + _waitd( + b2_ready.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 103 + ) + if cutlass.const_expr(TIMING): + acc_b2 = acc_b2 + (_gtime() - tw2) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 4) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_empty.subview(a2), + ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 104, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc2_ptr = cutlass.inttoptr( + tcol + ACC2_COL + a2 * cutlass.Int32(N2), 6, cutlass.Float32 + ) + for t2 in cutlass.range_constexpr(KT2): + tw1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 7) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 7) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw1) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b20 + cutlass.Int64( + (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) + * cutlass.Int32(NUM_BYTES_B2 >> 4) + ) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = ( + tcol + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + cutlass.Int32(t2 * NUM_KB * 4) + ) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(_nkb2(t2)): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc2_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc2, + cutlass.Boolean(t2 > 0 or kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc2_full.subview(a2)) + k2 = k2 + cutlass.Int32(1) + if cutlass.const_expr(TIMING): + if lane == 0: + tdbg.store(acc_b2, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(26)) + tdbg.store(acc_sf, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(27)) + + # ---- epilogue warps (0-3) + if warp < 4: + row_w = base_row + warp * 32 + for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): + qi = tidx + cutlass.Int32(r * EPI_THREADS) + sl_e = qi // cutlass.Int32(6) + if (qi < cutlass.Int32(6 * G_CAP)) & (sl_e < g_n): + fld = qi - sl_e * cutlass.Int32(6) + s_esc.store( + cutlass.Float32( + escal.load(idx=fld * cutlass.Int32(E_LOCAL) + s_el.load(idx=sl_e)) + ), + idx=sl_e * cutlass.Int32(8) + fld, + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + SH_ADD = SHARED + wys = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + pars = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + ard = (cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0)) + for p in range(ng + cutlass.Int32(LAG)): + if cutlass.const_expr(AR_PRE): + if (p == ng + cutlass.Int32(LAG - 1)) & ar_on: + # one 8-byte load each (the vectorizer merges adjacent arp.load()s into 16-byte reads + # at 8-byte offsets: misaligned) + ard = tuple( + cutlass.Int64(_ld_relaxed_u64(arp.subview(i_).data_ptr().toint())) + for i_ in (1, 3, 4, 5) + ) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_cl = cutlass.select_( + f_bb < nb, (blk0 + f_bb) * cutlass.Int32(R2), cutlass.Int32(0) + ) + pars[bt] = _ar_parity(arp, f_t, f_cl) + if cutlass.const_expr(SH_ADD): + if p == ng + cutlass.Int32(LAG - 1): + # this CTA's y_sh words, loaded while the last group's FC2 runs (re-polled in finalize) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_ch = cutlass.select_( + f_bb < nb, + f_t * cutlass.Int32(H) + + (blk0 + f_bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + cutlass.Int32(0), + ) + wys[bt] = _ld_relaxed_u32(dnx.subview(f_ch).data_ptr().toint()) + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + unit = slot * cutlass.Int32(U1) + cblk + _waitd( + acc_full.subview(ab), + (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 105, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 + ), + num=N1, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc_empty.subview(ab)) + pb = ab * cutlass.Int32(M_MAX * C1 * R1) + for t in cutlass.range_constexpr(M_MAX): + part = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + part = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part + ) + s_part.store(part, idx=pb + (t * C1 + warp) * R1 + lane) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + for tr in cutlass.range_constexpr((M_MAX + 3) // 4): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + xv = _add_rn( + _add_rn( + _add_rn( + s_part.load(idx=pb + (t * C1 + 0) * R1 + lane), + s_part.load(idx=pb + (t * C1 + 1) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 2) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 3) * R1 + lane), + ) + finish = cutlass.Int32(1) + if (j0 != cutlass.Int32(0)) | (j1 != cutlass.Int32(S1)): + paddr = ( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint() + ) + if j0 == cutlass.Int32(0): + _st_relaxed_u32(paddr, xv.bitcast(cutlass.Int32)) + finish = cutlass.Int32(0) + else: + pbits = _ld_relaxed_u32(paddr) + spins = cutlass.Int32(0) + while pbits == cutlass.Int32(PARTIAL_EMPTY): + pbits = _ld_relaxed_u32(paddr) + spins = spins + cutlass.Int32(1) + _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) + xv = _add_rn(pbits.bitcast(cutlass.Float32), xv) + cute.arch.sync_warp() + if cutlass.const_expr( + XVDUMP + ): # debug: FC1 sums into the (unused) partial slots + partials.store( + xv.bitcast(cutlass.Int32), + idx=(unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane, + ) + if finish == cutlass.Int32(1): + gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) + is_up = ( + (lane >> cutlass.Int32(3)) & cutlass.Int32(1) + ) == cutlass.Int32(0) + es = slot * cutlass.Int32(8) + qv, sfb = _swiglu_quant( + xv, + gate_raw, + is_up, + s_esc.load(idx=es), + s_esc.load(idx=es + cutlass.Int32(1)), + s_esc.load(idx=es + cutlass.Int32(2)), + s_esc.load(idx=es + cutlass.Int32(3)), + s_esc.load(idx=es + cutlass.Int32(4)), + ) + byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) + bt = [ + cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) + for i in range(8) + ] + rw0 = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | tagv + ) + rw1 = ( + bt[3] + | (bt[4] << cutlass.Int32(8)) + | (bt[5] << cutlass.Int32(16)) + | tagv + ) + rw2 = ( + bt[6] + | (bt[7] << cutlass.Int32(8)) + | ((sfb & cutlass.Int32(0xFF)) << cutlass.Int32(16)) + | tagv + ) + if cutlass.const_expr(B2IMG): + if lane == 0: + # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 + # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), + # row n at byte 16 n, byte c % 4 + n = (ug // cutlass.Int32(U1)) * cutlass.Int32(M_MAX) + t + t2i = cblk >> cutlass.Int32(4) + cci = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) + gbase = p * cutlass.Int32(IMG_BYTES) + lo = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | (bt[3] << cutlass.Int32(24)) + ) + hi = ( + bt[4] + | (bt[5] << cutlass.Int32(8)) + | (bt[6] << cutlass.Int32(16)) + | (bt[7] << cutlass.Int32(24)) + ) + img32.store( + (lo, hi), + idx=( + gbase + + t2i * cutlass.Int32(NUM_BYTES_B2) + + n * cutlass.Int32(KT_BYTES) + + (cci ^ (n & cutlass.Int32(7))) + * cutlass.Int32(16) + + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + img8 = cutlass.Array( + img32.data_ptr(0), + shape=(NG_CAP * IMG_BYTES,), + dtype=cutlass.Int8, + alignment=16, + ) + img8.store( + cutlass.Int8(sfb & cutlass.Int32(0xFF)), + idx=gbase + + cutlass.Int32(IMG_B2) + + ( + ( + t2i * cutlass.Int32(NUM_KB) + + ( + (cblk & cutlass.Int32(15)) + >> cutlass.Int32(2) + ) + ) + * cutlass.Int32(512) + ) + + n * cutlass.Int32(16) + + (cblk & cutlass.Int32(3)), + ) + else: + if cutlass.const_expr(HB_REPL == 1): + if lane == 0: + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4), + alignment=16, + ) + else: + if lane < HB_REPL: # copy `lane` + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4) + + lane * cutlass.Int32(HB_STRIDE), + alignment=16, + ) + k1 = k1 + cutlass.Int32(1) + prims.mbarrier_arrive( + fc1_done.subview(p & cutlass.Int32(3)) + ) # warp 5 may poll group p + if tidx == 0: + _stamp(tdbg, bx, 10) + if cutlass.const_expr(TIMING): + if p < cutlass.Int32(8): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(40) + p, + ) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_full.subview(a2), + (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 106, + ) + if cutlass.const_expr(TIMING): + if (tidx == 0) & (g == ng - cutlass.Int32(1)): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(8) + bb, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, + base_col + + ACC2_COL + + a2 * cutlass.Int32(N2) + + warp * cutlass.Int32(M_MAX), + cutlass.Float32, + ), + num=LD_POW2, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc2_empty.subview(a2)) + sl = g * cutlass.Int32(4) + warp + if sl < g_n: + a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) + for t in cutlass.range_constexpr(M_MAX): + kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + if kp >= cutlass.Int32(0): + ys.store( + _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)), + idx=( + (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + * cutlass.Int32(TOP_K) + + kp + ) + * cutlass.Int32(R2) + + lane, + ) + k2 = k2 + cutlass.Int32(1) + prims.mbarrier_arrive(b2_free.subview(g & cutlass.Int32(1))) + if tidx == 0: + _stamp(tdbg, bx, 6) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + if tidx == 0: + _stamp(tdbg, bx, 19) + # finalize: (block bb, token t) per warp round: the fma chain over t's top-4 in order + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + pair = warp + cutlass.Int32(bt * 4) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + if bb < nb: + acc_o = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(TOP_K): + acc_o = _fma_rn( + s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ), + acc_o, + ) + if cutlass.const_expr( + YDUMP + ): # debug: y of (token, k) into the partials workspace + partials.store( + cutlass.Float32( + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ) + ).bitcast(cutlass.Int32), + idx=(t * cutlass.Int32(TOP_K) + cutlass.Int32(k)) * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + ) + ch = ( + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(SH_ADD): + # + y_sh (this row's shared output, from its shared CTA), as shared_output.add_(routed_output) + wy = wys[bt] + while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: + wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + acc_o = _add_rn( + (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) + ) + out.store(cutlass.BFloat16(acc_o), idx=ch) + if ar_on: + cl = ( + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(AR_PRE): + _ar_push_pre( + ard, + pars[bt], + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + else: + _ar_push( + arp, + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + + # ---- teardown + if tidx == 0: + _stamp(tdbg, bx, 7) + epochs.store(ep + cutlass.Int32(1), idx=bx) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if ar_on | cutlass.Boolean(LATE_TRIGGER): + # the all-reduce polls the pushed slots without a grid wait: release it after every push of this CTA + cute.arch.griddepcontrol_launch_dependents() + if warp == MMA_WARP: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS + ) + + +@cute.jit +def m3_moe( + a1_tensor: cute.Tensor, + sfa1_tensor: cute.Tensor, + a2_tensor: cute.Tensor, + sfa2_tensor: cute.Tensor, + logits: cute.Tensor, + bias: cute.Tensor, + x32: cute.Tensor, + xsf32: cute.Tensor, + escal: cute.Tensor, + plan: cute.Tensor, + partials: cute.Tensor, + hbuf32: cute.Tensor, + epochs: cute.Tensor, + out: cute.Tensor, + route_dbg: cute.Tensor, + tdbg: cute.Tensor, + xb32: cute.Tensor, + rw: cute.Tensor, + sfsc: cute.Tensor, + lgbuf: cute.Tensor, + guw: cute.Tensor, + gusf: cute.Tensor, + dnw: cute.Tensor, + dnsf: cute.Tensor, + hx: cute.Tensor, + dnx: cute.Tensor, + arp: cute.Tensor, + s1w: cute.Tensor, + s2w: cute.Tensor, + img32: cute.Tensor, + b2cnt: cute.Tensor, + stream: cuda_driver.CUstream, +): + tma_a1_desc = cuda.create_tensor_map_tiled( + global_address=a1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H // 2, TWO_I, E_LOCAL], + global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], + box_dims=(KT_BYTES, R1, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_a2_desc = cuda.create_tensor_map_tiled( + global_address=a2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[I_TP // 2, H, E_LOCAL], + global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], + box_dims=(KT_BYTES, R2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_s1_desc = cuda.create_tensor_map_tiled( + global_address=sfa1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_s2_desc = cuda.create_tensor_map_tiled( + global_address=sfa2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_gu_desc = tma_a1_desc + tma_dn_desc = tma_a2_desc + if cutlass.const_expr(SHARED and not SH_PRETILE): + tma_gu_desc = cuda.create_tensor_map_tiled( + global_address=guw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H, 2 * SH_I], + global_strides=[H // 16], + box_dims=(KT_BYTES, 32), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_dn_desc = cuda.create_tensor_map_tiled( + global_address=dnw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[SH_I, H], + global_strides=[SH_I // 16], + box_dims=(KT_BYTES, 128), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + m3_moe_kernel( + tma_a1_desc, + tma_a2_desc, + tma_s1_desc, + tma_s2_desc, + tma_gu_desc, + tma_dn_desc, + logits, + bias, + x32, + xsf32, + escal, + plan, + partials, + hbuf32, + epochs, + out, + route_dbg, + tdbg, + xb32, + rw, + sfsc, + lgbuf, + guw, + gusf, + dnw, + dnsf, + hx, + dnx, + arp, + s1w, + s2w, + img32, + b2cnt, + ).launch( + grid=(NCTA + NSH, 1, 1), + block=(THREADS, 1, 1), + cluster=(1, 1, 1), + stream=stream, + use_pdl=True, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_moe/op.py new file mode 100644 index 000000000000..79d4e637c251 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_moe/op.py @@ -0,0 +1,683 @@ +# 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. +"""Host side of the MiniMax-M3 routed-expert kernel on one TP8 rank (routed and shared intermediate 384 per rank), +over the TRTLLM-Gen NVFP4 buffers read in place. + +``m3_moe`` (m3_moe_kernel.py, :class:`M3MoeState`): 1..8 decode tokens (decode batches, spec-decode verify). The call +takes the bf16 hidden states and folds in the router GEMV (bit-exact with m3_router_gemv_fp4) and the activation's +NVFP4 quantization. With ``shared=True`` it also runs the shared expert on dedicated CTAs and returns +bf16(shared + routed); otherwise it returns the rank's routed partial, bit for bit what the stock fused MoE returns. + +Usage (build before CUDA-graph capture; every call is capture-safe afterwards):: + + state = M3MoeState( + device, num_tokens=6, shared=True + ) # workspace + compiled kernel, shared by all layers + layer = state.layer( + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared_weights=(gate_up_w, gate_up_sf, down_w, down_sf), + ) + out = layer(hidden) # bf16 [M, 6144] +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +from typing import Dict, Optional, Tuple + +import torch + +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size +# The routed intermediate per TP rank. +INTER = 384 # capacity: fixed: the routed experts' width per tp8 rank +NUM_EXPERTS = 128 +# The shared expert's intermediate per TP rank (the routed one). +SHARED_INTER = 384 # capacity: fixed: the shared expert's width per tp8 rank +# M3_moe's shared-expert CTAs (beside num_ctas routed CTAs). +SHARED_CTAS = 24 # capacity: fixed: an SM split, not a size +TOP_K = 4 +U1 = 2 * INTER // 32 # FC1 units per expert +# The router GEMV takes at most 2 experts per CTA (router weights held in registers). +MIN_CTAS = 64 # capacity: fixed: an SM split, not a size + +_UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) # one FC1 unit's weight + scale bytes +_ROW2_BYTES = INTER // 2 + INTER // 16 # one FC2 weight row's bytes (values + scales) +_S1 = 6 # FC1 stages (items) per unit # capacity: fixed: pipeline stages +_ITEM_BYTES = _UNIT_BYTES // _S1 # capacity: fixed: a ring item's bytes +_modules: Dict[tuple, object] = {} +_compiled: Dict[tuple, object] = {} +# m3_moe trace-time options by token count 1..8 (index num_tokens - 1). +# ``sfa_cp``: 18 KB ring stages, the A scales copied by cp.async into the instruction layout (every token count). +SFA_CP_BY_M: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1, 1) +# ``b2img``: FC1 -> FC2 hand-off through per-group global B images + release counters (7 and 8 tokens). +B2IMG_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 1, 1) +# ``ar_pre``: the finalize's all-reduce push with its call parities / descriptor preloaded (all but 2 tokens). +AR_PRE_BY_M: Tuple[int, ...] = (1, 0, 1, 1, 1, 1, 1, 1) +# Ring ``stages`` (0 = the kernel's default, the most that fit). 6 at 8 tokens measured 0.4-0.5 us faster than the +# default 7 (same MMAs, bit-identical outputs). +STAGES_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 0, 6) +# ``hb_repl``: copies of the tagged FC1 -> FC2 records (applied when > 1); CTA bx polls copy bx % hb_repl, which +# spreads the polls of the few hot record lines over L2. 7 and 8 tokens hand off through the B images instead. +HB_REPL_BY_M: Tuple[int, ...] = (32, 32, 8, 8, 1, 1, 1, 1) +# ``lg_repl``: replicas of the router logit exchange at up to LG_REPL_MAX_TOKENS tokens (fewer stores before the +# exchange; bit-identical outputs). Above that the kernel's default applies. +LG_REPL = 32 +LG_REPL_MAX_TOKENS = 4 # capacity: fixed: measured exchange-replica crossover +# Byte-balanced FC1 ranges below one unit per CTA (:func:`_fine_ranges`) for groups too small to give every CTA a +# whole unit; False gives such groups whole units on the CTAs with the least FC2 work. +FINE_PLAN = True +# m3_moe phase lag: phase p streams FC1 of group p and FC2 of group p - M3_MOE_LAG (the kernel's default lag). +M3_MOE_LAG = 2 + + +_ENGINE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe_kernel.py") + + +def _kernel_module(config: dict, path: str): + key = (os.path.basename(path),) + tuple(sorted(config.items())) + mod = _modules.get(key) + if mod is None: + name = ( + f"{__name__}_" + + os.path.splitext(os.path.basename(path))[0] + + "_" + + "_".join(f"{k}{v}" for k, v in key[1:]) + ) + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + mod.M3_CONFIG = dict(config) + sys.modules[name] = mod + spec.loader.exec_module(mod) + _modules[key] = mod + return mod + + +def _view(t: torch.Tensor, align: int): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align) + + +def is_supported(device: Optional[torch.device] = None) -> tuple: + if not torch.cuda.is_available(): + return False, "no CUDA" + major, _ = torch.cuda.get_device_capability(device) + if major != 10: + return False, "needs the sm_100 family (tcgen05 block-scaled NVFP4 MMA)" + try: + import cutlass # noqa: F401 + except ImportError: + return False, "CuTe DSL (nvidia-cutlass-dsl) is not installed" + return True, "" + + +def ar_push_desc( + mc_ptr: int, flags_ptr: int, rank: int, world: int, max_tokens: int, device +) -> torch.Tensor: + """The ``arp`` descriptor that makes a call also push its bf16 output into a multicast all-reduce exchange (set A + of m3_mc_allreduce's workspace at ``mc_ptr``, slot half = the all-reduce CTA's call parity in ``flags_ptr``). + With it the kernel releases its PDL dependents only after every push; the default (enable 0) keeps the old + behaviour.""" + return torch.tensor( + [1, mc_ptr, flags_ptr, rank, world, max_tokens, 0, 0], dtype=torch.int64, device=device + ) + + +def _shared_views(shared_weights, who): + """(gate_up, its scales, down, its scales) of the shared expert as flat uint8 views, checked.""" + if shared_weights is None: + raise ValueError(f"{who} (shared): shared_weights are required") + gu_w, gu_sf, dn_w, dn_sf = shared_weights + want = {"gate_up": (gu_w, (2 * SHARED_INTER, HIDDEN)), "down": (dn_w, (HIDDEN, SHARED_INTER))} + for name, (t, shape) in want.items(): + if t.dtype != torch.float8_e4m3fn or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError(f"{who} (shared): {name} weight must be contiguous e4m3 {shape}") + if ( + gu_sf.numel() != 2 * SHARED_INTER * HIDDEN // 32 + or dn_sf.numel() != HIDDEN * SHARED_INTER // 32 + ): + raise ValueError( + f"{who} (shared): the UE8M0 scales must be the swizzled [N, K / 32] buffers" + ) + return ( + gu_w.view(torch.uint8).view(-1), + gu_sf.view(torch.uint8).reshape(-1), + dn_w.view(torch.uint8).view(-1), + dn_sf.view(torch.uint8).reshape(-1), + ) + + +def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24) -> torch.Tensor: + """The shared expert's weights as m3_moe's shared CTAs stream them (SH_PRETILE): per CTA s, per stage g, the 16 KB + shared-memory image (128-byte swizzle: 16-byte chunk j of row r at j ^ (r % 8)) of the stage's 128-row MMA tile and + its 512-byte instruction scale atom (row R's 4 codes at word 4 (R % 32) + R / 32). CTA s owns P = SHARED_INTER / 24 + = 16 positions in SPLIT = 128 / (2 P) = 4 K slices of GT = 48 / SPLIT = 12 k-tiles. Stages 0..GT-1: MMA row + R = P O + i (O = 2 q + u: slice q, u 0 gate / 1 up) = weight row u SHARED_INTER + P s + i at k-tile GT q + g; then + down tile j (rows 256 s + 128 j ..) at k-tile kt, j-major (SHARED_INTER / 128 k-tiles each). Built once per + layer.""" + dev = gu8.device + H_, I_ = HIDDEN, SHARED_INTER + P = I_ // n_cta + split = 128 // (2 * P) + gt = H_ // 128 // split + dkt = I_ // 128 + gu = gu8.view(2 * I_, H_) + dn = dn8.view(H_, I_) + gkb, dkb = H_ // 32, I_ // 32 + + def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast + return sf[ + ((n >> 7) * (kbn // 4) + (b >> 2)) * 512 + (n & 31) * 16 + ((n >> 5) & 3) * 4 + (b & 3) + ] + + r = torch.arange(128, device=dev) + q = r >> 5 + ri = r & 31 + octant = r // P + chunk = torch.arange(8, device=dev) + swz = chunk[None, :] ^ (r[:, None] & 7) # source chunk of image chunk j: j ^ (r % 8) + out = torch.empty( + n_cta, gt + 2 * dkt, 16384 + 512, dtype=torch.uint8, device=dev + ) # capacity: fixed: a workspace of the kernel's fixed geometry + for s in range(n_cta): + # gate_up: MMA row R -> weight row, slice + nrow = (octant & 1) * I_ + P * s + r % P + sl = octant >> 1 + for g in range(gt): + kt = sl * gt + g # per row + tile = gu[ + nrow[:, None], (kt[:, None] * 128 + torch.arange(128, device=dev)[None, :]) + ] # [128, 128] + img = tile.view(128, 8, 16)[r[:, None], swz] # image chunk j = tile chunk j ^ (r % 8) + out[s, g, :16384] = img.reshape(-1) + kb = torch.arange(4, device=dev) + codes = sf_bytes(gusf8, nrow[:, None], kt[:, None] * 4 + kb[None, :], gkb) # [128, 4] + atom = torch.empty(512, dtype=torch.uint8, device=dev) + atom[(ri[:, None] * 16 + q[:, None] * 4 + kb[None, :]).reshape(-1)] = codes.reshape(-1) + out[s, g, 16384:] = atom + for gd in range(2 * dkt): + jt, kt = gd // dkt, gd % dkt + row0 = 256 * s + 128 * jt + tile = dn[row0 : row0 + 128, kt * 128 : (kt + 1) * 128] + img = tile.reshape(128, 8, 16)[r[:, None], swz] + out[s, gt + gd, :16384] = img.reshape(-1) + out[s, gt + gd, 16384:] = dnsf8[ + ((row0 >> 7) * (dkb // 4) + kt) * 512 : ((row0 >> 7) * (dkb // 4) + kt) * 512 + 512 + ] + return out.view(-1) + + +# --------------------------------------------------------------------------------------------------------------------- +# m3_moe (m3_moe_kernel.py): 1..8 tokens per call, group-phased weight stream. +# --------------------------------------------------------------------------------------------------------------------- +# Pre-tiled shared weights per layer, keyed by the source buffers (which the engines read in place): every token-count +# state of a layer streams the same copy. An entry holds its sources, so their addresses are not reused while cached. +_PRETILED: Dict[tuple, tuple] = {} + + +def _pretiled(gu8, gusf8, dn8, dnsf8) -> torch.Tensor: + key = (gu8.device, gu8.data_ptr(), gusf8.data_ptr(), dn8.data_ptr(), dnsf8.data_ptr()) + hit = _PRETILED.get(key) + if hit is None: + hit = _PRETILED[key] = ((gu8, gusf8, dn8, dnsf8), _pretile_shared(gu8, gusf8, dn8, dnsf8)) + return hit[1] + + +def _fine_ranges(weights, total: int): + """Contiguous FC1 item ranges [s_c, e_c) over the CTAs, sized as close to ``weights`` (ideal real-valued item + counts summing to ``total``) as the split-unit protocol allows: every unit (S1 items) has at most one CTA boundary + strictly inside it (its head and tail parts, summed by the tail's holder), so a range that starts inside a unit + runs at least to that unit's end. Returns the item counts.""" + n = len(weights) + items = [0] * n + s = 0 + acc = 0.0 + for c in range(n): + acc += weights[c] + e = total if c == n - 1 else min(total, max(s, int(acc + 0.5))) + if s % _S1 and e < total: # started inside a unit: run to (at least) its end + e = max(e, min(total, (s // _S1 + 1) * _S1)) + items[c] = e - s + s = e + assert s == total + return items + + +def make_phase_plan( + num_ctas: int, + m_max: int, + lag: int = 2, + tail_w: float = 1.0, + early0: bool = False, + fine: bool = FINE_PLAN, +): + """The work split of ``m3_moe`` for each distinct-expert count G (4 .. 4 m_max): FC2 32-row output blocks in + contiguous per-CTA ranges (the same for every group), and per group p (slots 4 p ..) a range of its unit-major FC1 + stages per CTA, balancing the bytes each CTA streams in phase p (FC1 of group p plus FC2 of group p - lag, and the + trailing FC2s with the last group). Ranges are empty or >= 6 stages (a unit is split over at most two CTAs); a + group too small for every CTA gets whole units on the CTAs with the least FC2 work, or (``fine``) byte-balanced + ranges of any length with at most one CTA boundary inside a unit (:func:`_fine_ranges`). Returns int32 + [4 m_max + 1, num_ctas, m_max + 1] (per group i0 | i1 << 16, then b0 | nb << 16) and the largest block count.""" + gcap = TOP_K * m_max + nb2 = HIDDEN // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + blk_bytes = 32 * _ROW2_BYTES # one block, one slot + for g in range(TOP_K, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) + for p in range(ng): + n_p = min(4, g - 4 * p) + total = n_p * U1 * _S1 + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * min(4, g - 4 * (p - lag)) * blk_bytes + if p == ng - 1 and not ( + early0 and p == 0 and m_max >= 2 + ): # early0: group 0's ranges the same for every G + for q in range(max(0, ng - lag), ng): + fc2[c] += int(tail_w * nbs[c] * min(4, g - 4 * q) * blk_bytes) + items = [0] * num_ctas + if total >= _S1 * num_ctas: + + def counts_for(limit): + return [max(_S1, (limit - f) // _ITEM_BYTES) for f in fc2] + + lo, hi = 0, max(fc2) + total * _ITEM_BYTES + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > _S1), + key=lambda c: (-(fc2[c] + items[c] * _ITEM_BYTES), c), + ) + items[c] -= 1 + extra -= 1 + elif fine: + # water-fill: CTA c streams fc2[c] + w_c item bytes, equal where w_c > 0 + lo_f, hi_f = 0.0, float(max(fc2) + total * _ITEM_BYTES) + for _ in range(100): + mid = 0.5 * (lo_f + hi_f) + if sum(max(0.0, (mid - f) / _ITEM_BYTES) for f in fc2) >= total: + hi_f = mid + else: + lo_f = mid + wts = [max(0.0, (hi_f - f) / _ITEM_BYTES) for f in fc2] + sw = sum(wts) + items = _fine_ranges([w * total / sw for w in wts], total) + else: + chosen = sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // _S1]) + for c in chosen: + items[c] = _S1 + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | ((i0 + items[c]) << 16) + i0 += items[c] + assert i0 == total + return plan, max(nbs) + + +class M3MoeState: + """Workspace + compiled kernel of ``m3_moe`` (1..8 tokens per call) on one device, shared by the layers that run + on it in one stream order. The call takes the bf16 hidden states and computes the router logits (bit-exact with + m3_router_gemv_fp4) and the NVFP4 activation itself. ``shared``: the shared expert runs on SHARED_CTAS dedicated + CTAs beside the ``num_ctas`` routed ones and the call returns bf16(shared + routed). Build before CUDA-graph + capture; the first call compiles.""" + + fused_router = True + + def __init__( + self, + device: torch.device, + num_tokens: int, + num_ctas: Optional[int] = None, + shared: bool = False, + inter: int = INTER, + ): + """``inter``: the routed intermediate per rank; this engine is built for INTER (384) only.""" + if inter != INTER: + raise ValueError( + f"m3_moe (TP8): routed intermediate {inter} per rank is not supported (only {INTER})" + ) + self.inter = inter + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "M3MoeState allocates and compiles: build it before CUDA-graph capture" + ) + sms_all = torch.cuda.get_device_properties(device).multi_processor_count + if shared and (num_ctas or sms_all) + SHARED_CTAS > sms_all: + raise ValueError( + f"m3_moe (shared): num_ctas + {SHARED_CTAS} shared-expert CTAs must fit the {sms_all} SMs " + f"(all resident at once)" + ) + if not 1 <= num_tokens <= 8: + raise ValueError(f"m3_moe takes 1..8 tokens per call, not {num_tokens}") + self.device = torch.device(device) + ok, why = is_supported(self.device) + if not ok: + raise RuntimeError(f"m3_moe: {why}") + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.num_ctas = num_ctas or sms + if not MIN_CTAS <= self.num_ctas <= sms: + raise ValueError( + f"m3_moe runs {MIN_CTAS} .. {sms} CTAs (all resident at once), not {self.num_ctas}" + ) + self.m = num_tokens + # One token: the trailing FC2s weigh nothing in the last phase's balance. Group 0's FC1 ranges are the same + # for every distinct-expert count (early0). + plan, nb_cap = make_phase_plan( + self.num_ctas, num_tokens, M3_MOE_LAG, 0.0 if num_tokens == 1 else 1.0, True + ) + self.config = { + "m_max": num_tokens, + "num_ctas": self.num_ctas, + "nb2_cap": nb_cap, + "fused_router": 1, + "timing": 0, + "i_tp": INTER, + } + if HB_REPL_BY_M[num_tokens - 1] > 1: + self.config["hb_repl"] = HB_REPL_BY_M[num_tokens - 1] + if num_tokens <= LG_REPL_MAX_TOKENS: + self.config["lg_repl"] = LG_REPL + if shared: + self.config["sh_i"] = SHARED_INTER + self.config["shared"] = 1 + self.shared = shared + # 18 KB ring stages (A scales by cp.async into the instruction layout) per token count, unless configured + self.config["sfa_cp"] = SFA_CP_BY_M[num_tokens - 1] + # FC1 -> FC2 hand-off through per-group global B images + release counters, per token count + self.config["b2img"] = B2IMG_BY_M[num_tokens - 1] + # the finalize's all-reduce push with its call parities / descriptor preloaded, per token count + self.config["ar_pre"] = AR_PRE_BY_M[num_tokens - 1] + if STAGES_BY_M[num_tokens - 1] > 0: + self.config["stages"] = STAGES_BY_M[num_tokens - 1] + try: + self.mod = mod = _kernel_module(self.config, _ENGINE_PATH) + except AssertionError as err: + raise ValueError( + f"m3_moe: num_tokens={num_tokens}, num_ctas={self.num_ctas} does not fit the shared " + f"memory / TMEM budget (fewer CTAs means more FC2 blocks per CTA): {err}" + ) from err + kw = dict(device=self.device) + self.plan = plan.to(self.device) + self.partials = torch.full( + (mod.G_CAP * mod.U1 * num_tokens * 32,), -1, dtype=torch.int32, **kw + ) + # the FC1 -> FC2 records; with HB_REPL > 1 copies, HB_STRIDE int32 words apart + hb_words = mod.G_CAP * num_tokens * mod.U1 * 4 + if mod.HB_REPL > 1: + hb_words = mod.HB_REPL * mod.HB_STRIDE + self.hbuf32 = torch.zeros(hb_words, dtype=torch.int32, **kw) + self.epochs = torch.zeros(self.num_ctas + mod.NSH, dtype=torch.int32, **kw) + self.route_dbg = torch.zeros(num_tokens, 8, dtype=torch.int32, **kw) + self.tdbg = torch.zeros(self.num_ctas + mod.NSH, 64, dtype=torch.int64, **kw) + self.lgbuf = torch.zeros(mod.LG_REPL * mod.LG_STRIDE, dtype=torch.int64, **kw) + self._dummy = torch.zeros(4, dtype=torch.int32, **kw) + self._ar_off = torch.zeros(8, dtype=torch.int64, **kw) + # B2IMG: per-group FC2 B images and their unit counters (int32; every group gains CNT_PER_CALL per call) + self.img32 = torch.zeros(max(1, mod.NG_CAP * mod.IMG_BYTES // 4), dtype=torch.int32, **kw) + self.cnt = torch.zeros(mod.NG_CAP * 8 + 8, dtype=torch.int32, **kw) + # shared expert: h and y_sh as (call tag << 16 | bf16), tag 0 = never written + self.hx = torch.zeros(num_tokens * SHARED_INTER, dtype=torch.int32, **kw) + self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) + self._fn = None + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> "M3MoeLayer": + """One MoE layer's operands (the TRTLLM-Gen NVFP4 expert buffers, read in place, never copied). + ``shared_weights`` (a ``shared`` state): the shared expert's (gate_up e4m3 [768, 6144], its swizzled UE8M0 + scales, down e4m3 [6144, 384], its swizzled UE8M0 scales).""" + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + ) + if self.shared: + layer.shared_weights = _shared_views(shared_weights, "m3_moe") + layer.shared_mxfp8 = tuple(shared_weights) + if self.mod.SH_PRETILE: + # (gate_up, pre-tiled stages, down, dummy): the kernel streams the pre-tiled copy + gu8, _, dn8, _ = layer.shared_weights + layer.shared_weights = ( + gu8, + _pretiled(*layer.shared_weights), + dn8, + self._dummy.view(torch.uint8), + ) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError( + f"m3_moe (fused router): expects contiguous bf16 hidden states [{m}, {HIDDEN}]" + ) + fz = ( + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + ) + lgv, xq, xs = self._dummy.view(torch.float32), self._dummy, self._dummy + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + sh = layer.shared_weights if self.shared else (self._dummy.view(torch.uint8),) * 4 + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + lgv, + layer.bias, + xq, + xs, + layer.escal, + self.plan.view(-1), + self.partials, + self.hbuf32, + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + *fz, + *sh, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + layer.w3_w1_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + layer.w2_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + self.img32, + self.cnt, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_moe compiles on its first call: call it once before capture") + import cutlass.cute as cute + + aligns = [ + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + ] + sig = [_view(t, a) for t, a in zip(args, aligns)] + key = (_ENGINE_PATH,) + tuple(sorted(self.config.items())) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out + + +class M3MoeLayer: + """One MoE layer on a :class:`M3MoeState`: ``layer(hidden)``. The expert buffers are read in + place, never copied; the per-expert scalars are packed once (layer constants: rebuild the layer after reloading + weights).""" + + shared_weights = None # the kernel's view of the shared expert (a shared state) + shared_mxfp8 = None # the shared expert's (gate_up, its scales, down, its scales) as given + + def __init__( + self, + state, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + bias, + router_weight=None, + fc31_input_scale=None, + ): + e = NUM_EXPERTS + expected = { + "w3_w1_weight": (w3_w1_weight, (e, 2 * INTER, HIDDEN // 2)), + "w3_w1_weight_scale": (w3_w1_weight_scale, (e, 2 * INTER, HIDDEN // 16)), + "w2_weight": (w2_weight, (e, HIDDEN, INTER // 2)), + "w2_weight_scale": (w2_weight_scale, (e, HIDDEN, INTER // 16)), + } + for name, (t, shape) in expected.items(): + if t.element_size() != 1 or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError( + f"m3_moe: {name} is {t.dtype} {tuple(t.shape)}, expected contiguous 1-byte {shape}" + ) + self.state = state + self.w3_w1_weight = w3_w1_weight.view(torch.uint8) + self.w3_w1_weight_scale = w3_w1_weight_scale.view(torch.uint8) + self.w2_weight = w2_weight.view(torch.uint8) + self.w2_weight_scale = w2_weight_scale.view(torch.uint8) + self.escal = torch.stack( + [fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha] + ).float() + self.escal = self.escal.reshape(6, e).contiguous() + self.bias = bias.float().contiguous() + if ( + router_weight is None + or fc31_input_scale is None + or router_weight.dtype != torch.float32 + or tuple(router_weight.shape) != (e, HIDDEN) + or not router_weight.is_contiguous() + or fc31_input_scale.dtype != torch.float32 + or fc31_input_scale.numel() != 1 + ): + raise ValueError( + "m3_moe (fused router): needs router_weight f32 [128, 6144] and fc31_input_scale f32 [1]" + ) + self.router_weight = router_weight + self.fc31_input_scale = fc31_input_scale + + def __call__( + self, + hidden: torch.Tensor, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``hidden``: contiguous bf16 [M, 6144]. ``arp`` (int64 [8], see :func:`ar_push_desc`): also push the output + into an all-reduce exchange.""" + return self.state.run(self, hidden, out, arp=arp) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_prequant.py new file mode 100644 index 000000000000..48a7827f59b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_prequant.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hand-off of an MXFP8-quantized activation from the kernel that produced it to the next MXFP8 GEMM. + +A producer that already holds a BF16 row in registers (the multicast all-reduce + RMSNorm, the split-K attention +merge) also writes its MXFP8 form (E4M3 rows + UE8M0 scales in the 128x4-swizzled layout of +``trtllm::mxfp8_quantize(x, True)``) and publishes it here keyed by the BF16 tensor. The MXFP8 linear whose input +is that tensor takes it and skips its own quantize kernel. Every take clears the slot, so a hand-off is only ever +consumed by the first MXFP8 linear that runs after the producer. + +Module state shared across files (one instance each per process): + +* ``SLOT`` (PrequantSlot): published by ``m3_mc_allreduce`` (``quant=True``), the sparse-decode merge and the fused + decode attention (MXFP8 output); cleared by the sparse decode before it runs; taken by the MXFP8 linear dispatch + and by the one-launch QKV + attention + o_proj path (which also peeks ``pending``). +* ``AR_PUSH`` (ArPushSlot): armed by the fast forward around the attention call; the o_proj GEMV takes the target + and marks ``pushed``; the fast forward checks ``take_pushed`` before the attention boundary all-reduce. +* ``MOE_PUSH`` (MoePushSlot): armed by the fast forward before the MoE; the routed-expert engine takes ``engine`` + and the shared expert's down projection takes ``shared``, each marks its pushed output; the fast forward + checks ``take_pushed`` before the MoE boundary all-reduce. +""" + + +class PrequantSlot: + def __init__(self): + self.pending = None # (data_ptr, rows, width, e4m3 [rows, width], scales) + + def publish(self, x, xq, xsf): + self.pending = (x.data_ptr(), x.shape[0], x.shape[-1], xq, xsf) + + def clear(self): + self.pending = None + + def take(self, x): + """(e4m3 [x.shape[0], width], scales) if x is the published tensor, else None.""" + pending, self.pending = self.pending, None + if ( + pending is None + or x.dim() != 2 + or x.data_ptr() != pending[0] + or x.shape[0] != pending[1] + or x.shape[1] != pending[2] + ): + return None + return pending[3], pending[4] + + +SLOT = PrequantSlot() + + +class ArPushSlot: + """The attention boundary's all-reduce exchange for the o_proj GEMV to push into. + + The fast forward arms ``target`` (m3_mx_gemv ``ar`` tuple) around the attention call; the o_proj marks the output + it pushed in ``pushed``; the boundary then runs the all-reduce with ``prepushed`` = the ``ar`` tuple's mode.""" + + def __init__(self): + self.target = None + self.pushed = None # data_ptr of the pushed output + + def arm(self, target): + self.target = target + self.pushed = None + + def take_target(self): + target, self.target = self.target, None + return target + + def take_pushed(self, x): + pushed, self.pushed = self.pushed, None + return pushed is not None and x.data_ptr() == pushed + + +AR_PUSH = ArPushSlot() + + +class MoePushSlot: + """The MoE boundary's all-reduce exchange for its two producers to push into: the routed-expert engine (set A, + ``engine`` = its ``arp`` descriptor) and the shared expert's down projection (set B, ``shared`` = m3_mx_gemv + ``ar`` tuple). Each marks the output it pushed; the boundary then runs the all-reduce with ``prepushed`` = 3.""" + + def __init__(self): + self.engine = None + self.shared = None + self.engine_pushed = None # data_ptr of the pushed routed output + self.shared_pushed = None # data_ptr of the pushed shared output + + def arm(self, engine, shared): + self.engine, self.shared = engine, shared + self.engine_pushed = self.shared_pushed = None + + def disarm(self): + self.engine = self.shared = None + + def take_engine(self): + engine, self.engine = self.engine, None + return engine + + def take_shared(self): + shared, self.shared = self.shared, None + return shared + + def take_pushed(self, shared, routed): + """(routed pushed, shared pushed) for these outputs; clears the marks.""" + e = routed is not None and self.engine_pushed == routed.data_ptr() + s = shared is not None and self.shared_pushed == shared.data_ptr() + self.engine_pushed = self.shared_pushed = None + return e, s + + +MOE_PUSH = MoePushSlot() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_router_gemv.py new file mode 100644 index 000000000000..30855eb7f05d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_router_gemv.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-M fp32 router projection: ``logits[M, E] = float(x[M, K]) @ W[E, K]^T`` in one kernel. + +Replaces the bf16->fp32 cast, the TF32 GEMM and its split-K reduce on decode +steps. One CTA per ``ROWS`` experts: the fp32 weight rows are bulk-copied into +shared memory before ``griddepcontrol.wait`` (overlapping the producer kernel +under PDL), the bf16 activations after it. Products and sums are exact fp32 FMA. + +The kernels are ``torch.ops.trtllm.minimax_m3_router_gemv`` / ``_router_gemv_fp4`` / ``_router_tc``, reached +through the ``ops/moe`` entries. +""" + +from typing import Optional + +import torch + +from ..ops.moe.m3_router_gemv import m3_router_gemv as _router_gemv_op +from ..ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 as _router_gemv_fp4_op +from ..ops.moe.m3_router_tc import m3_router_tc as _router_tc_op + +# Opt-in dynamic shared memory left after the kernel's static reduction buffer (SM100: 227 KiB). +_MAX_DYNAMIC_SMEM = 227 * 1024 - 1024 # capacity: fixed: the device's shared-memory limit + + +def m3_router_gemv( + x: torch.Tensor, weight: torch.Tensor, rows: int = 1, use_pdl: bool = False +) -> Optional[torch.Tensor]: + """fp32 ``x @ weight.T`` for bf16 ``x`` with at most 16 rows; None when unsupported (only ``rows=1`` is built).""" + if ( + rows != 1 + or not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 16 + or weight.dtype != torch.float32 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or weight.shape[0] % rows != 0 + or x.data_ptr() % 16 != 0 + or (x.stride(0) * 2) % 16 != 0 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or rows * x.shape[1] * 4 + x.shape[0] * x.shape[1] * 2 > _MAX_DYNAMIC_SMEM + ): + return None + return _router_gemv_op(x, weight, rows, use_pdl) + + +def m3_router_gemv_fp4( + x: torch.Tensor, weight: torch.Tensor, sf_scale: torch.Tensor, use_pdl: bool = False +) -> Optional[tuple]: + """Router logits plus the NVFP4 activation the routed experts consume (bit-identical to + ``fp4_quantize(x, sf_scale, 16, False, False)``): (logits [M, E] fp32, x_fp4 [M, K / 2] uint8, + x_sf [M * K / 16] uint8 linear). None when unsupported.""" + if ( + not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 16 + or weight.dtype != torch.float32 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or x.data_ptr() % 16 != 0 + or (x.stride(0) * 2) % 16 != 0 + or sf_scale is None + or sf_scale.dtype != torch.float32 + or sf_scale.numel() != 1 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or x.shape[1] * 4 + x.shape[0] * x.shape[1] * 2 > _MAX_DYNAMIC_SMEM + ): + return None + return tuple(_router_gemv_fp4_op(x, weight, sf_scale.reshape(1), use_pdl)) + + +def m3_router_tc( + x: torch.Tensor, + weight: torch.Tensor, + sf_scale: Optional[torch.Tensor] = None, + use_pdl: bool = False, +) -> Optional[tuple]: + """Router logits (TF32 tensor cores, like the native TF32 GEMM) for 1..48 tokens with 128 experts, plus the + NVFP4 activation when ``sf_scale`` is given. Returns (logits,) or (logits, x_fp4, x_sf); None when unsupported.""" + if ( + not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 48 + or weight.dtype != torch.float32 + or weight.shape[0] != 128 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or (x.stride(0) * 2) % 16 != 0 + or x.data_ptr() % 16 != 0 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or 16 * (x.shape[1] // 8 + 4) * 4 + x.shape[0] * (x.shape[1] // 8 + 8) * 2 + > _MAX_DYNAMIC_SMEM + or (sf_scale is not None and (sf_scale.dtype != torch.float32 or sf_scale.numel() != 1)) + ): + return None + res = _router_tc_op(x, weight, None if sf_scale is None else sf_scale.reshape(1), use_pdl) + return tuple(res) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_sampled_accept.py new file mode 100644 index 000000000000..86b08a0b24fc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_sampled_accept.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). + +Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the +strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p +active must stay on the generic path; the caller decides. + +Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL +vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z +(the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it +is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens heavier +than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the nucleus +with probability >= top_p, so most rows decide on their second exchange. + +Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, +staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local +candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the +candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based Philox +keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one 4-word +record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and reads +the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the same +order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing and +two call halves keep a rank one call ahead from overwriting a slot still being read. +""" + +from typing import Optional + +import torch + +from ..ops.spec.m3_sampled_accept import m3_sampled_accept as _sampled_accept_op +from ..ops.spec.m3_sampled_accept_ws_words import m3_sampled_accept_ws_words as _ws_words_op + +# The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.cu) and reached through +# the ops/spec entry m3_sampled_accept. +MAX_ROWS = 48 # capacity: fixed: the kernel's built rows (decode steps) + + +class SampledAcceptWorkspace: + """Exchange slots behind one multicast mapping (two call halves), the (finished rows, call) counters, the per-row + token scratch and an extra-attempt counter. Created collectively, outside CUDA-graph capture.""" + + def __init__(self, uc, mc_ptr, rank, world, handle=None, comm=None): + self.uc, self.mc_ptr, self.rank, self.world = uc, mc_ptr, rank, world + self.handle, self.comm = handle, comm + dev = uc.device + self.counters = torch.zeros(2, dtype=torch.int32, device=dev) + self.tokens = torch.zeros(MAX_ROWS, dtype=torch.int32, device=dev) + self.extra_attempts = torch.zeros(1, dtype=torch.int32, device=dev) + + @classmethod + def create(cls, mapping, device) -> "SampledAcceptWorkspace": + words = int(_ws_words_op(max(mapping.tp_size, 1))) + if mapping.tp_size == 1: + uc = torch.zeros(words, dtype=torch.int64, device=device) + return cls(uc, uc.data_ptr(), 0, 1) + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_disabled + + comm = _get_mnnvl_workspace_comm(mapping) + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + if not mpi_disabled(): + comm.Free() + raise RuntimeError( + "SampledAcceptWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 8, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("SampledAcceptWorkspace.create failed on another rank") + return cls(uc, mc.data_ptr(), mapping.tp_rank, mapping.tp_size, handle, comm) + + +_WS: Optional[SampledAcceptWorkspace] = None + + +def workspace(mapping, device) -> Optional[SampledAcceptWorkspace]: + global _WS + if _WS is None: + if torch.cuda.is_current_stream_capturing(): + # A capture that finds it missing is the caller's miss (it is built on the eager steps). + return None + _WS = SampledAcceptWorkspace.create(mapping, device) + return _WS + + +def sampled_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + temperatures: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: Optional[torch.Tensor], + seeds: torch.Tensor, + offsets: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, + ws: Optional[SampledAcceptWorkspace] = None, + use_pdl: bool = True, +): + """Sampled verification (temperature + top-p) from this rank's bf16 logits shard [nreq * (draft_len + 1), V / tp]: + outputs as shard_accept (accepted [nreq, D+1], num_accepted [nreq], gather_ids [nreq]); None + when unavailable. Per-row sampling parameters and Philox (seed, offset) are the spec metadata's, one per row. + top_ks / min_ps are not applied: rows with top-k or min-p active must be routed to the generic path by the + caller.""" + del top_ks, min_ps + nreq, draft_len = draft_tokens.shape + rows = nreq * (draft_len + 1) + if logits.dim() != 2 or logits.shape[0] != rows or rows > MAX_ROWS or logits.shape[1] > 51200: + return None + if logits.dtype != torch.bfloat16 or logits.stride(1) != 1: + return None + ws = ws if ws is not None else workspace(mapping, logits.device) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + _sampled_accept_op( + logits, + int(vocab_offset), + temperatures[:rows].float().contiguous(), + top_ps[:rows].float().contiguous(), + seeds[:rows].long().contiguous(), + offsets[:rows].long().contiguous(), + ws.uc, + ws.mc_ptr, + ws.counters, + ws.tokens, + ws.rank, + ws.world, + draft_tokens.int().contiguous(), + batch_idx.int().contiguous(), + accepted, + num_accepted, + gather_ids, + int(force_total), + next_new, + force_vec, + ws.extra_attempts, + bool(use_pdl), + ) + return accepted, num_accepted, gather_ids diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_shard_accept.py new file mode 100644 index 000000000000..21521591bfa9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_shard_accept.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused greedy Eagle3 verification for the MiniMax-M3 target head (vocab-sharded TP head, verify rows <= 48). + +Replaces the all-gather of the target logits, the argmax and the strict-acceptance ops of a greedy verify step with +one kernel over this rank's BF16 logits shard (the target lm_head's own local GEMM output): + +* Every logit becomes a 64-bit key ``ordered(value) << 32 | ~global_index`` (ties pick the lowest index, as + ``torch.argmax``), max-reduced per row: warp shuffles, then ``atomicMax``. +* The last CTA to finish pushes the rank's per-row key into every rank's exchange slot through a multicast mapping + (Lamport: 0 = empty, the key is never 0), polls its own slots for all ranks, takes the global argmax, re-arms the + slots and runs the strict acceptance. Call n uses slot half n % 2. + +The argmax is taken over the BF16 logits without arithmetic, so the tokens are those of the in-tree path's argmax +over the gathered logits. The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/ +shardAccept.cu) and reached through the ``ops/spec`` entry ``m3_shard_accept``. +""" + +from dataclasses import dataclass +from typing import Any, Optional + +import torch + +from ..ops.spec.m3_shard_accept import m3_shard_accept as _shard_accept_op + +#: Verify rows per call, and exchange rows per call half of the workspace. +SHARD_ROWS = 48 # capacity: fixed: DECODE_MAX_TOKENS verify rows + + +@dataclass(eq=False) +class AcceptWorkspace: + """Per-rank exchange slots ([2 halves][SHARD_ROWS rows][world] uint64) behind one multicast mapping, the local + key accumulators and the (finished CTAs, call) counters.""" + + uc: torch.Tensor + mc_ptr: int + keys: torch.Tensor + counters: torch.Tensor + rank: int + world_size: int + handle: Any = None + comm: Any = None + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def create_workspace(mapping, device, slot_rows: int) -> AcceptWorkspace: + """Exchange slots ([2 halves][slot_rows rows][world] uint64) behind one multicast mapping over the TP group, the + local key accumulators [slot_rows] and the (finished CTAs, call) counters. Collective over the TP group; call + outside CUDA-graph capture, in the same order on every rank.""" + words = 2 * slot_rows * max(mapping.tp_size, 1) + if mapping.tp_size > 1: + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_disabled + + comm = _get_mnnvl_workspace_comm(mapping) + if not _mnnvl_workspace_all_succeeded(comm, True): + if not mpi_disabled(): + comm.Free() + raise RuntimeError("exchange workspace creation failed on another rank") + handle = _make_mnnvl_mcast_buffer( + comm, max(words * 8, 4096), mapping, mapping.is_multi_node() + ) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("exchange workspace creation failed on another rank") + mc_ptr, rank = mc.data_ptr(), mapping.tp_rank + else: + uc = torch.zeros(words, dtype=torch.int64, device=device) + mc_ptr, rank, handle, comm = uc.data_ptr(), 0, None, None + return AcceptWorkspace( + uc=uc, + mc_ptr=mc_ptr, + keys=torch.zeros(slot_rows, dtype=torch.int64, device=uc.device), + counters=torch.zeros(2, dtype=torch.int32, device=uc.device), + rank=rank, + world_size=max(mapping.tp_size, 1), + handle=handle, + comm=comm, + ) + + +def _workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace, created collectively outside CUDA-graph capture on first use.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + # A capture that finds it missing is the caller's miss (it is built on the eager steps). + return None + _WS = create_workspace(mapping, device, SHARD_ROWS) + return _WS + + +def shard_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, +): + """Greedy verification from this rank's BF16 logits shard. + + Args: + logits: [nreq * (draft_len + 1) <= SHARD_ROWS, V / tp] BF16, unit column stride: this rank's columns + (global ids ``vocab_offset ..``) of the verify rows, request-major. + draft_tokens: [nreq, draft_len] int32, the drafts being verified. + batch_idx: [nreq] int32, each request's row in the verify layout. + force_total / force_vec: forced acceptance (see ``m3_shard_accept``). + next_new: optional int32 [nreq, >= 1]; column 0 receives the last accepted token. + + Returns: + (accepted_tokens [nreq, draft_len + 1] int32, num_accepted [nreq] int32, first-pass gather ids [nreq] + int64), or None when the shape is not served or the workspace is missing during a capture. Collective over + the TP group. + """ + nreq, draft_len = draft_tokens.shape + if ( + logits.dim() != 2 + or logits.shape[0] != nreq * (draft_len + 1) + or logits.shape[0] > SHARD_ROWS + or logits.dtype != torch.bfloat16 + or logits.stride(1) != 1 + ): + return None + ws = _workspace(mapping, logits.device) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + _shard_accept_op( + logits, + vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_tokens.contiguous(), + batch_idx, + accepted, + num_accepted, + gather_ids, + force_total, + next_new, + _num_sms(), + force_vec, + ) + return accepted, num_accepted, gather_ids diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_spec_tail.py new file mode 100644 index 000000000000..0d4dc9cd1b64 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_spec_tail.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Small fused kernels of the MiniMax-M3 Eagle3 step (GB300, decode). + +``chunk_norm`` / ``capture_norm`` apply the drafter's per-layer fc_norm RMSNorms to the captured target hidden states +(the in-tree ``apply_eagle3_fc`` norms, same fp32 arithmetic); ``draft_input`` is a fused draft pass's input (the +embedding lookup, the drafter layer's input_layernorm / hidden_norm and their concat, same arithmetic); +``step_gather`` / ``slot_scatter`` move the step's integer outputs between the spec sampler's store and the engine's +input buffers. +""" + +from typing import Optional, Tuple + +import torch + +# The kernels are compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu) and reached through the +# ops/spec entries. +from ..ops.spec.m3_capture_norm import m3_capture_norm as _capture_norm_op +from ..ops.spec.m3_chunk_norm import m3_chunk_norm as _chunk_norm_op +from ..ops.spec.m3_draft_input import m3_draft_input as _draft_input_op +from ..ops.spec.m3_slot_scatter import m3_slot_scatter as _slot_scatter_op +from ..ops.spec.m3_step_gather import m3_step_gather as _step_gather_op + + +def draft_input( + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + eps_embed: float, + w_hidden: torch.Tensor, + eps_hidden: float, + use_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """(cat([rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden], -1) [M, 2H], hidden[rows] [M, H]) for + int32 ids [M] and int64 rows [M] (rows None: hidden[:M]); ids outside the table embed to zero.""" + m, h = ids.shape[0], table.shape[1] + out = torch.empty((m, 2 * h), dtype=table.dtype, device=table.device) + residual = torch.empty((m, h), dtype=table.dtype, device=table.device) + return _draft_input_op( + ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, use_pdl + ) + + +def chunk_norm( + src: torch.Tensor, weights: torch.Tensor, eps: float, use_pdl: bool = True +) -> torch.Tensor: + """Per-chunk RMSNorm of src [M, >= C*H] (row stride may exceed C*H) with weights [C, H] -> [M, C*H].""" + c, h = weights.shape + out = torch.empty((src.shape[0], c * h), dtype=src.dtype, device=src.device) + return _chunk_norm_op(src, weights, eps, out, use_pdl) + + +# The overlap scheduler's per-step gathers and the spec sampler's slot stores as one kernel each (step_gather / +# slot_scatter) instead of ~10 and 4 small torch ops; called by the target's SpecStepFold. + + +def step_gather( + new_tok, next_draft, lens, slots, ppi, P, T, D, ids, ids0, draft, d0, pos_off, kv_off, b0 +) -> bool: + """See step_gather_kernel; False when the inputs are outside what the kernel covers (caller falls back).""" + ts = (new_tok, next_draft, lens, slots, ppi, ids, draft, pos_off, kv_off) + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in ts) + or next_draft.dim() != 2 + or new_tok.dim() < 2 + ): + return False + if (new_tok.dim() == 3 and new_tok.shape[2] != 1) or not ( + lens.is_contiguous() + and ids.is_contiguous() + and draft.is_contiguous() + and pos_off.is_contiguous() + and kv_off.is_contiguous() + ): + return False + _step_gather_op( + new_tok, next_draft, lens, slots, ppi, P, T, D, ids, ids0, draft, d0, pos_off, kv_off, b0 + ) + return True + + +def slot_scatter( + slots, src_new, dst_new, src_next_new, dst_next_new, src_lens, dst_lens, src_draft, dst_draft +) -> bool: + """See slot_scatter_kernel (2-D int32 views, zero padding); False to fall back.""" + ts = (src_new, dst_new, src_next_new, dst_next_new, src_lens, dst_lens, src_draft, dst_draft) + if any(t.dtype != torch.int32 or not t.is_cuda for t in ts) or slots.dtype != torch.long: + return False + two = (src_new, dst_new, src_next_new, dst_next_new, src_draft, dst_draft) + if any(t.dim() != 2 for t in two) or src_lens.dim() != 1 or dst_lens.dim() != 1: + return False + _slot_scatter_op( + slots, + src_new, + dst_new, + src_next_new, + dst_next_new, + src_lens, + dst_lens, + src_draft, + dst_draft, + ) + return True + + +def capture_norm( + h: torch.Tensor, + residual: Optional[torch.Tensor], + weight: torch.Tensor, + eps: float, + dst: torch.Tensor, +) -> None: + """dst [M, H] (row stride may exceed H) = rmsnorm(h + residual) * weight with chunk_norm's rounding.""" + _capture_norm_op(h, residual, weight, eps, dst) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/__init__.py new file mode 100644 index 000000000000..77d7b84d9bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Mixture-of-experts entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_moe_engine.py new file mode 100644 index 000000000000..37d5ab5c6c46 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_moe_engine.py @@ -0,0 +1,325 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3's MoE layer for 1..8 decode tokens in one CuTe DSL launch (``m3_moe``, the TP8 build: routed and shared +expert width 384 per rank): the router, the routed NVFP4 experts and the shared MXFP8 expert, on caller-owned state +(an ``M3MoeState`` and one ``M3MoeLayer`` per MoE layer).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.quantization.utils.fp4_utils import ( + get_reorder_rows_for_gated_act_gemm_row_indices, + get_shuffle_matrix_a_row_indices, +) + +from ..m3_moe.op import ( + HIDDEN, + INTER, + SHARED_INTER, + TOP_K, + M3MoeLayer, + M3MoeState, + ar_push_desc, + is_supported, +) + +__all__ = [ + "M3MoeLayer", + "M3MoeState", + "ar_push_desc", + "is_supported", + "m3_moe_engine", + "moe_reference", + "moe_route", + "mxfp8_dequant", + "mxfp8_quant_dequant", + "unswizzle_sf", +] + +# The model's constants the kernels are built with: routed_scaling_factor, the shared expert's SwiGLU-OAI +# (swiglu_alpha, swiglu_limit, the (up + 1) offset). +ROUTED_SCALING = 2.0 +SHARED_ALPHA = 1.702 +SHARED_LIMIT = 7.0 + +_LOG2E = 1.4426950216293334961 +_RCP6 = 0.16666667163372039795 # fp32(1 / 6), the kernels' reciprocal of the E2M1 maximum +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) +_F32_MAX = 3.40282346638528859812e38 + + +def _f32(x: float) -> torch.Tensor: + return torch.tensor(x, dtype=torch.float32) + + +def _e2m1_decode(packed: torch.Tensor) -> torch.Tensor: + """fp32 values of packed E2M1 bytes [..., K / 2] (low nibble first) -> [..., K].""" + lut = torch.tensor(_E2M1, dtype=torch.float32, device=packed.device) + lo = lut[(packed & 0xF).long()] + hi = lut[(packed >> 4).long()] + return torch.stack((lo, hi), dim=-1).flatten(-2) + + +def _e2m1_round(x: torch.Tensor) -> torch.Tensor: + """x (fp32) to the nearest E2M1 value, ties to even, saturating at 6 (cvt.rn.satfinite.e2m1x2).""" + a = x.abs() + q = torch.full_like(a, 6.0) + for hi, value, inclusive in ( + (5.0, 4.0, True), + (3.5, 3.0, False), + (2.5, 2.0, True), + (1.75, 1.5, False), + (1.25, 1.0, True), + (0.75, 0.5, False), + (0.25, 0.0, True), + ): + q = torch.where(a <= hi if inclusive else a < hi, torch.full_like(a, value), q) + return torch.where(x < 0, -q, q) + + +def _e4m3_round(x: torch.Tensor) -> torch.Tensor: + """x (fp32) to E4M3, nearest even, saturating (cvt.rn.satfinite.e4m3x2), as fp32.""" + return x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() + + +def _ue8m0_up(x: torch.Tensor) -> torch.Tensor: + """UE8M0 code of x >= 0 rounded toward +inf, saturating (cvt.rp.satfinite.ue8m0x2), int32.""" + bits = x.contiguous().view(torch.int32) + code = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) + return code.clamp(0, 254) + + +def unswizzle_sf(sf: torch.Tensor, rows: int, blocks: int) -> torch.Tensor: + """The [rows, blocks] scale bytes of a 128x4-swizzled scale buffer (int64 indices, any byte dtype).""" + flat = sf.reshape(-1).view(torch.uint8) + r = torch.arange(rows, device=flat.device).view(-1, 1) + b = torch.arange(blocks, device=flat.device).view(1, -1) + padded_blocks = (blocks + 3) // 4 * 4 + idx = ( + (r // 128) * (128 * padded_blocks) + + (b // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (b % 4) + ) + return flat[idx] + + +def mxfp8_dequant(q: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: + """fp32 values of e4m3 ``q`` [R, K] with its 128x4-swizzled UE8M0 scales (one per 32 elements).""" + rows, k = q.shape + codes = unswizzle_sf(sf, rows, k // 32).float() + return q.view(torch.float8_e4m3fn).float() * torch.exp2(codes - 127.0).repeat_interleave( + 32, dim=1 + ) + + +def mxfp8_quant_dequant(x: torch.Tensor) -> torch.Tensor: + """x [R, K] through MXFP8 (one UE8M0 scale per 32, the smallest power of two >= amax / 448; e4m3 data), back to + fp32 -- mxfp8_quantize's recipe, which every MXFP8 activation here reproduces.""" + rows, k = x.shape + blocks = x.float().view(rows, k // 32, 32) + code = _ue8m0_up(blocks.abs().amax(dim=-1) * _f32(1.0 / 448.0).to(x.device)) + scale = torch.exp2(code.float() - 127.0).unsqueeze(-1) + return (_e4m3_round(blocks / scale) * scale).view(rows, k) + + +def _bf16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.bfloat16).float() + + +def _gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """a [M, K] @ b [N, K]^T, accumulated in fp64 and rounded once to fp32.""" + return (a.double() @ b.double().t()).float() + + +def _row_perm(rows: int, gated: bool) -> torch.Tensor: + """The TRTLLM-Gen row order: stored row j holds logical row perm[j] (the gated interleave, then the 32-row shuffle + of epilogue tile 128).""" + probe = torch.empty(rows, 1) + perm = get_shuffle_matrix_a_row_indices(probe, epilogue_tile_m=128) + if gated: + perm = get_reorder_rows_for_gated_act_gemm_row_indices(probe)[perm] + return perm + + +def _nvfp4_expert(w: torch.Tensor, sf: torch.Tensor, perm_inv: torch.Tensor) -> torch.Tensor: + """One expert's stored NVFP4 matrix (values [R, K / 2], swizzled e4m3 scales) as raw fp32 [R, K] in logical row + order (the global scale is not applied: it is the per-expert alpha).""" + rows, half = w.shape + codes = unswizzle_sf(sf, rows, 2 * half // 16).view(torch.float8_e4m3fn).float() + vals = _e2m1_decode(w) * codes.repeat_interleave(16, dim=1) + return vals[perm_inv] + + +def moe_route(x: torch.Tensor, router_weight: torch.Tensor, bias: torch.Tensor): + """MiniMax2 routing on the fp32 router GEMV: (top-4 expert ids [M, 4] int64, bf16-rounded weights [M, 4] fp32).""" + logits = _gemm(x, router_weight) + sig = (0.5 * torch.tanh(0.5 * logits).double() + 0.5).float() + keys = sig + bias + ids = torch.sort(keys, dim=-1, descending=True, stable=True).indices[:, :TOP_K] + sk = keys.gather(1, ids) - bias[ids] + total = (sk[:, 0] + sk[:, 2]) + (sk[:, 1] + sk[:, 3]) + w = (sk * _f32(ROUTED_SCALING).to(x.device)) / (total + _f32(1e-20).to(x.device)).unsqueeze(1) + return ids, _bf16(w) + + +def _nvfp4_act(x: torch.Tensor, sc: torch.Tensor) -> torch.Tensor: + """The activation's NVFP4 quantization (fp4_quantize(x, sc, 16)), as raw fp32 values (q x block scale).""" + rows, k = x.shape + blocks = x.view(rows, k // 16, 16) + vmax = blocks.abs().amax(dim=-1, keepdim=True) + sff = _e4m3_round(sc * (vmax * _f32(_RCP6).to(x.device))) + # The kernels' rcp.approx(sf x rcp.approx(sc)) lands where the correctly rounded sc / sf does on every element + # measured (the 1 / (sf x (1 / sc)) spelling differs on ~0.06% of them). + osc = torch.where(vmax != 0, sc / sff, torch.zeros_like(vmax)) + return (_e2m1_round(blocks * osc) * sff).view(rows, k) + + +def _fc1_epilogue(up: torch.Tensor, gate: torch.Tensor, esc: torch.Tensor) -> torch.Tensor: + """TRTLLM-Gen's SwiGLU-OAI FC1 epilogue on the raw FC1 sums (fp32 [I]) and its NVFP4 requantization: the raw FC2 + input values. ``esc`` = (fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha) of the expert, with + act_beta and act_clamp already divided by fc31_alpha (the loader's convention).""" + a, c, alpha, beta, lim = (esc[i] for i in range(5)) + gc = torch.where(gate > lim, lim, gate) + gc = torch.where(gate < -_F32_MAX, torch.full_like(gc, -_F32_MAX), gc) + uc = torch.minimum(torch.maximum(up, -lim), lim) + c1 = a * (alpha * _f32(_LOG2E).to(up.device)) + sgm = 1.0 / (torch.exp2(-(gc * c1)) + 1.0) + h = (uc.double() * a.double() + (a * beta).double()).float() * (gc * sgm) + blocks = h.view(-1, 16) + amax = (blocks.abs() * c).amax(dim=-1, keepdim=True) + sff = _e4m3_round(amax * _f32(_RCP6).to(up.device)) + osc = torch.where(sff != 0, 1.0 / sff, torch.zeros_like(sff)) + return (_e2m1_round((blocks * osc) * c) * sff).view(-1) + + +def _shared_expert(x: torch.Tensor, shared_mxfp8) -> torch.Tensor: + """The shared expert as the target runs it outside the engine: MXFP8 gate_up (bf16 out), SwiGLU-OAI with the + model's alpha / limit / (up + 1) in fp32 (bf16 out), MXFP8 down (bf16 out).""" + gu_w, gu_sf, dn_w, dn_sf = shared_mxfp8 + gu = _bf16(_gemm(mxfp8_quant_dequant(x), mxfp8_dequant(gu_w, gu_sf))) + gate, up = gu[:, :SHARED_INTER], gu[:, SHARED_INTER:] + a = torch.minimum(gate, torch.full_like(gate, SHARED_LIMIT)) + ub = up.clamp(-SHARED_LIMIT, SHARED_LIMIT) + 1.0 + sg = 1.0 / ( + torch.exp2((a * _f32(-SHARED_ALPHA).to(x.device)) * _f32(_LOG2E).to(x.device)) + 1.0 + ) + h = _bf16(ub * (sg * a)) + return _bf16(_gemm(mxfp8_quant_dequant(h), mxfp8_dequant(dn_w, dn_sf))) + + +def moe_reference(hidden: torch.Tensor, layer: M3MoeLayer, shared: bool) -> torch.Tensor: + """The MoE layer of the decode engines (both token ranges compute this), in plain torch: bf16 [M, 6144].""" + x = hidden.float() + m = x.shape[0] + ids, weights = moe_route(x, layer.router_weight, layer.bias) + xr = _nvfp4_act(x, layer.fc31_input_scale.reshape(()).float()) + inv31 = torch.argsort(_row_perm(2 * INTER, gated=True)).to(x.device) + inv2 = torch.argsort(_row_perm(HIDDEN, gated=False)).to(x.device) + y = torch.zeros(m, TOP_K, HIDDEN, dtype=torch.float32, device=x.device) + for e in torch.unique(ids).tolist(): + w31 = _nvfp4_expert(layer.w3_w1_weight[e], layer.w3_w1_weight_scale[e], inv31) + w2 = _nvfp4_expert(layer.w2_weight[e], layer.w2_weight_scale[e], inv2) + esc = layer.escal[:, e] + for t, k in (ids == e).nonzero().tolist(): + acc1 = _gemm(xr[t : t + 1], w31)[0] + # The buffer holds [w3 (up); w1 (gate)]. + h = _fc1_epilogue(acc1[:INTER], acc1[INTER:], esc) + y[t, k] = _bf16(_gemm(h.unsqueeze(0), w2)[0] * esc[5]) + acc = torch.zeros(m, HIDDEN, dtype=torch.float64, device=x.device) + for k in range(TOP_K): + acc = (acc + weights[:, k : k + 1].double() * y[:, k].double()).float().double() + routed = acc.float() + if not shared: + return routed.to(torch.bfloat16) + return (_shared_expert(x, layer.shared_mxfp8) + _bf16(routed)).to(torch.bfloat16) + + +class M3MoeEngine(OpWrapper): + """One MiniMax-M3 MoE layer for M (1..8) decode tokens: ``layer.state.run(layer, hidden, out, arp)``. + + ``hidden`` bf16 [M, 6144] contiguous, M = the state's token count. ``layer`` (from ``M3MoeState.layer``) carries the + TRTLLM-Gen NVFP4 expert buffers read in place, the per-expert scalars, the router weight / bias, the activation's + NVFP4 global scale and, on a ``shared`` state, the shared expert's MXFP8 weights; its state carries the compiled + kernel and the workspace every layer on it shares. Returns bf16 [M, 6144]: bf16(shared + routed) on a ``shared`` + state, else the routed partial. ``out``: write there instead of a new tensor. ``arp`` (``ar_push_desc``): also + push the rows into a multicast all-reduce exchange. + + The routed part is TRTLLM-Gen's NVFP4 MoE on the fp32 router logits: sigmoid (tanh form) + correction bias top-4 + of 128 (ties to the lower id), weights sigmoid / sum * 2.0 rounded to bf16; the activation quantized to NVFP4 with + the layer's global scale; per expert the raw FC1 sums, SwiGLU-OAI (alpha / beta / clamp), the NVFP4 requant with + fc31_scale_c, FC2, bf16(fc2_alpha x sum); the fma chain over the top-4 in order. The shared expert is the target's + MXFP8 SwiGLU-OAI MLP. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # One cell per token count: m_max is a trace-time constant, so each is its own build, and the target builds all + # eight (decode batches of 1..8 requests, spec-decode verify rows up to 8). The whys name what the build differs in. + CELLS: tuple[Cell, ...] = ( + Cell( + why="1 token: FC1 -> FC2 records replicated 32x, logits 32x, last phase without trailing FC2", + spec=dict(tokens=1), + ), + Cell( + why="2 tokens: records 32x, logits 32x, the one build without AR_PRE", + spec=dict(tokens=2), + ), + Cell(why="3 tokens: records 8x, logits 32x, 4 FC2-B builder warps", spec=dict(tokens=3)), + Cell( + why="4 tokens: one request verifying 3 drafts (the TP8 operating point), records 8x, logits 32x", + spec=dict(tokens=4), + ), + Cell( + why="5 tokens: single records, a routing warp takes two tokens, logits 32x", + spec=dict(tokens=5), + ), + Cell(why="6 tokens: one request verifying a 5-token draft", spec=dict(tokens=6)), + Cell(why="7 tokens: the B2IMG global FC1 -> FC2 hand-off", spec=dict(tokens=7)), + Cell( + why="8 tokens: B2IMG with a 6-stage ring, 32 distinct experts at most", + spec=dict(tokens=8), + ), + ) + + note = """ + The kernel's elementwise math uses the PTX approximations (rcp.approx, ex2.approx, tanh) where `reference` uses + torch's correctly rounded ones, and its GEMMs sum in tensor-core order where `reference` sums in fp64: a + quantization boundary (NVFP4 activation / intermediate code, routing near-tie) can fall on the other side, so the + two agree within the band the test derives, not bit for bit. The routed part is bit for bit the stock TRTLLM-Gen + NVFP4 MoE (fp4_quantize + fp4_block_scale_moe_runner) on the same routing; the test checks that too. + + Not covered: the all-reduce push (`arp`), which needs a multicast exchange (the decode forward's boundary all-reduce + exercises it); the non-shared state (`shared=False`), which this target never builds. Every layer on one state + must run in one stream order: the workspace is shared and tagged per call, never reset. + + No `is_valid`: the compiled call checks every buffer's dtype, shape and alignment against its signature and + raises on a mismatch (a misaligned `hidden`, an `out` that is not bf16 [M, 6144], a malformed `arp`). + """ + + def __call__( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return layer.state.run(layer, hidden, out, arp=arp) + + def reference( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Plain torch over the layer's buffers (see the class docstring); `out` and `arp` are ignored.""" + return moe_reference(hidden, layer, shared=layer.state.shared) + + +m3_moe_engine = M3MoeEngine() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_router_gemv.py new file mode 100644 index 000000000000..37870d0e6cd7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_router_gemv.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""fp32 router GEMV of the decode MoE: ``logits[M, E] = float(x) @ W^T`` at one token.""" + +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, +) + +_E4M3_MAX = 448.0 +_E2M1_MAX = 6.0 +#: Midpoints between consecutive e2m1 magnitudes; ``_TIE_UP[i]``: round-to-nearest-even takes the upper code. +_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) +_TIE_UP = (False, True, False, True, False, True, False) +#: The kernel builds each block's output scale through two ``rcp.approx.ftz.f32`` reciprocals (~2^-23 relative): +#: a value this close to an e2m1 midpoint may round to either neighbour, everything further out must match. +TIE_WINDOW = 2.0**-20 + + +def nvfp4_quantize( + x: torch.Tensor, sf_scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``fp4_quantize(x, sf_scale, 16, False, False)``: (packed e2m1 [M, K / 2], linear E4M3 scales [M * K / 16], + the pre-rounding scaled values [M, K] that ``near_tie`` reads). + + Per 16 columns: scale = e4m3(sf_scale * amax / 6), data = x * sf_scale / scale rounded to nearest-even e2m1. + """ + gs = sf_scale.float().reshape(()) + m, k = x.shape + blocks = x.float().view(m, k // 16, 16) + vecmax = blocks.abs().amax(-1) + sf = (gs * (vecmax / _E2M1_MAX)).clamp(max=_E4M3_MAX).to(torch.float8_e4m3fn) + out_scale = torch.where(vecmax != 0, gs / sf.float(), torch.zeros_like(vecmax)) + scaled = (blocks * out_scale.unsqueeze(-1)).view(m, k) + mid = torch.tensor(_E2M1_MIDPOINTS, device=x.device, dtype=torch.float32) + tie_up = torch.tensor(_TIE_UP, device=x.device, dtype=torch.bool) + a = scaled.abs().unsqueeze(-1) + code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8) + negative = (scaled < 0) | ((scaled == 0) & torch.signbit(scaled)) + code = code | (negative.to(torch.uint8) << 3) + packed = code[:, 0::2] | (code[:, 1::2] << 4) + return packed, sf.view(torch.uint8).reshape(-1), scaled + + +def near_tie(scaled: torch.Tensor) -> torch.Tensor: + """Elements whose e2m1 rounding the kernel's approximate reciprocals decide (see ``TIE_WINDOW``).""" + mid = torch.tensor(_E2M1_MIDPOINTS, device=scaled.device, dtype=torch.float32) + a = scaled.abs().unsqueeze(-1) + return ((a - mid).abs() <= TIE_WINDOW * a).any(-1) + + +def fp32_logits(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """``float(x) @ w^T`` as a true fp32 product (TF32 off).""" + with true_fp32_matmul(): + return x.float() @ w.t() + + +class M3RouterGemv(OpWrapper): + """fp32 router logits ``float(x[M, K]) @ w[E, K]^T`` with exact fp32 FMA, one CTA per expert. + + ``x`` is BF16 with 1..16 rows, unit column stride and 16-byte rows; ``w`` is the fp32 router weight; K a + multiple of 1024. ``rows`` experts per CTA; only 1 is built. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The target takes the GEMV at exactly one token (2..48 go to the tensor-core router): 128 experts, hidden 6144. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token, 128 experts x hidden 6144", + spec=dict(m=1, e=128, k=6144), + ), + ) + + note = """ + The op accepts up to 16 rows; the target sends one, so wider calls are not cells. The logits differ from a true + fp32 product only in accumulation order (per-thread FMA chains over K / 1024 strided column groups, a warp tree + and an 8-warp sum). + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_router_gemv(x, w, rows, pdl) + + def reference(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> torch.Tensor: + """True fp32 product; ``rows`` and ``pdl`` do not enter the arithmetic.""" + return fp32_logits(x, w) + + def is_valid(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> None: + # The activation rows are bulk-copied from x's base address; the op checks the row stride but not the base, + # and a misaligned bulk copy faults the context rather than raising. + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + + +m3_router_gemv = M3RouterGemv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_router_gemv_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_router_gemv_fp4.py new file mode 100644 index 000000000000..58f809e82cf6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_router_gemv_fp4.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""fp32 router GEMV plus the routed experts' NVFP4 activation, one token.""" + +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 + +from .m3_router_gemv import fp32_logits, nvfp4_quantize + + +class M3RouterGemvFp4(OpWrapper): + """``[logits fp32 [M, E], x_fp4 [M, K / 2] uint8, x_sf [M * K / 16] uint8]`` from one kernel. + + The logits are ``m3_router_gemv``'s; the NVFP4 activation is ``fp4_quantize(x, sf_scale, 16, False, False)`` + (16-element blocks, E4M3 scales in the linear layout, per-tensor ``sf_scale``) for the routed experts' FC1. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The fused-router decode step at one token: 128 experts, hidden 6144, the experts' input scale. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token: logits for 128 experts and the NVFP4 activation of hidden 6144", + spec=dict(m=1, e=128, k=6144, sf_scale=448.0 * 6.0 / 8.0), + ), + Cell( + why="a global scale that puts block scales near the E4M3 top, where the scale cast saturates", + spec=dict(m=1, e=128, k=6144, sf_scale=448.0 * 6.0 / 2.0), + ), + ) + + note = """ + The e2m1 codes are exact except within `TIE_WINDOW` of a rounding midpoint, where the kernel's approximate + reciprocals decide the direction. Up to 16 rows are accepted; the target sends one. + """ + + def __call__( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_router_gemv_fp4(x, w, sf_scale, pdl) + + def reference( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """True fp32 logits and the NVFP4 recipe; ``pdl`` does not enter the arithmetic.""" + packed, sf, _ = nvfp4_quantize(x, sf_scale) + return fp32_logits(x, w), packed, sf + + def is_valid(self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool) -> None: + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + # A zero or non-finite global scale makes every block scale 0 or NaN and the data codes garbage; the op + # reads the scalar on the device and cannot refuse it. + value = float(sf_scale.reshape(())) + assert value > 0.0 and value != float("inf"), ( + f"sf_scale must be positive and finite, got {value}" + ) + + +m3_router_gemv_fp4 = M3RouterGemvFp4() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_router_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_router_tc.py new file mode 100644 index 000000000000..e88af51f7fa2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_router_tc.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TF32 tensor-core router of the decode MoE for 2..48 tokens, optionally with the NVFP4 activation.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from .m3_router_gemv import nvfp4_quantize + + +def tf32_rna(w: torch.Tensor) -> torch.Tensor: + """fp32 -> TF32 as ``cvt.rna.tf32.f32``: keep 10 mantissa bits, round to nearest with ties away from zero.""" + bits = w.contiguous().view(torch.int32) + rounded = (bits + 0x1000) & ~0x1FFF + finite = (bits & 0x7F800000) != 0x7F800000 + return torch.where(finite, rounded, bits).view(torch.float32) + + +class M3RouterTc(OpWrapper): + """Router logits on TF32 tensor cores for 128 experts and 1..48 tokens: ``x @ tf32(w)^T``, fp32 accumulate. + + ``x`` is BF16 (exact in TF32), ``w`` the fp32 router weight, rounded to TF32 in the kernel. The 16 K slices' + partials are summed in slice order whatever cluster size the device co-schedules (16 or 8 CTAs), so the result + does not depend on it. With ``sf_scale`` also the NVFP4 activation as ``m3_router_gemv_fp4``: returns + ``[logits]`` or ``[logits, x_fp4, x_sf]``. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The target routes 2..48 tokens here: plain logits on the unfused path, logits + NVFP4 on the fused-router path. + CELLS: tuple[Cell, ...] = ( + Cell( + why="two tokens, logits only (the smallest call the target makes here)", + spec=dict(m=2, k=6144, fp4=False), + ), + Cell( + why="8 tokens with the NVFP4 activation: one 8-token MMA tile, every expert tile quantizes one row", + spec=dict(m=8, k=6144, fp4=True), + ), + Cell( + why="13 tokens: a partial second token tile", + spec=dict(m=13, k=6144, fp4=True), + ), + Cell( + why="the 48-token maximum, six token tiles, logits only", + spec=dict(m=48, k=6144, fp4=False), + ), + Cell( + why="the 48-token maximum with the NVFP4 activation", + spec=dict(m=48, k=6144, fp4=True), + ), + # Every other row count the router takes (dispatch.ROUTER_TC), with and without the NVFP4 activation. + *( + Cell( + why=f"{m} rows: a router row count between the ones above", + spec=dict(m=m, k=6144, fp4=m % 2 == 1), + ) + for m in range(2, 49) + if m not in (2, 8, 13, 48) + ), + ) + + note = """ + The kernel uses 16-CTA clusters when the device co-schedules eight of them and 8-CTA clusters otherwise; both + sum the same 16 slice partials in the same order. Only the one the test device picks is driven. + """ + + def __call__( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_router_tc(x, w, sf_scale, pdl) + + def reference( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> tuple[torch.Tensor, ...]: + """fp32 product of ``x`` and the TF32-rounded weight, plus the NVFP4 recipe with ``sf_scale``.""" + with true_fp32_matmul(): + logits = x.float() @ tf32_rna(w).t() + if sf_scale is None: + return (logits,) + packed, sf, _ = nvfp4_quantize(x, sf_scale) + return logits, packed, sf + + def is_valid( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> None: + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + if sf_scale is not None: + value = float(sf_scale.reshape(())) + assert value > 0.0 and value != float("inf"), ( + f"sf_scale must be positive and finite, got {value}" + ) + + +m3_router_tc = M3RouterTc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/msa_kv_update.py new file mode 100644 index 000000000000..8c2a67356d3f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/msa_kv_update.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One kernel for MSA's device-side KV-length correction (MsaAttentionMetadata.on_update_kv_lens). + +The torch formulation is ~23 one-CTA kernels (min, cumsum, gathers, div / clamp / remainder, the block-table lookup, +the valid-block counts) and runs at the head of every captured decode step, twice. This kernel computes the same +int32 results in one launch: + + kv_true[b] = min(kv_lens[b], staged[b]) -> seq_lens_out[b], cu_kv[b + 1] = prefix sum (cu_kv[0] = 0) + pos[t] = kv_true[r] - qo[r] + intra[t] (r = q_batch_row[t]) + out_loc[t] = table[r, min(max(pos, 0) // page, ncols - 1)] * page + max(pos, 0) % page + n_valid[t] = (max(pos + 1, 0) + page - 1) // page + +``msa_kv_update`` is the target's ``MsaKvUpdate`` (``tensorrt_llm._torch.pyexecutor.runtime_hooks``): the MSA +metadata calls it with its own buffers and falls back to the torch formulation when it returns False. The kernel is +``torch.ops.trtllm.minimax_m3_msa_kv_update``, called through this target's ``msa_kv_update`` entry. +""" + +from typing import Optional + +import torch + +from ..ops.attention.msa_kv_update import msa_kv_update as _kv_update_op + + +def msa_kv_update( + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], +) -> bool: + """The fused correction (see module doc); False, with nothing written, when an operand is outside what the + kernel covers. Arguments in ``MsaKvUpdate`` order; ``seq_lens`` are the per-request query lengths (``qo``).""" + ts = [ + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + msa_seq_lens, + out_cache_loc, + ] + ts += [t for t in (cu_kv_lens, n_valid_blocks) if t is not None] + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in ts) + or block_table.dim() != 2 + or block_table.stride(1) != 1 + ): + return False + if any(not t.is_contiguous() for t in ts if t is not block_table): + return False + if ( + kv_lens.numel() < batch + or kv_lens_staged.numel() < batch + or seq_lens.numel() < batch + or msa_seq_lens.numel() < batch + or q_batch_row.numel() < total_q + or q_intra.numel() < total_q + or out_cache_loc.numel() < total_q + or (cu_kv_lens is not None and cu_kv_lens.numel() < batch + 1) + or (n_valid_blocks is not None and n_valid_blocks.numel() < total_q) + ): + return False + _kv_update_op( + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + tokens_per_block, + batch, + total_q, + msa_seq_lens, + cu_kv_lens, + out_cache_loc, + n_valid_blocks, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/mx_gemv/__init__.py new file mode 100644 index 000000000000..7d2ec0224cf2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/mx_gemv/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) as a CuTe DSL kernel, configured for the TP8 +projection shapes (see op.py).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py new file mode 100644 index 000000000000..e3d7751231ee --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py @@ -0,0 +1,1238 @@ +# 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. +"""MXFP8 decode GEMV on tcgen05: ``y[M, N] = x[M, K] @ W[N, K]^T`` for M <= 8, block-scaled E4M3 x E4M3. + +The weight is the MMA's A operand (128 output features per CTA), the activation is B (8 token columns), +fp32 accumulation in TMEM. Both operands come in their MXFP8 storage: E4M3 values (K contiguous) and UE8M0 +scales in the 128x4-swizzled layout, whose 512-byte atom per (128 rows, 128 K) is exactly the source of a +``tcgen05.cp`` 32x128b.warpx4 scale copy, so both scale tensors are bulk-copied unchanged. For an activation of at +most 32 rows the atom of k-tile ``k`` sits at byte ``512 k``. + +A cluster of SPLIT CTAs shares one 128-row weight tile; rank r streams the interleaved k-tiles r, r + SPLIT, ... +through a RING-stage ring. Before the grid dependency each CTA issues its first RING weight tiles (and their scale +atoms) and prefetches the rest into L2; after it, the activation slice of the CTA (values and scale atoms) lands in +one buffer. Rank r owns rows [r * 128 / SPLIT, (r + 1) * 128 / SPLIT) of the tile: the other ranks store their fp32 +partials of those rows into its shared memory and arrive on its barrier (release, cluster scope); the owner adds the +SPLIT partials in rank order and stores bf16. + +Warps: 0 weight TMA, 1 activation TMA, 2 TMEM allocation + scale copies + MMA, 3 idle, 4-7 epilogue (warp 4 + i +reads TMEM lanes 32 i .. 32 i + 31). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass import dsl_user_op +from cutlass.experimental import primitives as prims + +CTA_M = 128 # weight rows per CTA = the MMA's M # capacity: fixed: a kernel tile constant +MMA_N = 8 # token columns +CTA_K = 128 # one k-tile: 128 E4M3 bytes per row = one 128-byte swizzle row # capacity: fixed: a kernel tile constant +MMA_INST_K = 32 +K_BLOCKS = ( + CTA_K // MMA_INST_K +) # 4 MMAs per k-tile, one per scale byte of the atom # capacity: fixed: a kernel tile constant +THREADS = 256 +SF_ATOM_BYTES = 512 +SF_ATOM_U16 = SF_ATOM_BYTES // 2 +SF_COLS = 4 # TMEM columns of one scale atom +W_TILE_BYTES = CTA_M * CTA_K +X_TILE_BYTES = MMA_N * CTA_K +EVICT_FIRST = 0x12F0000000000000 +A_COLS_PER_MMA = MMA_INST_K // 4 # TMEM columns of one 128 x K32 E4M3 granule of A + +a_dtype = cutlass.Float8E4M3FN +b_dtype = cutlass.Float8E4M3FN + + +def num_k_tiles(k_in: int) -> int: + return k_in // CTA_K + + +def acc_cols(nacc: int, mma_n: int = MMA_N) -> int: + # nacc independent accumulators of MMA_N columns: MMA j of the CTA accumulates into j % nacc, so the M128 x N8 + # MMAs do not form one serial dependency chain; the epilogue adds them. + return max(16, nacc * mma_n) + + +def a_col0(nacc: int, my_tiles: int, mma_n: int = MMA_N) -> int: + # a_tmem: [acc | SFA per k-tile | one SFB slot (reused in issue order) | A: 32 columns (128 x 128 E4M3) per k-tile] + return (acc_cols(nacc, mma_n) + (my_tiles + 1) * SF_COLS + 31) // 32 * 32 + + +def tmem_cols( + ring: int, nacc: int, a_tmem: bool = False, my_tiles: int = 0, mma_n: int = MMA_N +) -> int: + need = acc_cols(nacc, mma_n) + 2 * ring * SF_COLS + if a_tmem: + need = a_col0(nacc, my_tiles, mma_n) + my_tiles * K_BLOCKS * A_COLS_PER_MMA + cols = 32 + while cols < need: + cols *= 2 + return cols + + +@dsl_user_op +def read_globaltimer(*, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_async_v4(dst, w0, w1, w2, w3, mbar, *, loc=None, ip=None): + """st.async of four int32 words to a shared::cluster address, completing ``mbar`` there by 16 bytes.""" + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Int32(w0).ir_value(loc=loc, ip=ip), + cutlass.Int32(w1).ir_value(loc=loc, ip=ip), cutlass.Int32(w2).ir_value(loc=loc, ip=ip), + cutlass.Int32(w3).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "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, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _ld_shared_v4_f32(smem_ptr, *, loc=None, ip=None): + """ld.shared.v4.f32 of four fp32 at a 16-byte-aligned shared-memory pointer.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + out = _llvm.inline_asm( + _llvm.StructType.get_literal([_T.f32()] * 4), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", "=f,=f,=f,=f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + return [ + cutlass.Float32(_llvm.extractvalue(_T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4) + ] + + +@dsl_user_op +def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA completed by other CTAs' st.async.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@dsl_user_op +def _absmax(acc, v, *, loc=None, ip=None): + """max(acc, |v|).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(acc).ir_value(loc=loc, ip=ip), cutlass.Float32(v).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 a;\n\tabs.f32 a, $2;\n\tmax.f32 $0, $1, a;\n\t}", "=f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rcp_approx_ftz(x, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _e4m3x4(v0, v1, v2, v3, *, loc=None, ip=None): + """Four f32 to E4M3 (RN, satfinite), packed little-endian into one int32 word (v0 in the low byte).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n\t.reg .b16 lo, hi;\n\tcvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\tmov.b32 $0, {lo, hi};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).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, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16(x) (RN) as the fp32 bit pattern of the rounded value.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 h, z;\n\tcvt.rn.bf16.f32 h, $1;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _bf16_round(x): + return _bf16_bits(x).bitcast(cutlass.Float32) + + +@dsl_user_op +def _rope_bits(c, a, s, b, *, loc=None, ip=None): + """bf16(c * a + s * b) with separately rounded products (no FMA), as fp32 bits: the stock RoPE arithmetic.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (c, a, s, b)], + "{\n\t.reg .f32 p, q, r;\n\t.reg .b16 h, z;\n\tmul.rn.f32 p, $1, $2;\n\tmul.rn.f32 q, $3, $4;\n\t" + "add.rn.f32 r, p, q;\n\tcvt.rn.bf16.f32 h, r;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _din_load4(table, hidden, hidden_half, word_idx): + """Four int32 words at word_idx of the hidden rows (hidden_half == 1) or the embedding table: one CTA-uniform + base-address select instead of a branch.""" + addr = cutlass.select_( + hidden_half == cutlass.Int32(1), hidden.data_ptr(0).toint(), table.data_ptr(0).toint() + ) + ptr = cutlass.inttoptr( + addr + cutlass.Int64(word_idx) * cutlass.Int64(4), mem_space=1, dtype=cutlass.Int32 + ) + return ptr.load(count=4, alignment=16) + + +def _stamp(timers, timed: bool, bx, slot: int): + # Diagnostic: %globaltimer of milestone `slot` of CTA bx (timers [grid, 8] int64). + if timed: + timers.store(read_globaltimer(), idx=bx * cutlass.Int32(8) + cutlass.Int32(slot)) + + +def _k_tile(rank, i: int, split: int, my_tiles: int, contig: bool): + # contig: rank r streams k-tiles [r * my_tiles, (r + 1) * my_tiles); else the interleaved r, r + split, ... + if contig: + return rank * cutlass.Int32(my_tiles) + cutlass.Int32(i) + return rank + cutlass.Int32(i * split) + + +@cute.kernel +def m3_mx_gemv_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] E4M3 as bytes, box 128 x 128 + tma_desc_wsf: cutlass.GridConstant[ + cuda.TensorMap + ], # W scale atoms [N / 128 * K / 128, 256] u16 + tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] E4M3 as bytes, box 128 x 8 + tma_desc_xsf: cutlass.GridConstant[cuda.TensorMap], # x scale atoms [K / 128, 256] u16 + y: cutlass.Array, # [M * N] bf16, token-major + x_words: cutlass.Array, # x_bf16: the bf16 activation [M, K] as int32 words + timers: cutlass.Array, # [grid, 8] int64 milestones when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, # ar_push: the all-reduce workspace's multicast mapping (int32 words) + ar_flags: cutlass.Int64, # ar_push: its call flags (parity of the next call = flags[0] & 1) + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[ + int + ], # 1: push the output rows into the next all-reduce; 2: and trigger after it + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], # token columns of the MMA (8, 16 or 48): M <= mma_n + din_i: cutlass.Array, # din: int32 token ids [M] + din_table: cutlass.Array, # din: embedding [V, K / 2] bf16 as int32 words + din_hidden: cutlass.Array, # din: hidden rows [M, K / 2] bf16 as int32 words + din_w: cutlass.Array, # din: [embedding norm | hidden norm] weights [K] bf16 as int32 words + din_f: cutlass.Array, # din: fp32 [eps_e, eps_h] + din_vocab: cutlass.Int32, + din_g: cutlass.Array, # din, flag 1: int64 hidden row of each token, as int32 words + din_res: cutlass.Array, # din, flag 2: the gathered hidden rows out [M, K / 2] bf16 as int32 words + din_flags: cutlass.Int32, # din: 1 = gather the hidden rows, 2 = write them to din_res + rope_cs: cutlass.Array, # rope: fp32 (cos, sin) [positions][64][2] + rope_i: cutlass.Array, # rope: int32 cache lengths [M / T] + rope_bt: cutlass.Array, # rope: int32 K / V page tables [M / T][2][max_blocks] + rope_q: cutlass.Array, # rope: E4M3 Q out [M][hq][128] as uint8 + rope_pool: cutlass.Array, # rope: paged E4M3 cache [pages][hkv][tokens per page][128] as uint8 + rope_scale: cutlass.Array, # rope: fp32 [1] E4M3 quantization scale + rope_dims: cutlass.Array, # rope: int32 [max_blocks, tokens-per-page log2, hq, hkv, T] + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], +): + k_tiles = num_k_tiles(k_in) + x_tile_bytes = mma_n * CTA_K + my_tiles = k_tiles // split + rows_per_rank = CTA_M // split + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + rank = cutlass.Int32(0) + if cutlass.const_expr(split > 1): + rank = cute.arch.block_idx_in_cluster() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_wsf = tma_desc_wsf.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + tma_ptr_xsf = tma_desc_xsf.get_ptr() + n_tile = bx // cutlass.Int32(split) + m_offset = n_tile * cutlass.Int32(CTA_M) + + smem_a = cutlass.Array( + cutlass.Uint8, ring * W_TILE_BYTES, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + cutlass.Uint8, my_tiles * x_tile_bytes, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_sfa = cutlass.Array( + cutlass.Uint8, ring * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + smem_sfb = cutlass.Array( + cutlass.Uint8, my_tiles * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + w_full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + mma_done = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + act_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + reduce_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + # Partials of the rows this rank owns: [source rank][row][token] fp32 (its own slot stays unused). Rows are padded + # to mma_n + 4 words: the owner's 16-byte row reads (lane = row) then hit all 32 banks once per 8 lanes; at a + # stride of mma_n they were 8- to 16-way conflicted (4.3 us of the 48-token QKV tail). + mb_stride = mma_n + 4 + mailbox = cutlass.Array( + cutlass.Float32, + split * rows_per_rank * mb_stride, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + # din: per-token RMS factor of the half of x this CTA's k-tiles lie in. + din_rs = None + if cutlass.const_expr(din): + din_rs = cutlass.Array(cutlass.Float32, 8, space=cutlass.AddressSpace.smem, alignment=16) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_wsf) + prims.prefetch_tensormap(tma_ptr_x) + prims.prefetch_tensormap(tma_ptr_xsf) + if prims.elect_sync(): + for s in cutlass.range_constexpr(ring): + prims.mbarrier_init(w_full.subview(s), 1) + prims.mbarrier_init(mma_done.subview(s), 1) + prims.mbarrier_init(act_full, 1) + prims.mbarrier_init(acc_done, 1) + if cutlass.const_expr(split > 1 and st_async): + # Completed by the peers' st.async bytes; this CTA's own arrival carries the expected count. + prims.mbarrier_init(reduce_ready, 1) + prims.mbarrier_arrive_expect_tx( + reduce_ready, (split - 1) * rows_per_rank * mma_n * 4 + ) + elif cutlass.const_expr(split > 1): + prims.mbarrier_init(reduce_ready, (split - 1) * rows_per_rank) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n)) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + if cutlass.const_expr(split > 1): + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + if tx == 0: + _stamp(timers, timed, bx, 0) + + if warp_id == 0: + if prims.elect_sync(): + for i in cutlass.range_constexpr(ring, my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [k_global * cutlass.Int32(CTA_K), m_offset], [] + ) + if cutlass.const_expr((n_out // CTA_M) * split >= 80): + # Wide grids launch in part only after the predecessor (e.g. the fused attention) frees its SMs, which + # loses their pre-wait weight prefetch: the CTAs of the first half also stage their partner's + # (bx + grid / 2) weights into L2. + if bx < cutlass.Int32((n_out // CTA_M) * split // 2): + pb = bx + cutlass.Int32((n_out // CTA_M) * split // 2) + p_rank = pb % cutlass.Int32(split) + p_m_offset = (pb // cutlass.Int32(split)) * cutlass.Int32(CTA_M) + for i in cutlass.range_constexpr(my_tiles): + p_k = _k_tile(p_rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [p_k * cutlass.Int32(CTA_K), p_m_offset], [] + ) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + if cutlass.const_expr(i >= ring): + while not cute.arch.mbarrier_try_wait( + mma_done.subview(stage).data_ptr(), (i // ring - 1) % 2 + ): + pass + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.mbarrier_arrive_expect_tx(w_full.subview(stage), W_TILE_BYTES + SF_ATOM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * W_TILE_BYTES), + tma_ptr_w, + (k_global * cutlass.Int32(CTA_K), m_offset), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfa.subview(stage * SF_ATOM_BYTES), + tma_ptr_wsf, + (cutlass.Int32(0), n_tile * cutlass.Int32(k_tiles) + k_global), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + elif warp_id == 1: + prims.griddepcontrol(prims.GridDepAction.WAIT) + if cutlass.const_expr(trigger_early): + # Dependents launch once this grid's own dependency is satisfied (not at entry): consumers that read + # step metadata or caches before their own grid wait rely on every grid ahead of them having passed its + # wait, so the chain of releases never runs ahead of a producer still in flight. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if tx == 32: + _stamp(timers, timed, bx, 1) + if cutlass.const_expr(not x_bf16): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(act_full, my_tiles * (x_tile_bytes + SF_ATOM_BYTES)) + for i in cutlass.range_constexpr(my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(i * x_tile_bytes), + tma_ptr_x, + (k_global * cutlass.Int32(CTA_K), cutlass.Int32(0)), + act_full, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfb.subview(i * SF_ATOM_BYTES), + tma_ptr_xsf, + (cutlass.Int32(0), k_global), + act_full, + ) + elif warp_id == 2: + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + ring * SF_COLS + idesc = prims.Tcgen05MxInstrDesc.build( + a_dtype=a_dtype, b_dtype=b_dtype, scale_format=1, n_dim=mma_n, m_dim=CTA_M + ) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + if cutlass.const_expr(a_tmem): + # A staged in TMEM: before the grid wait every weight k-tile and its scale atom go on into TMEM as they + # land (the copies of a stage refilled later are committed to its mma_done); once the activation lands, + # its scale atoms and the MMAs, A read from TMEM. + base_row = tmem_raw >> 16 + base_col = tmem_raw & 0xFFFF + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + my_tiles * SF_COLS + a_base = tmem_raw + cutlass.Int32(a_col0(nacc, my_tiles, mma_n)) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + if prims.elect_sync(): + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_cp( + prims.Tcgen05CpShape.SHAPE_128X256B, + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), + 6, + cutlass.Int32, + ), + desc_a + ((MMA_INST_K >> 4) * kb), + group=prims.CTAGroup.CTA_1, + ) + prims.tcgen05_cp( + s2t_shape, + cutlass.inttoptr((base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32), + desc_sfa, multicast=s2t_multicast, + ) # fmt: skip + if cutlass.const_expr(i + ring < my_tiles): + prims.tcgen05_commit(mma_done.subview(stage)) + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + for i in cutlass.range_constexpr(my_tiles): + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + sfa_tmem = cutlass.inttoptr( + (base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32 + ) + sfb_tmem = cutlass.inttoptr((base_row << 16) | sfb_col0, 6, cutlass.Int32) + prims.tcgen05_cp(s2t_shape, sfb_tmem, desc_sfb, multicast=s2t_multicast) + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), 6, cutlass.Int32 + ), + desc_b + ((MMA_INST_K >> 4) * kb), idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + i > 0 or kb > 0, + sfa_tmem, sfb_tmem, + ) # fmt: skip + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + else: + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + sfa_tmem = (base_row << 16) | (sfa_col0 + stage * SF_COLS) + sfb_tmem = (base_row << 16) | (sfb_col0 + stage * SF_COLS) + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), desc_sfa, + multicast=s2t_multicast) # fmt: skip + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), desc_sfb, + multicast=s2t_multicast) # fmt: skip + for kb in cutlass.range_constexpr(K_BLOCKS): + inc = ((MMA_INST_K * 8 // 8) >> 4) * kb + j = i * K_BLOCKS + kb + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw + cutlass.Int32((j % nacc) * mma_n), 6, cutlass.Float32), + desc_a + inc, desc_b + inc, idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), j >= nacc, + cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), + ) # fmt: skip + prims.tcgen05_commit(mma_done.subview(stage)) + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + elif warp_id >= 4: + lane = tx % 32 + row = (warp_id % 4) * 32 + lane + if cutlass.const_expr(x_bf16): + # The bf16 activation of this CTA's k-tiles to MXFP8 in shared memory, as + # torch.ops.trtllm.mxfp8_quantize(x, True) does (UE8M0 = round-up(qz_amax * rcp(448)), + # E4M3 RN-satfinite): one 32-element block (token qz_t, k-tile qz_i, block qz_kb) per thread and round, + # the E4M3 bytes into the 128-byte-swizzled B tile, the qz_scale byte into the tile's qz_scale atom. + qz_tid = tx - cutlass.Int32(128) + prims.griddepcontrol(prims.GridDepAction.WAIT) + din_half_k = k_in // 2 + din_hidden_half = cutlass.Int32(0) + if cutlass.const_expr(din): + # Draft input x = [rmsnorm(table[ids]) * w_e, rmsnorm(hidden) * w_h] (an out-of-range id embeds to + # zero): this CTA's k-tiles lie in one half, so it needs that half's statistic of every token. 16 + # threads per token, 16-byte loads, all of a thread's loads in flight at once. + din_hidden_half = cutlass.select_( + rank * cutlass.Int32(my_tiles * CTA_K) >= cutlass.Int32(din_half_k), + cutlass.Int32(1), + cutlass.Int32(0), + ) + din_tok = qz_tid // cutlass.Int32(16) + din_part = qz_tid % cutlass.Int32(16) + din_tok_ld = cutlass.select_(din_tok < num_tokens, din_tok, cutlass.Int32(0)) + din_id = din_i.load(idx=din_tok_ld) + din_valid = (din_id >= cutlass.Int32(0)) & (din_id < din_vocab) + din_hrow = din_tok_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + din_hrow = din_g.load(idx=din_tok_ld * cutlass.Int32(2)) + din_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + din_hrow, + cutlass.select_(din_valid, din_id, cutlass.Int32(0)), + ) + # The first hidden-half CTA of tile 0 reads whole hidden rows here: it also writes them out. + din_copy = ( + ((din_flags & cutlass.Int32(2)) != cutlass.Int32(0)) + & (bx == cutlass.Int32(split // 2)) + & (din_hidden_half == cutlass.Int32(1)) + & (din_tok < num_tokens) + ) + din_base = din_row * cutlass.Int32(din_half_k // 2) + din_ss = cutlass.Float32(0.0) + din_vecs_per_part = din_half_k // 8 // 16 + for din_round in cutlass.range_constexpr(0, din_vecs_per_part, 16): + din_n = min(16, din_vecs_per_part - din_round) + din_loaded = [] + for din_j in cutlass.range_constexpr(din_n): + din_widx = din_base + ( + din_part + cutlass.Int32((din_round + din_j) * 16) + ) * cutlass.Int32(4) + din_loaded.append( + _din_load4(din_table, din_hidden, din_hidden_half, din_widx) + ) + if din_copy: + for din_j in cutlass.range_constexpr(din_n): + din_v = din_loaded[din_j] + din_res.store( + (din_v[0], din_v[1], din_v[2], din_v[3]), + idx=din_tok * cutlass.Int32(din_half_k // 2) + + (din_part + cutlass.Int32((din_round + din_j) * 16)) + * cutlass.Int32(4), + alignment=16, + ) + for din_j in cutlass.range_constexpr(din_n): + for din_e in cutlass.range_constexpr(4): + din_word = cutlass.Int32(din_loaded[din_j][din_e]) + din_lo = (din_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + din_hi = (din_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + din_ss = din_ss + din_lo * din_lo + din_hi * din_hi + for din_o in cutlass.range_constexpr(4): + din_ss = din_ss + cute.arch.shuffle_sync_bfly(din_ss, offset=1 << din_o) + if din_part == cutlass.Int32(0): + din_eps = din_f.load(idx=din_hidden_half) + din_r = cute.math.rsqrt( + din_ss / cutlass.Float32(din_half_k) + din_eps, fastmath=True + ) + din_keep = din_valid | (din_hidden_half == cutlass.Int32(1)) + din_rs.store( + cutlass.select_(din_keep, din_r, cutlass.Float32(0.0)), idx=din_tok + ) + prims.barrier_cta_sync(3, thread_count=128) + qz_b_words = cutlass.Array(smem_b.data_ptr(0), shape=(my_tiles * x_tile_bytes // 4,), dtype=cutlass.Int32, + alignment=16) # fmt: skip + qz_items = my_tiles * K_BLOCKS * mma_n + qz_rounds = (qz_items + 127) // 128 + # A group's loads first (clamped to a valid row), then its math: one L2 round trip per group of + # QZ_GROUP rounds instead of one per round, without holding every round's data in registers. + QZ_GROUP = qz_rounds if qz_rounds <= 8 else 4 + qz_loaded = [] + for qz_r in cutlass.range_constexpr(qz_rounds): + qz_item = qz_tid + cutlass.Int32(qz_r * 128) + qz_t = qz_item % cutlass.Int32(mma_n) + qz_kb = (qz_item // cutlass.Int32(mma_n)) % cutlass.Int32(K_BLOCKS) + qz_i = qz_item // cutlass.Int32(mma_n * K_BLOCKS) + if cutlass.const_expr(qz_r * 128 + 127 >= qz_items): + qz_i = cutlass.select_( + qz_item < cutlass.Int32(qz_items), qz_i, cutlass.Int32(0) + ) + qz_k_global = rank + qz_i * cutlass.Int32(split) + if cutlass.const_expr(contig): + qz_k_global = rank * cutlass.Int32(my_tiles) + qz_i + qz_t_ld = cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + qz_x_base = ( + qz_t_ld * cutlass.Int32(k_in // 2) + + qz_k_global * cutlass.Int32(CTA_K // 2) + + qz_kb * cutlass.Int32(MMA_INST_K // 2) + ) + qz_wvecs = None + if cutlass.const_expr(din): + # The draft input's raw values (table row of the id / hidden row) and its norm weights. + qz_col = qz_k_global * cutlass.Int32(CTA_K) + qz_kb * cutlass.Int32(MMA_INST_K) + qz_hcol = qz_col - din_hidden_half * cutlass.Int32(din_half_k) + qz_id = din_i.load(idx=qz_t_ld) + qz_hrow = qz_t_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + qz_hrow = din_g.load(idx=qz_t_ld * cutlass.Int32(2)) + qz_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + qz_hrow, + cutlass.select_( + (qz_id >= cutlass.Int32(0)) & (qz_id < din_vocab), + qz_id, + cutlass.Int32(0), + ), + ) + qz_src_w = qz_row * cutlass.Int32(din_half_k // 2) + qz_hcol // cutlass.Int32(2) + qz_vecs = [ + _din_load4( + din_table, + din_hidden, + din_hidden_half, + qz_src_w + cutlass.Int32(qz_v * 4), + ) + for qz_v in range(4) + ] + qz_wvecs = [ + din_w.load( + idx=qz_col // cutlass.Int32(2) + cutlass.Int32(qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + else: + qz_vecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(qz_v * 4), vector_size=4, alignment=16 + ) + for qz_v in range(4) + ] + qz_loaded.append((qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs)) + if cutlass.const_expr((qz_r + 1) % QZ_GROUP == 0 or qz_r + 1 == qz_rounds): + for qz_g in cutlass.range_constexpr(len(qz_loaded)): + qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs = qz_loaded[qz_g] + qz_vals = [] + qz_amax = cutlass.Float32(0.0) + if cutlass.const_expr(din): + qz_rs = din_rs.load( + idx=cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + ) + for qz_v in cutlass.range_constexpr(4): + for qz_e4 in cutlass.range_constexpr(4): + qz_word = cutlass.Int32(qz_vecs[qz_v][qz_e4]) + qz_lo = (qz_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + qz_hi = (qz_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + if cutlass.const_expr(din): + # bf16(x * rs * w), the draft-input kernel's rounding. + qz_wword = cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + qz_lo = _bf16_round( + qz_lo + * qz_rs + * (qz_wword << cutlass.Int32(16)).bitcast(cutlass.Float32) + ) + qz_hi = _bf16_round( + qz_hi + * qz_rs + * (qz_wword & cutlass.Int32(-65536)).bitcast( + cutlass.Float32 + ) + ) + qz_vals.append(qz_lo) + qz_vals.append(qz_hi) + qz_amax = _absmax(_absmax(qz_amax, qz_lo), qz_hi) + qz_bits = (qz_amax * _rcp_approx_ftz(cutlass.Float32(448.0))).bitcast( + cutlass.Int32 + ) + qz_e = ( + (qz_bits >> cutlass.Int32(23)) & cutlass.Int32(0xFF) + ) + cutlass.select_( + (qz_bits & cutlass.Int32(0x7FFFFF)) != cutlass.Int32(0), + cutlass.Int32(1), + cutlass.Int32(0), + ) + qz_e = cutlass.select_(qz_e > cutlass.Int32(254), cutlass.Int32(254), qz_e) + qz_scale = ((cutlass.Int32(254) - qz_e) << cutlass.Int32(23)).bitcast( + cutlass.Float32 + ) + qz_scale = cutlass.select_( + qz_amax != cutlass.Float32(0.0), qz_scale, cutlass.Float32(0.0) + ) + qz_q = [_e4m3x4(qz_vals[4 * qz_j] * qz_scale, qz_vals[4 * qz_j + 1] * qz_scale, + qz_vals[4 * qz_j + 2] * qz_scale, qz_vals[4 * qz_j + 3] * qz_scale) + for qz_j in range(8)] # fmt: skip + if qz_item < cutlass.Int32(qz_items): + if qz_t < num_tokens: + qz_row_w = qz_i * cutlass.Int32( + x_tile_bytes // 4 + ) + qz_t * cutlass.Int32(CTA_K // 4) + qz_c0 = qz_kb * cutlass.Int32(2) + qz_b_words.store( + (qz_q[0], qz_q[1], qz_q[2], qz_q[3]), + idx=qz_row_w + + (qz_c0 ^ (qz_t & cutlass.Int32(7))) * cutlass.Int32(4), + alignment=16, + ) + qz_b_words.store( + (qz_q[4], qz_q[5], qz_q[6], qz_q[7]), + idx=qz_row_w + + ((qz_c0 + cutlass.Int32(1)) ^ (qz_t & cutlass.Int32(7))) + * cutlass.Int32(4), + alignment=16, + ) + smem_sfb.store( + cutlass.Uint8(qz_e), + idx=qz_i * cutlass.Int32(SF_ATOM_BYTES) + + (qz_t % cutlass.Int32(32)) * cutlass.Int32(16) + + (qz_t // cutlass.Int32(32)) * cutlass.Int32(4) + + qz_kb, + ) + qz_loaded = [] + # Generic-proxy stores made visible to the MMA's async-proxy reads, then one arrival for all. + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.barrier_cta_sync(2, thread_count=128) + if qz_tid == cutlass.Int32(0): + prims.mbarrier_arrive(act_full) + if cutlass.const_expr(st_async): + while not cute.arch.mbarrier_test_wait(acc_done.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + # tcgen05.ld widths are powers of two: the accumulator columns in descending power-of-two pieces. + acc_all = [] + ld_col = 0 + for ld_w in (128, 64, 32, 16, 8): + if cutlass.const_expr((mma_n * nacc - ld_col) >= ld_w): + piece = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr( + tmem_ptr_i32.load() + cutlass.Int32(ld_col), 6, cutlass.Float32 + ), + num=ld_w, + ) + acc_all = acc_all + [piece[c] for c in range(ld_w)] + ld_col += ld_w + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + acc = [] + for t in cutlass.range_constexpr(mma_n): + v = cutlass.Float32(acc_all[t]) + for a in cutlass.range_constexpr(1, nacc): + v = v + cutlass.Float32(acc_all[a * mma_n + t]) + acc.append(v) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if tx == 128: + _stamp(timers, timed, bx, 4) + n = m_offset + row + if cutlass.const_expr(split == 1): + for t in cutlass.range_constexpr(mma_n): + if cutlass.Int32(t) < num_tokens: + y.store( + cutlass.BFloat16(cutlass.Float32(acc[t])), idx=cutlass.Int32(t * n_out) + n + ) + if cutlass.const_expr(ar_push > 0): + prims.griddepcontrol(prims.GridDepAction.WAIT) + # The parity of each (token, column slice) slot: the all-reduce's CTA t * 8 + slice flag. All loads are + # issued before the first store so they overlap (one L2 round trip, not one per token). + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_v = cutlass.Float32(acc[t]) + ar_hi = cute.arch.shuffle_sync_bfly(ar_v, offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), _bf16x2_word(ar_v, ar_hi) + ) + else: + owner = row // cutlass.Int32(rows_per_rank) + local = row % cutlass.Int32(rows_per_rank) + if cutlass.const_expr(rope): + # Rank r owns head dims [r * P, (r + 1) * P) and their rotation partners 64 + ..., P = rows / 2. + pairs = rows_per_rank // 2 + owner = (row % cutlass.Int32(64)) // cutlass.Int32(pairs) + local = (row // cutlass.Int32(64)) * cutlass.Int32(pairs) + ( + row % cutlass.Int32(64) + ) % cutlass.Int32(pairs) + if owner != rank: + base = (rank * cutlass.Int32(rows_per_rank) + local) * cutlass.Int32(mb_stride) + if cutlass.const_expr(st_async): + dst = _mapa_u32(mailbox.subview(base).data_ptr(), owner) + mb = _mapa_u32(reduce_ready.data_ptr(), owner) + words = [cutlass.Float32(acc[t]).bitcast(cutlass.Int32) for t in range(mma_n)] + for g in cutlass.range_constexpr(mma_n // 4): + _st_async_v4( + dst + cutlass.Int32(16 * g), + words[4 * g], + words[4 * g + 1], + words[4 * g + 2], + words[4 * g + 3], + mb, + ) + else: + for t in cutlass.range_constexpr(mma_n): + prims.mapa(mailbox.subview(base + cutlass.Int32(t)), owner).store( + cutlass.Float32(acc[t]) + ) + prims.mbarrier_arrive( + prims.mapa(reduce_ready, owner), scope=prims.MemScope.CLUSTER + ) + else: + if cutlass.const_expr(st_async): + while not _test_wait_cluster(reduce_ready.data_ptr(), 0): + pass + else: + while not prims.mbarrier_try_wait_parity( + reduce_ready, 0, scope=prims.MBarrierScope.CLUSTER + ): + pass + if lane == 0: + if row // cutlass.Int32(rows_per_rank) == rank: + if (row % cutlass.Int32(rows_per_rank)) == 0: + _stamp(timers, timed, bx, 5) + totals = [] + for g in cutlass.range_constexpr(mma_n // 4): + sums = [cutlass.Float32(0.0)] * 4 + for q in cutlass.range_constexpr(split): + parts = _ld_shared_v4_f32( + mailbox.subview( + (cutlass.Int32(q * rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g) + ).data_ptr() + ) + sums = [sums[j] + cutlass.select_(rank == cutlass.Int32(q), cutlass.Float32(acc[4 * g + j]), + parts[j]) for j in range(4)] # fmt: skip + for j in cutlass.range_constexpr(4): + totals.append(sums[j]) + if cutlass.const_expr(rope): + # The own-rank mailbox slot is unused: it carries the totals to the RoPE stage. + mailbox.store( + sums[j], + idx=(rank * cutlass.Int32(rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g + j), + ) + elif cutlass.Int32(4 * g + j) < num_tokens: + y.store( + cutlass.BFloat16(sums[j]), + idx=cutlass.Int32((4 * g + j) * n_out) + n, + ) + if cutlass.const_expr(ar_push > 0): + # Every lane of an owner warp is an owner (rows_per_rank >= 32), so lane pairs hold row pairs. + prims.griddepcontrol(prims.GridDepAction.WAIT) + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_hi = cute.arch.shuffle_sync_bfly(totals[t], offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), + _bf16x2_word(totals[t], ar_hi), + ) + prims.barrier_cta_sync(1, thread_count=128) + if cutlass.const_expr(rope): + # Token t of request t / T sits at cache position kv_len - T + t % T. Q / K heads: GPT-NeoX RoPE over + # the 128-wide head, bf16(cos * x -/+ sin * pair) on the bf16 projection; then E4M3 of x * scale, Q to + # rope_q, K / V into the page of that position. + pairs = rows_per_rank // 2 + rp_e = tx - cutlass.Int32(128) + rp_p = rp_e % cutlass.Int32(pairs) + rp_t = rp_e // cutlass.Int32(pairs) + if rp_t < num_tokens: + rp_max_blocks = rope_dims.load(idx=0) + rp_tpb_log2 = rope_dims.load(idx=1) + rp_hq = rope_dims.load(idx=2) + rp_hkv = rope_dims.load(idx=3) + rp_T = rope_dims.load(idx=4) + rp_d = rank * cutlass.Int32(pairs) + rp_p + rp_a = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + rp_p) * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_b = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + cutlass.Int32(pairs) + rp_p) + * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_req = rp_t // rp_T + rp_tok = rope_i.load(idx=rp_req) - rp_T + (rp_t - rp_req * rp_T) + rp_head = n_tile + if rp_head < rp_hq + rp_hkv: + rp_c = rope_cs.load(idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2)) + rp_s = rope_cs.load( + idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2) + + cutlass.Int32(1) + ) + rp_ra = _rope_bits(rp_c, rp_a, -rp_s, rp_b).bitcast(cutlass.Float32) + rp_rb = _rope_bits(rp_c, rp_b, rp_s, rp_a).bitcast(cutlass.Float32) + rp_a = rp_ra + rp_b = rp_rb + rp_qs = rope_scale.load(idx=0) + rp_w = _e4m3x4( + rp_a * rp_qs, rp_b * rp_qs, cutlass.Float32(0.0), cutlass.Float32(0.0) + ) + rp_lo = cutlass.Uint8(rp_w & cutlass.Int32(0xFF)) + rp_hi = cutlass.Uint8((rp_w >> cutlass.Int32(8)) & cutlass.Int32(0xFF)) + if rp_head < rp_hq: + rp_q = (rp_t * rp_hq + rp_head) * cutlass.Int32(128) + rp_d + rope_q.store(rp_lo, idx=rp_q) + rope_q.store(rp_hi, idx=rp_q + cutlass.Int32(64)) + else: + rp_kv = cutlass.select_( + rp_head < rp_hq + rp_hkv, cutlass.Int32(0), cutlass.Int32(1) + ) + rp_kh = rp_head - rp_hq - rp_kv * rp_hkv + rp_page = rope_bt.load( + idx=(rp_req * cutlass.Int32(2) + rp_kv) * rp_max_blocks + + (rp_tok >> rp_tpb_log2) + ) + # 64-bit byte offset: the pool view can exceed 2^31 bytes. + rp_dst = ( + (cutlass.Int64(rp_page) * cutlass.Int64(rp_hkv) + cutlass.Int64(rp_kh)) + << cutlass.Int64(rp_tpb_log2 + cutlass.Int32(7)) + ) + cutlass.Int64( + (rp_tok & ((cutlass.Int32(1) << rp_tpb_log2) - cutlass.Int32(1))) + * cutlass.Int32(128) + + rp_d + ) + rp_addr = rope_pool.data_ptr(0).toint() + rp_dst + cutlass.inttoptr(rp_addr, mem_space=1, dtype=cutlass.Uint8).store(rp_lo) + cutlass.inttoptr( + rp_addr + cutlass.Int64(64), mem_space=1, dtype=cutlass.Uint8 + ).store(rp_hi) + if tx == 128: + _stamp(timers, timed, bx, 6) + if cutlass.const_expr(ar_push == 2): + # The all-reduce polls the pushed words; it launches after this CTA's rows are out, not at grid end. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if warp_id == 4: + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32), + tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n), + ) + + +def _bytes_map(t, cols, rows, box_rows, swizzle): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[cols, rows], + global_strides=[cols // 16], + box_dims=[CTA_K, box_rows], + swizzle=swizzle, + ) + + +def _atom_map(t, atoms): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint16, + global_dims=[SF_ATOM_U16, atoms], + global_strides=[SF_ATOM_BYTES // 16], + box_dims=[SF_ATOM_U16, 1], + swizzle=cuda.TensorMapSwizzle.none, + ) + + +@cute.jit +def m3_mx_gemv( + w: cute.Tensor, # [N, K] E4M3 as uint8, K contiguous + wsf: cute.Tensor, # weight scales, 128x4-swizzled, as uint8 + x: cute.Tensor, # [M, K] E4M3 as uint8, M <= 8 + xsf: cute.Tensor, # activation scales, 128x4-swizzled (atom k at byte 512 k), as uint8 + y: cute.Tensor, # [M * N] bf16 + x_words: cute.Tensor, # x_bf16: bf16 x [M, K] as int32 words (x / xsf are then unused placeholders) + timers: cute.Tensor, # [grid * 8] int64, written when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[int], + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, + din_i: cute.Tensor, # din / rope inputs: see m3_mx_gemv_kernel (placeholders when off) + din_table: cute.Tensor, + din_hidden: cute.Tensor, + din_w: cute.Tensor, + din_f: cute.Tensor, + din_vocab: cutlass.Int32, + din_g: cute.Tensor, + din_res: cute.Tensor, + din_flags: cutlass.Int32, + rope_cs: cute.Tensor, + rope_i: cute.Tensor, + rope_bt: cute.Tensor, + rope_q: cute.Tensor, + rope_pool: cute.Tensor, + rope_scale: cute.Tensor, + rope_dims: cute.Tensor, + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], +) -> None: + k_tiles = num_k_tiles(k_in) + tma_desc_w = _bytes_map(w, k_in, n_out, CTA_M, cuda.TensorMapSwizzle.s128b) + tma_desc_wsf = _atom_map(wsf, (n_out // CTA_M) * k_tiles) + tma_desc_x = _bytes_map(x, k_in, num_tokens, mma_n, cuda.TensorMapSwizzle.s128b) + tma_desc_xsf = _atom_map(xsf, k_tiles) + m3_mx_gemv_kernel( + tma_desc_w, tma_desc_wsf, tma_desc_x, tma_desc_xsf, y, x_words, timers, num_tokens, n_out, k_in, split, ring, + trigger_early, contig, timed, nacc, a_tmem, st_async, x_bf16, ar_ws, ar_flags, ar_rank, ar_push, ar_world, + ar_maxt, mma_n, + din_i, din_table, din_hidden, din_w, din_f, din_vocab, din_g, din_res, din_flags, rope_cs, + rope_i, rope_bt, rope_q, rope_pool, rope_scale, rope_dims, din, rope, + ).launch( + grid=((n_out // CTA_M) * split, 1, 1), + block=(THREADS, 1, 1), + cluster=(split, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/mx_gemv/op.py new file mode 100644 index 000000000000..84808d7c61bf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/mx_gemv/op.py @@ -0,0 +1,472 @@ +# 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. +"""MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) at the TP8 per-rank projection shapes. + +Compiled on the first call per (N, K, split, ring, early trigger, PDL, push, MMA width), which must happen outside +CUDA-graph capture, or ahead of any call by ``precompile``; once ``freeze`` closes compiles, a call whose kernel is not +compiled declines instead. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +MMA_WIDTHS = ( + 8, + 16, + 32, + 48, +) # token columns of the compiled MMA; a call uses the smallest that holds M +_SMEM_LIMIT = 220 * 1024 # capacity: fixed: the device's shared-memory limit + + +def _mma_n(m: int) -> int: + return next(w for w in MMA_WIDTHS if m <= w) + + +def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: + return ( + ring * (128 * 128 + 512) + + tiles * (mma_n * 128 + 512) + + split * (128 // split) * (mma_n + 4) * 4 + + 1024 + ) + + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + +# (N, K) -> (split, ring[, contig, nacc, a_tmem, st_async]) of the TP8 per-rank projections: GB300 exposed time +# behind a PDL predecessor. The optional fields default to _CONTIG, _NACC, _A_TMEM, _ST_ASYNC. The sparse-layer QKV +# splits K over 8 CTAs (96 CTAs at 1 token; bit-identical to splits 4 and 2). +_CONFIGS: Dict[Tuple[int, int], tuple] = { + # sparse-layer QKV + index projection + (1536, 6144): (8, 8), + # dense-layer QKV + (1280, 6144): (4, 12), + # o_proj + (6144, 1024): (4, 8), + # dense-layer gate_up + (3072, 6144): (4, 12), + # dense-layer down + (6144, 1536): (2, 12), +} +# Per-token-count configs (N, K, M): at 2 tokens the sparse-layer QKV splits K over 4 CTAs, A operand from shared +# memory. +_CONFIGS_BY_M: Dict[Tuple[int, int, int], tuple] = { + (1536, 6144, 2): (4, 12, 1, 1, 0, 1), +} +# Above 8 tokens (spec-decode verify rows): the same 4-CTA split for the sparse-layer QKV. +_CONFIGS_WIDE: Dict[Tuple[int, int], tuple] = { + (1536, 6144): (4, 12, 1, 1, 0, 1), +} + +# Defaults of the optional config fields: contiguous k-tile ranges per split rank, one TMEM accumulator, the A +# operand staged in TMEM (when its columns fit), st.async split-K partial exchange. +_CONTIG = True +_NACC = 1 +_A_TMEM = True +_ST_ASYNC = True +_ensured = set() + + +_DUMMY: Dict[torch.device, torch.Tensor] = {} + + +def _dummy_timers(device) -> torch.Tensor: + """Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on.""" + t = _DUMMY.get(device) + if t is None: + t = _DUMMY[device] = torch.zeros(8, dtype=torch.int64, device=device) + return t + + +def _arg(t: torch.Tensor, align: int = 16): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def get_config(n: int, k: int, m: Optional[int] = None) -> Optional[tuple]: + cfg = _CONFIGS_BY_M.get((n, k, m)) if m is not None else None + if cfg is None and m is not None and m > 8: + cfg = _CONFIGS_WIDE.get((n, k)) + cfg = cfg or _CONFIGS.get((n, k)) + # The split-K epilogue gives each of the split CTAs 128 // split rows of a 128-row output tile. + if cfg is None or n % 128 != 0 or k % 128 != 0 or (k // 128) % cfg[0] != 0 or 128 % cfg[0] != 0: + return None + return cfg + + +def _declines(n: int, k: int, m: int, x_bf16: bool, cfg: tuple) -> bool: + """Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds.""" + split = cfg[0] + if not 0 < m <= MAX_TOKENS or 128 % split != 0 or (k // 128) % split != 0: + return True + tiles = k // 128 // split + mma_n = _mma_n(m) + # 17..32 rows at K < 1024 decline (an unsplit K below 1024 is slower than the FlashInfer GEMM there); no TP8 + # projection has K < 1024. + # The in-kernel bf16 -> MXFP8 quantization (x_bf16) costs one 32-element block per thread and round on 128 threads; + # beyond 8 rounds per CTA, mxfp8_quantize + this GEMV (or the fallback GEMM) is faster. + if (16 < m <= 32 and k < 1024) or (x_bf16 and tiles * 4 * mma_n > 8 * 128): + return True + return _smem_bytes(min(cfg[1], tiles, 2), tiles, split, mma_n) > _SMEM_LIMIT + + +def serves(n: int, k: int, m: int, x_bf16: bool) -> bool: + """Whether the GEMV with its default config takes an [M, K] activation (bf16 or MXFP8) against an [N, K] weight.""" + cfg = get_config(n, k, m) + return cfg is not None and not _declines(n, k, m, x_bf16, cfg) + + +def _resolve( + m: int, + n: int, + k: int, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple], +) -> Optional[tuple]: + """The compile key of an [M, K] x [N, K]^T call and its runtime push words, or None when the call is declined. + + ``ar`` is the push tuple of ``_run``; only its (world, bucket tokens, mode) enter the key, the addresses and the + rank are runtime arguments.""" + cfg = cfg or get_config(n, k, m) + if cfg is None or _declines(n, k, m, x_bf16, cfg): + return None + split, ring = cfg[0], cfg[1] + contig = bool(cfg[2]) if len(cfg) > 2 else _CONTIG + nacc = int(cfg[3]) if len(cfg) > 3 else _NACC + a_tmem = bool(cfg[4]) if len(cfg) > 4 else _A_TMEM + st_async = bool(cfg[5]) if len(cfg) > 5 else _ST_ASYNC + ring = min(ring, k // 128 // split) + tiles = k // 128 // split + mma_n = _mma_n(m) + while ring > 2 and _smem_bytes(ring, tiles, split, mma_n) > _SMEM_LIMIT: + ring -= 1 + if a_tmem: + from .m3_mx_gemv_kernel import tmem_cols + + if tmem_cols(ring, nacc, True, tiles, mma_n) > 512: + a_tmem = False + # ar = (multicast workspace address, flags address, rank, world, bucket tokens, mode): push the output rows into + # the next m3_mc_allreduce call's exchange slots (mode 2 also defers the dependents' launch until after the push). + ar_ws, ar_flags, ar_rank, ar_mode, ar_world, ar_maxt = 0, 0, 0, 0, 1, 8 + if ar is not None and split <= 4 and n % 2 == 0: + ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_mode = ar + if ar_mode == 2: + trigger_early = False + if use_pdl is None: + use_pdl = get_env_enable_pdl() + # The kernel's ``timed`` flag (per-CTA milestone timers) is always off. + key = ( + n, + k, + split, + ring, + trigger_early, + contig, + False, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + ) + return key, (ar_ws, ar_flags, ar_rank) + + +def _compile(key: tuple, args: tuple, m: int, runtime_ar: tuple, stream, extra: tuple): + """cute.compile the kernel of ``key`` (see ``_resolve``) on these argument views; cached.""" + import cutlass.cute as cute + + from . import m3_mx_gemv_kernel as kernel + + ( + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + ) = key + ar_ws, ar_flags, ar_rank = runtime_ar + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.m3_mx_gemv, + *args, + m, + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_ws, + ar_flags, + ar_rank, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + stream, + *extra, + False, + False, + ) + return fn + + +def _views( + x: torch.Tensor, xsf: Optional[torch.Tensor], weight, weight_scale, y, x_bf16: bool +) -> tuple: + k = x.shape[1] + if x_bf16: + x_bytes = x.view(torch.uint8) + xsf = _placeholder_scales(k, x.device) + x_words = x.view(torch.int32) + else: + x_bytes = x.view(torch.uint8) + x_words = _placeholder_words(x.device) + return ( + _arg(weight.view(torch.uint8)), + _arg(weight_scale.view(torch.uint8).reshape(-1)), + _arg(x_bytes), + _arg(xsf.view(torch.uint8).reshape(-1)), + _arg(y.view(-1) if y.numel() else _placeholder_words(x.device).view(torch.bfloat16)), + _arg(x_words.reshape(-1)), + _arg(_dummy_timers(x.device)), + ) + + +def _run( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight: torch.Tensor, + weight_scale: torch.Tensor, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + global last_pushed, last_uncompiled + last_pushed = False + last_uncompiled = None + import cuda.bindings.driver as cuda_driver + + m, k = x.shape + n = weight.shape[0] + if weight.shape[1] != k: + return None + resolved = _resolve(m, n, k, x_bf16, trigger_early, cfg, use_pdl, ar) + if resolved is None: + return None + key, runtime_ar = resolved + fn = _compiled.get(key) + if fn is None: + if _frozen[0]: + # Compiles are closed (the model built every kernel it dispatches to before warmup): decline, the caller + # reports the miss. + last_uncompiled = key + return None + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_mx_gemv must run once per shape outside CUDA-graph capture first" + ) + y = torch.empty(m, n, dtype=torch.bfloat16, device=x.device) + args = _views(x, xsf, weight, weight_scale, y, x_bf16) + extra = _extra_args(x.device) + stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + if fn is None: + fn = _compile(key, args, m, runtime_ar, stream, extra) + ar_ws, ar_flags, ar_rank = runtime_ar + fn(*args, m, ar_ws, ar_flags, ar_rank, stream, *extra) + last_pushed = key[11] > 0 + return y + + +def precompile( + weight: torch.Tensor, + weight_scale: torch.Tensor, + m: int, + x_bf16: bool, + use_pdl: Optional[bool] = None, + ar: Optional[Tuple[int, int, int]] = None, +) -> Optional[tuple]: + """Compile, without launching, the kernel an M-row call against ``weight`` takes; returns its key, or None when + such a call is declined. ``ar`` = (world, bucket tokens, mode) of a pushing call (its addresses and rank are + runtime arguments, so none is needed here). Runs outside capture, even while compiles are frozen.""" + import cuda.bindings.driver as cuda_driver + + n, k = weight.shape + push = None if ar is None else (0, 0, 0, *ar) + resolved = _resolve(m, n, k, x_bf16, True, None, use_pdl, push) + if resolved is None: + return None + key, runtime_ar = resolved + if key in _compiled: + return key + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_mx_gemv.precompile allocates and compiles: call it outside capture") + dev = weight.device + x = torch.zeros(m, k, dtype=torch.bfloat16 if x_bf16 else torch.float8_e4m3fn, device=dev) + xsf = None if x_bf16 else torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) + y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) + args = _views(x, xsf, weight, weight_scale, y, x_bf16) + stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) + _compile(key, args, m, runtime_ar, stream, _extra_args(dev)) + return key + + +def freeze(frozen: bool = True) -> None: + """Close (or reopen) compiles: while frozen, a call whose kernel is not compiled yet declines (returns None and + sets ``last_uncompiled``) instead of compiling. ``precompile`` and ``ensure_compiled`` still compile.""" + _frozen[0] = bool(frozen) + + +def frozen() -> bool: + return _frozen[0] + + +def is_compiled(key: tuple) -> bool: + return key in _compiled + + +def _extra_args(device) -> tuple: + """Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off).""" + i32 = _placeholder_words(device) + f32 = i32.view(torch.float32) + u8 = i32.view(torch.uint8) + d = (_arg(i32), _arg(i32), _arg(i32), _arg(i32), _arg(f32), 0, _arg(i32), _arg(i32), 0) + r = (_arg(f32), _arg(i32), _arg(i32), _arg(u8), _arg(u8), _arg(f32), _arg(i32)) + return d + r + + +def m3_mx_gemv( + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``xq`` E4M3 [M <= 8, K] with ``xsf`` its 128x4-swizzled UE8M0 scales; ``weight`` E4M3 [N, K] with + ``weight_scale`` swizzled. Returns bf16 [M, N], or None when the shape is not served.""" + return _run(xq, xsf, weight, weight_scale, False, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_bf16( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """The same GEMV for a bf16 ``x`` [M <= 8, K] (contiguous), quantized to MXFP8 inside the kernel exactly as + torch.ops.trtllm.mxfp8_quantize(x, True) does.""" + if x.dtype != torch.bfloat16 or x.dim() != 2 or not x.is_contiguous() or x.data_ptr() % 16 != 0: + return None + return _run(x, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar) + + +_PLACEHOLDERS: Dict[tuple, torch.Tensor] = {} +# Whether the last call pushed its output into an all-reduce (see _run's ``ar``). +last_pushed = False +# The key of the last call declined because compiles are frozen and its kernel was not compiled (see ``freeze``). +last_uncompiled: Optional[tuple] = None +_frozen = [False] + + +def _placeholder_scales(k: int, device) -> torch.Tensor: + key = ("sf", k, device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(k // 128 * 512, dtype=torch.uint8, device=device) + return t + + +def _placeholder_words(device) -> torch.Tensor: + key = ("w", device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(16, dtype=torch.int32, device=device) + return t + + +def ensure_compiled(weight: torch.Tensor, weight_scale: torch.Tensor, use_pdl: bool) -> None: + """Compiles the kernel of this weight's shape outside capture (a one-token dummy call).""" + if ( + weight.dim() != 2 + or weight.dtype != torch.float8_e4m3fn + or get_config(*weight.shape) is None + ): + return + n, k = weight.shape + if (n, k, use_pdl) in _ensured: + return + _ensured.add((n, k, use_pdl)) + xq = torch.zeros(1, k, dtype=torch.float8_e4m3fn, device=weight.device) + xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=weight.device) + m3_mx_gemv(xq, xsf, weight, weight_scale, use_pdl=use_pdl) + for width in MMA_WIDTHS[1:]: + m3_mx_gemv( + torch.zeros(width, k, dtype=torch.float8_e4m3fn, device=weight.device), + xsf, + weight, + weight_scale, + use_pdl=use_pdl, + ) + m3_mx_gemv_bf16( + torch.zeros(1, k, dtype=torch.bfloat16, device=weight.device), + weight, + weight_scale, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/__init__.py new file mode 100644 index 000000000000..acec8fdbb26f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Normalization entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/m3_embedding_firstnorm.py new file mode 100644 index 000000000000..1842171bde1a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/m3_embedding_firstnorm.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-step embedding gather and the first layer's input RMSNorm in one CuTe DSL launch, into caller buffers.""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm + +__all__ = ["DecodeEmbedding", "m3_embedding_firstnorm"] + +VOCAB = 200064 # MiniMax-M3's vocabulary; the decode table is the full (replicated) embedding + + +class M3EmbeddingFirstnorm(OpWrapper): + """Write ``table[ids]`` into ``raw`` (bf16 [M, 6144], a zero row for an id outside ``[0, V)``) and + ``raw * rsqrt(mean(raw^2) + eps) * weight`` into ``out``; returns ``(out, raw)`` -- the first layer's normed input + and its residual. ``ids`` int32 or int64 [M], M = 1..8; ``table`` bf16 [V, 6144] contiguous; ``weight`` bf16 + [6144]. The first call per (M, id dtype, table shape, PDL) compiles (``DecodeEmbedding`` does it at build).""" + + ARCHS = frozenset({Arch.SM_103}) + + # The target builds every (rows 1..8, id dtype) owner; the kernel is one CTA per row, so the rows differ only in + # grid size: the ends of the range and both id dtypes. + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, int32 ids", spec=dict(rows=1, id_dtype=torch.int32)), + Cell(why="eight requests, int32 ids", spec=dict(rows=8, id_dtype=torch.int32)), + Cell(why="one request, int64 ids", spec=dict(rows=1, id_dtype=torch.int64)), + Cell( + why="six rows (one request verifying a draft), int64 ids", + spec=dict(rows=6, id_dtype=torch.int64), + ), + # Every other row count the decode dispatch takes (dispatch.EMBED_FIRSTNORM), both id dtypes. + *( + Cell( + why=f"{rows} rows: a decode row count between the ones above", + spec=dict(rows=rows, id_dtype=torch.int32 if rows % 2 else torch.int64), + ) + for rows in (2, 3, 4, 5, 7) + ), + ) + + note = """ + The gathered rows are exact. The norm is flashinfer's RMSNorm arithmetic in its reduction order with a fast + rsqrt, where `reference` sums in fp64 and takes torch's rsqrt: within the band the test derives, not bit for bit. + An id outside [0, V) gives a zero row and a zero norm (torch's embedding would raise). No accepted input is known + to compute wrongly: the host side rejects any other dtype, shape or a strided buffer. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return m3_embed_norm(ids, table, weight, eps, raw, out, use_pdl) + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """New tensors (``raw`` / ``out`` are not written); `use_pdl` is ignored.""" + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] * valid.unsqueeze(1).to( + table.dtype + ) + x = rows.double() + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + return (x * rstd * weight.double()).to(torch.bfloat16), rows + + +m3_embedding_firstnorm = M3EmbeddingFirstnorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/__init__.py new file mode 100644 index 000000000000..d25c41dc9557 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Speculative-decoding entries this target alone calls: the Eagle3 verify ops, the capture / fc_norm RMSNorms and +the step-input staging.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_capture_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_capture_norm.py new file mode 100644 index 000000000000..131269ae7d07 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_capture_norm.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Eagle3 hidden-state capture with the drafter's fc_norm folded in (``m3_spec_tail.capture_norm``).""" + +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 + +from .m3_chunk_norm import rmsnorm_rows + + +class M3CaptureNorm(OpWrapper): + """``dst [M, H] = rmsnorm(bf16(h + r)) * w`` (``r`` optional), the rounding chunk_norm applies, into a capture + buffer slice whose row stride may exceed H. Returns ``dst``.""" + + ARCHS = frozenset({Arch.SM_103}) + + # The target's captured layers (hidden 6144) into the 3-layer capture buffer (row stride 3 x 6144), over the + # verify rows of draft_len 3 (4 .. 32 rows); the residual is present for every captured decoder layer. + CELLS: tuple[Cell, ...] = ( + Cell( + why="4 verify rows of one request, layer slot 0, residual added", + spec=dict(m=4, slot=0, residual=True), + ), + Cell( + why="cap 8 (32 rows) into the last layer slot", spec=dict(m=32, slot=2, residual=True) + ), + Cell( + why="no residual (a capture point without one): plain RMSNorm", + spec=dict(m=8, slot=1, residual=False), + ), + ) + + note = """ + fp32 statistic (block sum, rsqrtf), one bf16 rounding after the bf16 sum h + r; the band states the occasional + one-ulp move against torch's statistic. + """ + + def __call__( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_capture_norm(h, r, w, eps, dst) + return dst + + def reference( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + """bf16 sum, then the normalized row; ``dst`` only says where the op writes.""" + x = (h.float() + r.float()).to(torch.bfloat16) if r is not None else h + return rmsnorm_rows(x, w, eps) + + def is_valid( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> None: + # The norm weight is read as H dense bf16 values; a strided view (a column of a stacked weight) is read + # past its elements without an error. + assert w.is_contiguous(), "w must be dense [H]" + + +m3_capture_norm = M3CaptureNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_chunk_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_chunk_norm.py new file mode 100644 index 000000000000..f6adb2207a20 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_chunk_norm.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 fc_norm over the captured layers in one launch: per-chunk RMSNorm of the concatenated hidden states +(``m3_spec_tail.chunk_norm``).""" + +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 + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3ChunkNorm(OpWrapper): + """``dst[r, c*H:(c+1)*H] = rmsnorm(src[r, c*H:(c+1)*H]) * w[c]`` for the C chunks of ``w [C, H]``; ``src`` may + have a row stride wider than C * H (the capture buffer). Writes ``dst [M, C*H]`` and returns it.""" + + ARCHS = frozenset({Arch.SM_103}) + + # The Eagle3 capture buffer of the target (3 captured layers x hidden 6144, buffer rows wider than the layers + # when the final hidden is also captured), over the verify rows of draft_len 3: 4 / 8 / 32 rows. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request's 4 verify rows, dense [4, 3 x 6144]", + spec=dict(m=4, c=3, h=6144, extra_cols=0), + ), + Cell( + why="cap 8 (32 rows) read from the capture buffer with a fourth (final-hidden) slot in the row stride", + spec=dict(m=32, c=3, h=6144, extra_cols=6144), + ), + ) + + note = """ + fp32 statistic per chunk (block sum, rsqrtf), one bf16 rounding: an output moves by one bf16 ulp now and then + against torch's statistic; the cells' band states it. + """ + + def __call__( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + torch.ops.trtllm.minimax_m3_chunk_norm(src, w, eps, dst, pdl) + return dst + + def reference( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + """Each chunk normalized on its own; ``dst`` only says where the op writes, ``pdl`` shapes the launch.""" + c, h = w.shape + x = src[:, : c * h].reshape(src.shape[0], c, h) + return rmsnorm_rows(x, w, eps).reshape(src.shape[0], c * h) + + def is_valid( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ) -> None: + c, h = w.shape + assert w.is_contiguous(), "w must be dense [C, H] (chunk c read at c * H)" + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + + +m3_chunk_norm = M3ChunkNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_draft_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_draft_input.py new file mode 100644 index 000000000000..75f7a44dfcfb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_draft_input.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 draft pass's input in one launch: embedding lookup, both RMSNorms and the concat +(``m3_spec_tail.draft_input``).""" + +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 + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3DraftInput(OpWrapper): + """``out [M, 2H] = [rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden]`` and ``residual [M, H] = + hidden[rows]`` (rows default to 0..M-1) in one launch; ids outside the table embed to zero (the masked lookup of + a vocab-local embedding). Returns ``(out, residual)``.""" + + ARCHS = frozenset({Arch.SM_103}) + + # Every pass of the fused draft loop (dispatch.DRAFT_INPUT) at draft_len 3, drafter hidden 6144, eps 1e-6: the + # first pass of 1..8 requests (4..32 rows, the carried hidden read in place) and the later passes of 1..8 requests + # (one row each; the second pass gathers the last accepted rows from the first pass's 4 rows per request). + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"first pass of {n} request(s): {4 * n} rows, hidden read in place", + spec=dict(m=4 * n, h=6144, gather=False, invalid_id=False), + ) + for n in range(1, 9) + ), + *( + Cell( + why=f"second pass of {n} request(s): the last accepted rows gathered from {4 * n} rows", + spec=dict(m=n, h=6144, gather=True, invalid_id=False), + ) + for n in range(1, 9) + ), + Cell( + why="third pass of 3 requests: one row each, hidden read in place", + spec=dict(m=3, h=6144, gather=False, invalid_id=False), + ), + Cell( + why="one row with an out-of-table id: the embedding half is zero", + spec=dict(m=1, h=6144, gather=False, invalid_id=True), + ), + ) + + note = """ + The arithmetic is the in-tree RMSNorm's (fp32 statistic, (x * rsqrt) * w, one bf16 rounding); the statistic's + summation order differs from FlashInfer's, which moves an output by one bf16 ulp now and then (the test's band, + also held against the in-tree modules). The out-of-table id is the masked lookup of a vocab-local embedding; the + drafter's embedding is the whole table, so its ids are always in range. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + torch.ops.trtllm.minimax_m3_draft_input( + ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, pdl + ) + return out, residual + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """The two normalized halves (fp32 statistic, one bf16 rounding) and the gathered hidden rows. ``out`` / + ``residual`` only say where the op writes; ``pdl`` only shapes the launch.""" + valid = (ids >= 0) & (ids < table.shape[0]) + emb = rmsnorm_rows(table[torch.where(valid, ids, 0).long()], w_embed, eps_embed) + emb = torch.where(valid[:, None], emb, torch.zeros_like(emb)) + hid = hidden[rows.long()] if rows is not None else hidden[: ids.shape[0]] + ref = torch.cat([emb, rmsnorm_rows(hid, w_hidden, eps_hidden)], dim=-1) + return ref, (hid.clone() if residual is not None else None) + + def is_valid( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> None: + h = table.shape[1] + # Rows are read as hidden + rows[r] * H and table + id * H: both must be dense H-wide rows. + assert table.is_contiguous() and hidden.is_contiguous(), ( + "table / hidden must be dense [*, H]" + ) + if rows is not None: + assert bool(((rows >= 0) & (rows < hidden.shape[0])).all()), "rows must index hidden" + else: + assert hidden.shape[0] >= ids.shape[0], "hidden must hold a row per id" + # Vector loads of 16 bytes: H must be a multiple of 8 within the register tile, and the weights aligned. + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + assert residual is None or ( + residual.is_contiguous() and residual.shape == (ids.shape[0], h) + ), "residual must be contiguous [M, H]" + + +m3_draft_input = M3DraftInput() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_draft_step.py new file mode 100644 index 000000000000..3f00def536b6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_draft_step.py @@ -0,0 +1,359 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One pass of the fused greedy Eagle3 draft loop: the drafter's vocab-sharded BF16 head (streamed as BF16 fragments), +the cross-rank argmax and the loop's per-pass bookkeeping in one kernel (``m3_draft_step.draft_step``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from .m3_shard_accept import exchange_tokens, local_argmax, make_keys + +#: Requests (head rows) per pass, and exchange rows per call half of the workspace. +MAX_TOKENS = 8 +#: 16-row weight tiles one CTA can own (the kernel's shared-memory accumulator). +MAX_TILES_PER_CTA = 32 +#: Draft length of the fused loop, and the drafter's TP8 head shard (200064-id vocabulary) and hidden size. +DRAFT_LEN = 3 +DRAFT_SHARD = 25008 +#: Ranks of the TP group the argmax is exchanged over. +WORLD = 8 +HIDDEN = 6144 + + +def bf16_fragments(weight: torch.Tensor) -> torch.Tensor: + """BF16 [N, K] -> the same values in A-fragment order [N / 16, K / 16, 256]. + + One lane's 16 bytes of a 16-row tile's k16 step are the m16n8k16 A fragment: element (row h * 8 + g, k ks * 16 + + hk * 8 + jj * 2 + e) goes to [tile][ks][g][jj][hk][h][e] (lane = g * 4 + jj). + """ + n, k = weight.shape + assert n % 16 == 0 and k % 16 == 0 and weight.dtype == torch.bfloat16, (n, k, weight.dtype) + q = weight.view(n // 16, 2, 8, k // 16, 2, 4, 2).permute(0, 3, 2, 5, 4, 1, 6).contiguous() + return q.view(n // 16, k // 16, 256) + + +def from_fragments(wf: torch.Tensor) -> torch.Tensor: + """The inverse of ``bf16_fragments``: [N / 16, K / 16, 256] -> [N, K].""" + tiles, ks, _ = wf.shape + return ( + wf.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, ks * 16) + ) + + +def head_weight( + wf: torch.Tensor, + k: int, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The BF16 head shard [N, K] the kernel multiplies by: the fragment-ordered copy. The op also streams a 12-bit + coded head (``exc_off`` given); this target never builds one, and nothing here certifies it.""" + if exc_off is not None: + raise NotImplementedError("the tp8 target streams the BF16 fragments only") + return from_fragments(wf) + + +def shard_logits( + x: torch.Tensor, + wf: torch.Tensor, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The in-tree head's BF16 logits of the rows x [M, K] against the shard (``head_weight``): the fp32 product, + rounded once to BF16 (the BF16 GEMM's output), as fp32.""" + w = head_weight(wf, x.shape[1], exc_off, exc_pos, exc_w, exp_base) + with true_fp32_matmul(): + logits = x.float() @ w.float().t() + return logits.to(torch.bfloat16).float() + + +class M3DraftStep(OpWrapper): + """One fused greedy draft pass over ``rows <= 8`` requests (one head row each). + + Head input ``x[x_rows[b]]`` (pass 0: the last accepted verify row of each request) or ``x[b]``. The head is this + rank's BF16 lm_head shard in fragment order (``m3_draft_step.bf16_fragments``, global ids ``vocab_offset ..``), + (the op also takes a 12-bit coded stream with ``exc_off`` / ``exc_pos`` / ``exc_w`` / ``exp_base``, which this + target never builds); each logit is the fp32 dot product + rounded to BF16 and the row's token is the global argmax over all ranks' + shards (ties, -0 / +0 included, to the lowest global id), exchanged through the multicast slots + ``[2][8][world]``. The last CTA then writes ``draft_out[b, step]`` and ``ids_next[b]`` (the token), ``pos_next[b] + = pos_src[x_rows[b] or b] + 1``, and advances the draft KV lengths as the in-tree draft loop does between passes + (pass 0 with ``num_acc``: ``kv_lens -= draft_len - num_acc`` and ``seq_lens = 1``; later passes: ``kv_lens += + 1``). ``out_keys`` (optional) receives each row's winning key. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # Every pass of the fused draft loop (dispatch.DRAFT_STEP): draft_len 3 at 1..8 requests, the drafter's + # full-vocabulary BF16 head (25008 rows per rank of 8, hidden 6144). Pass 0 reads the verify rows' hidden (4 per + # request) through x_rows and rolls the KV back; the later passes read one row per request. + _PASSES: tuple[Cell, ...] = ( + *( + Cell( + why=f"pass 0 at {n} request(s): head rows gathered from {4 * n} verify rows, KV rollback, seq_lens = 1", + spec=dict( + nreq=n, + draft_len=DRAFT_LEN, + step=0, + rank=n % WORLD, + world=WORLD, + vs=DRAFT_SHARD, + k=HIDDEN, + ), + ) + for n in range(1, MAX_TOKENS + 1) + ), + Cell( + why="pass 1 at two requests: one row each, kv_lens + 1, rank 7's ids", + spec=dict( + nreq=2, draft_len=DRAFT_LEN, step=1, rank=7, world=WORLD, vs=DRAFT_SHARD, k=HIDDEN + ), + ), + Cell( + why="last pass (step 2) at 8 requests, rank 0", + spec=dict( + nreq=8, draft_len=DRAFT_LEN, step=2, rank=0, world=WORLD, vs=DRAFT_SHARD, k=HIDDEN + ), + ), + Cell( + why="pass 2 at 5 requests with exact cross-rank ties (the lower global id wins)", + spec=dict( + nreq=5, + draft_len=DRAFT_LEN, + step=2, + rank=5, + world=WORLD, + vs=DRAFT_SHARD, + k=HIDDEN, + ties=True, + ), + ), + ) + # The BF16 fragments at every pass shape. + CELLS: tuple[Cell, ...] = tuple( + Cell(why=f"bf16: {c.why}", spec=dict(c.spec, coded=False)) for c in _PASSES + ) + + note = """ + The head is the in-tree BF16 head's arithmetic (BF16 weights and activations, fp32 accumulation, a BF16 logit) + and the argmax is the in-tree one over BF16 logits, so a token differs from the in-tree path's only where the fp32 + sum, accumulated in another order, rounds to the other side of a BF16 boundary and that decides the row's maximum. + The cells' rows are decidable and the peers' keys sit 0.25 away (or tie exactly); the edge tests (forced BF16 + ties, -0 / +0, top-2 one fp32 step apart that round to one BF16 value, adjacent BF16 values) use integer data + whose fp32 sums are exact in any order and hold the kernel to the in-tree path bit for bit. The cross-rank + exchange is certified on one device with the peers' keys written into the slots (see m3_shard_accept). + """ + + def __call__( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + ) -> None: + torch.ops.trtllm.minimax_m3_draft_step( + x, + x_rows, + rows, + wf, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + num_ctas, + pdl, + out_keys, + exc_off, + exc_pos, + exc_w, + exp_base, + ) + + def reference( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + ) -> dict[str, torch.Tensor]: + """Every buffer the pass writes, computed on copies: "logits" (this rank's BF16 logits as fp32), "local_keys" + (this rank's key per row), "keys" (the winning key), "tokens", "draft_out", "ids_next", "pos_next", "kv_lens" + and "seq_lens". + + Reads ``ws_uc``, ``counters`` and ``kv_lens`` before the call. ``keys``, ``ws_mc_ptr``, ``num_ctas``, ``pdl`` + and ``out_keys`` do not change the result. + """ + h = x[x_rows.long()] if x_rows is not None else x[:rows] + logits = shard_logits(h, wf, exc_off, exc_pos, exc_w, exp_base) + best, idx = local_argmax(logits, vocab_offset) + local_keys = make_keys(best, idx) + tokens = exchange_tokens(local_keys, ws_uc, counters, rank, world, MAX_TOKENS) + half = int(counters[1].item()) & 1 + slots = ws_uc.view(torch.int64)[ + half * MAX_TOKENS * world : (half * MAX_TOKENS + rows) * world + ] + slots = slots.view(rows, world).clone() + slots[:, rank] = local_keys + sign = torch.iinfo(torch.int64).min + win = (slots ^ sign).max(dim=1).values ^ sign + d_out = draft_out.clone() + d_out[:rows, step] = tokens + src = x_rows.long() if x_rows is not None else torch.arange(rows, device=x.device) + lens = kv_lens.clone() + sl = seq_lens.clone() if seq_lens is not None else None + if step == 0 and num_acc is not None: + lens[:rows] -= draft_len - num_acc[:rows] + sl[:rows] = 1 + elif step > 0: + lens[:rows] += 1 + out = { + "logits": logits, + "local_keys": local_keys, + "keys": win, + "tokens": tokens, + "draft_out": d_out, + "ids_next": tokens.clone(), + "pos_next": pos_src[src] + 1, + "kv_lens": lens, + } + if sl is not None: + out["seq_lens"] = sl + return out + + def is_valid( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + ) -> None: + assert not bool(keys[:rows].any()), "keys must start zeroed" + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + # The kernel reads the head row x + x_rows[b] * K and the position pos_src[x_rows[b]] with no bound. + if x_rows is not None: + assert bool(((x_rows[:rows] >= 0) & (x_rows[:rows] < x.shape[0])).all()), ( + "x_rows must index x" + ) + assert bool((x_rows[:rows] < pos_src.shape[0]).all()), ( + "pos_src must cover the head rows" + ) + else: + assert pos_src.shape[0] >= rows, "pos_src must cover the head rows" + # Each CTA accumulates at most MAX_TILES_PER_CTA tiles in shared memory. + tiles = wf.shape[0] if exc_off is None else exc_off.numel() - 1 + assert -(-tiles // num_ctas) <= MAX_TILES_PER_CTA, "too many weight tiles per CTA" + if exc_off is not None: + # The exceptions are read per tile in list order with no bound against the coded stream. + assert int(exc_off[-1].item()) == exc_pos.numel(), ( + "exc_off must end at the exception count" + ) + assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" + # A NaN logit maps to a key above every finite one and wins the row silently. + h = x[x_rows.long()] if x_rows is not None else x[:rows] + assert not bool(torch.isnan(h).any()), "NaN head inputs win the argmax" + if num_acc is not None: + na = num_acc[:rows] + assert bool(((na >= 1) & (na <= draft_len + 1)).all()), ( + "num_acc counts must lie in [1, draft_len + 1]" + ) + + +m3_draft_step = M3DraftStep() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_force_vec.py new file mode 100644 index 000000000000..af4b06e3999f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_force_vec.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance, drawn as the in-tree worker draws them.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..m3_force_vec import force_vec_kernel + + +class M3ForceVec(OpWrapper): + """Bump ``counter`` (int64 [1]) by one, then for each of the ``batch_size`` requests ``b`` write + ``min(base_total + (pool[(c * counter_stride + b * slot_stride) & pool_mask] < frac), max_total)`` into ``out`` + (int32), ``c`` the bumped counter: the in-tree ``_apply_force_accepted_tokens`` draw over its fixed-seed fp32 + ``pool``. ``block`` is the launch width (a power of two >= batch_size; 16, 32 or 64 are compiled at warmup). + + Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) + runs it. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), + Cell(why="eight requests, fraction 0.3", spec=dict(batch=8, block=16, frac=0.3)), + ) + + note = """ + Integer arithmetic and one fp32 comparison: `reference` must match exactly. + """ + + def __call__( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + force_vec_kernel[(1,)]( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK=block, + ) + return out + + def reference( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + """The counts (a new tensor; ``counter`` and ``out`` are not written).""" + c = int(counter.item()) + 1 + slots = torch.arange(batch_size, dtype=torch.int64, device=pool.device) + idx = (c * counter_stride + slots * slot_stride) & pool_mask + r = pool[idx] + return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) + + +m3_force_vec = M3ForceVec() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_sampled_accept.py new file mode 100644 index 000000000000..d0ad11de2ce9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_sampled_accept.py @@ -0,0 +1,269 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sampled (temperature + top-p) Eagle3 verification over the TP-sharded vocabulary in one kernel +(``m3_sampled_accept.sampled_accept``).""" + +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 + +from .m3_shard_accept import strict_accept + +#: Verify rows per call. +MAX_ROWS = 48 +#: The fused sampling kernel's 1 / T guard. +TEMP_EPSILON = 1e-6 + + +def nucleus_probs( + logits: torch.Tensor, temperatures: torch.Tensor, top_ps: torch.Tensor +) -> torch.Tensor: + """Per-row sampling distribution [rows, Vs] (fp32): w = exp(x / (T + 1e-6) - max), kept iff the mass strictly + above its weight is at most top_p * sum(w) (ties included), renormalized over the kept tokens.""" + inv = 1.0 / (temperatures.float() + TEMP_EPSILON) + x = logits.float() * inv[:, None] + w = torch.exp(x - x.max(dim=1, keepdim=True).values).double() + z = w.sum(dim=1, keepdim=True) + order = torch.sort(w, dim=1, descending=True) + sw = order.values + # mass strictly above each sorted weight: the cumulative sum before the first entry of its tie group + csum = torch.cumsum(sw, dim=1) - sw + first = torch.ones_like(sw, dtype=torch.bool) + first[:, 1:] = sw[:, 1:] != sw[:, :-1] + group = torch.cumsum(first.long(), dim=1) - 1 + starts = torch.zeros_like(csum) + starts.scatter_reduce_(1, group, csum, reduce="amin", include_self=False) + above = torch.gather(starts, 1, group) + keep_sorted = above <= top_ps.double()[:, None] * z + keep = torch.zeros_like(keep_sorted) + keep.scatter_(1, order.indices, keep_sorted) + p = torch.where(keep, w, torch.zeros_like(w)) + return (p / p.sum(dim=1, keepdim=True)).float() + + +class M3SampledAccept(OpWrapper): + """Sampled verification of ``rows = nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard. + + Each row draws one token from its top-p nucleus of softmax(logits / T) over the full (all-rank) vocabulary by + distributed rejection sampling (Philox keyed by the row's seed / offset, the row, the attempt and the CTA), then + the strict acceptance runs against ``draft`` as in ``m3_shard_accept``: ``acc_tokens``, ``num_acc``, + ``gather_ids``, optional ``next_new[:, 0]``. ``tokens`` receives the sampled token per row; ``extra_attempts`` + accumulates exchange rounds beyond the second. The exchange goes through the multicast workspace (two call + halves of ``minimax_m3_sampled_accept_ws_words(world)`` words) with the (finished rows, call) ``counters``. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # draft_len 3 (T = 4 rows per request) at 1 / 2 / 8 requests (4 / 8 / 32 rows, 8 CTAs per row) and the 48-row + # maximum (4 CTAs per row); the 25008-column TP8 target shard. Single rank (see note). + CELLS: tuple[Cell, ...] = ( + Cell( + why="cap 1 (4 rows), top_p 0.8 over a 12-token head: the draws follow the nucleus distribution", + spec=dict( + nreq=1, vs=25008, vocab_offset=0, check="distribution", top_p=0.8, temperature=0.7 + ), + ), + Cell( + why="cap 8 (32 rows, 256 co-resident CTAs): top_p ~ 0 keeps only the argmax, token exact", + spec=dict( + nreq=8, + vs=25008, + vocab_offset=2 * 25008, + check="argmax", + top_p=1e-7, + temperature=1.0, + ), + ), + Cell( + why="cap 2 (8 rows) with top_p = 1 (no cut: the first draw stands), next_new and forced counts", + spec=dict( + nreq=2, + vs=25008, + vocab_offset=25008, + check="distribution", + top_p=1.0, + temperature=1.0, + force="vec", + ), + ), + Cell( + why="48 rows (4 CTAs per row): the second launch geometry, nucleus distribution at top_p 0.5", + spec=dict( + nreq=12, + vs=25008, + vocab_offset=3 * 25008, + check="distribution", + top_p=0.5, + temperature=1.3, + force="total", + ), + ), + ) + + note = """ + What is sampled is gated statistically, the rest exactly. Every drawn token must lie in the reference nucleus; + the empirical distribution over many calls (distinct Philox offsets) must sit within a stated total-variation + band of the reference's; the acceptance outputs must equal the strict acceptance of the drawn tokens exactly. + The Philox stream itself is not reproduced in torch: the kernel's draw uses the hardware exp / log + approximations (__expf, __logf) and a CTA-order inverse CDF, which torch does not reproduce bit for bit. + + Certified at world 1 only. The decision is replicated from every rank's per-round records, which another rank + computes from its own shard on the fly, so peers cannot be stood in for on one device; the cross-rank path + needs the multi-device job and is not covered here. Rows with top-k or min-p must not reach this op (it applies + neither); is_valid states it. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> None: + torch.ops.trtllm.minimax_m3_sampled_accept( + logits, + vocab_offset, + temperatures, + top_ps, + seeds, + offsets, + ws_uc, + ws_mc_ptr, + counters, + tokens, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + force_vec, + extra_attempts, + use_pdl, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> dict: + """{"probs": the per-row nucleus distribution over this shard [rows, Vs] (global id = vocab_offset + + column), "accept": tokens -> the strict acceptance outputs}. The draw itself is random; ``seeds`` / + ``offsets`` select it and do not change the distribution. World 1: this shard is the whole vocabulary.""" + rows = logits.shape[0] + probs = nucleus_probs(logits, temperatures[:rows], top_ps[:rows]) + + def accept(drawn: torch.Tensor) -> dict[str, torch.Tensor]: + return strict_accept(drawn, draft, batch_idx, force_total, next_new, force_vec) + + return {"probs": probs, "accept": accept} + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> None: + rows = logits.shape[0] + # top-k and min-p are not applied: a row that has them gets a token outside its real sampling set. + if top_ks is not None: + k = top_ks[:rows] + assert bool(((k <= 0) | (k >= logits.shape[1] * world)).all()), ( + "top-k rows must take the generic path" + ) + if min_ps is not None: + assert not bool((min_ps[:rows] > 0).any()), "min-p rows must take the generic path" + # Rows are request-major verify rows: tokens[b * (D + 1) + k]; a mismatched layout pairs drafts with the + # wrong rows. + assert rows == draft.shape[0] * (draft.shape[1] + 1), ( + "logits rows must be nreq * (draft_len + 1)" + ) + # Per-row parameters are read at the row index, unchecked. + for t in (temperatures, top_ps, seeds, offsets): + assert t.numel() >= rows and t.is_contiguous(), ( + "per-row parameters must cover every row" + ) + assert temperatures.dtype == torch.float32 and top_ps.dtype == torch.float32, ( + "fp32 sampling parameters" + ) + assert seeds.dtype == torch.int64 and offsets.dtype == torch.int64, ( + "int64 Philox seeds / offsets" + ) + # The finished-row counter elects the CTA that runs the acceptance. + assert int(counters[0].item()) == 0, ( + "counters[0] (finished rows) must be zero between calls" + ) + + +m3_sampled_accept = M3SampledAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_sampled_accept_ws_words.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_sampled_accept_ws_words.py new file mode 100644 index 000000000000..b0a17e38a983 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_sampled_accept_ws_words.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Size of the sampled-verify exchange workspace (``m3_sampled_accept.SampledAcceptWorkspace.create``).""" + +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 + +#: Exchange rounds per row, verify rows per call, CTAs per row (at most), uint64 words per record. +MAX_ATTEMPTS = 16 +MAX_ROWS = 48 +MAX_SPLITS = 8 +RECORD_WORDS = 4 + + +class M3SampledAcceptWsWords(OpWrapper): + """uint64 words of the multicast workspace ``minimax_m3_sampled_accept`` exchanges through for ``world`` ranks: + two call halves of ``[attempts][rows][world * splits][record words]``.""" + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="the tp8 target's workspace", spec=dict(world=8)), + Cell(why="a single rank (local workspace, no multicast)", spec=dict(world=1)), + ) + + note = """ + A host computation (no device work); it is an op so the workspace size and the kernel's indexing come from one + definition. + """ + + def __call__(self, world: int) -> int: + return torch.ops.trtllm.minimax_m3_sampled_accept_ws_words(world) + + def reference(self, world: int) -> int: + """Two halves of attempts x rows x (world x splits) records x record words.""" + return 2 * MAX_ATTEMPTS * MAX_ROWS * world * MAX_SPLITS * RECORD_WORDS + + +m3_sampled_accept_ws_words = M3SampledAcceptWsWords() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_shard_accept.py new file mode 100644 index 000000000000..500553c2bc88 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_shard_accept.py @@ -0,0 +1,284 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Greedy Eagle3 verification from this rank's logits shard: per-row argmax, the cross-rank (value, index) exchange +and the strict acceptance in one kernel (``m3_shard_accept.shard_accept``).""" + +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 + +#: Verify rows per call, and exchange rows per call half of its workspace. +SHARD_ROWS = 48 + +_LOW32 = 0xFFFFFFFF + + +def ordered_bits(v: torch.Tensor) -> torch.Tensor: + """The kernel's order-preserving uint32 image of fp32 values, as int64 (-0.0 maps to +0.0's image: they tie).""" + b = v.float().contiguous().view(torch.int32).to(torch.int64) & _LOW32 + b = torch.where(b == 0x80000000, torch.zeros_like(b), b) + return torch.where(b >= 0x80000000, _LOW32 - b, b | 0x80000000) + + +def make_keys(values: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """Exchange keys of (fp32 value, global token index) pairs: larger value wins, ties take the lower index.""" + return (ordered_bits(values) << 32) | (_LOW32 - index.to(torch.int64)) + + +def unsigned_max(keys: torch.Tensor, dim: int) -> torch.Tensor: + """Max of int64-stored uint64 keys (the kernel compares them unsigned).""" + sign = torch.iinfo(torch.int64).min + return (keys ^ sign).max(dim=dim).values ^ sign + + +def key_token(key: torch.Tensor) -> torch.Tensor: + return (_LOW32 - (key & _LOW32)).to(torch.int32) + + +def local_argmax(logits: torch.Tensor, vocab_offset: int) -> tuple[torch.Tensor, torch.Tensor]: + """(max value, lowest global index reaching it) per row of a [M, Vs] fp32 shard.""" + best = logits.max(dim=1).values + idx = (logits == best[:, None]).float().argmax(dim=1) # first (lowest) index at the max + return best, idx.to(torch.int64) + vocab_offset + + +def exchange_tokens( + own_keys: torch.Tensor, + ws_uc: torch.Tensor, + counters: torch.Tensor, + rank: int, + world: int, + slot_rows: int, +) -> torch.Tensor: + """Global tokens of the cross-rank pick: this rank's keys against the other ranks' keys already in its slots. + + Slot (half * slot_rows + t) * world + r of ``ws_uc`` holds rank r's key of row t; half = call counter & 1. A + single-device run stands in for the other ranks by writing their keys there before the call. + """ + m = own_keys.shape[0] + half = int(counters[1].item()) & 1 + slots = ( + ws_uc.view(torch.int64)[half * slot_rows * world : (half * slot_rows + m) * world] + .view(m, world) + .clone() + ) + slots[:, rank] = own_keys + return key_token(unsigned_max(slots, dim=1)) + + +def strict_accept( + tokens: torch.Tensor, + draft: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], +) -> dict[str, torch.Tensor]: + """The strict acceptance epilogue from the verify-row tokens [nreq * (D + 1)] (request-major). + + accepted[b] = tokens of request b's rows; num_accepted[b] = 1 + the matching prefix of draft[b] (overridden by + force_vec[b], else by force_total > 0); gather_ids[b] = batch_idx[b] * (D + 1) + num_accepted[b] - 1, and + next_new[b, 0] (when given) = the token at that row. Returns the outputs plus ``next_new`` as written. + """ + nreq, d = draft.shape + k1 = d + 1 + tokens = tokens.to(torch.int32) + accepted = tokens.view(nreq, k1).clone() + match = (draft.to(torch.int32) == accepted[:, :d]).to(torch.int32) + num = 1 + match.cumprod(dim=1).sum(dim=1) + if force_vec is not None: + num = force_vec[:nreq].to(num.dtype) + elif force_total > 0: + num = torch.full_like(num, force_total) + gather = batch_idx.to(torch.int64) * k1 + num.to(torch.int64) - 1 + out = {"accepted": accepted, "num_accepted": num.to(torch.int32), "gather_ids": gather} + if next_new is not None: + nn = next_new.clone() + nn[:, 0] = tokens[gather] + out["next_new"] = nn + return out + + +class M3ShardAccept(OpWrapper): + """Greedy verification of ``nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard ``[rows, Vs]`` + (global ids ``vocab_offset ..``). + + Each row's token is its global argmax over all ranks' shards: keys ``ordered(value) << 32 | ~index`` (ties take + the lowest index, as ``torch.argmax``) are max-reduced over the shard, pushed into every rank's exchange slot + through the multicast slots ``[2][48][world]`` (Lamport: 0 = empty) and max-reduced over the ranks by the last + CTA, which then runs the strict acceptance: ``acc_tokens [nreq, D + 1]``, ``num_acc [nreq]`` (``force_vec`` / + ``force_total`` override the counts), ``gather_ids [nreq]`` (the first draft pass's row per request) and, when + given, ``next_new[:, 0]`` (the last accepted token). + """ + + ARCHS = frozenset({Arch.SM_103}) + + # Every greedy verify step of the target (dispatch.SHARD_ACCEPT): draft_len 3 at 1..12 requests (4..48 rows), the + # target lm_head's BF16 shard logits (25008 columns per rank at TP8). The logits come from the head GEMM, so the + # op only reduces. + CELLS: tuple[Cell, ...] = ( + Cell( + why="cap 1: 4 verify rows (1 request x T = 4), rank 0 of 8, next_new staged", + spec=dict(nreq=1, draft_len=3, vs=25008, world=8, rank=0, force="none", strided=False), + ), + Cell( + why="cap 2: 8 verify rows, rank 2, global ids offset by 2 x 25008", + spec=dict(nreq=2, draft_len=3, vs=25008, world=8, rank=2, force="none", strided=False), + ), + Cell( + why="cap 4: 16 verify rows (4 requests x T = 4), rank 1 of 8, next_new staged", + spec=dict(nreq=4, draft_len=3, vs=25008, world=8, rank=1, force="none", strided=False), + ), + Cell( + why="cap 8: 32 verify rows, rank 7, global ids offset by 7 x 25008", + spec=dict(nreq=8, draft_len=3, vs=25008, world=8, rank=7, force="none", strided=False), + ), + Cell( + why="12 rows (3 requests) read from a row-strided view of a wider logits buffer, forced counts", + spec=dict(nreq=3, draft_len=3, vs=25008, world=8, rank=0, force="vec", strided=True), + ), + Cell( + why="the 48-row maximum (12 requests at draft 3), force_total, rows with exact cross-rank ties", + spec=dict( + nreq=12, + draft_len=3, + vs=25008, + world=8, + rank=2, + force="total", + strided=False, + ties=True, + ), + ), + # Every other request count of the verify range, on the 8-rank 25008-column shards. + *( + Cell( + why=f"{n} requests ({4 * n} verify rows): a request count between the ones above", + spec=dict( + nreq=n, draft_len=3, vs=25008, world=8, rank=n % 8, force="none", strided=False + ), + ) + for n in (5, 6, 7, 9, 10, 11) + ), + ) + + note = """ + The op reduces BF16 logits exactly, so the reference is bit-exact, including cross-rank ties (lowest global + index wins). The cross-rank exchange is certified on one device: the other ranks' keys are written into this + rank's slots before the call, which is what their multicast pushes would have put there; the multicast mapping + itself (ws_mc_ptr = this rank's slots, as at world 1) is not exercised across devices here. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + torch.ops.trtllm.minimax_m3_shard_accept( + logits, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + num_ctas, + force_vec, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> dict[str, torch.Tensor]: + """Lowest-index maximum of each row's BF16 logits, the largest key among it and the peers' keys in + ``ws_uc``, then the strict acceptance. Reads ``ws_uc`` / ``counters`` before the call; the scratch and + output buffers and ``num_ctas`` do not change the result.""" + best, idx = local_argmax(logits.float(), vocab_offset) + tokens = exchange_tokens(make_keys(best, idx), ws_uc, counters, rank, world, SHARD_ROWS) + out = strict_accept(tokens, draft, batch_idx, force_total, next_new, force_vec) + out["tokens"] = tokens + return out + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + assert not bool(keys[: logits.shape[0]].any()), "keys must start zeroed" + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + assert ws_uc.numel() >= 2 * SHARD_ROWS * world, "ws_uc must hold [2][48][world] slots" + # A NaN logit maps to a key above every finite one and wins the row silently. + assert not bool(torch.isnan(logits).any()), "NaN logits win the argmax" + nreq = draft.shape[0] + assert bool(((batch_idx[:nreq] >= 0) & (batch_idx[:nreq] < nreq)).all()), ( + "batch_idx must index the nreq requests" + ) + if force_vec is not None: + fv = force_vec[:nreq] + assert bool(((fv >= 1) & (fv <= draft.shape[1] + 1)).all()), ( + "force_vec counts must lie in [1, D + 1]" + ) + + +m3_shard_accept = M3ShardAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_slot_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_slot_scatter.py new file mode 100644 index 000000000000..470a6372f2b6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_slot_scatter.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The spec sampler's per-slot stores of one step's outputs, in one launch (``M3SpecStepFold.scatter`` -> +``m3_spec_tail.slot_scatter``).""" + +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 padded_rows(src: torch.Tensor, width: int) -> torch.Tensor: + """src rows zero-padded (or cut) to ``width`` columns.""" + out = torch.zeros(src.shape[0], width, dtype=src.dtype, device=src.device) + w = min(width, src.shape[1]) + out[:, :w] = src[:, :w] + return out + + +class M3SlotScatter(OpWrapper): + """For i < len(slots), all int32 2-D strided views: ``da[slots[i]] = sa[i]``, ``db[slots[i]] = sb[i]``, + ``dd[slots[i]] = sd[i]`` (each zero-padded / cut to the destination width) and ``dl[slots[i]] = sl[i]``. The + destinations are the store's transposed [slots, width] views.""" + + ARCHS = frozenset({Arch.SM_103}) + + # The sampler store at draft_len 3 (8 slots): new tokens / next new tokens [slots, 4] (transposed views of + # [4, slots, 1]), next draft tokens [slots, 3]; a worker output narrower than the store (dynamic draft length) + # is zero-padded. + CELLS: tuple[Cell, ...] = ( + Cell( + why="8 rows into shuffled slots, widths equal to the store's", + spec=dict(n=8, slots=8, src_w=4, src_d=3, store_w=4, store_d=3), + ), + Cell( + why="3 rows of a narrower step (runtime draft 2): zero padding to the store width", + spec=dict(n=3, slots=8, src_w=3, src_d=2, store_w=4, store_d=3), + ), + ) + + note = """ + Pure index copies: the reference is exact. Slots are not bounds-checked; is_valid states it. + """ + + def __call__(self, slots, sa, da, sb, db, sl, dl, sd, dd): + torch.ops.trtllm.minimax_m3_slot_scatter(slots, sa, da, sb, db, sl, dl, sd, dd) + + def reference(self, slots, sa, da, sb, db, sl, dl, sd, dd): + """The four destinations after the call, on copies.""" + s = slots.long() + out = {} + for name, src, dst in (("da", sa, da), ("db", sb, db), ("dd", sd, dd)): + r = dst.clone() + r[s] = padded_rows(src[: s.numel()], dst.shape[1]) + out[name] = r + lr = dl.clone() + lr[s] = sl[: s.numel()] + out["dl"] = lr + return out + + def is_valid(self, slots, sa, da, sb, db, sl, dl, sd, dd): + assert bool(((slots >= 0) & (slots < da.shape[0]) & (slots < dl.numel())).all()), ( + "slots must index the store" + ) + # Two rows into one slot race (which one lands is unordered). + assert slots.unique().numel() == slots.numel(), "slots must be distinct" + assert sa.shape[0] >= slots.numel() and sl.numel() >= slots.numel(), "a source row per slot" + + +m3_slot_scatter = M3SlotScatter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_step_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_step_gather.py new file mode 100644 index 000000000000..abea2cf620e9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_step_gather.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The overlap scheduler's per-step input gathers from the previous step's device outputs, in one launch +(``M3SpecStepFold.gather`` -> ``m3_spec_tail.step_gather``).""" + +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 + + +class M3StepGather(OpWrapper): + """For the ``P`` previous requests (sequence slots ``slots``, ``T`` tokens and ``D`` draft tokens each), all + int32:: + + ids[ids0 + p*T + t] = new_tok[t, slots[p]] (new_tok [T', slots, (1)]) + draft[d0 + p*D + d] = next_draft[slots[p], d] + pos_off[i] = lens[ppi[i - b0*T]] for i in [b0*T, (b0+P)*T), else 0 (every i < pos_off.numel()) + kv_off[i] = lens[slots[i - b0]] - T for i in [b0, b0+P), else 0 (every i < kv_off.numel()) + + Every other element of ``ids`` / ``draft`` is left as it was. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # draft_len 3: T = 4, D = 3; the store's [T, slots, 1] new tokens and [slots, D] next drafts (8 slots, the + # max batch); the engine's flat buffers sized for the max batch. + CELLS: tuple[Cell, ...] = ( + Cell( + why="an unchanged batch of 8 previous requests in shuffled slots, nothing before them (b0 = 0)", + spec=dict(p=8, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + Cell( + why="2 previous requests after 3 new ones: offsets ids0 / d0 / b0 nonzero, zero-filled offsets outside", + spec=dict(p=2, t=4, d=3, slots=8, b0=3, ids0=12, d0=9, extra_rows=3), + ), + Cell( + why="one request, the store's 2-D view of new tokens (beam dim squeezed)", + spec=dict(p=1, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0, two_d=True), + ), + ) + + note = """ + Pure index copies: the reference is exact. The op does not bound slots, ppi or the destination offsets; is_valid + states them. + """ + + def __call__( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + torch.ops.trtllm.minimax_m3_step_gather( + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ) + + def reference( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + """The four destinations after the call, on copies (``ids`` / ``draft`` keep their other elements).""" + nt = new_tok.reshape(new_tok.shape[0], new_tok.shape[1]) + s = slots[:P].long() + ids_r, draft_r = ids.clone(), draft.clone() + ids_r[ids0 : ids0 + P * T] = nt[:T, s].t().reshape(-1) + draft_r[d0 : d0 + P * D] = next_draft[s, :D].reshape(-1) + pos_r = torch.zeros_like(pos_off) + pos_r[b0 * T : (b0 + P) * T] = lens[ppi[: P * T].long()] + kv_r = torch.zeros_like(kv_off) + kv_r[b0 : b0 + P] = lens[s] - T + return {"ids": ids_r, "draft": draft_r, "pos_off": pos_r, "kv_off": kv_r} + + def is_valid( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + # Reads through slots / ppi and writes at the offsets are unchecked: any of them out of range reads or + # writes another buffer. + s = slots[:P] + assert bool( + ( + (s >= 0) & (s < next_draft.shape[0]) & (s < new_tok.shape[1]) & (s < lens.numel()) + ).all() + ), "slots must index the store" + assert bool(((ppi[: P * T] >= 0) & (ppi[: P * T] < lens.numel())).all()), ( + "ppi must index lens" + ) + assert new_tok.shape[0] >= T and next_draft.shape[1] >= D, "store narrower than T / D" + assert ids.numel() >= ids0 + P * T and draft.numel() >= d0 + P * D, "destinations too short" + assert pos_off.numel() >= (b0 + P) * T and kv_off.numel() >= b0 + P, ( + "offset buffers too short" + ) + + +m3_step_gather = M3StepGather() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/triton_sparse_decode.py new file mode 100644 index 000000000000..0e8dc710bc98 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/triton_sparse_decode.py @@ -0,0 +1,1662 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/common/ops/sparse_attn.py +"""Triton block-sparse GQA decode attention for MiniMax-M3. + +Flash-decoding over the blocks the indexer selected: one CTA per (query token, +top-k chunk, KV head) accumulates a partial output plus its log-sum-exp, and a +second kernel merges the chunks by LSE weight. That beats running the +context-schedule FMHA kernel at decode, where a single query token leaves most +of a 128-row Q tile idle. + +GB300/SM103 tuning targets GQA group 16, top-k 16 and four-token Eagle3 +verification. Split-K, query pairing and launch tables require four KV heads; +the linear-softmax policy also covers one and two KV heads. Other hardware +and shapes retain the default policies, including supported SM100 NVFP4. + +Vendored from the vLLM source linked in the file header (v0.26.1rc0-77-g6f91edf96). +Differences from upstream: + +* K and V are separate HND paged views ([num_pages, num_kv_heads, page_size, + head_dim]) with independent strides rather than one cache fused along the + last dim, because that is how the M3 KV pool is laid out here. +* Only the scalar KV-scale mode is kept; M3 stores unscaled E4M3 K/V. +* A packed NVFP4 cache is read in place, dequantized in-register from its E2M1 + nibbles and E4M3 block scales. +* Partial-output and LSE scratch come from the persistent buffer arena so their + addresses survive CUDA graph replay. +* A single top-k chunk writes the final output directly, avoiding scratch and + the otherwise redundant LSE merge. +* The merge kernel zeroes rows whose chunks are all empty instead of letting + them produce NaN, since CUDA-graph padding rows flow on into the rest of the + network here. + +This is the MiniMax-M3 NVFP4 / sm_103 / tp8 target's copy. On top of the +in-tree decode it routes long-context TP8 steps to the one-launch +``m3_attn_lc8`` kernel, short-sequence verify steps of at least +``MQ_MIN_TOKENS`` rows to the multi-query producer, releases PDL dependents +early, and has the merge emit the MXFP8 o_proj hand-off (``m3_prequant.SLOT``). +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + check_decode_span_shape, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers + +from . import m3_attn_lc8, m3_prequant + +# One sparse block is exactly one KV page. +SPARSE_BLOCK_SIZE = 128 # capacity: fixed: the MSA block size of the checkpoint + +_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) + +_TL_DTYPES = { + torch.bfloat16: tl.bfloat16, + torch.float16: tl.float16, + torch.float32: tl.float32, +} + +# Total CTAs the split-K partitioning aims for before it stops splitting. +_TARGET_GRID = 256 # capacity: fixed: a launch-grid target, not a size + +# NVFP4 quantizes in groups of 16 cache elements along head_dim, one E4M3 scale +# byte each; see the fp4_quantize call in msa_scatter.fused_write_layer_caches_nvfp4. +NVFP4_SF_VEC_SIZE = 16 + + +# Fewest query rows for which the merge writes the MXFP8 o_proj hand-off: every decode step, so the MXFP8 o_proj +# never runs its own quantize kernel after attention. +MERGE_MXFP8_MIN_TOKENS = 1 # capacity: fixed: a row lower bound + + +class MergeMxfp8: + """MXFP8 copy of the split-K merge output for the o_proj GEMM. + + From ``MERGE_MXFP8_MIN_TOKENS`` to ``MAX_ROWS`` rows the merge kernel writes the E4M3 rows and swizzled UE8M0 + scales next to the BF16 output and publishes them to ``m3_prequant.SLOT``, so the MXFP8 o_proj skips its own + quantization. One 128-row buffer per width serves every layer: each o_proj consumes it before the next layer's + merge rewrites it, in stream order. + """ + + MAX_ROWS = 128 # capacity: fixed: the merge hand-off's rows (decode steps) + + def __init__(self): + self._buffers = {} + + def buffers(self, width: int, device: torch.device): + key = (width, device) + if key not in self._buffers: + self._buffers[key] = ( + torch.empty((self.MAX_ROWS, width), dtype=torch.float8_e4m3fn, device=device), + torch.zeros((self.MAX_ROWS * (width // 32),), dtype=torch.uint8, device=device), + ) + return self._buffers[key] + + +MERGE_MXFP8 = MergeMxfp8() + + +@triton.jit +def _e2m1x2_to_f16x2(packed): + """Convert one packed E2M1 byte into its low/high FP16 values.""" + return tl.inline_asm_elementwise( + asm="""{ + .reg .b8 src; + .reg .b32 pair; + mov.b32 {src, _, _, _}, $2; + cvt.rn.f16x2.e2m1x2 pair, src; + mov.b32 {$0, $1}, pair; + }""", + constraints="=h,=h,r", + args=(packed,), + dtype=(tl.float16, tl.float16), + is_pure=True, + pack=1, + ) + + +@triton.jit +def _dequant_nvfp4_rows( + packed, # [BLOCK_N, D // 2] uint8, two E2M1 nibbles per byte + scale_bytes, # [BLOCK_N, SCALE_COLS] uint8, one E4M3 scale per 16 elements + BLOCK_N: tl.constexpr, + D: tl.constexpr, + SCALE_COLS: tl.constexpr, + SF_VEC: tl.constexpr, +): + """Dequantize one packed NVFP4 [BLOCK_N, D] tile to fp16. + + The caller gathers the block scales in whichever layout its role stores + them, so this handles only the value math and the scale broadcast. The + per-tensor scale is applied by the caller. + """ + # Blackwell converts both E2M1 nibbles with one instruction. Joining the + # low/high results along a new innermost axis and flattening it rebuilds + # the original element order without scalar exponent/sign reconstruction. + value_lo, value_hi = _e2m1x2_to_f16x2(packed) + values = tl.reshape(tl.join(value_lo, value_hi), (BLOCK_N, D)) + + scales = scale_bytes.to(tl.float8e4nv, bitcast=True).to(tl.float16) + scales = tl.reshape( + tl.broadcast_to(scales[:, :, None], (BLOCK_N, SCALE_COLS, SF_VEC)), (BLOCK_N, D) + ) + return values * scales + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max(16, triton.next_power_of_2(args["gqa_group_size"])), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +# Mixed ctx+gen steps launch on [num_ctx:] slices (unaligned seq_lens / block_table) with arbitrary +# total_q; specializing on those compiled a new variant mid-serving (~650 ms TTFT stall at c8). +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + out_ptr, # final [total_q, num_heads, head_dim] for FUSED_MERGE, else unused + counter_ptr, # int32 [total_q, num_kv_heads] arrival counters for FUSED_MERGE, zero at rest + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + stride_out_n, + stride_out_h, + stride_out_d, + num_kv_heads, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + FUSED_MERGE: tl.constexpr, # the last chunk to finish merges the partials in place + EARLY_TRIGGER: tl.constexpr, # release PDL dependents (weight prefetch) at entry +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query token, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + pid_b = pid_bc % total_q + pid_c = pid_bc // total_q + req_id = pid_b // decode_query_len + q_offset = pid_b - req_id * decode_query_len + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + if EARLY_TRIGGER: + # Dependents still wait for this grid to complete; triggering here only + # lets them start their input-independent prologue (e.g. weight loads). + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + req_id) + query_pos = seq_len - decode_query_len + q_offset + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + + # Bound the walk by the token's own valid block count rather than a + # sentinel: the selector emits ascending ids and only pads with -1 past + # min(topk, cdiv(kv_len, blk)) entries, so those are never dereferenced. + num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk = tl.minimum(max_topk, num_blocks) + chunk_end = tl.minimum(chunk_start + chunk_size_topk, real_topk) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + off_h = tl.arange(0, BLOCK_SIZE_H) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + h_mask = off_h < gqa_group_size + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + pid_b * stride_qn + + (pid_h + off_h[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so the result + # matches a caller that widened q itself before the call. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + cur_idx_ptr = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + for it in tl.range(chunk_start, chunk_end): + blk = tl.load(cur_idx_ptr) + cur_idx_ptr = cur_idx_ptr + stride_tk + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + page = tl.load(bt_row + blk).to(tl.int64) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len + if KV_NVFP4: + # K block scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + qk = tl.where(pos_mask[None, :], 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.exp2(qk - m_ij[:, None]) + l_ij = tl.sum(p, axis=1) + if USE_LINEAR_SOFTMAX: + alpha = tl.exp2(m_i - m_ij) + acc_o = acc_o * alpha[:, None] + else: + acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] + if KV_NVFP4: + # V block scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = l_i * alpha + l_ij + else: + lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + o_base = o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h + tl.store( + o_base + off_h[:, None] * stride_o_h + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + lse_base = lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h + tl.store(lse_base + off_h * stride_l_h, lse_out, mask=h_mask) + + if FUSED_MERGE: + # Split-K serial reduction: every chunk publishes its partial, and the + # chunk that arrives last for this (token, KV head) merges all of them + # exactly as _merge_topk_attn_out_kernel does, then re-arms the counter. + tl.debug_barrier() + counter = counter_ptr + pid_b * num_kv_heads + pid_kh + arrived = tl.atomic_add(counter, 1, sem="acq_rel", scope="gpu") + if arrived == NUM_TOPK_CHUNKS - 1: + tl.debug_barrier() + o_rows = ( + o_ptr + + pid_b * stride_o_b + + (pid_h + off_h[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d + ) + l_rows = lse_ptr + pid_b * stride_l_b + (pid_h + off_h) * stride_l_h + lse_max = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + lse_max = tl.maximum(lse_max, lse_c) + # All chunks empty (a CUDA-graph padding row): zero, never NaN. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + denom = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + denom += tl.exp2(lse_c - lse_max) + denom = tl.where(denom > 0, denom, 1.0) + merged = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + w_c = tl.exp2(lse_c - lse_max) / denom + o_c = tl.load( + o_rows + c * stride_o_c, mask=hd_mask, other=0.0, cache_modifier=".cg" + ).to(tl.float32) + merged += o_c * w_c[:, None] + tl.store( + out_ptr + + pid_b * stride_out_n + + (pid_h + off_h[:, None]) * stride_out_h + + off_d[None, :] * stride_out_d, + merged.to(out_ptr.dtype.element_ty), + mask=hd_mask, + ) + tl.atomic_xchg(counter, 0, sem="relaxed", scope="gpu") + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max( + 16, + triton.next_power_of_2(args["gqa_group_size"] * args["QUERY_GROUP_SIZE"]), + ), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_pair_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + QUERY_GROUP_SIZE: tl.constexpr, # consecutive Eagle queries sharing one block set + MAX_TOPK: tl.constexpr, + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query group, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + num_query_groups = total_q // QUERY_GROUP_SIZE + pid_g = pid_bc % num_query_groups + pid_c = pid_bc // num_query_groups + groups_per_req = decode_query_len // QUERY_GROUP_SIZE + req_id = pid_g // groups_per_req + q_group_offset = (pid_g - req_id * groups_per_req) * QUERY_GROUP_SIZE + pid_b = req_id * decode_query_len + q_group_offset + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + + seq_len = tl.load(seq_lens + req_id) + # Query rows are flattened as [query-in-group, GQA head]. + off_h = tl.arange(0, BLOCK_SIZE_H) + q_lane = off_h // gqa_group_size + h_lane = off_h - q_lane * gqa_group_size + h_mask = off_h < gqa_group_size * QUERY_GROUP_SIZE + q_token = pid_b + q_lane + query_pos = seq_len - decode_query_len + q_group_offset + q_lane + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + kv_len_max = tl.max(tl.where(h_mask, kv_len, 0), axis=0) + + # Each query can have a different valid-block count at a page boundary. + # The selector puts those valid entries first and pads only the tail. + kv_len_0 = tl.maximum(seq_len - decode_query_len + q_group_offset + 1, 0) + num_blocks_0 = (kv_len_0 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_0 = tl.minimum(max_topk, num_blocks_0) + chunk_end_0 = tl.minimum(chunk_start + chunk_size_topk, real_topk_0) + if QUERY_GROUP_SIZE == 2: + kv_len_1 = tl.maximum(seq_len - decode_query_len + q_group_offset + 2, 0) + num_blocks_1 = (kv_len_1 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_1 = tl.minimum(max_topk, num_blocks_1) + chunk_end_1 = tl.minimum(chunk_start + chunk_size_topk, real_topk_1) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + q_token[:, None] * stride_qn + + (pid_h + h_lane[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so this + # reproduces the values the caller's standalone widening kernel produced. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + idx_ptr_0 = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + if QUERY_GROUP_SIZE == 2: + idx_ptr_1 = idx_ptr_0 + stride_tn + idx_0 = chunk_start + idx_1 = chunk_start + count_0 = chunk_end_0 - chunk_start + count_1 = chunk_end_1 - chunk_start + chunk_capacity: tl.constexpr = (MAX_TOPK + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_offsets = tl.arange(0, chunk_capacity) + sentinel = 2147483647 + blocks_0 = tl.load( + idx_ptr_0 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_0, + other=sentinel, + ) + blocks_1 = tl.load( + idx_ptr_1 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_1, + other=sentinel, + ) + chunk_equal = tl.sum(blocks_0 != blocks_1, axis=0) == 0 + # Equal chunks take the measured shared-page loop. A divergent chunk + # walks the sorted union, bounded by the sum of both row lengths. + num_candidates = tl.where(chunk_equal, count_0, count_0 + count_1) + else: + num_candidates = chunk_end_0 - chunk_start + + for _ in tl.range(0, num_candidates): + if QUERY_GROUP_SIZE == 2: + if chunk_equal: + has_block = True + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + else: + valid_0 = idx_0 < chunk_end_0 + valid_1 = idx_1 < chunk_end_1 + blk_0 = tl.load(idx_ptr_0, mask=valid_0, other=sentinel) + blk_1 = tl.load(idx_ptr_1, mask=valid_1, other=sentinel) + use_0 = valid_0 & (blk_0 <= blk_1) + use_1 = valid_1 & (blk_1 <= blk_0) + has_block = use_0 | use_1 + blk = tl.minimum(blk_0, blk_1) + idx_0 += use_0.to(tl.int32) + idx_1 += use_1.to(tl.int32) + idx_ptr_0 += use_0.to(tl.int64) * stride_tk + idx_ptr_1 += use_1.to(tl.int64) * stride_tk + query_active = h_mask & (((q_lane == 0) & use_0) | ((q_lane == 1) & use_1)) + else: + has_block = True + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + + if has_block: + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + page = tl.load(bt_row + blk).to(tl.int64) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len_max + q_pos_mask = blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None] + if KV_NVFP4: + # K scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + active_pos = query_active[:, None] & q_pos_mask + qk = tl.where(active_pos, 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.where(query_active, tl.maximum(m_i, tl.max(qk, axis=1)), m_i) + p = tl.where(active_pos, tl.exp2(qk - m_ij[:, None]), 0.0) + l_ij = tl.sum(p, axis=1) + alpha = tl.where(query_active, tl.exp2(m_i - m_ij), 1.0) + acc_o = acc_o * alpha[:, None] + if KV_NVFP4: + # V scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = tl.where(query_active, l_i * alpha + l_ij, l_i) + else: + lse_ij = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + lse_i = tl.where(query_active, lse_ij, lse_i) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + tl.store( + o_ptr + + pid_c * stride_o_c + + q_token[:, None] * stride_o_b + + (pid_h + h_lane[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + tl.store( + lse_ptr + pid_c * stride_l_c + q_token * stride_l_b + (pid_h + h_lane) * stride_l_h, + lse_out, + mask=h_mask, + ) + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _rcp_approx_ftz(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 _store_mxfp8_head(x, q8_ptr, sf_ptr, row, head, num_sf_k_tiles, HEAD_DIM: tl.constexpr): + """Quantize one row's [head * HEAD_DIM, (head + 1) * HEAD_DIM) slice to MXFP8 as + torch.ops.trtllm.mxfp8_quantize(x, True) does: UE8M0 scale = round-up(amax * rcp(448)), E4M3 RN-satfinite, + scales in the 128x4-swizzled layout.""" + NUM_SF: tl.constexpr = HEAD_DIM // 32 + xr = tl.reshape(x.to(tl.float32), (NUM_SF, 32)) + amax = tl.max(tl.abs(xr), axis=1) + sfv = amax * _rcp_approx_ftz(tl.full((NUM_SF,), 448.0, tl.float32)) + bits = sfv.to(tl.int32, bitcast=True) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32) + e = tl.minimum(e, 254) + scale = ((254 - e) << 23).to(tl.float32, bitcast=True) + scale = tl.where(amax != 0, scale, 0.0) + q = (xr * scale[:, None]).to(tl.float8e4nv) + off_q = tl.reshape(tl.arange(0, HEAD_DIM), (NUM_SF, 32)) + tl.store(q8_ptr + row * (num_sf_k_tiles * 128) + head * HEAD_DIM + off_q, q) + k_sf = head * NUM_SF + tl.arange(0, NUM_SF) + off_sf = ( + (row // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + + k_sf % 4 + ) + tl.store(sf_ptr + off_sf, e.to(tl.uint8)) + + +@triton.heuristics({"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}) +@triton.jit +def _merge_topk_attn_out_kernel( + o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim] + lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads] + out_ptr, # merged out: [total_q, num_heads, head_dim] + head_dim, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_out_n, + stride_out_h, + stride_out_d, + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr = False, + q8_ptr=None, # MXFP8 copy of out: [total_q, num_heads * head_dim] E4M3 + sf_ptr=None, # its UE8M0 block scales, 128x4-swizzled + num_sf_k_tiles=0, # ceil(num_heads * head_dim / 128) + MXFP8_OUT: tl.constexpr = False, +): + pid_b, pid_h = tl.program_id(0), tl.program_id(1) + + if USE_PDL: + if EARLY_TRIGGER: + # Release the dependent (the o_proj GEMM) before waiting on the attention producer, so its weight + # prefetch overlaps the attention kernel as well as this merge; its griddepcontrol.wait still waits + # for this grid (and therefore the producer) to complete. + tl.extra.cuda.gdc_launch_dependents() + tl.extra.cuda.gdc_wait() + + off_c = tl.arange(0, NUM_TOPK_CHUNKS) + off_d = tl.arange(0, BLOCK_SIZE_D) + o = tl.load( + o_ptr + + pid_b * stride_o_b + + pid_h * stride_o_h + + off_c[:, None] * stride_o_c + + off_d[None, :] * stride_o_d, + mask=off_d[None, :] < head_dim, + other=0.0, + ).to(tl.float32) + # Empty chunks contribute -inf, hence weight 0. + lse = tl.load(lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c) + lse_max = tl.max(lse, axis=0) + # A row whose every chunk is empty (a CUDA-graph padding row) would give + # -inf - -inf here. Zero it instead: the padded output is discarded, but a + # NaN would survive into the residual stream and the all-reduce. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + weights = tl.exp2(lse - lse_max) + denom = tl.sum(weights, axis=0) + weights = weights / tl.where(denom > 0, denom, 1.0) + o_merged = tl.sum(o * weights[:, None], axis=0) + out_ptrs = out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d + tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim) + if MXFP8_OUT: + _store_mxfp8_head( + o_merged.to(out_ptr.dtype.element_ty), + q8_ptr, + sf_ptr, + pid_b, + pid_h, + num_sf_k_tiles, + BLOCK_SIZE_D, + ) + + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> int: + """Split-K factor over the top-k blocks, as a power of two. + + Depends only on shapes that are fixed for a captured batch size, so the + launch geometry is frozen inside a CUDA graph. + """ + target = max(1, min(max_topk, _TARGET_GRID // max(1, total_q * num_kv_heads))) + return 1 << (target.bit_length() - 1) + + +# NVFP4 instructions are supported on both datacenter Blackwell targets. +_SM100F_CAPABILITIES = ((10, 0), (10, 3)) +# Shape-specific tuning was measured on GB300 only. SM100 uses the defaults. +_NVFP4_TUNED_CAPABILITIES = ((10, 3),) + + +def _sm100f_nvfp4_num_topk_chunks( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> Optional[int]: + """Select GB300-measured split-K factors for M3's SM100-family decode shape.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return None + local_batch = total_q // decode_query_len + return { + 2: 8, + 3: 8, + 4: 8, + 5: 16, + 6: 4, + 7: 4, + 8: 2, + 9: 2, + 10: 8, + 11: 4, + 12: 2, + 13: 2, + 14: 2, + }.get(local_batch) + + +def _sm100f_nvfp4_use_linear_softmax( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> bool: + """Use the GB300-measured lower-register recurrence on validated M3 shapes.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads not in (1, 2, 4) + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return False + local_batch = total_q // decode_query_len + # The TP4/DEP32 single-request graph is 8.9% faster with the original + # log-domain recurrence; every other measured DEP8/16/32 bucket is neutral + # or faster with the linear denominator. + return not (num_kv_heads == 1 and local_batch == 1) + + +def _sm100f_nvfp4_query_group_size( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> int: + """Group Eagle queries using the GB300-measured shared-page policy.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return 1 + # The sorted-union kernel is 6.5%, 1.6%, and 24.2% faster at local batches + # 11, 12, and 14 with split-K 8. Other measured buckets are neutral or + # slower and retain the single-query producer. + return 2 if total_q // decode_query_len in (11, 12, 14) else 1 + + +def _sm100f_nvfp4_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select native-NVFP4 launch geometry for the SM100 family. + + The FP16 QK/PV path is register-sensitive. GB300 dense replay shows four + warps/two stages wins through local batch 28, while two warps/one stage + avoids the next residency cliff at local batches 29--32. Leave every + unmeasured architecture, model geometry, and larger batch on Triton's + established defaults. + """ + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + ): + return {} + local_batch = total_q // decode_query_len + if local_batch == 1 or local_batch in (2, 3, 5, 6, 7, 10, 11, 12): + return {"num_warps": 4, "num_stages": 1} + if local_batch == 4: + return {"num_warps": 2, "num_stages": 1} + if 8 <= local_batch <= 28: + return {"num_warps": 4, "num_stages": 2} + if 29 <= local_batch <= 32: + return {"num_warps": 2, "num_stages": 1} + return {} + + +def _sm100f_nvfp4_merge_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select GB300-measured one-warp merges for M3's SM100-family graphs.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability in _NVFP4_TUNED_CAPABILITIES + and num_kv_heads == 4 + and gqa_group_size == 16 + and max_topk == 16 + and decode_query_len == 4 + and total_q % decode_query_len == 0 + and total_q // decode_query_len in (2, 4, 5, 6, 7, 10, 11) + ): + return {"num_warps": 1} + return {} + + +def _check_nvfp4_inputs( + k_paged: torch.Tensor, + v_paged: torch.Tensor, + nvfp4_args: tuple[Optional[torch.Tensor], ...], + *, + head_dim: int, + scale_cols: int, +) -> None: + """Validate the packed-NVFP4 cache against what the kernel assumes. + + The kernel indexes block scales by a flat within-page offset and reshapes a + [., SCALE_COLS] scale tile up to head_dim, so it needs contiguous scale rows + and head_dim exactly BLOCK_SIZE_D. Both hold for MiniMax-M3; check them + rather than let a mismatch read neighboring pages. + """ + k_block_scale, v_block_scale, k_global_scale, v_global_scale = nvfp4_args + if any(arg is None for arg in nvfp4_args): + raise ValueError( + "MiniMax-M3 NVFP4 sparse decode needs all four of k_block_scale, " + "v_block_scale, k_global_scale and v_global_scale." + ) + if head_dim != triton.next_power_of_2(head_dim): + raise ValueError(f"NVFP4 sparse decode requires a power-of-two head_dim; got {head_dim}.") + for name, data in (("k_paged", k_paged), ("v_paged", v_paged)): + if int(data.shape[3]) * 2 != head_dim: + raise ValueError( + f"NVFP4 {name} must pack two elements per byte: expected " + f"head_dim/2 ({head_dim // 2}) columns, got {int(data.shape[3])}." + ) + for name, scale in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + if tuple(scale.shape) != tuple(k_paged.shape[:3]) + (scale_cols,): + raise ValueError( + f"NVFP4 {name} must be [num_pages, num_kv_heads, page_size, " + f"{scale_cols}] to match the cache; got {tuple(scale.shape)}." + ) + if scale.stride(3) != 1 or scale.stride(2) != scale_cols: + raise ValueError( + f"NVFP4 {name} needs contiguous per-page scale bytes; got strides " + f"{tuple(scale.stride())}." + ) + for name, gs in (("k_global_scale", k_global_scale), ("v_global_scale", v_global_scale)): + if gs.dtype != torch.float32 or gs.numel() != 1: + raise ValueError( + f"NVFP4 {name} must be a single-element FP32 tensor; got " + f"{gs.numel()} element(s) of {gs.dtype}." + ) + + +# Multi-query producer (triton_sparse_decode_mq) for identity-selection verify steps with at least this many query +# tokens; below it the per-token producer is as fast. +MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query producer's crossover +# Warps per multi-query producer program. +MQ_NUM_WARPS = 8 + + +def _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, +) -> None: + """minimax_m3_sparse_attn_decode through m3_attn_lc8: same table, output and MXFP8 hand-off contract.""" + use_pdl = get_env_enable_pdl() + total_q, num_heads, head_dim = q.shape + linear = _sm100f_nvfp4_use_linear_softmax( + total_q=total_q, + num_kv_heads=1, + gqa_group_size=num_heads, + max_topk=int(topk_idx.shape[-1]), + decode_query_len=decode_query_len, + ) + mxfp8 = None + m3_prequant.SLOT.clear() + if MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS: + mxfp8 = MERGE_MXFP8.buffers(num_heads * head_dim, output.device) + # The early trigger lets the o_proj launch and stream its weights while attention runs. + m3_attn_lc8.m3_lc8_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + linear_softmax=linear, + mxfp8=mxfp8, + use_pdl=use_pdl, + ) + if mxfp8 is not None: + m3_prequant.SLOT.publish( + output.view(total_q, num_heads * head_dim), mxfp8[0][:total_q], mxfp8[1] + ) + + +@torch.no_grad() +def minimax_m3_sparse_attn_decode( + q: torch.Tensor, # [total_q, num_heads, head_dim] + k_paged: torch.Tensor, # [num_pages, num_kv_heads, page_size, head_dim] + v_paged: torch.Tensor, # same layout as k_paged + topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk] + block_table: torch.Tensor, # [num_reqs, max_blocks] + seq_lens: torch.Tensor, # [num_reqs] int32 + *, + sm_scale: float, + output: torch.Tensor, # [total_q, num_heads, head_dim] + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, +) -> None: + """Block-sparse GQA decode attention, written into output in place. + + q may be FP8, as a fused QK-norm/RoPE producer emits it. The kernel widens + it to output.dtype in-register rather than making the caller materialize a + widened copy, which would cost a standalone kernel on every sparse layer. + The attention math is unchanged either way, since E4M3 widens exactly. + + kv_scale is an optional scalar dequantization factor for an FP8 cache; + MiniMax-M3 stores unscaled E4M3, so it is normally None. + + Passing the four NVFP4 scale tensors selects the packed-NVFP4 cache instead: + k_paged/v_paged then hold two E2M1 elements per byte, the block scales are + paged E4M3 bytes in the layouts msa_scatter writes, and the two per-tensor + dequant scales are single-element FP32 device tensors. + + num_topk_chunks overrides the split-K factor and exists for tests: the + merged result must not depend on it. + """ + total_q, num_heads, head_dim = q.shape + num_kv_heads = int(k_paged.shape[1]) + check_decode_span_shape( + "MiniMax-M3 Triton sparse decode", total_q, int(seq_lens.shape[0]), decode_query_len + ) + if int(k_paged.shape[2]) != SPARSE_BLOCK_SIZE: + raise ValueError( + f"MiniMax-M3 sparse decode requires page_size={SPARSE_BLOCK_SIZE}; " + f"got {int(k_paged.shape[2])}." + ) + if num_heads % num_kv_heads: + # The kernel covers heads pid_kh * gqa_group_size + [0, group), so a + # head count that is not a whole number of groups leaves the tail heads + # unwritten and the merge kernel reads uninitialized partials. + raise ValueError( + f"MiniMax-M3 sparse decode requires num_heads ({num_heads}) to be a " + f"multiple of num_kv_heads ({num_kv_heads})." + ) + nvfp4_args = (k_block_scale, v_block_scale, k_global_scale, v_global_scale) + kv_nvfp4 = any(arg is not None for arg in nvfp4_args) + scale_cols = head_dim // NVFP4_SF_VEC_SIZE + if kv_nvfp4: + capability = torch.cuda.get_device_capability(q.device) + if capability not in _SM100F_CAPABILITIES: + raise NotImplementedError( + "MiniMax-M3 NVFP4 sparse decode requires SM100/SM103; " + f"got SM{capability[0] * 10 + capability[1]}." + ) + _check_nvfp4_inputs(k_paged, v_paged, nvfp4_args, head_dim=head_dim, scale_cols=scale_cols) + # Two E2M1 elements per byte; Triton must see them unsigned so that the + # high-nibble shift does not sign-extend. + k_paged, v_paged = k_paged.view(torch.uint8), v_paged.view(torch.uint8) + k_block_scale, v_block_scale = ( + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + ) + max_topk = int(topk_idx.shape[-1]) + gqa_group_size = num_heads // num_kv_heads + use_scale = k_paged.dtype in _FP8_DTYPES and kv_scale is not None + # Triton needs a real pointer even for the unused arguments. + scale_arg = kv_scale if use_scale else output + k_sc_arg, v_sc_arg = (k_block_scale, v_block_scale) if kv_nvfp4 else (output, output) + k_gs_arg, v_gs_arg = (k_global_scale, v_global_scale) if kv_nvfp4 else (output, output) + + widen_q = q.dtype in _FP8_DTYPES + if widen_q and output.dtype not in _TL_DTYPES: + raise ValueError( + f"MiniMax-M3 sparse decode cannot widen FP8 q into {output.dtype}; " + f"supported compute dtypes are {sorted(d.__str__() for d in _TL_DTYPES)}." + ) + # Packed NVFP4 dequantizes to FP16. The production producer supplies FP8 + # q, which widens exactly to FP16, so keep both dot operands there instead + # of converting every K/V tile to BF16. Accumulation remains FP32. + if widen_q: + compute_dtype = tl.float16 if kv_nvfp4 else _TL_DTYPES[output.dtype] + else: + compute_dtype = tl.float32 + + # Long-context TP8 steps: attention, split-K merge and MXFP8 hand-off in one launch. A selected table on an NVFP4 + # cache with FP8 q at the TP8 heads is the kernel's (modeling's SPARSE_DECODE_LC8); operands it declines are a miss. + lc8_covered = ( + kv_nvfp4 + and widen_q + and not identity_blocks + and num_topk_chunks is None + and q.dim() == 3 + and int(q.shape[1]) == m3_attn_lc8.NUM_HEADS + ) + if lc8_covered and not m3_attn_lc8.lc8_applies( + q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output + ): + from .. import modeling + + modeling.missed( + modeling.SPARSE_DECODE_LC8.path, + total_q, + "the step's q / KV pages / scales / output are not the layout the kernel is built for", + ) + lc8_covered = False + if lc8_covered: + _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + ) + return + + query_group_size = ( + _sm100f_nvfp4_query_group_size( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else 1 + ) + if num_topk_chunks is None: + if query_group_size == 2: + measured_chunks = 8 + else: + measured_chunks = ( + _sm100f_nvfp4_num_topk_chunks( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else None + ) + num_topk_chunks = ( + measured_chunks + if measured_chunks is not None + else resolve_num_topk_chunks(total_q, num_kv_heads, max_topk) + ) + if kv_nvfp4 and total_q * num_kv_heads <= 64: + # GB300 TP4 Eagle verify (1 KV head, group 16, 6 query tokens): one block per program is fastest up to + # batch 8 (lowm5/bench_decode_chunks.py: b8 11.7 vs 14.8 us with the default split of 4). + num_topk_chunks = max_topk + elif num_topk_chunks & (num_topk_chunks - 1): + raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") + + # Every query token of a request reads the same blocks under identity selection (short-sequence mode): one + # program per (request, block) covers all of them and dequantizes each block once. + use_mq = ( + identity_blocks + and total_q >= MQ_MIN_TOKENS + and kv_nvfp4 + and num_kv_heads == 1 + and decode_query_len > 1 + and head_dim == 128 + and gqa_group_size * decode_query_len <= 128 + ) + if use_mq: + num_topk_chunks = max_topk + direct_output = num_topk_chunks == 1 + if direct_output: + # With one chunk, acc_o is already normalized and final. Point the + # producer at the caller's output and compile away the LSE store. + # Keeping a real dummy pointer preserves Triton's pointer contract. + o_arg = output + lse_arg = output + o_strides = (0, output.stride(0), output.stride(1), output.stride(2)) + lse_strides = (0, 0, 0) + else: + # Persistent arena rather than torch.empty, so the partials keep one + # address across CUDA graph replays. FP32 keeps the split-K factor from + # perturbing the merged result. + reserve = torch.cuda.is_current_stream_capturing() + o_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads, head_dim], + torch.float32, + buffer_name="m3_sparse_decode_o_partial", + reserve_buffer=reserve, + ) + lse_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads], + torch.float32, + buffer_name="m3_sparse_decode_lse_partial", + reserve_buffer=reserve, + ) + o_strides = tuple(o_arg.stride()) + lse_strides = tuple(lse_arg.stride()) + + use_pdl = get_env_enable_pdl() + pdl_launch = {"launch_pdl": True} if use_pdl else {} + use_linear_softmax = kv_nvfp4 and _sm100f_nvfp4_use_linear_softmax( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if query_group_size == 2: + launch_options = {"num_warps": 4, "num_stages": 1} + else: + launch_options = ( + _sm100f_nvfp4_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + + # The pair path merges the two sorted selector rows on device. Equal blocks + # share one K/V dequantization; divergent blocks update only their owning + # query, so the optimization does not rely on runtime selection equality. + if total_q % query_group_size or decode_query_len % query_group_size: + raise ValueError( + f"query group {query_group_size} must divide total_q={total_q} and " + f"decode_query_len={decode_query_len}." + ) + + producer_kernel = ( + _gqa_sparse_decode_pair_kernel if query_group_size == 2 else _gqa_sparse_decode_kernel + ) + single_kernel_args = ( + { + # No split-K arrival counters: the separate merge kernel combines the chunks. + "merge": (output, None), + "tail": (output.stride(0), output.stride(1), output.stride(2), num_kv_heads), + # An early release pays off when the dependent is the consumer of the + # final output (it can prefetch), or a merge grid that itself releases + # its dependent at entry: the o_proj then launches during this kernel and + # streams its weights while attention runs. Above 8 query rows the + # concurrent weight stream slows the larger attention grid more than the + # prefetch saves (GB300 TP4: -1.1% at 48 rows). + "options": {"FUSED_MERGE": False, "EARLY_TRIGGER": direct_output or total_q <= 8}, + } + if query_group_size == 1 + else {"merge": (), "tail": (), "options": {}} + ) + pair_kernel_options = ( + {"QUERY_GROUP_SIZE": query_group_size, "MAX_TOPK": max_topk} + if query_group_size == 2 + else {} + ) + if use_mq: + from .triton_sparse_decode_mq import launch_mq_producer + + launch_mq_producer( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + block_table, + seq_lens, + o_arg, + lse_arg, + sm_scale=sm_scale, + decode_query_len=decode_query_len, + use_pdl=use_pdl, + early_trigger=False, + num_warps=MQ_NUM_WARPS, + ) + else: + producer_kernel[ + (total_q // query_group_size) * num_topk_chunks, + num_kv_heads, + ]( + q, + k_paged, + v_paged, + scale_arg, + k_sc_arg, + v_sc_arg, + k_gs_arg, + v_gs_arg, + topk_idx, + o_arg, + lse_arg, + *single_kernel_args["merge"], + block_table, + seq_lens, + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + q.stride(0), + q.stride(1), + q.stride(2), + k_paged.stride(0), + k_paged.stride(1), + k_paged.stride(2), + k_paged.stride(3), + v_paged.stride(0), + v_paged.stride(1), + v_paged.stride(2), + v_paged.stride(3), + k_sc_arg.stride(0), + k_sc_arg.stride(1), + v_sc_arg.stride(0), + v_sc_arg.stride(1), + topk_idx.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + *o_strides, + *lse_strides, + block_table.stride(0), + *single_kernel_args["tail"], + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_SCALE=use_scale, + WIDEN_Q=widen_q, + COMPUTE_DTYPE=compute_dtype, + USE_PDL=use_pdl, + KV_NVFP4=kv_nvfp4, + USE_LINEAR_SOFTMAX=use_linear_softmax, + SCALE_COLS=scale_cols, + DIRECT_OUTPUT=direct_output, + **single_kernel_args["options"], + **pair_kernel_options, + **pdl_launch, + **launch_options, + ) + m3_prequant.SLOT.clear() + if not direct_output: + width = num_heads * head_dim + mxfp8 = ( + MERGE_MXFP8.buffers(width, output.device) + if MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS + and head_dim == 128 + and output.dim() == 3 + and output.is_contiguous() + and output.dtype == torch.bfloat16 + else None + ) + mxfp8_args = ( + { + "q8_ptr": mxfp8[0], + "sf_ptr": mxfp8[1], + "num_sf_k_tiles": width // 128, + "MXFP8_OUT": True, + } + if mxfp8 is not None + else {} + ) + merge_launch_options = ( + _sm100f_nvfp4_merge_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + _merge_topk_attn_out_kernel[(total_q, num_heads)]( + o_arg, + lse_arg, + output, + head_dim, + *o_strides, + *lse_strides, + output.stride(0), + output.stride(1), + output.stride(2), + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_PDL=use_pdl, + # Release the o_proj at the merge's entry so it streams its weights during the merge. + EARLY_TRIGGER=True, + **mxfp8_args, + **pdl_launch, + **merge_launch_options, + ) + if mxfp8 is not None: + m3_prequant.SLOT.publish(output.view(total_q, width), mxfp8[0][:total_q], mxfp8[1]) + + +__all__ = [ + "MERGE_MXFP8", + "MERGE_MXFP8_MIN_TOKENS", + "MergeMxfp8", + "NVFP4_SF_VEC_SIZE", + "SPARSE_BLOCK_SIZE", + "_sm100f_nvfp4_merge_launch_options", + "_sm100f_nvfp4_num_topk_chunks", + "_sm100f_nvfp4_query_group_size", + "_sm100f_nvfp4_use_linear_softmax", + "minimax_m3_sparse_attn_decode", + "resolve_num_topk_chunks", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/triton_sparse_decode_mq.py new file mode 100644 index 000000000000..048d31b82d68 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/triton_sparse_decode_mq.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multi-query block-sparse decode producer for MiniMax-M3 short-sequence verify steps. + +In short-sequence mode every request attends at most ``topk`` KV blocks and block selection is the identity, so +all ``decode_query_len`` query tokens of a request read the same blocks. The per-token producer +(``_gqa_sparse_decode_kernel``) then dequantizes every block once per query token. This producer runs one program +per (request, block) for all of the request's query tokens: rows are (token, head) pairs, each block is +dequantized once, and the causal limit is a per-row mask. + +It writes the same per-block partials (output scaled by 1/l and the V global scale, log2 LSE) as the per-token +producer with one block per chunk, so ``_merge_topk_attn_out_kernel`` merges them unchanged. +""" + +import triton +import triton.language as tl + +from .triton_sparse_decode import _dequant_nvfp4_rows + + +@triton.jit( + do_not_specialize=["num_reqs"], do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"] +) +def _gqa_sparse_decode_mq_kernel( + q_ptr, # [num_reqs * QLEN, num_heads, head_dim], FP8 or FP16/BF16 + k_ptr, # packed NVFP4 K: [num_pages, 1, page_size, head_dim / 2] uint8 + v_ptr, # packed NVFP4 V: same layout + k_sc_ptr, # K block scales [num_pages, 1, page_size, head_dim / 16], token-linear + v_sc_ptr, # V block scales, 4x4 token-quad order + k_gs_ptr, + v_gs_ptr, + t_ptr, # topk table [1, total_q, topk]; the request's last token's row is read + o_ptr, # partials [topk, total_q, num_heads, head_dim] fp32 + lse_ptr, # partials [topk, total_q, num_heads] fp32 (log2) + block_table_ptr, + seq_lens, + num_reqs, + sm_scale, + stride_qn, + stride_qh, + stride_k_blk, + stride_k_pos, + stride_v_blk, + stride_v_pos, + stride_ksc_blk, + stride_vsc_blk, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + QLEN: tl.constexpr, + GROUP: tl.constexpr, # query heads per KV head (16) + BLOCK_R: tl.constexpr, # next power of two of QLEN * GROUP + HEAD_DIM: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + SCALE_COLS: tl.constexpr, + USE_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr, +): + pid_r = tl.program_id(0) + pid_c = tl.program_id(1) # top-k position == block index under identity selection + sm_scale_log2e = sm_scale * 1.4426950409 * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + + off_r = tl.arange(0, BLOCK_R) + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, HEAD_DIM) + off_dp = tl.arange(0, HEAD_DIM // 2) + off_s = tl.arange(0, SCALE_COLS) + row_t = off_r // GROUP + row_h = off_r % GROUP + r_mask = off_r < QLEN * GROUP + tok = pid_r * QLEN + row_t + + if USE_PDL: + tl.extra.cuda.gdc_wait() + if EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + pid_r) + # Token t of the request sits at seq_len - QLEN + t and attends positions [0, kv_len_t). + kv_len = tl.maximum(seq_len - QLEN + row_t + 1, 0) + blk = tl.load(t_ptr + (pid_r * QLEN + QLEN - 1) * stride_tn + pid_c * stride_tk) + blk_ok = blk >= 0 + blk = tl.where(blk_ok, blk, 0) + page = tl.load(block_table_ptr + pid_r * stride_bt_b + blk, mask=blk_ok, other=0).to(tl.int64) + # Rows past the last token's position are never read: their cache bytes may be stale (0 * NaN in the PV dot). + n_valid = tl.where( + blk_ok, tl.minimum(tl.maximum(seq_len - blk * BLOCK_SIZE_K, 0), BLOCK_SIZE_K), 0 + ) + n_mask = off_n < n_valid + + q = tl.load( + q_ptr + tok[:, None] * stride_qn + row_h[:, None] * stride_qh + off_d[None, :], + mask=r_mask[:, None], + other=0.0, + ).to(tl.float16) + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + page * stride_k_blk + off_n[:, None] * stride_k_pos + off_dp[None, :], + mask=n_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + page * stride_ksc_blk + off_n[:, None] * SCALE_COLS + off_s[None, :], + mask=n_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=HEAD_DIM, + SCALE_COLS=SCALE_COLS, + SF_VEC=HEAD_DIM // SCALE_COLS, + ) + valid = ( + (blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None]) & n_mask[None, :] & r_mask[:, None] + ) + qk = tl.dot(q, tl.trans(k_deq).to(tl.float16)) * sm_scale_log2e + qk = tl.where(valid, qk, float("-inf")) + m = tl.max(qk, axis=1) + m_safe = tl.where(m == float("-inf"), 0.0, m) + p = tl.exp2(qk - m_safe[:, None]) + p_sum = tl.sum(p, axis=1) + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + page * stride_v_blk + off_n[:, None] * stride_v_pos + off_dp[None, :], + mask=n_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=n_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=HEAD_DIM, + SCALE_COLS=SCALE_COLS, + SF_VEC=HEAD_DIM // SCALE_COLS, + ).to(tl.float16) + acc = tl.dot(p.to(tl.float16), v) + # An empty partial stores zero with LSE -inf, so the merge gives it weight 0. + scale = tl.where(p_sum > 0.0, 1.0 / p_sum, 0.0) * v_global_scale + lse = tl.where(p_sum > 0.0, m + tl.log2(p_sum), float("-inf")) + acc = acc * scale[:, None] + tl.store( + o_ptr + + pid_c * stride_o_c + + tok[:, None] * stride_o_b + + row_h[:, None] * stride_o_h + + off_d[None, :], + acc, + mask=r_mask[:, None], + ) + tl.store(lse_ptr + pid_c * stride_l_c + tok * stride_l_b + row_h * stride_l_h, lse, mask=r_mask) + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +def launch_mq_producer( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + block_table, + seq_lens, + o_arg, + lse_arg, + *, + sm_scale, + decode_query_len, + use_pdl, + early_trigger, + num_warps=8, + num_stages=1, +): + """Launch the multi-query producer; arguments as minimax_m3_sparse_attn_decode after its NVFP4 views.""" + total_q, num_heads, head_dim = q.shape + num_reqs = total_q // decode_query_len + max_topk = int(topk_idx.shape[-1]) + rows = decode_query_len * num_heads + pdl_launch = {"launch_pdl": True} if use_pdl else {} + _gqa_sparse_decode_mq_kernel[(num_reqs, max_topk)]( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + o_arg, + lse_arg, + block_table, + seq_lens, + num_reqs, + sm_scale, + q.stride(0), + q.stride(1), + k_paged.stride(0), + k_paged.stride(2), + v_paged.stride(0), + v_paged.stride(2), + k_block_scale.stride(0), + v_block_scale.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + o_arg.stride(0), + o_arg.stride(1), + o_arg.stride(2), + lse_arg.stride(0), + lse_arg.stride(1), + lse_arg.stride(2), + block_table.stride(0), + QLEN=decode_query_len, + GROUP=num_heads, + BLOCK_R=triton.next_power_of_2(rows), + HEAD_DIM=head_dim, + BLOCK_SIZE_K=int(k_paged.shape[2]), + SCALE_COLS=head_dim // 16, + USE_PDL=use_pdl, + EARLY_TRIGGER=early_trigger, + num_warps=num_warps, + num_stages=num_stages, + **pdl_launch, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/weights.py new file mode 100644 index 000000000000..011445939d14 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/weights.py @@ -0,0 +1,883 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight tables and loaders: MiniMax-M3 NVFP4 / sm_103 / tp8. + +Two tables, one per checkpoint the target reads: + +* ``MODEL_WEIGHTS`` (``MiniMaxM3Tp8Weights``): the vision-language checkpoint's language model, + ``language_model.*``. Every rank is handed the whole checkpoint, so every tensor-parallel split is this + table's: the Q heads, the K/V/index-Q heads (4 each over 8 ranks: rank r holds head r // 2, as the built-in + loader duplicates them), the dense MLPs and the shared expert's + intermediate, the o_proj / down-projection input columns, every routed expert's intermediate (all 128 + experts live on every rank), and the vocabulary of the shell's ``lm_head``. The index-K head is replicated. +* ``DRAFT_WEIGHTS`` (``MiniMaxM3Tp8DraftWeights``): the one-layer BF16 Eagle3 drafter's own checkpoint + (``speculative_model``), loaded through the shell's ``load_draft_weights`` when speculation is on. + +Storage is HF ``[out, in]`` row-major. Five families need a transform, each a layout a kernel demands: + +* **MXFP8 block scales** (every attention projection, the dense MLPs and the shared expert): the checkpoint + stores UE8M0 bytes ``[N, K / 32]``; the GEMMs read the 128x4-swizzled order + (``trtllm::block_scale_interleave``). A fused projection concatenates its shards' scales first and swizzles + once. +* **The five-way sparse-layer projection** ``[Q | K | V | index-Q | index-K]`` and the three-way dense one + ``[Q | K | V]``: the per-rank shards concatenated along N. +* **Gated MLPs** (dense and shared expert): ``[gate ; up]`` rows. +* **NVFP4 routed experts** (layers 3..59): ``[up (w3) ; gate (w1)]`` rows, then the trtllm-gen row shuffle of + the weight and of its E4M3 scales (gated-activation reorder + 128-row epilogue tile shuffle, + ``trtllmgen_maybe_get_cached_*_permute_indices``) and the 128x4 swizzle of the scales. +* **Gemma boundary norms**: ``input_layernorm``, ``post_attention_layernorm`` and ``model.norm`` are stored as + Gemma norms, ``(1 + w) * x``; the fused all-reduce + RMSNorm kernels apply ``w * x``, so ``1`` is added + (in BF16) at load. The per-head Q/K/index norms stay Gemma: their kernels apply the offset themselves. + +The routed experts' per-tensor scalars (input scales and second-level weight scales, every expert) load raw; +``modeling`` derives the kernels' global scales, alphas and activation constants from them once the weights +are real. + +Loading contract: ``load(model, weights)`` consumes the engine-provided checkpoint dict, fills every declared +parameter exactly once and asserts coverage in both directions: every declared parameter written, and every +checkpoint key either consumed or in ``expected_unconsumed`` (the vision tower and its projector, which this +text-only target does not build, and the language model's multi-token-prediction head). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, Callable + +import torch +import torch.nn as nn + +from tensorrt_llm._torch.moe.fused_moe.quantization import ( + trtllmgen_maybe_get_cached_w2_permute_indices, + trtllmgen_maybe_get_cached_w3_w1_permute_indices, +) + +#: ``layers=ALL_LAYERS`` is the common case: a weight every decoder layer carries. +ALL_LAYERS = object() + +_LM = "language_model." +#: Checkpoint key families this target does not build: the vision tower and its two projectors (text only), and the +#: language model's multi-token-prediction head (speculation, when on, is the separate Eagle3 drafter). +_UNBUILT_PREFIXES = ( + "vision_tower.", + "multi_modal_projector.", + "patch_merge_mlp.", + _LM + "model.mtp.", +) + +MX_BLOCK = 32 # MXFP8: one UE8M0 scale per 32 elements along K +NV_BLOCK = 16 # NVFP4: one E4M3 scale per 16 elements along K +EPILOGUE_TILE_M = ( + 128 # the trtllm-gen MoE runner's row tile: fixes its weight and scale row shuffles +) + + +@dataclass(frozen=True) +class W: + """One weight role: what it is, where its bytes come from, what happens on the way. + + ``shape`` takes the dimension bundle and returns the declared shape. ``src`` is a callable + ``(d, i) -> [(key or key tuple, index into the param | None), ...]`` with keys relative to the checkpoint + root; a key tuple hands the transform those tensors together. ``layers`` is ``ALL_LAYERS``, ``None`` (not + per-layer) or ``d -> layer indices``. ``transform(src, d)`` runs on the destination's device before the copy. + """ + + name: str + shape: Callable[[Any], tuple[int, ...]] + src: Callable + dtype: torch.dtype | None = None + transform: Callable | None = None + layers: object | None = ALL_LAYERS + + def layer_indices(self, d: SimpleNamespace) -> list[int] | None: + if self.layers is None: + return None + if self.layers is ALL_LAYERS: + return list(range(d.num_layers)) + return list(self.layers(d)) + + +# ------------------------------------------------------------------------------------------------ shards + + +def _col( + t: torch.Tensor, d: SimpleNamespace, tp: int | None = None, rank: int | None = None +) -> torch.Tensor: + """This rank's rows (column-parallel output split), even split, sliced on the host and moved to the + destination device (only this rank's bytes are read).""" + tp = d.tp if tp is None else tp + rank = d.rank if rank is None else rank + n = t.shape[0] // tp + return _dev(t[rank * n : (rank + 1) * n], d) + + +def _heads(t: torch.Tensor, d: SimpleNamespace, head_dim: int) -> torch.Tensor: + """This rank's rows of a projection split by heads of ``head_dim`` rows: an even split when there are at least as + many heads as ranks, else the head of rank // (ranks per head), duplicated on those ranks.""" + n = t.shape[0] // head_dim + if n >= d.tp: + return _col(t, d) + h = d.rank // (d.tp // n) + return _dev(t[h * head_dim : (h + 1) * head_dim], d) + + +def _row(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """This rank's input columns (row-parallel input split), even split, on the destination device.""" + n = t.shape[1] // d.tp + return _dev(t[:, d.rank * n : (d.rank + 1) * n], d) + + +def _dev(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """``t`` on the destination device (``d.dev``, set per fill).""" + return t.to(d.dev, non_blocking=True) + + +def _swizzle(scale: torch.Tensor) -> torch.Tensor: + """128x4 swizzle of a row-major ``[N, K / 32]`` UE8M0 scale matrix, flat.""" + return torch.ops.trtllm.block_scale_interleave(scale.contiguous()) + + +# --------------------------------------------------------------------------------------------- transforms + + +def _gemma_fold(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """A Gemma boundary norm as the plain norm the fused all-reduce epilogue applies: ``w + 1`` in BF16.""" + return t + 1.0 + + +def _f32(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return t.to(torch.float32) + + +def _qkv_sparse_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[Q | K | V | index-Q | index-K]``: Q by attention heads, K/V/index-Q by their heads, index-K replicated.""" + q, k, v, iq, ik = t + hd, xd = d.head_dim, d.index_dim + return torch.cat( + [_col(q, d), _heads(k, d, hd), _heads(v, d, hd), _heads(iq, d, xd), _dev(ik, d)], dim=0 + ).contiguous() + + +def _qkv_sparse_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v, iq, ik = t + hd, xd = d.head_dim, d.index_dim + return _swizzle( + torch.cat( + [_col(q, d), _heads(k, d, hd), _heads(v, d, hd), _heads(iq, d, xd), _dev(ik, d)], dim=0 + ) + ) + + +def _qkv_dense_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + hd = d.head_dim + return torch.cat([_col(q, d), _heads(k, d, hd), _heads(v, d, hd)], dim=0).contiguous() + + +def _qkv_dense_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + hd = d.head_dim + return _swizzle(torch.cat([_col(q, d), _heads(k, d, hd), _heads(v, d, hd)], dim=0)) + + +def _index_qk_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[index-Q | index-K]`` in BF16: this rank's index-Q head (rank r holds head r // 2) and the replicated + index-K, the E4M3 values cast (the scale tensors, the last two sources, are consumed only).""" + iq, ik = t[0], t[1] + return torch.cat([_heads(iq, d, d.index_dim), _dev(ik, d)], dim=0).to(d.dtype).contiguous() + + +def _gate_up_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return torch.cat([_col(gate, d), _col(up, d)], dim=0).contiguous() + + +def _gate_up_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return _swizzle(torch.cat([_col(gate, d), _col(up, d)], dim=0)) + + +def _row_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _row(t, d).contiguous() + + +def _row_s(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _swizzle(_row(t, d)) + + +def _col_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _col(t, d).contiguous() + + +#: Permute indices of the trtllm-gen expert relayouts, keyed by (shape, family, scale group): one expert's are +#: every expert's, so each is computed once per load. +_PERMUTE_CACHE: dict = {} + + +def _expert_w31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 operand: ``[up ; gate]`` packed-FP4 rows of this rank, then the gated-activation row + reorder and the epilogue-tile shuffle.""" + up, gate = t + w = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + w, _PERMUTE_CACHE, EPILOGUE_TILE_M, is_gated_act_gemm=True + ) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 E4M3 scales: ``[up ; gate]`` rows of this rank, the same row reorder at scale + granularity, then the 128x4 swizzle.""" + up, gate = t + s = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK, is_gated_act_gemm=True + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled.view(torch.uint8).reshape(shape)) + .view(torch.float8_e4m3fn) + .reshape(shape) + ) + + +def _expert_w2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 operand: this rank's input columns, then the epilogue-tile row shuffle.""" + w = _row(t, d).contiguous() + perm = trtllmgen_maybe_get_cached_w2_permute_indices(w, _PERMUTE_CACHE, EPILOGUE_TILE_M) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 E4M3 scales: this rank's columns, the row shuffle at scale granularity, the swizzle.""" + s = _row(t, d).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w2_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled).view(torch.float8_e4m3fn).reshape(shape) + ) + + +# ------------------------------------------------------------------------------------------------ sources + + +def _one(template: str): + return lambda d, i, t=template: [(_LM + t.format(i=i), None)] + + +def _many(*templates: str): + return lambda d, i, ts=templates: [(tuple(_LM + t.format(i=i) for t in ts), None)] + + +def _per_expert(*suffixes: str): + """One source per routed expert, written at its expert index (all experts are local at tp8).""" + + def build(d, i, sx=suffixes): + out = [] + for e in range(d.num_experts): + keys = tuple(_LM + f"model.layers.{i}.block_sparse_moe.experts.{e}.{s}" for s in sx) + out.append((keys if len(keys) > 1 else keys[0], (e,))) + return out + + return build + + +def _dense_layers(d): + return range(d.first_sparse) + + +def _sparse_layers(d): + return range(d.first_sparse, d.num_layers) + + +def _three_way_qkv_layers(d): + """Layers whose QKV is the three-way [Q | K | V]: the dense layers, and every layer without the five-way packed + projection (``fuse_qkv_index_projection=False``).""" + return range(d.first_sparse) if d.fused else range(d.num_layers) + + +def _five_way_qkv_layers(d): + return range(d.first_sparse, d.num_layers) if d.fused else range(0) + + +def _unfused_sparse_layers(d): + return range(0) if d.fused else range(d.first_sparse, d.num_layers) + + +_ATTN = "model.layers.{i}.self_attn" +_MOE = "model.layers.{i}.block_sparse_moe" + + +class MiniMaxM3Tp8Weights: + """MiniMax-M3 NVFP4 / sm_103 / tp8: the language model's table and the passes over it. + + ``declare`` allocates every weight it names, ``manifest`` maps checkpoint keys onto them, ``load`` fills + them; ``_targets`` is the one place that decides which parameters exist. + """ + + #: Free the per-expert relayout scratch after each expert stack. + RELEASE_AFTER = "_e_s2" + + WEIGHTS: tuple[W, ...] = ( + # --- the two boundary norms of every layer (Gemma, folded) --- + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.input_layernorm.weight"), + transform=_gemma_fold, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.post_attention_layernorm.weight"), + transform=_gemma_fold, + ), + # --- attention: per-head Q/K norms (Gemma, unfolded), the packed QKV and the output projection --- + W("q_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".q_norm.weight")), + W("k_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".k_norm.weight")), + W( + "qkv_w", + shape=lambda d: (d.qkv_dense_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in "qkv")), + transform=_qkv_dense_w, + layers=_three_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_dense_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many(*(f"{_ATTN}.{p}_proj.weight_scale_inv" for p in "qkv")), + transform=_qkv_dense_s, + layers=_three_way_qkv_layers, + ), + W( + "qkv_w", + shape=lambda d: (d.qkv_sparse_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in ("q", "k", "v", "index_q", "index_k"))), + transform=_qkv_sparse_w, + layers=_five_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_sparse_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + *( + f"{_ATTN}.{p}_proj.weight_scale_inv" + for p in ("q", "k", "v", "index_q", "index_k") + ) + ), + transform=_qkv_sparse_s, + layers=_five_way_qkv_layers, + ), + # The separate BF16 index projection [index-Q | index-K] of the unfused sparse layers. The built-in model + # builds it as an unquantized Linear and fills it through the fused gate/up loader, which casts the + # checkpoint's E4M3 values to BF16 and does not apply their block scales: so does this table (the scale + # tensors are consumed, not applied). + W( + "idx_w", + shape=lambda d: (d.index_heads * d.index_dim + d.index_dim, d.hidden), + src=_many( + _ATTN + ".index_q_proj.weight", + _ATTN + ".index_k_proj.weight", + _ATTN + ".index_q_proj.weight_scale_inv", + _ATTN + ".index_k_proj.weight_scale_inv", + ), + transform=_index_qk_w, + layers=_unfused_sparse_layers, + ), + W( + "idx_q_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_q_norm.weight"), + layers=_sparse_layers, + ), + W( + "idx_k_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_k_norm.weight"), + layers=_sparse_layers, + ), + W( + "o_w", + shape=lambda d: (d.hidden, d.q_width), + dtype=torch.float8_e4m3fn, + src=_one(_ATTN + ".o_proj.weight"), + transform=_row_w, + ), + W( + "o_s", + shape=lambda d: (d.hidden * (d.q_width // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_ATTN + ".o_proj.weight_scale_inv"), + transform=_row_s, + ), + # --- dense SwiGLU-OAI MLP (layers 0..2) --- + W( + "mlp_gu_w", + shape=lambda d: (2 * d.dense_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight", "model.layers.{i}.mlp.up_proj.weight" + ), + transform=_gate_up_w, + layers=_dense_layers, + ), + W( + "mlp_gu_s", + shape=lambda d: (2 * d.dense_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight_scale_inv", + "model.layers.{i}.mlp.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_dense_layers, + ), + W( + "mlp_dn_w", + shape=lambda d: (d.hidden, d.dense_inter), + dtype=torch.float8_e4m3fn, + src=_one("model.layers.{i}.mlp.down_proj.weight"), + transform=_row_w, + layers=_dense_layers, + ), + W( + "mlp_dn_s", + shape=lambda d: (d.hidden * (d.dense_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one("model.layers.{i}.mlp.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_dense_layers, + ), + # --- MoE (layers 3..59): the fp32 router and its bias, the shared expert --- + W( + "router", + shape=lambda d: (d.num_experts, d.hidden), + dtype=torch.float32, + src=_one(_MOE + ".gate.weight"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "router_bias", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_one(_MOE + ".e_score_correction_bias"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "sh_gu_w", + shape=lambda d: (2 * d.shared_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + _MOE + ".shared_experts.gate_proj.weight", _MOE + ".shared_experts.up_proj.weight" + ), + transform=_gate_up_w, + layers=_sparse_layers, + ), + W( + "sh_gu_s", + shape=lambda d: (2 * d.shared_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + _MOE + ".shared_experts.gate_proj.weight_scale_inv", + _MOE + ".shared_experts.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_sparse_layers, + ), + W( + "sh_dn_w", + shape=lambda d: (d.hidden, d.shared_inter), + dtype=torch.float8_e4m3fn, + src=_one(_MOE + ".shared_experts.down_proj.weight"), + transform=_row_w, + layers=_sparse_layers, + ), + W( + "sh_dn_s", + shape=lambda d: (d.hidden * (d.shared_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_MOE + ".shared_experts.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_sparse_layers, + ), + # --- the routed experts, NVFP4 (trtllm-gen layouts), all 128 on every rank --- + W( + "e_w31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // 2), + dtype=torch.uint8, + src=_per_expert("w3.weight", "w1.weight"), + transform=_expert_w31, + layers=_sparse_layers, + ), + W( + "e_w2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2), + dtype=torch.uint8, + src=_per_expert("w2.weight"), + transform=_expert_w2, + layers=_sparse_layers, + ), + W( + "e_s31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w3.weight_scale", "w1.weight_scale"), + transform=_expert_s31, + layers=_sparse_layers, + ), + W( + "e_s2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w2.weight_scale"), + transform=_expert_s2, + layers=_sparse_layers, + ), + # The per-tensor scalars of every expert, raw; modeling derives the kernels' scalars from them. + *( + W( + f"e_{name}", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_per_expert(suffix), + layers=_sparse_layers, + ) + for name, suffix in ( + ("isc1", "w1.input_scale"), + ("isc3", "w3.input_scale"), + ("isc2", "w2.input_scale"), + ("ws2_1", "w1.weight_scale_2"), + ("ws2_3", "w3.weight_scale_2"), + ("ws2_2", "w2.weight_scale_2"), + ) + ), + # --- not per-layer --- + W( + "embed", + shape=lambda d: (d.vocab, d.hidden), + src=_one("model.embed_tokens.weight"), + layers=None, + ), + W( + "final_norm", + shape=lambda d: (d.hidden,), + src=_one("model.norm.weight"), + transform=_gemma_fold, + layers=None, + ), + ) + + def dims(self, core) -> SimpleNamespace: + cfg = core.model_config.pretrained_config + mapping = core.model_config.mapping + sparse = cfg.sparse_attention_config + freq = [int(f) for f in sparse["sparse_attention_freq"]] + tp = mapping.tp_size + head_dim = int(cfg.head_dim) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + index_heads = max(int(sparse["sparse_num_index_heads"]) // tp, 1) + index_dim = int(sparse["sparse_index_dim"]) + return SimpleNamespace( + fused=bool( + getattr( + core.model_config.sparse_attention_config, "fuse_qkv_index_projection", False + ) + ), + num_layers=cfg.num_hidden_layers, + first_sparse=freq.index(1), + hidden=cfg.hidden_size, + head_dim=head_dim, + heads=heads, + kv_heads=kv_heads, + index_heads=index_heads, + index_dim=index_dim, + q_width=heads * head_dim, + qkv_dense_rows=(heads + 2 * kv_heads) * head_dim, + qkv_sparse_rows=(heads + 2 * kv_heads) * head_dim + index_heads * index_dim + index_dim, + dense_inter=int(cfg.dense_intermediate_size) // tp, + shared_inter=int(cfg.shared_intermediate_size) * int(cfg.n_shared_experts) // tp, + moe_inter=int(cfg.intermediate_size) // tp, + num_experts=int(cfg.num_local_experts), + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + layers = entry.layer_indices(d) + if layers is None: + return [(entry.name, None)] + return [(f"l{i}_{entry.name}", i) for i in layers] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + """Allocate every weight the table declares (meta under the engine's meta init).""" + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, core) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + """param key -> [(checkpoint key or key tuple, index into the param | None, transform | None)].""" + d = self.dims(core) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, core, weights) -> set: + """The vision tower and its projectors, and the multi-token-prediction head: in the checkpoint, never built + by this target.""" + return {k for k in weights.keys() if k.startswith(_UNBUILT_PREFIXES)} + + def load(self, model, weights) -> None: + """Fill every declared weight and the shell's ``lm_head`` shard; assert coverage both ways.""" + core = model.model + d = self.dims(core) + manifest = self.manifest(core) + assert set(manifest) == set(core.w.keys()), ( + "manifest/parameter drift", + set(manifest) ^ set(core.w.keys()), + ) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(core, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(core.w[param_key], ckpt, index, transform, weights, d, consumed) + if param_key.endswith(self.RELEASE_AFTER): + torch.cuda.empty_cache() + # The shell owns the vocab-parallel head: this rank's rows of the full table. + _fill(model.lm_head.weight, _LM + "lm_head.weight", None, _col_w, weights, d, consumed) + shell = { + n for n, _ in model.named_parameters() if not n.startswith(("model.w.", "draft_model.")) + } + assert shell == {"lm_head.weight"}, f"unfed shell parameters: {sorted(shell)}" + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + "checkpoint coverage: leftover keys are not exactly the unbuilt families (vision, MTP head); " + f"unexpected {sorted(leftover - expected)[:8]}, missed {sorted(expected - leftover)[:8]}" + ) + + +def _materialize(entry) -> torch.Tensor: + """``[:]`` on a lazy safetensors slice; a 0-dim tensor (the per-tensor scalars) rejects that index.""" + return entry if getattr(entry, "ndim", 1) == 0 else entry[:] + + +def _fill(param, ckpt, index, transform, weights, d, consumed) -> None: + keys = ckpt if isinstance(ckpt, tuple) else (ckpt,) + for key in keys: + assert key in weights, f"checkpoint key missing: {key}" + dst = param.data if index is None else param.data[index] + d.dev = dst.device + # Host tensors: a transform slices this rank's part on the host and moves only that to the device. + src_all = tuple(_materialize(weights[key]) for key in keys) + if transform is not None: + src = transform(src_all if isinstance(ckpt, tuple) else src_all[0], d) + else: + assert len(src_all) == 1, "a multi-key row needs a transform" + src = src_all[0] + if src.dim() == 0 and dst.dim() == 0: + pass + elif src.numel() == 1 and dst.numel() == 1: + src = src.reshape(dst.shape) + src = src.to(dst.device, non_blocking=True) + assert tuple(dst.shape) == tuple(src.shape), (ckpt, tuple(dst.shape), tuple(src.shape)) + assert src.dtype == dst.dtype, (ckpt, src.dtype, dst.dtype) + dst.copy_(src, non_blocking=True) + consumed.update(keys) + _release(weights, keys) + + +def _release(weights, keys) -> None: + """Drop consumed checkpoint tensors from the engine's dict as the load goes (bounded host memory).""" + if hasattr(weights, "mark_consumed_keys"): + weights.mark_consumed_keys(keys) + else: + for key in keys: + weights.pop(key, None) + + +MODEL_WEIGHTS = MiniMaxM3Tp8Weights() + + +# -------------------------------------------------------------------------------------------- the drafter + + +def _draft_one(template: str): + return lambda d, i, t=template: [(t, None)] + + +def _draft_many(*templates: str): + return lambda d, i, ts=templates: [(tuple(ts), None)] + + +def _draft_fc_norms(d): + return range(d.capture_layers) + + +def _draft_qkv(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + hd = d.head_dim + return torch.cat([_col(q, d), _heads(k, d, hd), _heads(v, d, hd)], dim=0).contiguous() + + +class MiniMaxM3Tp8DraftWeights: + """The one-layer BF16 Eagle3 drafter (``speculative_model``), tp8. + + The fc over the three captured target layers and its per-layer fc_norms, the layer's three norms, its QKV + (input ``[embeds | hidden]``, so twice the hidden width; Q by heads, K/V by KV heads, each KV head on two + ranks) and o_proj, its gated MLP, the final norm and the full-vocabulary head (vocab-parallel). The + checkpoint's ``embed_tokens`` is the target's table (the drafter reads the target's); it is a predicted non-load. + """ + + WEIGHTS: tuple[W, ...] = ( + W( + "fc", + shape=lambda d: (d.hidden, d.capture_layers * d.hidden), + src=_draft_one("fc.weight"), + layers=None, + ), + W( + "fc_norm", + shape=lambda d: (d.hidden,), + src=lambda d, i: [(f"fc_norm.{i}.weight", None)], + layers=_draft_fc_norms, + ), + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.input_layernorm.weight"), + layers=None, + ), + W( + "hidden_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.hidden_norm.weight"), + layers=None, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.post_attention_layernorm.weight"), + layers=None, + ), + W( + "qkv", + shape=lambda d: (d.qkv_rows, 2 * d.hidden), + src=_draft_many(*(f"layers.0.self_attn.{p}_proj.weight" for p in "qkv")), + transform=_draft_qkv, + layers=None, + ), + W( + "o", + shape=lambda d: (d.hidden, d.q_width), + src=_draft_one("layers.0.self_attn.o_proj.weight"), + transform=_row_w, + layers=None, + ), + W( + "gu", + shape=lambda d: (2 * d.inter, d.hidden), + src=_draft_many("layers.0.mlp.gate_proj.weight", "layers.0.mlp.up_proj.weight"), + transform=_gate_up_w, + layers=None, + ), + W( + "dn", + shape=lambda d: (d.hidden, d.inter), + src=_draft_one("layers.0.mlp.down_proj.weight"), + transform=_row_w, + layers=None, + ), + W("final_norm", shape=lambda d: (d.hidden,), src=_draft_one("norm.weight"), layers=None), + ) + + def dims(self, draft) -> SimpleNamespace: + cfg = draft.draft_config.pretrained_config + mapping = draft.draft_config.mapping + tp = mapping.tp_size + head_dim = int(getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads)) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + return SimpleNamespace( + hidden=cfg.hidden_size, + capture_layers=draft.capture_layers, + head_dim=head_dim, + q_width=heads * head_dim, + qkv_rows=(heads + 2 * kv_heads) * head_dim, + inter=cfg.intermediate_size // tp, + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + if entry.layers is None: + return [(entry.name, None)] + return [(f"{entry.name}{i}", i) for i in entry.layers(d)] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, draft) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + d = self.dims(draft) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, draft, weights) -> set: + """The drafter checkpoint's own embedding: the drafter reads the target's table.""" + return {"embed_tokens.weight"} & set(weights.keys()) + + def load(self, draft, weights) -> None: + d = self.dims(draft) + manifest = self.manifest(draft) + assert set(manifest) == set(draft.w.keys()), set(manifest) ^ set(draft.w.keys()) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(draft, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(draft.w[param_key], ckpt, index, transform, weights, d, consumed) + _fill(draft.lm_head.weight, "lm_head.weight", None, _col_w, weights, d, consumed) + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + f"drafter checkpoint coverage: unexpected {sorted(leftover - expected)[:8]}, " + f"missed {sorted(expected - leftover)[:8]}" + ) + + +DRAFT_WEIGHTS = MiniMaxM3Tp8DraftWeights() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/routing.py new file mode 100644 index 000000000000..9b60e91f3165 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/routing.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Where a MiniMaxM3SparseForConditionalGeneration 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 MiniMax-M3 implementation; in ``require`` it raises that error. + +The KV cache the engine runs with is part of the decision: the targets are +written for an FP8 or NVFP4 main KV cache. The rest of their runtime (the MSA +sparse attention implementation with an FP8 indexer) is an engine argument +rather than a checkpoint fact, so the targets check it at construction +(``_validate_runtime``). + +The checkpoint declares the vision-language wrapper, and the language model's +shape lives in its ``text_config``, so that is what this tree reads. The +targets are text-only: they load no vision tower and refuse image input. + +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Optional, Tuple + +from tensorrt_llm.llmapi.llm_args import EagleDecodingConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +from ..._router_index import NoTarget, engine_kv_cache_quant_algo, routes + +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + +# The one GPU architecture these targets are written for: the decode kernels +# (the tcgen05 GEMVs, the fused MoE engines, the attention megakernels and the +# multicast all-reduce) are built and measured on GB300 only. +_SM = (10, 3) + +# The checkpoint's config shape, read from ``text_config``; see the note in the +# gpt_oss routing module for what shape-sniffing does and does not pin. +# +# (num_hidden_layers, hidden_size, num_local_experts, intermediate_size, +# num_attention_heads, num_key_value_heads). The expert width and the head +# counts fix the per-rank operand shapes the decode kernels are compiled for. +_SHAPE = (60, 6144, 128, 3072, 64, 4) + +# The architecture class and the shape do not pin everything the targets are built for: a newer +# checkpoint revision of this model family has the same architecture class and shapes but a different +# quant packing (compressed-tensors NVFP4 experts). The targets build a 5e6 RoPE table and dense +# attention on layers 0-2 (sparse on 3-59), so these two facts, with the quantization below, keep any +# other revision off the targets. +_ROPE_THETA = 5_000_000 +_SPARSE_ATTENTION_FREQ = (0,) * 3 + (1,) * 57 + + +def _field(config: Any, name: str) -> Any: + """Read ``name`` off a sub-config that may be an object or a plain dict.""" + if isinstance(config, dict): + return config.get(name) + return getattr(config, name, None) + + +def _layout_mismatch(text_config: Any) -> Optional[str]: + """Why the text model's RoPE base or sparse-attention layers are not the ones the targets build, or None.""" + if text_config is None: + return None # the shape check declines it + theta = _field(text_config, "rope_theta") + if theta is None: + rope = _field(text_config, "rope_parameters") or {} + theta = _field( + rope.get("full_attention", rope) if isinstance(rope, dict) else rope, "rope_theta" + ) + sparse = _field(text_config, "sparse_attention_config") + freq = None if sparse is None else _field(sparse, "sparse_attention_freq") + # The built-in model reads use_sparse_attention=False as dense attention on every layer. + enabled = sparse is not None and _field(sparse, "use_sparse_attention") is not False + reasons = [] + if theta != _ROPE_THETA: + reasons.append(f"rope_theta {theta}, not {_ROPE_THETA}") + if not enabled or freq is None or tuple(int(f) for f in freq) != _SPARSE_ATTENTION_FREQ: + reasons.append("sparse attention layers other than 3-59") + return "; ".join(reasons) if reasons else None + + +def _quant_algo(pretrained_config: Any) -> Optional[QuantAlgo]: + """The quantization the checkpoint's config declares, as the engine reads it. + + The NVFP4 checkpoint is a modelopt mixed-precision export (NVFP4 routed + experts, MXFP8 dense projections, bf16 router, embedding and head), declared + in ``config.json`` as ``MIXED_PRECISION``; the per-layer split lives in + ``hf_quant_config.json``, which ``route`` does not see. + """ + from tensorrt_llm._torch.model_config import ModelConfig + + declared = getattr(pretrained_config, "quantization_config", None) + if declared is None: + return None + try: + quant_config, _ = ModelConfig.load_hf_quant_config(declared, "TRTLLM") + except ValueError: + return None + return quant_config.quant_algo + + +def _parallel(m) -> Optional[str]: + """Name the parallel topology, or None if no target implements it. + + Attention is head-split and the routed experts are tensor-parallel with + every expert on every rank: the expert split decides which expert weights + each rank loads and into which shapes (per-rank expert width 768 at tp4, + 384 at tp8), so it selects a target rather than a runtime branch. + """ + for tp in (4, 8): + if ( + m.world_size == tp + and m.tp_size == tp + and m.pp_size == 1 + and m.cp_size == 1 + and m.moe_tp_size == tp + and m.moe_ep_size == 1 + and not m.enable_attention_dp + ): + return f"tp{tp}" + return None + + +def _eagle3_one_model_or_none(spec_config: Optional[Any]) -> bool: + """No speculation, or the one-model Eagle3 loop the target's worker extends.""" + if spec_config is None: + return True + return ( + isinstance(spec_config, EagleDecodingConfig) + and spec_config.spec_dec_mode.is_eagle3_one_model() + ) + + +@routes("MiniMaxM3SparseForConditionalGeneration") +def route( + model_config: ModelConfig, llm_args: TorchLlmArgs, compute_capability: Tuple[int, int] +) -> 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. + """ + pretrained_config = model_config.pretrained_config + mapping = model_config.mapping + spec_config = model_config.spec_config + kv_cache_quant_algo = engine_kv_cache_quant_algo(model_config, llm_args) + text_config = getattr(pretrained_config, "text_config", None) + shape = ( + None + if text_config is None + else tuple( + _field(text_config, name) + for name in ( + "num_hidden_layers", + "hidden_size", + "num_local_experts", + "intermediate_size", + "num_attention_heads", + "num_key_value_heads", + ) + ) + ) + quant_algo = _quant_algo(pretrained_config) + parallel = _parallel(mapping) + layout_miss = _layout_mismatch(text_config) + speculation = _eagle3_one_model_or_none(spec_config) + kv_cache = kv_cache_quant_algo in (QuantAlgo.FP8, QuantAlgo.NVFP4) + + # MiniMax-M3 NVFP4 on 4 GB300s (sm103), tensor-parallel attention and + # experts, one-model Eagle3 or no speculative decoding. + if ( + compute_capability == _SM + and shape == _SHAPE + and layout_miss is None + and quant_algo == QuantAlgo.MIXED_PRECISION + and parallel == "tp4" + and speculation + and kv_cache + ): + from .minimax_m3_nvfp4__sm_103__tp4.modeling import ModelingV2MiniMaxM3Nvfp4Sm103Tp4 + + return ModelingV2MiniMaxM3Nvfp4Sm103Tp4 + + # The same checkpoint on 8 GB300s over two trays (sm103), tensor-parallel + # attention and experts, one-model Eagle3 or no speculative decoding. + if ( + compute_capability == _SM + and shape == _SHAPE + and layout_miss is None + and quant_algo == QuantAlgo.MIXED_PRECISION + and parallel == "tp8" + and speculation + and kv_cache + ): + from .minimax_m3_nvfp4__sm_103__tp8.modeling import ModelingV2MiniMaxM3Nvfp4Sm103Tp8 + + return ModelingV2MiniMaxM3Nvfp4Sm103Tp8 + + raise NoTarget( + f"compute capability {compute_capability}, text_config shape {shape}" + f"{f' ({layout_miss})' if layout_miss else ''}, " + f"quantization {quant_algo.name if quant_algo is not None else 'none'}, parallel " + f"ws={mapping.world_size} tp={mapping.tp_size} pp={mapping.pp_size} cp={mapping.cp_size} " + f"moe_tp={mapping.moe_tp_size} moe_ep={mapping.moe_ep_size} " + f"attention_dp={mapping.enable_attention_dp}, speculative decoding " + f"{spec_config.spec_dec_mode.name if spec_config is not None else 'off'}" + f", KV cache {kv_cache_quant_algo or 'bf16'}; the " + "minimax_m3 targets are the MiniMax-M3 NVFP4 checkpoint on sm103 (GB300), 4 or 8 ranks of " + "tensor-parallel attention and experts, with one-model Eagle3 or no speculative decoding, and an FP8 or " + "NVFP4 KV cache" + ) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py index e465eda7d093..c2ac19e425a5 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py @@ -1207,7 +1207,6 @@ def run_generation( bmm1_scale if fp8_context_fmha and bmm1_scale is not None else bmm1_scale_static ) gen_bmm2_scale = bmm2_scale if fp8_context_fmha and bmm2_scale is not None else 1.0 - flashinfer.decode.trtllm_batch_decode_with_kv_cache( query=q_processed, kv_cache=(kv_pool, kv_pool), diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py index 120feaafeeee..bc1301e57c1b 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py @@ -228,24 +228,39 @@ def run_msa_nvfp4_sparse_gqa( ) cu_q = metadata.msa_cu_q_lens[: batch + 1] cu_kv = metadata.msa_cu_kv_lens[: batch + 1] - q2k = kv_block_indexes.permute(1, 0, 2).contiguous() - topk = int(q2k.shape[-1]) + topk = int(kv_block_indexes.shape[-1]) page_size = int(k_paged.shape[2]) if topk != MSA_REQUIRED_TOPK: raise ValueError(f"MiniMax-M3 MSA NVFP4 requires topK={MSA_REQUIRED_TOPK}, got {topk}") - k2q_row_ptr, k2q_q_indices, schedule = sparse.build_k2q_csr( - q2k, - cu_q, - cu_kv, - page_size, - total_k=total_k, - max_seqlen_k=max_kv_len, - max_seqlen_q=max_q_len, - total_rows=total_k_rows, - qhead_per_kv=int(q.shape[1]) // int(k_paged.shape[1]), - return_schedule=True, - ) + # A short pure-prefill step hands every layer the same identity selection (staged once per step by the + # backend), so its k2q CSR worklist is layer-invariant: build it for the first layer and reuse it. + identity = metadata._msa_prefill_identity + csr_key = (batch, int(kv_block_indexes.shape[0])) + cached = metadata._msa_identity_csr + if ( + identity is not None + and cached is not None + and cached[0] is identity + and cached[1] == csr_key + ): + k2q_row_ptr, k2q_q_indices, schedule = cached[2] + else: + q2k = kv_block_indexes.permute(1, 0, 2).contiguous() + k2q_row_ptr, k2q_q_indices, schedule = sparse.build_k2q_csr( + q2k, + cu_q, + cu_kv, + page_size, + total_k=total_k, + max_seqlen_k=max_kv_len, + max_seqlen_q=max_q_len, + total_rows=total_k_rows, + qhead_per_kv=int(q.shape[1]) // int(k_paged.shape[1]), + return_schedule=True, + ) + if identity is not None: + metadata._msa_identity_csr = (identity, csr_key, (k2q_row_ptr, k2q_q_indices, schedule)) result = sparse.sparse_atten_nvfp4_kv_func( q, k_paged.view(torch.uint8), diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index 7722f76ea664..e05523c35224 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -28,10 +28,11 @@ from __future__ import annotations -from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple import torch +from tensorrt_llm._device_copy import DeviceCopyBuffer from tensorrt_llm._torch.disaggregation.resource.page import MapperKind from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( KVCacheManagerV2, @@ -324,6 +325,10 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): :meth:`_prepare_page_table_tensor`. """ + # Decode-step shortcuts a model opts into (RuntimeHooks.msa); see configure_decode_shortcuts. + skip_unchanged_block_offsets = False + reuse_pool_views = False + _main_kv_mapper_kind = MapperKind.HND _main_kv_view_layout = "NHD" _dense_native_layer_ids: Tuple[int, ...] = () @@ -404,6 +409,8 @@ def __init__( self._kv_buffers: dict[tuple[int, str], torch.Tensor] = {} self._block_scale_buffers: dict[tuple[int, str], torch.Tensor] = {} super().__init__(*args, **kwargs) + self._unchanged_copies: Dict[int, DeviceCopyBuffer] = {} + self._pool_views: Dict[tuple, torch.Tensor] = {} if self.dtype == DataType.NVFP4: dense_target_layers = [ @@ -698,6 +705,29 @@ def copy_batch_block_offsets( ) -> None: if self.num_attention_op_pools == 0: return + guard = None + key = ( + dst_tensor.dtype, + dst_tensor.storage_offset(), + tuple(dst_tensor.shape), + tuple(dst_tensor.stride()), + tuple(request_ids), + beam_width, + num_seqs, + max_blocks, + ) + if self.skip_unchanged_block_offsets: + address = dst_tensor.untyped_storage().data_ptr() + guard = self._unchanged_copies.setdefault(address, DeviceCopyBuffer(dst_tensor)) + if num_contexts != 0 or ( + dst_tensor.is_cuda and torch.cuda.is_current_stream_capturing() + ): + guard.invalidate() + guard = None + elif guard.matches(self.host_kv_cache_block_offsets, key=key, exact=True): + return + else: + guard.invalidate() # The base materializer owns only the storage-pool output rows; virtual # draft pools are filled below with their own conversion metadata. base_dst = dst_tensor[: self.num_pools] if self._draft_op_pools else dst_tensor @@ -705,6 +735,8 @@ def copy_batch_block_offsets( base_dst, request_ids, beam_width, num_contexts, num_seqs, max_blocks=max_blocks ) if not self._draft_op_pools: + if guard is not None: + guard.record(self.host_kv_cache_block_offsets, key=key) return # Fill each virtual pool from its source pool's slot ids with the draft # layer's scale. @@ -724,6 +756,9 @@ def copy_batch_block_offsets( self._stream.cuda_stream, ) + if guard is not None: + guard.record(self.host_kv_cache_block_offsets, key=key) + def update_resources(self, scheduled_batch, attn_metadata=None, kv_cache_dtype_byte_size=None): # Only tree acceptance relocates draft KV, which M3's pool layout cannot do. if any(r.py_num_accepted_draft_tokens_indices for r in scheduled_batch.generation_requests): @@ -907,12 +942,71 @@ def get_buffers( addr_key, torch_dtype, num_slots, scale, page_shape = self._kv_slot_geometry( layer_idx, self._main_kv_layout_name() ) - full_slot_shape = [num_slots, scale, *page_shape] - full_view = convert_to_torch_tensor(TensorWrapper(addr_key, torch_dtype, full_slot_shape)) - kv = full_view[:, :2] # Sparse reference kernels use token-major indexing; preserve the # native head-major backing and expose a zero-copy strided view. - return kv if kv_layout == self._main_kv_layout_name() else kv.transpose(2, 3) + return self._pool_view( + "kv", + addr_key, + torch_dtype, + num_slots, + scale, + page_shape, + transpose=kv_layout != self._main_kv_layout_name(), + ) + + def configure_decode_shortcuts( + self, skip_unchanged_block_offsets: bool, reuse_pool_views: bool + ) -> None: + """Turn on the decode-step shortcuts a model opted into. + + ``skip_unchanged_block_offsets``: skip a decode-only block-table copy + whose inputs equal the previous copy's into the same destination. + ``reuse_pool_views``: hand out one pool view per geometry instead of + re-wrapping the pool memory on every attention call. + """ + self.skip_unchanged_block_offsets = skip_unchanged_block_offsets + self.reuse_pool_views = reuse_pool_views + + def _pool_view( + self, + kind: str, + addr_key: int, + torch_dtype: torch.dtype, + num_slots: int, + scale: int, + page_shape: List[int], + transpose: bool = False, + ) -> torch.Tensor: + """The ``[num_slots, 2, ...]`` K+V view of a pool rooted at ``addr_key``; ``transpose`` swaps the + page's head and token axes (the token-major view of head-major storage). + + The view is a pure function of the pool geometry, which is static once the + pool is laid out, so with ``reuse_pool_views`` it is built once per geometry + (a re-laid-out pool has another key). Inference-mode views are inference + tensors, so each mode keeps its own. + """ + key = ( + kind, + addr_key, + torch_dtype, + num_slots, + scale, + tuple(page_shape), + transpose, + torch.is_inference_mode_enabled(), + ) + view = self._pool_views.get(key) if self.reuse_pool_views else None + if view is None: + full_slot_shape = [num_slots, scale, *page_shape] + full_view = convert_to_torch_tensor( + TensorWrapper(addr_key, torch_dtype, full_slot_shape) + ) + view = full_view[:, :2] + if transpose: + view = view.transpose(2, 3) + if self.reuse_pool_views: + self._pool_views[key] = view + return view def _kv_scale_slot_geometry( self, layer_idx: int, kv_layout: Optional[str] @@ -983,9 +1077,7 @@ def get_block_scale_buffers( addr_key, torch_dtype, num_slots, scale, page_shape = self._kv_scale_slot_geometry( layer_idx, kv_layout ) - full_slot_shape = [num_slots, scale, *page_shape] - full_view = convert_to_torch_tensor(TensorWrapper(addr_key, torch_dtype, full_slot_shape)) - return full_view[:, :2] + return self._pool_view("scale", addr_key, torch_dtype, num_slots, scale, page_shape) def get_kv_subpage_pool( self, layer_idx: int, kv_layout: str = "HND" @@ -1147,6 +1239,7 @@ def __init__(self, manager: MiniMaxM3KVCacheManagerV2, draft_layer_ids: Sequence if manager.tokens_per_block in self.trtllm_gen_extra_tokens_per_block: _require_trtllm_gen_extra_page_route(manager.tokens_per_block) self._manager = manager + self._unchanged_copies: Dict[int, DeviceCopyBuffer] = {} layer_id = draft_layer_ids[0] local = int(manager.layer_offsets[layer_id]) self._local_layer_idx = local @@ -1217,19 +1310,44 @@ def copy_batch_block_offsets( """ assert beam_width == 1, "beam_width must be 1 for KVCacheManagerV2" manager = self._manager + guard = None + key = ( + dst_tensor.dtype, + dst_tensor.storage_offset(), + tuple(dst_tensor.shape), + tuple(dst_tensor.stride()), + tuple(request_ids), + beam_width, + num_seqs, + max_blocks, + ) + if manager.skip_unchanged_block_offsets: + address = dst_tensor.untyped_storage().data_ptr() + guard = self._unchanged_copies.setdefault(address, DeviceCopyBuffer(dst_tensor)) + if num_contexts != 0 or ( + dst_tensor.is_cuda and torch.cuda.is_current_stream_capturing() + ): + guard.invalidate() + guard = None + elif guard.matches(self.host_kv_cache_block_offsets, key=key, exact=True): + return + else: + guard.invalidate() copy_idx = manager.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) assert copy_idx.shape[0] == num_seqs if self._page_table_materializer.uses_device_expansion: self._page_table_materializer.copy_block_offsets_to(dst_tensor, copy_idx) - return - copy_batch_block_offsets_to_device( - self.host_kv_cache_block_offsets, - dst_tensor, - copy_idx, - self.index_scales, - self.kv_offset, - manager._stream.cuda_stream, - ) + else: + copy_batch_block_offsets_to_device( + self.host_kv_cache_block_offsets, + dst_tensor, + copy_idx, + self.index_scales, + self.kv_offset, + manager._stream.cuda_stream, + ) + if guard is not None: + guard.record(self.host_kv_cache_block_offsets, key=key) def get_minimax_m3_kv_cache_manager_cls(): diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py index b527fbbdced7..29a13c570d28 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py @@ -227,8 +227,8 @@ def from_sparse_params( class PagedKvSlotMapping(NamedTuple): """One step's paged-cache slot mapping (see build_paged_kv_slot_mapping).""" - req_to_token: torch.Tensor - slot_ids: torch.Tensor + req_to_token: Optional[torch.Tensor] + slot_ids: Optional[torch.Tensor] out_cache_loc: torch.Tensor block_ids_cpu: torch.Tensor @@ -240,6 +240,7 @@ def build_paged_kv_slot_mapping( qo_lens_cpu: torch.Tensor, qo_offset_cpu: torch.Tensor, device: torch.device, + device_tables: bool = True, ) -> PagedKvSlotMapping: """Build the backend-neutral paged-cache slot mapping. @@ -262,6 +263,9 @@ def build_paged_kv_slot_mapping( staged with non-blocking copies, so this neither syncs the host nor reads device memory back. It still allocates per step, so call it from prepare(), never from the forward path. + + With device_tables=False, req_to_token and slot_ids are not built (None) and out_cache_loc stays a pinned host + tensor, for callers that only stage out_cache_loc into their own persistent buffer. """ tokens_per_block = int(kv_cache_manager.tokens_per_block) # block_ids_per_seq is a [batch, max_blocks_per_seq] tensor; row b holds the @@ -273,13 +277,15 @@ def build_paged_kv_slot_mapping( max_blocks = int(block_ids.shape[1]) max_kv_len = max_blocks * tokens_per_block - # Expand block ids -> per-token slot ids. - block_ids_dev = block_ids.to(device, non_blocking=True).to(torch.int64) - within_block = torch.arange(tokens_per_block, device=device, dtype=torch.int64) - # Outer product per batch entry: [batch, max_blocks, tokens_per_block] - slot_grid = block_ids_dev.unsqueeze(-1) * tokens_per_block + within_block - req_to_token = slot_grid.reshape(batch, max_kv_len).to(torch.int32) - slot_ids = torch.arange(batch, device=device, dtype=torch.int32) + req_to_token = slot_ids = None + if device_tables: + # Expand block ids -> per-token slot ids. + block_ids_dev = block_ids.to(device, non_blocking=True).to(torch.int64) + within_block = torch.arange(tokens_per_block, device=device, dtype=torch.int64) + # Outer product per batch entry: [batch, max_blocks, tokens_per_block] + slot_grid = block_ids_dev.unsqueeze(-1) * tokens_per_block + within_block + req_to_token = slot_grid.reshape(batch, max_kv_len).to(torch.int32) + slot_ids = torch.arange(batch, device=device, dtype=torch.int32) # out_cache_loc: per-new-token slot ids, in flattened query-token order. # Expanding the per-request lengths on the host reproduces the same slot @@ -304,7 +310,10 @@ def build_paged_kv_slot_mapping( block_ids[row, block_col].to(torch.long) * tokens_per_block + (pos - block_col * tokens_per_block) ).to(torch.int32) - out_cache_loc = async_tensor_h2d(out_cache_loc_cpu, torch.int32, device) + if device_tables: + out_cache_loc = async_tensor_h2d(out_cache_loc_cpu, torch.int32, device) + else: + out_cache_loc = maybe_pin_memory(out_cache_loc_cpu) return PagedKvSlotMapping(req_to_token, slot_ids, out_cache_loc, block_ids) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py index dca8a6abd44c..7e28b67fef30 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py @@ -38,6 +38,9 @@ from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._torch.pyexecutor import h2d_cache, step_inputs +from tensorrt_llm._torch.pyexecutor.h2d_cache import copy_h2d_if_changed +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy from tensorrt_llm._utils import maybe_pin_memory from tensorrt_llm.bindings import DataType from tensorrt_llm.models.modeling_utils import QuantConfig @@ -102,6 +105,17 @@ def _worst_case_proxy_max_k_tiles( return int(proxy_plan[3]["max_k_tiles"]) +def _host_prefix_sum(lens_cpu: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """[0, cumsum(lens)...] as a pinned host tensor of ``dtype``.""" + sums = torch.cumsum(lens_cpu.to(torch.int64), 0) + return maybe_pin_memory(torch.cat([sums.new_zeros(1), sums]).to(dtype)) + + +def _fill_unused_slots(buf: torch.Tensor, start: int) -> None: + """buf[start:] = -1 after a write of the live slots buf[:start] (see h2d_cache.fill_tail_if_changed).""" + h2d_cache.fill_tail_if_changed(buf, start, -1) + + class MsaDecodeSpan(NamedTuple): """Generation rows served by MiniMax-M3's dedicated decode kernels. @@ -160,6 +174,13 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): msa_kv_indices: Optional[torch.Tensor] = None msa_max_score: Optional[torch.Tensor] = None msa_n_valid_blocks: Optional[torch.Tensor] = None + # Decode-step preparation the model opted into (RuntimeHooks.msa). + msa_policy: MsaRuntimePolicy = MsaRuntimePolicy() + # Short-sequence steps (MsaRuntimePolicy.short_sequence_identity): every query token has at most topk valid + # blocks, so the selection is the identity table [kv_heads, tokens, topk] = (i < n_valid ? i : -1), staged here + # by prepare(). _msa_short_ok says whether this step staged it. + msa_short_table: Optional[torch.Tensor] = None + _msa_short_ok: bool = False # The same page table and lengths as msa_kv_indices / msa_kv_lens, in the # per-request 2-D form the decode kernels index directly # (block_table[request, block] and seq_lens[request]). fmha_sm100 instead @@ -204,6 +225,12 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): # _msa_prefill_n_valid_buf is the persistent backing store for the view. _msa_prefill_n_valid_buf: Optional[torch.Tensor] = None _msa_prefill_n_valid_blocks: Optional[torch.Tensor] = None + # Pure-prefill step where every query token has at most topk valid blocks: selection is the identity, so the + # indexer returns this [total_q, topk] table (per KV head) instead of scoring (built once per step). + _msa_prefill_identity: Optional[torch.Tensor] = None + # (identity table, (batch, rows), k2q CSR worklist) the NVFP4 prefill built for it: the worklist is + # layer-invariant for an identity selection, so the later layers of the step reuse it. + _msa_identity_csr: Optional[tuple] = None # This step's per-request host lengths, staged by _stage_host_lengths. _msa_qo_lens_cpu: Optional[torch.Tensor] = None _msa_kv_lens_cpu: Optional[torch.Tensor] = None @@ -657,6 +684,40 @@ def _ensure_prefill_n_valid_buffer(self, total_q: int, device: torch.device) -> self._msa_prefill_n_valid_buf = buf return buf + def _stage_msa_short_table(self, n_valid: torch.Tensor, params) -> None: + """Stage the identity block table for short-sequence steps (see msa_short_table).""" + self._msa_short_ok = False + if ( + not self.msa_policy.short_sequence_identity + or n_valid.numel() == 0 + or self.num_contexts > 0 + ): + return + topk = int(MSA_REQUIRED_TOPK) + if int(n_valid.max()) > topk: + return + _, num_kv_heads = params.sharded_head_counts(self.mapping) + total_q = int(n_valid.shape[0]) + max_tokens = max(int(self.max_num_sequences), self._msa_max_decode_tokens(), total_q) + if ( + self.msa_short_table is None + or self.msa_short_table.numel() < num_kv_heads * max_tokens * topk + ): + self.msa_short_table = self.get_empty( + self.cuda_graph_buffers, + (num_kv_heads * max_tokens * topk,), + cache_name="msa_short_table", + dtype=torch.int32, + capture_graph=self.is_cuda_graph, + ) + ar = torch.arange(topk, dtype=torch.int32) + table = torch.where(ar[None, :] < n_valid.to(torch.int32)[:, None], ar[None, :], -1) + table = table.unsqueeze(0).expand(num_kv_heads, total_q, topk).contiguous() + copy_h2d_if_changed( + self.msa_short_table[: table.numel()], maybe_pin_memory(table.flatten()), pack=True + ) + self._msa_short_ok = True + def prepare(self) -> None: super().prepare() self._check_beam_width() @@ -854,6 +915,29 @@ def on_update_kv_lens(self) -> None: total_q = int(self.num_tokens) if batch <= 0 or total_q <= 0: return + page = int(self.kv_cache_manager.tokens_per_block) + n_valid_buf = ( + self._msa_prefill_n_valid_blocks + if self._msa_prefill_n_valid_blocks is not None + else self.msa_n_valid_blocks + ) + kv_update = self.msa_policy.kv_update + if kv_update is not None and kv_update( + self.kv_lens_cuda, + self.msa_kv_lens_staged, + self.seq_lens_cuda, + self.msa_q_batch_row, + self.msa_q_intra, + self.msa_block_table, + page, + batch, + total_q, + self.msa_seq_lens_cuda, + None if self.msa_cu_kv_lens is None else self.msa_cu_kv_lens[: batch + 1], + self.msa_out_cache_loc, + None if n_valid_buf is None else n_valid_buf[:total_q], + ): + return # kv_lens_cuda is the attended length (no num_extra_kv_tokens), the same # domain as the staged bound and msa_seq_lens_cuda. kv_true = torch.minimum(self.kv_lens_cuda[:batch], self.msa_kv_lens_staged[:batch]) @@ -915,6 +999,8 @@ def _build_step_plans(self) -> None: self._msa_prefill_gqa_plan = None self._msa_prefill_dense_plan = None self._msa_prefill_n_valid_blocks = None + self._msa_prefill_identity = None + self._msa_short_ok = False if not self._msa_fields_ready: return # Geometry is captured in __post_init__; skip when it is unavailable. @@ -945,7 +1031,9 @@ def _build_step_plans(self) -> None: qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size ) total_q = int(n_valid.shape[0]) - self.msa_n_valid_blocks[:total_q].copy_(n_valid.to(torch.int32), non_blocking=True) + # Packed but never skipped: on_update_kv_lens rewrites this buffer on the device. + step_inputs.copy_(self.msa_n_valid_blocks[:total_q], n_valid.to(torch.int32)) + self._stage_msa_short_table(n_valid, params) return fmha_sm100 = require_msa_module() @@ -976,10 +1064,21 @@ def plan_for(rows: Optional[Tuple[int, int]], **plan_kwargs) -> Optional[tuple]: **plan_kwargs, ) + n_valid_host = per_token_valid_blocks( + qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size + ) + total_q = int(n_valid_host.shape[0]) + short_prefill = ( + self.msa_policy.short_sequence_identity + and self._msa_decode_span is None + and total_q > 0 + and int(n_valid_host.max()) <= int(params.topk) + ) # Proxy plan: MQA (num_kv_heads=1) max-score pass over the index # branch; output_maxscore feeds the indexer's top-k block selection. - # Its rows are a batch prefix, so their pages prefix msa_kv_indices. - proxy_rows = self._msa_proxy_plan_rows() + # A short pure-prefill step selects every valid block, so it needs none. + # Otherwise its rows are a batch prefix, so their pages prefix msa_kv_indices. + proxy_rows = None if short_prefill else self._msa_proxy_plan_rows() segmented = proxy_rows is not None and segmented_proxy_plan( fmha_sm100, qo_lens_cpu[: proxy_rows[1]], @@ -1016,14 +1115,15 @@ def plan_for(rows: Optional[Tuple[int, int]], **plan_kwargs) -> Optional[tuple]: ) # Stage the valid-block count to the device once for the whole step # (see msa_prefill_n_valid_blocks). - n_valid_host = per_token_valid_blocks( - qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size - ) - total_q = int(n_valid_host.shape[0]) if total_q > 0: dev_buf = self._ensure_prefill_n_valid_buffer(total_q, _cache_device(self)) dev_buf[:total_q].copy_(n_valid_host.to(torch.int32), non_blocking=True) self._msa_prefill_n_valid_blocks = dev_buf[:total_q] + if short_prefill: + topk = int(params.topk) + ar = torch.arange(topk, dtype=torch.int32, device=dev_buf.device) + table = torch.where(ar[None, :] < dev_buf[:total_q, None], ar[None, :], -1) + self._msa_prefill_identity = table def _build_msa_fields(self) -> None: """Populate the MSA cache-write buffers for this step. @@ -1041,11 +1141,11 @@ def _build_msa_fields(self) -> None: kv_lens_cpu = self.msa_kv_lens_cpu qo_offset_cpu = self.msa_qo_offset_cpu if request_ids is None or qo_lens_cpu is None: - self.msa_out_cache_loc.fill_(-1) + _fill_unused_slots(self.msa_out_cache_loc, 0) return batch_size = int(qo_lens_cpu.shape[0]) if batch_size == 0: - self.msa_out_cache_loc.fill_(-1) + _fill_unused_slots(self.msa_out_cache_loc, 0) return kv_cache_manager = self.kv_cache_manager @@ -1062,6 +1162,8 @@ def _build_msa_fields(self) -> None: qo_lens_cpu=qo_lens_cpu, qo_offset_cpu=qo_offset_cpu, device=cache_device, + # Only out_cache_loc and the host block ids are read here; stage the slots straight from the host. + device_tables=not h2d_cache.enabled(), ) out_cache_loc = mapping.out_cache_loc # Only fmha_sm100 reads the flattened page table (the decode kernels @@ -1103,41 +1205,62 @@ def _build_msa_fields(self) -> None: f"smaller than the step's per-request page count ({block_table_cols})." ) - self.msa_out_cache_loc[:total_new_tokens].copy_(out_cache_loc, non_blocking=True) + step_inputs.copy_(self.msa_out_cache_loc[:total_new_tokens], out_cache_loc) # Captured producers also execute padded rows: the fused index producer # sits inside the captured region, so trimming it to a host-side count # would make its shape dynamic, and a negative slot is what makes those # rows cache-write no-ops instead. Invalidate the unwritten tail so they # cannot reuse the previous step's live slots, which address real pages. - if total_new_tokens < self.msa_out_cache_loc.shape[0]: - self.msa_out_cache_loc[total_new_tokens:].fill_(-1) + _fill_unused_slots(self.msa_out_cache_loc, total_new_tokens) if kv_indices is not None: - self.msa_kv_indices[: int(kv_indices.shape[0])].copy_(kv_indices, non_blocking=True) + step_inputs.copy_(self.msa_kv_indices[: int(kv_indices.shape[0])], kv_indices) # 2-D page table and per-request length for the decode kernels, # from the same host block ids the flat page table was built from. # Columns past a request's page count are left stale rather than # cleared: every consumer bounds its walk by seq_lens. - self.msa_block_table[:batch_size, :block_table_cols].copy_( - maybe_pin_memory(block_ids_cpu.to(torch.int32)), non_blocking=True + block_table_changed = copy_h2d_if_changed( + self.msa_block_table[:batch_size, :block_table_cols], + maybe_pin_memory(block_ids_cpu.to(torch.int32)), + # write_subpage_block_table below reads it on the device before the replay. + pack=self.msa_subpage_block_table is None, ) - self.msa_seq_lens_cuda[:batch_size].copy_(kv_lens_cpu, non_blocking=True) + # Packed only when the device scan below (h2d_cache off) does not read it before the replay. + if h2d_cache.enabled(): + step_inputs.copy_(self.msa_seq_lens_cuda[:batch_size], kv_lens_cpu) + else: + self.msa_seq_lens_cuda[:batch_size].copy_(kv_lens_cpu, non_blocking=True) if self.msa_cu_q_lens is not None: self._msa_live_batch = batch_size - self.msa_cu_q_lens[0].zero_() - self.msa_cu_kv_lens[0].zero_() qo_lens_cuda = self.msa_qo_lens_cuda[:batch_size] - qo_lens_cuda.copy_(qo_lens_cpu, non_blocking=True) - torch.cumsum( - qo_lens_cuda, - 0, - out=self.msa_cu_q_lens[1 : batch_size + 1], - ) - torch.cumsum( - self.msa_seq_lens_cuda[:batch_size], - 0, - out=self.msa_cu_kv_lens[1 : batch_size + 1], - ) + if h2d_cache.enabled(): + # Both prefix sums from the host lengths (the values the device scans would give), staged with + # their leading zero: copies instead of zero / scan kernels between the decode graphs. + copy_h2d_if_changed(qo_lens_cuda, qo_lens_cpu, pack=True) + copy_h2d_if_changed( + self.msa_cu_q_lens[: batch_size + 1], + _host_prefix_sum(qo_lens_cpu, self.msa_cu_q_lens.dtype), + pack=True, + ) + # Not through the skip cache: on_update_kv_lens rewrites this buffer on the device. + step_inputs.copy_( + self.msa_cu_kv_lens[: batch_size + 1], + _host_prefix_sum(kv_lens_cpu, self.msa_cu_kv_lens.dtype), + ) + else: + self.msa_cu_q_lens[0].zero_() + self.msa_cu_kv_lens[0].zero_() + qo_lens_cuda.copy_(qo_lens_cpu, non_blocking=True) + torch.cumsum( + qo_lens_cuda, + 0, + out=self.msa_cu_q_lens[1 : batch_size + 1], + ) + torch.cumsum( + self.msa_seq_lens_cuda[:batch_size], + 0, + out=self.msa_cu_kv_lens[1 : batch_size + 1], + ) self._msa_max_q_len = int(qo_lens_cpu.max().item()) self._msa_max_kv_len_all = int(kv_lens_cpu.max().item()) self._msa_total_k = int(kv_lens_cpu.to(torch.int64).sum().item()) @@ -1172,7 +1295,7 @@ def _build_msa_fields(self) -> None: # Sub-page expansion for the trtllm-gen dense layers, staged once here # instead of once per layer, outside capture into a graph-stable # buffer as with the slot table above. - if self.msa_subpage_block_table is not None: + if self.msa_subpage_block_table is not None and block_table_changed: write_subpage_block_table( self.msa_block_table[:batch_size], self._msa_subpages_per_slot, @@ -1196,14 +1319,16 @@ def _build_msa_fields(self) -> None: torch.arange(total_new_tokens, dtype=torch.int64) - torch.repeat_interleave(starts, qo_long) ).to(torch.int32) - self.msa_q_batch_row[:total_new_tokens].copy_( - maybe_pin_memory(batch_row_cpu), non_blocking=True + copy_h2d_if_changed( + self.msa_q_batch_row[:total_new_tokens], maybe_pin_memory(batch_row_cpu), pack=True + ) + copy_h2d_if_changed( + self.msa_q_intra[:total_new_tokens], maybe_pin_memory(intra_cpu), pack=True ) - self.msa_q_intra[:total_new_tokens].copy_(maybe_pin_memory(intra_cpu), non_blocking=True) # The staged lens are the upper bound on_update_kv_lens clamps to: the # same attended lengths msa_seq_lens_cuda was just filled from, kept # apart so the clamp stays fixed however often the hook runs. - self.msa_kv_lens_staged[:batch_size].copy_(kv_lens_cpu, non_blocking=True) + step_inputs.copy_(self.msa_kv_lens_staged[:batch_size], kv_lens_cpu) self._msa_fields_ready = True def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor: @@ -1496,6 +1621,11 @@ def run_indexer( # proxy plan scores the whole batch instead. gen_first is the span's # first query token: the scorer takes [gen_first, num_tokens) over rows # [ctx_rows, row_last), the proxy the context prefix ahead of both. + identity = metadata._msa_prefill_identity + if span is None and identity is not None and int(identity.shape[0]) == num_tokens: + # Short pure-prefill step: every token selects all its valid blocks (staged once per step). + num_kv_heads = int(self.indexer.config.num_kv_heads) + return identity.unsqueeze(1).expand(num_tokens, num_kv_heads, int(identity.shape[1])) block_table = None seq_lens_cuda = None decode_query_len = None diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 1a9121037a1d..1bd6764692ba 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -41,7 +41,9 @@ from tensorrt_llm.math_utils import ceil_div from tensorrt_llm.models.modeling_utils import QuantConfig +from ...pyexecutor import step_inputs from ...pyexecutor.config_utils import is_mla +from ...pyexecutor.h2d_cache import copy_h2d_if_changed from ...utils import (compute_swizzled_sf_shape, get_global_attrs, get_model_extra_attrs, helix_local_len_tensor) from ..kernels.sage_quant import SageQuantized, sage_quantize @@ -1012,8 +1014,8 @@ def prepare(self) -> None: device='cpu', ) self.prompt_lens_cpu[:self.num_seqs].copy_(prompt_lens) - self.prompt_lens_cuda[:self.num_seqs].copy_( - self.prompt_lens_cpu[:self.num_seqs], non_blocking=True) + copy_h2d_if_changed(self.prompt_lens_cuda[:self.num_seqs], + self.prompt_lens_cpu[:self.num_seqs]) # number of tokens in the kv cache for each sequence in the batch cached_token_lens = torch.tensor( @@ -1066,9 +1068,8 @@ def prepare(self) -> None: # the sequence length including the cached tokens and the input tokens. self.kv_lens[:self.num_seqs].copy_( kv_lens + self.kv_cache_params.num_extra_kv_tokens) - self.kv_lens_cuda[:self.num_seqs].copy_(maybe_pin_memory( - kv_lens[:self.num_seqs]), - non_blocking=True) + step_inputs.copy_(self.kv_lens_cuda[:self.num_seqs], + maybe_pin_memory(kv_lens[:self.num_seqs])) # total kv lens for context requests and generation requests, without extra tokens self.host_total_kv_lens[0] = kv_lens[:self.num_contexts].sum().item() self.host_total_kv_lens[1] = kv_lens[self.num_contexts:self. @@ -1217,7 +1218,7 @@ def prepare_encoder_decoder_from_precomputed_lengths( f"({self.kv_cache_manager.max_seq_len}).") num_seqs = self.num_seqs - self.prompt_lens_cuda[:num_seqs].copy_(prompt_lens, non_blocking=True) + copy_h2d_if_changed(self.prompt_lens_cuda[:num_seqs], prompt_lens) self.kv_lens_cuda[:num_seqs].copy_(kv_lens, non_blocking=True) self.host_total_kv_lens[0] = context_kv_tokens self.host_total_kv_lens[1] = generation_kv_tokens diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 54792dab7866..87b01bef8350 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -9445,8 +9445,8 @@ def is_supported( @classmethod def _compile(cls, q_dtype: torch.dtype, num_heads: int, - max_decode_query_len: int, head_dim: int): - key = (q_dtype, num_heads, max_decode_query_len, head_dim) + max_decode_query_len: int, head_dim: int, split_k: int): + key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k) if key in cls.kernel_cache: return @@ -9493,7 +9493,7 @@ def _sym_stride(): cute_dtype, num_heads, max_decode_query_len, - cls.SPLIT_K, + split_k, head_dim, ) cls.kernel_cache[key] = cute.compile( @@ -9518,19 +9518,24 @@ def forward( seq_lens: torch.Tensor, score: torch.Tensor, max_decode_query_len: int, + split_k: int = SPLIT_K, ) -> None: """Score one decode step, compiling the kernel on first use. + ``split_k`` is the number of CTAs per request; each block is scored + by exactly one of them, so the scores do not depend on it. + The compile allocates, so it must not land inside a CUDA graph capture. It does not: CUDAGraphRunner.capture runs eager warmup forwards first, and those cover every geometry the graphs it then captures will replay. """ _, num_heads, head_dim = idx_q.shape - key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim) + key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, + split_k) if key not in cls.kernel_cache: cls._compile(idx_q.dtype, num_heads, max_decode_query_len, - head_dim) + head_dim, split_k) cls.kernel_cache[key](idx_q, index_k_cache, block_table, score, seq_lens) @@ -9544,6 +9549,7 @@ def cute_dsl_minimax_m3_index_decode_score( seq_lens: torch.Tensor, score: torch.Tensor, max_decode_query_len: int, + split_k: int = 256, ) -> None: """Write per-block max index scores for one decode step, in place. @@ -9562,6 +9568,8 @@ def cute_dsl_minimax_m3_index_decode_score( passed without a copy. max_decode_query_len: compile-time bound on decode_query_len; num_index_heads * max_decode_query_len must not exceed 32. + split_k: CTAs per request (compile-time); each walks blocks + split, split + split_k, ... The scores do not depend on it. """ if not is_sm_100f(): raise ValueError( @@ -9586,6 +9594,7 @@ def cute_dsl_minimax_m3_index_decode_score( seq_lens, score, max_decode_query_len, + split_k, ) @torch.library.register_fake( @@ -9597,6 +9606,7 @@ def _( seq_lens: torch.Tensor, score: torch.Tensor, max_decode_query_len: int, + split_k: int = 256, ) -> None: return None diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index c3d6b5b26e01..2fc414e93a34 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -833,6 +833,19 @@ def _choose_mxfp8_tactic(custom_op: str, runner: TunableRunner, return tactic +def mxfp8_quantize_autotuned( + input: torch.Tensor, + tune: bool = False) -> Tuple[torch.Tensor, torch.Tensor]: + """The quantize half of ``mxfp8_quantize_gemm_autotuned``: MXFP8 activation and scales, + on this token bucket's tuned quantize backend.""" + quantize_runner = MXFP8QuantizeRunner(input.dtype) + quantize_inputs = [input] + return quantize_runner(quantize_inputs, + tactic=_choose_mxfp8_tactic( + _MXFP8_QUANTIZE_AUTOTUNED_OP, quantize_runner, + quantize_inputs, tune)) + + def mxfp8_quantize_gemm_autotuned( input: torch.Tensor, weight: torch.Tensor, @@ -847,6 +860,24 @@ def mxfp8_quantize_gemm_autotuned( quantize_inputs, tactic=_choose_mxfp8_tactic(_MXFP8_QUANTIZE_AUTOTUNED_OP, quantize_runner, quantize_inputs, tune)) + return mxfp8_gemm_autotuned(act, act_scale, weight, weight_scale, + output_dtype, tune) + + +def mxfp8_gemm_autotuned( + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool = False, +) -> torch.Tensor: + """The GEMM half of mxfp8_quantize_gemm_autotuned for an activation already in MXFP8. + + Same runner, op key and tuning-cache entries as the decode-graph MXFP8 GEMM: + the tactic per token bucket comes from the AutoTuner (merged across ranks + when the runtime asks for it), never from a per-rank backend sweep. + """ gemm_runner = FlashInferMXFP8GemmRunner(output_dtype) gemm_inputs = [act, act_scale, weight, weight_scale] return gemm_runner(gemm_inputs, @@ -855,6 +886,43 @@ def mxfp8_quantize_gemm_autotuned( gemm_inputs, tune)) +def flashinfer_mxfp8_gemm_cute_dsl( + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + sweep_configs: bool, +) -> torch.Tensor: + """FlashInfer's CuTeDSL MXFP8 GEMM for an activation its producer already quantized. + + Unlike ``mxfp8_quantize_gemm_autotuned``, the backend is fixed rather than + chosen per token bucket: the AutoTuner profiles with the weight L2-resident, + where CUTLASS can edge out CuTeDSL, while a decode step streams it cold. + + ``sweep_configs`` runs FlashInfer's own CuTeDSL config sweep: it profiles + the configs for this shape while ``flashinfer_mxfp8_autotune`` is active + (the warmup tuning pass) and replays the cached best one afterwards; + without it the heuristic config runs. + """ + if sweep_configs: + from ..modules.linear import flashinfer_mxfp8_autotune_active + context = _flashinfer_autotune( + tune_mode=flashinfer_mxfp8_autotune_active()) + else: + context = _flashinfer_autotune(tune_mode=False, skip_ops="mxfp8_gemm") + with context: + return _flashinfer_mm_mxfp8( + act, + weight.t(), + act_scale, + weight_scale, + out_dtype=output_dtype, + use_8x4_sf_layout=False, + backend="cute-dsl", + ) + + class FP4GemmRunner(TunableRunner): runner_dict = dict() tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( diff --git a/tensorrt_llm/_torch/models/modeling_auto.py b/tensorrt_llm/_torch/models/modeling_auto.py index 0592317d293a..370fa5626590 100644 --- a/tensorrt_llm/_torch/models/modeling_auto.py +++ b/tensorrt_llm/_torch/models/modeling_auto.py @@ -1,7 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from dataclasses import replace from typing import Generic, Optional, Type +from tensorrt_llm.logger import logger + +from .._experimental.modeling_v2._router_index import (ModelingV2Mode, NoTarget, + resolved_mode) from ..model_config import ModelConfig from ..utils import model_extra_attrs from .modeling_utils import (DecoderModelForCausalLM, TConfig, TModel, @@ -74,7 +79,29 @@ def from_config( config.skip_create_weights_in_init = True config._frozen = True extra_attrs = config.extra_attrs + # What the loader put there, before a target that then refuses could + # add to it. + loader_attrs = dict(extra_attrs) with model_extra_attrs(extra_attrs): - model = cls(config) + try: + model = cls(config) + except NoTarget as why: + # A modeling_v2 target refuses, at construction, a runtime that + # routing does not see (set by engine arguments). Under auto + # that falls back to the built-in implementation, + # as a refusal in routing does; under require it is the error. + target = config.modeling_v2_target + if cls is not target or resolved_mode( + target) is ModelingV2Mode.REQUIRE: + raise + logger.info( + f"modeling_v2: {target.__name__} refused this runtime ({why}); " + "using the built-in implementation") + fallback = replace(config, modeling_v2_target=None) + # replace() starts init=False fields afresh; the built-in model + # reads the loader's settings (GEMM backends, allreduce sizes) + # from extra_attrs. + fallback.extra_attrs.update(loader_attrs) + return AutoModelForCausalLM.from_config(fallback) model.extra_attrs = extra_attrs return model diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index a4818fd02a71..923cd807e2fc 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -3479,6 +3479,16 @@ def _mxfp8_cutlass_op_available() -> bool: "flashinfer_mxfp8_decode_graph_capture_active", default=False) +def flashinfer_mxfp8_autotune_active() -> bool: + """Whether the current forward runs inside ``flashinfer_mxfp8_autotune``.""" + return _FLASHINFER_MXFP8_AUTOTUNE_ACTIVE.get() + + +def flashinfer_mxfp8_decode_graph_capture_active() -> bool: + """Whether the current forward is a decode CUDA-graph warmup or capture.""" + return _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE.get() + + @contextmanager def flashinfer_mxfp8_autotune() -> Iterator[None]: """Tune FlashInfer MXFP8 tactics while enabling auto-dispatched calls.""" diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 151c4d84ddcb..cc44a83f3e11 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -54,6 +54,7 @@ from ..speculative import (draft_prompt_lookahead, get_num_extra_kv_tokens, get_num_spec_layers, get_spec_decoder, should_use_separate_draft_kv_cache) +from ..speculative.spec_sampler_base import SpecSampler from . import config_utils from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer, extract_mamba_kv_cache_params, @@ -78,6 +79,7 @@ from .resource_manager import (KVCacheCompressionManager, KVCacheManager, PeftCacheManager, ResourceManager, ResourceManagerType) +from .runtime_hooks import runtime_hooks_of from .sampler import EarlyStopSampler, EarlyStopWithMMResult, TorchSampler from .sampler.penalties import PenaltyStore from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler, @@ -4536,7 +4538,11 @@ def instantiate_sampler( ) if engine.spec_config is not None and engine.spec_config.spec_dec_mode.has_spec_decoder( ): - return get_spec_decoder(sampler_args, engine.spec_config) + sampler = get_spec_decoder(sampler_args, engine.spec_config) + spec_step_fold = runtime_hooks_of(engine.model).spec_step_fold + if spec_step_fold is not None and isinstance(sampler, SpecSampler): + sampler.bind_spec_step_fold(spec_step_fold) + return sampler if mm_encoder_only: # NOTE: handle model outputs specially for mm encoder executor/engine. diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 992f5a20da61..79acf13fef59 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -30,6 +30,7 @@ from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..speculative.utils import get_draft_kv_cache_manager from ..utils import make_weak_ref, piecewise_cuda_graph +from . import step_inputs from .llm_request import LlmRequest, get_draft_token_length from .resource_manager import (BaseResourceManager, NoFreeSlotsError, ResourceManager, ResourceManagerType) @@ -216,6 +217,15 @@ class CUDAGraphRunnerConfig: dist: Optional[Distributed] dynamic_draft_len_mapping: Optional[Dict[int, int]] = None sparse_attention_config: Optional[BaseSparseAttentionConfig] = None + # Total sequence length at or below which a generation batch replays the + # short-sequence graph family; None captures one family. Derived from the + # sparse attention config when it asks for the split, else from the model's + # runtime hooks. + short_long_seq_len_threshold: Optional[int] = None + # Give the graph metadata a buffer pool of its own instead of the + # process-wide one, so no other engine writes its metadata buffers (an + # engine that skips unchanged uploads into them relies on that). + private_metadata_buffers: bool = False enable_encoder_decoder_mixed_cuda_graph: bool = False enable_in_graph_sampling: bool = False static_input_ids: Optional[torch.Tensor] = None @@ -236,6 +246,8 @@ class CUDAGraphRunner: multiple graphs, keyed by batch shape and execution-path specializations. """ WARMUP_STEPS = 1 + # CUDAGraphRunnerConfig.short_long_seq_len_threshold, when the config set one. + _configured_seq_len_threshold: Optional[int] = None def __init__(self, config: CUDAGraphRunnerConfig): self.config = config @@ -248,6 +260,7 @@ def __init__(self, config: CUDAGraphRunnerConfig): self.max_beam_width = config.max_beam_width self.spec_config = config.spec_config self.sparse_config = config.sparse_attention_config + self._configured_seq_len_threshold = config.short_long_seq_len_threshold self.is_encoder_decoder = config.is_encoder_decoder self.enable_encoder_decoder_mixed_cuda_graph = ( config.enable_encoder_decoder_mixed_cuda_graph) @@ -277,7 +290,9 @@ def __init__(self, config: CUDAGraphRunnerConfig): self.shared_static_tensors: Dict[str, torch.Tensor] = {} if self.enabled: self._create_shared_static_tensors() - self.cuda_graph_meta_buffers = get_memory_buffers() + self.cuda_graph_meta_buffers = (Buffers() + if config.private_metadata_buffers else + get_memory_buffers()) # On-the-fly capture is disabled by default to prevent workspace # tensor reallocation from invalidating addresses baked into existing @@ -394,6 +409,20 @@ def _gather_adp_graph_batch_info(self, batch: ScheduledRequests, tiers = gathered[:, 2].tolist() return list(zip(flags, sizes, tiers)) + @property + def short_long_seq_len_threshold(self) -> Optional[int]: + """The total sequence length splitting the short and long graph families, or None for one family. + + The runner config's value when it set one, else the sparse attention + config's when that asks for the split. + """ + if self._configured_seq_len_threshold is not None: + return self._configured_seq_len_threshold + if (isinstance(self.sparse_config, SeqLenAwareSparseAttentionConfig) + and self.sparse_config.needs_separate_short_long_cuda_graphs()): + return self.sparse_config.seq_len_threshold + return None + def _get_seq_len_mode( self, batch: ScheduledRequests, @@ -406,8 +435,7 @@ def _get_seq_len_mode( that the decoder runner temporarily placed in the generation list. It is empty for the existing generation-only path. """ - if (isinstance(self.sparse_config, SeqLenAwareSparseAttentionConfig) - and self.sparse_config.needs_separate_short_long_cuda_graphs()): + if self.short_long_seq_len_threshold is not None: # Some sparse attention algorithms need to use different forward paths for short and long sequences. # For example, the DSA can skip the MQA and Top-K in the indexer for short sequences to reduce the # computational overhead. To support this feature, we need to capture separate CUDA graphs for short @@ -444,7 +472,7 @@ def _get_seq_len_mode( from ..speculative import get_num_extra_kv_tokens num_extra_kv_tokens = get_num_extra_kv_tokens(self.spec_config) max_seq_len = max(total_seq_lens) - if max_seq_len <= self.sparse_config.seq_len_threshold - num_extra_kv_tokens: + if max_seq_len <= self.short_long_seq_len_threshold - num_extra_kv_tokens: short_seq_len_mode = True else: short_seq_len_mode = False @@ -920,7 +948,7 @@ def replay( # the source is the static view itself, and there is nothing to copy. static_input_ids = static_tensors["input_ids"][:seqlen] if input_ids.data_ptr() != static_input_ids.data_ptr(): - static_input_ids.copy_(input_ids) + step_inputs.copy_(static_input_ids, input_ids) position_ids = current_inputs["position_ids"] if self.config.use_mrope: @@ -969,7 +997,7 @@ def replay( "which would corrupt the static input buffer.") static_position_ids = static_tensors["position_ids"][:, :seqlen] if position_ids.data_ptr() != static_position_ids.data_ptr(): - static_position_ids.copy_(position_ids) + step_inputs.copy_(static_position_ids, position_ids) num_encoder_tokens = key.num_encoder_tokens if num_encoder_tokens: @@ -993,6 +1021,9 @@ def replay( encoder_hidden_states) static_encoder_hidden_states[actual_num_encoder_tokens:].zero_() + # Issue the step's packed input copies (no-op without an open pack) before + # anything below reads the static inputs on the device. + step_inputs.flush() if self._pre_replay_hooks: hook_tensors = static_tensors if position_id_offsets is not None: diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py index a85547af6ea8..d34b0d43c1c8 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py @@ -34,6 +34,7 @@ ResourceManager, ResourceManagerType, ) +from tensorrt_llm._torch.pyexecutor.runtime_hooks import runtime_hooks_of from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors from tensorrt_llm._torch.pyexecutor.sampler.sampler_common import SampleType from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests @@ -46,11 +47,7 @@ from tensorrt_llm._torch.speculative.utils import get_static_draft_len from tensorrt_llm._torch.utils import get_per_request_prefill_cuda_graph_flag from tensorrt_llm.inputs.multimodal import _has_mm_payload_keys -from tensorrt_llm.llmapi.llm_args import ( - BaseSparseAttentionConfig, - PrefillCudaGraphBackend, - SeqLenAwareSparseAttentionConfig, -) +from tensorrt_llm.llmapi.llm_args import BaseSparseAttentionConfig, PrefillCudaGraphBackend from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -720,6 +717,13 @@ def run_cuda_graph_warmup(self, resource_manager: ResourceManager) -> None: ) with flashinfer_autotune_context, flashinfer_mxfp8_decode_graph_capture(): self._capture_generation_cuda_graphs(resource_manager) + if ( + is_warmup_only + and runtime_hooks_of(self.model).merge_tuned_tactics_across_ranks + ): + # Tactics tuned while warming the decode graphs agree across + # ranks before the capture pass reads them. + AutoTuner.get().post_tune_merge_tactics() self._capture_additional_cuda_graphs(resource_manager) # Piecewise graphs have separate capture machinery and do not use the # whole-model attention workspace. Capture them only on the second pass. @@ -751,11 +755,8 @@ def _capture_generation_cuda_graphs(self, resource_manager: ResourceManager) -> if self.mapping.has_cp_helix(): effective_max_seq_len = self._config.max_seq_len // self.mapping.cp_size - sparse_config = self.sparse_attention_config - if ( - isinstance(sparse_config, SeqLenAwareSparseAttentionConfig) - and sparse_config.needs_separate_short_long_cuda_graphs() - ): + short_long_seq_len_threshold = self.cuda_graph_runner.short_long_seq_len_threshold + if short_long_seq_len_threshold is not None: # For short sequences, subtract the maximum runtime tokens consumed # by a generation step so all current-step tokens stay within the # sequence length threshold. PARD uses 2K tokens here, not K+1. @@ -763,7 +764,7 @@ def _capture_generation_cuda_graphs(self, resource_manager: ResourceManager) -> self._config.max_draft_len ) # For long sequences, use the default maximum sequence length. - max_seq_len = sparse_config.seq_len_threshold - max_runtime_tokens_per_gen_step + max_seq_len = short_long_seq_len_threshold - max_runtime_tokens_per_gen_step if max_seq_len < effective_max_seq_len: max_seq_len_list = [effective_max_seq_len, max_seq_len] else: diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index 7bb868bdf4ed..7a103241bace 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -17,6 +17,7 @@ from tensorrt_llm._torch.kernels.spec_step_copies import op as spec_step_copies from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.models.modeling_multimodal_mixin import _build_request_multimodal_input +from tensorrt_llm._torch.pyexecutor import step_inputs from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import get_mrope_dummy_seq_slot from tensorrt_llm._torch.pyexecutor.guided_decoder import CapturableGuidedDecoder from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 @@ -30,7 +31,12 @@ ResourceManager, ResourceManagerType, ) -from tensorrt_llm._torch.pyexecutor.runtime_hooks import DecodePrologue +from tensorrt_llm._torch.pyexecutor.runtime_hooks import ( + DEFAULT_RUNTIME_HOOKS, + DecodePrologue, + RuntimeHooks, + runtime_hooks_of, +) from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import RemoteTailPhase @@ -325,6 +331,8 @@ class DecoderInputPreparer: path overrides ``_new_extra_inputs`` and ``_prepare_inputs_fast_path``. """ + # Step-loop behaviour the model opted into (RuntimeHooks); all off unless it declares otherwise. + runtime_hooks: RuntimeHooks = DEFAULT_RUNTIME_HOOKS # The model's decode prologue (RuntimeHooks.decode_prologue), set by the runner. decode_prologue: DecodePrologue | None = None @@ -351,6 +359,7 @@ def __init__( self.mapping = mapping self.dist = dist self.model = model + self.runtime_hooks = runtime_hooks_of(model) self.get_runtime_tokens_per_gen_step = get_runtime_tokens_per_gen_step self.iter_states = iter_states self._lora = lora @@ -1894,6 +1903,16 @@ def previous_seq_slots_device() -> torch.Tensor: self._staged_previous_batch_indices = list(previous_batch_indices) return previous_slots + def previous_pos_indices_device(previous_batch_tokens: int) -> None: + if previous_pos_indices != self._staged_previous_pos_indices: + previous_pos_indices_host = torch.tensor( + previous_pos_indices, dtype=torch.int, pin_memory=prefer_pinned() + ) + self._buffers.previous_pos_indices_cuda[0:previous_batch_tokens].copy_( + previous_pos_indices_host, non_blocking=True + ) + self._staged_previous_pos_indices = list(previous_pos_indices) + if scheduled_requests.is_attention_dp_phase_idle: # Graph padding normally reuses an existing input buffer. An idle # rank has no real row to initialize that buffer on its first step. @@ -1915,16 +1934,60 @@ def previous_seq_slots_device() -> torch.Tensor: self._buffers.draft_tokens_cuda[: len(draft_tokens)].copy_( draft_tokens, non_blocking=True ) + step_fold = ( + self.runtime_hooks.spec_step_fold + if self._config.is_spec_decode and not self._dspark_trims_submitted_tokens + else None + ) + fold_rows = None + fold_staged = False + if step_fold is not None: + # A decode-only step's generation rows; the model's forward may stage the next + # step's inputs for exactly these rows (see SpecStepFold). + fold_rows = ( + [request.py_seq_slot for request in extend_requests] + if scheduled_requests.num_context_requests == 0 + else None + ) + if fold_rows is not None and any(slot is None for slot in fold_rows): + fold_rows = None + fold_staged = step_fold.begin_step(fold_rows, len(extend_requests)) if next_draft_tokens_device is not None: 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 num_extend_reqeust_wo_dummy = len(extend_requests) - len(extend_dummy_requests) + # The previous forward already staged this step's inputs when it ran for exactly + # these rows and this step adds nothing of its own. + inputs_staged = ( + fold_staged + and len(extend_dummy_requests) == 0 + and num_tokens == 0 + and num_draft_tokens == 0 + and previous_batch_len == len(extend_requests) + and previous_batch_indices == fold_rows + ) + if not inputs_staged and step_fold is not None and previous_batch_len > 0: + previous_slots = previous_seq_slots_device() + previous_pos_indices_device(previous_batch_len * runtime_tokens_per_gen_step) + inputs_staged = step_fold.gather( + new_tokens_device, + next_draft_tokens_device, + new_tokens_lens_device, + previous_slots, + self._buffers.previous_pos_indices_cuda, + previous_batch_len, + runtime_tokens_per_gen_step, + runtime_draft_token_buffer_width, + num_tokens, + num_draft_tokens, + num_extend_reqeust_wo_dummy - previous_batch_len, + ) # The offsets of requests without a previous batch and of dummy # requests must read 0. On a CUDA graph step (no token padding) # where every row is a request with a previous batch, the gathers # below overwrite all rows _preprocess_inputs reads, so the buffers # need no zeroing first. - if not ( + if not inputs_staged and not ( not self._dspark_trims_submitted_tokens and attn_metadata.is_cuda_graph and previous_batch_len == num_extend_reqeust_wo_dummy @@ -1935,21 +1998,14 @@ def previous_seq_slots_device() -> torch.Tensor: self._buffers.previous_pos_id_offsets_cuda.zero_() self._buffers.previous_kv_lens_offsets_cuda.zero_() - if previous_batch_len > 0: + if previous_batch_len > 0 and not inputs_staged: previous_slots = previous_seq_slots_device() is_ragged_gen = host.previous_batch_tokens_per_request is not None and any( tokens != runtime_tokens_per_gen_step for tokens in host.previous_batch_tokens_per_request ) previous_batch_tokens = len(previous_pos_indices) - if previous_pos_indices != self._staged_previous_pos_indices: - previous_pos_indices_host = torch.tensor( - previous_pos_indices, dtype=torch.int, pin_memory=prefer_pinned() - ) - self._buffers.previous_pos_indices_cuda[0:previous_batch_tokens].copy_( - previous_pos_indices_host, non_blocking=True - ) - self._staged_previous_pos_indices = list(previous_pos_indices) + previous_pos_indices_device(previous_batch_tokens) # The order of requests in a batch: [context requests, generation requests] # generation requests: ['requests that do not have previous batch', 'requests that @@ -2157,9 +2213,7 @@ def _copy_position_ids_to_device( generation_positions.view(1, 1, -1).expand(3, 1, -1), ) else: - self._buffers.position_ids_cuda[:total_num_tokens].copy_( - host_position_ids, non_blocking=True - ) + step_inputs.copy_(self._buffers.position_ids_cuda[:total_num_tokens], host_position_ids) host.host_position_ids = host_position_ids diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py index 266d51671be1..677009cfd412 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py @@ -18,6 +18,12 @@ KVCacheManager, ResourceManagerType, ) +from tensorrt_llm._torch.pyexecutor.runtime_hooks import ( + DEFAULT_RUNTIME_HOOKS, + MsaRuntimePolicy, + RuntimeHooks, + runtime_hooks_of, +) from tensorrt_llm._torch.speculative import SpecMetadata, get_spec_metadata from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -46,6 +52,9 @@ def resolve_mamba_metadata_cls(model: torch.nn.Module) -> type[Mamba2Metadata]: class DecoderMetadata: """Attention and speculative metadata reused across decoder iterations.""" + # Step-loop behaviour the model opted into (RuntimeHooks); all off unless it declares otherwise. + runtime_hooks: RuntimeHooks = DEFAULT_RUNTIME_HOOKS + def __init__( self, *, @@ -58,6 +67,7 @@ def __init__( self._config = config self.mapping = mapping self.model = model + self.runtime_hooks = runtime_hooks_of(model) self._buffers = buffers self._attention_runtime_features = attention_runtime_features # This field is initialized lazily on the first forward pass. @@ -104,9 +114,36 @@ def setup_attention( config, "chunk_size", self.attn_metadata.mamba_chunk_size ) self.attn_metadata.mamba_metadata_cls = resolve_mamba_metadata_cls(self.model) + self._arm_msa_policy(kv_cache_manager) return self.attn_metadata + def _arm_msa_policy(self, kv_cache_manager: KVCacheManager | KVCacheManagerV2 | None) -> None: + """Hand the model's MiniMax-M3 MSA preparation choices to the metadata and the KV cache.""" + msa = self.runtime_hooks.msa + if msa == MsaRuntimePolicy(): + return + # Imported here: only a model that armed an MSA policy needs the sparse backend. + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import ( + MiniMaxM3KVCacheManagerV2, + ) + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, + ) + + if not isinstance( + self.attn_metadata, MiniMaxM3MsaSparseAttentionMetadata + ) or not isinstance(kv_cache_manager, MiniMaxM3KVCacheManagerV2): + raise ValueError( + "RuntimeHooks.msa requires the MiniMax-M3 MSA sparse attention backend and KV cache, " + f"got {type(self.attn_metadata).__name__} and {type(kv_cache_manager).__name__}" + ) + self.attn_metadata.msa_policy = msa + kv_cache_manager.configure_decode_shortcuts( + skip_unchanged_block_offsets=msa.skip_unchanged_block_offsets, + reuse_pool_views=msa.reuse_pool_views, + ) + def setup_speculative( self, spec_resource_manager: BaseResourceManager | None ) -> SpecMetadata | None: @@ -115,6 +152,7 @@ def setup_speculative( self.spec_metadata = create_spec_metadata( self._config, self.model.config, spec_resource_manager ) + self.spec_metadata.skip_greedy_rng_upload = self.runtime_hooks.skip_greedy_rng_upload return self.spec_metadata @property diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 03242760a2a8..54aa2d9d8a60 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -26,6 +26,7 @@ from tensorrt_llm._torch.moe.workspace import CutlassWorkspaceReclaimer from tensorrt_llm._torch.peft.lora.config import LoraConfig from tensorrt_llm._torch.peft.lora.manager import LoraModelConfig +from tensorrt_llm._torch.pyexecutor import h2d_cache, step_inputs from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner import BreakableCUDAGraphRunner from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( CUDAGraphRunner, @@ -34,6 +35,7 @@ KeyType, ) from tensorrt_llm._torch.pyexecutor.guided_decoder import CapturableGuidedDecoder +from tensorrt_llm._torch.pyexecutor.h2d_cache import H2DSkipCache from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType from tensorrt_llm._torch.pyexecutor.runtime_hooks import ( @@ -42,8 +44,14 @@ runtime_hooks_of, ) from tensorrt_llm._torch.pyexecutor.sampler.beam_search import PendingBeamHistories +from tensorrt_llm._torch.pyexecutor.sampler.sampler import ( + GREEDY_HIDDEN_STATES_KEY, + GREEDY_TOKENS_KEY, + greedy_tokens_usable, +) from tensorrt_llm._torch.pyexecutor.sampler.sampler_common import SampleType from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer from tensorrt_llm._torch.pyexecutor.workspace import EagerWorkspaceReclaimer from tensorrt_llm._torch.speculative import SpecMetadata, get_draft_kv_cache_manager @@ -53,7 +61,12 @@ with_model_extra_attrs, ) from tensorrt_llm._utils import nvtx_range -from tensorrt_llm.llmapi.llm_args import BaseSparseAttentionConfig, PrefillCudaGraphBackend +from tensorrt_llm.llmapi.llm_args import ( + BaseSparseAttentionConfig, + MiniMaxM3SparseAttentionConfig, + PrefillCudaGraphBackend, + SeqLenAwareSparseAttentionConfig, +) from tensorrt_llm.mapping import Mapping from ...lora import LoraParamBuilder, make_cuda_graph_lora_manager @@ -134,6 +147,12 @@ class DecoderRunner(ScheduledModelRunner): # Step-loop behaviour the model opted into (RuntimeHooks); all off unless it declares otherwise. runtime_hooks: RuntimeHooks = DEFAULT_RUNTIME_HOOKS + # This engine's skip cache and step-input pack (None: the hook is off), active only during its own forward. + _h2d_cache: H2DSkipCache | None = None + _step_pack: StepInputPack | None = None + # Whether a batch may keep the greedy tokens a forward sampled in place of the logits. False when the executor + # runs a guided decoder, which reads the logits of every batch. + keep_greedy_only_outputs = True # Optional execution paths that a subclass can turn off. _context_warmups_supported = True @@ -180,6 +199,13 @@ def __init__( self._metrics = metrics self._config = config + self.runtime_hooks = runtime_hooks_of(model) + self._h2d_cache = H2DSkipCache() if self.runtime_hooks.skip_unchanged_h2d else None + self._step_pack = ( + StepInputPack(torch.device("cuda", torch.cuda.current_device())) + if self.runtime_hooks.pack_step_inputs + else None + ) self.forward_pass_callable = None # Optional in-graph sampling hook, registered by PyExecutor. Called at @@ -202,6 +228,14 @@ def __init__( self._moe_workspace_reclaimer: CutlassWorkspaceReclaimer | None = None self._buffers = DecoderInputBuffers.for_config(self._config, self._use_beam_search) + if self._config.is_spec_decode and self.runtime_hooks.spec_step_fold is not None: + self.runtime_hooks.spec_step_fold.bind_engine_buffers( + self._buffers.input_ids_cuda, + self._buffers.draft_tokens_cuda, + self._buffers.previous_pos_id_offsets_cuda, + self._buffers.previous_kv_lens_offsets_cuda, + self._config.max_batch_size, + ) self._metadata = DecoderMetadata( config=self._config, mapping=mapping, @@ -212,7 +246,6 @@ def __init__( self._dummy_requests = self._make_dummy_requests() self._init_decoder_state() self._inputs = self._make_input_preparer(mm_item_scheduler) - self.runtime_hooks = runtime_hooks_of(self._model) self._inputs.decode_prologue = self.runtime_hooks.decode_prologue self._graphs = self._make_cuda_graphs() self._inputs.cuda_graph_runner = self.cuda_graph_runner @@ -335,7 +368,15 @@ def forward( resource_manager: ResourceManager, is_dummy: bool = False, ) -> Any: - return self._forward_decoder(inputs, resource_manager, is_dummy=is_dummy) + with h2d_cache.active(self._h2d_cache), step_inputs.active(self._step_pack): + try: + return self._forward_decoder(inputs, resource_manager, is_dummy=is_dummy) + except BaseException: + # A failed step issues none of its packed copies, so forget the shadows recorded for them too. + step_inputs.discard() + if self._h2d_cache is not None: + self._h2d_cache.clear() + raise def _initialize_moe_workspace_reclaimer(self) -> None: """Create the native workspace owner for this runner's lifetime.""" @@ -433,6 +474,8 @@ def _cuda_graph_runner_config(self) -> CUDAGraphRunnerConfig: mapping=self._mapping, dist=self._dist, sparse_attention_config=self._sparse_attention_config, + short_long_seq_len_threshold=self._short_long_seq_len_threshold(), + private_metadata_buffers=self.runtime_hooks.skips_unchanged_metadata_uploads, enable_in_graph_sampling=self._config.enable_in_graph_sampling, # Speculative steps rewrite every input token and position, so # captured graphs can read the persistent buffers directly. @@ -448,6 +491,28 @@ def _cuda_graph_runner_config(self) -> CUDAGraphRunnerConfig: ), ) + def _short_long_seq_len_threshold(self) -> int | None: + """The total sequence length that splits the short and long decode graph families, or None for one family. + + Set by a sparse attention config that asks for the split, or by a model whose MSA + policy stages identity selections for short sequences: a sequence within the + top-k blocks attends every valid block. + """ + sparse_config = self._sparse_attention_config + if ( + isinstance(sparse_config, SeqLenAwareSparseAttentionConfig) + and sparse_config.needs_separate_short_long_cuda_graphs() + ): + return sparse_config.seq_len_threshold + if self.runtime_hooks.msa.short_sequence_identity: + if not isinstance(sparse_config, MiniMaxM3SparseAttentionConfig): + raise ValueError( + "MsaRuntimePolicy.short_sequence_identity requires a MiniMax-M3 sparse " + f"attention config, got {type(sparse_config).__name__}" + ) + return sparse_config.sparse_topk_blocks * sparse_config.sparse_block_size + return None + def _initialize_breakable_cuda_graph_runner(self) -> BreakableCUDAGraphRunner | None: if self._config.prefill_cuda_graph_backend != PrefillCudaGraphBackend.BREAKABLE: return None @@ -624,6 +689,16 @@ def _forward_decoder( execution_requests.generation_requests, ) + # Pack the step's small input copies into one upload and one scatter at the + # graph replay. Only a replay of an already captured graph, which always + # reaches the flush, may open a pack. + if ( + self.runtime_hooks.pack_step_inputs + and can_run_graph + and not self.cuda_graph_runner.is_warmup_only + and not self.cuda_graph_runner.needs_capture(key) + ): + step_inputs.begin() with self._cuda_graph_metadata_prepare_scope( attn_metadata, key, @@ -699,12 +774,16 @@ def postprocess_inputs(model_inputs: dict[str, Any]) -> None: else None ), ) + # No-op unless a pack was opened above and the replay did not issue it. + step_inputs.flush() if captures_model_body: return outputs if self.forward_pass_callable is not None: self.forward_pass_callable() + if isinstance(outputs, dict) and outputs.get(GREEDY_TOKENS_KEY) is not None: + outputs = self._logits_for_greedy_tokens(semantic_requests, outputs) self._execute_logit_post_processors(semantic_requests, outputs, pending_beam_histories) if not isinstance(outputs, dict): @@ -832,6 +911,39 @@ def _apply_logits_processors( # back would be a self-assignment, which torch rejects for the # non-contiguous slices a TP-padded vocab produces. + def _logits_for_greedy_tokens( + self, scheduled_requests: ScheduledRequests, outputs: dict + ) -> dict: + """Replace greedy tokens sampled by the forward with logits when a request needs logits. + + A forward that samples greedily in place of the LM head returns the head's input + alongside the tokens. When any generation request needs more than its greedy token, + the logits come from that input eagerly, followed by the in-graph sampling and the + logits post-processing the forward skipped. Returns the outputs to use: ``outputs`` + itself, or a new dict (a replayed graph's outputs dict is reused by later replays). + """ + if ( + self.keep_greedy_only_outputs + and self._config.max_beam_width == 1 + and scheduled_requests.num_context_requests == 0 + # An idle attention-DP rank runs a dummy row for an empty batch, whose outputs are restored from logits. + and len(scheduled_requests.generation_requests) > 0 + and all( + greedy_tokens_usable(request) for request in scheduled_requests.generation_requests + ) + ): + return outputs + outputs = dict(outputs) + outputs.pop(GREEDY_TOKENS_KEY) + hidden_states = outputs.pop(GREEDY_HIDDEN_STATES_KEY) + # One row per generation request: every row is a last token. + outputs["logits"] = self._model.logits_processor.forward( + hidden_states, self._model.lm_head, None, True + ) + if self.sample_in_graph_callable is not None and self._mapping.is_last_pp_rank(): + self.sample_in_graph_callable(outputs) + return outputs + def _start_beam_histories( self, scheduled_requests: ScheduledRequests ) -> tuple[list[LlmRequest], PendingBeamHistories] | None: @@ -877,6 +989,11 @@ def _execute_logit_post_processors( return logits_tensor = outputs["logits"] + if logits_tensor is None: + # Greedy tokens sampled by the forward; _logits_for_greedy_tokens supplied + # logits instead whenever a request needed them. + return + if outputs.get("logits_vocab_shard", False): # The spec worker kept the logits vocabulary-sharded (this TP # rank's columns). A post-processor sees the whole vocabulary, so diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py index 4b5f99651907..90daafe169eb 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py @@ -23,6 +23,7 @@ from tensorrt_llm._torch.nccl_window_tensor_scope import discard_nccl_window_tensor_outputs from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType +from tensorrt_llm._torch.pyexecutor.runtime_hooks import runtime_hooks_of from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import warmup_sampling_module from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.pyexecutor.trace_log_utils import log_mem_snapshot @@ -1164,6 +1165,13 @@ def run_autotuner_pass(autotune_context: Any, synchronize_trtllm_cache: bool) -> ran_flashinfer_forward = run_autotuner_pass( flashinfer_mxfp8_autotune(), synchronize_trtllm_cache=False ) + # The model's declared hooks (the all-off default when it declares none): the warmup passes also run on + # engines built without the runner's constructor. + if runtime_hooks_of(self.model).merge_tuned_tactics_across_ranks: + # Rank-local profiling noise must not leave tensor-parallel peers + # serving different tactics for one shape: keep the fastest everywhere. + # Collective; before the native MXFP8 cache below registers the tactics. + AutoTuner.get().post_tune_merge_tactics() if enable_flashinfer_mxfp8_autotuner: if ran_flashinfer_forward: diff --git a/tensorrt_llm/_torch/pyexecutor/h2d_cache.py b/tensorrt_llm/_torch/pyexecutor/h2d_cache.py new file mode 100644 index 000000000000..2a7a4d715447 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/h2d_cache.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Skip host-to-device metadata copies whose values the destination already holds. + +Steady decode re-uploads many small per-step metadata tensors whose values do not change (slot indices, query +lengths, prompt lengths, page tables between page boundaries). Each upload is a separate copy on the compute stream, +serialized between the decode graphs. An ``H2DSkipCache`` keeps one host shadow per destination buffer (the last +values written through it, and the view they were written to) and skips the copy when the same view receives the same +values again. + +Each engine that armed ``RuntimeHooks.skip_unchanged_h2d`` owns one cache and makes it the active one for its forward +(``active``); ``copy_h2d_if_changed`` acts on the calling context's active cache and, with none, is +``dst.copy_(src, non_blocking=True)``. The attention metadata that calls it has no engine to ask, which is why the +cache travels in the context rather than as an argument. + +Route a copy through it only for a destination that no other code writes (no device kernel, no other copy path, no other +engine): the shadow describes the device contents only under that invariant. Pooled metadata buffers satisfy it because +an engine that arms the skip gives its CUDA graph metadata a buffer pool of its own +(``RuntimeHooks.skips_unchanged_metadata_uploads``); eager metadata allocates its own buffers. A write through a +different view of the same buffer replaces the shadow, so overlapping views never skip on stale state. +""" + +from __future__ import annotations + +import contextlib +from contextvars import ContextVar +from typing import Iterator, Optional + +import torch + +from tensorrt_llm._device_copy import DeviceCopyBuffer + +from . import step_inputs + + +class H2DSkipCache: + """One engine's host shadows of its skip-routed destinations (see the module docstring).""" + + def __init__(self) -> None: + # Device storage address -> (the storage, kept alive so the address is not reused; the view key; the host + # values). Keyed by storage rather than by view base: tensors made under inference mode do not record their + # base. + self._shadow: dict[int, tuple[object, tuple, DeviceCopyBuffer]] = {} + # Device storage address -> (the storage, kept alive as above; the view key; (first row, value) of the view's + # last tail fill). + self._tails: dict[int, tuple[object, tuple, tuple[int, int]]] = {} + + def clear(self) -> None: + """Forget every shadow, so the next copy to each destination is issued (e.g. after a discarded step).""" + self._shadow.clear() + self._tails.clear() + + def copy(self, dst: torch.Tensor, src: torch.Tensor, pack: bool) -> bool: + storage = dst.untyped_storage() + addr = storage.data_ptr() + if src.is_cuda or not dst.is_cuda or torch.cuda.is_current_stream_capturing(): + self._shadow.pop(addr, None) + dst.copy_(src, non_blocking=True) + return True + key = (dst.storage_offset(), tuple(dst.shape), tuple(dst.stride()), dst.dtype) + entry = self._shadow.get(addr) + if entry is not None and entry[1] == key and entry[2].matches(src, exact=True): + return False + buffer = entry[2] if entry is not None and entry[1] == key else DeviceCopyBuffer(dst) + buffer.invalidate() + self._shadow.pop(addr, None) + if pack: + step_inputs.copy_(dst, src) + else: + dst.copy_(src, non_blocking=True) + buffer.record(src) + self._shadow[addr] = (storage, key, buffer) + return True + + def fill_tail(self, buf: torch.Tensor, start: int, value: int) -> None: + storage = buf.untyped_storage() + addr = storage.data_ptr() + if start >= buf.shape[0]: + # The live rows cover the whole view: the previous tail fill is overwritten. + self._tails.pop(addr, None) + return + if torch.cuda.is_current_stream_capturing(): + self._tails.pop(addr, None) + buf[start:].fill_(value) + return + key = (buf.storage_offset(), tuple(buf.shape), tuple(buf.stride()), buf.dtype) + entry = self._tails.get(addr) + if entry is not None and entry[1] == key and entry[2] == (start, value): + return + buf[start:].fill_(value) + self._tails[addr] = (storage, key, (start, value)) + + +_active: ContextVar[Optional[H2DSkipCache]] = ContextVar("trtllm_h2d_skip_cache", default=None) + + +@contextlib.contextmanager +def active(cache: Optional[H2DSkipCache]) -> Iterator[None]: + """Make ``cache`` the calling context's cache (None: no skipping) for the duration.""" + token = _active.set(cache) + try: + yield + finally: + _active.reset(token) + + +def enabled() -> bool: + """Whether the calling context skips unchanged copies.""" + return _active.get() is not None + + +def copy_h2d_if_changed(dst: torch.Tensor, src: torch.Tensor, pack: bool = False) -> bool: + """``dst.copy_(src, non_blocking=True)`` unless the active cache knows ``dst`` holds ``src``; True if copied. + + ``pack``: the copy may join an open ``step_inputs`` pack, i.e. land only at the pack's flush (the graph replay). + Pass it only for destinations no device work reads before that replay. + """ + cache = _active.get() + if cache is None: + dst.copy_(src, non_blocking=True) + return True + return cache.copy(dst, src, pack) + + +def fill_tail_if_changed(buf: torch.Tensor, start: int, value: int) -> None: + """``buf[start:].fill_(value)`` unless the active cache's last tail fill of ``buf`` was the same. + + Call it after every write of the live rows ``buf[:start]``, including one that covers the whole of ``buf`` + (``start == len(buf)``, which fills nothing): a live-row write past the recorded start overwrites part of the + tail. The skip is sound only while nothing else writes ``buf`` past ``start``. + """ + cache = _active.get() + if cache is None: + if start < buf.shape[0]: + buf[start:].fill_(value) + return + cache.fill_tail(buf, start, value) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 2e70a897803e..8723bd0fda7f 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -824,15 +824,13 @@ def cuda_graph_padding_enabled(self) -> bool: @property def _dspark_confidence_enabled(self) -> bool: - return isinstance( - self._runner, - DecoderRunner) and self._runner._dspark_confidence_enabled + return isinstance(self._runner, DecoderRunner) and self._runner._dspark_confidence_enabled @property def _dspark_trims_submitted_tokens(self) -> bool: - return isinstance( - self._runner, - DecoderRunner) and self._runner._dspark_trims_submitted_tokens + return ( + isinstance(self._runner, DecoderRunner) and self._runner._dspark_trims_submitted_tokens + ) @property def _dspark_sps_cost_table(self): @@ -849,7 +847,7 @@ def _dspark_exact_candidate_cells(self): @property def _dspark_exact_identity_words(self): if not isinstance(self._runner, DecoderRunner): - return (0, ) * 8 + return (0,) * 8 return self._runner._dspark_exact_identity_words @property @@ -924,6 +922,11 @@ def init_cuda_graph_lora_manager(self, lora_config: LoraConfig) -> None: assert self._lora_model is not None self._runner.init_cuda_graph_lora_manager(lora_config, self._lora_model) + def drop_greedy_only_outputs(self) -> None: + """Have every batch return logits, for an executor-side guided decoder that reads them.""" + if isinstance(self._runner, DecoderRunner): + self._runner.keep_greedy_only_outputs = False + def set_guided_decoder(self, guided_decoder: CapturableGuidedDecoder) -> bool: if hasattr(self.model, "set_guided_decoder"): success = self.model.set_guided_decoder(guided_decoder) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index d319f73e439d..52a65e4a99df 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -9170,9 +9170,11 @@ def _sample_async(self, scheduled_batch, beam_width = self.sampler.beam_width( scheduled_batch.all_requests()) + # A forward that sampled greedy tokens in place of the logits + # returns no logits; no request of such a batch asks for them. HandleLogits()(scheduled_batch.context_requests, scheduled_batch.generation_requests, - batch_outputs["logits"], beam_width, + batch_outputs.get("logits"), beam_width, num_context_logits_prefix_sum, self.sampler.is_generation_model()) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index f0e465944c5a..3fe94785b4a9 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -989,6 +989,7 @@ def allocation_scope(current_stage: ExecutorMemoryType): ): # GuidedDecoder is applicable to non-speculative decoding and drafter-based speculative decoding. guided_decoder = GuidedDecoder(**kwargs) + model_engine.drop_greedy_only_outputs() elif spec_config.spec_dec_mode.support_capturable_guided_decoder( ): # CapturableGuidedDecoder is applicable to one-model speculative decoding. diff --git a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py index 772845919949..3409d72cb772 100644 --- a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py +++ b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py @@ -1,22 +1,149 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Step-loop behaviour a model may opt into, read through ``runtime_hooks_of``. +"""Step-loop behaviour a model may opt into, read once by the engine. -A model declares its hooks as ``runtime_hooks`` (an instance attribute, a class -attribute or a property). Every field defaults to off, and a model that declares -nothing gets ``RuntimeHooks()``: the step loop then runs exactly as it does -without this module. Every engine call site -of a hook runs its original statements when the hook is off. +A model declares its hooks as ``runtime_hooks`` (an instance attribute or a class +attribute). Every field defaults to off, and a model that declares nothing gets +``RuntimeHooks()``: the step loop, the samplers and the attention metadata then +run exactly as they do without this module. The hooks change *where* and *how* the step loop moves data, never what a -forward computes. The protocols below state only the methods an engine call -site needs; the objects that implement them are owned by the model. +forward computes. Each one is armed only by a model that has measured it on its +own deployment; none is a general default. """ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Optional, Protocol +from typing import TYPE_CHECKING, Any, Optional, Protocol + +if TYPE_CHECKING: + import torch + + +class SpecStepFold(Protocol): + """Speculative step-input staging folded into the previous step's forward. + + A one-model speculative step ends by storing its accepted tokens, draft + tokens and lengths into the sampler's per-slot store, and the next step + begins by gathering those values back into the engine's flat input buffers. + A model that implements this protocol does both inside its last draft pass: + it writes the sampler store and stages the next step's inputs directly, and + the engine and the sampler skip their copies on the steps where it did. + + The object is owned by the model; the engine and the sampler only call it. + Every method is host-side bookkeeping or a single kernel launch on the + current stream. + """ + + def bind_engine_buffers( + self, + input_ids: "torch.Tensor", + draft_tokens: "torch.Tensor", + position_id_offsets: "torch.Tensor", + kv_lens_offsets: "torch.Tensor", + max_batch_size: int, + ) -> None: + """Receive the engine's persistent speculative input buffers, once, at construction.""" + + def bind_sampler_store(self, store: object) -> None: + """Receive the speculative sampler's persistent per-slot store, once, at construction.""" + + def begin_step(self, seq_slots: Optional[list[int]], num_rows: int) -> bool: + """Start a step of ``num_rows`` generation rows. + + ``seq_slots`` are those rows' sequence slots on a decode-only step whose + rows all have one, else None: the forward then stores nothing for them. + Returns whether the previous step's forward staged its inputs for + ``seq_slots``; the engine additionally checks that this step adds no + inputs of its own before it skips its gathers. + """ + + def gather( + self, + new_tokens: "torch.Tensor", + next_draft_tokens: "torch.Tensor", + new_tokens_lens: "torch.Tensor", + previous_slots: "torch.Tensor", + previous_pos_indices: "torch.Tensor", + previous_batch_len: int, + tokens_per_gen_step: int, + draft_token_buffer_width: int, + num_tokens: int, + num_draft_tokens: int, + first_previous_row: int, + ) -> bool: + """Do the engine's previous-batch gathers in one launch. + + Writes the engine buffers exactly as the engine's own gathers would, + including zeroing both offset buffers outside the gathered range. + Returns False, having written nothing, when it cannot take the step. + """ + + def sampler_rows_stored(self, outputs: dict, num_skip: int, seq_slots: list[int]) -> bool: + """Whether the forward already stored this step's outputs for ``seq_slots`` into the sampler store.""" + + def scatter( + self, + seq_slots: "torch.Tensor", + new_tokens: "torch.Tensor", + next_new_tokens: "torch.Tensor", + new_tokens_lens: "torch.Tensor", + next_draft_tokens: "torch.Tensor", + ) -> bool: + """Store the step's outputs into the sampler store in one launch. + + Returns False, having written nothing, when it cannot take the step. + """ + + +class MsaKvUpdate(Protocol): + """Advance the MSA decode kernels' per-request state after the KV lengths moved. + + Called by ``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` with the + metadata's own buffers; recomputes the attended lengths, the new tokens' + cache slots and their valid-block counts in one launch, exactly as the + metadata's own update does. Returns False, having written nothing, when it + cannot take the step. + """ + + def __call__( + self, + kv_lens: "torch.Tensor", + kv_lens_staged: "torch.Tensor", + seq_lens: "torch.Tensor", + q_batch_row: "torch.Tensor", + q_intra: "torch.Tensor", + block_table: "torch.Tensor", + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: "torch.Tensor", + cu_kv_lens: Optional["torch.Tensor"], + out_cache_loc: "torch.Tensor", + n_valid_blocks: Optional["torch.Tensor"], + ) -> bool: ... + + +@dataclass(frozen=True) +class MsaRuntimePolicy: + """MiniMax-M3 MSA metadata preparation a model may opt into. + + Every option changes how a step's selection and cache-write tables reach + the device, not which blocks are attended. + """ + + # A step whose every query token has at most top-k valid blocks attends all + # of them, so its selection is the identity table: prepare stages that table + # instead of a scoring plan, and decode CUDA graphs are keyed by short and + # long contexts. + short_sequence_identity: bool = False + # The model's one-launch replacement for the update after the KV lengths move. + kv_update: Optional[MsaKvUpdate] = None + # Skip a block-table copy whose inputs equal the previous step's. + skip_unchanged_block_offsets: bool = False + # Build each KV pool view once per pool geometry instead of on every attention call. + reuse_pool_views: bool = False class DecodePrologue(Protocol): @@ -51,13 +178,35 @@ def at_shutdown(self) -> None: class RuntimeHooks: """The step-loop hooks a model opts into. All off by default.""" + # Pack a decode step's small host-to-device input copies into one upload and + # one scatter launch, issued just before the CUDA graph replay. + pack_step_inputs: bool = False + decode_prologue: Optional[DecodePrologue] = None + executor_events: Optional[ExecutorEvents] = None + # Skip host-to-device metadata copies whose values the destination already holds. + skip_unchanged_h2d: bool = False # Run the overlap executor loop with the execution stream current, so the # loop's own stream work and the forward are ordered on one stream. loop_on_execution_stream: bool = False - # The decode step's input prologue in one launch (see DecodePrologue). - decode_prologue: Optional[DecodePrologue] = None - # Actions after warmup, before the overlap loop and at shutdown (see ExecutorEvents). - executor_events: Optional[ExecutorEvents] = None + # After autotuning, keep the fastest profiled tactic per shape on every + # tensor-parallel rank, so the ranks capture the same kernels. + merge_tuned_tactics_across_ranks: bool = False + # Skip the per-step Philox seed and offset uploads on all-greedy + # one-model speculative batches (the argmax path never reads them). + skip_greedy_rng_upload: bool = False + # Speculative step-input staging folded into the forward (see SpecStepFold). + spec_step_fold: Optional[SpecStepFold] = None + # MiniMax-M3 MSA metadata preparation. + msa: MsaRuntimePolicy = MsaRuntimePolicy() + + @property + def skips_unchanged_metadata_uploads(self) -> bool: + """Whether a hook skips uploads into graph metadata buffers that already hold the values. + + Such a skip is sound only while no other engine writes those buffers, so the engine + then gives its CUDA graph metadata a buffer pool of its own. + """ + return self.skip_unchanged_h2d or self.msa.skip_unchanged_block_offsets DEFAULT_RUNTIME_HOOKS = RuntimeHooks() diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index b1f256e560c7..f66c80cef78d 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -79,7 +79,7 @@ meet_stop_token_criteria, scatter_new_tokens, ) -from .sampling_resolution import SamplingMode, SlotSamplingParams +from .sampling_resolution import SamplingMode, SlotSamplingParams, request_row_params from .seed_manager import _SeedManager from .token_ban import OverlappedTokenBanHandler, SynchronousTokenBanHandler, TokenBanHandler from .top_p_decay import TopPDecayHandler @@ -119,6 +119,34 @@ class SampleState(Generic[GenericSampleStateTensorsHost, GenericSampleStateTenso # Generic bounds not supported, https://github.com/python/typing/issues/548 GenericSampleState = TypeVar("GenericSampleState", bound=SampleState) # type: ignore +# Model output keys of a forward that samples greedily in place of its LM head: per-row greedy token ids, and the +# head's input rows so the logits can still be computed for a batch that needs them. Such a forward returns no +# logits. +GREEDY_TOKENS_KEY = "greedy_tokens" +GREEDY_HIDDEN_STATES_KEY = "greedy_hidden_states" + + +def greedy_tokens_usable(request: LlmRequest) -> bool: + """Whether a generation request is served by the forward's greedy token alone, without logits. + + The request conditions of TorchSampler's stable greedy path, plus every other reader of the logits. A dummy + request is not usable: the stable greedy path takes no batch that carries one. + """ + return not request.is_dummy and ( + get_draft_token_length(request) == 0 + and request.py_embedding_bias is None + and not getattr(request, "py_bad_words", None) + and not getattr(request, "py_no_repeat_ngram_size", None) + and not has_occurrence_penalty(request) + and not request.py_min_length + and not request.py_return_log_probs + and not request.py_return_generation_logits + and not request.py_logits_post_processors + and request.py_guided_decoding_params is None + and not check_stop_words_length(request) + and request_row_params(request).mode is SamplingMode.GREEDY + ) + @dataclass(frozen=True, kw_only=True) class _FusedRowParamsBuffers: @@ -2345,7 +2373,15 @@ def _process_requests( torch.Tensor, bool, ]: - raw_logits_cuda = model_outputs["logits"] + # A forward that sampled greedy tokens in place of the logits (GREEDY_TOKENS_KEY) may return no logits. + raw_logits_cuda: Optional[torch.Tensor] = model_outputs.get("logits") + greedy_tokens: Optional[torch.Tensor] = model_outputs.get(GREEDY_TOKENS_KEY) + if greedy_tokens is not None: + rows_cuda = greedy_tokens + elif raw_logits_cuda is not None: + rows_cuda = raw_logits_cuda + else: + raise RuntimeError("the forward returned neither logits nor greedy tokens") generation_requests = scheduled_requests.generation_requests request_ids = [request.py_request_id for request in generation_requests] @@ -2358,7 +2394,7 @@ def _process_requests( bool(generation_requests) and self.max_beam_width == 1 and scheduled_requests.num_context_requests == 0 - and len(generation_requests) <= raw_logits_cuda.shape[0] + and len(generation_requests) <= rows_cuda.shape[0] and all( not request.is_dummy and get_draft_token_length(request) == 0 for request in generation_requests @@ -2395,6 +2431,9 @@ def _process_requests( live_requests = [request for request in generation_requests if not request.is_dummy] staged_matches_batch = ( self._in_graph_staged + # A forward that sampled greedily in place of its LM head skipped the + # in-graph sampling, so the staged buffers were never written. + and greedy_tokens is None and self._in_graph_dest_indices is not None and [request.py_request_id for request in live_requests] == self._in_graph_request_ids @@ -2484,12 +2523,20 @@ def _process_requests( self._stable_greedy_seq_slots_host = seq_slots_host self._stable_greedy_seq_slots_cuda = seq_slots_cuda - next_tokens = fast_greedy_sample_kernel( - raw_logits_cuda[: len(generation_requests)], - new_tokens_cuda, - seq_slots_cuda, - self.max_beam_width, - ) + if greedy_tokens is not None: + next_tokens = greedy_tokens[: len(generation_requests)].to(new_tokens_cuda.dtype) + scatter_new_tokens( + next_tokens, new_tokens_cuda, seq_slots_cuda, self.max_beam_width + ) + else: + # Without greedy tokens, rows_cuda is the logits. + assert raw_logits_cuda is not None + next_tokens = fast_greedy_sample_kernel( + raw_logits_cuda[: len(generation_requests)], + new_tokens_cuda, + seq_slots_cuda, + self.max_beam_width, + ) new_tokens_host = self._copy_to_host(next_tokens) return ( generation_requests, @@ -2504,6 +2551,10 @@ def _process_requests( self._stable_greedy_request_ids = [] self._stable_greedy_seq_slots = [] + if raw_logits_cuda is None: + raise RuntimeError( + "the forward returned greedy tokens without logits for a batch that needs logits" + ) sampling_requests, sampling_requests_metadata, logits_cuda = self._select_generated_logits( scheduled_requests, raw_logits_cuda, diff --git a/tensorrt_llm/_torch/pyexecutor/step_inputs.py b/tensorrt_llm/_torch/pyexecutor/step_inputs.py new file mode 100644 index 000000000000..c538896c04da --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/step_inputs.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pack a decode step's small per-step input copies into one upload and one scatter launch. + +Between two decode-graph replays, the next step's metadata reaches the device as a chain of tiny copies (positions, +KV lengths, cache slots, and the graph runner's static-input copies). Each is a separate operation on the stream the +next graph waits for, so the GPU idles for the whole chain. + +While a pack is open (``begin`` ... ``flush``), ``copy_(dst, src)`` records a host-to-device copy by snapshotting the +host bytes into a pinned staging buffer, and a device-to-device copy as a (dst, src) descriptor. ``flush`` issues one +host-to-device copy of [descriptors | data] into a device staging buffer and one ``trtllm::step_inputs_scatter`` +launch, which writes every staged host copy and then every device copy (in that order, so a device copy may read a +destination a host copy just filled). With no pack open, ``copy_`` is ``dst.copy_(src, non_blocking=True)``. + +Each engine that armed ``RuntimeHooks.pack_step_inputs`` owns one ``StepInputPack`` (on its device) and makes it the +active pack for its forward (``active``); the module functions act on the active pack of the calling context, so +executors on other threads or devices never share one. The engine opens it only when the step will reach the flush (a +replay of an already captured graph). Route a copy through ``copy_`` only when no device work reads or writes its +destination between the copy site and the flush: the pack moves the copy to the flush point. A copy that cannot be +packed (non-contiguous, device-side dtype change) flushes the pack first and then runs as a plain copy, which keeps +the original order. A pack holds every tensor it staged until it is issued, so a temporary source outlives its copy. +""" + +from __future__ import annotations + +import contextlib +from contextvars import ContextVar +from typing import Iterator, Optional + +import torch + +from tensorrt_llm._utils import prefer_pinned + + +@torch.library.register_fake("trtllm::step_inputs_scatter") +def _step_inputs_scatter_fake(stage: torch.Tensor) -> None: + # Writes the staged rows into their destinations; returns nothing. + return None + + +class StepInputPack: + """One engine's pack of deferred step-input copies (see the module docstring).""" + + CAPACITY = 1 << 16 + MAX_DESC = 256 + # Pinned staging buffers in rotation: a forced flush reopens the pack without waiting on the upload just issued. + RING = 4 + HEADER_BYTES = 16 + DESC_BYTES = 24 + ALIGN = 16 + + def __init__(self, device: torch.device) -> None: + self.host = [ + torch.empty(self.CAPACITY, dtype=torch.uint8, pin_memory=prefer_pinned()) + for _ in range(self.RING) + ] + self.events: list[Optional[torch.cuda.Event]] = [None] * self.RING + self.dev = torch.empty(self.CAPACITY, dtype=torch.uint8, device=device) + self.data_start = self.HEADER_BYTES + self.DESC_BYTES * self.MAX_DESC + self.parity = 0 + self.open = False + self.h2d: list[tuple[int, int, int]] = [] + self.d2d: list[tuple[int, int, int]] = [] + # Every tensor a staged copy reads or writes, kept alive until the pack is issued. + self.staged: list[torch.Tensor] = [] + self.used = self.data_start + + def begin(self) -> None: + if self.open: + self.flush() + event = self.events[self.parity] + if event is not None: + # The upload that last read this pinned buffer has run. + event.synchronize() + self._clear() + self.open = True + + def _clear(self) -> None: + self.h2d.clear() + self.d2d.clear() + self.staged.clear() + self.used = self.data_start + + def stage(self, dst: torch.Tensor, src: torch.Tensor) -> bool: + if not dst.is_cuda or dst.numel() != src.numel(): + return False + nbytes = dst.numel() * dst.element_size() + if src.device.type == "cpu": + # A contiguous destination is one descriptor; a 2-D slice with contiguous rows (a page table's leading + # rows and columns) is one per row. + if dst.is_contiguous(): + rows, row_bytes, row_stride = 1, nbytes, 0 + elif dst.dim() == 2 and dst.stride(1) == 1: + rows = dst.shape[0] + row_bytes = dst.shape[1] * dst.element_size() + row_stride = dst.stride(0) * dst.element_size() + else: + return False + offset = (self.used + self.ALIGN - 1) // self.ALIGN * self.ALIGN + if ( + offset + nbytes > self.CAPACITY + or len(self.h2d) + len(self.d2d) + rows > self.MAX_DESC + ): + return False + # A dtype change converts on the host, as copy_ would on the device. + src = src.to(dst.dtype).contiguous() + self.host[self.parity][offset : offset + nbytes].copy_(src.view(-1).view(torch.uint8)) + base = dst.data_ptr() + for row in range(rows): + self.h2d.append((base + row * row_stride, offset + row * row_bytes, row_bytes)) + self.staged.append(dst) + self.used = offset + nbytes + return True + if len(self.h2d) + len(self.d2d) >= self.MAX_DESC: + return False + if src.is_cuda and src.dtype == dst.dtype and src.is_contiguous() and dst.is_contiguous(): + self.d2d.append((dst.data_ptr(), src.data_ptr(), nbytes)) + self.staged += (dst, src) + return True + return False + + def flush(self) -> None: + if not self.open: + return + self.open = False + if not self.h2d and not self.d2d: + return + host = self.host[self.parity] + desc = [len(self.h2d), len(self.d2d)] + [ + value for copy in self.h2d + self.d2d for value in copy + ] + host[: 8 * len(desc)].view(torch.int64).copy_(torch.tensor(desc, dtype=torch.int64)) + # Descriptors and data in one upload; the gap between them is unused padding. + self.dev[: self.used].copy_(host[: self.used], non_blocking=True) + torch.ops.trtllm.step_inputs_scatter(self.dev) + event = self.events[self.parity] + if event is None: + event = self.events[self.parity] = torch.cuda.Event() + event.record() + self.parity = (self.parity + 1) % self.RING + # The scatter is queued on the current stream: the caching allocator keeps the staged tensors' memory for it. + self._clear() + + def discard(self) -> None: + """Close the pack without issuing it (its step failed).""" + self.open = False + self._clear() + + +_active: ContextVar[Optional[StepInputPack]] = ContextVar("trtllm_step_input_pack", default=None) + + +@contextlib.contextmanager +def active(pack: Optional[StepInputPack]) -> Iterator[None]: + """Make ``pack`` the calling context's pack (None: no packing) for the duration.""" + token = _active.set(pack) + try: + yield + finally: + _active.reset(token) + + +def begin() -> None: + """Open the active pack for the coming copies. Call only when the step will reach ``flush``.""" + pack = _active.get() + if pack is not None: + pack.begin() + + +def is_open() -> bool: + pack = _active.get() + return pack is not None and pack.open + + +def copy_(dst: torch.Tensor, src: torch.Tensor) -> torch.Tensor: + """``dst.copy_(src, non_blocking=True)``, deferred into the active open pack when there is one.""" + pack = _active.get() + if pack is not None and pack.open and not torch.cuda.is_current_stream_capturing(): + if pack.stage(dst, src): + return dst + # Keep the order: everything staged so far lands before this copy, and the copies after it pack again. + pack.flush() + out = dst.copy_(src, non_blocking=True) + pack.begin() + return out + return dst.copy_(src, non_blocking=True) + + +def discard() -> None: + """Close the active pack without issuing it (its step failed).""" + pack = _active.get() + if pack is not None: + pack.discard() + + +def flush() -> None: + """Issue the active open pack (one upload and one scatter launch on the current stream) and close it.""" + pack = _active.get() + if pack is not None: + pack.flush() diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index bac560ef8f36..4143466e8320 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -17,6 +17,7 @@ from ..attention.backends.flashinfer import FlashInferAttentionMetadata from ..attention.backends.sparse.params import MTPIndexShareMetadata from ..model_config import ModelConfig +from ..pyexecutor.h2d_cache import copy_h2d_if_changed from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import LlmRequest from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager @@ -346,8 +347,7 @@ def prepare(self): dtype=torch.int, device='cpu', pin_memory=prefer_pinned()) - self.batch_indices_cuda[:num_seqs].copy_(batch_indices, - non_blocking=True) + copy_h2d_if_changed(self.batch_indices_cuda[:num_seqs], batch_indices) # 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 @@ -368,8 +368,7 @@ def prepare(self): eagle_slot_ids_tensor = torch.tensor(eagle_slot_ids, dtype=torch.int, pin_memory=prefer_pinned()) - self.slot_ids[:num_seqs].copy_(eagle_slot_ids_tensor, - non_blocking=True) + copy_h2d_if_changed(self.slot_ids[:num_seqs], eagle_slot_ids_tensor) sa_manager = getattr(self.spec_resource_manager, 'sa_manager', None) if sa_manager is not None: diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 97b67a555523..3846c3df36ee 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -571,6 +571,10 @@ def is_ragged_verify(self) -> bool: group_all_greedy_sample: Optional[bool] = None # Whether to use rejection sampling for one-model speculative decoding. use_rejection_sampling: bool = False + # Skip the Philox seed and offset uploads on all-greedy batches (deploy-time; + # from the model's RuntimeHooks). The argmax path never reads them, and the + # next sampled batch rewrites every row it reads. + skip_greedy_rng_upload: bool = False # Advanced-sampling specialization (deploy-time; from DecodingBaseConfig.advanced_sampling_mode). advanced_sampling_mode: AdvancedSamplingMode = AdvancedSamplingMode.FULL # Whether the occurrence penalties are enabled (deploy-time; from diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index ccb3a06073e5..8e71b8fe089c 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -33,6 +33,7 @@ from ..kernels.spec_step_copies import op as spec_step_copies from ..pyexecutor.llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..pyexecutor.resource_manager import BaseResourceManager +from ..pyexecutor.runtime_hooks import SpecStepFold from ..pyexecutor.sampler import ( DEFAULT_BEAM_IDX, AsyncWorkerMixin, @@ -168,6 +169,9 @@ class SpecSampler(Sampler[SampleStateSpec], AsyncWorkerMixin): draft length every mode used to compute for itself. """ + # The model's speculative step-input staging, when it opted into one (bind_spec_step_fold). + spec_step_fold: Optional[SpecStepFold] = None + SampleState = SampleStateSpec def is_generation_model(self) -> bool: @@ -327,6 +331,11 @@ def __init__( new_tokens_lens=int_tensor((seq_slots,)), ) + def bind_spec_step_fold(self, spec_step_fold: SpecStepFold) -> None: + """Let the model's forward store into this sampler's per-slot store (see SpecStepFold).""" + self.spec_step_fold = spec_step_fold + spec_step_fold.bind_sampler_store(self.store) + def _request_common_handling( self, request: LlmRequest, @@ -425,6 +434,12 @@ def _scatter_to_stores( ] o_next_new_tokens = outputs["next_new_tokens"][num_skip : num_skip + num_sampling_requests] + step_fold = self.spec_step_fold + if step_fold is not None and step_fold.scatter( + slots_device, o_new_tokens, o_next_new_tokens, o_new_tokens_lens, o_next_draft_tokens + ): + return + # Pad or truncate to match fixed-size store buffers for index_copy_. # The worker output width tracks runtime_draft_len, which dynamic draft # length shrinks below the statically allocated store width. @@ -543,8 +558,12 @@ def sample_async( # The previous step's host copies may still be reading the stores. torch.cuda.current_stream().wait_event(self._store_copies_done) slots = [r.py_seq_slot for r in sampling_requests] - scattered = False - if self._store_scatter is not None: + step_fold = self.spec_step_fold + scattered = step_fold is not None and step_fold.sampler_rows_stored( + outputs, num_skip, slots + ) + # A model's step fold scatters the stores itself, in _scatter_to_stores. + if not scattered and step_fold is None and self._store_scatter is not None: # A step whose batch is unchanged leaves the slot table as it is. self._slot_table_buffer.update(torch.tensor(slots, dtype=torch.int32)) scattered = self._store_scatter.scatter( diff --git a/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py b/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py new file mode 100644 index 000000000000..a5c777b5cf7c --- /dev/null +++ b/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Whole-model gate for the MiniMax-M3 modeling_v2 targets. + +Every test here passes ``modeling_v2="require"`` to ``LLM``: under ``"auto"`` a configuration that missed a target's +criteria would quietly fall back to the built-in implementation, pass, and report the built-in's numbers as the +target's. + +The deployment is the one the targets route (``models/minimax_m3/routing.py``): the NVFP4 checkpoint on sm_103, +tensor-parallel attention and experts on 4 or 8 ranks, the MiniMax-M3 MSA sparse attention with the fused five-way +QKV/index projection, an FP8 indexer and an NVFP4 KV cache. Each target runs GSM8K without speculation and with +one-model Eagle3, where a chat-GSM8K acceptance probe also guards the drafter (accuracy alone does not notice a +corrupted draft KV). +""" + +import json + +import pytest + +from tensorrt_llm import LLM +from tensorrt_llm._torch._experimental.modeling_v2 import MODELING_V2_VALIDATE_ENV +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.llmapi import ( + CudaGraphConfig, + Eagle3DecodingConfig, + KvCacheConfig, + MiniMaxM3SparseAttentionConfig, + MoeConfig, + SamplingParams, +) +from tensorrt_llm.quantization import QuantAlgo + +from ...conftest import llm_models_root +from ..accuracy_core import GSM8K, LlmapiAccuracyTestHarness + +# The targets assert their own SM at construction: certification is per GPU architecture. +skip_not_sm103 = pytest.mark.skipif( + get_sm_version() != 103, + reason="the MiniMax-M3 modeling_v2 targets are certified on sm_103 only", +) + +# The chat-GSM8K acceptance floors of the in-tree test on an NVFP4 KV cache (test_llm_api_pytorch.TestMiniMaxM3). +MIN_ACCEPTANCE_RATE = 0.78 +MIN_ACCEPTANCE_LENGTH = 3.3 + + +def _drain_spec_stats(llm): + drafted = accepted = steps = 0 + for s in llm.get_stats(timeout=2): + s = json.loads(s) if isinstance(s, str) else s + sd = s.get("specDecodingStats") or {} + drafted += sd.get("numDraftTokens", 0) + accepted += sd.get("numAcceptedTokens", 0) + steps += sd.get("numRequestsWithDraftTokens", 0) + return drafted, accepted, steps + + +class _MiniMaxM3Target(LlmapiAccuracyTestHarness): + MODEL_NAME = "nvidia/MiniMax-M3-NVFP4" + MODEL_PATH = f"{llm_models_root()}/MiniMax-M3-NVFP4" + EAGLE_PATH = f"{llm_models_root()}/MiniMax-M3-EAGLE3-GQA" + TP = 4 + + def _llm(self, eagle3: bool) -> LLM: + spec = ( + Eagle3DecodingConfig(max_draft_len=3, speculative_model=self.EAGLE_PATH) + if eagle3 + else None + ) + return LLM( + self.MODEL_PATH, + tensor_parallel_size=self.TP, + moe_expert_parallel_size=1, + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.6, enable_block_reuse=False, dtype="nvfp4" + ), + sparse_attention_config=MiniMaxM3SparseAttentionConfig( + implementation="msa", indexer_kv_dtype="fp8", fuse_qkv_index_projection=True + ), + moe_config=MoeConfig(backend="TRTLLM"), + max_seq_len=4096, + max_batch_size=64, + cuda_graph_config=CudaGraphConfig(enable_padding=True, max_batch_size=64), + speculative_config=spec, + enable_iter_perf_stats=eagle3, + max_stats_len=5000, + trust_remote_code=True, + modeling_v2="require", + env_overrides={MODELING_V2_VALIDATE_ENV: "1"}, + ) + + def _gsm8k(self, eagle3: bool) -> None: + with self._llm(eagle3) as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4 + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + if not eagle3: + return + from datasets import load_dataset + + questions = [ + r["question"] for r in load_dataset(GSM8K.DATASET_DIR, "main", split="test") + ][:200] + prompts = [ + llm.tokenizer.apply_chat_template( + [{"role": "user", "content": q}], tokenize=False, add_generation_prompt=True + ) + for q in questions + ] + _drain_spec_stats(llm) + llm.generate(prompts, SamplingParams(max_tokens=512, temperature=0)) + drafted, accepted, steps = _drain_spec_stats(llm) + assert steps > 0, "no speculative iterations recorded" + rate, length = accepted / drafted, 1 + accepted / steps + print( + f"MiniMax-M3 tp{self.TP} Eagle3 chat-GSM8K acceptance: rate={rate:.3f}, length={length:.3f}" + ) + assert rate > MIN_ACCEPTANCE_RATE, ( + f"acceptance rate {rate:.3f} <= {MIN_ACCEPTANCE_RATE}" + ) + assert length > MIN_ACCEPTANCE_LENGTH, ( + f"acceptance length {length:.3f} <= {MIN_ACCEPTANCE_LENGTH}" + ) + + +class TestModelingV2MiniMaxM3Nvfp4Sm103Tp4(_MiniMaxM3Target): + """MiniMax-M3 NVFP4 / sm_103 / tp4: four GB300s.""" + + TP = 4 + + @skip_not_sm103 + @pytest.mark.skip_less_device(4) + def test_gsm8k(self): + self._gsm8k(eagle3=False) + + @skip_not_sm103 + @pytest.mark.skip_less_device(4) + def test_eagle3_gsm8k_and_acceptance(self): + self._gsm8k(eagle3=True) + + +class TestModelingV2MiniMaxM3Nvfp4Sm103Tp8(_MiniMaxM3Target): + """MiniMax-M3 NVFP4 / sm_103 / tp8: eight GB300s over two trays (an 8-rank MPI world across the two nodes).""" + + TP = 8 + + @skip_not_sm103 + @pytest.mark.skip_less_mpi_world_size(8) + def test_gsm8k(self): + self._gsm8k(eagle3=False) + + @skip_not_sm103 + @pytest.mark.skip_less_mpi_world_size(8) + def test_eagle3_gsm8k_and_acceptance(self): + self._gsm8k(eagle3=True) diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index ffe4b3e8255b..982fe8086061 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -141,7 +141,7 @@ l0_b300: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] # TIMEOUT measured: 3122 cases in 13m54s on one GB300 (the cpu_only files run on the CPU # stage instead), so 30 is 2.2x headroom. - - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm TIMEOUT (30) + - unittest/_torch/modeling_v2 --import-mode=importlib --ignore=unittest/_torch/modeling_v2/comm --ignore-glob=unittest/_torch/modeling_v2/**/comm/* TIMEOUT (30) - accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[tp1] # qwen4_exp target, pre-merge: gsm8k; one MTP engine for gsm8k with the acceptance-length # floor (the only gate that sees a miscomputed draft layer), mmmu (the vision tower) and diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index a5655df8c510..02313f3afd24 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -12,6 +12,8 @@ l0_cpu: backend: generic orchestrator: mpi tests: + # MiniMax-M3 target unit tests. + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py - unittest/_torch/modules/test_linear_output_slice.py::test_fp8_weight_slice_cache - unittest/_torch/modules/test_linear_output_slice.py::test_prepare_output_slice_rejects_unsupported_projection_before_forward - unittest/_torch/modules/test_linear_output_slice.py::test_prepare_output_slice_dispatches_to_linear_method diff --git a/tests/integration/test_lists/test-db/l0_gb300.yml b/tests/integration/test_lists/test-db/l0_gb300.yml index ce21a01c9ddc..17a72724ca1c 100644 --- a/tests/integration/test_lists/test-db/l0_gb300.yml +++ b/tests/integration/test_lists/test-db/l0_gb300.yml @@ -18,6 +18,29 @@ l0_gb300: stage: pre_merge backend: pytorch tests: + # MiniMax-M3 target unit tests. + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/activation + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/embedding + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/norm + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_capture_boundary.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_dispatch.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_rank_agreement.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/activation + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/embedding + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/norm + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_capture_boundary.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_dispatch.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py # ------------- PyTorch tests --------------- # MLA context head slices, including FP8 scale-block boundaries. - unittest/_torch/modules/test_linear_output_slice.py diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index 58d14ea21ba6..e12c1c358c82 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -15,6 +15,9 @@ l0_gb300_multi_gpus: stage: post_merge backend: pytorch tests: + # MiniMax-M3 target unit tests. + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm # One ordered sequence in one job: the cases share module-global rank state. - unittest/_torch/modeling_v2/comm - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k_identity_vs_mtp3[dep4] @@ -28,6 +31,8 @@ l0_gb300_multi_gpus: - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k[dep4_moetp4] - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k_chunked_prefill[tp4_ep4] - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k_chunked_prefill[dep4] + - accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_gsm8k + - accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_eagle3_gsm8k_and_acceptance # Covers the GPU cases of tests/unittest/_torch/attention/ on GB300, kernels/ and # multi_gpu/ included; its cpu_only cases run on CPU (l0_cpu). The GPU cases of # kernels/parallel_hw_agnostic/ run here too, although l0_b200 and l0_b300 ignore that diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml index 6f14f710ea35..522774da2c26 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml @@ -17,6 +17,9 @@ l0_gb300_multi_nodes_node2_gpu8: # Kimi K3 GSM8K accuracy plus an acceptance-length gate on the TEP8 DSpark # recipe. Measured 25m45s on 2x4 GB300. - accuracy/test_kimi3.py::TestKimiK3DSpark::test_gsm8k_tep8 TIMEOUT (60) + # The MiniMax-M3 modeling_v2 tp8 target (two trays): GSM8K, and with Eagle3 GSM8K plus a chat acceptance gate. + - accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp8::test_gsm8k TIMEOUT (60) + - accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp8::test_eagle3_gsm8k_and_acceptance TIMEOUT (60) - accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[tp8] - accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[dep8] - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k[tp8_moetp2_ep4] diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py index c9945d39d48c..61026b6a5f65 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py @@ -175,11 +175,13 @@ def test_index_decode_score_matches_reference(dtype, num_heads, decode_query_len @skip_not_sm100 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) -def test_index_decode_score_multi_round_split_k(dtype): +@pytest.mark.parametrize("split_k", [None, 256, 512], ids=["default", "split256", "split512"]) +def test_index_decode_score_multi_round_split_k(dtype, split_k): """A request longer than split_k pages makes the per-CTA block loop iterate. - split_k is 256, so 33000 tokens is 258 blocks and CTAs 0 and 1 each handle - two blocks while the rest handle one. + The default split_k is 256, so 33000 tokens is 258 blocks and CTAs 0 and 1 + each handle two blocks while the rest handle one. With 512 every block has + its own CTA and the trailing CTAs have none. """ seq_lens = [33000, 32768 + 1] idx_q, k_cache, block_table, seq_lens_dev, _, score, expected = _make_inputs( @@ -187,12 +189,39 @@ def test_index_decode_score_multi_round_split_k(dtype): ) assert (max(seq_lens) + PAGE_SIZE - 1) // PAGE_SIZE > 256 + split_k_args = () if split_k is None else (split_k,) torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( - idx_q, k_cache, block_table, seq_lens_dev, score, 1 + idx_q, k_cache, block_table, seq_lens_dev, score, 1, *split_k_args ) _assert_scores_close(score, expected) +@skip_not_sm100 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize( + ("num_heads", "decode_query_len", "seq_lens"), + [(4, 1, [300, 1500, 2049, 33000]), (4, 4, [4097, 70000])], + ids=["decode", "spec_verify"], +) +def test_index_decode_score_does_not_depend_on_split_k( + dtype, num_heads, decode_query_len, seq_lens +): + """Each block is scored by exactly one CTA, so the scores are identical across split_k.""" + idx_q, k_cache, block_table, seq_lens_dev, _, score, expected = _make_inputs( + seq_lens, dtype=dtype, num_heads=num_heads, decode_query_len=decode_query_len, seed=5 + ) + scores = {} + for split_k in (256, 512): + score.fill_(-float("inf")) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, k_cache, block_table, seq_lens_dev, score, decode_query_len, split_k + ) + scores[split_k] = score.clone() + + assert torch.equal(scores[256], scores[512]) + _assert_scores_close(scores[512], expected) + + @skip_not_sm100 def test_index_decode_score_transposed_view_matches_contiguous(): """The transposed selector view must produce the same values as a direct write. diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index 6f611923aa2c..31912c481f51 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -42,6 +42,7 @@ ) from tensorrt_llm._torch.attention.backends.sparse.registry import _resolve_minimax_m3_backend_cls from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy from tensorrt_llm.bindings import DataType from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig @@ -831,6 +832,7 @@ class FakeMetadata: # indexer takes the proxy path and needs no decode span. msa_decode_span = None msa_prefill_proxy_plan = (False, 0, 2, {}, None) + _msa_prefill_identity = None msa_prefill_n_valid_blocks = torch.ones(2, dtype=torch.int32, device="cuda") msa_prefill_proxy_kv_indices = torch.arange(2, dtype=torch.int32, device="cuda") msa_qo_lens_cpu = torch.ones(2, dtype=torch.int32) @@ -948,6 +950,8 @@ def select_blocks(self, *args: object, **kwargs: object) -> torch.Tensor: return torch.zeros(num_tokens, 1, 16, dtype=torch.int32) class FakeMetadata: + _msa_prefill_identity = None + msa_prefill_n_valid_blocks = torch.ones(num_tokens, dtype=torch.int32) msa_n_valid_blocks = torch.ones(num_tokens, dtype=torch.int32) msa_worst_case_max_k_tiles = 8 @@ -1978,6 +1982,122 @@ def test_on_update_kv_lens_is_a_noop_without_speculative_decoding(monkeypatch): assert metadata.msa_n_valid_blocks.tolist() == [0] * 5 +class _RecordingKvUpdate: + """A model's one-launch KV update (MsaKvUpdate) that takes or declines every step.""" + + def __init__(self, taken: bool) -> None: + self._taken = taken + self.calls = [] + + def __call__( + self, + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + tokens_per_block, + batch, + total_q, + msa_seq_lens, + cu_kv_lens, + out_cache_loc, + n_valid_blocks, + ) -> bool: + self.calls.append( + dict( + kv_lens=kv_lens, + kv_lens_staged=kv_lens_staged, + seq_lens=seq_lens, + q_batch_row=q_batch_row, + q_intra=q_intra, + block_table=block_table, + tokens_per_block=tokens_per_block, + batch=batch, + total_q=total_q, + msa_seq_lens=msa_seq_lens, + cu_kv_lens=cu_kv_lens, + out_cache_loc=out_cache_loc, + n_valid_blocks=n_valid_blocks, + ) + ) + return self._taken + + +def test_default_msa_policy_has_no_kv_update(monkeypatch): + metadata = _kv_lens_update_metadata( + monkeypatch, qo_lens=(1,), kv_staged=(5,), kv_corrected=(5,), page_size=4 + ) + + assert metadata.msa_policy == MsaRuntimePolicy() + assert metadata.msa_policy.kv_update is None + + +@pytest.mark.parametrize("prefill_counts", [False, True]) +def test_a_taken_kv_update_replaces_the_rederivation(monkeypatch, prefill_counts): + metadata = _kv_lens_update_metadata( + monkeypatch, + qo_lens=(2, 3), + kv_staged=(9, 12), + kv_corrected=(8, 12), + page_size=4, + prefill_counts=prefill_counts, + ) + kv_update = _RecordingKvUpdate(taken=True) + metadata.msa_policy = MsaRuntimePolicy(kv_update=kv_update) + + metadata.on_update_kv_lens() + + # The hook owns the update: the metadata's own re-derivation does not run. + assert metadata.msa_seq_lens_cuda.tolist() == [9, 12] + assert metadata.msa_out_cache_loc.tolist() == [-1] * 8 + assert metadata.msa_n_valid_blocks.tolist() == [0] * 8 + assert len(kv_update.calls) == 1 + call = kv_update.calls[0] + assert call["kv_lens"] is metadata.kv_lens_cuda + assert call["kv_lens_staged"] is metadata.msa_kv_lens_staged + assert call["q_batch_row"] is metadata.msa_q_batch_row + assert call["q_intra"] is metadata.msa_q_intra + assert call["block_table"] is metadata.msa_block_table + assert call["msa_seq_lens"] is metadata.msa_seq_lens_cuda + assert call["out_cache_loc"] is metadata.msa_out_cache_loc + assert call["seq_lens"].tolist() == [2, 3] + assert (call["tokens_per_block"], call["batch"], call["total_q"]) == (4, 2, 5) + assert call["cu_kv_lens"] is None + counts = metadata._msa_prefill_n_valid_blocks if prefill_counts else metadata.msa_n_valid_blocks + assert call["n_valid_blocks"].data_ptr() == counts.data_ptr() + assert call["n_valid_blocks"].numel() == 5 + + +def test_a_declined_kv_update_falls_back_to_the_rederivation(monkeypatch): + metadata = _kv_lens_update_metadata( + monkeypatch, qo_lens=(2, 3), kv_staged=(9, 12), kv_corrected=(8, 12), page_size=4 + ) + kv_update = _RecordingKvUpdate(taken=False) + metadata.msa_policy = MsaRuntimePolicy(kv_update=kv_update) + + metadata.on_update_kv_lens() + + assert len(kv_update.calls) == 1 + assert metadata.msa_seq_lens_cuda.tolist() == [8, 12] + assert metadata.msa_out_cache_loc[:5].tolist() == [6, 7, 109, 110, 111] + assert metadata.msa_n_valid_blocks[:5].tolist() == [2, 2, 3, 3, 3] + + +def test_kv_update_is_not_called_without_dynamic_kv_lens(monkeypatch): + metadata = _kv_lens_update_metadata( + monkeypatch, qo_lens=(1, 1), kv_staged=(4, 6), kv_corrected=(3, 5), page_size=4 + ) + metadata._msa_kv_lens_dynamic = False + kv_update = _RecordingKvUpdate(taken=True) + metadata.msa_policy = MsaRuntimePolicy(kv_update=kv_update) + + metadata.on_update_kv_lens() + + assert kv_update.calls == [] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize( ("src_dtype", "cache_dtype"), diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_runtime_policy.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_runtime_policy.py new file mode 100644 index 000000000000..0e022aaa28f2 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_runtime_policy.py @@ -0,0 +1,399 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 MSA decode-step preparation a model opts into (RuntimeHooks.msa). + +Covers the short-sequence identity table the metadata stages, and the KV cache +manager's block-table and pool-view shortcuts. With the default policy every +path runs as before; the existing test_msa_backend.py tests cover that. +""" + +from unittest.mock import create_autospec + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3KVCacheManagerV2 +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.common import ( + MiniMaxM3SparseMetadataParams, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + MSA_REQUIRED_TOPK, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, + _fill_unused_slots, +) +from tensorrt_llm._torch.pyexecutor import h2d_cache +from tensorrt_llm._torch.pyexecutor.h2d_cache import H2DSkipCache +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy +from tensorrt_llm.bindings import DataType +from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig +from tensorrt_llm.mapping import Mapping + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + +NUM_KV_HEADS = 2 +PARAMS = MiniMaxM3SparseMetadataParams( + global_num_q_heads=8, global_num_kv_heads=NUM_KV_HEADS, num_index_heads=4, topk=16 +) + + +@pytest.fixture +def manager(monkeypatch) -> MiniMaxM3KVCacheManagerV2: + """An M3 cache manager whose V2 base allocates nothing (as in test_msa_backend.py).""" + + def fake_base_init(self, *args, **kwargs) -> None: + del args, kwargs + self.is_disagg = False + self.dtype = DataType.BF16 + self.layer_offsets = {3: 0} + # _prepare_page_table_tensor's result for the two-pool host table the copy tests build. + self.num_attention_op_pools = 2 + + def no_index_k_buffer(self, layer_idx, **kwargs): + del self, layer_idx, kwargs + return None + + monkeypatch.setattr(KVCacheManagerV2, "__init__", fake_base_init) + monkeypatch.setattr(KVCacheManagerV2, "get_index_k_buffer", no_index_k_buffer) + monkeypatch.setattr(MiniMaxM3KVCacheManagerV2, "_compute_num_total_slots", lambda self: 0) + # The constructor resolves every layer's K/V view, which needs a real pool. Stub it there + # only: the pool-view tests below call get_buffers against their own geometry. + with monkeypatch.context() as construction: + construction.setattr(MiniMaxM3KVCacheManagerV2, "get_buffers", lambda self, *args: None) + return MiniMaxM3KVCacheManagerV2( + num_layers=4, sparse_attention_config=MiniMaxM3SparseAttentionConfig() + ) + + +# --------------------------------------------------------------------------- +# Short-sequence identity table +# --------------------------------------------------------------------------- + + +def _short_table_metadata( + policy: MsaRuntimePolicy, manager: MiniMaxM3KVCacheManagerV2, num_contexts: int = 0 +) -> MiniMaxM3MsaSparseAttentionMetadata: + metadata = object.__new__(MiniMaxM3MsaSparseAttentionMetadata) + metadata.msa_policy = policy + metadata._num_contexts = num_contexts + metadata.kv_cache_manager = manager + metadata.max_num_sequences = 4 + metadata.max_num_tokens = 16 + metadata.cuda_graph_buffers = None + metadata.is_cuda_graph = False + metadata.mapping = Mapping() + return metadata + + +def _identity(n_valid: list[int]) -> torch.Tensor: + ar = torch.arange(MSA_REQUIRED_TOPK, dtype=torch.int32) + valid = torch.tensor(n_valid, dtype=torch.int32) + table = torch.where(ar[None, :] < valid[:, None], ar[None, :], -1) + return table.unsqueeze(0).expand(NUM_KV_HEADS, len(n_valid), MSA_REQUIRED_TOPK) + + +def _staged(metadata: MiniMaxM3MsaSparseAttentionMetadata, total_q: int) -> torch.Tensor: + torch.cuda.synchronize() + size = NUM_KV_HEADS * total_q * MSA_REQUIRED_TOPK + return metadata.msa_short_table[:size].view(NUM_KV_HEADS, total_q, MSA_REQUIRED_TOPK).cpu() + + +@requires_cuda +@pytest.mark.parametrize("n_valid", [[1], [1, 16, 5], [16, 16, 16, 16]]) +def test_short_steps_stage_the_identity_table(manager, n_valid): + metadata = _short_table_metadata(MsaRuntimePolicy(short_sequence_identity=True), manager) + + metadata._stage_msa_short_table(torch.tensor(n_valid, dtype=torch.int32), PARAMS) + + assert metadata._msa_short_ok + assert metadata.msa_short_table.is_cuda + assert torch.equal(_staged(metadata, len(n_valid)), _identity(n_valid)) + + +@requires_cuda +def test_the_default_policy_stages_nothing(manager): + metadata = _short_table_metadata(MsaRuntimePolicy(), manager) + + metadata._stage_msa_short_table(torch.tensor([1, 2], dtype=torch.int32), PARAMS) + + assert not metadata._msa_short_ok + assert metadata.msa_short_table is None + + +@requires_cuda +@pytest.mark.parametrize( + ("n_valid", "num_contexts"), + [([1, MSA_REQUIRED_TOPK + 1], 0), ([], 0), ([1, 2], 1)], + ids=["more_than_topk_blocks", "empty_step", "context_step"], +) +def test_ineligible_steps_do_not_use_the_identity_table(manager, n_valid, num_contexts): + metadata = _short_table_metadata(MsaRuntimePolicy(short_sequence_identity=True), manager) + # A previous eligible step staged the table. + metadata._stage_msa_short_table(torch.tensor([3], dtype=torch.int32), PARAMS) + assert metadata._msa_short_ok + + metadata._num_contexts = num_contexts + metadata._stage_msa_short_table(torch.tensor(n_valid, dtype=torch.int32), PARAMS) + + assert not metadata._msa_short_ok + + +@requires_cuda +@pytest.mark.parametrize("skip_cache", [False, True], ids=["plain", "skip_cache"]) +def test_the_identity_table_follows_each_step(manager, skip_cache): + """Also with the engine's skip cache active, which skips the repeated [2, 4] upload.""" + metadata = _short_table_metadata(MsaRuntimePolicy(short_sequence_identity=True), manager) + + with h2d_cache.active(H2DSkipCache() if skip_cache else None): + for n_valid in ([2, 3], [2, 4], [2, 4], [16], [2, 3]): + metadata._stage_msa_short_table(torch.tensor(n_valid, dtype=torch.int32), PARAMS) + assert torch.equal(_staged(metadata, len(n_valid)), _identity(n_valid)) + + +@requires_cuda +def test_the_identity_table_grows_past_its_first_bound(manager): + metadata = _short_table_metadata(MsaRuntimePolicy(short_sequence_identity=True), manager) + metadata._stage_msa_short_table(torch.tensor([1], dtype=torch.int32), PARAMS) + first_numel = metadata.msa_short_table.numel() + n_valid = [i % MSA_REQUIRED_TOPK + 1 for i in range(first_numel // MSA_REQUIRED_TOPK)] + + metadata._stage_msa_short_table(torch.tensor(n_valid, dtype=torch.int32), PARAMS) + + assert metadata._msa_short_ok + assert metadata.msa_short_table.numel() >= NUM_KV_HEADS * len(n_valid) * MSA_REQUIRED_TOPK + assert torch.equal(_staged(metadata, len(n_valid)), _identity(n_valid)) + + +# --------------------------------------------------------------------------- +# Block-table copy shortcut +# --------------------------------------------------------------------------- + + +@pytest.fixture +def base_copy(monkeypatch): + """The V2 base copy, recorded instead of launched.""" + copy = create_autospec(KVCacheManagerV2.copy_batch_block_offsets, spec_set=True) + monkeypatch.setattr(KVCacheManagerV2, "copy_batch_block_offsets", copy) + return copy.mock + + +def _block_table_host(manager: MiniMaxM3KVCacheManagerV2) -> torch.Tensor: + manager.host_kv_cache_block_offsets = torch.arange(2 * 4 * 2 * 8, dtype=torch.int32).view( + 2, 4, 2, 8 + ) + return manager.host_kv_cache_block_offsets + + +def _decode_copy(manager, dst, request_ids=(0, 1), num_contexts=0, max_blocks=8): + manager.copy_batch_block_offsets( + dst, list(request_ids), 1, num_contexts, len(request_ids), max_blocks=max_blocks + ) + + +def test_block_offsets_shortcut_is_off_by_default(manager, base_copy): + _block_table_host(manager) + dst = torch.zeros(2, 4, 2, 8, dtype=torch.int32) + + _decode_copy(manager, dst) + _decode_copy(manager, dst) + + assert not manager.skip_unchanged_block_offsets + assert base_copy.call_count == 2 + + +def test_an_unchanged_decode_copy_is_skipped(manager, base_copy): + _block_table_host(manager) + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=True, reuse_pool_views=False) + dst = torch.zeros(2, 4, 2, 8, dtype=torch.int32) + + _decode_copy(manager, dst) + _decode_copy(manager, dst) + + assert base_copy.call_count == 1 + + +def test_a_changed_host_table_is_copied(manager, base_copy): + host = _block_table_host(manager) + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=True, reuse_pool_views=False) + dst = torch.zeros(2, 4, 2, 8, dtype=torch.int32) + + _decode_copy(manager, dst) + host[0, 1, 0, 3] += 100 # a request grew into a new page + _decode_copy(manager, dst) + _decode_copy(manager, dst) + + assert base_copy.call_count == 2 + + +@pytest.mark.parametrize( + "change", + [ + dict(request_ids=(1, 0)), + dict(request_ids=(0, 2)), + dict(max_blocks=4), + ], + ids=["batch_order", "request_set", "max_blocks"], +) +def test_a_changed_batch_is_copied(manager, base_copy, change): + _block_table_host(manager) + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=True, reuse_pool_views=False) + dst = torch.zeros(2, 4, 2, 8, dtype=torch.int32) + + _decode_copy(manager, dst) + _decode_copy(manager, dst, **change) + + assert base_copy.call_count == 2 + + +def test_each_destination_keeps_its_own_record(manager, base_copy): + _block_table_host(manager) + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=True, reuse_pool_views=False) + first = torch.zeros(2, 4, 2, 8, dtype=torch.int32) + second = torch.zeros(2, 4, 2, 8, dtype=torch.int32) + + _decode_copy(manager, first) + _decode_copy(manager, second) + _decode_copy(manager, first) + _decode_copy(manager, second) + + assert base_copy.call_count == 2 + + +def test_a_step_with_contexts_always_copies_and_resets_the_record(manager, base_copy): + _block_table_host(manager) + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=True, reuse_pool_views=False) + dst = torch.zeros(2, 4, 2, 8, dtype=torch.int32) + + _decode_copy(manager, dst) + _decode_copy(manager, dst, num_contexts=1) + _decode_copy(manager, dst, num_contexts=1) + # The context steps wrote dst through the base copy: the next decode copy must run. + _decode_copy(manager, dst) + + assert base_copy.call_count == 4 + + +def test_a_failed_block_copy_does_not_leave_a_valid_record(manager, base_copy): + host = _block_table_host(manager) + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=True, reuse_pool_views=False) + dst = torch.zeros_like(host) + _decode_copy(manager, dst) + base_copy.side_effect = RuntimeError("copy failed") + with pytest.raises(RuntimeError, match="copy failed"): + _decode_copy(manager, dst, request_ids=(1, 0)) + base_copy.side_effect = None + _decode_copy(manager, dst) + assert base_copy.call_count == 3 + + +def test_a_different_destination_view_requires_a_new_copy(manager, base_copy): + host = _block_table_host(manager) + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=True, reuse_pool_views=False) + storage = torch.zeros(2 * host.numel(), dtype=host.dtype) + first = storage[: host.numel()].view_as(host) + second = storage[host.numel() :].view_as(host) + for dst in (first, second, first): + _decode_copy(manager, dst) + assert base_copy.call_count == 3 + + +# --------------------------------------------------------------------------- +# Pool-view shortcut +# --------------------------------------------------------------------------- + +NUM_SLOTS, SCALE, PAGE_SHAPE = 3, 4, [2, 8, 16] + + +@pytest.fixture +def pools(monkeypatch): + """Two pools the manager's slot geometry resolves to: one for K/V, one for block scales.""" + numel = NUM_SLOTS * SCALE * 2 * 8 * 16 + kv_pool = torch.zeros(numel, dtype=torch.bfloat16, device="cuda") + scale_pool = torch.zeros(numel, dtype=torch.uint8, device="cuda") + + def kv_slot_geometry(self, layer_idx, kv_layout): + del self, layer_idx, kv_layout + return kv_pool.data_ptr(), torch.bfloat16, NUM_SLOTS, SCALE, list(PAGE_SHAPE) + + def kv_scale_slot_geometry(self, layer_idx, kv_layout): + del self, layer_idx, kv_layout + return scale_pool.data_ptr(), torch.uint8, NUM_SLOTS, SCALE, list(PAGE_SHAPE) + + monkeypatch.setattr(MiniMaxM3KVCacheManagerV2, "_kv_slot_geometry", kv_slot_geometry) + monkeypatch.setattr( + MiniMaxM3KVCacheManagerV2, "_kv_scale_slot_geometry", kv_scale_slot_geometry + ) + return kv_pool, scale_pool + + +def _assert_kv_view( + view: torch.Tensor, pool: torch.Tensor, manager: MiniMaxM3KVCacheManagerV2 +) -> None: + # The pool is laid out in the manager's physical layout; a view in the other layout swaps the + # page's head and token axes over the same memory. + page = list(PAGE_SHAPE) + if manager._main_kv_view_layout != manager._main_kv_layout_name(): + page[0], page[1] = page[1], page[0] + assert view.data_ptr() == pool.data_ptr() + assert list(view.shape) == [NUM_SLOTS, 2, *page] + assert view.stride(0) == SCALE * 2 * 8 * 16 + + +@requires_cuda +def test_pool_views_are_rebuilt_on_every_call_by_default(manager, pools): + kv_pool, _ = pools + + first = manager.get_buffers(3) + second = manager.get_buffers(3) + + assert not manager.reuse_pool_views + assert first is not second + _assert_kv_view(first, kv_pool, manager) + _assert_kv_view(second, kv_pool, manager) + + +@requires_cuda +def test_reused_pool_views_are_built_once_per_geometry(manager, pools): + kv_pool, scale_pool = pools + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=False, reuse_pool_views=True) + + view = manager.get_buffers(3) + scale_view = manager.get_block_scale_buffers(3) + + assert manager.get_buffers(3) is view + assert manager.get_block_scale_buffers(3) is scale_view + assert scale_view is not view + _assert_kv_view(view, kv_pool, manager) + assert scale_view.data_ptr() == scale_pool.data_ptr() + assert scale_view.dtype == torch.uint8 + + +@requires_cuda +def test_reused_pool_views_are_kept_per_inference_mode(manager, pools): + kv_pool, _ = pools + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=False, reuse_pool_views=True) + + view = manager.get_buffers(3) + with torch.inference_mode(): + inference_view = manager.get_buffers(3) + assert manager.get_buffers(3) is inference_view + + assert inference_view is not view + assert inference_view.is_inference() + assert not view.is_inference() + _assert_kv_view(inference_view, kv_pool, manager) + + +@requires_cuda +def test_unused_slots_are_filled_after_a_full_buffer_step(): + buf = torch.full((8,), 7, dtype=torch.int32, device="cuda") + with h2d_cache.active(H2DSkipCache()): + _fill_unused_slots(buf, 3) + # A step whose live rows cover the whole buffer overwrites the filled tail. + buf.copy_(torch.arange(8, dtype=torch.int32, device="cuda")) + _fill_unused_slots(buf, buf.shape[0]) + buf[:3] = 0 + _fill_unused_slots(buf, 3) + assert buf.tolist() == [0, 0, 0, -1, -1, -1, -1, -1] diff --git a/tests/unittest/_torch/compilation/test_phase_selective_forward.py b/tests/unittest/_torch/compilation/test_phase_selective_forward.py index 80c6c69ba206..6b073786eb0d 100644 --- a/tests/unittest/_torch/compilation/test_phase_selective_forward.py +++ b/tests/unittest/_torch/compilation/test_phase_selective_forward.py @@ -168,7 +168,9 @@ def test_decoder_runner_validates_compiled_body_at_initialization( torch.nn.Module.__init__(model) wrapper, _, _ = _create_prefill_compiled_model() model.model = wrapper if has_wrapper else torch.nn.Identity() - config = SimpleNamespace(compile_only_piecewise_graphs=enabled, max_beam_width=1) + config = SimpleNamespace( + compile_only_piecewise_graphs=enabled, max_beam_width=1, is_spec_decode=False + ) from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as runner_module for component in ("DecoderInputBuffers", "DecoderMetadata", "DecoderWarmup"): diff --git a/tests/unittest/_torch/distributed/test_mnnvl_workspace_comm.py b/tests/unittest/_torch/distributed/test_mnnvl_workspace_comm.py index b4314d1657e6..2c97485b9ba1 100644 --- a/tests/unittest/_torch/distributed/test_mnnvl_workspace_comm.py +++ b/tests/unittest/_torch/distributed/test_mnnvl_workspace_comm.py @@ -21,8 +21,13 @@ to make the MNNVL workspace allocation fail and fall back to NCCL. """ +import importlib +from unittest.mock import Mock + import pytest +import torch +import tensorrt_llm._utils as utils from tensorrt_llm._torch.distributed import ops @@ -279,3 +284,70 @@ def test_explicit_request_still_respects_hardware(mnnvl_capable_hardware, monkey mapping = _FakeMapping() mapping.multi_node = True assert not ops.MNNVLAllReduce.is_mnnvl(mapping, torch.int8, explicitly_requested=True) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("tp_size", [4, 8]) +@pytest.mark.parametrize("kind", ["mc_allreduce", "shard_accept", "sampled_accept"]) +@pytest.mark.parametrize("setup_succeeds", [True, False]) +def test_minimax_workspace_constructors_use_tp_process_group( + ray_mode, monkeypatch, tp_size: int, kind: str, setup_succeeds: bool +) -> None: + """Run the production constructors without MPI; only GPU allocation is substituted.""" + module = importlib.import_module( + "tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3." + f"minimax_m3_nvfp4__sm_103__tp{tp_size}.ops.m3_{kind}" + ) + pg = _FakeProcessGroup() + pg.GROUP_SIZE = tp_size + pg.backend.group_size = tp_size if setup_succeeds else 0 + mapping = _FakeMapping(pg) + mapping.tp_size = tp_size + no_mpi = Mock(side_effect=AssertionError("MPI must not be used by Ray workers")) + monkeypatch.setattr(ops, "mpi_comm", no_mpi) + monkeypatch.setattr(utils, "mpi_comm", no_mpi) + monkeypatch.setattr(utils, "mpi_disabled", lambda: True) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + synchronize = Mock() + monkeypatch.setattr(torch.cuda, "synchronize", synchronize) + + # CPU buffers exercise tensor initialization and constructor outputs without CUDA. + uc = torch.empty(4096, dtype=torch.int32 if kind == "mc_allreduce" else torch.int64) + mc = torch.empty_like(uc) + handle = Mock(spec_set=["get_uc_buffer", "get_mc_buffer"]) + handle.get_uc_buffer.return_value = uc + handle.get_mc_buffer.return_value = mc + allocate = Mock(return_value=handle) + monkeypatch.setattr(ops, "_make_mnnvl_mcast_buffer", allocate) + if kind == "sampled_accept": + # This host-only C++ op supplies a size; communicator selection stays real. + monkeypatch.setattr(module, "_ws_words_op", lambda world: 4096) + + def create() -> object: + if kind == "mc_allreduce": + return module.McAllReduceWorkspace.create(mapping, hidden=6144) + if kind == "shard_accept": + return module.create_workspace(mapping, torch.device("cpu"), slot_rows=4) + return module.SampledAcceptWorkspace.create(mapping, torch.device("cpu")) + + if not setup_succeeds: + # A borrowed ProcessGroup has no MPI Free(); rejection must preserve the intended error. + with pytest.raises(RuntimeError, match="outside CUDA-graph capture|creation failed"): + create() + allocate.assert_not_called() + synchronize.assert_not_called() + else: + workspace = create() + assert workspace.comm is pg + assert workspace.rank == mapping.tp_rank + if kind == "sampled_accept": + assert workspace.world == tp_size + else: + assert workspace.world_size == tp_size + assert workspace.handle is handle + allocate.assert_called_once() + assert allocate.call_args.args[0] is pg + assert pg.backend.allreduces == 2 + synchronize.assert_called_once() + no_mpi.assert_not_called() + assert pg.requested_devices == [torch.device("cpu")] * pg.backend.allreduces diff --git a/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup.py b/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup.py index 2ca5fbd39a61..6686ad989095 100644 --- a/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup.py +++ b/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup.py @@ -43,6 +43,7 @@ from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup import DecoderWarmup from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType +from tensorrt_llm._torch.pyexecutor.runtime_hooks import DEFAULT_RUNTIME_HOOKS from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer from tensorrt_llm._torch.utils import is_torch_compiling, torch_compiling @@ -368,6 +369,7 @@ def test_compiled_mxfp8_warmup_backend_selection( monkeypatch.setitem(sys.modules, "flashinfer", SimpleNamespace(autotune=flashinfer_tune)) method = MXFP8LinearMethod() engine = SimpleNamespace( + runtime_hooks=DEFAULT_RUNTIME_HOOKS, _warmup_timer=_WarmupTimer(rank=0), _config=SimpleNamespace( enable_autotuner=True, @@ -489,6 +491,7 @@ def test_autotuner_warmup_skips_only_completed_token_shape( ) -> None: """A repeated warmup with an already tuned token shape must not re-run tuning.""" engine = SimpleNamespace( + runtime_hooks=DEFAULT_RUNTIME_HOOKS, _config=SimpleNamespace( enable_autotuner=True, torch_compile_enabled=False, diff --git a/tests/unittest/_torch/executor/engine/test_decoder_forward_scopes.py b/tests/unittest/_torch/executor/engine/test_decoder_forward_scopes.py new file mode 100644 index 000000000000..d3c72fd0a563 --- /dev/null +++ b/tests/unittest/_torch/executor/engine/test_decoder_forward_scopes.py @@ -0,0 +1,242 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DecoderRunner.forward scopes its own step-input pack and H2D skip cache to the step.""" + +from typing import Callable, Optional +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import h2d_cache, step_inputs +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as decoder_runner_module +from tensorrt_llm._torch.pyexecutor.h2d_cache import H2DSkipCache, copy_h2d_if_changed +from tensorrt_llm._torch.pyexecutor.sampler.sampler import ( + GREEDY_HIDDEN_STATES_KEY, + GREEDY_TOKENS_KEY, +) +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +class _ModelWithoutWeights(torch.nn.Module): + """What DecoderRunner.forward reads from the model before it runs a step.""" + + extra_attrs: dict = {} + + +class _StepFailed(RuntimeError): + pass + + +def _runner( + step: Callable[..., object], + skip_cache: Optional[H2DSkipCache], + pack: Optional[StepInputPack], +) -> DecoderRunner: + # forward reads only the model's extra attributes, the two per-engine objects and the step. + runner = object.__new__(DecoderRunner) + runner._model = _ModelWithoutWeights() + runner._h2d_cache = skip_cache + runner._step_pack = pack + runner._forward_decoder = step + return runner + + +def _host(values): + return torch.tensor(values, dtype=torch.int32).pin_memory() + + +def test_forward_activates_the_runner_cache_and_pack_for_the_step_only(): + skip_cache = H2DSkipCache() + pack = StepInputPack(torch.device("cuda", torch.cuda.current_device())) + dst = torch.zeros(3, dtype=torch.int32, device="cuda") + values = _host([1, 2, 3]) + seen = {} + + def step(inputs, resource_manager, is_dummy=False): + seen["skip_enabled"] = h2d_cache.enabled() + step_inputs.begin() + seen["pack_open"] = pack.open + seen["copies"] = [copy_h2d_if_changed(dst, values) for _ in range(2)] + step_inputs.flush() + return "outputs" + + runner = _runner(step, skip_cache, pack) + + assert runner.forward(object(), resource_manager=None) == "outputs" + + assert seen == {"skip_enabled": True, "pack_open": True, "copies": [True, False]} + # Outside the step neither is active any more. + assert not h2d_cache.enabled() + assert not step_inputs.is_open() + # The shadow the step recorded belongs to the runner's own cache. + with h2d_cache.active(skip_cache): + assert copy_h2d_if_changed(dst, values) is False + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3] + + +def test_forward_without_armed_hooks_activates_nothing(): + seen = {} + + def step(inputs, resource_manager, is_dummy=False): + step_inputs.begin() + seen["skip_enabled"] = h2d_cache.enabled() + seen["pack_open"] = step_inputs.is_open() + + _runner(step, None, None).forward(object(), resource_manager=None) + + assert seen == {"skip_enabled": False, "pack_open": False} + + +def test_forward_hides_an_outer_context_cache_and_pack(): + """Another engine's objects, active around this call, are not used by this runner's step.""" + outer_cache = H2DSkipCache() + outer_pack = StepInputPack(torch.device("cuda", torch.cuda.current_device())) + seen = {} + + def step(inputs, resource_manager, is_dummy=False): + seen["skip_enabled"] = h2d_cache.enabled() + seen["pack_open"] = step_inputs.is_open() + + with h2d_cache.active(outer_cache), step_inputs.active(outer_pack): + outer_pack.begin() + _runner(step, None, None).forward(object(), resource_manager=None) + assert step_inputs.is_open() + outer_pack.discard() + + assert seen == {"skip_enabled": False, "pack_open": False} + + +def test_a_failed_step_discards_its_pack_and_clears_its_cache(): + skip_cache = H2DSkipCache() + pack = StepInputPack(torch.device("cuda", torch.cuda.current_device())) + dst = torch.zeros(3, dtype=torch.int32, device="cuda") + values = _host([7, 8, 9]) + + def failing_step(inputs, resource_manager, is_dummy=False): + step_inputs.begin() + assert copy_h2d_if_changed(dst, values, pack=True) is True + assert pack.staged + raise _StepFailed("forward failed after staging its inputs") + + runner = _runner(failing_step, skip_cache, pack) + + with pytest.raises(_StepFailed): + runner.forward(object(), resource_manager=None) + + assert not pack.open + assert pack.staged == [] and pack.h2d == [] and pack.d2d == [] + assert not h2d_cache.enabled() and not step_inputs.is_open() + torch.cuda.synchronize() + assert dst.tolist() == [0, 0, 0] + + # The next step re-sends the values the failed one never delivered. + def next_step(inputs, resource_manager, is_dummy=False): + step_inputs.begin() + sent = copy_h2d_if_changed(dst, values, pack=True) + step_inputs.flush() + return sent + + runner._forward_decoder = next_step + if hasattr(torch.ops.trtllm, "step_inputs_scatter"): + assert runner.forward(object(), resource_manager=None) is True + torch.cuda.synchronize() + assert dst.tolist() == [7, 8, 9] + else: + # Without the scatter op the pack cannot be issued; the cache decision is what matters. + with h2d_cache.active(skip_cache): + assert copy_h2d_if_changed(dst, values) is True + + +class _LogitsProcessor: + def forward(self, hidden_states, lm_head, attn_metadata, return_context_logits): + return lm_head(hidden_states) * 2 + + +class _GreedyFallbackModel(torch.nn.Module): + """The LM head and logits processor the greedy fallback runs eagerly.""" + + def __init__(self) -> None: + super().__init__() + self.lm_head = torch.nn.Identity() + self.logits_processor = _LogitsProcessor() + + +class _ScheduledWithContext: + """A batch with a context request: greedy tokens cannot serve it.""" + + num_context_requests = 1 + generation_requests = () + + +def test_greedy_fallback_leaves_the_replayed_graph_outputs_intact(): + runner = object.__new__(DecoderRunner) + runner._model = _GreedyFallbackModel() + # max_beam_width is read from the runner config; the fallback reads nothing else from it. + runner._config = Mock(spec_set=["max_beam_width"], max_beam_width=1) + runner.sample_in_graph_callable = None + hidden = torch.ones(2, 4, device="cuda") + tokens = torch.zeros(2, dtype=torch.int32, device="cuda") + # The dict a CUDA graph runner returns on every replay of its key. + graph_outputs = {GREEDY_TOKENS_KEY: tokens, GREEDY_HIDDEN_STATES_KEY: hidden} + + outputs = runner._logits_for_greedy_tokens(_ScheduledWithContext(), graph_outputs) + + assert GREEDY_TOKENS_KEY not in outputs and GREEDY_HIDDEN_STATES_KEY not in outputs + assert torch.equal(outputs["logits"], hidden * 2) + # The next replay of the same graph still sees its greedy outputs and takes the fallback again. + assert graph_outputs == {GREEDY_TOKENS_KEY: tokens, GREEDY_HIDDEN_STATES_KEY: hidden} + + +class _ScheduledGenerationOnly: + """A generation-only batch every request of which greedy tokens could serve (greedy_tokens_usable patched).""" + + num_context_requests = 0 + generation_requests = (object(),) + + +class _ScheduledIdle: + """An idle attention-DP rank's empty batch.""" + + num_context_requests = 0 + generation_requests = () + + +@pytest.mark.parametrize("keep", [True, False]) +def test_greedy_only_outputs_follow_the_guided_decoder_switch(keep, monkeypatch): + """With an executor-side guided decoder (keep_greedy_only_outputs False) the batch gets logits; otherwise the + forward's greedy-only outputs are returned as they are.""" + runner = object.__new__(DecoderRunner) + runner._model = _GreedyFallbackModel() + runner._config = Mock(spec_set=["max_beam_width"], max_beam_width=1) + runner.sample_in_graph_callable = None + runner.keep_greedy_only_outputs = keep + monkeypatch.setattr(decoder_runner_module, "greedy_tokens_usable", lambda request: True) + hidden = torch.ones(2, 4, device="cuda") + tokens = torch.zeros(2, dtype=torch.int32, device="cuda") + graph_outputs = {GREEDY_TOKENS_KEY: tokens, GREEDY_HIDDEN_STATES_KEY: hidden} + + outputs = runner._logits_for_greedy_tokens(_ScheduledGenerationOnly(), graph_outputs) + + if keep: + assert outputs is graph_outputs and "logits" not in outputs + else: + assert torch.equal(outputs["logits"], hidden * 2) + + +def test_an_idle_batch_gets_logits_for_its_dummy_row(): + """An empty batch (an idle attention-DP rank's dummy row) never keeps greedy-only outputs.""" + runner = object.__new__(DecoderRunner) + runner._model = _GreedyFallbackModel() + runner._config = Mock(spec_set=["max_beam_width"], max_beam_width=1) + runner.sample_in_graph_callable = None + hidden = torch.ones(1, 4, device="cuda") + tokens = torch.zeros(1, dtype=torch.int32, device="cuda") + outputs = runner._logits_for_greedy_tokens( + _ScheduledIdle(), {GREEDY_TOKENS_KEY: tokens, GREEDY_HIDDEN_STATES_KEY: hidden} + ) + assert torch.equal(outputs["logits"], hidden * 2) diff --git a/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py b/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py new file mode 100644 index 000000000000..3bfbaee8c913 --- /dev/null +++ b/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DecoderRunner and DecoderMetadata wiring of the step-loop hooks a model declares.""" + +import pytest + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import ( + MiniMaxM3KVCacheManagerV2, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.metadata import DecoderMetadata +from tensorrt_llm._torch.pyexecutor.runtime_hooks import ( + DEFAULT_RUNTIME_HOOKS, + MsaRuntimePolicy, + RuntimeHooks, +) +from tensorrt_llm.llmapi.llm_args import ( + BaseSparseAttentionConfig, + DeepSeekSparseAttentionConfig, + MiniMaxM3SparseAttentionConfig, +) + +pytestmark = pytest.mark.cpu_only + +_SHORT_IDENTITY = RuntimeHooks(msa=MsaRuntimePolicy(short_sequence_identity=True)) + + +def _runner( + sparse_attention_config: BaseSparseAttentionConfig | None, hooks: RuntimeHooks +) -> DecoderRunner: + # Only the two attributes the threshold reads. + runner = object.__new__(DecoderRunner) + runner._sparse_attention_config = sparse_attention_config + runner.runtime_hooks = hooks + return runner + + +def _metadata(hooks: RuntimeHooks, attn_metadata: object) -> DecoderMetadata: + # Only the two attributes the MSA arming reads. + metadata = object.__new__(DecoderMetadata) + metadata.runtime_hooks = hooks + metadata.attn_metadata = attn_metadata + return metadata + + +def test_no_threshold_without_a_split_or_a_short_identity_policy(): + assert _runner(None, DEFAULT_RUNTIME_HOOKS)._short_long_seq_len_threshold() is None + assert ( + _runner( + MiniMaxM3SparseAttentionConfig(), DEFAULT_RUNTIME_HOOKS + )._short_long_seq_len_threshold() + is None + ) + + +def test_a_sparse_config_that_asks_for_the_split_sets_the_threshold(): + config = DeepSeekSparseAttentionConfig(index_topk=512, skip_indexer_for_short_seqs=True) + + # The sparse config wins over the hooks. + assert _runner(config, _SHORT_IDENTITY)._short_long_seq_len_threshold() == 512 + + +def test_short_identity_threshold_is_the_topk_span(): + config = MiniMaxM3SparseAttentionConfig(sparse_topk_blocks=16, sparse_block_size=128) + + assert _runner(config, _SHORT_IDENTITY)._short_long_seq_len_threshold() == 16 * 128 + + +@pytest.mark.parametrize( + "config", + [None, DeepSeekSparseAttentionConfig(index_topk=512, skip_indexer_for_short_seqs=False)], + ids=["dense", "dsa_without_split"], +) +def test_short_identity_requires_the_minimax_m3_sparse_config(config): + with pytest.raises(ValueError, match="requires a MiniMax-M3 sparse attention config"): + _runner(config, _SHORT_IDENTITY)._short_long_seq_len_threshold() + + +def test_a_default_msa_policy_arms_nothing(): + # A non-MSA metadata and no cache manager are fine when nothing is armed. + metadata = _metadata(DEFAULT_RUNTIME_HOOKS, object.__new__(TrtllmAttentionMetadata)) + + metadata._arm_msa_policy(None) + + +def test_an_msa_policy_requires_the_msa_backend_and_cache(): + metadata = _metadata( + RuntimeHooks(msa=MsaRuntimePolicy(reuse_pool_views=True)), + object.__new__(TrtllmAttentionMetadata), + ) + + with pytest.raises(ValueError, match="requires the MiniMax-M3 MSA sparse attention backend"): + metadata._arm_msa_policy(object.__new__(MiniMaxM3KVCacheManagerV2)) + + +def test_an_msa_policy_reaches_the_metadata_and_the_cache(): + policy = MsaRuntimePolicy( + short_sequence_identity=True, + skip_unchanged_block_offsets=True, + reuse_pool_views=False, + ) + metadata = _metadata( + RuntimeHooks(msa=policy), object.__new__(MiniMaxM3MsaSparseAttentionMetadata) + ) + manager = object.__new__(MiniMaxM3KVCacheManagerV2) + + metadata._arm_msa_policy(manager) + + assert metadata.attn_metadata.msa_policy is policy + assert manager.skip_unchanged_block_offsets is True + assert manager.reuse_pool_views is False diff --git a/tests/unittest/_torch/executor/test_cuda_graph_runner.py b/tests/unittest/_torch/executor/test_cuda_graph_runner.py index 1fa9c38800ad..85f953cf8aef 100644 --- a/tests/unittest/_torch/executor/test_cuda_graph_runner.py +++ b/tests/unittest/_torch/executor/test_cuda_graph_runner.py @@ -15,8 +15,10 @@ from _torch.executor.engine_test_utils import _make_request_stub import tensorrt_llm._torch.pyexecutor.cuda_graph_runner as cuda_graph_runner_module +from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( CUDAGraphRunner, + CUDAGraphRunnerConfig, EncoderCUDAGraphRunner, EncoderCUDAGraphRunnerConfig, KeyType, @@ -24,8 +26,15 @@ _restore_spec_decode_capture_state, _save_spec_decode_capture_state, ) +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests -from tensorrt_llm.llmapi.llm_args import SeqLenAwareSparseAttentionConfig +from tensorrt_llm.bindings import SamplingConfig +from tensorrt_llm.llmapi.llm_args import ( + BaseSparseAttentionConfig, + DeepSeekSparseAttentionConfig, + SeqLenAwareSparseAttentionConfig, +) +from tensorrt_llm.mapping import Mapping if TYPE_CHECKING: from collections.abc import Iterator @@ -461,6 +470,110 @@ def test_encoder_decoder_fixed_slots_restore_source_order() -> None: torch.testing.assert_close(restored_output, expected_output) +def _decoder_runner( + *, + short_long_seq_len_threshold: int | None = None, + sparse_attention_config: BaseSparseAttentionConfig | None = None, + private_metadata_buffers: bool = False, +) -> CUDAGraphRunner: + """A generation-graph runner with capture disabled, so it allocates no device buffers.""" + config = CUDAGraphRunnerConfig( + use_cuda_graph=False, + cuda_graph_padding_enabled=False, + cuda_graph_batch_sizes=[1, 2], + max_cuda_graph_batch_size=2, + batch_size=2, + max_beam_width=1, + max_num_tokens=2, + use_mrope=False, + spec_config=None, + enable_attention_dp=False, + original_max_total_draft_tokens=0, + is_encoder_decoder=False, + mapping=Mapping(), + dist=None, + sparse_attention_config=sparse_attention_config, + short_long_seq_len_threshold=short_long_seq_len_threshold, + private_metadata_buffers=private_metadata_buffers, + ) + return CUDAGraphRunner(config) + + +def _dsa_config(skip_indexer_for_short_seqs: bool) -> DeepSeekSparseAttentionConfig: + return DeepSeekSparseAttentionConfig( + index_topk=512, skip_indexer_for_short_seqs=skip_indexer_for_short_seqs + ) + + +def test_short_long_threshold_from_the_runner_config_wins() -> None: + runner = _decoder_runner( + short_long_seq_len_threshold=2_048, sparse_attention_config=_dsa_config(True) + ) + + assert runner.short_long_seq_len_threshold == 2_048 + + +def test_short_long_threshold_is_derived_from_a_sparse_config_that_asks_for_the_split() -> None: + runner = _decoder_runner(sparse_attention_config=_dsa_config(True)) + + assert runner.short_long_seq_len_threshold == 512 + + +@pytest.mark.parametrize( + "sparse_attention_config", + [None, _dsa_config(False)], + ids=["dense", "dsa_without_split"], +) +def test_short_long_threshold_is_none_for_a_single_graph_family(sparse_attention_config) -> None: + runner = _decoder_runner(sparse_attention_config=sparse_attention_config) + + assert runner.short_long_seq_len_threshold is None + + +def test_graph_metadata_uses_the_process_buffer_pool_by_default() -> None: + assert _decoder_runner().cuda_graph_meta_buffers is get_memory_buffers() + + +def test_private_metadata_buffers_give_each_runner_its_own_pool() -> None: + first = _decoder_runner(private_metadata_buffers=True) + second = _decoder_runner(private_metadata_buffers=True) + + assert first.cuda_graph_meta_buffers is not get_memory_buffers() + assert first.cuda_graph_meta_buffers is not second.cuda_graph_meta_buffers + + +def _generation_batch(*num_tokens: int) -> ScheduledRequests: + batch = ScheduledRequests() + batch.generation_requests = [ + LlmRequest( + request_id=request_id, + max_new_tokens=4, + input_tokens=list(range(tokens)), + sampling_config=SamplingConfig(), + seq_slot=request_id, + is_streaming=False, + ) + for request_id, tokens in enumerate(num_tokens) + ] + return batch + + +@pytest.mark.parametrize( + ("num_tokens", "short"), + [((100,), True), ((100, 2_048), True), ((100, 2_049), False), ((4_000,), False)], +) +def test_seq_len_mode_splits_at_the_configured_threshold(num_tokens, short) -> None: + runner = _decoder_runner(short_long_seq_len_threshold=2_048) + + assert runner._get_seq_len_mode(_generation_batch(*num_tokens)) is short + + +def test_seq_len_mode_is_long_without_a_threshold() -> None: + runner = _decoder_runner() + + assert runner._get_seq_len_mode(_generation_batch(1)) is False + + class SparseGraphKeyTestCase(unittest.TestCase): def test_sparse_sequence_mode_uses_promoted_context_cursor(self) -> None: sparse_config = Mock(spec=SeqLenAwareSparseAttentionConfig) diff --git a/tests/unittest/_torch/executor/test_device_copy_buffer.py b/tests/unittest/_torch/executor/test_device_copy_buffer.py index 6aa126d67b6c..829fb57bb28b 100644 --- a/tests/unittest/_torch/executor/test_device_copy_buffer.py +++ b/tests/unittest/_torch/executor/test_device_copy_buffer.py @@ -203,3 +203,36 @@ def counted(dst, src, *args, **kwargs): metadata.update(_host([1, 2])) assert storage.tolist() == [1, 2] assert writes == [[1, 2], [1, 2], [9, 2], [1, 2]] + + +def test_specialized_copy_records_host_inputs_and_layout_key() -> None: + buffer = DeviceCopyBuffer(torch.zeros(4, dtype=torch.int32)) + host = _host([1, 2, 3, 4]).view(2, 2) + key = ((7, 9), 1, 2, 8) + assert not buffer.matches(host, key=key, exact=True) + buffer.record(host, key=key) + assert buffer.matches(host, key=key, exact=True) + assert not buffer.matches(host, key=((9, 7), 1, 2, 8), exact=True) + assert not buffer.matches(host.flatten(), key=key, exact=True) + assert not buffer.matches(host.to(torch.int64), key=key, exact=True) + host[0, 0] = 99 + assert not buffer.matches(host, key=key, exact=True) + + +def test_invalidation_discards_a_recorded_copy() -> None: + buffer = DeviceCopyBuffer(torch.zeros(4, dtype=torch.int32)) + host = _host([1, 2]) + buffer.record(host) + assert buffer.matches(host) + buffer.invalidate() + assert not buffer.matches(host) + buffer.update(host) + assert buffer.tensor.tolist() == [1, 2, 0, 0] + + +def test_new_layout_does_not_retain_the_previous_prefix_tail() -> None: + buffer = DeviceCopyBuffer(torch.zeros(4, dtype=torch.int32)) + buffer.record(_host([1, 2, 3, 4]), key="first") + buffer.record(_host([1, 2]), key="second") + assert buffer.matches(_host([1, 2]), key="second") + assert not buffer.matches(_host([1, 2, 3, 4]), key="second") diff --git a/tests/unittest/_torch/executor/test_h2d_cache.py b/tests/unittest/_torch/executor/test_h2d_cache.py new file mode 100644 index 000000000000..bd86fc1d1fe1 --- /dev/null +++ b/tests/unittest/_torch/executor/test_h2d_cache.py @@ -0,0 +1,333 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Skipping host-to-device metadata copies whose values the destination already holds.""" + +import threading + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import h2d_cache +from tensorrt_llm._torch.pyexecutor.h2d_cache import ( + H2DSkipCache, + copy_h2d_if_changed, + fill_tail_if_changed, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture +def cache(): + """A fresh cache, active for the test body.""" + skip_cache = H2DSkipCache() + with h2d_cache.active(skip_cache): + yield skip_cache + torch.cuda.synchronize() + + +def _host(values, dtype=torch.int32): + return torch.tensor(values, dtype=dtype).pin_memory() + + +def _device_zeros(n, dtype=torch.int32): + return torch.zeros(n, dtype=dtype, device="cuda") + + +def test_without_an_active_cache_every_copy_is_issued_like_copy_(): + assert not h2d_cache.enabled() + dst = _device_zeros(4) + reference = _device_zeros(4) + src = _host([3, 1, 4, 1]) + + for _ in range(3): + assert copy_h2d_if_changed(dst, src) is True + reference.copy_(src, non_blocking=True) + torch.cuda.synchronize() + assert torch.equal(dst, reference) + + +def test_active_scopes_the_cache_to_the_with_block(): + skip_cache = H2DSkipCache() + dst = _device_zeros(2) + values = _host([1, 2]) + + with h2d_cache.active(skip_cache): + assert h2d_cache.enabled() + assert copy_h2d_if_changed(dst, values) is True + assert copy_h2d_if_changed(dst, values) is False + assert not h2d_cache.enabled() + assert copy_h2d_if_changed(dst, values) is True + # Re-activating the same cache finds its shadow again. + with h2d_cache.active(skip_cache): + assert copy_h2d_if_changed(dst, values) is False + + +def test_active_none_disables_skipping_inside_an_outer_cache(cache): + dst = _device_zeros(2) + values = _host([1, 2]) + assert copy_h2d_if_changed(dst, values) is True + + with h2d_cache.active(None): + assert not h2d_cache.enabled() + assert copy_h2d_if_changed(dst, values) is True + assert copy_h2d_if_changed(dst, values) is False + + +def test_two_caches_do_not_share_shadows(cache): + dst = _device_zeros(3) + values = _host([7, 8, 9]) + other = H2DSkipCache() + + assert copy_h2d_if_changed(dst, values) is True + with h2d_cache.active(other): + # The other cache has never seen dst. + assert copy_h2d_if_changed(dst, values) is True + assert copy_h2d_if_changed(dst, values) is False + assert copy_h2d_if_changed(dst, values) is False + + +def test_a_cache_active_on_one_thread_is_invisible_to_another(cache): + dst = _device_zeros(3) + values = _host([4, 5, 6]) + assert copy_h2d_if_changed(dst, values) is True + seen = {} + + def other_thread(): + seen["enabled"] = h2d_cache.enabled() + seen["copied"] = [copy_h2d_if_changed(dst, values) for _ in range(2)] + torch.cuda.synchronize() + + thread = threading.Thread(target=other_thread) + thread.start() + thread.join() + + assert seen == {"enabled": False, "copied": [True, True]} + assert copy_h2d_if_changed(dst, values) is False + + +def test_threads_with_their_own_caches_keep_them_apart(): + dst = [_device_zeros(2) for _ in range(2)] + values = _host([1, 1]) + caches = [H2DSkipCache(), H2DSkipCache()] + barrier = threading.Barrier(2) + results = {} + + def worker(index): + with h2d_cache.active(caches[index]): + first = copy_h2d_if_changed(dst[index], values) + barrier.wait() + second = copy_h2d_if_changed(dst[index], values) + # The other thread's destination was never shadowed by this cache. + other = copy_h2d_if_changed(dst[1 - index], values) + torch.cuda.synchronize() + results[index] = (first, second, other) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert results == {0: (True, False, True), 1: (True, False, True)} + + +def test_a_repeated_identical_copy_is_skipped(cache): + dst = _device_zeros(4) + + assert copy_h2d_if_changed(dst, _host([5, 6, 7, 8])) is True + assert copy_h2d_if_changed(dst, _host([5, 6, 7, 8])) is False + torch.cuda.synchronize() + assert dst.tolist() == [5, 6, 7, 8] + + +def test_clear_makes_the_next_identical_copy_issue(cache): + """The runner clears the cache after a discarded step, whose copies never landed.""" + dst = _device_zeros(3) + values = _host([1, 2, 3]) + assert copy_h2d_if_changed(dst, values) is True + + cache.clear() + + assert copy_h2d_if_changed(dst, values) is True + assert copy_h2d_if_changed(dst, values) is False + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3] + + +def test_changed_values_are_copied(cache): + dst = _device_zeros(4) + + assert copy_h2d_if_changed(dst, _host([1, 2, 3, 4])) is True + assert copy_h2d_if_changed(dst, _host([1, 2, 3, 5])) is True + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3, 5] + assert copy_h2d_if_changed(dst, _host([1, 2, 3, 5])) is False + + +def test_the_shadow_is_a_snapshot_of_the_source(cache): + """Mutating the host source after the copy must not make a later copy look unchanged.""" + dst = _device_zeros(3) + src = _host([1, 2, 3]) + + assert copy_h2d_if_changed(dst, src) is True + torch.cuda.synchronize() + src[0] = 9 + assert copy_h2d_if_changed(dst, src) is True + torch.cuda.synchronize() + assert dst.tolist() == [9, 2, 3] + + +def test_a_source_dtype_change_copies(cache): + dst = _device_zeros(3) + + assert copy_h2d_if_changed(dst, _host([1, 2, 3], torch.int32)) is True + assert copy_h2d_if_changed(dst, _host([1, 2, 3], torch.int64)) is True + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3] + + +def test_a_copy_through_another_view_replaces_the_shadow(cache): + buf = _device_zeros(8) + head, tail = buf[:4], buf[4:] + values = _host([1, 2, 3, 4]) + + assert copy_h2d_if_changed(head, values) is True + assert copy_h2d_if_changed(tail, _host([5, 6, 7, 8])) is True + # The shadow now describes the tail, so the head copy must not be skipped. + assert copy_h2d_if_changed(head, values) is True + assert copy_h2d_if_changed(head, values) is False + torch.cuda.synchronize() + assert buf.tolist() == [1, 2, 3, 4, 5, 6, 7, 8] + + +def test_an_overlapping_view_never_skips_on_stale_state(cache): + buf = _device_zeros(6) + values = _host([1, 1, 1, 1]) + + assert copy_h2d_if_changed(buf[:4], values) is True + assert copy_h2d_if_changed(buf[2:6], _host([2, 2, 2, 2])) is True + # buf[2:4] was overwritten through the other view: same values, same view, must copy. + assert copy_h2d_if_changed(buf[:4], values) is True + torch.cuda.synchronize() + assert buf.tolist() == [1, 1, 1, 1, 2, 2] + + +def test_a_view_with_another_shape_of_the_same_bytes_copies(cache): + buf = _device_zeros(4) + + assert copy_h2d_if_changed(buf, _host([1, 2, 3, 4])) is True + assert copy_h2d_if_changed(buf.view(2, 2), _host([[1, 2], [3, 4]])) is True + assert copy_h2d_if_changed(buf.view(2, 2), _host([[1, 2], [3, 4]])) is False + + +def test_a_cuda_source_always_copies_and_drops_the_shadow(cache): + dst = _device_zeros(4) + values = _host([1, 2, 3, 4]) + device_src = torch.arange(4, dtype=torch.int32, device="cuda") + + assert copy_h2d_if_changed(dst, values) is True + assert copy_h2d_if_changed(dst, device_src) is True + assert copy_h2d_if_changed(dst, device_src) is True + # The device-side write replaced what the shadow described. + assert copy_h2d_if_changed(dst, values) is True + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3, 4] + + +def test_a_host_destination_always_copies(cache): + dst = torch.zeros(3, dtype=torch.int32) + + assert copy_h2d_if_changed(dst, _host([1, 2, 3])) is True + assert copy_h2d_if_changed(dst, _host([1, 2, 3])) is True + assert dst.tolist() == [1, 2, 3] + + +def test_capture_always_copies_and_drops_the_shadow(cache): + dst = _device_zeros(4) + values = _host([4, 3, 2, 1]) + assert copy_h2d_if_changed(dst, values) is True + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + results = [] + with torch.cuda.graph(graph): + results.append(copy_h2d_if_changed(dst, values)) + assert results == [True] + + dst.zero_() + graph.replay() + torch.cuda.synchronize() + assert dst.tolist() == [4, 3, 2, 1] + # The captured copy reruns on every replay, so the shadow cannot vouch for the buffer. + assert copy_h2d_if_changed(dst, values) is True + + +def test_without_an_active_cache_every_tail_fill_runs(): + buf = _device_zeros(6) + fill_tail_if_changed(buf, 2, -1) + buf[3] = 5 + fill_tail_if_changed(buf, 2, -1) + fill_tail_if_changed(buf, 6, -1) + assert buf.tolist() == [0, 0, -1, -1, -1, -1] + + +def test_a_repeated_tail_fill_is_skipped_and_a_moved_start_refills(cache): + buf = _device_zeros(6) + fill_tail_if_changed(buf, 2, -1) + # Writes nothing else is allowed to make; they show which fills ran. + buf[3] = 5 + fill_tail_if_changed(buf, 2, -1) + assert buf[3].item() == 5 + fill_tail_if_changed(buf, 4, -1) + buf[3] = 7 + fill_tail_if_changed(buf, 2, -1) + assert buf.tolist() == [0, 0, -1, -1, -1, -1] + + +def test_a_whole_buffer_write_drops_the_tail_record(cache): + buf = _device_zeros(6) + fill_tail_if_changed(buf, 2, -1) + buf.copy_(torch.arange(6, dtype=torch.int32, device="cuda")) + fill_tail_if_changed(buf, 6, -1) + fill_tail_if_changed(buf, 2, -1) + assert buf.tolist() == [0, 1, -1, -1, -1, -1] + + +def test_clear_drops_the_tail_records(cache): + buf = _device_zeros(6) + fill_tail_if_changed(buf, 2, -1) + buf[4] = 9 + cache.clear() + fill_tail_if_changed(buf, 2, -1) + assert buf[2:].eq(-1).all() + + +def test_a_tail_record_keeps_its_storage_so_a_new_buffer_cannot_inherit_it(cache): + buf = _device_zeros(1 << 20) + fill_tail_if_changed(buf, 3, -1) + address = buf.data_ptr() + del buf + # Without the record holding the storage, the allocator hands this same block back, and the fill below would be + # skipped on the freed buffer's record. + fresh = _device_zeros(1 << 20) + assert fresh.data_ptr() != address + fill_tail_if_changed(fresh, 3, -1) + assert fresh[3:].eq(-1).all() + + +def test_failed_copy_through_another_view_invalidates_the_old_record(cache, monkeypatch): + dst = _device_zeros(4) + host = _host([1, 2, 3, 4]) + assert cache.copy(dst, host, pack=False) + + def partial_copy_then_fail(target, source): + target.copy_(source, non_blocking=True) + raise RuntimeError("copy submission failed") + + monkeypatch.setattr(h2d_cache.step_inputs, "copy_", partial_copy_then_fail) + with pytest.raises(RuntimeError, match="copy submission failed"): + cache.copy(dst[1:], _host([9, 9, 9]), pack=True) + assert cache.copy(dst, host, pack=False) + torch.cuda.synchronize() + assert dst.tolist() == host.tolist() diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py index d2459e6804de..d15b77c4fa10 100644 --- a/tests/unittest/_torch/executor/test_runtime_hooks.py +++ b/tests/unittest/_torch/executor/test_runtime_hooks.py @@ -24,6 +24,7 @@ from tensorrt_llm._torch.pyexecutor.runtime_hooks import ( DEFAULT_RUNTIME_HOOKS, ExecutorEvents, + MsaRuntimePolicy, RuntimeHooks, runtime_hooks_of, ) @@ -63,7 +64,13 @@ def test_a_model_without_hooks_gets_the_all_off_default(model): def test_every_field_defaults_to_off(): for field in dataclasses.fields(RuntimeHooks): - assert getattr(DEFAULT_RUNTIME_HOOKS, field.name) in (False, None), field.name + value = getattr(DEFAULT_RUNTIME_HOOKS, field.name) + if dataclasses.is_dataclass(value): + assert all( + getattr(value, item.name) in (False, None) for item in dataclasses.fields(value) + ) + else: + assert value in (False, None), field.name def test_hooks_declared_as_none_mean_the_default(): @@ -347,6 +354,12 @@ def value(self, node: ast.AST, env: dict): if node.attr in _SEAM_ATTRS: return _SEAM_ATTRS[node.attr] return _UNKNOWN + if isinstance(node, ast.IfExp): + condition = self.value(node.test, env) + if condition is not _UNKNOWN: + return self.value(node.body if condition else node.orelse, env) + body, other = self.value(node.body, env), self.value(node.orelse, env) + return body if body is not _UNKNOWN and body is other else _UNKNOWN if self._is_seam_call(node): return DEFAULT_RUNTIME_HOOKS if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not): @@ -483,6 +496,18 @@ def step(self, x): ] +@pytest.mark.parametrize( + "expression, expected", + [ + ("self.runtime_hooks.spec_step_fold if enabled else None", None), + ("self.runtime_hooks.spec_step_fold if enabled else object()", _UNKNOWN), + ], +) +def test_conditional_hook_binding_requires_both_branches_to_agree(expression, expected): + walk = _AllOffWalk("sample.py") + assert walk.value(ast.parse(expression, mode="eval").body, {}) is expected + + def test_the_decode_prologue_call_site_uses_its_protocol(): """The input preparer calls the prologue with the protocol's positional and keyword parameters.""" from tensorrt_llm._torch.pyexecutor.runtime_hooks import DecodePrologue @@ -683,3 +708,97 @@ def test_executor_loop_cleanup_calls_at_shutdown_after_a_completed_loop(complete executor._executor_loop_cleanup() assert executor.is_shutdown assert events.calls == (["at_shutdown"] if completed else []) + + +class _M3PlainModel(torch.nn.Module): + """A model that declares no hooks, like every model outside the opt-in.""" + + +class _M3ClassHooksModel(torch.nn.Module): + runtime_hooks = RuntimeHooks(pack_step_inputs=True, skip_unchanged_h2d=True) + + +def test_m3_model_without_hooks_gets_the_all_off_default(): + hooks = runtime_hooks_of(_M3PlainModel()) + + assert hooks is DEFAULT_RUNTIME_HOOKS + assert hooks == RuntimeHooks() + assert not hooks.pack_step_inputs + assert not hooks.skip_unchanged_h2d + assert not hooks.loop_on_execution_stream + assert not hooks.merge_tuned_tactics_across_ranks + assert not hooks.skip_greedy_rng_upload + assert hooks.spec_step_fold is None + assert hooks.msa == MsaRuntimePolicy() + assert not hooks.msa.short_sequence_identity + assert hooks.msa.kv_update is None + assert not hooks.msa.skip_unchanged_block_offsets + assert not hooks.msa.reuse_pool_views + + +def test_m3_hooks_declared_as_none_mean_the_default(): + model = _M3PlainModel() + model.runtime_hooks = None + + assert runtime_hooks_of(model) is DEFAULT_RUNTIME_HOOKS + + +def test_m3_class_attribute_hooks_are_returned_as_declared(): + model = _M3ClassHooksModel() + + assert runtime_hooks_of(model) is _M3ClassHooksModel.runtime_hooks + + +def test_m3_instance_attribute_hooks_are_returned_as_declared(): + hooks = RuntimeHooks( + loop_on_execution_stream=True, + msa=MsaRuntimePolicy(short_sequence_identity=True, reuse_pool_views=True), + ) + model = _M3PlainModel() + model.runtime_hooks = hooks + + assert runtime_hooks_of(model) is hooks + + +def test_m3_instance_attribute_overrides_class_attribute(): + hooks = RuntimeHooks(skip_greedy_rng_upload=True) + model = _M3ClassHooksModel() + model.runtime_hooks = hooks + + assert runtime_hooks_of(model) is hooks + + +@pytest.mark.parametrize( + "declared", + [True, {"pack_step_inputs": True}, MsaRuntimePolicy(), "RuntimeHooks()"], + ids=["bool", "dict", "msa_policy", "str"], +) +def test_m3_hooks_of_the_wrong_type_raise(declared): + model = _M3PlainModel() + model.runtime_hooks = declared + + with pytest.raises(TypeError, match=r"_M3PlainModel\.runtime_hooks must be a RuntimeHooks"): + runtime_hooks_of(model) + + +def test_m3_hooks_are_immutable_once_declared(): + hooks = RuntimeHooks() + + with pytest.raises(dataclasses.FrozenInstanceError): + hooks.pack_step_inputs = True + with pytest.raises(dataclasses.FrozenInstanceError): + hooks.msa.reuse_pool_views = True + + +@pytest.mark.parametrize( + "hooks, skips", + [ + (RuntimeHooks(), False), + (RuntimeHooks(pack_step_inputs=True, skip_greedy_rng_upload=True), False), + (RuntimeHooks(skip_unchanged_h2d=True), True), + (RuntimeHooks(msa=MsaRuntimePolicy(skip_unchanged_block_offsets=True)), True), + ], + ids=["default", "no_skip_of_metadata", "h2d_skip", "msa_block_offset_skip"], +) +def test_m3_every_metadata_upload_skip_asks_for_private_metadata_buffers(hooks, skips): + assert hooks.skips_unchanged_metadata_uploads is skips diff --git a/tests/unittest/_torch/executor/test_step_inputs.py b/tests/unittest/_torch/executor/test_step_inputs.py new file mode 100644 index 000000000000..c1b4b95a02d4 --- /dev/null +++ b/tests/unittest/_torch/executor/test_step_inputs.py @@ -0,0 +1,604 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Packing a decode step's small input copies into one upload and one scatter launch.""" + +import threading + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import h2d_cache, step_inputs +from tensorrt_llm._torch.pyexecutor.h2d_cache import H2DSkipCache, copy_h2d_if_changed +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack + + +def _scatter_op_registered() -> bool: + return torch.cuda.is_available() and hasattr(torch.ops.trtllm, "step_inputs_scatter") + + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.skipif( + not _scatter_op_registered(), + reason="torch.ops.trtllm.step_inputs_scatter is not registered (built from " + "cpp/tensorrt_llm/thop/minimaxM3/stepInputsScatterOp.cpp)", + ), +] + +# About half a second of GPU time: long enough that work queued behind it is still pending when +# the host checks, short enough for a unit test. +_SLEEP_CYCLES = 1_000_000_000 + + +def _device(): + return torch.device("cuda", torch.cuda.current_device()) + + +@pytest.fixture +def pack(): + """A fresh pack, active for the test body.""" + step_pack = StepInputPack(_device()) + with step_inputs.active(step_pack): + yield step_pack + step_inputs.flush() + torch.cuda.synchronize() + + +def _host(values, dtype=torch.int32): + return torch.tensor(values, dtype=dtype).pin_memory() + + +def _zeros(*shape, dtype=torch.int32): + return torch.zeros(*shape, dtype=dtype, device="cuda") + + +# --------------------------------------------------------------------------- +# No active pack +# --------------------------------------------------------------------------- + + +def test_without_an_active_pack_every_call_is_a_plain_copy(): + dst = _zeros(4) + d2d_dst = _zeros(4) + src_dev = torch.arange(4, dtype=torch.int32, device="cuda") + + step_inputs.begin() + assert not step_inputs.is_open() + assert step_inputs.copy_(dst, _host([1, 2, 3, 4])) is dst + step_inputs.copy_(d2d_dst, src_dev) + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3, 4] + assert d2d_dst.tolist() == [0, 1, 2, 3] + step_inputs.flush() + step_inputs.discard() + + +def test_an_inactive_pack_is_not_used(): + idle = StepInputPack(_device()) + dst = _zeros(2) + + step_inputs.copy_(dst, _host([3, 4])) + torch.cuda.synchronize() + + assert dst.tolist() == [3, 4] + assert not idle.open and idle.staged == [] + + +# --------------------------------------------------------------------------- +# Deferral and ordering +# --------------------------------------------------------------------------- + + +def test_packed_host_copy_lands_only_at_flush(pack): + dst = _zeros(4) + src = _host([1, 2, 3, 4]) + + step_inputs.begin() + assert step_inputs.is_open() and pack.open + assert step_inputs.copy_(dst, src) is dst + # The pack snapshots the host bytes at the copy site. + src.fill_(7) + torch.cuda.synchronize() + assert dst.tolist() == [0, 0, 0, 0] + + step_inputs.flush() + assert not step_inputs.is_open() + assert pack.staged == [] + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3, 4] + + +def test_packed_host_copy_into_a_row_strided_2d_slice(pack): + table = torch.full((4, 6), -1, dtype=torch.int32, device="cuda") + dst = table[1:3, 2:5] + assert not dst.is_contiguous() and dst.stride(1) == 1 + expected = table.clone() + expected[1:3, 2:5] = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32) + + step_inputs.begin() + step_inputs.copy_(dst, _host([[1, 2, 3], [4, 5, 6]])) + torch.cuda.synchronize() + assert torch.equal(table, torch.full_like(table, -1)) + + step_inputs.flush() + torch.cuda.synchronize() + assert torch.equal(table, expected) + + +def test_packed_host_copy_converts_dtype_on_the_host(pack): + dst = _zeros(3) + + step_inputs.begin() + step_inputs.copy_(dst, _host([7, -8, 9], dtype=torch.int64)) + step_inputs.flush() + torch.cuda.synchronize() + + assert dst.tolist() == [7, -8, 9] + + +def test_packed_device_copy_lands_only_at_flush(pack): + dst = _zeros(5) + src = torch.arange(5, dtype=torch.int32, device="cuda") + + step_inputs.begin() + step_inputs.copy_(dst, src) + torch.cuda.synchronize() + assert dst.tolist() == [0] * 5 + + step_inputs.flush() + torch.cuda.synchronize() + assert dst.tolist() == [0, 1, 2, 3, 4] + + +def test_device_copy_reads_what_a_host_copy_of_the_same_pack_wrote(pack): + """The graph runner copies the engine's freshly uploaded buffers into its static inputs.""" + staged = _zeros(4) + static = _zeros(1, 8) + + step_inputs.begin() + step_inputs.copy_(staged, _host([10, 11, 12, 13])) + step_inputs.copy_(static[:, :4], staged.unsqueeze(0)) + step_inputs.flush() + torch.cuda.synchronize() + + assert static.tolist() == [[10, 11, 12, 13, 0, 0, 0, 0]] + + +# The scatter's CTA copies descriptor i with thread t on bytes t, t + 256, ... of it. When the +# second copy starts TAIL bytes before the end of the first (TAIL not a multiple of the warp +# size), its first bytes are the ones the first copy writes last, from other warps: without a +# barrier between descriptors, a warp that finished the first copy early writes (or reads) them +# before a slower warp has written its part of the first copy. +_TAILS = [239, 1000, 4073] + + +@pytest.mark.parametrize("tail", _TAILS) +@pytest.mark.parametrize("kind", ["host", "device"]) +def test_overlapping_destinations_land_in_descriptor_order(pack, kind, tail): + """The later copy wins on the overlap, as it would eagerly.""" + # Host copies share the 64 KiB staging buffer; device copies are bounded only by memory, and a + # long first copy lets the warps drift apart. + n = 24 * 1024 if kind == "host" else 1 << 20 + m = 8 * 1024 + first = torch.full((n,), 1, dtype=torch.uint8) + second = torch.full((m,), 2, dtype=torch.uint8) + if kind == "host": + first, second = first.pin_memory(), second.pin_memory() + else: + first, second = first.cuda(), second.cuda() + start = n - tail + expected = torch.ones(start + m, dtype=torch.uint8) + expected[start:] = 2 + + mismatches = [] + for trial in range(50): + buf = torch.zeros(start + m, dtype=torch.uint8, device="cuda") + step_inputs.begin() + step_inputs.copy_(buf[:n], first) + step_inputs.copy_(buf[start:], second) + step_inputs.flush() + torch.cuda.synchronize() + wrong = int((buf.cpu() != expected).sum()) + if wrong: + mismatches.append((trial, wrong)) + + assert not mismatches, f"(trial, overlap bytes where the earlier copy won): {mismatches}" + + +@pytest.mark.parametrize("tail", _TAILS) +def test_a_device_copy_reads_the_destination_of_an_earlier_device_copy(pack, tail): + """A chain of device copies: the second reads the bytes the first writes last.""" + n = 1 << 20 + source = torch.randint(1, 256, (n,), dtype=torch.uint8, device="cuda") + expected = source[n - tail :].clone() + + mismatches = [] + for trial in range(50): + middle = torch.zeros(n, dtype=torch.uint8, device="cuda") + out = torch.zeros(tail, dtype=torch.uint8, device="cuda") + step_inputs.begin() + step_inputs.copy_(middle, source) + step_inputs.copy_(out, middle[n - tail :]) + step_inputs.flush() + torch.cuda.synchronize() + wrong = int((out != expected).sum()) + if wrong: + mismatches.append((trial, wrong)) + + assert not mismatches, ( + f"(trial, stale bytes read before the first copy wrote them): {mismatches}" + ) + + +# --------------------------------------------------------------------------- +# Staged tensors stay alive +# --------------------------------------------------------------------------- + + +def test_a_temporary_device_source_outlives_its_copy(pack): + """The pack holds the source until the scatter is queued, so its memory is not reused.""" + numel = 1 << 16 + expected = torch.arange(numel, dtype=torch.int32) * 3 + 1 + dst = _zeros(numel) + + step_inputs.begin() + # The only reference to this source is the one the pack keeps. + step_inputs.copy_(dst, torch.arange(numel, dtype=torch.int32, device="cuda") * 3 + 1) + assert len(pack.staged) == 2 + # Churn the allocator with same-sized blocks: a freed source block would be handed out here. + garbage = [torch.full((numel,), -7, dtype=torch.int32, device="cuda") for _ in range(16)] + torch.cuda.synchronize() + step_inputs.flush() + del garbage + # Blocks allocated after the flush are ordered behind the scatter on this stream. + after = [torch.full((numel,), -9, dtype=torch.int32, device="cuda") for _ in range(16)] + torch.cuda.synchronize() + + assert torch.equal(dst.cpu(), expected) + del after + + +def test_a_temporary_destination_view_outlives_its_copy(pack): + owner = _zeros(8) + + step_inputs.begin() + step_inputs.copy_(owner.view(2, 4)[1], _host([5, 6, 7, 8])) + assert len(pack.staged) == 1 + step_inputs.flush() + torch.cuda.synchronize() + + assert owner.tolist() == [0, 0, 0, 0, 5, 6, 7, 8] + + +# --------------------------------------------------------------------------- +# Fallbacks +# --------------------------------------------------------------------------- + + +def test_an_unpackable_copy_flushes_first_and_keeps_order(pack): + buf = _zeros(8) + other = _zeros(4) + later = _zeros(2) + + step_inputs.begin() + step_inputs.copy_(buf, _host([1] * 8)) + step_inputs.copy_(other, _host([3, 3, 3, 3])) + # A strided 1-D destination cannot be packed: the pack lands first, then this copy. + step_inputs.copy_(buf[::2], _host([2, 2, 2, 2])) + assert step_inputs.is_open() + torch.cuda.synchronize() + assert buf.tolist() == [2, 1, 2, 1, 2, 1, 2, 1] + assert other.tolist() == [3, 3, 3, 3] + + # The pack reopened: later copies are deferred again. + step_inputs.copy_(later, _host([5, 6])) + torch.cuda.synchronize() + assert later.tolist() == [0, 0] + step_inputs.flush() + torch.cuda.synchronize() + assert later.tolist() == [5, 6] + + +def test_a_device_copy_with_a_dtype_change_is_not_packed(pack): + dst = _zeros(3, dtype=torch.int64) + src = torch.tensor([4, 5, 6], dtype=torch.int32, device="cuda") + + step_inputs.begin() + step_inputs.copy_(dst, src) + torch.cuda.synchronize() + + assert dst.tolist() == [4, 5, 6] + assert step_inputs.is_open() + + +def test_a_copy_larger_than_the_staging_buffer_falls_back_in_order(pack): + numel = StepInputPack.CAPACITY // 4 + big = _zeros(numel) + big_values = torch.arange(numel, dtype=torch.int32).pin_memory() + + step_inputs.begin() + step_inputs.copy_(big[:4], _host([-1, -1, -1, -1])) + step_inputs.copy_(big, big_values) + torch.cuda.synchronize() + + # The staged prefix landed before the full copy, which therefore wins. + assert torch.equal(big.cpu(), big_values) + assert step_inputs.is_open() + + +def test_copies_that_together_overflow_the_staging_buffer_all_land(pack): + # Each copy is a quarter of the staging buffer, so the fourth one does not fit. + chunk = StepInputPack.CAPACITY // 16 + dsts = [_zeros(chunk) for _ in range(6)] + values = [torch.full((chunk,), i + 1, dtype=torch.int32).pin_memory() for i in range(6)] + + step_inputs.begin() + for dst, value in zip(dsts, values): + step_inputs.copy_(dst, value) + step_inputs.flush() + torch.cuda.synchronize() + + for i, dst in enumerate(dsts): + assert torch.equal(dst.cpu(), values[i]), f"copy {i} did not land" + + +def test_more_copies_than_descriptors_all_land(pack): + count = StepInputPack.MAX_DESC + 40 + dst = _zeros(count) + + step_inputs.begin() + for i in range(count): + step_inputs.copy_(dst[i : i + 1], _host([i + 1])) + step_inputs.flush() + torch.cuda.synchronize() + + assert dst.tolist() == list(range(1, count + 1)) + + +def test_a_row_strided_copy_that_needs_more_descriptors_than_remain_falls_back(pack): + rows = 8 + table = torch.full((rows, 4), -1, dtype=torch.int32, device="cuda") + filler = _zeros(StepInputPack.MAX_DESC - 2) + rows_values = torch.arange(rows * 3, dtype=torch.int32).reshape(rows, 3) + + step_inputs.begin() + for i in range(filler.numel()): + step_inputs.copy_(filler[i : i + 1], _host([i + 1])) + # Two descriptors remain; the slice needs one per row, so it flushes and copies directly. + step_inputs.copy_(table[:, :3], rows_values.pin_memory()) + step_inputs.flush() + torch.cuda.synchronize() + + assert filler.tolist() == list(range(1, filler.numel() + 1)) + assert torch.equal(table[:, :3].cpu(), rows_values) + assert table[:, 3].tolist() == [-1] * rows + + +def test_mixed_device_copies_past_the_descriptor_limit_all_land(pack): + count = StepInputPack.MAX_DESC + 10 + src = torch.arange(1, count + 1, dtype=torch.int32, device="cuda") + dst = _zeros(count) + + step_inputs.begin() + for i in range(count): + step_inputs.copy_(dst[i : i + 1], src[i : i + 1]) + step_inputs.flush() + torch.cuda.synchronize() + + assert torch.equal(dst, src) + + +def test_begin_on_an_open_pack_issues_it_first(pack): + dst = _zeros(2) + + step_inputs.begin() + step_inputs.copy_(dst, _host([8, 9])) + step_inputs.begin() + torch.cuda.synchronize() + + assert dst.tolist() == [8, 9] + + +def test_reusing_a_pinned_staging_buffer_waits_for_its_upload(pack): + """With every staging buffer in the ring used once, the next pack reuses the first one. + + Its upload is queued behind a long kernel, so rewriting the buffer without waiting would + upload the newer bytes in place of the older ones. + """ + ring = StepInputPack.RING + dsts = [_zeros(4) for _ in range(ring + 1)] + values = [[i * 10 + j for j in range(4)] for i in range(ring + 1)] + + step_inputs.begin() + torch.cuda._sleep(_SLEEP_CYCLES) + step_inputs.copy_(dsts[0], _host(values[0])) + step_inputs.flush() + first_upload = pack.events[0] + for i in range(1, ring): + step_inputs.begin() + step_inputs.copy_(dsts[i], _host(values[i])) + step_inputs.flush() + assert not first_upload.query(), "the queued kernel finished too early for this check" + + step_inputs.begin() + assert first_upload.query() + step_inputs.copy_(dsts[ring], _host(values[ring])) + step_inputs.flush() + torch.cuda.synchronize() + + for i, dst in enumerate(dsts): + assert dst.tolist() == values[i], f"pack {i} uploaded the wrong bytes" + + +def test_a_copy_while_capturing_is_not_deferred(pack): + dst = _zeros(3) + src = torch.tensor([1, 2, 3], dtype=torch.int32, device="cuda") + + step_inputs.begin() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step_inputs.copy_(dst, src) + step_inputs.flush() + torch.cuda.synchronize() + assert dst.tolist() == [0, 0, 0] + + graph.replay() + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3] + + +# --------------------------------------------------------------------------- +# Discard +# --------------------------------------------------------------------------- + + +def test_discard_drops_the_staged_copies(pack): + dst = _zeros(3) + + step_inputs.begin() + step_inputs.copy_(dst, _host([1, 2, 3])) + step_inputs.copy_(dst, torch.ones(3, dtype=torch.int32, device="cuda")) + step_inputs.discard() + assert not pack.open and pack.staged == [] and pack.h2d == [] and pack.d2d == [] + step_inputs.flush() + torch.cuda.synchronize() + + assert dst.tolist() == [0, 0, 0] + + +def test_after_a_discarded_step_identical_values_are_sent_again(pack): + """The runner's recovery: discard the pack and clear the skip cache.""" + skip_cache = H2DSkipCache() + dst = _zeros(3) + values = _host([4, 5, 6]) + + with h2d_cache.active(skip_cache): + step_inputs.begin() + assert copy_h2d_if_changed(dst, values, pack=True) is True + step_inputs.discard() + torch.cuda.synchronize() + assert dst.tolist() == [0, 0, 0] + # Without clearing, the cache would vouch for values that never landed. + assert copy_h2d_if_changed(dst, values, pack=True) is False + skip_cache.clear() + + step_inputs.begin() + assert copy_h2d_if_changed(dst, values, pack=True) is True + step_inputs.flush() + torch.cuda.synchronize() + + assert dst.tolist() == [4, 5, 6] + + +def test_skip_cache_copies_join_the_active_pack(pack): + skip_cache = H2DSkipCache() + dst = _zeros(3) + values = _host([4, 5, 6]) + + with h2d_cache.active(skip_cache): + step_inputs.begin() + assert copy_h2d_if_changed(dst, values, pack=True) is True + torch.cuda.synchronize() + assert dst.tolist() == [0, 0, 0] + # Already staged with these values: the second copy is skipped. + assert copy_h2d_if_changed(dst, values, pack=True) is False + step_inputs.flush() + torch.cuda.synchronize() + + assert dst.tolist() == [4, 5, 6] + + +# --------------------------------------------------------------------------- +# Isolation between packs +# --------------------------------------------------------------------------- + + +def test_a_nested_pack_does_not_see_the_outer_pack(pack): + outer_dst = _zeros(2) + inner_dst = _zeros(2) + other = StepInputPack(_device()) + + step_inputs.begin() + step_inputs.copy_(outer_dst, _host([1, 2])) + with step_inputs.active(other): + assert not step_inputs.is_open() + step_inputs.copy_(inner_dst, _host([3, 4])) # plain copy: the inner pack is closed + step_inputs.begin() + step_inputs.copy_(inner_dst, _host([5, 6])) + torch.cuda.synchronize() + assert inner_dst.tolist() == [3, 4] + assert outer_dst.tolist() == [0, 0] + step_inputs.flush() + assert step_inputs.is_open() and len(pack.h2d) == 1 + step_inputs.flush() + torch.cuda.synchronize() + + assert inner_dst.tolist() == [5, 6] + assert outer_dst.tolist() == [1, 2] + + +def test_active_none_disables_packing_inside_an_outer_pack(pack): + dst = _zeros(2) + + step_inputs.begin() + with step_inputs.active(None): + assert not step_inputs.is_open() + step_inputs.copy_(dst, _host([7, 8])) + torch.cuda.synchronize() + assert dst.tolist() == [7, 8] + assert pack.h2d == [] + + +def test_an_open_pack_on_one_thread_is_invisible_to_another(pack): + mine = _zeros(2) + theirs = _zeros(2) + seen = {} + + step_inputs.begin() + step_inputs.copy_(mine, _host([1, 1])) + + def other_thread(): + seen["open"] = step_inputs.is_open() + step_inputs.copy_(theirs, _host([2, 2])) + step_inputs.flush() + torch.cuda.synchronize() + seen["theirs"] = theirs.tolist() + seen["mine"] = mine.tolist() + + thread = threading.Thread(target=other_thread) + thread.start() + thread.join() + + assert seen == {"open": False, "theirs": [2, 2], "mine": [0, 0]} + assert step_inputs.is_open() and len(pack.h2d) == 1 + step_inputs.flush() + torch.cuda.synchronize() + assert mine.tolist() == [1, 1] + + +def test_threads_with_their_own_packs_keep_them_apart(): + packs = [StepInputPack(_device()) for _ in range(2)] + dsts = [_zeros(4) for _ in range(2)] + barrier = threading.Barrier(2) + results = {} + + def worker(index): + with step_inputs.active(packs[index]): + step_inputs.begin() + step_inputs.copy_(dsts[index], _host([index + 1] * 4)) + barrier.wait() + # Both packs are open with one staged copy each. + results[index] = (packs[index].open, len(packs[index].h2d)) + barrier.wait() + step_inputs.flush() + torch.cuda.synchronize() + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert results == {0: (True, 1), 1: (True, 1)} + assert dsts[0].tolist() == [1] * 4 + assert dsts[1].tolist() == [2] * 4 diff --git a/tests/unittest/_torch/kernels/parallel_hw_agnostic/test_custom_ops.py b/tests/unittest/_torch/kernels/parallel_hw_agnostic/test_custom_ops.py index 51d10786cd08..c39ee3ada042 100644 --- a/tests/unittest/_torch/kernels/parallel_hw_agnostic/test_custom_ops.py +++ b/tests/unittest/_torch/kernels/parallel_hw_agnostic/test_custom_ops.py @@ -18,6 +18,10 @@ import torch._library.utils as library_utils import tensorrt_llm # noqa: F401 + +# The MiniMax-M3 ops' fake impls are registered by their modeling_v2 package. +import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.fake_ops # noqa: F401 +import tensorrt_llm._torch.pyexecutor.step_inputs # noqa: F401 from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops import ( _cpp_fakes, # noqa: F401 ) 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 index d3ac1acda903..980cd15a09ef 100644 --- 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 @@ -7,6 +7,7 @@ need the 4-node launch its docstring describes. """ +__extra_import_path__ = ["."] import _glm_rank_job import pytest import torch 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 index b536ec7d560e..af00750b1940 100644 --- 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 @@ -7,6 +7,7 @@ 16-rank run over four nodes. """ +__extra_import_path__ = ["."] import _draft_step_rank_job import pytest 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 index faf00183c753..d784f9480daf 100644 --- 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 @@ -7,6 +7,7 @@ 16-rank run over four nodes. """ +__extra_import_path__ = ["."] import _draft_step_rank_job import pytest diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/activation/test_modeling_v2_m3_silu_and_mul.py new file mode 100644 index 000000000000..ea8d241d40c9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/activation/test_modeling_v2_m3_silu_and_mul.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_silu_and_mul entry of the MiniMax-M3 tp4 target: SwiGLU-OAI and plain SwiGLU.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.activation.m3_silu_and_mul import ( # noqa: E501 + m3_silu_and_mul as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_silu_and_mul requires a CUDA device" + +OAI = (7.0, 1.702, 1.0) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + # Gate and up spread past the clamp (|x| up to ~12) so both clamps engage. + x = (torch.randn(spec["m"], 2 * spec["inter"], generator=g, device="cuda") * 4.0).to( + torch.bfloat16 + ) + return (x, *OAI) if spec["oai"] else (x, None, None, None) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """fp32 arithmetic and one bf16 rounding on both sides: the kernel's fast sigmoid moves the fp32 value by a few + fp32 ulps, which flips the bf16 rounding of an element by at most one step (1 unit of ULP[bf16] of the row); + measured on GB300: worst element 0.5 unit, relative RMS 0.01.""" + args = _build(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(*args) + assert out.shape == (cell.spec["m"], cell.spec["inter"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=1.0, rms_ulp=0.1) + + +def test_the_gate_sees_a_missing_clamp() -> None: + """Dropping the SwiGLU-OAI clamps moves every element whose gate or up passed the limit far outside the band.""" + spec = op.CELLS[0].spec + x, *oai = _build(dict(spec, m=64), 3) + out = op(x, *oai) + with pytest.raises(AssertionError): + assert_within_ulp(out, op.reference(x, None, oai[1], oai[2]), element_ulp=1.0, rms_ulp=0.1) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/_m3_attn_runtime_tp4.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/_m3_attn_runtime_tp4.py new file mode 100644 index 000000000000..7bfd8de384f9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/_m3_attn_runtime_tp4.py @@ -0,0 +1,274 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real MiniMax-M3 attention runtime state at the TP4 per-rank geometry, for the backend-step entries' GPU tests. + +One hybrid NVFP4 KV cache manager (``MiniMaxM3KVCacheManagerV2``) laid out as the target's: a dense layer in the FP8 +half, a sparse MSA layer in the NVFP4 half with its FP8 index-K cache, and (``draft=True``) the one-model Eagle3 +drafter's shared layer, read through the manager's FP8 draft view as the runtime's draft swap reads it. Metadata is +the runtime's own type, filled by its own ``prepare()``; the backends are built the way the target builds them. + +The geometry is one rank's: 16 query heads over one KV head and one index head of dim 128. The mapping is a single +rank holding exactly those heads, so nothing is sharded again. + +History rows are written by the in-tree scatters (``fused_write_layer_caches`` for the FP8 layers, +``fused_write_layer_caches_nvfp4`` for the NVFP4 layer and its index-K), at the slots the manager assigned, which are +not in request order. +""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import ( + MiniMaxM3KVCacheManagerV2, + MiniMaxM3MsaSparseAttention, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches, + fused_write_layer_caches_nvfp4, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.llmapi.llm_args import ( + Eagle3DecodingConfig, + KvCacheConfig, + MiniMaxM3SparseAttentionConfig, +) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +HEADS = 16 +KV_HEADS = 1 +INDEX_HEADS = 1 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +DENSE, SPARSE, DRAFT = 0, 1, 2 +MAX_BATCH = 8 +MAX_SEQ_LEN = 8192 +MAX_NUM_TOKENS = 8192 +#: The Eagle3 drafter checkpoint's RoPE: NeoX over the whole head, theta 5e6. +DRAFT_ROPE_THETA = 5.0e6 + +#: E4M3-exact planted index levels, 1/8 apart in relative terms (0.25 .. 15): 48 distinct block scores. +LEVELS = tuple(2.0**e * (1 + m / 8) for e in range(-2, 4) for m in range(8)) +#: The index value of a block's non-peak rows; its score (8) is below every level's (>= 16). +LOW = 0.125 + + +def mapping() -> Mapping: + return Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + + +def sparse_attention_config() -> MiniMaxM3SparseAttentionConfig: + """The target's runtime: MSA, the FP8 indexer, the fused five-way projection (M3's checkpoint selection rule: + no forced initial blocks, the query's own block forced).""" + return MiniMaxM3SparseAttentionConfig( + implementation="msa", + indexer_kv_dtype="fp8", + fuse_qkv_index_projection=True, + sparse_num_index_heads=INDEX_HEADS, + num_attention_heads=HEADS, + num_key_value_heads=KV_HEADS, + ) + + +def make_manager(*, draft: bool) -> MiniMaxM3KVCacheManagerV2: + torch.cuda.init() + return MiniMaxM3KVCacheManagerV2( + KvCacheConfig( + max_tokens=4 * MAX_SEQ_LEN, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + host_cache_size=0, + event_buffer_max_size=0, + ), + CacheType.SELF, + num_layers=2, + num_kv_heads=[KV_HEADS] * (3 if draft else 2), + head_dim=HEAD_DIM, + tokens_per_block=PAGE, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + mapping=mapping(), + dtype=DataType.NVFP4, + vocab_size=16, + spec_config=( + Eagle3DecodingConfig(max_draft_len=3, speculative_model="unused") if draft else None + ), + sparse_layer_ids=[SPARSE], + sparse_attention_config=sparse_attention_config(), + ) + + +def make_msa_backend(layer: int) -> MiniMaxM3MsaSparseAttention: + """One layer's MSA backend as ``MiniMaxM3Tp4ModelingV2`` builds it (NVFP4 KV quant config on every layer).""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.NVFP4) + backend = create_attention( + "TRTLLM", + layer, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=None, + quant_config=quant, + q_scaling=1.0, + sparse_params=sparse_attention_config().to_sparse_params(layer_idx=layer), + ) + backend.update_quant_config(quant) + assert isinstance(backend, MiniMaxM3MsaSparseAttention), type(backend).__name__ + return backend + + +def make_draft_backend() -> TrtllmAttention: + """The drafter layer's attention as ``M3Eagle3Draft`` builds it: fused NeoX RoPE, an FP8 KV cache (the target + overrides the drafter's KV dtype to FP8), no other quantization.""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + backend = create_attention( + "TRTLLM", + DRAFT, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=RopeParams(dim=HEAD_DIM, theta=DRAFT_ROPE_THETA, max_positions=MAX_SEQ_LEN), + ), + quant_config=quant, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + backend.update_quant_config(quant) + assert type(backend) is TrtllmAttention, type(backend).__name__ + return backend + + +def make_msa_metadata(manager) -> MiniMaxM3MsaSparseAttentionMetadata: + metadata = MiniMaxM3MsaSparseAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + max_num_sequences=MAX_BATCH, + kv_cache_manager=manager, + mapping=mapping(), + sparse_metadata_params=sparse_attention_config().to_sparse_metadata_params(), + ) + # The target's MSA preparation choice that changes what prepare() stages (the identity tables of short steps). + metadata.msa_policy = MsaRuntimePolicy(short_sequence_identity=True) + return metadata + + +def make_draft_metadata(manager) -> TrtllmAttentionMetadata: + """Drafter metadata over the shared manager's FP8 draft view, the manager the runtime's draft swap installs.""" + view = manager.get_draft_kv_cache_view() + assert view is not None and view.dtype == DataType.FP8 + return TrtllmAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + kv_cache_manager=view, + mapping=mapping(), + ) + + +def allocate(manager, request_ids: Sequence[int], qlens: Sequence[int], cached: Sequence[int]): + manager.add_dummy_requests(list(request_ids), token_nums=[c + n for c, n in zip(cached, qlens)]) + + +def prepare(metadata, request_ids, qlens, cached, num_contexts: int) -> None: + metadata.seq_lens = torch.tensor(list(qlens), dtype=torch.int) + metadata.num_contexts = num_contexts + metadata.request_ids = list(request_ids) + metadata.prompt_lens = [ + c + n if row < num_contexts else c for row, (c, n) in enumerate(zip(cached, qlens)) + ] + metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=list(cached)) + metadata.prepare() + + +def flat_slots(manager, layer: int, request_id: int, positions: torch.Tensor) -> torch.Tensor: + """The cache slots (page slot id * PAGE + offset) of a request's positions in ``layer``.""" + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0], + dtype=torch.int64, + device="cuda", + ) + positions = positions.to(device="cuda", dtype=torch.int64) + return (pages[positions // PAGE] * PAGE + positions % PAGE).to(torch.int32) + + +def write_fp8_rows(manager, layer, request_id, positions, k, v) -> None: + """K / V rows [n, 128] (E4M3) into an FP8 layer's cache at ``positions``.""" + view = manager.get_buffers(layer, "HND") + assert view.dtype == torch.float8_e4m3fn + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches(view[:, 0], view[:, 1], None, slots, k, v, None) + + +def write_nvfp4_rows(manager, metadata, layer, request_id, positions, k, v, idx_k, quant) -> None: + """BF16 K / V rows quantized into the NVFP4 layer's cache, and index-K rows into its FP8 index cache.""" + data = manager.get_buffers(layer, "HND") + scales = manager.get_block_scale_buffers(layer, "HND") + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches_nvfp4( + data[:, 0], + data[:, 1], + scales[:, 0], + scales[:, 1], + metadata.msa_idx_k_cache(layer), + slots, + k, + v, + idx_k, + quant, + ) + + +def index_directions(gen: torch.Generator) -> tuple[torch.Tensor, torch.Tensor]: + """Two orthogonal ±1 index directions (the halves of the 128 dims), E4M3-exact.""" + signs = torch.randint(0, 2, (HEAD_DIM,), generator=gen, device="cuda").float() * 2 - 1 + u0 = torch.zeros(HEAD_DIM, device="cuda") + u1 = torch.zeros(HEAD_DIM, device="cuda") + u0[: HEAD_DIM // 2] = signs[: HEAD_DIM // 2] + u1[HEAD_DIM // 2 :] = signs[HEAD_DIM // 2 :] + return u0, u1 + + +def planted_index_keys(length: int, u0, u1, gen: torch.Generator) -> torch.Tensor: + """Index-K rows [length, 128] (bf16, E4M3-exact) whose per-block maximum score is unambiguous. + + Every block has one peak row at a random position, ``a_b u0 + c_b u1`` with ``a`` and ``c`` two independent + permutations of ``LEVELS``; its other rows are ``LOW (u0 + u1)``. A query along ``u0`` then ranks the blocks by + ``a``, one along ``u1`` by ``c``, every score 1/8 or more apart from the next.""" + blocks = (length + PAGE - 1) // PAGE + assert blocks <= len(LEVELS) + levels = torch.tensor(LEVELS, device="cuda") + a = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + c = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + rows = (LOW * (u0 + u1)).expand(length, HEAD_DIM).clone() + for b in range(blocks): + span = min(PAGE, length - b * PAGE) + peak = b * PAGE + int(torch.randint(0, span, (1,), generator=gen, device="cuda")) + rows[peak] = a[b] * u0 + c[b] * u1 + return rows.to(torch.bfloat16) + + +def nvfp4_quant(*values: torch.Tensor) -> torch.Tensor: + """The (Q, K, V) quantization scales ``fused_write_layer_caches_nvfp4`` takes: each role's amax at the top of the + E2M1 x E4M3 range.""" + quant = torch.ones(3, dtype=torch.float32, device="cuda") + for role, x in enumerate(values, start=1): + quant[role] = (448.0 * 6.0) / x.abs().max().float() + return quant diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/_m3_fused_attn_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/_m3_fused_attn_step.py new file mode 100644 index 000000000000..85eddd0cddc4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/_m3_fused_attn_step.py @@ -0,0 +1,255 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A short-sequence decode step of one MiniMax-M3 TP4 sparse layer, for the fused attention entries' tests +(``m3_fused_decode_attn``, ``m3_attn_mega``, ``m3_attn_c1``). + +The history is written into the paged NVFP4 pool by the in-tree scatter (``fused_write_layer_caches_nvfp4``, which +quantizes with ``torch.ops.trtllm.fp4_quantize``) at the same ``kv_quant_scale`` the step inserts with. The new +tokens' slots and every row past a request's length hold random bytes, so a kernel that skipped the insert or +read past the length is caught. Pages are shuffled so a kernel that ignored the block table reads the wrong ones. +""" + +from __future__ import annotations + +import math +from types import SimpleNamespace +from typing import Sequence + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_fused_decode_attn import ( # noqa: E501 + attend_cached, + nvfp4_dequantize, + query_rows, + v_scale_rows, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches_nvfp4, +) + +NUM_HEADS = 16 +HEAD_DIM = 128 +PAGE = 128 +PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM +HIDDEN = 6144 +#: The MXFP8 hand-off buffers' rows (the target's MERGE_MXFP8 buffer). +MXFP8_ROWS = 128 +#: NVFP4 second-level scales: block scales stay inside E4M3 for |x| up to 16 (K) and 8 (V, N(0, 1) rows). +K_QUANT = 448.0 * 6.0 / 16.0 +V_QUANT = 448.0 * 6.0 / 8.0 +ROPE_THETA = 5.0e6 +MAX_POSITION = 2304 + + +def rotary_cos_sin(device) -> torch.Tensor: + """fp32 [max_pos, 2, 32]: NeoX cos / sin of the 64 rotary dims at every position.""" + inv_freq = 1.0 / ( + ROPE_THETA ** (torch.arange(0, 64, 2, dtype=torch.float64, device=device) / 64) + ) + ang = ( + torch.arange(MAX_POSITION, dtype=torch.float64, device=device)[:, None] * inv_freq[None, :] + ) + return torch.stack([ang.cos(), ang.sin()], dim=1).float().contiguous() + + +def build_step( + seq_lens: Sequence[int], qlen: int, seed: int, packed: torch.Tensor = None +) -> SimpleNamespace: + """Caches, metadata and weights of one step; ``packed`` (bf16 [T, 2560]) defaults to N(0, 1) rows.""" + dev = torch.device("cuda") + g = torch.Generator(device=dev).manual_seed(seed) + batch = len(seq_lens) + t_total = batch * qlen + max_blocks = max(math.ceil(s / PAGE) for s in seq_lens) + num_pages = batch * (max_blocks + 1) + 3 + # One spare column: the kernel must read the columns the length needs, not the table's width. + block_table = torch.randperm(num_pages, device=dev, generator=g)[: batch * (max_blocks + 1)] + block_table = block_table.to(torch.int32).reshape(batch, max_blocks + 1).contiguous() + kv_data = torch.randint( + 0, 256, (num_pages, 2, 1, PAGE, HEAD_DIM // 2), dtype=torch.uint8, device=dev, generator=g + ) + kv_scale = torch.randint( + 0, 256, (num_pages, 2, 1, PAGE, HEAD_DIM // 16), dtype=torch.uint8, device=dev, generator=g + ) + idx_k = torch.randint( + 0, 256, (num_pages, 1, PAGE, HEAD_DIM), dtype=torch.uint8, device=dev, generator=g + ).view(torch.float8_e4m3fn) + kv_quant_scale = torch.tensor([1.0, K_QUANT, V_QUANT], dtype=torch.float32, device=dev) + + hist_slots = [] + for r, n in enumerate(seq_lens): + for pos in range(n - qlen): + hist_slots.append(int(block_table[r, pos // PAGE]) * PAGE + pos % PAGE) + if hist_slots: + h = len(hist_slots) + k_hist = (torch.randn(h, HEAD_DIM, device=dev, generator=g) * 1.5).to(torch.bfloat16) + v_hist = torch.randn(h, HEAD_DIM, device=dev, generator=g).to(torch.bfloat16) + i_hist = torch.randn(h, HEAD_DIM, device=dev, generator=g).to(torch.bfloat16) + assert fused_write_layer_caches_nvfp4( + kv_data[:, 0], + kv_data[:, 1], + kv_scale[:, 0], + kv_scale[:, 1], + idx_k, + torch.tensor(hist_slots, dtype=torch.int32, device=dev), + k_hist, + v_hist, + i_hist, + kv_quant_scale, + ) + + positions, slots = [], [] + for r, n in enumerate(seq_lens): + for t in range(qlen): + pos = n - qlen + t + positions.append(pos) + slots.append(int(block_table[r, pos // PAGE]) * PAGE + pos % PAGE) + if packed is None: + packed = torch.randn(t_total, PACKED_WIDTH, device=dev, generator=g).to(torch.bfloat16) + + def norm_weight(): + return (torch.randn(HEAD_DIM, device=dev, generator=g) * 0.1).to(torch.bfloat16) + + return SimpleNamespace( + packed=packed, + kv_data=kv_data, + kv_scale=kv_scale, + idx_k=idx_k, + out_cache_loc=torch.tensor(slots, dtype=torch.int32, device=dev), + position_ids=torch.tensor(positions, dtype=torch.int32, device=dev), + block_table=block_table, + seq_lens=torch.tensor(list(seq_lens), dtype=torch.int32, device=dev), + kv_quant_scale=kv_quant_scale, + k_global_scale=torch.tensor([1.0 / K_QUANT], dtype=torch.float32, device=dev), + v_global_scale=torch.tensor([1.0 / V_QUANT], dtype=torch.float32, device=dev), + q_norm_weight=norm_weight(), + k_norm_weight=norm_weight(), + index_k_norm_weight=norm_weight(), + rotary_cos_sin=rotary_cos_sin(dev), + eps=1e-6, + sm_scale=HEAD_DIM**-0.5, + decode_query_len=qlen, + num_tokens=t_total, + ) + + +def attention_args(s: SimpleNamespace) -> tuple: + """The attention arguments every entry takes, in schema order (``kv_data`` .. ``rotary_cos_sin``, the two + scalars and ``decode_query_len``).""" + return ( + s.kv_data, + s.kv_scale, + s.idx_k, + s.out_cache_loc, + s.position_ids, + s.block_table, + s.seq_lens, + s.kv_quant_scale, + s.k_global_scale, + s.v_global_scale, + s.q_norm_weight, + s.k_norm_weight, + s.index_k_norm_weight, + s.rotary_cos_sin, + s.eps, + s.sm_scale, + s.decode_query_len, + ) + + +def swizzle(rows_exp: torch.Tensor, padded_rows: int) -> torch.Tensor: + """[rows, k_blocks] UE8M0 bytes -> the 128x4-swizzled buffer of ``padded_rows`` rows (filler 127).""" + rows, kb = rows_exp.shape + k_tiles = (kb + 3) // 4 + out = torch.full((padded_rows * k_tiles * 4,), 127, dtype=torch.uint8, device=rows_exp.device) + r = torch.arange(rows, device=rows_exp.device)[:, None] + k = torch.arange(kb, device=rows_exp.device)[None, :] + off = ( + (r // 128) * k_tiles * 512 + (k // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + k % 4 + ) + out[off.reshape(-1)] = rows_exp.reshape(-1) + return out + + +def mxfp8_operand( + rows: int, cols: int, exp_lo: int, exp_hi: int, g: torch.Generator, padded_rows: int +): + """A random MXFP8 operand: E4M3 values from N(0, 1) and UE8M0 exponents in [exp_lo, exp_hi], swizzled.""" + dev = torch.device("cuda") + q = torch.randn(rows, cols, device=dev, generator=g).to(torch.float8_e4m3fn) + ex = torch.randint( + exp_lo, exp_hi + 1, (rows, cols // 32), dtype=torch.uint8, device=dev, generator=g + ) + return q, swizzle(ex, padded_rows) + + +def new_slot_mask(s: SimpleNamespace, num_pages: int) -> torch.Tensor: + """bool [pages, 128]: the rows the step inserts.""" + mask = torch.zeros(num_pages, PAGE, dtype=torch.bool, device=s.kv_data.device) + for slot in s.out_cache_loc.tolist(): + mask[slot // PAGE, slot % PAGE] = True + return mask + + +# Band of the attention output, in bf16 ulps (2^-8) of each [128] head row's largest magnitude, and of the relative +# RMS, against the float64 decode over the caches the op wrote. Both sides read the same operands (E4M3 q, E2M1 x +# E4M3 K/V, all exact in fp16), so the gap is the bf16 rounding on both sides (half an ulp each), the fp16 +# probabilities before PV (2^-11 relative), fp32 accumulation and approximate exp2/log2 in the split-K partials and +# their LSE merge. Measured over 46 steps: at most 2.0 ulps per element, 0.2 ulp RMS. +ATTN_ELEMENT_ULP = 3.0 +ATTN_RMS_ULP = 0.5 +# Inserted rows against the reference's NVFP4 / E4M3 recipes: the kernel's RMSNorm sums the squares in another order +# and uses approximate rsqrt / reciprocals, so an element on a rounding tie can land one code apart. Measured over +# 46 steps: no differing scale byte, at most 3 differing value codes in a step (1 in 300 of the inserted elements). +# Such a code moves one key's (or value's) element by an E2M1 step, up to 44 ulps of an output row, which is why +# the output is gated against the decode over the op's own caches rather than the reference's. +CODE_MISMATCH_MIN = 4 +CODE_MISMATCH_FRACTION = 0.005 + + +def check_insert(s: SimpleNamespace, before: tuple, ref) -> None: + """Every row the step does not insert is untouched; inserted rows match the reference's recipe.""" + kv_data0, kv_scale0, idx0 = before + pages = s.kv_data.shape[0] + mask = new_slot_mask(s, pages) + keep = ~mask + for plane in (0, 1): + assert torch.equal(s.kv_data[:, plane, 0][keep], kv_data0[:, plane, 0][keep]), ( + "a row past the insert moved" + ) + assert torch.equal(s.kv_scale[:, 0, 0][keep], kv_scale0[:, 0, 0][keep]) + assert torch.equal(s.idx_k.view(torch.uint8)[:, 0][keep], idx0.view(torch.uint8)[:, 0][keep]) + v_now = torch.stack([v_scale_rows(s.kv_scale[p, 1, 0]) for p in range(pages)]) + v_was = torch.stack([v_scale_rows(kv_scale0[p, 1, 0]) for p in range(pages)]) + v_ref = torch.stack([v_scale_rows(ref.kv_scale[p, 1, 0]) for p in range(pages)]) + assert torch.equal(v_now[keep], v_was[keep]), "a V scale outside the inserted rows moved" + assert torch.equal(s.kv_scale[:, 0, 0][mask], ref.kv_scale[:, 0, 0][mask]), "K block scales" + assert torch.equal(v_now[mask], v_ref[mask]), "V block scales (4x4 token-quad order)" + n = int(mask.sum()) * HEAD_DIM + for plane, scales in ((0, s.kv_scale[:, 0, 0][mask]), (1, v_now[mask])): + got = nvfp4_dequantize(s.kv_data[:, plane, 0][mask], scales) + want = nvfp4_dequantize(ref.kv_data[:, plane, 0][mask], scales) + assert int((got != want).sum()) <= max(CODE_MISMATCH_MIN, CODE_MISMATCH_FRACTION * n), ( + f"plane {plane} E2M1 codes" + ) + got = s.idx_k.view(torch.uint8)[:, 0][mask] + want = ref.idx_k.view(torch.uint8)[:, 0][mask] + assert int((got != want).sum()) <= max(CODE_MISMATCH_MIN, CODE_MISMATCH_FRACTION * n), ( + "index-K E4M3 codes" + ) + + +def decode_over_written_caches(s: SimpleNamespace, packed: torch.Tensor) -> torch.Tensor: + """The reference's decode (bf16 [T, 2048]) of ``packed``'s queries over the caches as the op left them.""" + q = query_rows(packed, s.position_ids, s.q_norm_weight, s.eps, s.rotary_cos_sin) + return attend_cached( + q, + s.kv_data, + s.kv_scale, + s.block_table, + s.seq_lens, + s.k_global_scale, + s.v_global_scale, + s.sm_scale, + s.decode_query_len, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/_m3_nvfp4_cache.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/_m3_nvfp4_cache.py new file mode 100644 index 000000000000..2254f6c0b68c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/_m3_nvfp4_cache.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A MiniMax-M3 TP4 sparse-layer decode step over a packed NVFP4 cache, for the attention entries' tests. + +The cache is written by the in-tree scatter (``fused_write_layer_caches_nvfp4``, +which quantizes with ``torch.ops.trtllm.fp4_quantize``), into one +``[pages, 2 (K, V), 1, 128, .]`` pool per role exactly as ``msa_paged_kv`` and +``get_block_scale_buffers`` hand the K and V views over. Pages are shuffled so +a kernel that ignored the block table would read the wrong ones. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches_nvfp4, +) + +NUM_HEADS = 16 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 + + +def token_kv_len(seq_len: int, decode_query_len: int, offset: int) -> int: + """The KV extent query token ``offset`` of a request attends, as every decode kernel derives it.""" + return max(seq_len - decode_query_len + offset + 1, 0) + + +def build_step( + seq_lens: Sequence[int], decode_query_len: int, *, identity: bool, seed: int +) -> SimpleNamespace: + """One decode step: E4M3 q [T, 16, 128], the packed NVFP4 K/V views and their scales, a [1, T, 16] + head-major top-k table (the identity table, or a random ascending selection with a -1 tail), the block + table and per-request lengths.""" + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + total_q = batch * decode_query_len + max_blocks = max(1, max((s + PAGE - 1) // PAGE for s in seq_lens)) + num_pages = batch * max_blocks + num_slots = num_pages * PAGE + block_table = ( + torch.randperm(num_pages, device="cuda", generator=g) + .to(torch.int32) + .reshape(batch, max_blocks) + ) + data_pool = torch.zeros(num_pages, 2, 1, PAGE, HEAD_DIM // 2, dtype=torch.uint8, device="cuda") + scale_pool = torch.zeros( + num_pages, 2, 1, PAGE, HEAD_DIM // 16, dtype=torch.uint8, device="cuda" + ) + quant = torch.ones(3, dtype=torch.float32, device="cuda") + values = [] + for role in (0, 1): + x = torch.randn(num_slots, HEAD_DIM, device="cuda", generator=g, dtype=torch.float32).to( + torch.bfloat16 + ) + quant[role + 1] = (448.0 * 6.0) / x.abs().max().float().item() + values.append(x) + assert fused_write_layer_caches_nvfp4( + data_pool[:, 0], + data_pool[:, 1], + scale_pool[:, 0], + scale_pool[:, 1], + None, + torch.arange(num_slots, device="cuda", dtype=torch.int32), + values[0], + values[1], + None, + quant, + ) + q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, device="cuda", generator=g).to( + torch.float8_e4m3fn + ) + topk = torch.full((1, total_q, TOPK), -1, dtype=torch.int32, device="cuda") + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + nblocks = (token_kv_len(seq_lens[req], decode_query_len, offset) + PAGE - 1) // PAGE + if identity: + assert nblocks <= TOPK, "an identity step has every request within 16 blocks" + chosen = torch.arange(nblocks, device="cuda") + else: + chosen = torch.randperm(nblocks, device="cuda", generator=g)[:TOPK].sort().values + topk[0, tok, : chosen.numel()] = chosen.to(torch.int32) + return SimpleNamespace( + q=q, + k_paged=data_pool[:, 0], + v_paged=data_pool[:, 1], + k_block_scale=scale_pool[:, 0], + v_block_scale=scale_pool[:, 1], + k_global_scale=(1.0 / quant[1:2]).contiguous(), + v_global_scale=(1.0 / quant[2:3]).contiguous(), + topk=topk, + block_table=block_table, + seq_lens=torch.tensor(list(seq_lens), dtype=torch.int32, device="cuda"), + decode_query_len=decode_query_len, + sm_scale=HEAD_DIM**-0.5, + source=values, + quant=quant, + num_slots=num_slots, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_c1.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_c1.py new file mode 100644 index 000000000000..202a2e2a8543 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_c1.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_attn_c1 entry (fused decode attention + o_proj GEMV in one 96-CTA launch).""" + +import inspect +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_attn_c1 import ( # noqa: E501 + m3_attn_c1 as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_fused_decode_attn import ( # noqa: E501 + mx_gemm, + mxfp8_rows, + unswizzle_scales, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.m3_attn_mega_c1 import ( # noqa: E501 + MegaWorkspace, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_fused_attn_step import ( # noqa: E402 + ATTN_ELEMENT_ULP, + ATTN_RMS_ULP, + attention_args, + build_step, + check_insert, + decode_over_written_caches, + mxfp8_operand, +) +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_attn_c1 requires a CUDA device" + +QLEN = 6 + +# Attention and the cache insert: the bands of m3_fused_decode_attn (the same kernel code). +# o_proj GEMV, on the op's own MXFP8 hand-off: exact operands, fp32 split-K accumulation over 2048 products against +# an fp64 sum, one bf16 rounding on each side: at most one bf16 step apart, which is 2 of these ulps (2^-8 of the row +# scale) in the row maximum's binade. Measured: at most 1.12 ulps. +_GEMV_ELEMENT_ULP = 2.0 +_GEMV_RMS_ULP = 0.5 + + +def _args(s, ws, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + wo, wosf = mxfp8_operand(6144, 2048, 120, 122, g, 6144) + out = torch.empty(s.num_tokens, 2048, dtype=torch.bfloat16, device="cuda") + y = torch.empty(s.num_tokens, 6144, dtype=torch.bfloat16, device="cuda") + return [ + s.packed, + ws.flags, + *attention_args(s), + 1, + 1, + out, + ws.q8, + ws.sf, + True, + wo, + wosf, + y, + 0, + 0, + 0, + 1, + 8, + 0, + ] + + +def _arg(name): + """Position of a parameter in the op's argument list (``is_valid`` mirrors ``__call__``).""" + return list(inspect.signature(type(op).is_valid).parameters)[1:].index(name) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- the fixture from _validating + certifying(op) + seed = zlib.crc32(cell.why.encode()) % 2**31 + s = build_step([cell.spec["seq_len"]], QLEN, seed=seed) + ws = MegaWorkspace(torch.device("cuda")) + for _ in range(cell.spec["calls"]): + args = _args(s, ws, seed + 1) + out, y = args[_arg("output")], args[_arg("y")] + wo, wosf = args[_arg("o_weight")], args[_arg("o_weight_scale")] + parity = int(ws.flags[32]) & 1 + before = (s.kv_data.clone(), s.kv_scale.clone(), s.idx_k.clone()) + ref = op.reference(*args) + op(*args) + check_insert(s, before, ref) + want = decode_over_written_caches(s, s.packed) + assert_within_ulp(out.view(-1, 128), want.view(-1, 128), ATTN_ELEMENT_ULP, ATTN_RMS_ULP) + # The o_proj units read the attention's MXFP8 hand-off, which is the merge's recipe on the bf16 output. + q8 = ws.q8.view(2, 8, 2048)[parity, :QLEN] + ex = unswizzle_scales(ws.sf.view(2, -1)[parity], QLEN, 64) + q8_want, ex_want = mxfp8_rows(out) + assert torch.equal(q8, q8_want) and torch.equal(ex, ex_want), "MXFP8 hand-off" + y_staged = mx_gemm(q8, ex, wo.view(torch.uint8), unswizzle_scales(wosf, 6144, 64)) + assert_within_ulp(y.view(-1, 128), y_staged.view(-1, 128), _GEMV_ELEMENT_ULP, _GEMV_RMS_ULP) + assert bool((ws.q8.view(2, -1)[parity ^ 1] == 255).all()), ( + "the next call's buffer is at rest" + ) + assert int(ws.flags[32]) & 1 == parity ^ 1, "the call counter advanced" + + +def test_guard_refuses_a_handoff_not_at_rest() -> None: + """The o_proj units poll the hand-off for codes other than 0xFF: stale codes read as this call's.""" + s = build_step([300], QLEN, seed=5) + ws = MegaWorkspace(torch.device("cuda")) + ws.q8.zero_() + args = _args(s, ws, 6) + with pytest.raises(AssertionError, match="not at rest"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_slot_the_block_table_does_not_give() -> None: + s = build_step([300], QLEN, seed=5) + args = _args(s, MegaWorkspace(torch.device("cuda")), 6) + loc = s.out_cache_loc.clone() + loc[2] = (loc[2] + 7) % (s.kv_data.shape[0] * 128) + args[_arg("out_cache_loc")] = loc + with pytest.raises(AssertionError, match="out_cache_loc must be the slot"): + with validating(op): + op(*args) + + +def test_op_rejects_another_query_length() -> None: + """Compiled for one request of 6 query tokens (96 CTAs, one o_proj unit each) only.""" + s = build_step([300], 4, seed=5) + args = _args(s, MegaWorkspace(torch.device("cuda")), 6) + args[_arg("decode_query_len")] = 4 + with pytest.raises(RuntimeError, match="six query tokens"): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_mega.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_mega.py new file mode 100644 index 000000000000..d470e2c52847 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_mega.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_attn_mega entry (QKV GEMV + fused decode attention + o_proj GEMV in one launch).""" + +import inspect +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_attn_mega import ( # noqa: E501 + m3_attn_mega as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_fused_decode_attn import ( # noqa: E501 + mx_gemm, + mxfp8_rows, + reference_attention, + unswizzle_scales, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.m3_attn_mega import ( # noqa: E501 + MegaWorkspace, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_fused_attn_step import ( # noqa: E402 + ATTN_ELEMENT_ULP, + ATTN_RMS_ULP, + attention_args, + build_step, + check_insert, + decode_over_written_caches, + mxfp8_operand, +) +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_attn_mega requires a CUDA device" + +# Each stage is gated against its own reference fed the op's own input to that stage, so a band states one +# stage's error, not the propagation of an earlier stage's rounding through a quantizer. +# +# QKV GEMV and o_proj GEMV: exact MXFP8 operands, fp32 accumulation over 6144 / 2048 products in split-K partials +# against an fp64 sum, one bf16 rounding on each side: at most one bf16 step apart, which is 2 of these ulps (2^-8 +# of the row scale) in the row maximum's binade. Measured: bit-exact but for a few elements (0.43 ulp at most). +_GEMV_ELEMENT_ULP = 2.0 +_GEMV_RMS_ULP = 0.5 +# Attention and the cache insert: the bands of m3_fused_decode_attn (the same kernel code), on the op's own rows. + + +def _operands(seed): + g = torch.Generator(device="cuda").manual_seed(seed) + xq, xsf = mxfp8_operand(1, 6144, 126, 128, g, 128) + w, wsf = mxfp8_operand(2560, 6144, 119, 121, g, 2560) + wo, wosf = mxfp8_operand(6144, 2048, 120, 122, g, 6144) + return xq, xsf, w, wsf, wo, wosf + + +def _args(s, ops, ws): + xq, xsf, w, wsf, wo, wosf = ops + out = torch.empty(1, 2048, dtype=torch.bfloat16, device="cuda") + y = torch.empty(1, 6144, dtype=torch.bfloat16, device="cuda") + return [ + xq, + xsf, + w, + wsf, + ws.qkv_out, + ws.flags, + *attention_args(s), + 1, + 1, + out, + ws.q8, + ws.sf, + True, + wo, + wosf, + y, + 0, + 0, + 0, + 1, + 8, + 0, + ] + + +def _arg(name): + """Position of a parameter in the op's argument list (``is_valid`` mirrors ``__call__``).""" + return list(inspect.signature(type(op).is_valid).parameters)[1:].index(name) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- the fixture from _validating + certifying(op) + seed = zlib.crc32(cell.why.encode()) % 2**31 + s = build_step([cell.spec["seq_len"]], 1, seed=seed) + ops = _operands(seed + 1) + wo, wosf = ops[4], ops[5] + ws = MegaWorkspace(torch.device("cuda")) + for _ in range(cell.spec["calls"]): + args = _args(s, ops, ws) + out, y = args[_arg("output")], args[_arg("y")] + parity = int(ws.flags[32]) & 1 + caches = (s.kv_data.clone(), s.kv_scale.clone(), s.idx_k.clone()) + ref = op.reference(*args) + op(*args) + packed = ws.qkv_out[parity, :1] + assert_within_ulp( + packed.view(-1, 128), ref.packed.view(-1, 128), _GEMV_ELEMENT_ULP, _GEMV_RMS_ULP + ) + check_insert(s, caches, reference_attention(packed, *caches, *attention_args(s)[3:])) + want = decode_over_written_caches(s, packed) + assert_within_ulp(out.view(-1, 128), want.view(-1, 128), ATTN_ELEMENT_ULP, ATTN_RMS_ULP) + # The o_proj units read the attention's MXFP8 hand-off, which is the merge's recipe on the bf16 output. + q8 = ws.q8.view(2, 8, 2048)[parity, :1] + ex = unswizzle_scales(ws.sf.view(2, -1)[parity], 1, 64) + q8_want, ex_want = mxfp8_rows(out) + assert torch.equal(q8, q8_want) and torch.equal(ex, ex_want), "MXFP8 hand-off" + y_staged = mx_gemm(q8, ex, wo.view(torch.uint8), unswizzle_scales(wosf, 6144, 64)) + assert_within_ulp(y.view(-1, 128), y_staged.view(-1, 128), _GEMV_ELEMENT_ULP, _GEMV_RMS_ULP) + # The next call's buffers are back at rest. + assert bool((ws.qkv_out.view(torch.int16)[parity ^ 1] == -1).all()) + assert bool((ws.q8.view(2, -1)[parity ^ 1] == 255).all()) + assert int(ws.flags[32]) & 1 == parity ^ 1, "the call counter advanced" + + +def test_guard_refuses_a_workspace_not_at_rest() -> None: + """The attention CTAs poll the QKV rows for the 0xFFFF marker: stale rows read as this call's.""" + s = build_step([300], 1, seed=5) + ws = MegaWorkspace(torch.device("cuda")) + ws.qkv_out.zero_() + args = _args(s, _operands(6), ws) + with pytest.raises(AssertionError, match="not at rest"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_request_past_the_identity_window() -> None: + s = build_step([300], 1, seed=5) + args = _args(s, _operands(6), MegaWorkspace(torch.device("cuda"))) + seq_lens = s.seq_lens.clone() + seq_lens[0] = 2049 + args[_arg("seq_lens")] = seq_lens + with pytest.raises(AssertionError, match="identity selection covers 16 blocks"): + with validating(op): + op(*args) + + +def test_op_rejects_more_than_one_token() -> None: + """Compiled for one request with one query token only.""" + s = build_step([300, 400], 1, seed=5) + args = _args(s, _operands(6), MegaWorkspace(torch.device("cuda"))) + args[_arg("xq")] = torch.cat([args[_arg("xq")]] * 2) + args[_arg("output")] = torch.empty(2, 2048, dtype=torch.bfloat16, device="cuda") + args[_arg("y")] = torch.empty(2, 6144, dtype=torch.bfloat16, device="cuda") + with pytest.raises(RuntimeError, match="one request with one query token"): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_dense_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_dense_decode.py new file mode 100644 index 000000000000..f080ec79d446 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_dense_decode.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dense decode against the fp32 reference and between valid sequence-length bounds.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_dense_decode import ( # noqa: E501 + HEAD_DIM, + HEADS, + KV_HEADS, + MODEL_MAX, + SHORT_BOUND, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_dense_decode import ( # noqa: E501 + m3_dense_decode as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_dense_decode requires a CUDA device" + +PAGE = 32 +SUBPAGES = 2 +SM_SCALE = HEAD_DIM**-0.5 +# The FP8-KV band of the in-tree dense decode test (test_minimax_m3_dense_decode.test_matches_reference). +FP8_TOL = 6e-2 + + +class _Pool: + """The one method the kernel asks of a KV cache manager.""" + + def __init__(self, pool: torch.Tensor): + self.pool = pool + + def get_kv_subpage_pool(self, layer_idx: int, kv_layout: str = "HND"): + assert kv_layout == "HND" + return self.pool, SUBPAGES + + +def _inputs(kv_lens, qlen: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(kv_lens) + pages = max((n + PAGE - 1) // PAGE for n in kv_lens) + slots = batch * pages + 3 + pool = torch.randn( + (slots * SUBPAGES, KV_HEADS, PAGE, HEAD_DIM), device="cuda", generator=g + ).clamp(-4, 4) + # Slots handed out non-contiguously, as the block manager does. + perm = torch.randperm(slots, generator=g, device="cuda")[: batch * pages] + block_table = perm.reshape(batch, pages).to(torch.int32) + # The last third of every request's values is shifted by +1: a third of each row's softmax mass sits there, far + # outside the band, so a decode that missed the tail of a context fails. + for b, n in enumerate(kv_lens): + start = 2 * n // 3 + first_full = (start + PAGE - 1) // PAGE + last = (n + PAGE - 1) // PAGE + if first_full < last: + v_idx = block_table[b, first_full:last].long() * SUBPAGES + 1 + pool[v_idx] += 1.0 + if start % PAGE: + v0 = int(block_table[b, start // PAGE]) * SUBPAGES + 1 + pool[v0, :, start % PAGE :] += 1.0 + q = torch.randn((batch * qlen, HEADS, HEAD_DIM), device="cuda", generator=g).to(torch.bfloat16) + return ( + q, + _Pool(pool.to(torch.float8_e4m3fn)), + block_table, + torch.tensor(kv_lens, device="cuda", dtype=torch.int32), + ) + + +def _args(q, manager, block_table, seq_lens, qlen, max_seq_len): + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + return ( + q, + manager, + 0, + block_table, + seq_lens, + SM_SCALE, + out, + qlen, + max_seq_len, + seq_lens.shape[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 — pytest resolves the fixture by this name + spec = cell.spec + q, manager, block_table, seq_lens = _inputs( + spec["kv_lens"], spec["qlen"], zlib.crc32(cell.why.encode()) % 2**31 + ) + args = _args(q, manager, block_table, seq_lens, spec["qlen"], spec["max_seq_len"]) + certifying(op) + ref = op.reference(*args) + out = op(*args).clone() + torch.testing.assert_close(out.float(), ref, rtol=FP8_TOL, atol=FP8_TOL) + if spec["tight_bound"]: + model_max_output = op( + *_args(q, manager, block_table, seq_lens, spec["qlen"], MODEL_MAX) + ).clone() + # The split reorders only the merge of the CTAs' fp32 partial sums before the one BF16 rounding: every element + # is within 4 of its own BF16 spacing (2^(e - 8) for |x| = m * 2^e, m in [0.5, 1)). + ref_bf16 = model_max_output.float() + _, exp = torch.frexp(ref_bf16.abs()) + spacing = torch.where( + ref_bf16 == 0, + torch.full_like(ref_bf16, 2.0**-133), + torch.ldexp(torch.ones_like(ref_bf16), exp - 8), + ) + assert bool(((out.float() - ref_bf16).abs() <= 4 * spacing).all()) + assert_within_ulp(out, model_max_output, element_ulp=4.0, rms_ulp=1.0) + + +def test_guard_refuses_more_requests_than_the_counters() -> None: + q, manager, block_table, seq_lens = _inputs((3000, 3000), 1, 3) + args = list(_args(q, manager, block_table, seq_lens, 1, SHORT_BOUND)) + args[9] = 1 + with pytest.raises(AssertionError, match="counters cover"): + with validating(op): + op(*args) + + +def test_guard_refuses_an_understated_sequence_length_bound() -> None: + q, manager, block_table, seq_lens = _inputs((20000,), 1, 7) + args = _args(q, manager, block_table, seq_lens, 1, SHORT_BOUND) + with pytest.raises(AssertionError, match="sequence length exceeds its bound"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_draft_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_draft_attention.py new file mode 100644 index 000000000000..222c0aa4fcef --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_draft_attention.py @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_attention entry: the Eagle3 drafter layer's attention against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp4``): the target's hybrid NVFP4 +``MiniMaxM3KVCacheManagerV2`` with its one-model Eagle3 draft layer, the requests allocated in it, drafter +``TrtllmAttentionMetadata`` over the manager's FP8 draft view (the manager the runtime's draft swap installs) filled +by its ``prepare()``, and the drafter's ``TrtllmAttention`` built as ``M3Eagle3Draft`` builds it (fused NeoX RoPE at +the drafter checkpoint's theta, FP8 KV). Earlier rows go into the draft layer's cache through the in-tree scatter at +the slots the manager assigned; the call rotates and appends the step's own. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_draft_attention import ( # noqa: E501 + m3_draft_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp4 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_attention requires a CUDA device" + +# Band, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS. +# The reference reads the same FP8 cache rows the call attends, so what is left is the FP8 FMHA's own arithmetic: +# it takes the rotated Q in E4M3 (rounding the reference's Q to E4M3 takes the RMS gap from ~9 to ~7 ulps) and rounds +# the softmax probabilities before the PV product. Measured on GB300 over these cells (max element / RMS): 13.6 / 8.6 +# (4 x 4 rows), 9.9 / 7.2 (8 x 1), 20.4 / 7.7 (512 context rows). The band leaves about 1.5x; a reference shifted +# by one row is far outside it. +_ELEMENT_ULP = 30.0 +_RMS_ULP = 12.0 + +# Cached tokens per request of each cell; the new rows per request come from the cell. +_CACHED = { + (4, 4): (200, 129, 1000, 64), + (8, 1): (5, 127, 128, 300, 1000, 2047, 64, 900), + (1, None): (0,), +} +#: Rows past the live ones in the fused QKV, as a padded step hands them over. +_PAD = 3 + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _rope(x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """NeoX RoPE over the whole head at the drafter's theta, float64.""" + half = rt.HEAD_DIM // 2 + inv_freq = rt.DRAFT_ROPE_THETA ** ( + -torch.arange(0, rt.HEAD_DIM, 2, dtype=torch.float64, device=x.device) / rt.HEAD_DIM + ) + angles = positions.to(torch.float64)[:, None] * inv_freq[None, :] + x = x.to(torch.float64) + x1, x2 = x[..., :half], x[..., half:] + return torch.cat( + [x1 * angles.cos() - x2 * angles.sin(), x2 * angles.cos() + x1 * angles.sin()], -1 + ) + + +@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 this name + spec = cell.spec + context = spec.get("context") + cached = _CACHED[spec["reqs"], spec.get("qlen")] + qlens = (context,) if context else (spec["qlen"],) * spec["reqs"] + assert len(cached) == spec["reqs"] + contexts = 1 if context else 0 + rids = list(range(spec["reqs"])) + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=True) + try: + backend = rt.make_draft_backend() + metadata = rt.make_draft_metadata(manager) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.HEAD_DIM, generator=gen, device="cuda") + hist = hist.to(torch.float8_e4m3fn) + rt.write_fp8_rows(manager, rt.DRAFT, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + assert int(metadata.num_tokens) == total + width = (rt.HEADS + 2 * rt.KV_HEADS) * rt.HEAD_DIM + qkv = torch.randn(total + _PAD, width, generator=gen, device="cuda").to(torch.bfloat16) + + certifying(op) + out = op(backend, qkv, metadata) + ref = op.reference(backend, qkv, metadata) + assert out.shape == (total, rt.HEADS * rt.HEAD_DIM) and out.dtype == torch.bfloat16 + assert torch.isfinite(out.float()).all() + element, rms = _ulp_stats(out, ref) + print(f"\n[{cell.why}] max element {element:.2f} ulp, rms {rms:.3f} ulp") + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + # The call's append: each live row's rotated K and its V at the row's position, E4M3 (unit KV scale), to + # within one E4M3 step of the float64 rotation (the kernel rotates in FP32). + pool = manager.get_buffers(rt.DRAFT, "HND") + q_size = rt.HEADS * rt.HEAD_DIM + first = 0 + for rid, c, n in zip(rids, cached, qlens): + positions = torch.arange(c, c + n, device="cuda") + slots = rt.flat_slots(manager, rt.DRAFT, rid, positions).long() + rows = qkv[first : first + n] + want_k = _rope(rows[:, q_size : q_size + rt.HEAD_DIM].float(), positions) + want_v = rows[:, q_size + rt.HEAD_DIM :].to(torch.float8_e4m3fn) + got_k = pool[slots // rt.PAGE, 0, 0, slots % rt.PAGE].to(torch.float64) + got_v = pool[slots // rt.PAGE, 1, 0, slots % rt.PAGE] + step = want_k.abs().clamp_min(2.0**-6) * 2.0**-3 + assert bool(((got_k - want_k).abs() <= step).all()), "K is appended rotated, in E4M3" + assert torch.equal(got_v.view(torch.uint8), want_v.view(torch.uint8)) + first += n + + # The band discriminates: a reference whose rows are shifted by one token is outside it. + shifted = op.reference(backend, qkv.roll(1, 0), metadata) + with pytest.raises(AssertionError): + assert_within_ulp(out, shifted, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..867eb5089c25 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp8_indexer_qk_norm_rope entry of the MiniMax-M3 tp4 target: index-Q out and index-K in the +cache against the in-place fused norm + RoPE followed by the E4M3 cast, byte for byte up to the 0.1% budget the +in-tree test of the kernel allows.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_fp8_indexer_qk_norm_rope import ( # noqa: E501 + m3_fp8_indexer_qk_norm_rope as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp8_indexer_qk_norm_rope requires a CUDA device" + +DIM, ROTARY, EPS, THETA, PAGE = 128, 64, 1e-6, 5e6, 128 + + +def _bytes_close(a: torch.Tensor, b: torch.Tensor) -> None: + assert a.shape == b.shape and a.dtype == b.dtype == torch.float8_e4m3fn + differ = int((a.view(torch.uint8) != b.view(torch.uint8)).sum()) + assert differ <= (a.numel() - 1) // 1000, f"{differ} of {a.numel()} E4M3 bytes differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + m, heads = cell.spec["m"], cell.spec["heads"] + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + qk = torch.randn(m, (heads + 1) * DIM, generator=g, device="cuda").to(torch.bfloat16) + qw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (m,), generator=g, device="cuda", dtype=torch.int32) + pages = (m + PAGE - 1) // PAGE + 1 + cache = torch.zeros(pages, 1, PAGE, DIM, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(pages * PAGE, generator=g, device="cuda")[:m].to(torch.int32) + before = qk.clone() + args = (qk, cache, slots, heads, DIM, ROTARY, EPS, qw, kw, THETA, pos) + with validating(op): + q = op(*args) + assert torch.equal(qk, before), "the BF16 input is left untouched" + q_ref, k_ref = op.reference(*args) + _bytes_close(q.reshape(m, heads, DIM), q_ref) + flat = cache.view(-1, DIM) + _bytes_close(flat[slots.long()], k_ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_fused_decode_attn.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_fused_decode_attn.py new file mode 100644 index 000000000000..f5fd205acd48 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_fused_decode_attn.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fused_decode_attn entry (fused norm/RoPE/KV-insert + identity sparse decode + merge).""" + +import inspect +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_fused_decode_attn import ( # noqa: E501 + PAGE, + mxfp8_rows, + unswizzle_scales, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_fused_decode_attn import ( # noqa: E501 + m3_fused_decode_attn as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.m3_fused_attn import ( # noqa: E501 + default_geometry, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_fused_attn_step import ( # noqa: E402 + ATTN_ELEMENT_ULP, + ATTN_RMS_ULP, + MXFP8_ROWS, + attention_args, + build_step, + check_insert, + decode_over_written_caches, +) +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_fused_decode_attn requires a CUDA device" + + +def _buffers(s): + t = s.num_tokens + out = torch.empty(t, 2048, dtype=torch.bfloat16, device="cuda") + q8 = torch.zeros(MXFP8_ROWS, 2048, dtype=torch.float8_e4m3fn, device="cuda") + sf = torch.zeros(MXFP8_ROWS * 64, dtype=torch.uint8, device="cuda") + return out, q8, sf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- the fixture from _validating + """The insert against the reference's recipe, the output against the reference's decode over the caches the + op wrote (see ATTN_ELEMENT_ULP / CODE_MISMATCH_FRACTION), the MXFP8 hand-off bit-exactly.""" + certifying(op) + spec = cell.spec + qg, bpc = spec["geometry"] + assert default_geometry(spec["qlen"], len(spec["seq_lens"])) == (qg, bpc), ( + "the cell's geometry is the target's" + ) + s = build_step(spec["seq_lens"], spec["qlen"], seed=zlib.crc32(cell.why.encode()) % 2**31) + out, q8, sf = _buffers(s) + args = (s.packed, *attention_args(s), qg, bpc, out, q8, sf, True) + before = (s.kv_data.clone(), s.kv_scale.clone(), s.idx_k.clone()) + ref = op.reference(*args) + op(*args) + check_insert(s, before, ref) + want = decode_over_written_caches(s, s.packed) + assert_within_ulp(out.view(-1, 128), want.view(-1, 128), ATTN_ELEMENT_ULP, ATTN_RMS_ULP) + # The MXFP8 hand-off is the merge's recipe applied to the op's own bf16 rows: bit-exact. + t = s.num_tokens + q8_want, ex_want = mxfp8_rows(out) + assert torch.equal(q8[:t].view(torch.uint8), q8_want), "MXFP8 hand-off values" + assert torch.equal(unswizzle_scales(sf, t, 64), ex_want), "MXFP8 hand-off scales" + + +def _cell_inputs(index=1): + spec = op.CELLS[index].spec + s = build_step(spec["seq_lens"], spec["qlen"], seed=7) + out, q8, sf = _buffers(s) + qg, bpc = spec["geometry"] + return s, [s.packed, *attention_args(s), qg, bpc, out, q8, sf, True] + + +def _arg(args, name): + """Position of a parameter in the op's argument list (``is_valid`` mirrors ``__call__``).""" + return list(inspect.signature(type(op).is_valid).parameters)[1:].index(name) + + +def test_guard_refuses_a_request_past_the_identity_window() -> None: + """16 blocks of 128: a 2049-token request would have its new token placed in block 15 and its tail unread.""" + s, args = _cell_inputs() + seq_lens = s.seq_lens.clone() + seq_lens[0] = 16 * PAGE + 1 + args[_arg(args, "seq_lens")] = seq_lens + with pytest.raises(AssertionError, match="identity selection covers 16 blocks"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_slot_the_block_table_does_not_give() -> None: + """The attention reads the inserted page from the cache it loaded: a slot elsewhere is never attended.""" + s, args = _cell_inputs() + loc = s.out_cache_loc.clone() + loc[0] = (loc[0] + 1) % (s.kv_data.shape[0] * PAGE) + args[_arg(args, "out_cache_loc")] = loc + with pytest.raises(AssertionError, match="out_cache_loc must be the slot"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_block_table_narrower_than_the_length() -> None: + s, args = _cell_inputs() + args[_arg(args, "block_table")] = s.block_table[:, :1].clone() + with pytest.raises(AssertionError, match="block_table must hold every attended block"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_strided_kv_pool() -> None: + """The kernel reads each [128, 64] head block as contiguous bytes; a transposed view is read as if it were.""" + s, args = _cell_inputs() + args[_arg(args, "kv_data")] = s.kv_data.transpose(3, 4).contiguous().transpose(3, 4) + with pytest.raises(AssertionError, match="kv_data must be"): + with validating(op): + op(*args) + + +def test_op_rejects_an_uncompiled_geometry() -> None: + """Only the target's geometries are compiled; anything else is refused by the op itself.""" + s, args = _cell_inputs(index=0) + args[_arg(args, "tokens_per_cluster")] = 2 + with pytest.raises(RuntimeError, match="is not compiled"): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_index_score.py new file mode 100644 index 000000000000..341fb87a4993 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_index_score.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_index_score entry: the MSA index scorer with its pre-wait schedule. + +Every cell is held to the fp32 reference (a few fp32 ulps: the kernel sums the 128 exact products in its k-split order) +and to the in-tree scorer (``cute_dsl_minimax_m3_index_decode_score``) bit for bit, which the pre-wait prefetch and the +early tile issue must not change. +""" + +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.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_index_score import ( # noqa: E501 + PAGE, + SPLIT_K, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_index_score import ( # noqa: E501 + m3_index_score as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_index_score requires a CUDA device" + +HEADS = 1 # index heads per rank at TP4 +HEAD_DIM = 128 + + +def _inputs(seq_lens, qlen: int, seed: int, dtype=torch.float8_e4m3fn): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + max_blocks = (max(seq_lens) + PAGE - 1) // PAGE + stride = ((max_blocks + 15) // 16) * 16 + pages = batch * max_blocks + block_table = ( + torch.randperm(pages, device="cuda", generator=g).to(torch.int32).view(batch, max_blocks) + ) + idx_q = torch.randn(batch * qlen, HEADS, HEAD_DIM, device="cuda", generator=g).to(dtype) + index_k = torch.randn(pages, PAGE, HEAD_DIM, device="cuda", generator=g).to(dtype) + seq = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + return idx_q, index_k, block_table, seq, stride + + +def _score(stride: int, total_q: int): + backing = torch.full( + (HEADS, stride, total_q), -float("inf"), device="cuda", dtype=torch.float32 + ) + return backing.transpose(1, 2) + + +@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 this name + spec = cell.spec + qlen = spec["qlen"] + idx_q, index_k, block_table, seq, stride = _inputs( + spec["seq_lens"], qlen, zlib.crc32(cell.why.encode()) + ) + score = _score(stride, idx_q.shape[0]) + args = (idx_q, index_k, block_table, seq, score, qlen, SPLIT_K, spec["early"]) + certifying(op) + ref = op.reference(*args) + op(*args) + stock = _score(stride, idx_q.shape[0]) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen, SPLIT_K + ) + torch.cuda.synchronize() + torch.testing.assert_close(score, ref, rtol=1e-5, atol=1e-4) + assert torch.equal(score, stock), "the scores differ from the in-tree scorer" + + +@pytest.mark.parametrize("early", [True, False], ids=["early", "prefetch"]) +def test_bf16_index_is_the_in_tree_scorer(early) -> None: + idx_q, index_k, block_table, seq, stride = _inputs((5000,), 4, 11, dtype=torch.bfloat16) + score, stock = _score(stride, 4), _score(stride, 4) + op(idx_q, index_k, block_table, seq, score, 4, SPLIT_K, early) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, 4, SPLIT_K + ) + assert torch.equal(score, stock) + + +def test_guard_refuses_a_request_past_its_block_table() -> None: + idx_q, index_k, block_table, seq, stride = _inputs((1000,), 1, 5) + seq = seq + 4 * PAGE + with pytest.raises(AssertionError, match="past its block-table row"): + with validating(op): + op(idx_q, index_k, block_table, seq, _score(stride + 16, 1), 1, SPLIT_K, True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_msa_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_msa_attention.py new file mode 100644 index 000000000000..4b96f18d5c09 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_msa_attention.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_msa_attention entry: one layer's MSA backend step against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp4``): a hybrid NVFP4 ``MiniMaxM3KVCacheManagerV2`` +with the requests allocated in it, ``MiniMaxM3MsaSparseAttentionMetadata`` filled by its ``prepare()``, and the layer's +``MiniMaxM3MsaSparseAttention`` built as the target builds it. Earlier rows (and, on a sparse layer, the step's own +rows, which the target's producer writes before the call) go into the cache through the in-tree scatters, at the +slots the manager assigned. The sparse layers' index-K rows are planted so every block's score is 1/8 or more apart +from every other's: the top-k selection the reference recomputes is then unambiguous. + +The cells are the target's calls (``_attention_core`` in modeling.py): a dense layer with live bf16 K / V in a +prefill chunk and live FP8 K / V in a decode step (the dense layers sit in the hybrid cache's FP8 half), a sparse +layer whose producer already wrote the caches (``k = v = None``) in a prefill chunk beyond 16 blocks and in a mixed +step, both with the NVFP4 scales the target passes. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_msa_attention import ( # noqa: E501 + m3_msa_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp4 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_msa_attention requires a CUDA device" + +# Bands, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS, +# per kernel the step reaches. The reference reads the same quantized cache rows (and casts Q to E4M3 where the FP8 +# FMHA does), so what is left is each FMHA's own arithmetic, measured on GB300 over these cells (max element / RMS): +# * fmha_sm100 over the FP8 pool (dense context rows): 10.1 / 6.1. It runs E4M3 x E4M3 and its softmax +# probabilities are rounded before the PV product. +# * trtllm-gen over the FP8 pool (dense generation rows): 7.0 / 6.1. It rounds the probabilities to E4M3 (scaled by +# 448 against the running maximum): emulating that brings a single-tile request to 0.6 RMS. +# * the NVFP4 CSR prefill (sparse context rows): 34.4 / 11.8 at 4096 rows, 26.8 / 9.1 at 512. It stages the +# reconstructed K / V in E4M3 (the in-tree A/B test, test_minimax_m3_sparse_attn_decode, notes it and holds it +# to 3e-2). +# * the Triton NVFP4 sparse decode (sparse generation rows): 1.0 / 0.25 -- the band of its own entry +# (m3_sparse_decode): exact operands, FP16 probabilities, BF16 rounding. +# Each band leaves about 1.5x the measured maximum; a shifted or wrongly selected reference is far outside them. +_FP8_CONTEXT = dict(element_ulp=16.0, rms_ulp=9.0) +_FP8_DECODE = dict(element_ulp=12.0, rms_ulp=9.0) +_NVFP4_CONTEXT = dict(element_ulp=48.0, rms_ulp=16.0) +_NVFP4_DECODE = dict(element_ulp=2.0, rms_ulp=0.5) +_BANDS = { + ("dense", "context"): (_FP8_CONTEXT, None), + ("dense", "generation"): (None, _FP8_DECODE), + ("sparse", "context"): (_NVFP4_CONTEXT, None), + ("sparse", "mixed"): (_NVFP4_CONTEXT, _NVFP4_DECODE), +} + + +def _assert_in_bands(output, ref, bands, context_rows: int) -> None: + """Gate the context rows and the generation rows each on their own kernel's band.""" + context_band, decode_band = bands + if context_rows: + assert_within_ulp(output[:context_rows], ref[:context_rows], **context_band) + if context_rows < output.shape[0]: + assert_within_ulp(output[context_rows:], ref[context_rows:], **decode_band) + + +# (request rows, cached tokens per request, new tokens per request, context rows) per cell. +_SHAPES = { + ("dense", "context"): ((384,), (1024,), 1), + ("dense", "generation"): ((127, 128, 1000, 2500), (1, 1, 1, 1), 0), + ("sparse", "context"): ((0,), (4096,), 1), + ("sparse", "mixed"): ((0, 2999, 4500), (512, 1, 1), 1), +} + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + """The two quantities ``assert_within_ulp`` gates, in BF16 ulps.""" + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _shift(x: torch.Tensor) -> torch.Tensor: + """``x``'s rows moved down by one (byte-wise, so E4M3 rows move too).""" + return x.view(torch.uint8).roll(1, 0).view(x.dtype) if x.element_size() == 1 else x.roll(1, 0) + + +def _shape(cell): + spec = cell.spec + cached, qlens, contexts = _SHAPES[spec["layer"], spec["phase"]] + if spec["phase"] == "context": + assert sum(qlens) == spec["tokens"] + elif spec["phase"] == "generation": + assert len(qlens) == spec["reqs"] + else: + assert sum(qlens[:contexts]) == spec["tokens"] and len(qlens) - contexts == spec["reqs"] + return cached, qlens, contexts + + +def _run_dense(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + rt.write_fp8_rows(manager, rt.DENSE, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + # A prefill chunk arrives in BF16 (the hybrid cache's dense layers take FP8 Q / K / V only on decode steps). + dtype = torch.bfloat16 if contexts else torch.float8_e4m3fn + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + k = torch.randn(total, rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + v = torch.randn(total, rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + # The NVFP4 target passes its layer scales on every layer; an FP8 layer's write and FMHA take none. + inv = torch.tensor([1.0, 0.5, 2.0], device="cuda") + fwd = 1.0 / inv + return (q, k, v, None, inv, fwd), dict(rids=rids, cached=cached, qlens=qlens, contexts=contexts) + + +def _run_sparse(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + u0, u1 = rt.index_directions(gen) + lengths = [c + n for c, n in zip(cached, qlens)] + keys = [torch.randn(n, rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths] + values = [torch.randn(n, rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths] + quant = rt.nvfp4_quant(torch.cat(keys), torch.cat(values)) + for rid, n, kk, vv in zip(rids, lengths, keys, values): + # Every row, the step's own included: the producer writes a sparse layer's caches before the call. + idx_k = rt.planted_index_keys(n, u0, u1, gen) + rt.write_nvfp4_rows( + manager, + metadata, + rt.SPARSE, + rid, + torch.arange(n), + kk.to(torch.bfloat16), + vv.to(torch.bfloat16), + idx_k, + quant, + ) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + # Query tokens alternate between the two planted directions, so neighbouring tokens select different blocks. + idx_q = torch.stack([u0, u1])[torch.arange(total, device="cuda") % 2].to(torch.float8_e4m3fn) + return (q, None, None, idx_q, quant, 1.0 / quant), dict( + rids=rids, cached=cached, qlens=qlens, contexts=contexts + ) + + +@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 this name + spec = cell.spec + sparse = spec["layer"] == "sparse" + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=False) + try: + backend = rt.make_msa_backend(rt.SPARSE if sparse else rt.DENSE) + metadata = rt.make_msa_metadata(manager) + run = _run_sparse if sparse else _run_dense + (q, k, v, idx_q, inv, fwd), step = run(cell, gen, manager, backend, metadata) + output = torch.full( + (q.shape[0], rt.HEADS * rt.HEAD_DIM), float("nan"), dtype=torch.bfloat16, device="cuda" + ) + certifying(op) + result = op(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert result is output + ref = op.reference(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert torch.isfinite(output.float()).all(), "every live row is written" + bands = _BANDS[spec["layer"], spec["phase"]] + context_rows = sum(step["qlens"][: step["contexts"]]) + for name, rows in ( + ("context", slice(0, context_rows)), + ("generation", slice(context_rows, None)), + ): + if output[rows].numel(): + element, rms = _ulp_stats(output[rows], ref[rows]) + print( + f"\n[{cell.why}] {name} rows: max element {element:.2f} ulp, rms {rms:.3f} ulp" + ) + _assert_in_bands(output, ref, bands, context_rows) + + if k is not None: + # The call's own cache write: the live K / V rows land, as a plain E4M3 cast, at their positions. + pool = manager.get_buffers(rt.DENSE, "HND") + first = 0 + for rid, c, n in zip(step["rids"], step["cached"], step["qlens"]): + slots = rt.flat_slots(manager, rt.DENSE, rid, torch.arange(c, c + n)).long() + for role, src in ((0, k), (1, v)): + got = pool[slots // rt.PAGE, role, 0, slots % rt.PAGE] + want = src[first : first + n].to(torch.float8_e4m3fn) + assert torch.equal(got.view(torch.uint8), want.view(torch.uint8)) + first += n + + # The band discriminates: a reference whose query rows are shifted by one token is outside it. + shifted = op.reference(backend, _shift(q), k, v, idx_q, metadata, output, sparse, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, shifted, bands, context_rows) + if sparse: + # And the selection is load-bearing: attending every causal block instead is outside it too. + dense = op.reference(backend, q, k, v, idx_q, metadata, output, False, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, dense, bands, context_rows) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_producer.py new file mode 100644 index 000000000000..803ee09dfee3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_producer.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_producer entry: the pre-wait-staged NVFP4 QKV + indexer producer. + +Every cell is held to the native producer (``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert``) bit for bit on q, +index-q and every K / V / scale / index-K cache byte, and to the fp32 torch reference: the FP8 outputs within one E4M3 +step, the decoded NVFP4 K / V within one E2M1 step of their block scale. +""" + +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.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_producer import ( # noqa: E501 + FP4_BLOCK, + HEADS_Q, + PAGE, + decode_nvfp4_token, + scale_index, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_producer import ( # noqa: E501 + m3_producer as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_producer requires a CUDA device" + +PAGES = 1024 +MAX_POSITIONS = 262144 +EPS = 1e-6 +KV_QUANT = (1.0, 0.37, 0.61) + + +def _caches(): + return ( + torch.zeros(PAGES, 2, 1, PAGE, 64, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 2, 1, PAGE, 8, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 1, PAGE, 128, dtype=torch.float8_e4m3fn, device="cuda"), + ) + + +def _inputs(tokens: int, scale: float, bad_slots: bool, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + rotary = torch.randn(MAX_POSITIONS, 2, 32, device="cuda", generator=g) + weights = [(torch.randn(128, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + packed = ( + torch.randn(tokens, (HEADS_Q + 4) * 128, device="cuda", generator=g) * scale + ).bfloat16() + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + if bad_slots: + slots[1] = -1 + slots[2] = PAGES * PAGE + 5 + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + quant = torch.tensor(KV_QUANT, device="cuda") + return packed, slots, quant, weights, rotary, positions + + +def _args(inputs, caches, early: bool = True): + packed, slots, quant, weights, rotary, positions = inputs + return (packed, *caches, slots, quant, EPS, *weights, rotary, positions, HEADS_Q, True, early) + + +def _native(inputs, caches): + packed, slots, quant, weights, rotary, positions = inputs + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, *caches, slots, quant, HEADS_Q, 1, 1, 128, 64, EPS, *weights, rotary, positions + ) + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +@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 this name + spec = cell.spec + early = spec.get("early", True) + inputs = _inputs( + spec["tokens"], spec["scale"], spec.get("bad_slots", False), zlib.crc32(cell.why.encode()) + ) + stock_caches, caches = _caches(), _caches() + stock = _native(inputs, stock_caches) + args = _args(inputs, caches, early) + certifying(op) + ref = op.reference(*args) + out = op(*args) + torch.cuda.synchronize() + names = ("q", "index_q", "kv_data", "kv_scale", "index_k") + for name, s, o in zip(names, (*stock, *stock_caches), (*out, *caches)): + assert torch.equal(s.view(torch.uint8), o.view(torch.uint8)), ( + f"{name} differs from the native producer" + ) + + _within_fp8(out[0], ref["q"], "q") + _within_fp8(out[1][:, 0], ref["index_q"], "index_q") + kv_data, kv_scale, index_k = caches + slots = inputs[1] + checked = 0 + for t in range(spec["tokens"]): + slot = int(slots[t]) + if slot < 0 or slot >= PAGES * PAGE: + continue + page, within = divmod(slot, PAGE) + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") + for plane, name in ((0, "k"), (1, "v")): + scales = kv_scale[page, plane, 0].reshape(-1)[scale_index(within, plane).cuda()] + dq = decode_nvfp4_token(kv_data[page, plane, 0, within], scales, KV_QUANT[1 + plane]) + r = ref[name][t].bfloat16().float() + # Round to nearest E2M1 under the block's E4M3-rounded scale s (in value units s / quant_scale): at most + # half the widest E2M1 step (one unit), plus the clipping at 6 units when s rounded down; a scale that + # underflows E4M3 zeroes the block. A slack of one E4M3 step covers the kernel's approximate reciprocals. + qs = KV_QUANT[1 + plane] + amax = r.view(-1, FP4_BLOCK).abs().amax(dim=1) + unit = (qs * amax / 6.0).to(torch.float8_e4m3fn).float() / qs + bound = torch.where( + unit > 0, 1.15 * unit + (amax - 6.0 * unit / 1.15).clamp_min(0.0), amax + ) + bound = bound.repeat_interleave(FP4_BLOCK) + 1e-12 + assert bool(((dq - r).abs() <= bound).all()), (name, t) + checked += 1 + assert checked >= spec["tokens"] - 2 + + +def test_early_trigger_does_not_change_the_bits() -> None: + inputs = _inputs(48, 1.0, False, 7) + a, b = _caches(), _caches() + out_a = op(*_args(inputs, a, early=True)) + out_b = op(*_args(inputs, b, early=False)) + for x, y in zip((*out_a, *a), (*out_b, *b)): + assert torch.equal(x.view(torch.uint8), y.view(torch.uint8)) + + +def test_guard_refuses_a_repeated_slot() -> None: + inputs = list(_inputs(4, 1.0, False, 3)) + inputs[1] = inputs[1].clone() + inputs[1][1] = inputs[1][0] + with pytest.raises(AssertionError, match="must not repeat a slot"): + with validating(op): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_operand_off_the_device() -> None: + inputs = list(_inputs(4, 1.0, False, 9)) + inputs[5] = inputs[5].cpu() # position ids on the host + with pytest.raises(RuntimeError, match="packed's CUDA device"): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_unaligned_index_k_page_stride() -> None: + """The index-K rows are written as 32-bit words: a page stride that is not a multiple of 4 is refused.""" + kv_data, kv_scale, _ = _caches() + backing = torch.zeros(PAGES * (128 * 128 + 1), dtype=torch.float8_e4m3fn, device="cuda") + index_k = backing.as_strided((PAGES, 1, PAGE, 128), (128 * 128 + 1, 128 * 128, 128, 1)) + with pytest.raises(RuntimeError, match="index_k"): + op(*_args(_inputs(4, 1.0, False, 13), (kv_data, kv_scale, index_k))) + + +def test_op_refuses_unaligned_weights_or_rope_table() -> None: + """The norm weights are read as 8-byte words and the RoPE rows as float4: offset views are refused.""" + packed, slots, quant, weights, rotary, positions = _inputs(4, 1.0, False, 17) + offset_weight = torch.zeros(129, dtype=torch.bfloat16, device="cuda")[1:] + with pytest.raises(RuntimeError, match="norm weights"): + op( + *_args( + (packed, slots, quant, [offset_weight, *weights[1:]], rotary, positions), _caches() + ) + ) + flat = torch.zeros(rotary.numel() + 1, device="cuda") + offset_rope = flat[1:].view(rotary.shape) + with pytest.raises(RuntimeError, match="rotary_cos_sin"): + op(*_args((packed, slots, quant, weights, offset_rope, positions), _caches())) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..2e80fa0e980b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_qk_norm_rope_fp8 entry of the MiniMax-M3 tp4 target: the out-of-place FP8 form of the fused +per-head Gemma RMSNorm + partial NeoX RoPE against the in-place kernel followed by the E4M3 cast. V is a cast, so it +matches bit for bit; Q and K round once from fp32 where the composition rounds through bf16 first, so each element is +within one E4M3 step of it.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_qk_norm_rope_fp8 import ( # noqa: E501 + m3_qk_norm_rope_fp8 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_qk_norm_rope_fp8 requires a CUDA device" + +HEAD_DIM, ROTARY_DIM, EPS, THETA = 128, 64, 1e-6, 5e6 + + +def _args(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + q, kv = spec["q"], spec["kv"] + qkv = torch.randn(spec["m"], (q + 2 * kv) * HEAD_DIM, generator=g, device="cuda").to( + torch.bfloat16 + ) + qw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (spec["m"],), generator=g, device="cuda", dtype=torch.int32) + return ( + qkv, + q, + kv, + kv, + HEAD_DIM, + ROTARY_DIM, + EPS, + qw, + kw, + THETA, + True, + pos, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + args = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + before = args[0].clone() + with validating(op): + out = op(*args) + assert torch.equal(args[0], before), "the FP8 form must leave qkv untouched" + ref = op.reference(*args) + assert out.dtype == torch.float8_e4m3fn and out.shape == ref.shape + v_start = (cell.spec["q"] + cell.spec["kv"]) * HEAD_DIM + assert torch.equal(out[:, v_start:].view(torch.uint8), ref[:, v_start:].view(torch.uint8)) + assert _e4m3_steps(out[:, :v_start], ref[:, :v_start]) <= 1 + + +def _e4m3_steps(a, b) -> int: + """The largest distance between two E4M3 tensors in representable values (E4M3 codes order by magnitude).""" + + def signed(t): + c = t.view(torch.uint8).int() + # Sign-magnitude to a line: -0 and +0 both map to 0. + return torch.where(c >= 128, 128 - c, c) + + return int((signed(a) - signed(b)).abs().max()) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_qkv_insert.py new file mode 100644 index 000000000000..c570f15f772a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_qkv_insert.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_qkv_insert entry: the generic horizontal producer on an FP8 main cache. + +Q, index-Q and every written K / V / index-K cache row are held within one E4M3 step of the fp32 reference. +""" + +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.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_qkv_insert import ( # noqa: E501 + m3_qkv_insert as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_qkv_insert requires a CUDA device" + +PAGES, PAGE, HEADS, KV, INDEX, D, ROTARY = 256, 128, 16, 1, 1, 128, 64 +MAX_POSITIONS = 262144 +EPS = 1e-6 + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +def _args(tokens: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + packed = torch.randn( + tokens, (HEADS + 2 * KV + INDEX + 1) * D, device="cuda", generator=g + ).bfloat16() + kv = torch.zeros(PAGES, 2, KV, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + index_k = torch.zeros(PAGES, 1, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + weights = [(torch.randn(D, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + rotary = torch.randn(MAX_POSITIONS, 2, ROTARY // 2, device="cuda", generator=g) + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + return ( + packed, + kv, + None, + index_k, + slots, + None, + HEADS, + KV, + INDEX, + D, + ROTARY, + EPS, + *weights, + rotary, + positions, + ) + + +@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 this name + args = _args(cell.spec["m"], zlib.crc32(cell.why.encode())) + certifying(op) + ref = op.reference(*args) + q, index_q = op(*args) + torch.cuda.synchronize() + _within_fp8(q.view(ref["q"].shape), ref["q"], "q") + _within_fp8(index_q.view(ref["index_q"].shape), ref["index_q"], "index_q") + kv, index_k, slots = args[1], args[3], args[4] + for t in range(cell.spec["m"]): + page, within = divmod(int(slots[t]), PAGE) + _within_fp8(kv[page, 0, :, within], ref["k"][t], f"k token {t}") + _within_fp8(kv[page, 1, :, within], ref["v"][t], f"v token {t}") + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_select_blocks.py new file mode 100644 index 000000000000..77583bdb979e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_select_blocks.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_select_blocks entry of the MiniMax-M3 tp4 target: MSA top-16 block selection with the forced +init / local blocks and per-row valid-block masks.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_select_blocks import ( # noqa: E501 + m3_select_blocks as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_blocks requires a CUDA device" + +TOPK, INIT, LOCAL = 16, 0, 1 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + n_blocks, total_q = spec["n_blocks"], spec["total_q"] + # Distinct scores (a permutation) so the selection has no ties to break. + scores = ( + torch.randperm(n_blocks * total_q, generator=g, device="cuda") + .float() + .view(1, n_blocks, total_q) + ) + if spec["mixed"]: + n_valid = torch.randint( + 17, n_blocks + 1, (total_q,), generator=g, device="cuda", dtype=torch.int32 + ) + else: + n_valid = torch.full((total_q,), n_blocks, dtype=torch.int32, device="cuda") + with validating(op): + out = op(scores, TOPK, n_valid, INIT, LOCAL, True) + ref = op.reference(scores, TOPK, n_valid, INIT, LOCAL, True) + assert out.shape == ref.shape + assert torch.equal(out, ref) + assert out.permute(1, 0, 2).is_contiguous(), "head-major backing" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_select_long.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_select_long.py new file mode 100644 index 000000000000..42f804302ef7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_select_long.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_select_long entry (top-16 MSA block selection for long decode rows).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_select_long import ( # noqa: E501 + m3_select_long as op, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + select_blocks_from_maxscore, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_long requires a CUDA device" + +# Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 2-, 4- and 9-key +# widths, past the row (clamped), and zero (a CUDA-graph padding row). +_MIXED_VALID = (0, 5, 16, 17, 200, 2048, 2049, 4096, 4097, 8193, 9000, 1, 300, 6000, 129, 3000) + + +def _scores(spec, g): + nb, tq = spec["num_blocks"], spec["total_q"] + s = torch.randn(spec["num_kv_heads"], nb, tq, generator=g, device="cuda", dtype=torch.float32) + kind = spec["scores"] + if kind == "ties": + # Even tokens: a flat row with ten winners, so thousands of keys tie at the 16th-largest score (the + # re-read path). Odd tokens: scores on a coarse grid, so ties are broken by block index. + s[:, :, 0::2] = 0.25 + winners = torch.randint(0, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda") + s[:, :, 0::2].scatter_(1, winners[:, :, 0::2], 3.0) + s[:, :, 1::2] = (s[:, :, 1::2] * 2).round() / 2 + elif kind == "nonfinite": + s[:, 7::50, :] = float("inf") + s[:, 11::70, :] = float("nan") + s[:, 3::40, :] = -0.0 + s[:, 4::40, :] = 0.0 + # Odd tokens: almost everything -inf, so picks run past the finite scores. + s[:, :, 1::2] = float("-inf") + s[:, 0:5, 1::2] = 1.0 + return s + + +def _valid(spec, g): + tq, nb = spec["total_q"], spec["num_blocks"] + if spec["valid"] == "mixed": + vals = [_MIXED_VALID[i % len(_MIXED_VALID)] for i in range(tq)] + else: + vals = (nb - torch.randint(0, 12, (tq,), generator=g, device="cuda")).tolist() + vals[0] = nb + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _case(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + scores = _scores(spec, g) + nvb = _valid(spec, g) + out = torch.empty( + spec["num_kv_heads"], spec["total_q"], TOPK, dtype=torch.int32, device="cuda" + ).permute(1, 0, 2) + return scores, nvb, out, spec["init_blocks"], spec["local_blocks"], False + + +def _mass_tie_rows(scores, nvb, init_blocks, local_blocks): + """[total_q, num_kv_heads] mask of rows where more than 1024 valid blocks share the 16th-largest score.""" + num_kv_heads, num_blocks, total_q = scores.shape + blocks = torch.arange(num_blocks, device=scores.device) + raw = nvb[:total_q].long() + valid = raw.clamp(0, num_blocks) + forced = scores.permute(2, 0, 1).clone() # [total_q, heads, blocks] + forced = torch.where(blocks < init_blocks, 1.0e30, forced) + forced = torch.where(blocks >= (raw - local_blocks).clamp_min(0)[:, None, None], 1.0e29, forced) + in_row = (blocks[None, :] < valid[:, None])[:, None, :] + ranked = torch.where(in_row, forced, float("-inf")).sort(dim=-1, descending=True).values + kth = ranked[..., min(TOPK, num_blocks) - 1 : min(TOPK, num_blocks)] + return ((forced == kth) & in_row).sum(dim=-1) > 1024 + + +@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 this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + + # Selection is an exact claim: the same blocks or a different attention. + assert out.dtype is torch.int32 + assert torch.equal(out, ref), "selected blocks differ from the reference" + # And it is the in-tree selector's answer, bit for bit (the op the target falls back to) -- except on rows + # where more than 1024 blocks tie at the 16th-largest score: there the in-tree op does not break the tie to + # the smallest blocks, while this op and the reference do. + scores, nvb, _, init_blocks, local_blocks, _ = args + in_tree = select_blocks_from_maxscore( + scores, + topk=TOPK, + n_valid_blocks=nvb, + init_blocks=init_blocks, + local_blocks=local_blocks, + head_major_output=True, + ) + mass_tie = _mass_tie_rows(scores, nvb, init_blocks, local_blocks) + assert torch.equal(out[~mass_tie], in_tree[~mass_tie]), ( + "selected blocks differ from the in-tree minimax_m3_select_blocks" + ) + + +def test_pdl_launch_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=1) + pdl_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, False) + op(scores, nvb, pdl_out, init_blocks, local_blocks, True) + assert torch.equal(out, pdl_out) + + +def test_early_trigger_selects_the_same() -> None: + """Releasing the dependent grid right after the select's own wait changes only when the next kernel may launch.""" + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=2) + early_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False) + op(scores, nvb, early_out, init_blocks, local_blocks, True, True) + assert torch.equal(out, early_out) + + +def test_op_refuses_a_row_past_its_width() -> None: + scores = torch.zeros(1, 9217, 1, device="cuda") + nvb = torch.ones(1, dtype=torch.int32, device="cuda") + out = torch.empty(1, 1, TOPK, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="at most 9216"): + op(scores, nvb, out, 0, 1, False) + + +def test_guard_refuses_negative_forcing() -> None: + args = list(_case(op.CELLS[0].spec, seed=2)) + args[4] = -1 + with pytest.raises(AssertionError, match="non-negative"): + with validating(op): + op(*args) + + +def test_guard_refuses_out_aliasing_scores() -> None: + scores, nvb, _, init_blocks, local_blocks, _ = _case(op.CELLS[0].spec, seed=3) + aliased = ( + scores.view(torch.int32) + .reshape(-1)[: scores.shape[2] * TOPK] + .view(scores.shape[2], 1, TOPK) + ) + with pytest.raises(AssertionError, match="alias"): + with validating(op): + op(scores, nvb, aliased, init_blocks, local_blocks, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_sparse_decode.py new file mode 100644 index 000000000000..a57127e9e44d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_sparse_decode.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sparse_decode entry (the target's Triton block-sparse decode and its routes).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_sparse_decode import ( # noqa: E501 + m3_sparse_decode as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_nvfp4_cache import ( # noqa: E402 — needs the path declared above + HEAD_DIM, + NUM_HEADS, + build_step, +) +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_sparse_decode requires a CUDA device" + +# Band, in BF16 ulps (2^-8) of each [16, 128] head row's largest magnitude, and of the relative RMS. Every route +# computes from exact operands (E4M3 q and E2M1 x E4M3 K/V are exact in FP16), so the gap to the float64 reference +# is the BF16 rounding on both sides (one ulp of the row at most), the FP16 probabilities before PV, the split-K +# partials kept in FP32 and merged by their log2-sum-exp, and approximate exp2/log2: 2 ulps per element, the two +# roundings (each at most half an ulp) dominating the RMS, so half an ulp. Measured on GB300 over the cells: at most +# 1.80 ulps per element and 0.20 ulps RMS. +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + + +def _call(case, output, identity): + return op( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + identity, + ) + + +@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 this name + spec = cell.spec + case = build_step( + spec["seq_lens"], + spec["decode_query_len"], + identity=spec["identity"], + seed=zlib.crc32(cell.why.encode()) % 2**31, + ) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + certifying(op) + _call(case, output, spec["identity"]) + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + spec["identity"], + ) + assert torch.isfinite(output.float()).all(), "padding rows must be zeros, not NaN" + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_identity_flag_does_not_change_the_answer() -> None: + """On an identity table the identity routes (the producer, or the multi-query producer from 17 rows) and the + selected-table route (the long-context kernel) compute the same attention.""" + case = build_step((2048, 130, 5, 0, 1777, 900, 2000, 0), 4, identity=True, seed=17) + outs = [] + for identity in (True, False): + output = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + _call(case, output, identity) + outs.append(output) + assert_within_ulp(outs[0], outs[1], element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_guard_refuses_a_padded_selection_slot() -> None: + """The Triton producers index the block table with every one of a token's first min(16, blocks) ids, + -1 included.""" + case = build_step((1500,), 1, identity=True, seed=23) + case.topk[0, 0, 3] = -1 + output = torch.empty(1, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="blocks of the request"): + with validating(op): + _call(case, output, True) + + +def test_guard_refuses_a_false_identity_promise() -> None: + """identity_blocks lets the multi-query producer read one table for the whole request.""" + case = build_step((1500, 1600, 1700, 1800, 1500, 1600, 1700, 1800), 4, identity=True, seed=29) + case.topk[0, 5, :2] = case.topk[0, 5, :2].flip(0) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="identity_blocks promises"): + with validating(op): + _call(case, output, True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py new file mode 100644 index 000000000000..21fec5fc2607 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sparse_decode_lc entry (one-launch long-context sparse decode).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_sparse_decode_lc import ( # noqa: E501 + WIDTH, + dequant_nvfp4_pages, + mxfp8_handoff_reference, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_sparse_decode_lc import ( # noqa: E501 + m3_sparse_decode_lc as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_nvfp4_cache import build_step # noqa: E402 — needs the path declared above +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_sparse_decode_lc requires a CUDA device" + +#: Hand-off buffers are sized as MERGE_MXFP8 sizes them. +_HANDOFF_ROWS = 128 + +# Band, in BF16 ulps (2^-8) of each output row's largest magnitude, and of the relative RMS. +# +# The reference is float64 attention over the exactly dequantized blocks, rounded once to BF16. The kernel's +# operands are exact too (E4M3 q and E2M1 x E4M3 K/V are exact in FP16), so what separates the two is: the final +# BF16 rounding on both sides (up to one ulp of the element, at most one of the row's scale), the softmax +# probabilities rounded to FP16 before PV (2^-11 relative each, averaging down over the block), ex2.approx / +# lg2.approx (~2^-22) and FP32 accumulation order. Hence 2 ulps per element; the RMS is dominated by the two +# roundings (each at most half an ulp, about 0.3 in aggregate), so half an ulp. Measured on GB300 over the cells: +# at most 1.91 ulps per element and 0.25 ulps RMS. +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + + +def _call(case, spec, out, q8, sf): + return op( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + WIDTH // 128, + case.decode_query_len, + case.sm_scale, + spec["linear_softmax"], + False, + True, + ) + + +def _buffers(total_q, mxfp8): + out = torch.empty(total_q, WIDTH, dtype=torch.bfloat16, device="cuda") + if not mxfp8: + return out, None, None + q8 = torch.zeros(_HANDOFF_ROWS, WIDTH, dtype=torch.uint8, device="cuda") + sf = torch.zeros(_HANDOFF_ROWS * (WIDTH // 32), dtype=torch.uint8, device="cuda") + return out, q8, sf + + +@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 this name + spec = cell.spec + case = build_step( + spec["seq_lens"], + spec["decode_query_len"], + identity=False, + seed=zlib.crc32(cell.why.encode()) % 2**31, + ) + total_q = case.q.shape[0] + out, q8, sf = _buffers(total_q, spec["mxfp8"]) + certifying(op) + _call(case, spec, out, q8, sf) + ref = op.reference( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + WIDTH // 128, + case.decode_query_len, + case.sm_scale, + spec["linear_softmax"], + False, + True, + ) + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + # Padding rows attend nothing and must be exact zeros, not NaN. + assert torch.isfinite(out.float()).all() + if spec["mxfp8"]: + # The hand-off is a deterministic function of the op's own BF16 rows: exact. + ref_q8, ref_sf = mxfp8_handoff_reference(out, _HANDOFF_ROWS) + assert torch.equal(q8[:total_q], ref_q8), ( + "MXFP8 rows differ from the recipe applied to the BF16 rows" + ) + assert torch.equal(sf, ref_sf), "UE8M0 scales differ from the recipe / swizzle" + + +def test_reference_dequant_matches_the_quantizer() -> None: + """The reference reads the paged bytes through its own layout claims (K scales token-major, V scales in + 4x4 token quads). Hold them to the quantizer's own flat output: every slot was written in order, so slot + ``s`` is page ``s // 128``, token ``s % 128``, and the flat scales are token-major by construction.""" + case = build_step((3000,), 1, identity=False, seed=5) + pages = torch.arange(case.k_paged.shape[0], device="cuda") + views = ((case.k_paged, case.k_block_scale, False), (case.v_paged, case.v_block_scale, True)) + for role, (data, scale, quad) in enumerate(views): + flat_data, flat_scale = torch.ops.trtllm.fp4_quantize( + case.source[role].reshape(case.num_slots, 1, -1).contiguous(), + case.quant[role + 1 : role + 2], + 16, + False, + False, + ) + expected = dequant_nvfp4_pages( + flat_data.view(torch.uint8).reshape(-1, 1, 128, 64), + flat_scale.view(torch.uint8).reshape(-1, 1, 128, 8), + pages, + quad_scales=False, + ) + got = dequant_nvfp4_pages(data, scale, pages, quad_scales=quad) + assert torch.equal(got, expected), ( + f"role {role}: paged dequantization disagrees with the quantizer" + ) + + +def test_the_variants_agree() -> None: + """The two softmax recurrences and the hand-off switch must not change the BF16 rows beyond rounding.""" + spec = dict(seq_lens=(5000, 12345), decode_query_len=4) + case = build_step(spec["seq_lens"], 4, identity=False, seed=11) + outs = [] + for linear in (False, True): + for mxfp8 in (False, True): + out, q8, sf = _buffers(case.q.shape[0], mxfp8) + _call(case, dict(linear_softmax=linear), out, q8, sf) + outs.append(out) + for other in outs[1:]: + assert_within_ulp(other, outs[0], element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + torch.testing.assert_close(outs[1], outs[0], rtol=0, atol=0) # hand-off on/off: same rows + + +def test_op_refuses_the_uninstantiated_late_trigger() -> None: + case = build_step((3000,), 1, identity=False, seed=3) + out, q8, sf = _buffers(1, True) + with pytest.raises(RuntimeError, match="early_trigger"): + op( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + WIDTH // 128, + 1, + case.sm_scale, + True, + False, + False, + ) + + +def test_guard_refuses_a_block_past_the_request() -> None: + """A selected id at or past the request's block count reads a page the token cannot attend; with every + position masked the CTA's softmax is NaN.""" + case = build_step((3000,), 1, identity=False, seed=4) + case.topk[0, 0, 0] = 40 + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="block count"): + with validating(op): + _call(case, dict(linear_softmax=False), out, q8, sf) + + +def test_guard_refuses_strided_page_rows() -> None: + """The kernel walks a page's token rows at a fixed 64-byte pitch; a padded row stride is misread.""" + case = build_step((3000,), 1, identity=False, seed=6) + wide = torch.zeros(*case.k_paged.shape[:3], 80, dtype=torch.uint8, device="cuda")[..., :64] + wide.copy_(case.k_paged) + case.k_paged = wide + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="64-byte rows"): + with validating(op): + _call(case, dict(linear_softmax=False), out, q8, sf) + + +def test_guard_refuses_a_narrow_handoff() -> None: + case = build_step((3000,), 1, identity=False, seed=7) + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="num_sf_k_tiles"): + with validating(op): + op( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + 8, + 1, + case.sm_scale, + False, + False, + True, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_msa_kv_update.py new file mode 100644 index 000000000000..e1146b79a8b3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_msa_kv_update.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the msa_kv_update entry (the one-launch MSA KV-length correction).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.msa_kv_update import ( # noqa: E501 + TOKENS_PER_BLOCK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.msa_kv_update import ( # noqa: E501 + msa_kv_update as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.msa_kv_update import ( # noqa: E501 + msa_kv_update as hook, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "msa_kv_update requires a CUDA device" + +#: Spare capacity past the live rows, as the metadata's buffers have: the op must leave it alone. +_SPARE = 3 + + +def _case(spec, seed): + """The metadata's buffers for one step, with stale values in every output so a skipped write shows.""" + g = torch.Generator(device="cpu").manual_seed(seed) + batch, qlen = spec["batch"], spec["query_len"] + total_q = batch * qlen + staged = torch.randint(qlen, 5000, (batch,), generator=g, dtype=torch.int32) + accept = spec["accept"] + if accept == "full": + kv = staged.clone() + elif accept == "page_edge": + staged = torch.tensor([130, 257, 2 * TOKENS_PER_BLOCK + qlen][:batch], dtype=torch.int32) + kv = staged - torch.tensor([0, 1, qlen - 1][:batch], dtype=torch.int32) + else: + kv = staged - torch.randint(0, qlen, (batch,), generator=g, dtype=torch.int32) + if accept == "padded": + kv[batch // 2 :] = 0 + staged[batch // 2 :] = 0 + ncols = (5000 + qlen) // TOKENS_PER_BLOCK + 2 + table = ( + torch.randperm(batch * ncols + 7, generator=g)[: batch * ncols] + .to(torch.int32) + .reshape(batch, ncols) + ) + dev = dict(device="cuda", dtype=torch.int32) + stale = -12345 + return ( + torch.cat([kv, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.cat([staged, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.full((batch + _SPARE,), qlen, **dev), + torch.arange(batch, **dev).repeat_interleave(qlen), + torch.arange(qlen, **dev).repeat(batch), + table.cuda(), + TOKENS_PER_BLOCK, + batch, + total_q, + torch.full((batch + _SPARE,), stale, **dev), + torch.full((batch + 1 + _SPARE,), stale, **dev) if spec["cu_kv"] else None, + torch.full((total_q + _SPARE,), stale, **dev), + torch.full((total_q + _SPARE,), stale, **dev) if spec["n_valid"] else None, + ) + + +def _check(args, ref): + seq, cu_kv, loc, n_valid = args[9], args[10], args[11], args[12] + batch, total_q = args[7], args[8] + ref_seq, ref_cu, ref_loc, ref_nv = ref + # Integer bookkeeping: exact, and nothing past the live rows is touched. + assert torch.equal(seq[:batch], ref_seq) + assert torch.equal(loc[:total_q], ref_loc) + assert bool((seq[batch:] == -12345).all()) and bool((loc[total_q:] == -12345).all()) + assert (cu_kv is None) == (ref_cu is None) and (n_valid is None) == (ref_nv is None) + if cu_kv is not None: + assert torch.equal(cu_kv[: batch + 1], ref_cu) + assert bool((cu_kv[batch + 1 :] == -12345).all()) + if n_valid is not None: + assert torch.equal(n_valid[:total_q], ref_nv) + assert bool((n_valid[total_q:] == -12345).all()) + + +@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 this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + op(*args) + _check(args, op.reference(*args)) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_hook_takes_every_cell(cell) -> None: + """The target's hook (what the metadata calls) takes each certified step and writes what the op does.""" + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31 + 1) + assert hook(*args) is True + _check(args, op.reference(*args)) + + +def test_the_hook_declines_what_the_op_would_misread() -> None: + """A strided operand is read wrongly by the op, so the hook declines it (the metadata then runs its own + torch update) and writes nothing.""" + args = list(_case(op.CELLS[2].spec, seed=5)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + before = args[11].clone() + assert hook(*args) is False + assert torch.equal(args[11], before) + + +def test_guard_refuses_a_strided_operand() -> None: + args = list(_case(op.CELLS[2].spec, seed=6)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_request_index_past_the_batch() -> None: + args = list(_case(op.CELLS[2].spec, seed=7)) + args[3] = args[3].clone() + args[3][-1] = args[7] + with pytest.raises(AssertionError, match="q_batch_row"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_allreduce_op_matrix.py new file mode 100644 index 000000000000..39e0e20b4138 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_allreduce_op_matrix.py @@ -0,0 +1,442 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the ``m3_allreduce`` entry, 4 ranks: its cells against the entry's reference, and the +fused draft loop's layer boundaries built on it (``M3Eagle3Draft.ar_norm``: the plain all-reduce, then FlashInfer's +fused add + RMSNorm) against stock Eagle3's, bit for bit. + +Stock ``Eagle3DecoderLayer`` all-reduces each projection plainly (BF16 out) and runs the ``RMSNorm`` module with the +residual: FlashInfer's ``fused_add_rmsnorm``, which adds in FP32 and normalizes the FP32 sum. The draft loop's +boundaries must be that path exactly, under every all-reduce strategy the single-node drafter can run (AUTO, +ONESHOT, TWOSHOT, NCCL; MNNVL is multi-node only), at tp4 shapes (H = 6144, eps = 1e-6, 1 to 32 rows). The inputs are +random rows and rows built so a BF16 rounding of the residual sum would move the norm the most, including the +review's case: unit weights, sums 1 + 2^-8 on 6143 coordinates and 2 - 2^-8 on the last, both BF16 rounding +midpoints that round to even in opposite directions. + +Asserts the normalized rows and the residual sum bitwise equal to the stock path's, and reports the max ulp of each +against an fp64 reference (ulps of BF16 at the larger magnitude; a correctly rounded output is within 0.5). + +Its own 4-rank launcher, as ``_m3_mc_allreduce_norm_op_matrix.py`` is; the collected entry point is +``test_modeling_v2_m3_allreduce.py``. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_allreduce_op_matrix.py +""" + +import math +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +TOKENS = (1, 4, 8, 32) +RESULT_TAG = "[ulp]" +"""Prefix of rank 0's result lines; the pytest entry reports them.""" + +RANK = 0 +WORLD = 1 + +Inputs = Tuple[List[torch.Tensor], torch.Tensor, torch.Tensor] +"""(every rank's [m, H] BF16 partial rows, the [m, H] BF16 residual, the [H] BF16 norm weight).""" + + +# ---------------------------------------------------------------- inputs (every rank builds every rank's rows) + + +def _gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _random(seed: int, m: int, x_scale: float, r_scale: float, unit_w: bool) -> Inputs: + g = _gen(seed) + xs = [ + (torch.randn(m, HIDDEN, generator=g, device="cuda") * x_scale).bfloat16() + for _ in range(WORLD) + ] + residual = (torch.randn(m, HIDDEN, generator=g, device="cuda") * r_scale).bfloat16() + if unit_w: + w = torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16) + else: + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return xs, residual, w + + +def _on_rank0(m: int, x: torch.Tensor) -> List[torch.Tensor]: + """``x`` as rank 0's partial, zeros elsewhere (the cross-rank sum is exact).""" + return [x.bfloat16()] + [ + torch.zeros(m, HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in range(WORLD - 1) + ] + + +def _review_case(m: int) -> Inputs: + """Sums 1 + 2^-8 on the first H - 1 coordinates and 2 - 2^-8 on the last, unit weights. + + Both sums are BF16 midpoints: 1 + 2^-8 rounds (to even) down to 1 and 2 - 2^-8 up to 2, so the fused path's + norm sees a residual moved in opposite directions on the bulk and on the last coordinate. + """ + residual = torch.ones(m, HIDDEN, device="cuda") + residual[:, -1] = 2.0 - 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _review_case_flipped(m: int) -> Inputs: + """The mirror: the bulk on odd-mantissa midpoints (1 + 3 * 2^-8 rounds up), the last one rounding down.""" + residual = torch.full((m, HIDDEN), 1.0 + 2.0**-7, device="cuda") + residual[:, -1] = 2.0 - 2.0 * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _midpoints(seed: int, m: int, outliers: int) -> Inputs: + """Every sum a BF16 midpoint: residual r random in [1, 2) and x = half its ulp (+2^-8), split over the ranks. + + The bulk has an odd mantissa (rounds up); ``outliers`` coordinates per row carry an even one (rounds down) and + a larger magnitude, so they weigh in the RMS while rounding the other way. + """ + g = _gen(seed) + mant = ( + torch.randint(0, 64, (m, HIDDEN), generator=g, device="cuda") * 2 + 1 + ) # odd 7-bit mantissa + r = 1.0 + mant.float() * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + if outliers: + cols = torch.randperm(HIDDEN, generator=g, device="cuda")[:outliers] + even = torch.randint(0, 64, (m, outliers), generator=g, device="cuda") * 2 + r[:, cols] = 8.0 * (1.0 + even.float() * 2.0**-7) + x[:, cols] = 8.0 * 2.0**-8 + # Split x over the ranks in exact BF16 pieces: x/2 on rank 0, x/4 on ranks 1 and 2, 0 on rank 3. + parts = [x / 2, x / 4, x / 4] + [torch.zeros_like(x)] * (WORLD - 3) + return ( + [p.bfloat16() for p in parts[:WORLD]], + r.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _cancelling(seed: int, m: int) -> Inputs: + """Large partials that cancel across the ranks (+A, -A, +B, -B + d) onto a unit residual.""" + g = _gen(seed) + a = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + b = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + d = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 2.0**-6).bfloat16().float() + parts = [a, -a, b, (-b + d)] + residual = (1.0 + torch.rand(m, HIDDEN, generator=g, device="cuda")).bfloat16() + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return [p.bfloat16() for p in parts[:WORLD]], residual, w + + +def _cases() -> List[Tuple[str, Callable[[int], Inputs]]]: + cases: List[Tuple[str, Callable[[int], Inputs]]] = [ + ("review: 1+2^-8 x6143, 2-2^-8 x1", _review_case), + ("review mirrored", _review_case_flipped), + ("all midpoints, no outliers", lambda m: _midpoints(11, m, 0)), + ("all midpoints, 1 outlier", lambda m: _midpoints(12, m, 1)), + ("all midpoints, 64 outliers", lambda m: _midpoints(13, m, 64)), + ("cancelling partials", lambda m: _cancelling(14, m)), + ] + for i, (xs, rs, unit) in enumerate( + ((1.0, 4.0, False), (1.0, 1.0, True), (0.05, 8.0, False), (8.0, 0.5, False)) + ): + cases.append( + ( + f"random x{xs} r{rs}{' unit-w' if unit else ''}", + lambda m, i=i, xs=xs, rs=rs, unit=unit: _random(100 + i, m, xs, rs, unit), + ) + ) + return cases + + +# ---------------------------------------------------------------- the three paths + + +def _fp64(xs: List[torch.Tensor], residual: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + s = residual.double() + for x in xs: + s = s + x.double() + rms = torch.rsqrt(s.square().mean(-1, keepdim=True) + EPS) + return s * rms * w.double() + + +def _ulp_at(v: torch.Tensor) -> torch.Tensor: + """BF16's ulp at |v| (fp64): 2^(e - 8) for |v| in [2^(e-1), 2^e), floored at the smallest normal's.""" + _, e = torch.frexp(v.abs().double()) + return torch.ldexp(torch.ones_like(v, dtype=torch.float64), torch.clamp(e - 8, min=-133)) + + +def _ulps(out: torch.Tensor, ref: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the larger magnitude of the two.""" + o, r = out.double(), ref.double() + return ((o - r).abs() / _ulp_at(torch.maximum(o.abs(), r.abs()))).max().item() + + +def _ulps_at(out: torch.Tensor, ref: torch.Tensor, scale: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the magnitude ``scale`` (the operands'), not at the result's.""" + return ((out.double() - ref.double()).abs() / _ulp_at(scale.abs())).max().item() + + +def _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """The draft loop's boundary (``M3Eagle3Draft.ar_norm``): ``m3_allreduce`` plain, then the catalog's + ``flashinfer_fused_add_rmsnorm``.""" + from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 + flashinfer_fused_add_rmsnorm, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + + return flashinfer_fused_add_rmsnorm(m3_allreduce(x, ar, None), residual.clone(), w, EPS) + + +def _check_cells(mapping: Any) -> None: + """Every cell of ``m3_allreduce.CELLS`` (AUTO strategy, as the target builds it) against the entry's reference: + rank-order fp32 sums with one BF16 rounding of the sum and one of the norm. A strategy may round the cross-rank sum + to BF16 before adding the residual, so the residual sum is bounded at its operands' magnitude + u = ulp(max(|residual|, |sum|)): half an ulp for that rounding plus half an output ulp for each of the two final + roundings, an output ulp being at most 2u, so 2.5u (where the two cancel, the output's own ulp is far smaller). + The normed rows are the norm of the op's own residual sum, within 2 ulps of the reference norm of that sum.""" + import sys as _sys + from pathlib import Path as _Path + + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams + from tensorrt_llm.functional import AllReduceStrategy + + _sys.path.insert(0, str(_Path(__file__).resolve().parents[4])) + from _validating import validating + + ar = AllReduce(mapping=mapping, strategy=AllReduceStrategy.AUTO, dtype=torch.bfloat16) + for i, cell in enumerate(m3_allreduce.CELLS): + spec = cell.spec + xs, residual, w = _random(200 + i, spec["m"], 1.0, 4.0, False) + params = ( + None + if spec.get("plain") + else AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual.clone(), + norm_weight=w, + eps=EPS, + trigger_completion_at_end=False, + ) + ) + # The reference first: the fused path may write the updated residual into ``params.residual``. + ref = m3_allreduce.reference(xs, ar, params) + with validating(m3_allreduce): + out = m3_allreduce(xs[RANK], ar, params) + torch.cuda.synchronize() + total = sum(x.double() for x in xs) + if params is None: + err = _ulps_at(out, ref, total) + assert err <= 1.0, (cell.why, err) + else: + err = _ulps_at(out[1], ref[1], torch.maximum(total.abs(), residual.double().abs())) + assert err <= 2.5, (cell.why, "residual", err) + h = out[1].double() + own = h * torch.rsqrt(h.square().mean(-1, keepdim=True) + EPS) * w.double() + assert _ulps(out[0], own) <= 2.0, (cell.why, "normed", _ulps(out[0], own)) + dist.barrier() + if RANK == 0: + print(f"{RESULT_TAG} cell {cell.why} passed", flush=True) + + +def _stock(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """Stock Eagle3: ``o_proj``'s plain all-reduce (BF16 out), then ``RMSNorm(h, residual)``'s FlashInfer + ``fused_add_rmsnorm`` (FP32 add, norm of the FP32 sum; in place on both).""" + h = ar(x) + res = residual.clone() + torch.ops.trtllm.flashinfer_fused_add_rmsnorm(h, res, w, EPS) + return h, res + + +# ---------------------------------------------------------------- the matrix + + +def _single_node_cluster_info() -> None: + """``TorchDist`` (the communicator under ``TLLM_DISABLE_MPI``) asks Ray for each rank's node and GPU to build its + node-local group. These ranks are spawned without Ray on one node, so every rank reports the same node and its + own device.""" + from tensorrt_llm._torch.distributed.communicator import TorchDist + + def cluster_info(self): + if self.cluster_info is None: + self.cluster_info = [("127.0.0.1", r) for r in range(dist.get_world_size())] + return self.cluster_info + + TorchDist._get_cluster_info = cluster_info + + +def _run_matrix() -> None: + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers trtllm::flashinfer_fused_add_rmsnorm + from tensorrt_llm._torch.distributed import AllReduce + from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE + from tensorrt_llm.functional import AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + assert IS_FLASHINFER_AVAILABLE, "stock Eagle3's add + RMSNorm is FlashInfer's fused_add_rmsnorm" + _single_node_cluster_info() + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK, gpus_per_node=WORLD) + _check_cells(mapping) + strategies = ( + AllReduceStrategy.AUTO, + AllReduceStrategy.ONESHOT, + AllReduceStrategy.TWOSHOT, + AllReduceStrategy.NCCL, + ) + worst = 0.0 + mismatches: List[str] = [] + for strategy in strategies: + ar = AllReduce(mapping=mapping, strategy=strategy, dtype=torch.bfloat16) + assert ar.mnnvl_allreduce is None and ar.symm_mem_allreduce is None, strategy + for name, build in _cases(): + ulp, res_ulp = 0.0, 0.0 + for m in TOKENS: + xs, residual, w = build(m) + target_out, target_res = _target(ar, xs[RANK], residual, w) + stock_out, stock_res = _stock(ar, xs[RANK], residual, w) + torch.cuda.synchronize() + for what, a, b in ( + ("normed", target_out, stock_out), + ("residual", target_res, stock_res), + ): + bad = int((a.view(torch.int16) != b.view(torch.int16)).sum().item()) + if bad: + mismatches.append( + f"{strategy.name} {name} m={m}: {bad} {what} elements differ from stock" + ) + ref = _fp64(xs, residual, w) + ref_res = residual.double() + sum(x.double() for x in xs) + u, ur = _ulps(target_out, ref), _ulps(target_res, ref_res) + assert math.isfinite(u) and math.isfinite(ur), (strategy, name, m) + ulp, res_ulp = max(ulp, u), max(res_ulp, ur) + dist.barrier() + worst = max(worst, ulp) + if RANK == 0: + print( + f"{RESULT_TAG} {strategy.name:7s} | {name:34s} | target = stock vs fp64: normed {ulp:.3f}" + f" | residual {res_ulp:.3f}", + flush=True, + ) + if RANK == 0: + scope = f"all cases, m in {TOKENS}" + print( + f"{RESULT_TAG} MAX | {scope:34s} | target = stock vs fp64: normed {worst:.3f}", + flush=True, + ) + assert not mismatches, "\n".join(mismatches) + + +def _run_one_rank() -> int: + global RANK, WORLD + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 4, f"the drafter runs at tp4, got {WORLD} ranks" + try: + _run_matrix() + except BaseException: + import traceback + + print(f"[rank {RANK}] FAILED", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + os._exit(13) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker per claimed device, each in its own session; on the first nonzero exit or the deadline every + session is killed (as ``_m3_mc_allreduce_norm_op_matrix.py`` does).""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 4 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 4, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 4" + ) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + 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: Optional[str] = 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/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..1b3888c2a638 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,663 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``, 4 ranks. + +Its own launcher, as the catalog's collective matrices are: run plainly it spawns +one rank process per device in CUDA_VISIBLE_DEVICES and kills them all on the +first failure or on the deadline. A rank whose peers never push polls forever, +so the deadline is what keeps a broken exchange from wedging the calling run. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_mc_allreduce_norm_op_matrix.py + +Each rank builds the three multicast workspaces (8 / 16 / 48 tokens) the target +builds, over the torch.distributed group (the non-MPI path of +``_make_mnnvl_mcast_buffer``). Every rank can regenerate every rank's inputs from +the shared seed convention, which is what keeps each reference arithmetic rather +than a second collective. Producer pushes (``prepushed`` 2 / 3) are emulated by +storing the canonicalized rows through the multicast mapping at the call parity +the flags select, which is exactly what the producing kernels store. + +Not a pytest module, despite the ``check_*`` bodies: they are one fixed sequence +inside a single 4-rank job whose workspace state (call parities) carries from +one to the next. The collected entry point is +``test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py``. +""" + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Dict, List, Optional + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +CLUSTER = 8 +BUCKETS = (8, 16, 48) +EMPTY_WORD = -2147483648 + +# Bound inside the rank body: importing the entries pulls in tensorrt_llm. +allreduce_entry: Any = None +finalize_entry: Any = None +ar_module: Any = None +validating: Any = None +RANK = 0 +WORLD = 1 +WORKSPACES: Dict[int, Any] = {} + + +# ---------------------------------------------------------------- inputs + + +def _gen(seed: int, rank: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed * 1009 + rank + 1) + + +def _rows(seed: int, rank: int, m: int, scale: float = 1.0) -> torch.Tensor: + """One rank's [m, H] BF16 rows; ``rank`` -1 for tensors every rank holds alike.""" + g = _gen(seed, rank + 1) + return (torch.randn(m, HIDDEN, generator=g, device="cuda") * scale).to(torch.bfloat16) + + +def _residual(seed: int, m: int) -> torch.Tensor: + return _rows(seed + 7, -1, m, 4.0) + + +def _norm_w(seed: int) -> torch.Tensor: + g = _gen(seed + 11, 0) + return (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + +def _finalize_inputs(seed: int, rank: int, m: int, topk: int, dropped: int): + """One rank's unfinalized MoE outputs: fc2 [m * topk, H], perm [m, topk], fp32 scales, shared [m, H].""" + g = _gen(seed + 23, rank) + rows = m * topk + fc2 = (torch.randn(rows, HIDDEN, generator=g, device="cuda") * 0.5).to(torch.bfloat16) + perm = torch.randperm(rows, generator=g, device="cuda").to(torch.int32) + if dropped: + drop = torch.randperm(rows, generator=g, device="cuda")[:dropped] + perm[drop] = -1 + perm = perm.view(m, topk).contiguous() + scales = torch.rand(m, topk, generator=g, device="cuda", dtype=torch.float32) + shared = _rows(seed + 29, rank, m, 0.5) + return fc2, perm, scales, shared + + +# ---------------------------------------------------------------- workspace / producer emulation + + +def _make_workspace(maxt: int): + """One bucket's exchange, as ``McAllReduceWorkspace.create`` builds it, over the torch.distributed group.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + from tensorrt_llm.mapping import Mapping + + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK) + words = 4 * maxt * WORLD * (HIDDEN // 2) + group = dist.group.WORLD + handle = _make_mnnvl_mcast_buffer(group, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar_module.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=group, + max_tokens=maxt, + ) + ws.flags_for(CLUSTER) + return ws + + +def _push(ws, rows: torch.Tensor, exchange_set: int) -> None: + """Store this rank's rows into every rank's exchange, as a producer kernel does before a prepushed call. + + Slot of (token t, cluster CTA c): set * 2 * half + parity(t, c) * half + (t * world + rank) * H / 2 + c * H / 16 + words, parity read from the CTA's flag; -0.0 halves are stored as +0.0 (0x80000000 is the empty word). + """ + m = rows.shape[0] + torch.cuda.synchronize() + parity = (ws.flags_for(CLUSTER).view(ws.max_tokens, CLUSTER)[:m] & 1).long() + half = ws.max_tokens * WORLD * (HIDDEN // 2) + chunk = HIDDEN // 2 // CLUSTER + h16 = rows.contiguous().view(torch.int16) + h16 = torch.where(h16 == -32768, torch.zeros_like(h16), h16) + words = h16.view(torch.int32).view(m, CLUSTER, chunk) + t = torch.arange(m, device="cuda").view(m, 1, 1) + c = torch.arange(CLUSTER, device="cuda").view(1, CLUSTER, 1) + w = torch.arange(chunk, device="cuda").view(1, 1, chunk) + index = ( + exchange_set * 2 * half + + parity.view(m, CLUSTER, 1) * half + + (t * WORLD + RANK) * (HIDDEN // 2) + + c * chunk + + w + ) + ws.mc[index.reshape(-1)] = words.reshape(-1) + torch.cuda.synchronize() + dist.barrier() + + +# ---------------------------------------------------------------- gates + + +def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + assert out.dtype == ref.dtype and out.shape == ref.shape, ( + where, + out.dtype, + ref.dtype, + out.shape, + ) + a, b = out.contiguous().view(torch.uint8), ref.contiguous().view(torch.uint8) + bad = (a != b).sum().item() + assert bad == 0, f"{where}: {bad} of {a.numel()} bytes differ" + + +#: Share of normed elements one BF16 step off the float64 reference: measured at most 1.1e-4 per call over +#: every call of this matrix on GB300 (see ``_assert_normed``). The gate sits ~18x above that and far below what a +#: wrong statistic produces (most elements move). +NORMED_STEP_FRACTION = 0.002 +NORMED_STEP_SEEN: List[float] = [] + + +def _assert_normed(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + """At most one BF16 rounding step from the float64 reference, and rarely that. + + The kernel's RMS statistic is a shuffle / warp / cluster tree of FP32 partial sums and its ``rsqrtf`` is + approximate (~2^-22 relative): the scaled value moves by ~1e-6 relative, which changes the BF16 rounding only + for values within that distance of a rounding midpoint -- one step, never two. The fraction bound catches a + statistic that is systematically off (a missing rank, a wrong H), which moves most elements. + """ + assert out.dtype == ref.dtype == torch.bfloat16 and out.shape == ref.shape, where + o, r = out.float(), ref.float() + _, exponent = torch.frexp(torch.maximum(o.abs(), r.abs())) + step = torch.ldexp(torch.ones_like(o), exponent - 8) + diff = (o - r).abs() + assert bool((diff <= step).all()), f"{where}: max diff {diff.max().item()} beyond one BF16 step" + frac = (diff != 0).float().mean().item() + NORMED_STEP_SEEN.append(frac) + assert frac <= NORMED_STEP_FRACTION, f"{where}: {frac:.4f} of normed elements off by one step" + + +def _assert_quant(outs, ref, m: int, where: str) -> None: + """The MXFP8 copy: bit-exact against the recipe applied to the op's own normed rows (a pure function of + them), and equal to the reference's wherever the reference's normed 32-column block is the op's.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + mxfp8_rows, + sf_written, + ) + + q, sf = outs[2], outs[3] + own_q, own_sf = mxfp8_rows(outs[0]) + written = sf_written(m, HIDDEN // 32, sf.device) + _assert_bitwise(q, own_q, f"{where} e4m3 rows") + _assert_bitwise(sf[written], own_sf[written], f"{where} UE8M0 scales") + same_block = (outs[0].view(m, -1, 32) == ref[0].view(m, -1, 32)).all(-1) + qb = q.view(torch.uint8).view(m, -1, 32) + rqb = ref[2].view(torch.uint8).view(m, -1, 32) + assert bool((qb == rqb).all(-1)[same_block].all()), ( + f"{where}: e4m3 differs on a block whose normed rows agree" + ) + + +def _check_outputs(outs, ref, m: int, quant: bool, where: str) -> None: + assert len(outs) == (4 if quant else 2), (where, len(outs)) + _assert_bitwise(outs[1], ref[1], f"{where} updated residual") + _assert_normed(outs[0], ref[0], f"{where} normed") + if quant: + _assert_quant(outs, ref, m, where) + + +# ---------------------------------------------------------------- calls + + +def _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, opts): + return allreduce_entry.m3_mc_allreduce_norm( + x, + x2, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + ws.max_tokens, + quant, + prepushed, + opts, + ) + + +def _plain_case(spec: Dict[str, Any], seed: int, where: str) -> None: + m, maxt, quant, prepushed = spec["m"], spec["maxt"], spec["quant"], spec["prepushed"] + ws = WORKSPACES[maxt] + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] if spec["x2"] else None + residual, norm_w = _residual(seed, m), _norm_w(seed) + x, x2 = xs[RANK], (None if x2s is None else x2s[RANK]) + if prepushed >= 2: + _push(ws, x, 0) + if prepushed == 3: + _push(ws, x2, 1) + with validating(allreduce_entry.m3_mc_allreduce_norm): + outs = _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, spec["opts"]) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference(xs, x2s, residual, norm_w, EPS, quant) + _check_outputs(outs, ref, m, quant, where) + dist.barrier() + + +# ---------------------------------------------------------------- checks + + +def check_plain_cells() -> None: + """Every cell of ``m3_mc_allreduce_norm.CELLS`` through ``__call__`` under ``validating``, gated on its + ``reference``.""" + for i, cell in enumerate(allreduce_entry.m3_mc_allreduce_norm.CELLS): + _plain_case(cell.spec, 1000 + 17 * i, cell.why) + + +def check_finalize_cells() -> None: + """Every cell of ``m3_mc_finalize_allreduce_norm.CELLS``, gated on its ``reference``.""" + op = finalize_entry.m3_mc_finalize_allreduce_norm + for i, cell in enumerate(op.CELLS): + spec = cell.spec + m, maxt, quant, topk = spec["m"], spec["maxt"], spec["quant"], spec["topk"] + seed = 2000 + 17 * i + ws = WORKSPACES[maxt] + per_rank = [_finalize_inputs(seed, r, m, topk, spec["dropped"]) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + with validating(op): + outs = op( + fc2, + perm, + scales, + shared, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + maxt, + quant, + spec["opts"], + ) + torch.cuda.synchronize() + ref = op.reference( + [p[0] for p in per_rank], + [p[1] for p in per_rank], + [p[2] for p in per_rank], + [p[3] for p in per_rank], + residual, + norm_w, + EPS, + quant, + ) + _check_outputs(outs, ref, m, quant, cell.why) + dist.barrier() + + +def check_the_target_wrappers_select_bucket_and_options() -> None: + """The target's own wrappers (bucket by token count, option table) on top of the entries, every token count of + the 8-token bucket plus one of each larger bucket, and the finalize wrapper; the MXFP8 publication is the + prequant slot's business and is left off.""" + buckets = ar_module.McBuckets([WORKSPACES[b] for b in BUCKETS]) + for m in (*range(1, 9), 11, 40): + seed = 3000 + m + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + normed, updated = ar_module.m3_mc_allreduce_norm( + xs[RANK], residual, norm_w, EPS, buckets, x2=x2s[RANK] + ) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference(xs, x2s, residual, norm_w, EPS, False) + _check_outputs([normed, updated], ref, m, False, f"wrapper m={m}") + dist.barrier() + for m in (9, 33): + seed = 3100 + m + per_rank = [_finalize_inputs(seed, r, m, 4, 2) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + normed, updated = ar_module.m3_mc_finalize_allreduce_norm( + fc2, perm, scales, shared, residual, norm_w, EPS, buckets + ) + torch.cuda.synchronize() + ref = finalize_entry.m3_mc_finalize_allreduce_norm.reference( + *[[p[j] for p in per_rank] for j in range(4)], residual, norm_w, EPS, False + ) + _check_outputs([normed, updated], ref, m, False, f"finalize wrapper m={m}") + dist.barrier() + + +def check_cuda_graph_replays() -> None: + """Two boundaries captured in one graph (8- and 16-token buckets), replayed with fresh inputs. + + The decode graph holds every boundary call; each replay rotates the same per-CTA parities the eager calls do, + so a replay that reused a stale half or skipped a re-arm reads the previous replay's rows and fails the + bitwise residual gate (every replay gets new inputs). + """ + shapes = ((3, 8, 15), (12, 16, 15)) + norm_w = _norm_w(4000) + bufs = [] + for m, maxt, opts in shapes: + bufs.append( + dict( + m=m, + maxt=maxt, + opts=opts, + x=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + x2=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + residual=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + ) + ) + + def body(): + return [ + _call_plain( + WORKSPACES[b["maxt"]], b["x"], b["x2"], b["residual"], norm_w, False, 0, b["opts"] + ) + for b in bufs + ] + + def fill(seed: int): + refs = [] + for j, b in enumerate(bufs): + s = seed + 31 * j + xs = [_rows(s, r, b["m"]) for r in range(WORLD)] + x2s = [_rows(s + 3, r, b["m"]) for r in range(WORLD)] + b["x"].copy_(xs[RANK]) + b["x2"].copy_(x2s[RANK]) + b["residual"].copy_(_residual(s, b["m"])) + refs.append( + allreduce_entry.m3_mc_allreduce_norm.reference( + xs, x2s, b["residual"].clone(), norm_w, EPS, False + ) + ) + torch.cuda.synchronize() + dist.barrier() + return refs + + refs = fill(4100) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + warm = body() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(warm[j], refs[j], b["m"], False, f"graph warm-up m={b['m']}") + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = body() + for rep in range(4): + refs = fill(4200 + 100 * rep) + graph.replay() + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(captured[j], refs[j], b["m"], False, f"replay {rep} m={b['m']}") + dist.barrier() + del graph + dist.barrier() + + +def check_the_gate_sees_a_missing_rank() -> None: + """The residual gate is not blind to the reduction: dropping any one rank's partial from the reference + moves most elements of the updated residual (each rank's rows are distinct by construction).""" + m, seed = 4, 5000 + xs = [_rows(seed, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + full = allreduce_entry.m3_mc_allreduce_norm.reference(xs, None, residual, norm_w, EPS, False) + for r in range(WORLD): + partial = [x if i != r else torch.zeros_like(x) for i, x in enumerate(xs)] + wrong = allreduce_entry.m3_mc_allreduce_norm.reference( + partial, None, residual, norm_w, EPS, False + ) + frac = (wrong[1] != full[1]).float().mean().item() + assert frac > 0.9, f"dropping rank {r} changes only {frac:.3f} of the residual" + dist.barrier() + + +def _expect_rejected(fn, match: str, where: str) -> None: + try: + fn() + except AssertionError as exc: + assert match in str(exc), f"{where}: rejected for the wrong reason: {exc}" + else: + raise AssertionError(f"{where}: the guard admitted it") + + +def check_guards() -> None: + """``is_valid`` refuses what the op accepts and answers wrongly. Nothing here reaches the kernel: the guard + raises first, on every rank alike, so no rank is left polling.""" + m, seed = 4, 6000 + ws8, ws16 = WORKSPACES[8], WORKSPACES[16] + x, residual, norm_w = _rows(seed, RANK, m), _residual(seed, m), _norm_w(seed) + op = allreduce_entry.m3_mc_allreduce_norm + + def call(ws_uc, flags, norm): + with validating(op): + op( + x, + None, + residual, + norm, + EPS, + ws_uc, + ws8.mc.data_ptr(), + flags, + RANK, + WORLD, + CLUSTER, + False, + 8, + False, + 0, + 15, + ) + + strided = torch.empty(2 * HIDDEN, dtype=torch.bfloat16, device="cuda")[::2] + strided.copy_(norm_w) + _expect_rejected( + lambda: call(ws8.uc, ws8.flags_for(CLUSTER), strided), "contiguous", "strided norm_w" + ) + _expect_rejected( + lambda: call(ws8.uc[: ws8.uc.numel() // 2], ws8.flags_for(CLUSTER), norm_w), + "words", + "short workspace", + ) + _expect_rejected( + lambda: call(ws8.uc, ws16.flags_for(CLUSTER), norm_w), "counters", "another bucket's flags" + ) + + fin = finalize_entry.m3_mc_finalize_allreduce_norm + fc2, perm, scales, shared = _finalize_inputs(seed, RANK, 9, 4, 0) + bad = perm.clone() + bad[0, 0] = fc2.shape[0] + + def call_fin(): + with validating(fin): + fin( + fc2, + bad, + scales, + shared, + _residual(seed, 9), + norm_w, + EPS, + ws16.uc, + ws16.mc.data_ptr(), + ws16.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + 16, + False, + 15, + ) + + _expect_rejected(call_fin, "perm spans", "perm past the FC2 rows") + dist.barrier() + + +CHECKS = ( + check_plain_cells, + check_finalize_cells, + check_the_target_wrappers_select_bucket_and_options, + check_cuda_graph_replays, + check_the_gate_sees_a_missing_rank, + check_guards, +) + + +def _run_one_rank() -> int: + """Body of one spawned rank: run every check; a failing rank exits nonzero and the launcher kills its + siblings.""" + global RANK, WORLD, allreduce_entry, finalize_entry, ar_module, validating + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 4, f"the entries are certified at 4 ranks, got {WORLD}" + + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.m3_mc_allreduce as _ar # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.comm import ( # noqa: E501 + m3_mc_allreduce_norm as _plain, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.comm import ( # noqa: E501 + m3_mc_finalize_allreduce_norm as _finalize, + ) + + allreduce_entry, finalize_entry, ar_module = _plain, _finalize, _ar + # Spawned by file path, not collected by pytest: reach the sibling helper the plain way (confined to the + # rank processes). /comm -> tests/unittest/_torch/modeling_v2. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + validating = _validating_cm + for maxt in BUCKETS: + WORKSPACES[maxt] = _make_workspace(maxt) + + 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 early wedges its peers' polls; the launcher kills them on this exit. + os._exit(13) + print(f"[rank {RANK}] {check.__name__} passed", flush=True) + dist.barrier() + torch.cuda.synchronize() + print( + f"[rank {RANK}] {len(CHECKS)} checks passed; normed one-step fraction max " + f"{max(NORMED_STEP_SEEN):.5f} over {len(NORMED_STEP_SEEN)} calls", + flush=True, + ) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker process per claimed device, each in its own session; on the first nonzero exit or on the + deadline every session is killed.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 4 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 4, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 4" + ) + 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", + ) + # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed + # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. + 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: Optional[str] = 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/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_rank_job.py new file mode 100644 index 000000000000..eeeb202ae093 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_rank_job.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run this target's collective rank job from pytest, and report what it did. + +The same contract as ``tests/unittest/_torch/modeling_v2/comm/_rank_job.py``: the +launcher beside this file owns the per-device rank processes and the deadline +that kills a wedged collective; this only selects the devices, starts it in a +fresh interpreter (by file path: this tree is not a package) and turns its exit +code into an assertion. +""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +WORLD_SIZE = 4 +"""tp4's world size -- the topology these entries are certified for.""" + +_LAUNCHER_GRACE_S = 300 +"""Headroom over the launcher'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 launcher's module-level budget without importing it by name (it imports only torch at + module scope).""" + 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) -> str: + """Run ``__op_matrix``'s launcher over ``WORLD_SIZE`` devices; returns its stdout.""" + launcher = Path(__file__).resolve().with_name(f"_{entry}_op_matrix.py") + env = dict(os.environ, CUDA_VISIBLE_DEVICES=_devices()) + + # The ranks import tensorrt_llm absolutely, and a source checkout is not necessarily installed. + # tests/unittest/_torch/modeling_v2/models/minimax_m3//comm -> repo root. + repo_root = Path(__file__).resolve().parents[8] + assert (repo_root / "tensorrt_llm").is_dir(), ( + f"expected the repo root at {repo_root}, found no tensorrt_llm/ there; " + f"this file moved without its parents[] index following" + ) + env["PYTHONPATH"] = os.pathsep.join(p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p) + + constants = _load_launcher_constants(launcher) + timeout = constants.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}" + ) + return completed.stdout diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_allreduce.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_allreduce.py new file mode 100644 index 000000000000..4f34ae908c81 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_allreduce.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the ``m3_allreduce`` certification matrix. + +The matrix is ``_m3_allreduce_op_matrix.py`` beside this file, its own 4-rank launcher; see ``_rank_job``. It drives +every cell of the entry against its reference, then requires the fused draft loop's boundaries built on it (plain +all-reduce, then the fused add + RMSNorm) to equal stock Eagle3's bit for bit under every strategy, and reports both +against an fp64 reference as a warning, one line per strategy and case. +""" + +import warnings + +import _rank_job +import pytest +import torch + +assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" + +_RESULT_TAG = "[ulp]" + + +# Spawns 4 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_allreduce_op_matrix() -> None: + out = _rank_job.run("m3_allreduce") + lines = [line for line in out.splitlines() if line.startswith(_RESULT_TAG)] + assert lines, f"the launcher reported no measurements:\n{out}" + warnings.warn( + "draft boundary (= stock Eagle3) vs fp64, max BF16 ulps\n" + "\n".join(lines), stacklevel=1 + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..31a054642e89 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the multicast all-reduce + RMSNorm entries' certification matrix. + +Covers both ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm`` (one +kernel, one exchange protocol). The matrix is ``_m3_mc_allreduce_norm_op_matrix.py`` +beside this file, its own 4-rank launcher; see ``_rank_job`` for why. +""" + +import _rank_job +import pytest +import torch + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + + +# Spawns 4 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_mc_allreduce_norm_op_matrix() -> None: + _rank_job.run("m3_mc_allreduce_norm") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..4d7b4f235c59 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests of the ``m3_mc_finalize_allreduce_norm`` entry that need no collective. + +Its cells are certified, with ``m3_mc_allreduce_norm``'s (one kernel, one exchange protocol), by the 4-rank matrix +``_m3_mc_allreduce_norm_op_matrix.py`` that ``test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py`` runs. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.comm.m3_mc_finalize_allreduce_norm import ( # noqa: E501 + finalize_partial, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.comm.m3_mc_finalize_allreduce_norm import ( # noqa: E501 + m3_mc_finalize_allreduce_norm as op, +) + + +def test_finalize_partial_rounds_the_routed_sum_before_the_shared_expert() -> None: + """As stock: the finalized routed output is BF16, and the shared expert is added to it (shared.add_(routed)). + Four routed rows of 1.0 at FP32 routing weight 0.2509765625 sum to 1.00390625, which rounds to BF16 1.0; plus the + shared -1 that is 0, where adding before the rounding would leave 0.00390625.""" + fc2 = torch.ones((4, 8), dtype=torch.bfloat16) + perm = torch.arange(4, dtype=torch.int32).view(1, 4) + scales = torch.full((1, 4), 0.2509765625, dtype=torch.float32) + shared = torch.full((1, 8), -1.0, dtype=torch.bfloat16) + out = finalize_partial(fc2, perm, scales, shared) + routed = finalize_partial(fc2, perm, scales, None) + assert torch.equal(out, shared.clone().add_(routed)) + assert torch.equal(out, torch.zeros(1, 8, dtype=torch.bfloat16)) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_each_cells_partial_is_the_stock_finalize_then_shared_add(cell) -> None: + """Per cell shape (rows, top-k, dropped slots), every rank's partial in the entry's reference is the stock MoE's: + the finalized routed rows rounded to BF16, then ``shared.add_(routed)``. The 4-rank matrix drives the cells on the + GPU against this reference.""" + spec = cell.spec + g = torch.Generator().manual_seed(spec["m"] * 31 + spec["dropped"]) + m, topk, hidden = spec["m"], spec["topk"], 64 + fc2 = torch.randn(m * topk, hidden, generator=g).to(torch.bfloat16) + perm = torch.randperm(m * topk, generator=g).view(m, topk).to(torch.int32) + perm.view(-1)[: spec["dropped"]] = -1 + scales = torch.rand(m, topk, generator=g) + shared = torch.randn(m, hidden, generator=g).to(torch.bfloat16) + routed = finalize_partial(fc2, perm, scales, None) + assert torch.equal(finalize_partial(fc2, perm, scales, shared), shared.clone().add_(routed)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/embedding/test_modeling_v2_m3_embedding.py new file mode 100644 index 000000000000..39cedf0fff84 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/embedding/test_modeling_v2_m3_embedding.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_embedding entry of the MiniMax-M3 tp4 target: an exact gather with zero rows outside the +table.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.embedding.m3_embedding import ( # noqa: E501 + m3_embedding as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_embedding requires a CUDA device" + +VOCAB, HIDDEN = 4096, 6144 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + table = torch.randn(VOCAB, HIDDEN, generator=g, device="cuda").to(torch.bfloat16) + ids = torch.randint( + 0, VOCAB, (spec["rows"],), generator=g, device="cuda", dtype=spec["id_dtype"] + ) + if spec["pad"]: + ids[1::3] = -1 + with validating(op): + out = op(ids, table) + assert torch.equal(out, op.reference(ids, table)) + if spec["pad"]: + assert not out[1::3].any(), "padded ids must embed to zero rows" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/_m3_mxfp8_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/_m3_mxfp8_operands.py new file mode 100644 index 000000000000..e44ad56d743d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/_m3_mxfp8_operands.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 operands for the MXFP8 GEMM entries' tests: E4M3 values and 128x4-swizzled UE8M0 block scales.""" + +import torch + + +def swizzle(byte: torch.Tensor) -> torch.Tensor: + """[N, K / 32] scale bytes into the 128x4-swizzled buffer of ``mxfp8_quantize(w, True)``.""" + rows, cols = byte.shape + padded_cols, padded_rows = (cols + 3) // 4 * 4, (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + + (c // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + out = torch.zeros(padded_rows * padded_cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def mxfp8(rows: int, k: int, g: torch.Generator, lo: int = 118, hi: int = 125): + """An MXFP8 operand: E4M3 values of a unit normal, UE8M0 block scales 2^(lo-127) .. 2^(hi-128).""" + v = torch.randn(rows, k, generator=g, device="cuda").to(torch.float8_e4m3fn) + sf = torch.randint(lo, hi, (rows, k // 32), generator=g, device="cuda", dtype=torch.int32) + return v, swizzle(sf.to(torch.uint8)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_bf16_linear.py new file mode 100644 index 000000000000..08349a6dbe71 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_bf16_linear.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_bf16_linear entry of the MiniMax-M3 tp4 target (the drafter's projections).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm.m3_bf16_linear import ( # noqa: E501 + m3_bf16_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_bf16_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Accumulation order (cuBLAS split-K) then one bf16 rounding on both sides: at most 2 units of ULP[bf16] of the + row, RMS well under one; measured on GB300: worst 0.5 unit.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w = (torch.randn(spec["n"], spec["k"], generator=g, device="cuda") * 0.02).to(torch.bfloat16) + with validating(op): + out = op(x, w) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(x, w), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_fp32_linear.py new file mode 100644 index 000000000000..a008ead936da --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_fp32_linear.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp32_linear entry of the MiniMax-M3 tp4 target (the router beyond the decode routers' rows).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm.m3_fp32_linear import ( # noqa: E501 + m3_fp32_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp32_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """torch's default fp32 matmul precision on this hardware may be TF32 (~2^-11 relative per product): the band is + stated in fp32 ulps of each row's largest logit and covers it (K = 6144 TF32 products: ~2^13 ulps worst element); + measured on GB300 with the engine's defaults: worst element ~1700 ulps, RMS ~300.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = (torch.randn(spec["m"], spec["k"], generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(spec["e"], spec["k"], generator=g, device="cuda") * 0.02 + with validating(op): + out = op(x, w) + assert out.dtype == torch.float32 and out.shape == (spec["m"], spec["e"]) + assert_within_ulp(out, op.reference(x, w), element_ulp=8192.0, rms_ulp=2048.0) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mx_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mx_gemv.py new file mode 100644 index 000000000000..88e3d562d94a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mx_gemv.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv entry: the decode MXFP8 GEMV on a producer-quantized activation. + +Operands as the target holds them: the weight through mxfp8_quantize with the swizzled scale layout (the MXFP8 linear +method's storage), the activation through mxfp8_quantize with the swizzled layout (what the boundary all-reduce and +the attention merge publish). +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm.m3_mx_gemv import ( # noqa: E501 + m3_mx_gemv as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mx_gemv requires a CUDA device" + + +# Element band, in ULP units of the row scale: see compare. +ELEMENT_ULP = 1.5 + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Element ELEMENT_ULP of the row scale, relative RMS 0.1 ulp. + + Kernel and reference multiply bit-identical dequantized operands; the kernel sums in fp32 (tensor-core blocks, + then the split-K partials), the reference in fp64. Over K <= 6144 that order difference is ~1e-6 relative, so the + two bf16 roundings disagree only where the sum sits within that of a rounding boundary: one bf16 ulp of that + element, which for an element in its row maximum's binade is up to two of ``assert_within_ulp``'s units of the row + scale, and a few in ten thousand elements in aggregate. Any fp32 sum shows it: torch's IEEE fp32 matmul of the + same dequantized operands measured at most 1.22 units against the fp64 reference over every row count 1..48 on + the o_proj and the QKV (6 seeds each, GB300), the kernel at most 1.25, both at most 0.009 ulp RMS. The band + leaves the margin between that and a full extra ulp. + """ + assert_within_ulp(out, ref, element_ulp=ELEMENT_ULP, rms_ulp=0.1) + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + return x, xq.view(m, k), xsf, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + _, xq, xsf, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(xq, xsf, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + compare(out, op.reference(xq, xsf, w, wsf)) + + +def test_an_unserved_shape_returns_none() -> None: + """Shapes outside the tuned table fall back at the call site; the op says so rather than computing.""" + _, xq, xsf, w, wsf = build(256, 512, 1, seed=1) + assert op(xq, xsf, w, wsf) is None + + +def test_guard_refuses_linear_activation_scales() -> None: + """A linear [M, K / 32] scale buffer is read as the swizzled one.""" + x, _, _, w, wsf = build(6144, 2048, 4, seed=2) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, False) + with pytest.raises(AssertionError, match="swizzled"): + with validating(op): + op(xq.view(4, 2048), xsf, w, wsf) + + +def test_guard_refuses_a_strided_weight() -> None: + _, xq, xsf, w, wsf = build(6144, 2048, 1, seed=3) + wide = torch.zeros(6144, 4096, dtype=w.dtype, device="cuda") + wide[:, :2048] = w + with pytest.raises(AssertionError, match="contiguous e4m3"): + with validating(op): + op(xq, xsf, wide[:, :2048], wsf) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..8ed3fc6df4a3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv_bf16 entry: the decode MXFP8 GEMV quantizing a bf16 activation in its prologue. + +The weight as the MXFP8 linear method stores it (mxfp8_quantize, swizzled scales). Besides the reference band, the +prologue's quantization is held to mxfp8_quantize itself: the GEMV on a bf16 x must equal, bit for bit, the +prequantized GEMV on mxfp8_quantize(x) -- the claim the target relies on when it mixes the two paths. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm.m3_mx_gemv import ( # noqa: E501 + m3_mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm.m3_mx_gemv_bf16 import ( # noqa: E501 + m3_mx_gemv_bf16 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mx_gemv_bf16 requires a CUDA device" + +# The element band of test_modeling_v2_m3_mx_gemv.compare (derived there): the same fp32 sums against fp64. +ELEMENT_ULP = 1.5 + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + return x, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + x, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(x, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + # The band of test_modeling_v2_m3_mx_gemv.compare: the quantized operands are the reference's exactly (checked + # below), the sums differ only in order. + assert_within_ulp(out, op.reference(x, w, wsf), element_ulp=ELEMENT_ULP, rms_ulp=0.1) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + pre = m3_mx_gemv(xq.view(m, k), xsf, w, wsf, use_pdl=True) + assert pre is not None and torch.equal(out.view(torch.int16), pre.view(torch.int16)), ( + "the prologue's MXFP8 quantization differs from mxfp8_quantize" + ) + + +def test_a_strided_activation_is_declined() -> None: + """The prologue reads x densely; the op declines a strided x (None) rather than computing on it.""" + x, w, wsf = build(6144, 2048, 2, seed=1) + wide = torch.zeros(2, 4096, dtype=torch.bfloat16, device="cuda") + wide[:, :2048] = x + assert op(wide[:, :2048], w, wsf) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py new file mode 100644 index 000000000000..043b4da71ff9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_gemm entry of the MiniMax-M3 tp4 target (the decode-graph GEMM on a prequantized +activation).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm.m3_mxfp8_gemm import ( # noqa: E501 + m3_mxfp8_gemm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_gemm requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Same E4M3 operands and power-of-two scales on both sides; accumulation order then one bf16 rounding: at most + 2 units of ULP[bf16] of the row (the low-M GEMM's band), RMS well under one.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + act, act_sf = mxfp8(spec["m"], spec["k"], g) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + with validating(op): + out = op(act, act_sf, w, w_sf, torch.bfloat16, False) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp( + out, + op.reference(act, act_sf, w, w_sf, torch.bfloat16, False), + element_ulp=2.0, + rms_ulp=0.25, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py new file mode 100644 index 000000000000..221cee463aa1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_linear entry of the MiniMax-M3 tp4 target: MXFP8 quantize + GEMM on each backend the +engine selects (native CUTLASS, FlashInfer, the per-bucket tuned quantize + GEMM).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm.m3_mxfp8_linear import ( # noqa: E501 + m3_mxfp8_linear as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_linear requires a CUDA device" + +_FLAGS = { + "native": (False, False, False, False), + "flashinfer": (False, False, True, False), + "tuned": (True, False, False, False), +} + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """The activation's MXFP8 round trip is exact on both sides (same block rule); then the accumulation order and + one bf16 rounding: at most 2 units of ULP[bf16] of the row, RMS well under one.""" + spec = cell.spec + if ( + spec["path"] == "flashinfer" + and getattr(torch.ops.trtllm, "flashinfer_mm_mxfp8", None) is None + ): + pytest.skip("this build has no trtllm::flashinfer_mm_mxfp8") + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + alpha = torch.ones(1, dtype=torch.float32, device="cuda") + args = (x, w, w_sf, alpha, torch.bfloat16, *_FLAGS[spec["path"]]) + with validating(op): + out = op(*args) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_target_head.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_target_head.py new file mode 100644 index 000000000000..69608913f82e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_target_head.py @@ -0,0 +1,248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_target_head entry: the target lm_head shard GEMV over the lossless 12-bit coding of its +BF16 weight, against its reference, the in-tree BF16 head GEMM (``F.linear``) and fp64. + +* The coding round-trips bit for bit (code values and exception lists), and a corrupted coding fails the check. +* Integer-valued data (fp32 sums exact in any order): the kernel's logits equal the in-tree GEMM's bit for bit. +* Random data: the kernel's error against fp64 is the in-tree GEMM's (the same products, another fp32 order, one + BF16 rounding), its argmax is the in-tree one and its result is the same run to run. +""" + +import zlib +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops import ( # noqa: E501 + m3_target_head as target_head_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm.m3_target_head import ( # noqa: E501 + CTA_WARPS, + HIDDEN, + TARGET_SHARD, + decode, + encode_chunked, + round_trip_ok, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm.m3_target_head import ( # noqa: E501 + m3_target_head as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_target_head requires a CUDA device" + + +def _seed(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def _num_sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def _weight(n: int, k: int, seed: int, exceptions: bool = False) -> torch.Tensor: + """A [n, k] head of trained-weight scale; with ``exceptions``, values on both sides of the exponent window, exact + zeros and a large outlier in many tiles.""" + g = torch.Generator(device="cuda").manual_seed(seed) + w = torch.randn(n, k, device="cuda", generator=g) * 0.02 + if exceptions: + w[::97, ::131] = 1e-9 + w[::89, 7] = 0.0 + w[3, 5] = 7.5 + w[n - 1, k - 1] = -3.0 + return w.to(torch.bfloat16) + + +_CODED: dict = {} + + +def _coded(n: int, k: int, seed: int, exceptions: bool): + """(weight, encode_chunked(weight)), built once per shape.""" + key = (n, k, seed, exceptions) + if key not in _CODED: + _CODED.clear() + w = _weight(n, k, seed, exceptions) + _CODED[key] = (w, encode_chunked(w)) + return _CODED[key] + + +def _args(coded, x: torch.Tensor, n: int, grid: int): + packed, base, exc_off, exc_pos, exc_w = coded + out = torch.full((x.shape[0], n), float("nan"), dtype=torch.bfloat16, device="cuda") + part = torch.zeros(grid * CTA_WARPS * 256, dtype=torch.float32, device="cuda") + counter = torch.zeros((n + 15) // 16, dtype=torch.int32, device="cuda") + return (packed, x, out, part, counter, exc_off, exc_pos, exc_w.view(-1), base, n, grid, True) + + +@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 this name + """The kernel against its reference and the in-tree BF16 GEMM: the error against fp64 is the in-tree GEMM's + (both are one BF16 rounding of an fp32 sum of the same exact products), the argmax is the in-tree one, nearly + every logit is the in-tree one bit for bit, and a second call gives the same bits.""" + spec = cell.spec + m, n, k = spec["m"], spec["n"], spec["k"] + exceptions = spec.get("exceptions", False) + w, coded = _coded(n, k, 1, exceptions) + if exceptions: + assert coded[3].numel() > 0, "the cell must exercise the exception lists" + g = torch.Generator(device="cuda").manual_seed(_seed(cell.why)) + x = torch.randn(m, k, device="cuda", generator=g).to(torch.bfloat16) + args = _args(coded, x, n, _num_sms()) + certifying(op) + ref = op.reference(*args) + out = op(*args).clone() + again = op(*_args(coded, x, n, _num_sms())).clone() + torch.cuda.synchronize() + assert not bool(args[4].any()), "the split-tile counters are not re-armed" + assert torch.equal(out, again), "the split-tile reduction is not deterministic" + stock = torch.nn.functional.linear(x, w) + exact = x.double() @ w.double().t() + # Each side: one BF16 rounding (<= half an ulp of the logit) of an fp32 sum whose order error is far below it. + err_out = (out.double() - exact).abs().max() + err_stock = (stock.double() - exact).abs().max() + err_ref = (ref.double() - exact).abs().max() + assert err_out <= 1.01 * max(err_stock, err_ref) + 1e-6, (err_out, err_stock, err_ref) + assert (out == stock).float().mean() > 0.99 + assert (out == ref).float().mean() > 0.99 + assert torch.equal(out.float().argmax(-1), stock.float().argmax(-1)) + + +@pytest.mark.parametrize("m", [1, 4, 8]) +def test_integer_data_is_the_bf16_head_bit_for_bit(m) -> None: + """Small-integer weights and activations: every fp32 sum is exact in any order, so the coded head's logits are + the in-tree BF16 GEMM's bit for bit (exceptions included: zeros and integers outside the exponent window).""" + g = torch.Generator(device="cuda").manual_seed(_seed(f"integer {m}")) + n, k = 4096, HIDDEN + w = torch.randint(-3, 4, (n, k), generator=g, device="cuda").to(torch.bfloat16) + w[::7, ::5] = 96.0 # far above the window of the {1, 2, 3} exponents: exceptions + x = torch.randint(-2, 3, (m, k), generator=g, device="cuda").to(torch.bfloat16) + coded = encode_chunked(w) + assert coded[3].numel() > 0 + assert round_trip_ok(*coded, w) + with validating(op): + out = op(*_args(coded, x, n, _num_sms())) + assert torch.equal(out, torch.nn.functional.linear(x, w)) + + +def test_the_head_writes_only_the_shards_columns() -> None: + """Integer data over a 4000-row shard: every logit is the in-tree GEMM's bit for bit, and the columns past n of a + wider output are left untouched.""" + n, k, m = 4000, 1024, 2 + g = torch.Generator(device="cuda").manual_seed(_seed("partial last tile")) + # Small integers: every fp32 sum is exact in any order, so the logits are comparable bit for bit. + w = torch.randint(-3, 4, (n, k), generator=g, device="cuda").to(torch.bfloat16) + w[::7, ::5] = 96.0 # exceptions, including in the partial tile + x = torch.randint(-2, 3, (m, k), generator=g, device="cuda").to(torch.bfloat16) + coded = encode_chunked(w) + assert round_trip_ok(*coded, w) + args = list(_args(coded, x, n, _num_sms())) + args[2] = torch.full((m, n + 16), float("nan"), dtype=torch.bfloat16, device="cuda") + with validating(op): + op(*args) + out = args[2] + assert torch.isnan(out[:, n:]).all(), "the head wrote past the shard's last column" + stock = torch.nn.functional.linear(x, w) + assert torch.equal(out[:, :n], stock) + + +def test_a_shard_of_partial_tiles_keeps_the_bf16_head() -> None: + """The kernel streams whole 16-row tiles: a shard whose rows are not a multiple of 16 is not taken (the fused + verify keeps the in-tree GEMM), and the op refuses it.""" + w = torch.zeros(4004, 1024, dtype=torch.bfloat16, device="cuda") + assert not target_head_kernels.head_supported(SimpleNamespace(weight=w, bias=None)) + assert target_head_kernels.head_supported(SimpleNamespace(weight=w[:4000], bias=None)) + + +@pytest.mark.parametrize("n,k", [(TARGET_SHARD, HIDDEN), (4004, 1024)]) +def test_round_trip_is_bitwise(n, k) -> None: + w, coded = _coded(n, k, 1, True) + assert coded[3].numel() > 0 + assert round_trip_ok(*coded, w) + assert torch.equal(decode(*coded, n, k).view(torch.int16), w.view(torch.int16)) + + +def test_a_corrupted_coding_fails_closed() -> None: + """A flipped sign/mantissa bit, a wrong exception value or a wrong base each fail the bitwise round trip.""" + w = _weight(4004, 1024, 2, exceptions=True) + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + assert round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w) + bad = packed.clone() + bad[100] ^= 1 + assert not round_trip_ok(bad, base, exc_off, exc_pos, exc_w, w) + bad_w = exc_w.clone() + bad_w[0, 0] += 1e-3 + assert not round_trip_ok(packed, base, exc_off, exc_pos, bad_w, w) + assert not round_trip_ok(packed, base + 1, exc_off, exc_pos, exc_w, w) + + +def test_prepare_head_keeps_only_a_bitwise_coding(monkeypatch) -> None: + """prepare_head keeps the coded shard when it round-trips and runs it as the entry; a coding that does not + round-trip is dropped (the in-tree GEMM stays the path).""" + w = _weight(4000, 1024, 3, exceptions=True) + head = target_head_kernels.prepare_head(SimpleNamespace(weight=w, bias=None)) + assert head.coded and head.weight is w + x = torch.randn(4, 1024, device="cuda").to(torch.bfloat16) + out = target_head_kernels.logits(head, x) + stock = torch.nn.functional.linear(x, w) + assert torch.equal(out.float().argmax(-1), stock.float().argmax(-1)) + assert (out == stock).float().mean() > 0.99 + + def corrupt(weight, *args, **kwargs): + packed, base, exc_off, exc_pos, exc_w = encode_chunked(weight) + packed = packed.clone() + packed[7] ^= 0x40 + return packed, base, exc_off, exc_pos, exc_w + + monkeypatch.setattr(target_head_kernels, "encode_chunked", corrupt) + dropped = target_head_kernels.prepare_head(SimpleNamespace(weight=w, bias=None)) + assert not dropped.coded and dropped.w is None + + +def test_guard_refuses_a_dirty_counter() -> None: + w, coded = _coded(4000, 1024, 1, True) + args = list(_args(coded, torch.zeros(2, 1024, dtype=torch.bfloat16, device="cuda"), 4000, 8)) + args[4][3] = 1 + with pytest.raises(AssertionError, match="counter must be zero"): + with validating(op): + op(*args) + + +def _tiny_exceptions_head(n: int = 4096, k: int = 1024) -> torch.Tensor: + """Weights of 1 (exponent window base 113) with column 0 zero except 2^-30 / 2^-29 on rows 0 / 1: exceptions far + below the window whose exact products are the only non-zero logits.""" + w = torch.ones((n, k), dtype=torch.bfloat16, device="cuda") + w[:, 0] = 0.0 + w[0, 0] = 2.0**-30 + w[1, 0] = 2.0**-29 + return w + + +@pytest.mark.parametrize("m", [1, 4, 8]) +def test_tiny_exception_products_are_exact(m) -> None: + """An exception adds exactly x * w: with a one-hot activation on column 0 the logits are 0, 2^-30 and 2^-29 bit + for bit as the BF16 head's, and row 1 wins (a correction against a non-zero placeholder cancels them to 0).""" + w = _tiny_exceptions_head() + coded = encode_chunked(w) + assert coded[1] == 113 and round_trip_ok(*coded, w) + x = torch.zeros((m, w.shape[1]), dtype=torch.bfloat16, device="cuda") + x[:, 0] = 1.0 + with validating(op): + out = op(*_args(coded, x, w.shape[0], _num_sms())) + stock = torch.nn.functional.linear(x, w) + assert torch.equal(out.view(torch.int16), stock.view(torch.int16)) + assert out.float().argmax(-1).tolist() == [1] * m + + +def test_a_shard_smaller_than_the_grid_is_stored() -> None: + """A shard with fewer coded units than warps of the requested grid: the launch narrows the grid so every warp owns + a unit, and every logit is stored.""" + w = _tiny_exceptions_head(n=256) + coded = encode_chunked(w) + x = torch.zeros((2, w.shape[1]), dtype=torch.bfloat16, device="cuda") + x[:, 0] = 1.0 + out = op(*_args(coded, x, w.shape[0], _num_sms())) + assert torch.equal(out.view(torch.int16), torch.nn.functional.linear(x, w).view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_mxfp8_low_m_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_mxfp8_low_m_gemm.py new file mode 100644 index 000000000000..ed99da197a8a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_mxfp8_low_m_gemm.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 mxfp8_low_m_gemm entry of the MiniMax-M3 tp4 target.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm.mxfp8_low_m_gemm import ( # noqa: E501 + mxfp8_low_m_gemm as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "mxfp8_low_m_gemm requires a CUDA device" + + +def _swizzle(byte: torch.Tensor) -> torch.Tensor: + """[N, K / 32] scale bytes into the 128x4-swizzled buffer of ``mxfp8_quantize(w, True)``.""" + rows, cols = byte.shape + padded_cols, padded_rows = (cols + 3) // 4 * 4, (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + + (c // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + out = torch.zeros(padded_rows * padded_cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def _build(spec, seed): + """A BF16 activation and an MXFP8 weight: E4M3 values of a unit normal, UE8M0 block scales 2^-9 .. 2^-3.""" + g = torch.Generator(device="cuda").manual_seed(seed) + m, n, k = spec["m"], spec["n"], spec["k"] + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + w = torch.randn(n, k, generator=g, device="cuda").to(torch.float8_e4m3fn) + sf = torch.randint(118, 125, (n, k // 32), generator=g, device="cuda", dtype=torch.int32) + return x, w, _swizzle(sf.to(torch.uint8)) + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Accumulation order only, then one BF16 rounding. + + Both sides multiply the same E4M3 operands and power-of-two scales; the kernel sums each 32-wide block in the + MMA, rescales it exactly and accumulates blocks per warp, warps in order and split-K ranks in order, the + reference sums in a single fp32 GEMM. The fp32 results differ by a few fp32 ulps of the row's magnitude, which + moves the BF16 rounding by at most one step of the element -- at most 2 units of ``ULP[bf16]`` against the row + scale. Measured on GB300 over every cell: worst element 0.32 unit, relative RMS 0.006 unit. + """ + assert_within_ulp(out, ref, element_ulp=2.0, rms_ulp=0.25) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + x, w, wsf = _build(spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + args = (x, w, wsf, spec["warps"], spec["cluster_k"], spec["row_tiles"], False) + with validating(op): + out = op(*args) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + _compare(out, op.reference(*args)) + + +def test_the_gate_sees_a_wrong_scale_layout() -> None: + """Block scales read in the wrong order land far outside the band: most outputs pick up other blocks' + scales, which a layout bug in either the op or the reference would do.""" + spec = op.CELLS[1].spec + g = torch.Generator(device="cuda").manual_seed(7) + x, w, wsf = _build(spec, seed=7) + scrambled = wsf[torch.randperm(wsf.numel(), generator=g, device="cuda")] + args = (x, w, wsf, spec["warps"], spec["cluster_k"], spec["row_tiles"], False) + out = op(*args) + with pytest.raises(AssertionError): + _compare(out, op.reference(x, w, scrambled, *args[3:])) + + +def test_guard_refuses_a_short_scale_buffer() -> None: + """The kernel bulk-copies scale chunks by computed offset; a linear-length buffer is read past its end.""" + spec = op.CELLS[0].spec + x, w, wsf = _build(spec, seed=11) + n, k = w.shape + with pytest.raises(AssertionError, match="swizzled UE8M0 buffer"): + with validating(op): + op( + x, + w, + wsf[: n * k // 32 - 512], + spec["warps"], + spec["cluster_k"], + spec["row_tiles"], + False, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/_m3_moe_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/_m3_moe_operands.py new file mode 100644 index 000000000000..98a592728241 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/_m3_moe_operands.py @@ -0,0 +1,170 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One tp4 rank's MiniMax-M3 MoE layer operands for the decode-engine entries' tests. + +The routed experts are random checkpoint-format NVFP4 tensors ([w3 (up); w1 (gate)] rows, e4m3 block scales) put +through TRT-LLM's own TRTLLM-Gen NVFP4 loader shuffle (``NVFP4TRTLLMGenFusedMoEMethod._shuffle_all_experts``), which +writes the buffers the engines read in place. The per-expert scalars follow the loader's conventions (act_beta and +act_clamp divided by fc31_alpha, fc31_scale_c = fc2_input_scale x fc31_alpha) and are sized so the FC1 sums land +around the SwiGLU-OAI clamp and the outputs near unit scale. The shared expert's MXFP8 weights come from +mxfp8_quantize with the swizzled scale layout, as the MXFP8 linear method stores them. +""" + +import functools +from types import SimpleNamespace + +import torch + +H, I_TP, NUM_EXPERTS, TOP_K = 6144, 768, 128, 4 +SHARED_I = 768 +SWIGLU_ALPHA, SWIGLU_LIMIT = 1.702, 7.0 +ROUTED_SCALING = 2.0 +MINIMAX2_ROUTING = 5 # RoutingMethodType.MiniMax2 +# NVFP4 global scales: 448 x 6 over the amax each side is sized for. +FC31_INPUT_SCALE = 448.0 * 6.0 / 8.0 +FC2_INPUT_SCALE = 448.0 * 6.0 / 24.0 + + +def _rand_e4m3_scales(shape, gen) -> torch.Tensor: + """E4M3 block scales in [0.5, 2): biased exponents 6 and 7, any mantissa.""" + return torch.randint(0x30, 0x40, shape, dtype=torch.uint8, device="cuda", generator=gen).view( + torch.float8_e4m3fn + ) + + +@functools.lru_cache(maxsize=None) +def experts(seed: int = 20261013): + """The routed experts and the router of one MoE layer, in the layout and conventions of the TRTLLM-Gen loader. + + Returns a SimpleNamespace: ``tensors`` (the ten expert operands in ``M3MoeState.layer``'s order), ``bias``, + ``router_weight``, ``fc31_input_scale`` and ``shared`` (gate_up, its scales, down, its scales). + """ + from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4TRTLLMGenFusedMoEMethod + + gen = torch.Generator(device="cuda").manual_seed(seed) + e = NUM_EXPERTS + w31 = torch.randint( + 0, 256, (e, 2 * I_TP, H // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + w2 = torch.randint(0, 256, (e, H, I_TP // 2), dtype=torch.uint8, device="cuda", generator=gen) + w31s = _rand_e4m3_scales((e, 2 * I_TP, H // 16), gen) + w2s = _rand_e4m3_scales((e, H, I_TP // 16), gen) + module = SimpleNamespace( + w3_w1_weight=w31, + w2_weight=w2, + w3_w1_weight_scale=w31s, + w2_weight_scale=w2s, + is_gated_activation=True, + bias=False, + ) + NVFP4TRTLLMGenFusedMoEMethod()._shuffle_all_experts(module) + torch.cuda.synchronize() + + # Raw FC1 sums: rms(x) x FC31_INPUT_SCALE x rms(E2M1 value) x rms(block scale) x sqrt(H) ~ 1.0e5 at rms(x) = 1; + # fc31_alpha brings them to ~1.5 in the dequantized domain, the FC2 alpha the outputs to ~1. + base = 1.5 / (FC31_INPUT_SCALE * 2.93 * 1.3 * H**0.5) + fc31_alpha = base * (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) + fc2_alpha = (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) / ( + FC2_INPUT_SCALE * 2.93 * 1.3 * I_TP**0.5 + ) + tensors = ( + w31, + w31s, + w2, + w2s, + fc31_alpha, + FC2_INPUT_SCALE * fc31_alpha, # fc31_scale_c + torch.full((e,), SWIGLU_ALPHA, device="cuda"), # act_alpha + 1.0 / fc31_alpha, # act_beta: the (up + 1) offset over fc31_alpha + SWIGLU_LIMIT / fc31_alpha, # act_clamp over fc31_alpha + fc2_alpha, + ) + bias = (0.02 * torch.randn(e, device="cuda", generator=gen)).float() + router_weight = (torch.randn(e, H, device="cuda", generator=gen) / H**0.5).float().contiguous() + fc31_input_scale = torch.tensor([FC31_INPUT_SCALE], dtype=torch.float32, device="cuda") + + gu = (0.02 * torch.randn(2 * SHARED_I, H, device="cuda", generator=gen)).to(torch.bfloat16) + dn = (0.05 * torch.randn(H, SHARED_I, device="cuda", generator=gen)).to(torch.bfloat16) + gu_q, gu_sf = torch.ops.trtllm.mxfp8_quantize(gu, True) + dn_q, dn_sf = torch.ops.trtllm.mxfp8_quantize(dn, True) + shared = (gu_q.view(2 * SHARED_I, H), gu_sf, dn_q.view(H, SHARED_I), dn_sf) + return SimpleNamespace( + tensors=tensors, + bias=bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared=shared, + ) + + +# Routing cases: random (up to 4 M distinct experts), and every token on the same four experts (the bias outweighs +# any sigmoid difference), which leaves later slot groups empty. +CONCENTRATED = (5, 40, 77, 126) + + +def layer_for(state, case: str, shared: bool = True): + """One layer on ``state`` over :func:`experts` for a routing case. ``shared=False``: the shared expert's weights + are zero, so the engine returns the routed part alone (bf16(0 + bf16(routed))).""" + ops = experts() + bias = ops.bias.clone() + if case == "concentrated": + bias[list(CONCENTRATED)] += 1.0 + shared_weights = ops.shared + if not shared: + gu, gu_sf, dn, dn_sf = ops.shared + shared_weights = (torch.zeros_like(gu), gu_sf, torch.zeros_like(dn), dn_sf) + return state.layer( + *ops.tensors, + bias, + router_weight=ops.router_weight, + fc31_input_scale=ops.fc31_input_scale, + shared_weights=shared_weights, + ) + + +def stock_routed( + hidden: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor +) -> torch.Tensor: + """The stock TRTLLM-Gen NVFP4 MoE (fp4_quantize, then fp4_block_scale_moe_runner with SwiGLU-OAI) on the given + routing, over :func:`experts`: what the built-in MiniMaxM3MoE computes for the routed part.""" + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers the stock ops + + ops = experts() + w31, w31s, w2, w2s, fc31_alpha, scale_c, act_alpha, act_beta, act_clamp, fc2_alpha = ops.tensors + xq, xsf = torch.ops.trtllm.fp4_quantize(hidden, ops.fc31_input_scale, 16, False, False) + out = torch.ops.trtllm.fp4_block_scale_moe_runner( + None, None, xq, xsf.view(torch.float8_e4m3fn).flatten(), + w31, w31s, None, act_alpha, act_beta, act_clamp, w2, w2s, None, + scale_c, fc31_alpha, fc2_alpha, + NUM_EXPERTS, TOP_K, None, None, I_TP, 0, NUM_EXPERTS, ROUTED_SCALING, MINIMAX2_ROUTING, True, 0, + topk_weights=topk_weights.to(torch.bfloat16), topk_ids=topk_ids.to(torch.int32), + ) # fmt: skip + return out[0] if isinstance(out, (list, tuple)) else out + + +def hidden_states(m: int, seed: int) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(m, H, device="cuda", generator=gen).to(torch.bfloat16) + + +def check_routed_part(op, state, m: int, seed: int) -> None: + """Exact, independent of the reference: the engine's routing (its route_dbg record: ids, then the bf16 weights' + fp32 bits) is the reference routing bit for bit, and its routed part is the stock TRTLLM-Gen MoE's output on that + routing bit for bit -- the claim the engines are built on.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.moe.m3_moe_engine import ( # noqa: E501 + moe_route, + ) + + layer = layer_for(state, "random", shared=False) + hidden = hidden_states(m, seed=seed) + out = op(hidden, layer) + ids = state.route_dbg[:, :4].long() + weights = state.route_dbg[:, 4:].contiguous().view(torch.float32) + ref_ids, ref_weights = moe_route(hidden.float(), layer.router_weight, layer.bias) + assert torch.equal(ids, ref_ids) and torch.equal(weights, ref_weights), ( + "routing differs from the reference" + ) + stock = stock_routed(hidden, ids, weights) + assert torch.equal(out.view(torch.int16), stock.view(torch.int16)), ( + "routed part differs from TRTLLM-Gen's" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/_m3_router.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/_m3_router.py new file mode 100644 index 000000000000..0eefbc3aa8ea --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/_m3_router.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared operands and gates of the router entries' tests (m3_router_gemv, m3_router_gemv_fp4, m3_router_tc: one +kernel source).""" + +import zlib + +import torch + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp # noqa: E402 — needs the path declared above + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.moe.m3_router_gemv import ( # noqa: E501 + m3_router_gemv as gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.moe.m3_router_gemv import ( # noqa: E501 + near_tie, + nvfp4_quantize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.moe.m3_router_gemv_fp4 import ( # noqa: E501 + m3_router_gemv_fp4 as gemv_fp4, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.moe.m3_router_tc import ( # noqa: E501 + m3_router_tc as tc, +) + +__all__ = [ + "SF_SCALE", + "_build", + "_compare_fp4", + "_compare_logits", + "_seed", + "gemv", + "gemv_fp4", + "tc", +] + +#: The experts' NVFP4 input scale the tc cells use: 448 * 6 / 8, block scales well inside E4M3. +SF_SCALE = 336.0 + + +def _build(m: int, e: int, k: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(m, k, generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(e, k, generator=g, device="cuda") * 0.02 + return x, w + + +def _seed(cell) -> int: + return zlib.crc32(cell.why.encode()) % 2**31 + + +def _compare_logits(out: torch.Tensor, ref: torch.Tensor) -> None: + """fp32 logits that differ from a true fp32 product only in summation order. + + Over K = 6144 products a reordered fp32 sum moves by a few ulps of the sum of magnitudes, which near + cancellation is ~sqrt(K) times the row's largest logit's ulp at most. Gated in fp32 ulps of each row's largest + logit: 64 element-wise, 8 in relative RMS. Measured on GB300: GEMV 1.6 / 1.4 ulp, tensor-core router 3.8 / 2.0 + ulp at 2..48 tokens; the tensor-core logits against a true fp32 product (the TF32 rounding the band must see) + sit at 2200+ / 1700 ulp. + """ + assert out.dtype == ref.dtype == torch.float32 and out.shape == ref.shape + assert_within_ulp(out, ref, element_ulp=64.0, rms_ulp=8.0) + + +def _codes(packed: torch.Tensor) -> torch.Tensor: + """[M, K / 2] packed e2m1 bytes -> [M, K] 4-bit codes, low nibble first.""" + lo, hi = packed & 0xF, packed >> 4 + return torch.stack((lo, hi), dim=-1).reshape(packed.shape[0], -1) + + +def _compare_fp4( + q: torch.Tensor, sf: torch.Tensor, x: torch.Tensor, sf_scale: torch.Tensor +) -> None: + """Scales bit-exact; codes bit-exact except within ``TIE_WINDOW`` of an e2m1 midpoint, where the kernel's + approximate reciprocals choose the side (and then only by one code).""" + ref_q, ref_sf, scaled = nvfp4_quantize(x, sf_scale) + assert q.dtype == torch.uint8 and q.shape == ref_q.shape + assert torch.equal(sf, ref_sf), f"{(sf != ref_sf).sum().item()} block scales differ" + got, want = _codes(q), _codes(ref_q) + differ = got != want + assert not bool((differ & ~near_tie(scaled)).any()), ( + f"{(differ & ~near_tie(scaled)).sum().item()} e2m1 codes differ away from a rounding tie" + ) + magnitude = ((got & 7).int() - (want & 7).int()).abs() + assert bool((magnitude[differ] <= 1).all()), "a tie resolved more than one code away" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_moe16_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_moe16_engine.py new file mode 100644 index 000000000000..e80edf3bda3a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_moe16_engine.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the moe/m3_moe16_engine entry: the 9..13-row MoE engine on its caller-owned state. + +One GB300 rank of the tp4 deployment, the target's build (128 routed CTAs + 24 shared-expert CTAs). Operands and the +band: ``_m3_moe_operands`` and ``test_modeling_v2_m3_moe_engine`` (the same math). +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.m3_moe.op16 import ( # noqa: E501 + M16_NUM_CTAS, + SHARED_CTAS, + M3Moe16State, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.moe.m3_moe16_engine import ( # noqa: E501 + m3_moe16_engine as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_moe_operands import H, check_routed_part, hidden_states, layer_for # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_moe16_engine requires a CUDA device" + +# As test_modeling_v2_m3_moe_engine.compare (same arithmetic, same quantization boundaries). +BAND_ELEMENT = 4.0 +BAND_RMS = 0.5 + +_STATES = {} + + +def _state(m: int) -> M3Moe16State: + sms = torch.cuda.get_device_properties(torch.device("cuda")).multi_processor_count + if sms < M16_NUM_CTAS + SHARED_CTAS: + pytest.skip(f"m3_moe16 needs {M16_NUM_CTAS + SHARED_CTAS} resident SMs; device has {sms}") + if m not in _STATES: + _STATES[m] = M3Moe16State(torch.device("cuda"), num_tokens=m) + return _STATES[m] + + +@pytest.mark.parametrize("case", ("random", "concentrated")) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, case) -> None: + m = cell.spec["tokens"] + layer = layer_for(_state(m), case) + hidden = hidden_states(m, seed=zlib.crc32((cell.why + case).encode()) % 2**31) + with validating(op): + out = op(hidden, layer) + again = op(hidden, layer) + assert out.shape == (m, H) and out.dtype is torch.bfloat16 + assert torch.equal(out.view(torch.int16), again.view(torch.int16)), ( + "a repeated call changed the output" + ) + assert_within_ulp(out, op.reference(hidden, layer), element_ulp=BAND_ELEMENT, rms_ulp=BAND_RMS) + + +@pytest.mark.parametrize("m", (9, 13)) +def test_routed_part_is_trtllm_gen_bit_for_bit(m) -> None: + check_routed_part(op, _state(m), m, seed=100 + m) + + +def test_insufficient_sm_budget_is_rejected() -> None: + device = torch.device("cuda") + if torch.cuda.get_device_properties(device).multi_processor_count >= M16_NUM_CTAS + SHARED_CTAS: + pytest.skip("device has enough SMs for the MoE16 schedule") + with pytest.raises(ValueError, match="routed CTAs"): + M3Moe16State(device, num_tokens=9) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_moe_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_moe_engine.py new file mode 100644 index 000000000000..0b54f5fb2700 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_moe_engine.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the moe/m3_moe_engine entry: the 1..8-token MoE engine on its caller-owned state. + +One GB300 rank of the tp4 deployment: 128 experts (all local, intermediate 768 per rank), top 4, hidden 6144, the +shared expert folded in, ``num_ctas`` = SMs - 24 as the target builds it. Operands: ``_m3_moe_operands``. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.moe.m3_moe_engine import ( # noqa: E501 # noqa: E501 + M3MoeState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_moe_operands import H, check_routed_part, hidden_states, layer_for # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" + +_STATES = {} + + +def _state(m: int) -> M3MoeState: + if m not in _STATES: + dev = torch.device("cuda") + sms = torch.cuda.get_device_properties(dev).multi_processor_count + _STATES[m] = M3MoeState(dev, num_tokens=m, num_ctas=sms - 24, shared=True) + return _STATES[m] + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """The engine against the plain-torch MoE: BAND_ELEMENT bf16 ulp of the row's largest magnitude per element, + BAND_RMS ulp relative RMS. + + Both sides quantize the same operands with the same recipes (the activation's and the intermediate's NVFP4 codes + and block scales, the shared expert's MXFP8); they differ in summation order (tensor-core fp32 against fp64) and in + the PTX approximations (rcp / ex2 / tanh). Such a difference changes the output only where it moves a value across + a rounding boundary: one bf16 ulp of a slot output or of the sum, or, rarely, one intermediate NVFP4 code, which + reaches the output as one weight column's share. Measured over every cell, both routing cases and 3 seeds each + (GB300): at most 1.48 ulp element, 0.18 ulp RMS. The band leaves ~3x on both; a misread recipe (the activation's + 1 / (sf x (1 / sc)) spelling of the reciprocal on ~0.06% of its codes) measured 12 ulp element, 7 ulp RMS. + """ + assert_within_ulp(out, ref, element_ulp=BAND_ELEMENT, rms_ulp=BAND_RMS) + + +BAND_ELEMENT = 4.0 +BAND_RMS = 0.5 + + +@pytest.mark.parametrize("case", ("random", "concentrated")) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, case) -> None: + m = cell.spec["tokens"] + layer = layer_for(_state(m), case) + hidden = hidden_states(m, seed=zlib.crc32((cell.why + case).encode()) % 2**31) + with validating(op): + out = op(hidden, layer) + again = op(hidden, layer) + assert out.shape == (m, H) and out.dtype is torch.bfloat16 + # The workspace is tagged per call and never reset: a repeated call is the same call. + assert torch.equal(out.view(torch.int16), again.view(torch.int16)), ( + "a repeated call changed the output" + ) + compare(out, op.reference(hidden, layer)) + + +@pytest.mark.parametrize("m", (1, 2, 8)) +def test_routed_part_is_trtllm_gen_bit_for_bit(m) -> None: + check_routed_part(op, _state(m), m, seed=100 + m) + + +def test_layers_on_one_state_do_not_leak_into_each_other() -> None: + """Two layers share one state's workspace, in stream order: A, B, A gives A's output twice, bit for bit, and the + caller's ``out`` receives the same rows as a fresh tensor.""" + state = _state(4) + a, b = layer_for(state, "random"), layer_for(state, "concentrated") + hidden = hidden_states(4, seed=7) + first = op(hidden, a) + op(hidden, b) + out = torch.empty(4, H, dtype=torch.bfloat16, device="cuda") + second = op(hidden, a, out=out) + assert second.data_ptr() == out.data_ptr() + assert torch.equal(first.view(torch.int16), out.view(torch.int16)) + + +@pytest.mark.parametrize( + "bad", + ("misaligned hidden", "fp32 out", "short out", "long out", "int32 arp"), +) +def test_malformed_buffers_are_rejected_by_the_call(bad) -> None: + """Why the entry has no is_valid: the compiled call checks each buffer against its signature and raises.""" + state = _state(2) + layer = layer_for(state, "random") + hidden, kw = hidden_states(2, seed=3), {} + if bad == "misaligned hidden": + hidden = torch.randn(2 * H + 8, device="cuda").to(torch.bfloat16)[2 : 2 + 2 * H].view(2, H) + elif bad == "fp32 out": + kw["out"] = torch.empty(2, H, dtype=torch.float32, device="cuda") + elif bad == "short out": + kw["out"] = torch.empty(1, H, dtype=torch.bfloat16, device="cuda") + elif bad == "long out": + kw["out"] = torch.empty(3, H, dtype=torch.bfloat16, device="cuda") + else: + kw["arp"] = torch.zeros(8, dtype=torch.int32, device="cuda") + with pytest.raises(ValueError): + op(hidden, layer, **kw) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv.py new file mode 100644 index 000000000000..2da1ac3c6470 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_router_gemv entry of the MiniMax-M3 tp4 target.""" + +import pytest +import torch + +__extra_import_path__ = ["../../../..", "."] +from _m3_router import ( # noqa: E402, F401 + SF_SCALE, + _build, + _compare_fp4, + _compare_logits, + _seed, + gemv, + gemv_fp4, + tc, +) +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "the router entries require a CUDA device" + + +@pytest.mark.parametrize("cell", gemv.CELLS, ids=[c.why[:48] for c in gemv.CELLS]) +def test_gemv_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], spec["e"], spec["k"], _seed(cell)) + with validating(gemv): + out = gemv(x, w, 1, False) + _compare_logits(out, gemv.reference(x, w, 1, False)) + + +def test_the_gemv_is_a_true_fp32_product() -> None: + """And the GEMV, by contrast, is no TF32 product: the TF32 reference rejects it.""" + x, w = _build(1, 128, 6144, 6) + out = gemv(x, w, 1, False) + with pytest.raises(AssertionError): + _compare_logits(out, tc.reference(x, w, None, False)[0]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv_fp4.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv_fp4.py new file mode 100644 index 000000000000..036161f44f43 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv_fp4.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_router_gemv_fp4 entry of the MiniMax-M3 tp4 target.""" + +import pytest +import torch + +__extra_import_path__ = ["../../../..", "."] +from _m3_router import ( # noqa: E402, F401 + SF_SCALE, + _build, + _compare_fp4, + _compare_logits, + _seed, + gemv, + gemv_fp4, + tc, +) +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "the router entries require a CUDA device" + + +@pytest.mark.parametrize("cell", gemv_fp4.CELLS, ids=[c.why[:48] for c in gemv_fp4.CELLS]) +def test_gemv_fp4_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], spec["e"], spec["k"], _seed(cell)) + sf_scale = torch.tensor([spec["sf_scale"]], dtype=torch.float32, device="cuda") + with validating(gemv_fp4): + logits, q, sf = gemv_fp4(x, w, sf_scale, False) + ref_logits, _, _ = gemv_fp4.reference(x, w, sf_scale, False) + _compare_logits(logits, ref_logits) + _compare_fp4(q, sf, x, sf_scale) + + +def test_guard_refuses_a_zero_global_scale() -> None: + x, w = _build(1, 128, 6144, 4) + with pytest.raises(AssertionError, match="positive and finite"): + with validating(gemv_fp4): + gemv_fp4(x, w, torch.zeros(1, dtype=torch.float32, device="cuda"), False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_tc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_tc.py new file mode 100644 index 000000000000..03ecdb80ffcd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_tc.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_router_tc entry of the MiniMax-M3 tp4 target.""" + +import pytest +import torch + +__extra_import_path__ = ["../../../..", "."] +from _m3_router import ( # noqa: E402, F401 + SF_SCALE, + _build, + _compare_fp4, + _compare_logits, + _seed, + gemv, + gemv_fp4, + tc, +) +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "the router entries require a CUDA device" + + +@pytest.mark.parametrize("cell", tc.CELLS, ids=[c.why[:48] for c in tc.CELLS]) +def test_tc_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], 128, spec["k"], _seed(cell)) + sf_scale = torch.tensor([SF_SCALE], dtype=torch.float32, device="cuda") if spec["fp4"] else None + with validating(tc): + outs = tc(x, w, sf_scale, False) + ref = tc.reference(x, w, sf_scale, False) + assert len(outs) == len(ref) == (3 if spec["fp4"] else 1) + _compare_logits(outs[0], ref[0]) + if spec["fp4"]: + _compare_fp4(outs[1], outs[2], x, sf_scale) + + +def test_tc_reference_is_tf32_and_the_gate_sees_it() -> None: + """The tensor-core router rounds the weight to TF32: against a true fp32 product its logits sit far outside + the band, so the TF32 rounding in ``reference`` is load-bearing rather than decoration.""" + x, w = _build(16, 128, 6144, 5) + out = tc(x, w, None, False)[0] + with pytest.raises(AssertionError): + _compare_logits(out, gemv.reference(x, w, 1, False)) + + +def test_guard_refuses_a_misaligned_activation() -> None: + """A row view starting 8 bytes into the buffer keeps a legal row stride but not a legal bulk-copy source.""" + x, w = _build(2, 128, 6144, 3) + flat = torch.empty(2 * 6144 + 4, dtype=torch.bfloat16, device="cuda") + shifted = flat[4:].view(2, 6144) + shifted.copy_(x) + with pytest.raises(AssertionError, match="16-byte aligned"): + with validating(tc): + tc(shifted, w, None, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py new file mode 100644 index 000000000000..8778dbe2805d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the norm/m3_embedding_firstnorm entry: the decode embedding gather + the first layer's RMSNorm. + +The table is the model's full bf16 [200064, 6144] embedding shape (random values at the checkpoint's scale), eps the +checkpoint's 1e-6. +""" + +import functools +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + VOCAB, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + m3_embedding_firstnorm as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_embedding_firstnorm requires a CUDA device" + +H = 6144 +EPS = 1e-6 + + +@functools.lru_cache(maxsize=None) +def _operands(): + g = torch.Generator(device="cuda").manual_seed(20261013) + table = (0.02 * torch.randn(VOCAB, H, generator=g, device="cuda")).to(torch.bfloat16) + weight = (1.0 + 0.1 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + return table, weight + + +def _buffers(rows): + return (torch.empty(rows, H, dtype=torch.bfloat16, device="cuda"), + torch.empty(rows, H, dtype=torch.bfloat16, device="cuda")) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + rows, id_dtype = cell.spec["rows"], cell.spec["id_dtype"] + table, weight = _operands() + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + ids = torch.randint(0, VOCAB, (rows,), generator=g, device="cuda").to(id_dtype) + raw, out = _buffers(rows) + with validating(op): + normed, residual = op(ids, table, weight, EPS, raw, out) + assert normed.data_ptr() == out.data_ptr() and residual.data_ptr() == raw.data_ptr() + ref_normed, ref_raw = op.reference(ids, table, weight, EPS, raw, out) + # The gather is a copy: exact. + assert torch.equal(raw.view(torch.int16), ref_raw.view(torch.int16)) + # The norm: the kernel's fp32 sum of squares (warp butterflies, then across 4 warps) and fast rsqrt against an + # fp64 reference, then x * rstd * w rounded once to bf16. The fp32 side is ~1e-6 relative off, so the two roundings + # disagree only next to a bf16 rounding boundary: one ulp at most per element, a small fraction in aggregate. + # Measured over these cells (3 seeds each, GB300): at most 0.51 ulp element, 0.026 ulp RMS. + assert_within_ulp(out, ref_normed, element_ulp=1.0, rms_ulp=0.1) + + +def test_ids_outside_the_table_give_zero_rows() -> None: + """Documented, not refused: an id outside [0, V) gathers a zero row whose norm is zero.""" + table, weight = _operands() + ids = torch.tensor([-1, 5, VOCAB], dtype=torch.int64, device="cuda") + raw, out = _buffers(3) + op(ids, table, weight, EPS, raw, out) + assert not raw[0].any() and not raw[2].any() and not out[0].any() and not out[2].any() + assert torch.equal(raw[1], table[5]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/_m3_spec_inputs.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/_m3_spec_inputs.py new file mode 100644 index 000000000000..5c5d4161aa78 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/_m3_spec_inputs.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Inputs shared by the tp4 target's spec entry tests: TP4 dimensions, single-device stand-ins for the cross-rank +exchanges, and the comparison bands.""" + +from __future__ import annotations + +import zlib + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + make_keys, +) + +#: MiniMax-M3 at TP4: hidden size, vocabulary, the target head's shard and RMSNorm epsilon. +HIDDEN = 6144 +VOCAB = 200064 +WORLD = 4 +TARGET_SHARD = VOCAB // WORLD +EPS = 1e-6 + + +def seed_of(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def num_sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def randn(shape, g: torch.Generator, std: float = 1.0, dtype=torch.bfloat16) -> torch.Tensor: + return (torch.randn(shape, generator=g, device="cuda") * std).to(dtype) + + +def head_workspace(world: int, call: int = 0, slot_rows: int = SHARD_ROWS): + """(ws_uc int64 [2 * slot_rows * world], keys int64 [slot_rows], counters int32 [2]) of a fresh call ``call``.""" + uc = torch.zeros(2 * slot_rows * world, dtype=torch.int64, device="cuda") + keys = torch.zeros(slot_rows, dtype=torch.int64, device="cuda") + counters = torch.tensor([0, call], dtype=torch.int32, device="cuda") + return uc, keys, counters + + +def plant_peers( + uc: torch.Tensor, + counters: torch.Tensor, + own_best: torch.Tensor, + rank: int, + world: int, + shard: int, + g: torch.Generator, + slot_rows: int = SHARD_ROWS, + ties: bool = False, +) -> None: + """Write the other ranks' keys for each row into this rank's slots, as their multicast pushes would. + + Each peer's maximum is the own row maximum shifted by +-0.25 (well clear of the accumulation-order band), so + some rows are won by a peer and some by this rank; with ``ties`` a peer on an even row reports exactly the own + maximum, and the lower global index must win. + """ + if world == 1: + return + m = own_best.shape[0] + half = int(counters[1].item()) & 1 + for r in range(world): + if r == rank: + continue + sign = torch.where(torch.rand(m, generator=g, device="cuda") < 0.5, -1.0, 1.0) + value = own_best.float() + 0.25 * sign + if ties: + value = torch.where(torch.arange(m, device="cuda") % 2 == 0, own_best.float(), value) + index = r * shard + torch.randint(0, shard, (m,), generator=g, device="cuda") + keys = make_keys(value, index) + for t in range(m): + uc[(half * slot_rows + t) * world + r] = keys[t] + + +def assert_bf16_close(out: torch.Tensor, ref: torch.Tensor, max_frac: float = 1e-3) -> None: + """At most one bf16 ulp per element, and only on a ``max_frac`` share of the elements. + + Kernel and reference compute the same fp32 expression and round once; they differ only where a different + summation order (or rsqrtf) puts the fp32 value on the other side of a bf16 rounding boundary. + """ + assert out.shape == ref.shape and out.dtype == ref.dtype == torch.bfloat16 + o, r = out.float(), ref.float() + ulp = torch.ldexp(torch.ones_like(r), torch.frexp(r).exponent - 8) + ulp = torch.where(r == 0, torch.full_like(r, 2.0**-133), ulp) + diff = (o - r).abs() + assert bool((diff <= ulp * 1.0001).all()), f"max |diff| / ulp = {(diff / ulp).max().item():.2f}" + frac = (diff > 0).float().mean().item() + assert frac <= max_frac, f"{frac:.2e} of the elements moved by one ulp (allowed {max_frac})" + + +def craft_drafts(tokens: torch.Tensor, nreq: int, d: int) -> torch.Tensor: + """Drafts matching request b's verify tokens for its first (b + 2) % (d + 1) positions and differing from + there on, so the acceptance counts vary across requests.""" + t = tokens.view(nreq, d + 1)[:, :d].clone() + for b in range(nreq): + n = (b + 2) % (d + 1) + t[b, n:] += 1 + return t.contiguous() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_capture_norm.py new file mode 100644 index 000000000000..52fdac496253 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_capture_norm.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_capture_norm entry (Eagle3 capture with fc_norm folded in).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_capture_norm import ( # noqa: E501 + m3_capture_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_capture_norm requires a CUDA device" + +LAYERS = 3 + + +def _case(spec, seed): + g = gen(seed) + m, slot = spec["m"], spec["slot"] + h = randn((m, HIDDEN), g, 2.0) + r = randn((m, HIDDEN), g, 2.0) if spec["residual"] else None + w = (1.0 + randn((HIDDEN,), g, 0.1).float()).to(torch.bfloat16) + buf = torch.full((m, LAYERS * HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda") + dst = buf[:, slot * HIDDEN : (slot + 1) * HIDDEN] + return (h, r, w, EPS, dst), buf + + +@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 this name + """Within one bf16 ulp per element on at most 0.1% of the elements; the other layer slots of the buffer + untouched.""" + args, buf = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + slot = cell.spec["slot"] + others = torch.cat([buf[:, : slot * HIDDEN], buf[:, (slot + 1) * HIDDEN :]], dim=1) + assert bool((others == 7.0).all()) + + +def test_guard_refuses_a_strided_weight() -> None: + (h, r, w, eps, dst), _ = _case(op.CELLS[0].spec, 1) + w2 = torch.stack([w, w], dim=1)[:, 0] + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(h, r, w2, eps, dst) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_chunk_norm.py new file mode 100644 index 000000000000..985b51e8650c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_chunk_norm.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_chunk_norm entry (the Eagle3 fc_norm over the captured layers).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_chunk_norm import ( # noqa: E501 + m3_chunk_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_chunk_norm requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, c, h = spec["m"], spec["c"], spec["h"] + buf = randn((m, c * h + spec["extra_cols"]), g, 2.0) + src = buf[:, : c * h + spec["extra_cols"]] + w = (1.0 + randn((c, h), g, 0.1).float()).to(torch.bfloat16) + dst = torch.full((m, c * h), 7.0, dtype=torch.bfloat16, device="cuda") + return (src, w, EPS, dst, True) + + +@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 this name + """Within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic, rsqrtf).""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + + +def test_guard_refuses_a_strided_weight() -> None: + src, w, eps, dst, pdl = _case(op.CELLS[0].spec, 1) + wt = torch.empty((w.shape[1], w.shape[0]), dtype=w.dtype, device="cuda").t() + wt.copy_(w) + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(src, wt, eps, dst, pdl) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_input.py new file mode 100644 index 000000000000..7863cf2e33cb --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_input.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_input entry (embedding + both RMSNorms + concat in one launch), against its reference and +against the in-tree drafter's modules.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_draft_input import ( # noqa: E501 + m3_draft_input as op, +) +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.rms_norm import RMSNorm + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + VOCAB, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_input requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, h = spec["m"], spec["h"] + table = randn((VOCAB, h), g) + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + if spec["invalid_id"]: + ids[0] = -1 + hidden = randn((4 * m if spec["gather"] else m, h), g) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 2 if spec["gather"] else None + w_e = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + w_h = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + out = torch.full((m, 2 * h), 7.0, dtype=torch.bfloat16, device="cuda") + residual = torch.full((m, h), 7.0, dtype=torch.bfloat16, device="cuda") + return (ids, table, hidden, rows, w_e, w_h, EPS, EPS, out, residual, True) + + +@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 this name + """Normalized halves within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic in a + different summation order, rsqrtf); the gathered residual and the zero embedding of an invalid id exact.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out, residual = op(*args) + ref_out, ref_res = op.reference(*args) + torch.cuda.synchronize() + assert_bf16_close(out, ref_out) + assert torch.equal(residual, ref_res) + if cell.spec["invalid_id"]: + assert not bool(out[0, : cell.spec["h"]].float().any()) + + +def test_guard_refuses_rows_past_hidden() -> None: + gather = next(c for c in op.CELLS if c.spec["gather"]) + args = list(_case(gather.spec, 1)) + args[3] = args[3] + 100 + with pytest.raises(AssertionError, match="rows must index hidden"): + with validating(op): + op(*args) + + +@pytest.mark.parametrize("m,gather", [(4, False), (32, False), (8, True)], ids=str) +def test_the_in_tree_draft_input(m, gather) -> None: + """Against the in-tree drafter's own modules: the Embedding lookup, then the layer's input_layernorm and + hidden_norm (in-tree RMSNorm) and the concat of Eagle3DecoderLayer.forward. Within one bf16 ulp on at most 0.1% of + the elements (the RMS statistic summed in another order); the residual (the carried hidden rows) exact.""" + g = gen(seed_of(f"in-tree draft input {m} {gather}")) + h = HIDDEN + embed = Embedding(VOCAB, h, dtype=torch.bfloat16).cuda() + embed.weight.data = randn((VOCAB, h), g) + norms = [] + for _ in range(2): + norm = RMSNorm(hidden_size=h, eps=EPS, dtype=torch.bfloat16).cuda() + norm.weight.data = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + norms.append(norm) + input_layernorm, hidden_norm = norms + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + hidden = randn((4 * m if gather else m, h), g, 2.0) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 3 if gather else None + carried = hidden[rows] if gather else hidden + stock = torch.cat([input_layernorm(embed(ids)), hidden_norm(carried)], dim=-1) + out, residual = m3_spec_tail.draft_input( + ids, embed.weight, hidden, rows, input_layernorm.weight, EPS, hidden_norm.weight, EPS + ) + torch.cuda.synchronize() + assert_bf16_close(out, stock) + assert torch.equal(residual, carried) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_step.py new file mode 100644 index 000000000000..e2a419a1d3a1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_step.py @@ -0,0 +1,687 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_step entry (one fused greedy draft pass: the BF16 draft head, streamed from its lossless +12-bit coding or as BF16 fragments, the cross-rank argmax and the draft loop's bookkeeping), and its tokens against the +in-tree drafter's head and greedy draft sampling. + +The in-tree path: each rank's LM head shard GEMM (BF16 out), ``LogitsProcessor``'s float cast, then either the argmax +of the gathered logits (``torch.argmax``) or ``greedy_sample_draft_with_tp_gather``'s per-rank ``torch.max`` and the +pick among the ranks' (index, value) pairs. The edge cases use integer data, whose fp32 sums are exact in any order, +so the kernel and the in-tree path see the same BF16 logits and must agree bit for bit; every in-tree comparison runs +with the coded head (the default) and with the BF16 fragments, and on integer data the coded head's winning keys are +the BF16 fragments' bit for bit. +""" + +import warnings +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops import ( # noqa: E501 + m3_draft_step as draft_step_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm.m3_target_head import ( # noqa: E501 + encode_chunked, + round_trip_ok, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_draft_step import ( # noqa: E501 + DRAFT_LEN, + DRAFT_SHARD, + HIDDEN, + MAX_TOKENS, + bf16_fragments, + from_fragments, + head_weight, + shard_logits, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_draft_step import ( # noqa: E501 + m3_draft_step as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + WORLD, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_step requires a CUDA device" + +T = DRAFT_LEN + 1 +_LOW32 = 0xFFFFFFFF + +# ------------------------------------------------------------------ the per-row rule against the in-tree token +# +# The kernel and the in-tree head (a cuBLAS BF16 GEMM) both round an FP32 sum of the same exact BF16 products to BF16, +# in different summation orders. A row's token can differ from the in-tree one only where some candidate's FP32 sum +# can land on either side of a BF16 rounding midpoint. Each row is held to: +# +# kernel token == in-tree token, or the row is provably undecidable: with the fp64 reference logit L_i (exact BF16 +# products, fp64 sum) and the FP32-order bound E_i = gamma_K * sum_k |w_ik * x_k|, the BF16 value of any FP32 sum of +# candidate i lies in [lo_i, hi_i] = [bf16(L_i - E_i), bf16(L_i + E_i)]; the kernel's token must win for some +# choice of values within those intervals (larger value, then the lower global id), and a BF16 rounding midpoint +# must lie within E of the kernel's or the in-tree token's L. +# +# gamma_K = K * 2^-22 is la-xh2's bound (4x the round-to-nearest unit roundoff 2^-24 per addition: any summation +# order, and a truncating tensor-core accumulator); K counts the row's non-zero products (adding an exact zero does +# not round). lo / hi take one more FP32 step outward before the BF16 rounding, which covers rounding L -+ E to FP32. +# Each undecidable row is reported as a warning with its midpoint distances, and each test reports its row count; any +# other mismatch fails. +GAMMA_UNIT = 2.0**-22 + + +def _bf16_outward(v: torch.Tensor, direction: float) -> torch.Tensor: + f = v.float() + f = torch.nextafter(f, torch.full_like(f, direction)) + b = f.to(torch.bfloat16).float() + return torch.where(b == 0, torch.zeros_like(b), b) + + +def bounds(h: torch.Tensor, w: torch.Tensor): + """(fp64 logits L, bound E, lo, hi) [rows, N] of the head input rows h against the head rows w.""" + cols = (h != 0).any(dim=0) # zero columns add exact zeros + hd, wd = h[:, cols].double(), w[:, cols].double() + logits = hd @ wd.t() + bound = (h != 0).sum(dim=1, keepdim=True).double() * GAMMA_UNIT * (hd.abs() @ wd.abs().t()) + del wd + inf = float("inf") + return logits, bound, _bf16_outward(logits - bound, -inf), _bf16_outward(logits + bound, inf) + + +def midpoint_distance(v: float) -> float: + """Distance of v to the nearest BF16 rounding midpoint.""" + b = torch.tensor(v, dtype=torch.float64).float().to(torch.bfloat16).float() + if float(b) == 0.0: + return abs(v) + ulp = 2.0 ** (int(torch.frexp(b).exponent) - 8) + return min(abs(v - (float(b) + ulp / 2)), abs(v - (float(b) - ulp / 2))) + + +def _winnable(tok: int, ids: torch.Tensor, lo: torch.Tensor, hi: torch.Tensor) -> bool: + """Whether candidate ``tok`` wins the row for some values within the candidates' [lo, hi].""" + k = int((ids == tok).nonzero()[0, 0]) + others = ids != tok + beat = (hi[k] > lo[others]) | ((hi[k] == lo[others]) & (tok < ids[others])) + return bool(beat.all()) + + +def check_row(name: str, kernel_tok: int, stock_tok: int, ids, logits, bound, lo, hi) -> bool: + """The per-row rule for one row over the candidates ``ids`` [C] (logits / bound / lo / hi [C]); True when the + row is undecidable (reported as a warning), False when the tokens match; a mismatch outside the bound fails.""" + if kernel_tok == stock_tok: + return False + assert _winnable(kernel_tok, ids, lo, hi), ( + f"{name}: token {kernel_tok} differs from the in-tree {stock_tok} outside the FP32-order bound" + ) + ik = int((ids == kernel_tok).nonzero()[0, 0]) + i_stock = int((ids == stock_tok).nonzero()[0, 0]) + lk, ls, ek, es = ( + float(logits[ik]), + float(logits[i_stock]), + float(bound[ik]), + float(bound[i_stock]), + ) + dk, ds = midpoint_distance(lk), midpoint_distance(ls) + assert dk <= ek or ds <= es, f"{name}: neither logit is within its bound of a BF16 midpoint" + warnings.warn( + f"[draft-step undecidable] {name}: kernel {kernel_tok} (L {lk:.9g}, midpoint distance {dk:.3e}, bound " + f"{ek:.3e}) vs in-tree {stock_tok} (L {ls:.9g}, midpoint distance {ds:.3e}, bound {es:.3e})" + ) + return True + + +def report(name: str, rows: int, undecidable: int) -> None: + warnings.warn(f"[draft-step rows] {name}: {rows} rows, {undecidable} undecidable") + + +def key_value(keys: torch.Tensor) -> torch.Tensor: + """The fp32 value a key orders by (the inverse of the kernel's order-preserving image).""" + u = (keys >> 32) & _LOW32 + bits = torch.where(u >= 0x80000000, u - 0x80000000, _LOW32 - u) + bits = torch.where(bits >= 2**31, bits - 2**32, bits) # as a signed 32-bit pattern + return bits.to(torch.int32).view(torch.float32) + + +def key_index(keys: torch.Tensor) -> torch.Tensor: + return _LOW32 - (keys & _LOW32) + + +def _bookkeeping(nreq: int, step: int, n_x_rows: int, g: torch.Generator): + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if step == 0 + else None + ) + seq_lens = torch.full((nreq,), T, dtype=torch.int32, device="cuda") if step == 0 else None + next_new = torch.full((nreq, T), -5, dtype=torch.int32, device="cuda") + pos_src = torch.randint(100, 3000, (n_x_rows,), generator=g, device="cuda", dtype=torch.int32) + kv_lens = torch.randint(200, 3000, (nreq,), generator=g, device="cuda", dtype=torch.int32) + ids_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + pos_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + return num_acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next + + +def _head(w: torch.Tensor, coded: bool): + """(wf, (exc_off, exc_pos, exc_w, exp_base)) of a BF16 shard: its 12-bit coding (checked to round-trip) or its + BF16 fragments.""" + if not coded: + return bf16_fragments(w), (None, None, None, 0) + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + assert round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w) + return packed, (exc_off, exc_pos, exc_w.view(-1), base) + + +def _args( + x, + x_rows, + nreq, + wf, + rank, + world, + step, + g, + uc=None, + keys=None, + counters=None, + num_acc=None, + coded=(None, None, None, 0), +): + """The op's argument tuple for one rank (fresh workspace and bookkeeping buffers unless given); ``coded`` = the + coded head's (exc_off, exc_pos, exc_w, exp_base).""" + acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next = _bookkeeping( + nreq, step, x.shape[0], g + ) + if num_acc is None: + num_acc = acc + if uc is None: + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + out_keys = torch.zeros(nreq, dtype=torch.int64, device="cuda") + tiles = wf.shape[0] if coded[0] is None else coded[0].numel() - 1 + return (x, x_rows, nreq, wf, rank * tiles * 16, keys, counters, uc, uc.data_ptr(), rank, world, + next_new[:, 1:], step, ids_next, pos_src, pos_next, kv_lens, num_acc, seq_lens, DRAFT_LEN, num_sms(), + True, out_keys, *coded) # fmt: skip + + +def _case(spec, seed): + g = gen(seed) + nreq, step, rank, world, vs, k = ( + spec["nreq"], + spec["step"], + spec["rank"], + spec["world"], + spec["vs"], + spec["k"], + ) + pass0 = step == 0 + x = randn((nreq * T if pass0 else nreq, k), g) + wf, coded = _head(randn((vs, k), g, 0.02), spec["coded"]) + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if pass0 + else None + ) + x_rows = (torch.arange(nreq, device="cuda") * T + num_acc.long() - 1) if pass0 else None + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + h = x[x_rows] if pass0 else x + own = shard_logits(h, wf, *coded).max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=MAX_TOKENS, ties=spec.get("ties", False) + ) + return _args(x, x_rows, nreq, wf, rank, world, step, g, uc, keys, counters, num_acc, coded) + + +@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 this name + """The bookkeeping is exact on every row (next positions, KV and sequence lengths). Each row's token is the + in-tree one (this rank's BF16 shard logits from the head's GEMM, their first maximum, against the peers' keys) or + the row is provably undecidable (check_row); the winning key's index is the token, and its value lies in the + token's [lo, hi] (or is the peer's pushed key when a peer wins).""" + args = _case(cell.spec, seed_of(cell.why)) + x, x_rows, nreq, wf, offset, _, counters, uc = args[:8] + rank, world = args[9], args[10] + half = int(counters[1].item()) & 1 + pushed = ( + uc[half * MAX_TOKENS * world : (half * MAX_TOKENS + nreq) * world].view(nreq, world).clone() + ) + h = x[x_rows] if x_rows is not None else x[:nreq] + w = head_weight(wf, x.shape[1], *args[23:27]) + best, idx = local_argmax(torch.nn.functional.linear(h, w).float(), offset) + stock = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, MAX_TOKENS) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert args[6].tolist() == [0, 1] and not bool(args[5].any()), "counters / keys not re-armed" + step = args[12] + others = [c for c in range(args[11].shape[1]) if c != step] + assert torch.equal(args[11][:, others], ref["draft_out"][:, others]) + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + assert torch.equal(args[15], ref["pos_next"]) + assert torch.equal(args[16], ref["kv_lens"]) + if args[18] is not None: + assert torch.equal(args[18], ref["seq_lens"]) + # Candidates per row: this rank's head rows (fp64 bounds) and each peer's pushed key (an exact value). + logits, bound, lo, hi = bounds(h, w) + own_ids = offset + torch.arange(w.shape[0], device="cuda") + peers = [p for p in range(world) if p != rank] + undecidable = 0 + for r in range(nreq): + pv = key_value(pushed[r, peers]).double() + ids = torch.cat([own_ids, key_index(pushed[r, peers])]) + undecidable += check_row( + f"{cell.why[:40]} row {r}", + int(args[13][r]), + int(stock[r]), + ids, + torch.cat([logits[r], pv]), + torch.cat([bound[r], torch.zeros_like(pv)]), + torch.cat([lo[r], pv.float()]), + torch.cat([hi[r], pv.float()]), + ) + report(cell.why[:40], nreq, undecidable) + out_keys = args[22] + assert torch.equal(key_index(out_keys), args[13].long()), "the winning key's index is the token" + got = key_value(out_keys) + for r in range(nreq): + c = int(args[13][r]) - offset + if 0 <= c < w.shape[0]: + assert float(lo[r, c]) <= float(got[r]) <= float(hi[r, c]), ( + f"row {r}: winning value outside [lo, hi]" + ) + else: + assert int(out_keys[r]) in pushed[r].tolist(), ( + f"row {r}: a peer won with a key it did not push" + ) + + +def test_fragments_round_trip() -> None: + g = gen(seed_of("fragments")) + w = randn((64, 96), g) + wf = bf16_fragments(w) + assert wf.shape == (4, 6, 256) + assert torch.equal(from_fragments(wf), w) + head = SimpleNamespace(weight=w, bias=None) + # K 96: never coded, the BF16 fragments. + prepared = draft_step_kernels.prepare_head(head, 3, True) + assert not prepared.coded + assert torch.equal(prepared.wf, wf) and prepared.vocab_offset == 3 * 64 and prepared.weight is w + + +# ------------------------------------------------------------------ against the in-tree drafter, bit for bit + + +def _stock_tokens(shards: list, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """The in-tree greedy draft tokens of the rows x over the vocab-sharded BF16 head: (argmax of the gathered float + logits, greedy_sample_draft_with_tp_gather's per-rank (index, value) pick).""" + logits = [torch.nn.functional.linear(x, w).float() for w in shards] + gathered = torch.argmax(torch.cat(logits, dim=-1), dim=-1).to(torch.int32) + combined = [] + for rank, shard_logit in enumerate(logits): + worker = SimpleNamespace(mapping=SimpleNamespace(tp_rank=rank)) + combined.append(Eagle3OneModelWorker._get_local_max_and_combined(worker, shard_logit)) + pick = Eagle3OneModelWorker._get_draft_tokens_from_gathered + pick = getattr(pick, "_torchdynamo_orig_callable", pick) + tp_gather = pick(None, torch.cat(combined, dim=-1)).to(torch.int32) + return gathered, tp_gather + + +def _fused_tokens(shards: list, x: torch.Tensor, coded: bool, step: int = 1) -> list: + """Each rank's fused draft tokens: every rank's own key from a one-rank run, then each rank's four-rank pass with + its peers' keys in its slots, as their multicast pushes would put them.""" + g = gen(seed_of("fused tokens")) + nreq = x.shape[0] + heads = [_head(w, coded) for w in shards] + own_keys = [] + for rank, (wf, cd) in enumerate(heads): + args = list(_args(x, None, nreq, wf, rank, 1, step, g, coded=cd)) + args[9] = 0 # one-rank run: this rank's slot only, global ids of `rank` + op(*args) + own_keys.append(args[22].clone()) + tokens = [] + for rank, (wf, cd) in enumerate(heads): + uc, keys, counters = head_workspace(WORLD, slot_rows=MAX_TOKENS) + for peer in range(WORLD): + if peer != rank: + for t in range(nreq): + uc[t * WORLD + peer] = own_keys[peer][t] + args = _args(x, None, nreq, wf, rank, WORLD, step, g, uc, keys, counters, coded=cd) + with validating(op): + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + tokens.append(args[13].clone()) + return tokens + + +def _integer_head(g: torch.Generator, rows: int = WORLD * DRAFT_SHARD) -> torch.Tensor: + """A [rows, 6144] head of small integers over the first 4 columns (zero elsewhere): its logits against integer + rows are exact integers in any summation order, many of which share a BF16 value.""" + w = torch.zeros((rows, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, :4] = torch.randint(-60, 61, (rows, 4), generator=g, device="cuda").to(torch.bfloat16) + return w + + +def _shards(w: torch.Tensor) -> list: + return [w[r * DRAFT_SHARD : (r + 1) * DRAFT_SHARD].contiguous() for r in range(WORLD)] + + +def _assert_stock(w: torch.Tensor, x: torch.Tensor, coded: bool, expect=None) -> None: + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + if expect is not None: + assert gathered.tolist() == expect, (gathered.tolist(), expect) + for rank, tokens in enumerate(_fused_tokens(shards, x, coded)): + assert torch.equal(tokens, gathered), ( + f"rank {rank}: {tokens.tolist()} vs stock {gathered.tolist()}" + ) + + +def _ones_rows(nreq: int) -> torch.Tensor: + x = torch.zeros((nreq, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[:, :4] = 1.0 + return x + + +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_forced_bf16_ties_take_the_lowest_global_id(coded) -> None: + """A head row copied onto every rank that wins every row of ones (logit 244, every other at most 240): the lowest + global id must win, inside a shard and across ranks. Then integer rows in [-2, 2], whose logits (up to 480 in + magnitude, BF16 spacing 2 above 256) collide often, exactly or after BF16 rounding.""" + g = gen(seed_of("forced ties")) + w = _integer_head(g) + win = torch.tensor([61.0, 61.0, 61.0, 61.0], device="cuda").to(torch.bfloat16) + for gid in (3 * DRAFT_SHARD + 5, DRAFT_SHARD + 77, DRAFT_SHARD + 9000, 2 * DRAFT_SHARD + 1): + w[gid, :4] = win + _assert_stock(w, _ones_rows(5), coded, expect=[DRAFT_SHARD + 77] * 5) + # Random integer rows: ties after BF16 rounding happen naturally. + x = torch.randint(-2, 3, (8, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + x[:, 4:] = 0 + _assert_stock(w, x, coded) + + +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_signed_zero_maxima_take_the_lowest_global_id(coded) -> None: + """Rows whose maximum is zero: an all-zero head input (every logit +-0) and logits <= 0 with zeros by cancellation + on several ranks. -0 and +0 compare equal, so the lowest global id reaching zero wins.""" + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 0] = -1.0 + w[:, 1] = -1.0 + for gid in (2 * DRAFT_SHARD + 5, DRAFT_SHARD + 40000, 3 * DRAFT_SHARD + 7, 49999): + w[gid, 0] = 1.0 + w[gid, 1] = -1.0 + x = torch.zeros((3, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[1, :2] = 1.0 + x[2, :2] = 2.0 + x[2, 2] = -1.0 # w[:, 2] is zero: a -0 product on every row + _assert_stock(w, x, coded, expect=[0, 49999, 49999]) + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + (1000, 1001, "lower"), # 1001 rounds to 1000 in BF16: a tie, the lower id wins + (1001, 1002, "lower"), # 1001 and the tie-to-even midpoint 1002 both round to 1000 + (1000, 1004, "higher"), # adjacent BF16 values: the larger wins + (1003, 1000, "lower"), # 1003 rounds up to 1004 > 1000 + ], + ids=["one_fp32_step", "midpoint", "adjacent_bf16", "rounds_up"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_near_equal_top2_follow_the_bf16_logits(lower, higher, winner, where, coded) -> None: + """Top-2 logits one to four fp32 steps apart: the winner is the one the BF16 logits decide (the in-tree logits' + dtype), not the fp32 sum. The pair sits inside rank 3's shard or on ranks 1 and 2; every other logit is at most + 240 in magnitude.""" + g = gen(seed_of(f"near top2 {lower} {higher} {where}")) + w = _integer_head(g) + if where == "in_shard": + lo_id, hi_id = 3 * DRAFT_SHARD + 100, 3 * DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD + 49000, 2 * DRAFT_SHARD + 3 + for gid, value in ((lo_id, lower), (hi_id, higher)): + q, r = divmod(value, 4) # four integers <= 251, exact in BF16 + w[gid, :4] = torch.tensor([q + (i < r) for i in range(4)], device="cuda").to(torch.bfloat16) + expect = lo_id if winner == "lower" else hi_id + _assert_stock(w, _ones_rows(2), coded, expect=[expect, expect]) + + +# ------------------------------------------------------------------ non-integer rows at a BF16 rounding midpoint + +#: A BF16 rounding midpoint (between 6.0 and 6.03125; a tie there rounds to the even 6.0) and the FP32 spacing there. +MIDPOINT = 6.0 + 2.0**-6 +FP32_ULP = 2.0**-21 + + +def _near_midpoint_head(g: torch.Generator, targets: dict): + """A full-vocabulary head [4 * 50016, 6144] and one head-input row x with 65 non-zero entries: 1, 2^-12, 2^-20 + and 2^-28 on columns 0..3 and 61 random non-integer BF16 values. Every head row carries random non-integer BF16 + values on those 61 columns (logits about 0.2 in magnitude); each row of ``targets`` (global id -> fp64 value) also + gets columns 0..3 set so its fp64 logit is the target to within 2^-30 (a coarse value, then three finer + corrections, each BF16-rounded).""" + x = torch.zeros((1, HIDDEN), dtype=torch.bfloat16, device="cuda") + scales = (1.0, 2.0**-12, 2.0**-20, 2.0**-28) + for col, scale in enumerate(scales): + x[0, col] = scale + x[0, 4:65] = randn((61,), g, 0.5) + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 4:65] = randn((WORLD * DRAFT_SHARD, 61), g, 0.05) + xd = x[0].double() + for gid, target in targets.items(): + rest = float(target - (w[gid].double() * xd).sum()) + for col, scale in enumerate(scales): + v = torch.tensor(rest / scale, dtype=torch.float64).to(torch.bfloat16) + w[gid, col] = v + rest -= float(v) * scale + assert abs(rest) <= 2.0**-30, (gid, rest) + return w, x + + +def _check_fused_against_stock( + name: str, w: torch.Tensor, x: torch.Tensor, coded: bool, expect=None +) -> int: + """Each rank's fused tokens against the in-tree tokens under the per-row rule (full-vocabulary candidates); + ``expect`` (a token per row) additionally pins rows the bound decides. Returns the undecidable rows.""" + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + fused = _fused_tokens(shards, x, coded) + for tokens in fused[1:]: + assert torch.equal(tokens, fused[0]), "the ranks disagree on the global token" + logits, bound, lo, hi = bounds(x, w) + ids = torch.arange(w.shape[0], device="cuda") + undecidable = 0 + for r in range(x.shape[0]): + undecidable += check_row( + f"{name} row {r}", + int(fused[0][r]), + int(gathered[r]), + ids, + logits[r], + bound[r], + lo[r], + hi[r], + ) + if expect is not None: + assert int(gathered[r]) == expect[r] and int(fused[0][r]) == expect[r], ( + name, + int(gathered[r]), + int(fused[0][r]), + expect[r], + ) + others = ids != expect[r] + assert _winnable(expect[r], ids, lo[r], hi[r]), ( + f"{name}: the bound does not let {expect[r]} win" + ) + rivals = [int(i) for i in ids[others][hi[r][others] >= lo[r][expect[r]]]] + assert not any(_winnable(i, ids, lo[r], hi[r]) for i in rivals), ( + f"{name}: the row is not decided" + ) + report(name, x.shape[0], undecidable) + return undecidable + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + ( + -2, + 2, + None, + ), # straddling the midpoint by 2 FP32 ulps: either side is in-tree up to the order + (2, -2, None), + (0, -2, None), # exactly on the midpoint (ties to the even 6.0) against just below it + ( + -2048, + 2048, + "higher", + ), # ~1e-3 either side, well outside the bound: decided by BF16 rounding + ( + -4096, + -2048, + "lower", + ), # both below the midpoint, ~1e-3 apart: a BF16 tie the lower id wins + ], + ids=["straddle", "straddle_reversed", "on_midpoint", "decided_up", "decided_tie"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_top2_at_a_bf16_midpoint(lower, higher, winner, where, coded) -> None: + """Non-integer rows whose top-2 fp64 logits sit the given number of FP32 ulps from a BF16 rounding midpoint, inside + rank 3's shard or on ranks 1 and 2, every other logit far below. The per-row rule holds on every row; outside the + bound the token is the one the BF16 rounding decides, bit for bit.""" + g = gen(seed_of(f"midpoint {lower} {higher} {where}")) + if where == "in_shard": + lo_id, hi_id = 3 * DRAFT_SHARD + 100, 3 * DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD + 49000, 2 * DRAFT_SHARD + 3 + w, x = _near_midpoint_head( + g, {lo_id: MIDPOINT + lower * FP32_ULP, hi_id: MIDPOINT + higher * FP32_ULP} + ) + expect = None if winner is None else [lo_id if winner == "lower" else hi_id] + _check_fused_against_stock(f"midpoint {where} {lower:+d}/{higher:+d} ulp", w, x, coded, expect) + + +# ------------------------------------------------------------------ the coded head against the BF16 fragments + + +@pytest.mark.parametrize("nreq", [1, 2, 8]) +def test_coded_head_is_the_bf16_head_bit_for_bit(nreq) -> None: + """Integer data (fp32 sums exact in any order) with exceptions on both sides of the exponent window: the coded + head's winning keys (value and global id) and its bookkeeping are the BF16 fragments' bit for bit, on one rank and + across the ranks.""" + g = gen(seed_of(f"coded vs bf16 {nreq}")) + sign = torch.randint(0, 2, (DRAFT_SHARD, HIDDEN), generator=g, device="cuda") * 2 - 1 + w = (torch.randint(1, 41, (DRAFT_SHARD, HIDDEN), generator=g, device="cuda") * sign).to( + torch.bfloat16 + ) + # Exceptions on both sides of the window of 1..40: exact zeros, and 2^15 (fp32 sums stay below 2^24). + w[::5, ::512] = 0.0 + w[::13, 7::600] = 32768.0 + x = torch.randint(-2, 3, (nreq, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + wf_c, cd = _head(w, True) + assert cd[1].numel() > 0, "the shard must carry exceptions" + wf_b, cb = _head(w, False) + outs = [] + for wf, c in ((wf_c, cd), (wf_b, cb)): + args = _args(x, None, nreq, wf, 2, 1, 1, gen(seed_of("bookkeeping")), coded=c) + args = list(args) + args[9] = 0 + with validating(op): + op(*args) + torch.cuda.synchronize() + outs.append(args) + coded_args, bf16_args = outs + for i in (11, 13, 15, 16, 22): # draft_out, ids_next, pos_next, kv_lens, out_keys + assert torch.equal(coded_args[i], bf16_args[i]), i + best, idx = local_argmax(torch.nn.functional.linear(x, w).float(), 2 * DRAFT_SHARD) + assert torch.equal(coded_args[22], make_keys(best, idx)), ( + "the winning keys are not the in-tree ones" + ) + + +@pytest.mark.parametrize("nreq", [1, 2, 8]) +def test_tiny_exception_products_are_exact(nreq) -> None: + """Weights of 1 (exponent window base 113) with column 0 zero except 2^-30 / 2^-29 on rows 0 / 1, and a one-hot + activation on column 0: the exceptions' exact products are the only non-zero logits, so the coded head must pick + row 1 with the BF16 head's key bit for bit (a correction against a non-zero placeholder cancels them to 0).""" + n, k = 256, 1024 + w = torch.ones((n, k), dtype=torch.bfloat16, device="cuda") + w[:, 0] = 0.0 + w[0, 0] = 2.0**-30 + w[1, 0] = 2.0**-29 + x = torch.zeros((nreq, k), dtype=torch.bfloat16, device="cuda") + x[:, 0] = 1.0 + wf_c, cd = _head(w, True) + assert cd[3] == 113 + keys = [] + for wf, c in ((wf_c, cd), _head(w, False)): + args = _args(x, None, nreq, wf, 0, 1, 1, gen(seed_of("tiny exceptions")), coded=c) + with validating(op): + op(*args) + torch.cuda.synchronize() + assert args[13].tolist() == [1] * nreq + keys.append(args[22]) + best, idx = local_argmax(torch.nn.functional.linear(x, w).float(), 0) + assert torch.equal(keys[0], keys[1]) and torch.equal(keys[0], make_keys(best, idx)) + + +def test_prepare_head_codes_the_shard_and_fails_closed(monkeypatch) -> None: + """prepare_head keeps the coded shard when it decodes to the weight bit for bit; a coding that does not is + dropped for the BF16 fragments.""" + g = gen(seed_of("prepare coded")) + w = randn((4096, 1024), g, 0.02) + head = SimpleNamespace(weight=w, bias=None) + prepared = draft_step_kernels.prepare_head(head, 2, True) + assert prepared.coded and prepared.vocab_offset == 2 * 4096 and prepared.weight is w + assert torch.equal(head_weight(prepared.wf, 1024, prepared.exc_off, prepared.exc_pos, prepared.exc_w, + prepared.exp_base), w) # fmt: skip + + def corrupt(weight, *args, **kwargs): + packed, base, exc_off, exc_pos, exc_w = encode_chunked(weight) + packed = packed.clone() + packed[11] ^= 0x10 + return packed, base, exc_off, exc_pos, exc_w + + monkeypatch.setattr(draft_step_kernels, "encode_chunked", corrupt) + dropped = draft_step_kernels.prepare_head(head, 2, True) + assert not dropped.coded and dropped.exc_off is None + assert torch.equal(dropped.wf, bf16_fragments(w)) + + +def test_prepare_head_keeps_bf16_when_coding_is_off(monkeypatch) -> None: + """``coded`` False (``DRAFT_HEAD_CODED`` off) keeps the BF16 fragments even for a shard whose coding would decode + bit for bit.""" + g = gen(seed_of("prepare bf16")) + w = randn((4096, 1024), g, 0.02) + prepared = draft_step_kernels.prepare_head(SimpleNamespace(weight=w, bias=None), 2, False) + assert not prepared.coded and prepared.exc_off is None + assert torch.equal(prepared.wf, bf16_fragments(w)) + + +def test_guard_refuses_head_rows_past_x() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[1] = args[1] + 1000 + with pytest.raises(AssertionError, match="x_rows must index x"): + with validating(op): + op(*args) + + +def test_guard_refuses_too_few_ctas() -> None: + args = list(_case(op.CELLS[1].spec, 2)) + args[20] = 8 + with pytest.raises(AssertionError, match="tiles per CTA"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py new file mode 100644 index 000000000000..d88cea0ad9f5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py @@ -0,0 +1,289 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp4 target's speculative step against the in-tree (stock) Eagle3 numerics, at the real TP4 shapes. + +* Verify: the target logits are the in-tree BF16 LM head's own shard GEMM, and the fused shard argmax + acceptance + gives the in-tree greedy acceptance of the gathered logits, bit for bit, on every rank (peers' keys are the real + per-shard maxima, computed by the same kernel). +* Drafter: the worker runs the in-tree drafter's modules and weights; its own drafter arithmetic is the first pass's + fc_norm (chunk_norm, or capture_norm folded into the capture), held here to the in-tree ``apply_eagle3_fc`` RMSNorm + modules within one bf16 ulp (its effect on the fc output and on the full-vocabulary BF16 head's draft tokens is + measured too), and the fused draft loop's draft input and BF16 head, held to the in-tree modules and greedy draft + sampling by their own tests (test_modeling_v2_m3_draft_input.py, test_modeling_v2_m3_draft_step.py). + +Synthetic weights: hidden 6144, vocabulary 200064 (4 x 50016 shards), 3 captured layers, draft length 3. +""" + +import importlib.util +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + make_keys, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + m3_shard_accept as shard_accept_op, +) +from tensorrt_llm._torch.models.modeling_speculative import Eagle3ForCausalLM +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 + EPS, + HIDDEN, + TARGET_SHARD, + VOCAB, + WORLD, + assert_bf16_close, + craft_drafts, + gen, + head_workspace, + num_sms, + randn, + seed_of, +) + +assert torch.cuda.is_available(), "the exact-numerics checks require a CUDA device" + +DRAFT_LEN = 3 +CAPTURED_LAYERS = 3 +#: Largest share of fc_norm outputs one bf16 ulp off the in-tree RMSNorm (fp32 statistic summed in another order). +NORM_ULP_FRAC = 1e-3 + + +def _stock_accept(logits: torch.Tensor, draft: torch.Tensor): + """The in-tree greedy strict acceptance (``SpecWorkerBase._sample_and_accept_draft_tokens_base``) of the gathered + float logits of a decode-only batch, no forcing.""" + md = SimpleNamespace(is_all_greedy_sample=True, is_ragged_verify=False, is_cuda_graph=False) + worker = SimpleNamespace(force_num_accepted_tokens=0.0) + worker._sample_tokens_for_batch = lambda *a: Eagle3OneModelWorker._sample_tokens_for_batch( + worker, *a + ) + worker._apply_force_accepted_tokens = lambda *a, **k: ( + Eagle3OneModelWorker._apply_force_accepted_tokens(worker, *a, **k) + ) + nreq = draft.shape[0] + return Eagle3OneModelWorker._sample_and_accept_draft_tokens_base( + worker, logits, draft, 0, nreq, md + ) + + +def _lm_head_shard(weight: torch.Tensor, rank: int) -> LMHead: + """Rank ``rank``'s local part of the target's vocab-sharded BF16 LM head at TP4: an LM head over that rank's + 50016 rows (the TP4 module's local GEMM is this module's; its all-gather is the caller's).""" + head = LMHead(num_embeddings=TARGET_SHARD, embedding_dim=HIDDEN, dtype=torch.bfloat16) + head.weight.data = weight[rank * TARGET_SHARD : (rank + 1) * TARGET_SHARD].contiguous() + return head + + +def _run_shard_accept(logits, rank, world, draft, keys_of_peers=None, call=0): + """The fused verify of one rank: its shard logits, the peers' keys planted in its slots.""" + nreq = draft.shape[0] + rows = logits.shape[0] + uc, keys, counters = head_workspace(world, call=call, slot_rows=SHARD_ROWS) + half = call & 1 + for peer, peer_keys in (keys_of_peers or {}).items(): + for t in range(rows): + uc[(half * SHARD_ROWS + t) * world + peer] = peer_keys[t] + acc = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + shard_accept_op( + logits, rank * TARGET_SHARD if world > 1 else 0, keys, counters, uc, uc.data_ptr(), rank, world, + draft, batch_idx, acc, num, gather, 0, next_new, num_sms(), None, + ) # fmt: skip + torch.cuda.synchronize() + return acc, num, gather, next_new + + +def _local_keys(logits: torch.Tensor, rank: int) -> torch.Tensor: + """What rank ``rank`` pushes: the key of its own shard maximum per row, found by the same kernel at world 1.""" + rows = logits.shape[0] + nreq = rows // (DRAFT_LEN + 1) + dummy = torch.full((nreq, DRAFT_LEN), -1, dtype=torch.int32, device="cuda") + acc, _, _, _ = _run_shard_accept(logits, 0, 1, dummy) + local = acc.view(-1).to(torch.int64) + value = logits.float().gather(1, local[:, None]).squeeze(1) + return make_keys(value, local + rank * TARGET_SHARD) + + +@pytest.mark.parametrize( + "nreq,ties", [(1, False), (2, False), (4, False), (8, True), (12, False)], ids=str +) +def test_fused_verify_is_the_stock_greedy_acceptance(nreq, ties) -> None: + """Target logits (in-tree BF16 head) and verify outputs bit-identical to the in-tree path, on all four ranks.""" + g = gen(seed_of(f"verify {nreq} {ties}")) + rows = nreq * (DRAFT_LEN + 1) + x = randn((rows, HIDDEN), g, 1.0) + weight = randn((VOCAB, HIDDEN), g, 0.02) + if ties: + # Two equal head rows on ranks 1 and 3 that win every row: the lower global id must be the token. + win = (x.float().sum(0) / x.float().sum(0).norm() * 4.0).to(torch.bfloat16) + weight[TARGET_SHARD + 77] = win + weight[3 * TARGET_SHARD + 5] = win + heads = [_lm_head_shard(weight, r) for r in range(WORLD)] + # Stock: each rank's LMHead forward, all-gathered, then LogitsProcessor's float cast. + shard_logits = [head(x) for head in heads] + for head, stock in zip(heads, shard_logits): + assert torch.equal(head.apply_linear(x, None), stock), ( + "the verify's shard GEMM is not the head's" + ) + gathered = torch.cat(shard_logits, dim=1).float() + target_tokens = gathered.argmax(dim=-1).to(torch.int32) + if ties: + assert bool((target_tokens == TARGET_SHARD + 77).all()) + draft = craft_drafts(target_tokens, nreq, DRAFT_LEN) + stock_acc, stock_num = _stock_accept(gathered, draft) + peer_keys = {r: _local_keys(shard_logits[r], r) for r in range(WORLD)} + for rank in range(WORLD): + peers = {r: k for r, k in peer_keys.items() if r != rank} + acc, num, gather, next_new = _run_shard_accept( + shard_logits[rank], rank, WORLD, draft, peers + ) + assert torch.equal(acc, stock_acc.to(torch.int32)), f"rank {rank}: accepted tokens" + assert torch.equal(num, stock_num.to(torch.int32)), f"rank {rank}: accepted counts" + expect_gather = torch.arange(nreq, device="cuda") * (DRAFT_LEN + 1) + num.long() - 1 + assert torch.equal(gather, expect_gather), f"rank {rank}: first-pass gather rows" + assert torch.equal(next_new[:, 0], target_tokens[expect_gather]), f"rank {rank}: next_new" + + +class _Drafter: + """The in-tree Eagle3 first-pass pieces at the target's shapes: per-layer fc_norm RMSNorms and the fc.""" + + def __init__(self, g: torch.Generator): + self.norms = [] + for _ in range(CAPTURED_LAYERS): + norm = RMSNorm(hidden_size=HIDDEN, eps=EPS, dtype=torch.bfloat16).cuda() + norm.weight.data = (1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to( + torch.bfloat16 + ) + self.norms.append(norm) + self.fc = Linear(HIDDEN * CAPTURED_LAYERS, HIDDEN, bias=False, dtype=torch.bfloat16).cuda() + self.fc.weight.data = randn((HIDDEN, HIDDEN * CAPTURED_LAYERS), g, 0.02) + self.model = SimpleNamespace( + dtype=torch.bfloat16, + hidden_size=HIDDEN, + fc_norm=torch.nn.ModuleList(self.norms), + _norm_before_fc=False, + fc=self.fc, + ) + self.stacked = torch.stack([n.weight.detach() for n in self.norms]).contiguous() + + def stock_normed(self, hs: torch.Tensor) -> torch.Tensor: + """``apply_eagle3_fc``'s fc input: each captured layer through its own in-tree RMSNorm.""" + chunks = hs.chunk(CAPTURED_LAYERS, dim=-1) + return torch.cat([norm(c) for norm, c in zip(self.norms, chunks)], dim=-1) + + def stock_fc(self, hs: torch.Tensor) -> torch.Tensor: + return Eagle3ForCausalLM.apply_eagle3_fc(self, hs) + + +def _captures(rows: int, g: torch.Generator): + """Per captured layer (hidden, residual) and the in-tree capture buffer of bf16(hidden + residual).""" + layers = [(randn((rows, HIDDEN), g, 1.5), randn((rows, HIDDEN), g, 3.0)) for _ in range(3)] + buf = torch.cat([h + r for h, r in layers], dim=-1) + return layers, buf + + +def _row_identical(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return (a == b).all(dim=-1) + + +@pytest.mark.parametrize("rows", (4, 8, 32, 48), ids=lambda r: f"rows{r}") +def test_first_drafter_inputs_are_the_stock_fc_norm(rows) -> None: + """chunk_norm and capture_norm against ``apply_eagle3_fc``'s in-tree RMSNorms: within one bf16 ulp on at most + NORM_ULP_FRAC of the elements; the fc (the in-tree module in both paths) then differs only on rows whose input + did.""" + g = gen(seed_of(f"fc_norm {rows}")) + drafter = _Drafter(g) + layers, buf = _captures(rows, g) + stock_in = drafter.stock_normed(buf) + stock_out = drafter.stock_fc(buf) + assert torch.equal(drafter.fc(stock_in), stock_out) + + chunk_in = m3_spec_tail.chunk_norm(buf, drafter.stacked, EPS) + capture_in = torch.empty_like(buf) + for i, (h, r) in enumerate(layers): + m3_spec_tail.capture_norm( + h, r, drafter.stacked[i], EPS, capture_in[:, i * HIDDEN : (i + 1) * HIDDEN] + ) + torch.cuda.synchronize() + for name, ours in (("chunk_norm", chunk_in), ("capture_norm", capture_in)): + assert_bf16_close(ours, stock_in, max_frac=NORM_ULP_FRAC) + out = drafter.fc(ours) + same_rows = _row_identical(ours, stock_in) + assert torch.equal(out[same_rows], stock_out[same_rows]), ( + f"{name}: fc moved an unchanged row" + ) + moved = (ours != stock_in).float().mean().item() + print( + f"[exact-numerics] rows={rows} {name}: {moved:.2e} of fc inputs one ulp off, " + f"{int((~same_rows).sum())}/{rows} rows differ, fc max |diff| " + f"{(out.float() - stock_out.float()).abs().max().item():.3e}" + ) + assert torch.equal(chunk_in, capture_in), "capture_norm and chunk_norm disagree" + + +@pytest.mark.parametrize("rows", (4, 32), ids=lambda r: f"rows{r}") +def test_draft_tokens_through_the_full_vocab_head(rows) -> None: + """The fc_norm's ulp-level differences carried through the full-vocabulary BF16 head (a proxy of the draft head + input): the greedy draft token equals the stock one on every row whose stock top-2 margin exceeds the logit + perturbation.""" + g = gen(seed_of(f"draft head {rows}")) + drafter = _Drafter(g) + _, buf = _captures(rows, g) + weight = randn((VOCAB, HIDDEN), g, 0.02) + head = LMHead(num_embeddings=VOCAB, embedding_dim=HIDDEN, dtype=torch.bfloat16) + head.weight.data = weight + stock_logits = head(drafter.stock_fc(buf)).float() + ours_logits = head(drafter.fc(m3_spec_tail.chunk_norm(buf, drafter.stacked, EPS))).float() + perturb = (ours_logits - stock_logits).abs().max(dim=-1).values + top2 = stock_logits.topk(2, dim=-1).values + decidable = (top2[:, 0] - top2[:, 1]) > 2 * perturb + stock_tok, ours_tok = stock_logits.argmax(-1), ours_logits.argmax(-1) + assert torch.equal(ours_tok[decidable], stock_tok[decidable]) + print( + f"[exact-numerics] rows={rows} draft head: {int((ours_tok == stock_tok).sum())}/{rows} tokens equal, " + f"{int(decidable.sum())} decidable, max logit |diff| {perturb.max().item():.3e}" + ) + + +def test_the_drafter_is_the_in_tree_drafter() -> None: + """Outside the fused draft loop's steps the worker drafts with the in-tree loop, forward and sampling, and nothing + re-quantizes the drafter's linears or the LM heads: no hook of a quantized / pruned head or drafter projection is + left in the target, and the fused head streams the BF16 head weight itself.""" + assert ( + modeling.M3Eagle3Worker._forward_draft_loop is not Eagle3OneModelWorker._forward_draft_loop + ) + for name in ( + "_forward_linear_draft_loop", + "_run_draft_forward", + "sample_draft_tokens", + "sample_and_accept_draft_tokens", + ): + assert getattr(modeling.M3Eagle3Worker, name) is getattr(Eagle3OneModelWorker, name), name + for module, name in ( + (modeling, "M3DraftGemvLinearMethod"), + (modeling, "adopt_draft_linears"), + (modeling, "greedy_head_outputs"), + (modeling, "M3FP8LMHeadMethod"), + (modeling, "adopt_fp8_lm_heads"), + ): + assert not hasattr(module, name), f"{module.__name__}.{name}" + pkg = modeling.__name__.rsplit(".", 1)[0] + for gone in ("ops.m3_draft_head", "ops.spec.m3_draft_head", "ops.gemm.m3_draft_gemv"): + assert importlib.util.find_spec(f"{pkg}.{gone}") is None, gone diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_force_vec.py new file mode 100644 index 000000000000..d8808030aba3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_force_vec.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_force_vec entry of the MiniMax-M3 tp4 target: the fractional forced-acceptance draw.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_force_vec requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(5) + pool = torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, generator=g, device="cuda") + counter = torch.tensor([41], dtype=torch.int64, device="cuda") + out = torch.empty(spec["batch"], dtype=torch.int32, device="cuda") + args = ( + pool, + counter, + out, + spec["batch"], + 2, + 4, + spec["frac"], + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + spec["block"], + ) + ref = op.reference(*args) + with validating(op): + op(*args) + assert int(counter.item()) == 42, "the draw bumps the step counter once" + assert torch.equal(out, ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_sampled_accept.py new file mode 100644 index 000000000000..189d28ac638f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_sampled_accept.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 m3_sampled_accept entry (fused temperature + top-p verify) and its workspace size.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_sampled_accept import ( # noqa: E501 + m3_sampled_accept as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_sampled_accept_ws_words import ( # noqa: E501 + m3_sampled_accept_ws_words as ws_words, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_sampled_accept requires a CUDA device" + +DRAFT_LEN = 3 +#: Hot tokens of the distribution cells; every other logit sits 22 below the top (mass ~ 1e-5 of the total). +HEAD = 12 +#: Draws per distribution cell, and the total-variation band against the reference distribution. With 12 +#: categories and >= 4000 draws the expected TV of an exact sampler is ~0.02 (sd ~0.006); 0.06 is ~6 sd above it +#: while a wrong nucleus cut (one category too many or too few, >= 1 / 12 of the mass moved) lands at >= 0.08. +DRAWS = 4000 +TV_BAND = 0.06 + + +def _logits(spec, g, rows): + vs = spec["vs"] + if spec["check"] == "flat": + return (torch.randn((rows, vs), generator=g, device="cuda") * 0.3).to(torch.bfloat16) + if spec["check"] == "argmax": + return (torch.randn((rows, vs), generator=g, device="cuda") * 3).to(torch.bfloat16) + row = torch.full((vs,), -20.0, device="cuda") + hot = torch.randperm(vs, generator=g, device="cuda")[:HEAD] + row[hot] = torch.linspace(2.0, 0.2, HEAD, device="cuda") + return row.to(torch.bfloat16).expand(rows, vs).contiguous() + + +def _case(spec, seed): + g = gen(seed) + nreq = spec["nreq"] + rows = nreq * (DRAFT_LEN + 1) + logits = _logits(spec, g, rows) + words = ws_words(1) + uc = torch.zeros(words, dtype=torch.int64, device="cuda") + counters = torch.zeros(2, dtype=torch.int32, device="cuda") + tokens = torch.zeros(48, dtype=torch.int32, device="cuda") + extra = torch.zeros(1, dtype=torch.int32, device="cuda") + temperatures = torch.full((rows,), spec["temperature"], dtype=torch.float32, device="cuda") + top_ps = torch.full((rows,), spec["top_p"], dtype=torch.float32, device="cuda") + seeds = torch.randint(0, 2**62, (rows,), generator=g, device="cuda", dtype=torch.int64) + offsets = torch.zeros(rows, dtype=torch.int64, device="cuda") + # Drafts drawn from the likeliest tokens, so some prefixes match. + top = logits[0].float().topk(3).indices.to(torch.int32) + spec["vocab_offset"] + draft = top[torch.randint(0, 3, (nreq, DRAFT_LEN), generator=g, device="cuda")].contiguous() + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + force = spec.get("force", "none") + force_total = 3 if force == "total" else 0 + force_vec = (torch.randint(1, DRAFT_LEN + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if force == "vec" else None) # fmt: skip + return [logits, spec["vocab_offset"], temperatures, top_ps, seeds, offsets, uc, uc.data_ptr(), counters, tokens, + 0, 1, draft, batch_idx, acc, num, gather, force_total, next_new, force_vec, extra, True] # fmt: skip + + +@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 this name + """Statistical for the draw, exact for everything derived from it. + + Every drawn token has nonzero reference probability (lies in the nucleus); the acceptance outputs equal the + reference's strict acceptance of the drawn tokens, every call; the call counter advances once per call; and, + for the distribution cells, the draws over many calls (fresh Philox offsets) are within TV_BAND of the + reference distribution (see DRAWS / TV_BAND). + """ + spec = cell.spec + args = _case(spec, seed_of(cell.why)) + logits, offset, rows = args[0], spec["vocab_offset"], args[0].shape[0] + certifying(op) + ref = op.reference(*args) + probs = ref["probs"] + calls = 1 if spec["check"] == "argmax" else -(-DRAWS // rows) + counts = torch.zeros(logits.shape[1], dtype=torch.float64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + drawn = args[9][:rows].clone() + local = drawn.long() - offset + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + assert bool((probs.gather(1, local[:, None]) > 0).all()), ( + "a drawn token outside the nucleus" + ) + acc = ref["accept"](drawn) + assert torch.equal(args[14], acc["accepted"]) + assert torch.equal(args[15], acc["num_accepted"]) + assert torch.equal(args[16], acc["gather_ids"]) + assert torch.equal(args[18], acc["next_new"]) + assert args[8].tolist() == [0, c + 1] + counts.index_add_(0, local, torch.ones(rows, dtype=torch.float64, device="cuda")) + if spec["check"] == "distribution": + emp = counts / counts.sum() + tv = 0.5 * (emp - probs[0].double()).abs().sum().item() + assert tv <= TV_BAND, f"total variation {tv:.3f} against the reference nucleus distribution" + + +@pytest.mark.parametrize("top_p", [0.001, 0.0002]) +def test_long_rejection_chains_stay_in_the_nucleus(top_p) -> None: + """A tiny nucleus over a flat tail: rows that need more exchange rounds than the workspace has slots. + + Each round's candidate is drawn over the tokens heavier than the last one, so a row takes ~ln(1 / top_p) + 1 + rounds on average, with a tail past the 16 exchange slots (about 1% of the rows at top_p 0.001). Rounds reuse the + slots cyclically, told apart by a lap-parity bit in the wire tags, and no round accepts unconditionally: every + draw, however many rounds it took, lies in the reference nucleus. World 1 (one rank's 8 CTAs per row exchange + through the same slot ring); the cross-rank path needs the multi-device job. + """ + spec = dict(nreq=1, vs=50016, vocab_offset=50016, check="flat", top_p=top_p, temperature=1.0) + args = _case(spec, seed_of(f"flat tail top_p {top_p}")) + logits, rows = args[0], args[0].shape[0] + probs = op.reference(*args)["probs"] + calls = -(-DRAWS // rows) + outside = torch.zeros((), dtype=torch.int64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + local = args[9][:rows].long() - spec["vocab_offset"] + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + outside += (probs.gather(1, local[:, None]) == 0).sum() + assert args[8].tolist() == [0, c + 1] + draws = calls * rows + rounds = 2.0 + args[20].item() / draws + kept = int((probs > 0).sum(1).min().item()) + assert outside.item() == 0, ( + f"{outside.item()} of {draws} draws outside the nucleus (kept >= {kept} of {logits.shape[1]}, " + f"{rounds:.2f} rounds per row on average)" + ) + assert rounds > 4.0, ( + f"only {rounds:.2f} rounds per row: the rows no longer stress the slot ring" + ) + + +@pytest.mark.parametrize("cell", ws_words.CELLS, ids=[c.why[:44] for c in ws_words.CELLS]) +def test_workspace_words(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(ws_words) + world = cell.spec["world"] + assert ws_words(world) == ws_words.reference(world) + + +def test_guard_refuses_top_k_rows() -> None: + args = _case(op.CELLS[0].spec, 1) + top_ks = torch.full((args[0].shape[0],), 20, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="top-k rows"): + with validating(op): + op(*args, top_ks=top_ks) + + +def test_guard_refuses_fp64_temperatures() -> None: + args = _case(op.CELLS[0].spec, 2) + args[2] = args[2].double() + with pytest.raises(AssertionError, match="fp32 sampling parameters"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_shard_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_shard_accept.py new file mode 100644 index 000000000000..a984ca817a2e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_shard_accept.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_shard_accept entry (greedy verify from this rank's logits shard).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + m3_shard_accept as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 + craft_drafts, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_shard_accept requires a CUDA device" + + +def _case(spec, seed, edit_logits=None): + g = gen(seed) + nreq, d, vs, world, rank = ( + spec["nreq"], + spec["draft_len"], + spec["vs"], + spec["world"], + spec["rank"], + ) + m = nreq * (d + 1) + if spec["strided"]: + logits = randn((m, vs + 72), g, 3.0)[:, :vs] + else: + logits = randn((m, vs), g, 3.0) + if edit_logits is not None: + edit_logits(logits) + uc, keys, counters = head_workspace(world, slot_rows=SHARD_ROWS) + own = logits.float().max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=SHARD_ROWS, ties=spec.get("ties", False) + ) + offset = rank * vs + best, idx = local_argmax(logits.float(), offset) + tokens = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, SHARD_ROWS) + draft = craft_drafts(tokens, nreq, d) + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + force_total = 3 if spec["force"] == "total" else 0 + force_vec = (torch.randint(1, d + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if spec["force"] == "vec" else None) # fmt: skip + return (logits, offset, keys, counters, uc, uc.data_ptr(), rank, world, draft, batch_idx, acc, num, gather, + force_total, next_new, num_sms(), force_vec) # fmt: skip + + +@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 this name + """Bit-exact: the op reduces BF16 logits without arithmetic, ties (within the shard and across ranks) go to + the lowest global index on both sides.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) + assert args[3].tolist() == [0, 1] and not bool(args[2].any()) + + +def test_guard_refuses_nan_logits() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0][1, 7] = float("nan") + with pytest.raises(AssertionError, match="NaN"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_batch_index_past_the_requests() -> None: + args = list(_case(op.CELLS[0].spec, 2)) + args[9] = args[9] + 1 + with pytest.raises(AssertionError, match="batch_idx"): + with validating(op): + op(*args) + + +def _signed_zero_maxima(logits: torch.Tensor) -> None: + """Every row's maximum is zero, reached as -0.0 at a lower index and +0.0 at a higher one.""" + logits.fill_(-1.0) + logits[:, 3] = -0.0 + logits[:, 9] = 0.0 + + +@pytest.mark.parametrize("ties", [False, True], ids=["peer_rows", "peer_ties"]) +def test_signed_zero_maxima_take_the_lower_index(ties) -> None: + """-0.0 and +0.0 compare equal, so the lower index wins as with torch.argmax (also across ranks).""" + spec = dict(op.CELLS[1].spec, ties=ties) + args = _case(spec, seed_of(f"signed-zero-{ties}"), edit_logits=_signed_zero_maxima) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_slot_scatter.py new file mode 100644 index 000000000000..ec99ade9b043 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_slot_scatter.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_slot_scatter entry (the spec sampler's per-slot stores in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_slot_scatter import ( # noqa: E501 + m3_slot_scatter as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_slot_scatter requires a CUDA device" + +STALE = -4 + + +def _case(spec, seed): + """Worker outputs and the store's transposed [slots, width] views of its [width, slots, 1] buffers.""" + g = gen(seed) + n, nslots = spec["n"], spec["slots"] + + def rnd(shape): + return torch.randint(0, VOCAB, shape, generator=g, device="cuda", dtype=torch.int32) + + store_new = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_next = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_lens = torch.full((nslots,), STALE, dtype=torch.int32, device="cuda") + store_draft = torch.full((nslots, spec["store_d"]), STALE, dtype=torch.int32, device="cuda") + slots = torch.randperm(nslots, generator=g, device="cuda")[:n] + lens = torch.randint(1, 5, (n,), generator=g, device="cuda", dtype=torch.int32) + return (slots, rnd((n, spec["src_w"])), store_new.squeeze(-1).T, rnd((n, spec["src_w"])), + store_next.squeeze(-1).T, lens, store_lens, rnd((n, spec["src_d"])), store_draft) # fmt: skip + + +@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 this name + """Index copies: exact, zero padding included, other slots untouched.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in (("da", args[2]), ("db", args[4]), ("dl", args[6]), ("dd", args[8])): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_repeated_slots() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0] = args[0].clone() + args[0][1] = args[0][0] + with pytest.raises(AssertionError, match="distinct"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_step_gather.py new file mode 100644 index 000000000000..d75053d8c674 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_step_gather.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_step_gather entry (the overlap scheduler's step-input gathers in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_step_gather import ( # noqa: E501 + m3_step_gather as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_step_gather requires a CUDA device" + +#: Rows of the engine's flat buffers (the max batch), and their stale fill. +MAX_ROWS = 8 +STALE = -4 + + +def _case(spec, seed): + g = gen(seed) + p, t, d, nslots, b0 = spec["p"], spec["t"], spec["d"], spec["slots"], spec["b0"] + i32 = dict(dtype=torch.int32, device="cuda") + new_tok = torch.randint(0, VOCAB, (t, nslots, 1), generator=g, device="cuda", dtype=torch.int32) + if spec.get("two_d"): + new_tok = new_tok.squeeze(-1) + next_draft = torch.randint(0, VOCAB, (nslots, d), generator=g, device="cuda", dtype=torch.int32) + lens = torch.randint(1, t + 1, (nslots,), generator=g, device="cuda", dtype=torch.int32) + slots = torch.randperm(nslots, generator=g, device="cuda")[:p].to(torch.int32) + ppi = slots.repeat_interleave(t) + ids = torch.full((MAX_ROWS * t,), STALE, **i32) + draft = torch.full((MAX_ROWS * d,), STALE, **i32) + pos_off = torch.full((MAX_ROWS * t,), STALE, **i32) + kv_off = torch.full((MAX_ROWS,), STALE, **i32) + return ( + new_tok, + next_draft, + lens, + slots, + ppi, + p, + t, + d, + ids, + spec["ids0"], + draft, + spec["d0"], + pos_off, + kv_off, + b0, + ) + + +@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 this name + """Index copies: exact, including the zero-filled offsets and the untouched rest of ids / draft.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in ( + ("ids", args[8]), + ("draft", args[10]), + ("pos_off", args[12]), + ("kv_off", args[13]), + ): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_a_slot_past_the_store() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[3] = args[3].clone() + args[3][0] = 8 + with pytest.raises(AssertionError, match="slots must index the store"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_capture_boundary.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_capture_boundary.py new file mode 100644 index 000000000000..3fc723b67501 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_capture_boundary.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The generic layer's feed-forward exit on Eagle3 capture layers, in both targets' bodies (``PrefillTarget._layer``, +``DecodeTarget._generic_layer``). + +Every layer's MoE / MLP output goes unreduced to the fused boundary (all-reduce + residual + the next layer's input +norm, the final norm after the last layer), and a capture layer captures the reduced residual sum that boundary +writes, with no residual of its own. CPU only: the attention, the feed-forward and the boundaries are stand-ins. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 + modeling, +) + +LAYERS = 4 + + +class _Spec: + def __init__(self, layers): + self.layers = set(layers) + + def is_layer_capture(self, i): + return i in self.layers + + +@pytest.fixture(params=["prefill", "decode"]) +def layer(request, monkeypatch): + """(core, the target's generic layer bound to a target over ``core``).""" + calls = [] + w = {f"l{i}_{n}": f"{n}{i}" for i in range(LAYERS) for n in ("in_norm", "post_norm")} + w["final_norm"] = "final" + c = SimpleNamespace( + w=w, num_layers=LAYERS, moe_layers=frozenset({2, 3}), rms_eps=1e-6, calls=calls + ) + monkeypatch.setattr(modeling, "_can_mc", lambda x, folded=False: False) + monkeypatch.setattr( + modeling, + "_native_allreduce_norm", + lambda core, x, r, norm: (("normed", x, norm), ("updated", x, r)), + ) + + def boundary(self, x, x2, residual, norm_w): + calls.append((x, x2, norm_w)) + return ("next", norm_w), ("sum", x, x2) + + if request.param == "prefill": + cls = modeling.PrefillTarget + monkeypatch.setattr(cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i))) + monkeypatch.setattr(cls, "_swiglu_mlp", lambda self, h, prefix, low_m: ("mlp", prefix)) + else: + cls = modeling.DecodeTarget + monkeypatch.setattr( + cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i), False, False) + ) + monkeypatch.setattr( + cls, "_swiglu_mlp", lambda self, h, prefix, low_m: (("mlp", prefix), False) + ) + monkeypatch.setattr(cls, "_attention", lambda self, i, *a, **k: (("attn", i), False)) + monkeypatch.setattr(cls, "_generic_boundary", boundary) + target = cls(c) + return c, target._layer if request.param == "prefill" else target._generic_layer + + +@pytest.mark.parametrize("i", range(LAYERS)) +def test_a_capture_layer_captures_the_boundary_residual_sum(layer, i) -> None: + core, run = layer + captured = [] + out = run( + i, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([i]), + lambda spec, layer, h, r: captured.append((layer, h, r)), + ) + norm = core.w[f"l{i + 1}_in_norm"] if i + 1 < LAYERS else "final" + assert core.calls[-1][2] == norm, ( + "the boundary applies the next layer's norm (the final norm after the last)" + ) + assert captured == [(i, out[1], None)] + if i in core.moe_layers: + assert core.calls[-1][1] == ("routed", i), "the routed part folds into the boundary" + + +def test_a_layer_without_capture_captures_nothing(layer) -> None: + _, run = layer + captured = [] + run( + 1, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([]), + lambda *a: captured.append(a), + ) + assert not captured diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_dispatch.py new file mode 100644 index 000000000000..450acab7c7d4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_dispatch.py @@ -0,0 +1,902 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode target's dispatch table (``modeling.py``, imported here as ``dispatch``): its row ranges against the +catalog cells that certify them and the kernels' own limits, the selection for every decode row count, and the miss +/ uncovered policy. + +CPU only: nothing here launches a kernel. +""" + +import ast +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops import ( # noqa: E501 + m3_attn_mega_c1, + m3_draft_step, + m3_fused_attn, + m3_mc_allreduce, + m3_shard_accept, + m3_target_head, + triton_sparse_decode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops import ( # noqa: E501 + mxfp8_low_m_gemm as low_m, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm import ( # noqa: E501 + m3_target_head as target_head_entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm import ( # noqa: E501 + mxfp8_low_m_gemm as low_m_entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm.m3_mx_gemv import ( # noqa: E501 + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + SHARED_DOWN, + SHARED_GATE_UP, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.mx_gemv import ( # noqa: E501 + op as mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec import ( # noqa: E501 + m3_draft_step as draft_step_entry, +) +from tensorrt_llm._torch.modules.linear import TensorParallelMode + +PROJECTIONS = (QKV_SPARSE, QKV, O_PROJ, SHARED_GATE_UP, SHARED_DOWN, DENSE_GATE_UP, DENSE_DOWN) +ROWS = range(1, dispatch.DECODE_MAX_TOKENS + 2) # every decode row count and the first one past it + + +@pytest.fixture(autouse=True) +def _policy(): + """Every test starts outside require mode with nothing logged, and leaves it so.""" + dispatch.set_require(False) + dispatch.reset_logged() + yield + dispatch.set_require(False) + dispatch.reset_logged() + + +class _Log: + def __init__(self): + self.info, self.warning = [], [] + + +@pytest.fixture +def log(monkeypatch) -> _Log: + """The dispatch's log lines, by level.""" + rec = _Log() + monkeypatch.setattr(dispatch.logger, "info", rec.info.append) + monkeypatch.setattr(dispatch.logger, "warning", rec.warning.append) + return rec + + +# ------------------------------------------------------------------ the ranges against their certificates + + +@pytest.mark.parametrize("fast_path", dispatch.FAST_PATHS, ids=lambda r: r.path) +def test_every_covered_shape_is_a_certified_cell(fast_path) -> None: + uncertified = fast_path.keys - fast_path.certified() + assert not uncertified, ( + f"{fast_path.path} covers {sorted(uncertified, key=str)} but " + f"{type(fast_path.entry).__name__}.CELLS certifies none of them: add a cell or narrow the range" + ) + + +def test_the_fast_path_names_are_unique() -> None: + names = [r.path for r in dispatch.FAST_PATHS] + assert len(names) == len(set(names)) + + +# ------------------------------------------------------------------ the ranges against the kernels' own limits + + +@pytest.mark.parametrize("nk", PROJECTIONS, ids=str) +def test_gemv_rows_are_the_rows_the_kernel_serves(nk) -> None: + n, k = nk + served = {m for m in range(1, mx_gemv.MAX_TOKENS + 1) if mx_gemv.serves(n, k, m, False)} + assert dispatch.MX_GEMV_ROWS[nk] == served + served_bf16 = [m for m in range(1, mx_gemv.MAX_TOKENS + 1) if mx_gemv.serves(n, k, m, True)] + assert served_bf16 == list(range(1, dispatch.MX_GEMV_BF16_MAX_TOKENS[nk] + 1)) + + +@pytest.mark.parametrize("nk", sorted(dispatch.LOW_M_GEMM_MAX_TOKENS), ids=str) +def test_low_m_gemm_rows_are_its_configuration_limit(nk) -> None: + warps, cluster_k, max_m, row_tiles = low_m.get_config(*nk) + assert dispatch.LOW_M_GEMM_MAX_TOKENS[nk] == max_m + for m in range(1, max_m + 1): + assert low_m._smem_bytes(m, nk[1], warps, cluster_k, row_tiles) <= low_m._MAX_DYNAMIC_SMEM + + +def test_the_low_m_gemm_has_a_range_for_every_configuration() -> None: + assert set(dispatch.LOW_M_GEMM_MAX_TOKENS) == set(low_m._CONFIGS) + + +def test_c1_takes_exactly_its_range() -> None: + for qlen in range(1, dispatch.FUSED_ATTN_MAX_QUERY_LEN + 1): + for reqs in range(1, dispatch.FUSED_ATTN_MAX_TOKENS + 1): + assert m3_attn_mega_c1.c1_applies(qlen, reqs) == dispatch.ATTN_C1.covers((qlen, reqs)) + + +def test_every_fused_attention_shape_has_a_compiled_geometry() -> None: + for qlen, reqs in dispatch.FUSED_DECODE_ATTN.keys: + assert (qlen, *m3_fused_attn.default_geometry(qlen, reqs)) in m3_fused_attn.GEOMETRIES + + +def test_the_plans_are_the_all_reduce_buckets() -> None: + assert tuple(p.max_tokens for p in dispatch.DECODE_PLANS) == m3_mc_allreduce.BUCKETS + assert dispatch.DECODE_MAX_TOKENS == m3_mc_allreduce.MAX_TOKENS == mx_gemv.MAX_TOKENS + + +def test_the_thresholds_are_the_kernels() -> None: + assert dispatch.SPARSE_DECODE_MQ_MIN_TOKENS == triton_sparse_decode.MQ_MIN_TOKENS + assert dispatch.MOE16_MAX_TOKENS < dispatch.ATTN_PUSH_MAX_TOKENS <= dispatch.DECODE_MAX_TOKENS + verify_rows = dispatch.SHARD_ACCEPT_MAX_REQUESTS * (dispatch.SHARD_ACCEPT_DRAFT_LEN + 1) + assert verify_rows == dispatch.DECODE_MAX_TOKENS == m3_shard_accept.SHARD_ROWS + + +def test_no_fast_path_quantizes_a_bf16_head_or_the_drafter() -> None: + """The target carries the exact-numerics configuration only: the lm_head and the Eagle3 drafter run at the + checkpoint's precision, so no fast path (nor a range constant) of an FP8 / NVFP4 / pruned head or of a quantized + drafter projection is left to select.""" + names = {r.path for r in dispatch.FAST_PATHS} + assert not names & { + "fp8_lm_head", + "draft_head", + "m3_draft_gemv", + "m3_draft_in_gemv", + "m3_gemv_ar", + } + for constant in ( + "FP8_LM_HEAD", + "FP8_LM_HEAD_MAX_TOKENS", + "DRAFT_HEAD", + "DRAFT_HEAD_MAX_TOKENS", + ): + assert not hasattr(dispatch, constant), constant + + +# ------------------------------------------------------------------ selection, every decode row count + + +def _expected_moe(rows: int): + if rows <= 8: + return "m3_moe_engine" + if rows <= 13: + return "m3_moe16_engine" + if rows <= 48: + return "mc_finalize_allreduce_norm" + return None + + +def _expected_router(rows: int): + if rows == 1: + return "m3_router_gemv" + if rows <= 48: + return "m3_router_tc" + return None + + +@pytest.mark.parametrize("rows", ROWS) +def test_selection(rows) -> None: + assert dispatch.is_decode_step(rows) == (rows <= 48) + if rows <= 48: + assert dispatch.plan_for(rows).name == ( + "tiny" if rows <= 8 else "small" if rows <= 16 else "verify" + ) + else: + with pytest.raises(AssertionError, match="no decode plan"): + dispatch.plan_for(rows) + assert dispatch.moe_path(rows) == _expected_moe(rows) + assert dispatch.router_path(rows) == _expected_router(rows) + assert dispatch.EMBED_FIRSTNORM.covers(rows) == (rows <= 8) + assert dispatch.ATTN_PUSH.covers(rows) == (rows <= 16) + assert dispatch.MOE_PUSH.covers(rows) == (rows <= 8) + assert dispatch.MC_ALLREDUCE.covers(rows) == (rows <= 48) + + +@pytest.mark.parametrize("rows", ROWS) +def test_linear_selection(rows) -> None: + bf16_max = { + QKV_SPARSE: 16, + QKV: 16, + O_PROJ: 32, + SHARED_GATE_UP: 16, + SHARED_DOWN: 16, + DENSE_GATE_UP: 8, + DENSE_DOWN: 16, + } + for nk in PROJECTIONS: + prequant_gemv = rows <= 48 and not (nk == SHARED_DOWN and 17 <= rows <= 32) + assert dispatch.linear_path(*nk, rows, True, True) == ( + "m3_mx_gemv" if prequant_gemv else "flashinfer_mxfp8_gemm" + ) + for low_m_gemm in (False, True): + assert dispatch.linear_path(*nk, rows, False, low_m_gemm) == ( + "m3_mx_gemv_bf16" if rows <= bf16_max[nk] else "mxfp8_gemm" + ) + # A projection outside the table takes the generic GEMM at every row count. + assert dispatch.linear_path(4096, 4096, rows, True, True) == "flashinfer_mxfp8_gemm" + assert dispatch.linear_path(4096, 4096, rows, False, True) == "mxfp8_gemm" + + +def test_draft_step_selection() -> None: + """The fused draft loop takes the configured draft length at 1..8 requests (one head row per request, the + kernel's limit); its draft input every pass's rows: the first pass's 4 per request, then one per request.""" + assert ( + dispatch.DRAFT_STEP_MAX_REQUESTS + == draft_step_entry.MAX_TOKENS + == m3_draft_step.MAX_TOKENS + == 8 + ) + assert dispatch.DRAFT_STEP_DRAFT_LEN == dispatch.SHARD_ACCEPT_DRAFT_LEN + for draft_len in range(1, 6): + for nreq in range(1, 14): + covered = draft_len == 3 and nreq <= 8 + assert dispatch.DRAFT_STEP.covers((draft_len, nreq)) == covered + assert dispatch.DRAFT_INPUT.keys == frozenset(range(1, 9)) | {12, 16, 20, 24, 28, 32} + + +def test_verify_selection() -> None: + for draft_len in range(1, 6): + for nreq in range(1, 14): + covered = draft_len == 3 and nreq * (draft_len + 1) <= 48 + assert dispatch.SHARD_ACCEPT.covers((draft_len, nreq)) == covered + + +def test_target_head_selection() -> None: + """The coded target head takes every verify of at most 8 rows (one or two requests at draft length 3); the early + fc fork the same verifies.""" + assert dispatch.TARGET_HEAD_MAX_ROWS == 8 == dispatch.FC_EARLY_MAX_ROWS + for rows in ROWS: + assert dispatch.TARGET_HEAD.covers(rows) == (rows <= 8) + + +@pytest.mark.parametrize("rows", ROWS) +def test_attention_shapes(rows) -> None: + for qlen in range(1, 9): + if rows % qlen: + continue + reqs = rows // qlen + assert dispatch.FUSED_DECODE_ATTN.covers((qlen, reqs)) == (rows <= 16) + assert dispatch.ATTN_MEGA.covers((qlen, reqs)) == (qlen == 1 and reqs == 1) + assert dispatch.ATTN_C1.covers((qlen, reqs)) == (qlen == 6 and reqs == 1) + + +# ------------------------------------------------------------------ require mode and fallback logging + + +def test_a_miss_raises_under_require() -> None: + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_moe_engine'.* 5 rows.*not built"): + dispatch.missed("m3_moe_engine", 5, "the engine is not built") + + +def test_a_miss_logs_once_per_reason_otherwise(log) -> None: + for _ in range(3): + dispatch.missed("m3_moe_engine", 5, "the engine is not built") + dispatch.missed("m3_moe_engine", 6, "the engine failed to build") + assert len(log.warning) == 2 + assert "m3_moe_engine at 5 rows: the engine is not built" in log.warning[0] + + +# ------------------------------------------------------------------ the phase split + + +def test_the_core_selects_the_target_on_the_context_count_alone() -> None: + """Every pure generation step is decode's, whatever its width; a step with context rows (a mixed one too) is + prefill's. Nothing but ``num_contexts`` is read: the metadata stand-in carries no other field.""" + core = SimpleNamespace(_prefill="prefill", _decode="decode") + select = dispatch.MiniMaxM3Tp4ModelingV2._select_target + assert select(core, SimpleNamespace(num_contexts=0)) == "decode" + assert select(core, SimpleNamespace(num_contexts=1)) == "prefill" + assert select(core, SimpleNamespace(num_contexts=7)) == "prefill" + + +_DECODE_FIELDS = ( + "num_tokens", + "num_seqs", + "kv_cache_manager", + "msa_idx_k_cache", + "msa_out_cache_loc", + "msa_decode_span", + "msa_block_table", + "msa_seq_lens_cuda", + "_msa_short_ok", + "msa_short_table", + "msa_proxy_max_score_view", + "msa_worst_case_max_k_tiles", + "msa_n_valid_blocks", +) + + +def _decode_metadata(num_tokens: int) -> SimpleNamespace: + """The decode projection's fields and nothing else (no ``num_contexts`` / ``num_ctx_tokens``).""" + md = SimpleNamespace(**{name: object() for name in _DECODE_FIELDS}) + md.num_tokens = num_tokens + return md + + +def test_the_decode_projection_states_the_phase_and_reads_no_count_back() -> None: + md = _decode_metadata(4) + step = dispatch.DecodeTarget(SimpleNamespace()).step_args(md) + assert step["num_contexts"] == 0 and step["num_ctx_tokens"] == 0 + assert step["metadata"] is md and step["msa_short_ok"] is md._msa_short_ok + + +@pytest.mark.parametrize( + "rows,embeds", [(49, False), (64, False), (4, True)], ids=["49_rows", "64_rows", "embeds"] +) +def test_a_wide_generation_step_runs_every_layer_generically_in_decode( + monkeypatch, rows, embeds +) -> None: + """Past the decode kernels' rows (or with embeddings) the decode target runs its own generic layer for every + layer, with its own projection, and never the decode-kernel layer.""" + layers, seen = 3, [] + core = SimpleNamespace(num_layers=layers, w={"embed": None}) + monkeypatch.setattr(dispatch, "m3_embedding", lambda ids, table: ("embedded", int(ids.numel()))) + + def generic(self, i, position_ids, hidden, residual, step, spec_metadata, capture): + seen.append((i, step["num_contexts"], step["metadata"])) + return ("hidden", i), ("residual", i) + + def kernels(self, *args): + raise AssertionError("the decode-kernel layer ran on a wide step") + + monkeypatch.setattr(dispatch.DecodeTarget, "_generic_layer", generic) + monkeypatch.setattr(dispatch.DecodeTarget, "_layer", kernels) + md = _decode_metadata(rows) + ids = torch.arange(rows, dtype=torch.int32) + out = dispatch.DecodeTarget(core).forward( + md, None if embeds else ids, None, ("embeds",) if embeds else None, None, None + ) + assert out == ("hidden", layers - 1) + assert seen == [(i, 0, md) for i in range(layers)] + + +def test_uncovered_never_raises_and_logs_once(log) -> None: + dispatch.set_require(True) + for _ in range(3): + dispatch.uncovered("fused_decode_attn", "a request beyond 16 KV blocks") + assert len(log.info) == 1 and not log.warning + + +class _Core: + """The parts of the core ``DecodeTarget``'s embedding and push targets read: no embedding owners built.""" + + def __init__(self): + self.embed_owners = {} + self.w = {"embed": torch.arange(64 * 4, dtype=torch.float32).view(64, 4)} + self.mapping = None + self.hidden = 6144 + + +@pytest.mark.parametrize("rows", (1, 8)) +def test_embedding_without_its_owner_raises_under_require(rows) -> None: + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=rf"'embed_firstnorm'.* {rows} rows.*owner" + ): + dispatch.DecodeTarget(_Core())._embed(torch.arange(rows, dtype=torch.int32), rows) + + +def test_embedding_without_its_owner_falls_back_and_logs_once(log) -> None: + target = dispatch.DecodeTarget(_Core()) + for _ in range(2): + hidden, residual = target._embed(torch.arange(3, dtype=torch.int32), 3) + assert residual is None and hidden.shape == (3, 4) + assert len(log.warning) == 1 and "embed_firstnorm at 3 rows" in log.warning[0] + + +def test_embedding_past_its_rows_is_uncovered_in_every_mode(log) -> None: + dispatch.set_require(True) + hidden, residual = dispatch.DecodeTarget(_Core())._embed(torch.arange(9, dtype=torch.int32), 9) + assert residual is None and hidden.shape == (9, 4) + assert len(log.info) == 1 and "embed_firstnorm" in log.info[0] and not log.warning + + +def test_a_push_without_its_workspace_raises_under_require(monkeypatch) -> None: + monkeypatch.setattr(m3_mc_allreduce, "get_workspace", lambda mapping, hidden: None) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'attn_push'.* 4 rows"): + dispatch.DecodeTarget(_Core())._attn_push_target(4) + + +def test_a_push_without_its_workspace_falls_back_and_logs_once(monkeypatch, log) -> None: + monkeypatch.setattr(m3_mc_allreduce, "get_workspace", lambda mapping, hidden: None) + target = dispatch.DecodeTarget(_Core()) + assert target._attn_push_target(4) is None + assert target._moe_push_targets(4) is None + assert target._attn_push_target(4) is None + assert len(log.warning) == 2 + assert "attn_push at 4 rows" in log.warning[0] and "moe_push at 4 rows" in log.warning[1] + + +def test_the_target_has_no_dispatch_module_left() -> None: + """The dispatch tables live in ``modeling.py``: no module of the target still reaches a ``dispatch`` module.""" + target_dir = Path(dispatch.__file__).parent + assert not (target_dir / "dispatch.py").exists() + stale = [] + for path in sorted(target_dir.rglob("*.py")): + for node in ast.walk(ast.parse(path.read_text())): + if isinstance(node, ast.ImportFrom) and any(a.name == "dispatch" for a in node.names): + stale.append(f"{path.relative_to(target_dir)}:{node.lineno}") + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "dispatch" + ): + stale.append(f"{path.relative_to(target_dir)}:{node.lineno}: dispatch.{node.attr}") + assert not stale, stale + + +def test_low_m_gemm_cells_follow_its_configurations() -> None: + for (n, k), (warps, cluster_k, row_tiles, max_m) in low_m_entry.CONFIGS.items(): + assert low_m.get_config(n, k) == (warps, cluster_k, max_m, row_tiles) + assert dispatch.LOW_M_GEMM_MAX_TOKENS[(n, k)] == max_m + + +# ------------------------------------------------------------------ the fused draft loop's preconditions + +_HEAD_ROWS = 64 +_HIDDEN = 32 + + +def _drafter(head_weight=None): + """A drafter with only what M3Eagle3Worker._m3_drafter_supported reads: its vocab-sharded BF16 head (the + drafter's one-layer shape is validated when ``M3Eagle3Draft`` is built).""" + weight = ( + head_weight + if head_weight is not None + else torch.zeros((_HEAD_ROWS, _HIDDEN), dtype=torch.bfloat16) + ) + lm_head = SimpleNamespace( + weight=weight, bias=None, tp_mode=TensorParallelMode.COLUMN, padding_size=0, tp_size=4 + ) + return SimpleNamespace(lm_head=lm_head) + + +def _worker(): + """An M3Eagle3Worker past the step's verify: the fused verify left the first-pass gather rows and the fused + first draft inputs; nothing else of the in-tree worker is read.""" + worker = dispatch.M3Eagle3Worker.__new__(dispatch.M3Eagle3Worker) + worker.mapping = SimpleNamespace(tp_size=4, tp_rank=1, enable_attention_dp=False) + worker._d2t = None + worker._m3_draft_head = None + worker._m3_first_gather_ids = torch.zeros(1, dtype=torch.int64) + worker._m3_fused_first_inputs = True + return worker + + +_ATTN_MD = SimpleNamespace(kv_lens_cuda=None) +_SPEC_MD = SimpleNamespace(runtime_draft_len=3) + + +@pytest.fixture +def no_gpu_draft_head(monkeypatch): + """The fused head's CUDA-only lookups replaced: a 148-SM device, and a built exchange workspace.""" + monkeypatch.setattr(m3_draft_step, "_num_sms", lambda: 148) + monkeypatch.setattr(m3_draft_step, "_WS", object()) + + +def test_draft_loop_without_the_prebuilt_head_raises_under_require(no_gpu_draft_head) -> None: + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_draft_step'.* 4 rows.*prebuild"): + _worker()._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, _drafter()) + + +def test_draft_loop_without_the_prebuilt_head_falls_back_and_logs_once( + no_gpu_draft_head, log +) -> None: + worker = _worker() + for _ in range(2): + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, _drafter()) + assert len(log.warning) == 1 and "m3_draft_step at 4 rows" in log.warning[0] + + +def test_draft_loop_with_a_head_built_for_another_weight_raises_under_require( + no_gpu_draft_head, +) -> None: + worker = _worker() + assert worker.prebuild_draft_head(_drafter()) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_draft_step'"): + worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 2, _drafter()) + + +def test_draft_loop_takes_a_covered_step_once_the_head_is_built(no_gpu_draft_head) -> None: + dispatch.set_require(True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + head = worker._m3_draft_head + assert head.weight is drafter.lm_head.weight and head.vocab_offset == _HEAD_ROWS + assert torch.equal(draft_step_entry.from_fragments(head.wf), drafter.lm_head.weight) + for nreq in range(1, 9): + assert worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, nreq, drafter) + + +def test_draft_loop_streams_the_coded_head_once_built(no_gpu_draft_head) -> None: + """A drafter head whose K the coding covers is prebuilt in its 12-bit coding (decoding to the weight bit for bit) + and serves every covered step; one whose coding fails its round trip keeps the BF16 fragments (fail closed).""" + dispatch.set_require(True) + g = torch.Generator().manual_seed(11) + weight = (torch.randn((_HEAD_ROWS, 128), generator=g) * 0.02).to(torch.bfloat16) + drafter = _drafter(weight) + worker = _worker() + assert worker.prebuild_draft_head(drafter) + head = worker._m3_draft_head + assert head.coded and head.weight is weight and head.vocab_offset == _HEAD_ROWS + decoded = draft_step_entry.head_weight( + head.wf, 128, head.exc_off, head.exc_pos, head.exc_w, head.exp_base + ) + assert torch.equal(decoded.view(torch.int16), weight.view(torch.int16)) + for nreq in range(1, 9): + assert worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, nreq, drafter) + + +def test_draft_head_that_fails_its_round_trip_keeps_the_bf16_fragments( + no_gpu_draft_head, monkeypatch +) -> None: + real = target_head_entry.encode_chunked + + def corrupt(weight, *args, **kwargs): + packed, base, exc_off, exc_pos, exc_w = real(weight) + packed = packed.clone() + packed[5] ^= 2 + return packed, base, exc_off, exc_pos, exc_w + + monkeypatch.setattr(m3_draft_step, "encode_chunked", corrupt) + dispatch.set_require(True) + weight = torch.ones((_HEAD_ROWS, 128), dtype=torch.bfloat16) + drafter = _drafter(weight) + worker = _worker() + assert worker.prebuild_draft_head(drafter) + head = worker._m3_draft_head + assert not head.coded and head.exc_off is None + assert torch.equal(draft_step_entry.from_fragments(head.wf), weight) + assert worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + + +def test_draft_loop_past_its_requests_is_uncovered_in_every_mode(no_gpu_draft_head, log) -> None: + dispatch.set_require(True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 9, drafter) + assert not worker._m3_fused_loop_ok(_ATTN_MD, SimpleNamespace(runtime_draft_len=2), 4, drafter) + assert len(log.info) == 2 and not log.warning + + +def test_draft_loop_needs_the_fused_first_inputs(no_gpu_draft_head, log) -> None: + """A step whose first draft inputs came from the in-tree path (or that the fused verify did not prepare) drafts + with the in-tree loop, silently: it is not a draft-step shape.""" + dispatch.set_require(True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + worker._m3_fused_first_inputs = False + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + worker._m3_fused_first_inputs, worker._m3_first_gather_ids = True, None + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + assert not log.info and not log.warning + + +def test_an_unsupported_drafter_is_uncovered(no_gpu_draft_head, log) -> None: + dispatch.set_require(True) + worker = _worker() + padded = _drafter() + padded.lm_head.padding_size = 16 + for drafter in (padded,): + assert not worker.prebuild_draft_head(drafter) + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + assert log.info and not log.warning + + +def test_a_capture_without_the_exchange_workspace_raises_under_require(monkeypatch) -> None: + monkeypatch.setattr(m3_draft_step, "_num_sms", lambda: 148) + monkeypatch.setattr(m3_draft_step, "_WS", None) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_draft_step'.*workspace.*captured"): + worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + + +# ------------------------------------------------------------------ the coded target head's preconditions + +_TARGET_ROWS = 64 +_TARGET_K = 128 + + +def _target_lm_head(weight=None): + if weight is None: + g = torch.Generator().manual_seed(7) + weight = (torch.randn(_TARGET_ROWS, _TARGET_K, generator=g) * 0.02).to(torch.bfloat16) + return SimpleNamespace( + weight=weight, bias=None, tp_mode=TensorParallelMode.COLUMN, padding_size=0, tp_size=4 + ) + + +@pytest.fixture +def no_gpu_target_head(monkeypatch): + monkeypatch.setattr(m3_target_head, "_num_sms", lambda: 148) + + +def _target_worker(): + worker = dispatch.M3Eagle3Worker.__new__(dispatch.M3Eagle3Worker) + worker._m3_target_head = None + return worker + + +def test_target_head_without_the_prebuilt_copy_raises_under_require() -> None: + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_target_head'.* 4 rows.*prebuild"): + _target_worker()._m3_coded_head(_target_lm_head(), 4) + + +def test_target_head_without_the_prebuilt_copy_falls_back_and_logs_once(log) -> None: + worker = _target_worker() + for _ in range(2): + assert worker._m3_coded_head(_target_lm_head(), 8) is None + assert len(log.warning) == 1 and "m3_target_head at 8 rows" in log.warning[0] + + +def test_target_head_takes_a_covered_verify_once_built(no_gpu_target_head) -> None: + """The prebuilt copy decodes to the head's weight bit for bit and serves every covered row count; a copy built + for another weight is a miss.""" + dispatch.set_require(True) + lm_head = _target_lm_head() + worker = _target_worker() + assert worker.prebuild_target_head(lm_head) + head = worker._m3_target_head + assert head.coded and head.weight is lm_head.weight and head.grid == 148 + decoded = target_head_entry.decode( + head.w, head.base, head.exc_off, head.exc_pos, head.exc_w, _TARGET_ROWS, _TARGET_K + ) + assert torch.equal(decoded.view(torch.int16), lm_head.weight.view(torch.int16)) + for rows in range(1, 9): + assert worker._m3_coded_head(lm_head, rows) is head + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_target_head'"): + worker._m3_coded_head(_target_lm_head(lm_head.weight.clone()), 4) + + +def test_target_head_past_its_rows_is_uncovered_in_every_mode(no_gpu_target_head, log) -> None: + dispatch.set_require(True) + lm_head = _target_lm_head() + worker = _target_worker() + assert worker.prebuild_target_head(lm_head) + for rows in (9, 12, 32, 48): + assert worker._m3_coded_head(lm_head, rows) is None + assert log.info and not log.warning + + +def test_target_head_that_fails_its_round_trip_keeps_the_bf16_gemm( + no_gpu_target_head, monkeypatch, log +) -> None: + """Fail closed: a coding that does not decode to the weight bit for bit is dropped at prebuild (a warning), and the + verify runs the in-tree GEMM without a miss, in every mode.""" + real = target_head_entry.encode_chunked + + def corrupt(weight, *args, **kwargs): + packed, base, exc_off, exc_pos, exc_w = real(weight) + packed = packed.clone() + packed[3] ^= 1 + return packed, base, exc_off, exc_pos, exc_w + + monkeypatch.setattr(m3_target_head, "encode_chunked", corrupt) + dispatch.set_require(True) + lm_head = _target_lm_head() + worker = _target_worker() + assert not worker.prebuild_target_head(lm_head) + assert not worker._m3_target_head.coded + assert len(log.warning) == 1 and "bit for bit" in log.warning[0] + assert worker._m3_coded_head(lm_head, 4) is None + assert any("round trip" in line for line in log.info) + + +def test_an_unsupported_target_head_is_not_coded(no_gpu_target_head) -> None: + worker = _target_worker() + padded = _target_lm_head() + padded.padding_size = 16 + narrow = _target_lm_head(torch.zeros((_TARGET_ROWS, 96), dtype=torch.bfloat16)) + for lm_head in (padded, narrow): + assert not worker.prebuild_target_head(lm_head) + assert worker._m3_target_head is None + + +# ------------------------------------------------------------------ the early fc fork + + +class _CaptureMetadata: + def __init__(self, nreq: int): + self.layers_to_capture = (1, 29, 56) + self.num_capture_tokens = nreq * 4 + self.runtime_draft_len = 3 + self.captured = [] + + def maybe_capture_hidden_states(self, layer_id, hidden_states, residual) -> None: + self.captured.append(layer_id) + + +def _fc_early_worker(monkeypatch, eligible: bool = True): + worker = dispatch.M3Eagle3Worker.__new__(dispatch.M3Eagle3Worker) + worker._m3_capture_norm = None + worker._m3_fc_norm_w = None + worker._m3_fc_early_drafter = None + worker._m3_fc_early_stream = None + forks = [] + monkeypatch.setattr(worker, "tail_eligible", lambda *args: eligible) + + def fork(md, num_tokens, draft_model, after=None): + forks.append((num_tokens, draft_model, after)) + return f"stream-{len(forks)}" + + monkeypatch.setattr(worker, "_m3_fc_fork", fork) + return worker, forks + + +@pytest.mark.parametrize("nreq", [1, 2]) +def test_the_fc_is_forked_at_the_last_captured_layer(monkeypatch, nreq) -> None: + """A fused-verify step of at most 8 verify rows forks the drafter's fc once, right after the last captured layer + (every captured layer still captures), and the fork is left for the verify to join.""" + worker, forks = _fc_early_worker(monkeypatch) + md = _CaptureMetadata(nreq) + drafter = object() + worker.arm_capture_norm(SimpleNamespace(num_seqs=nreq), md, object(), drafter) + for layer_id in md.layers_to_capture[:-1]: + worker.capture_hidden_states(md, layer_id, None) + assert not forks + worker.capture_hidden_states(md, md.layers_to_capture[-1], None) + worker.capture_hidden_states(md, dispatch.FINAL_HIDDEN_LAYER_ID, None) + assert md.captured == [1, 29, 56, dispatch.FINAL_HIDDEN_LAYER_ID] + assert forks == [(4 * nreq, drafter, None)] + assert worker._m3_fc_early_stream == "stream-1" and worker._m3_fc_early_drafter is None + + +def test_no_early_fork_with_a_final_state_capture(monkeypatch) -> None: + """The final-state column is written after the model returns, so captures that include it keep the fc at its + stock place instead of forking at the last captured decoder layer.""" + worker, forks = _fc_early_worker(monkeypatch) + md = _CaptureMetadata(1) + md.layers_to_capture = (dispatch.FINAL_HIDDEN_LAYER_ID, 29, 56) + worker.arm_capture_norm(SimpleNamespace(num_seqs=1), md, object(), object()) + for layer_id in (29, 56, dispatch.FINAL_HIDDEN_LAYER_ID): + worker.capture_hidden_states(md, layer_id, None) + assert not forks and worker._m3_fc_early_stream is None + + +@pytest.mark.parametrize("nreq,eligible", [(3, True), (8, True), (1, False)]) +def test_no_early_fork_past_its_rows_or_off_the_fused_verify(monkeypatch, nreq, eligible) -> None: + worker, forks = _fc_early_worker(monkeypatch, eligible) + md = _CaptureMetadata(nreq) + worker._m3_fc_early_stream = "stale" + worker.arm_capture_norm(SimpleNamespace(num_seqs=nreq), md, object(), object()) + for layer_id in md.layers_to_capture: + worker.capture_hidden_states(md, layer_id, None) + assert not forks and worker._m3_fc_early_stream is None + + +# ------------------------------------------------------------------ the NVFP4 producer's coverage + + +def _producer_core(nvfp4: bool = True, heads: int = 16): + """What the producer dispatch reads off the core before its first precondition.""" + return SimpleNamespace(main_kv_is_nvfp4=nvfp4, heads=heads, kv_heads=1, index_heads=1) + + +def test_the_producer_covers_the_tp4_nvfp4_layers_only() -> None: + """``PRODUCER`` covers a sparse layer on an NVFP4 cache at the TP4 heads; elsewhere the generic producer runs + without a miss.""" + covered = dispatch._producer_covered + assert covered(_producer_core()) + assert not covered(_producer_core(nvfp4=False)) + assert not covered(_producer_core(heads=8)) + assert not covered(SimpleNamespace(main_kv_is_nvfp4=True, heads=16, kv_heads=2, index_heads=2)) + assert dispatch.PRODUCER.covers("nvfp4_tp4") + + +def _produce(core): + packed = torch.zeros(4, 20 * 128, dtype=torch.bfloat16) + return dispatch._sparse_qkv_producer(core, 3, packed, None, None, None, None) + + +def test_a_covered_producer_layer_missing_a_precondition_raises_under_require() -> None: + """A TP4 NVFP4 sparse layer whose step lacks a producer precondition (here the position ids) is a miss.""" + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=r"'m3_producer' covers this step at 4 rows" + ): + _produce(_producer_core()) + + +def _runtime(fused: bool, indexer: str = "fp8", kv: str = "fp8"): + """A model config with the runtime fields ``_validate_runtime`` reads.""" + sparse = SimpleNamespace( + algorithm="minimax_m3", + implementation="msa", + fuse_qkv_index_projection=fused, + indexer_kv_dtype=indexer, + ) + mode = SimpleNamespace( + has_fp8_kv_cache=lambda: kv == "fp8", has_fp4_kv_cache=lambda: kv == "nvfp4" + ) + return SimpleNamespace( + sparse_attention_config=sparse, quant_config=SimpleNamespace(quant_mode=mode) + ) + + +@pytest.mark.parametrize("fused", [True, False], ids=["five_way", "unfused"]) +def test_the_target_takes_both_projection_layouts(fused) -> None: + """The packed five-way projection and the separate QKV + BF16 index projections + (fuse_qkv_index_projection=False) are both this target's; a BF16 indexer or a BF16 KV cache is not.""" + dispatch._validate_runtime(_runtime(fused)) + dispatch._validate_runtime(_runtime(fused, kv="nvfp4")) + with pytest.raises(ValueError, match="indexer_kv_dtype='fp8'"): + dispatch._validate_runtime(_runtime(fused, indexer="bf16")) + with pytest.raises(ValueError, match="FP8 or NVFP4 KV cache"): + dispatch._validate_runtime(_runtime(fused, kv="bf16")) + + +@pytest.mark.parametrize("target", ["PrefillTarget", "DecodeTarget"]) +def test_an_unfused_layer_the_producer_covers_is_a_miss(target) -> None: + """As the built-in target: a layer the NVFP4 producer covers needs the five-way projection, so the unfused path is + a miss there (raised under require) before anything runs, in either target.""" + dispatch.set_require(True) + attention = getattr(dispatch, target)(_producer_core())._unfused_sparse_attention + with pytest.raises(dispatch.FastPathUnavailable, match="fuse_qkv_index_projection is disabled"): + if target == "PrefillTarget": + attention(3, None, None, None) + else: + attention(3, None, None, None, None) + + +def test_a_covered_producer_miss_falls_back_and_logs_once(log) -> None: + core = _producer_core() + assert _produce(core) is None and _produce(core) is None + assert len(log.warning) == 1 and "m3_producer at 4 rows" in log.warning[0] + + +@pytest.mark.parametrize("nvfp4,heads", [(False, 16), (True, 8)], ids=["fp8_cache", "tp8_heads"]) +def test_an_uncovered_producer_layer_falls_back_silently(log, nvfp4, heads) -> None: + dispatch.set_require(True) + assert _produce(_producer_core(nvfp4, heads)) is None + assert not log.warning + + +def test_a_covered_producer_layer_with_the_wrong_pool_layout_is_not_supported() -> None: + """The pool-layout check the covered branch takes before the generic producer (then a miss).""" + pool = torch.zeros(4, 2, 1, 128, 32, dtype=torch.uint8) + scales = torch.zeros(4, 2, 1, 128, 8, dtype=torch.uint8) + assert not dispatch._producer_layout(pool, scales) + good = torch.zeros(4, 2, 1, 128, 64, dtype=torch.uint8) + assert dispatch._producer_layout(good, scales) + assert not dispatch._producer_layout(good, scales.transpose(3, 4).contiguous().transpose(3, 4)) + + +# ------------------------------------------------------------------ the shell's speculative plumbing + + +def test_the_shell_installs_guided_decoding_on_the_worker() -> None: + """With speculation the guided decoder goes to the worker (the in-tree one-engine shell's delegation); without a + worker the shell reports it cannot take one.""" + shell_cls = dispatch.ModelingV2MiniMaxM3Nvfp4Sm103Tp4 + shell = shell_cls.__new__(shell_cls) + torch.nn.Module.__init__(shell) + installed = [] + shell.spec_worker = SimpleNamespace(set_guided_decoder=lambda g: installed.append(g) or True) + decoder = object() + assert shell.set_guided_decoder(decoder) and installed == [decoder] + shell.spec_worker = None + assert shell.set_guided_decoder(decoder) is False diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py new file mode 100644 index 000000000000..df4afe2f15b3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py @@ -0,0 +1,302 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode kernels the core builds before the executor's warmup (``prebuild_kernels``), and that a forward after it +never compiles one: the MoE engines of every ``MOE_ENGINE`` / ``MOE16_ENGINE`` token count, the decode GEMVs with +their push variants, and the index scorer with and without its early-issue schedule at every decode query length. + +The engine builds are stubbed (a real one needs a loaded checkpoint's expert buffers; the engine kernels have their +own tests); the GEMVs and the scorer compile for real on the GPU. +""" + +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers the trtllm ops +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops import ( # noqa: E501 + m3_index_score as index_score_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_index_score import ( # noqa: E501 + m3_index_score, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm.m3_mx_gemv import ( # noqa: E501 + O_PROJ, + SHARED_DOWN, + SHARED_GATE_UP, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.mx_gemv import ( # noqa: E501 + op as mx_gemv, +) + +dispatch = modeling + +assert torch.cuda.is_available(), "the prebuild checks require a CUDA device" + +INDEX_HEADS = 4 +INDEX_DIM = 128 +DRAFT_LEN = 3 +_CORE_METHODS = ( + "prebuild_kernels", + "_prebuild_engines", + "_try_build", + "engine", + "moe16_engine", + "_lookup", + "_scorer_variants", + "_prebuild_index_scorer", + "decode_query_lens", +) + + +@pytest.fixture(autouse=True) +def _restore_modes(): + require, frozen = dispatch.require_mode(), mx_gemv.frozen() + yield + dispatch.set_require(require) + mx_gemv.freeze(frozen) + dispatch.reset_logged() + + +class _FakeSparseAttn: + def __init__(self, indexer_kv_dtype: str): + self.indexer_kv_dtype = indexer_kv_dtype + self.m3_config = SimpleNamespace(num_index_heads=INDEX_HEADS, sparse_index_dim=INDEX_DIM) + + +def _core(num_layers: int = 3, indexer_kv_dtype: str = "fp8", built=None, fail=False): + """The parts of ``MiniMaxM3Tp4ModelingV2`` the prebuild reads: one rank (the build agreement takes no + collective), MoE + sparse layers 1.. (layer 0 dense), no projections (the GEMV prebuild has its own test), the + speculative config; the engine builds stubbed.""" + core = SimpleNamespace( + mapping=SimpleNamespace(tp_size=1, enable_attention_dp=False), + moe_layers=frozenset(range(1, num_layers)), + sparse_layers=frozenset(range(1, num_layers)), + attn=[_FakeSparseAttn(indexer_kv_dtype) for _ in range(num_layers)], + w={"embed": torch.empty(1, device="cuda")}, + mxfp8=SimpleNamespace(quant_method=SimpleNamespace(use_cutlass=True)), + model_config=SimpleNamespace( + spec_config=SimpleNamespace(max_draft_len=DRAFT_LEN, max_total_draft_tokens=DRAFT_LEN) + ), + moe_engines={}, + moe16_engines={}, + projections=lambda: [], + ) + for name in _CORE_METHODS: + setattr(core, name, getattr(modeling.MiniMaxM3Tp4ModelingV2, name).__get__(core)) + built = [] if built is None else built + + def build(i, num_tokens, shared=True): + if fail: + raise ValueError("the routed experts are not TRTLLM-Gen NVFP4 SwiGLU-OAI") + built.append(("moe", i, num_tokens, shared)) + return ("engine", i, num_tokens) + + def build16(i, num_tokens): + if fail: + raise ValueError("the routed experts are not TRTLLM-Gen NVFP4 SwiGLU-OAI") + built.append(("moe16", i, num_tokens)) + return ("engine16", i, num_tokens) + + core._build_engine, core._build_moe16_engine = build, build16 + return core + + +def _forbid_compiles(monkeypatch, core=None) -> None: + """Every compile / build entry point a forward could reach raises.""" + import cutlass.cute as cute + + def boom(*args, **kwargs): + raise AssertionError("a compile entry point ran after prebuild") + + monkeypatch.setattr(cute, "compile", boom) + monkeypatch.setattr(mx_gemv, "ensure_compiled", boom) + monkeypatch.setattr(mx_gemv, "precompile", boom) + monkeypatch.setattr(index_score_kernel.IndexScoreRunner, "_compile", boom) + if core is not None: + core._build_engine = core._build_moe16_engine = boom + + +def test_decode_query_lens_follow_the_speculative_config() -> None: + core = _core() + assert core.decode_query_lens() == [1, 1 + DRAFT_LEN] + core.model_config.spec_config = None + assert core.decode_query_lens() == [1] + core.model_config.spec_config = SimpleNamespace( + max_draft_len=4, max_total_draft_tokens=4, draft_len_schedule={1: 4, 8: 2} + ) + assert core.decode_query_lens() == [1, 3, 5] + + +def test_prebuild_builds_every_engine_count_and_scorer_variant() -> None: + built = [] + core = _core(built=built) + times = core.prebuild_kernels() + + assert mx_gemv.frozen(), "prebuild must close the GEMV compiles" + + # Layer 0 is dense; every MoE layer gets every count, in layer order. + def per_layer(i): + engines = [("moe", i, n, True) for n in range(1, dispatch.MOE_ENGINE_MAX_TOKENS + 1)] + return engines + [ + ("moe16", i, n) for n in range(dispatch.MOE16_MIN_TOKENS, dispatch.MOE16_MAX_TOKENS + 1) + ] + + assert built == per_layer(1) + per_layer(2) + assert set(core.moe_engines[2]) == {(n, True, True) for n in dispatch.MOE_ENGINE.keys} + assert set(core.moe16_engines[2]) == set(dispatch.MOE16_ENGINE.keys) + + want = { + (torch.float8_e4m3fn, INDEX_HEADS, qlen, INDEX_DIM, modeling.INDEX_SCORE_SPLIT_K, early) + for qlen in (1, 1 + DRAFT_LEN) + for early in (True, False) + } + assert set(times["scorer_keys"]) == want + assert want <= set(index_score_kernel.IndexScoreRunner.kernel_cache) + for key in ("gemv_s", "moe_s", "scorer_s"): + assert times[key] >= 0.0 + + # A second build builds no engine again. + built.clear() + core.prebuild_kernels() + assert built == [] + + +def test_engine_lookups_after_prebuild_never_build(monkeypatch) -> None: + core = _core() + core._prebuild_engines(5) + _forbid_compiles(monkeypatch, core) + for require in (False, True): + dispatch.set_require(require) + for n in sorted(dispatch.MOE_ENGINE.keys): + assert core.engine(5, n, shared=True) == ("engine", 5, n) + for n in sorted(dispatch.MOE16_ENGINE.keys): + assert core.moe16_engine(5, n) == ("engine16", 5, n) + + +def test_missing_engine_is_a_miss_not_a_build(monkeypatch) -> None: + core = _core() + core.moe_engines[7], core.moe16_engines[7] = {}, {} + _forbid_compiles(monkeypatch, core) + dispatch.set_require(False) + assert core.engine(7, 3, shared=True) is None + assert core.moe16_engine(7, 10) is None + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="not built before warmup"): + core.engine(7, 3, shared=True) + with pytest.raises(dispatch.FastPathUnavailable, match="not built before warmup"): + core.moe16_engine(7, 10) + + +def test_failed_engine_build_is_logged_then_missed() -> None: + core = _core(fail=True) + core._prebuild_engines(9) + assert all(v is None for v in core.moe_engines[9].values()) and len(core.moe_engines[9]) == 8 + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="build failed"): + core.engine(9, 1, shared=True) + + +def _projection(nk, push): + n, k = nk + w = (torch.randn(n, k, device="cuda") / k**0.5).to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + return SimpleNamespace(weight=wq.view(n, k), weight_scale=wsf, push=push) + + +@pytest.mark.parametrize( + "nk,push", + ((O_PROJ, "attn"), (SHARED_DOWN, "moe"), (SHARED_GATE_UP, None)), + ids=("o_proj", "shared_down", "shared_gate_up"), +) +def test_gemvs_after_prebuild_never_compile(monkeypatch, nk, push) -> None: + use_pdl = True + module = _projection(nk, push) + variants = modeling._gemv_variants(*nk, push, 4) + pushes = [v for v in variants if v[2] is not None] + if push is None: + assert not pushes + else: + covered = dispatch.ATTN_PUSH.keys if push == "attn" else dispatch.MOE_PUSH.keys + assert {rows for rows, _, _ in pushes} == set(covered) + assert {p[1] for _, _, p in pushes} == {b for b in (8, 16) if b <= max(covered)} + mx_gemv.freeze(False) + # As prebuild_kernels compiles one projection. + mx_gemv.ensure_compiled(module.weight, module.weight_scale, use_pdl) + for rows, x_bf16, ar in variants: + mx_gemv.precompile(module.weight, module.weight_scale, rows, x_bf16, use_pdl=use_pdl, ar=ar) + mx_gemv.freeze(True) + _forbid_compiles(monkeypatch) + + n, k = nk + for rows, x_bf16, ar in variants: + if ar is not None: + # A push needs a multicast exchange to launch: its kernel is compiled (the forward finds it). + resolved = mx_gemv._resolve(rows, n, k, x_bf16, True, None, use_pdl, (0, 0, 0, *ar)) + if resolved is not None: + assert mx_gemv.is_compiled(resolved[0]), (rows, x_bf16, ar) + continue + if x_bf16: + x = torch.randn(rows, k, device="cuda").to(torch.bfloat16) + out = mx_gemv.m3_mx_gemv_bf16(x, module.weight, module.weight_scale, use_pdl=use_pdl) + else: + xq, xsf = torch.ops.trtllm.mxfp8_quantize( + torch.randn(rows, k, device="cuda").to(torch.bfloat16), True + ) + out = mx_gemv.m3_mx_gemv( + xq.view(rows, k), xsf, module.weight, module.weight_scale, use_pdl=use_pdl + ) + served = mx_gemv.serves(n, k, rows, x_bf16) + assert (out is not None) == served, (rows, x_bf16) + assert mx_gemv.last_uncompiled is None + torch.cuda.synchronize() + + +def test_frozen_gemv_declines_an_uncompiled_kernel(monkeypatch) -> None: + module = _projection(SHARED_DOWN, None) + mx_gemv.freeze(True) + _forbid_compiles(monkeypatch) + # A launch option nothing prebuilt (no early trigger, without a push): declined, never compiled. + xq, xsf = torch.ops.trtllm.mxfp8_quantize( + torch.randn(3, 768, device="cuda").to(torch.bfloat16), True + ) + key = mx_gemv._resolve(3, 6144, 768, False, False, None, True, None)[0] + if mx_gemv.is_compiled(key): + pytest.skip("this process already compiled the probe kernel") + out = mx_gemv.m3_mx_gemv( + xq.view(3, 768), xsf, module.weight, module.weight_scale, trigger_early=False, use_pdl=True + ) + assert out is None and mx_gemv.last_uncompiled == key + + +@pytest.mark.parametrize("indexer_kv_dtype", ("fp8", "bf16")) +def test_scorer_after_prebuild_never_compiles(monkeypatch, indexer_kv_dtype) -> None: + core = _core(indexer_kv_dtype=indexer_kv_dtype) + keys = core._prebuild_index_scorer(core.decode_query_lens(), torch.device("cuda")) + assert len(keys) == 4 + _forbid_compiles(monkeypatch, core) + dtype = torch.float8_e4m3fn if indexer_kv_dtype == "fp8" else torch.bfloat16 + for _, heads, qlen, dim, split_k, early in keys: + rows = 1 if early else 2 + blocks = 3 + idx_q = torch.randn(rows * qlen, heads, dim, device="cuda").to(dtype) + idx_k = torch.randn(rows * blocks, 1, 128, dim, device="cuda").to(dtype) + block_table = torch.arange(rows * blocks, dtype=torch.int32, device="cuda").view( + rows, blocks + ) + seq_lens = torch.full((rows,), blocks * 128 - 5, dtype=torch.int32, device="cuda") + score = torch.empty(heads, blocks, rows * qlen, dtype=torch.float32, device="cuda") + m3_index_score( + idx_q, + idx_k.squeeze(1), + block_table, + seq_lens, + score.transpose(1, 2), + qlen, + split_k, + early, + ) + torch.cuda.synchronize() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_rank_agreement.py new file mode 100644 index 000000000000..ec23eb22fc9d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_rank_agreement.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode robustness without a GPU kernel: the multicast all-reduce epilogue options keep the cluster-barrier bit, +an engine build that fails on one TP rank drops that engine on every rank (two ranks simulated with threads), and +the Eagle3 fc fork's output is recorded on the stream that reads it.""" + +import threading +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops import ( # noqa: E501 + m3_mc_allreduce as mc, +) + + +@pytest.fixture(autouse=True) +def _restore_modes(): + require = modeling.require_mode() + yield + modeling.set_require(require) + modeling.reset_logged() + + +@pytest.mark.parametrize("finalize", [False, True]) +@pytest.mark.parametrize("bucket", [8, 16, 32, 48]) +def test_ar_opts_keep_cluster_barrier_bit(monkeypatch, bucket, finalize): + # Tables without bit 2 and with buckets missing: every lookup, fallback included, keeps the barrier. + monkeypatch.setattr(mc, "AR_OPTS_BY_M", (0, 1, 4, 5, 8, 9, 12, 13)) + monkeypatch.setattr(mc, "AR_OPTS_BY_BUCKET", {}) + monkeypatch.setattr(mc, "AR_FIN_OPTS_BY_BUCKET", {}) + for m in range(1, min(bucket, 8) + 1): + opts = mc._ar_opts(m, bucket, finalize=finalize) + assert opts & mc._OPT_CLUSTER_BARRIER and 0 <= opts <= 15, (m, bucket, finalize, opts) + + +def test_ar_opts_shipped_tables_unchanged(): + for m in range(1, 9): + assert mc._ar_opts(m, 8) == mc.AR_OPTS_BY_M[m - 1] + for b, o in mc.AR_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b) == o + for b, o in mc.AR_FIN_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b, finalize=True) == o + + +class _FakeComm: + """An MPI-like TP communicator for ``world`` threads: allreduce sums the ranks' ints.""" + + def __init__(self, world): + self.world = world + self.barrier = threading.Barrier(world) + self.lock = threading.Lock() + self.vals = [] + self.calls = 0 + + def allreduce(self, v): + with self.lock: + self.vals.append(v) + self.calls += 1 + self.barrier.wait() + total = sum(self.vals) + self.barrier.wait() + with self.lock: + self.vals.clear() + self.barrier.wait() + return total + + def Get_size(self): # noqa: N802 (MPI name) + return self.world + + +def _mapping(tp_size=2, adp=False): + return SimpleNamespace(tp_size=tp_size, enable_attention_dp=adp, tp_group=list(range(tp_size))) + + +def _run_ranks(fns): + out, errs = [None] * len(fns), [] + + def run(i): + try: + out[i] = fns[i]() + except BaseException as e: # noqa: BLE001 + errs.append(e) + + threads = [threading.Thread(target=run, args=(i,)) for i in range(len(fns))] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + assert not errs, errs + return out + + +@pytest.fixture +def two_rank_comm(monkeypatch): + from tensorrt_llm._torch.distributed import ops + + comm = _FakeComm(2) + monkeypatch.setattr(mc, "_tp_comm", lambda mapping: comm) + monkeypatch.setattr(ops, "mpi_disabled", lambda: False) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + return comm + + +def test_tp_agree_single_rank_and_adp_take_no_collective(monkeypatch): + monkeypatch.setattr(mc, "_tp_comm", mock.Mock(side_effect=AssertionError("no collective"))) + assert mc.tp_agree(_mapping(1), True) is True + assert mc.tp_agree(_mapping(1), False) is False + assert mc.tp_agree(_mapping(4, adp=True), False) is False + + +def test_tp_agree_refuses_capture(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError): + mc.tp_agree(_mapping(2), True) + + +@pytest.mark.parametrize("oks", [(True, True), (True, False), (False, True), (False, False)]) +def test_tp_agree_built_two_ranks(two_rank_comm, oks): + built = [object() if ok else None for ok in oks] + out = _run_ranks([lambda i=i: mc.tp_agree_built(_mapping(2), built[i], "t") for i in range(2)]) + assert out == (built if all(oks) else [None, None]) + + +def _core(fail: bool) -> SimpleNamespace: + """The engine-building half of the core (``MiniMaxM3Tp4ModelingV2``) on one of two simulated TP ranks, its + engine builds stubbed (and failing on one rank).""" + core = SimpleNamespace(mapping=_mapping(2), moe_engines={}, moe16_engines={}) + cls = modeling.MiniMaxM3Tp4ModelingV2 + + def build(i, num_tokens, shared=True): + if fail: + raise RuntimeError("compile failed on this rank") + return ("engine", num_tokens) + + core._build_engine = build + core._build_moe16_engine = lambda i, num_tokens: build(i, num_tokens) + for name in ("_prebuild_engines", "_try_build", "engine", "moe16_engine", "_lookup"): + setattr(core, name, getattr(cls, name).__get__(core)) + return core + + +@pytest.mark.parametrize("fail", [(False, False), (False, True)]) +def test_engine_build_agreed_across_ranks(two_rank_comm, fail): + cores = [_core(f) for f in fail] + _run_ranks([lambda c=c: c._prebuild_engines(3) for c in cores]) + keys = sorted(cores[0].moe_engines[3]) + sorted(cores[0].moe16_engines[3]) + assert keys and two_rank_comm.calls == 2 * len(keys) + for engines in ("moe_engines", "moe16_engines"): + a, b = getattr(cores[0], engines)[3], getattr(cores[1], engines)[3] + assert a.keys() == b.keys() + for key in a: + assert (a[key] is None) == (b[key] is None) == any(fail), (engines, key) + # Both ranks then take the same path in the forward: the engine, or (after a miss) the generic MoE. + modeling.set_require(False) + for n in sorted(modeling.MOE_ENGINE.keys): + got = [c.engine(3, n, shared=True) for c in cores] + assert (got[0] is None) == (got[1] is None) + + +@pytest.mark.parametrize("capturing", [False, True]) +def test_fc_join_records_output_on_main(monkeypatch, capturing): + M3Eagle3Worker = modeling.M3Eagle3Worker + main, fork = mock.Mock(name="main"), mock.Mock(name="fork") + hidden = mock.Mock(name="hidden", is_cuda=True) + monkeypatch.setattr(torch.cuda, "current_stream", lambda *a: main) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: capturing) + worker = SimpleNamespace(_m3_fc_hidden=hidden) + M3Eagle3Worker._m3_fc_join(worker, fork) + main.wait_stream.assert_called_once_with(fork) + if capturing: + hidden.record_stream.assert_not_called() + else: + hidden.record_stream.assert_called_once_with(main) + worker._m3_fc_hidden = None + M3Eagle3Worker._m3_fc_join(worker, fork) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/activation/test_modeling_v2_m3_silu_and_mul.py new file mode 100644 index 000000000000..470ed6d7cbd1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/activation/test_modeling_v2_m3_silu_and_mul.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_silu_and_mul entry of the MiniMax-M3 tp8 target: SwiGLU-OAI and plain SwiGLU.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.activation.m3_silu_and_mul import ( # noqa: E501 + m3_silu_and_mul as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_silu_and_mul requires a CUDA device" + +OAI = (7.0, 1.702, 1.0) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + # Gate and up spread past the clamp (|x| up to ~12) so both clamps engage. + x = (torch.randn(spec["m"], 2 * spec["inter"], generator=g, device="cuda") * 4.0).to( + torch.bfloat16 + ) + return (x, *OAI) if spec["oai"] else (x, None, None, None) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """fp32 arithmetic and one bf16 rounding on both sides: the kernel's fast sigmoid moves the fp32 value by a few + fp32 ulps, which flips the bf16 rounding of an element by at most one step (1 unit of ULP[bf16] of the row); + measured on GB300: worst element 0.5 unit, relative RMS 0.01.""" + args = _build(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(*args) + assert out.shape == (cell.spec["m"], cell.spec["inter"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=1.0, rms_ulp=0.1) + + +def test_the_gate_sees_a_missing_clamp() -> None: + """Dropping the SwiGLU-OAI clamps moves every element whose gate or up passed the limit far outside the band.""" + spec = op.CELLS[0].spec + x, *oai = _build(dict(spec, m=64), 3) + out = op(x, *oai) + with pytest.raises(AssertionError): + assert_within_ulp(out, op.reference(x, None, oai[1], oai[2]), element_ulp=1.0, rms_ulp=0.1) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/_m3_attn_runtime_tp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/_m3_attn_runtime_tp8.py new file mode 100644 index 000000000000..e9aaad5c65fa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/_m3_attn_runtime_tp8.py @@ -0,0 +1,274 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real MiniMax-M3 attention runtime state at the TP8 per-rank geometry, for the backend-step entries' GPU tests. + +One hybrid NVFP4 KV cache manager (``MiniMaxM3KVCacheManagerV2``) laid out as the target's: a dense layer in the FP8 +half, a sparse MSA layer in the NVFP4 half with its FP8 index-K cache, and (``draft=True``) the one-model Eagle3 +drafter's shared layer, read through the manager's FP8 draft view as the runtime's draft swap reads it. Metadata is +the runtime's own type, filled by its own ``prepare()``; the backends are built the way the target builds them. + +The geometry is one rank's: 8 query heads over one KV head (M3's 4 KV heads, each duplicated on a rank pair) and +one index head of dim 128. The mapping is a single rank holding exactly those heads, so nothing is sharded again. + +History rows are written by the in-tree scatters (``fused_write_layer_caches`` for the FP8 layers, +``fused_write_layer_caches_nvfp4`` for the NVFP4 layer and its index-K), at the slots the manager assigned, which are +not in request order. +""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import ( + MiniMaxM3KVCacheManagerV2, + MiniMaxM3MsaSparseAttention, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches, + fused_write_layer_caches_nvfp4, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.llmapi.llm_args import ( + Eagle3DecodingConfig, + KvCacheConfig, + MiniMaxM3SparseAttentionConfig, +) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +HEADS = 8 +KV_HEADS = 1 +INDEX_HEADS = 1 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +DENSE, SPARSE, DRAFT = 0, 1, 2 +MAX_BATCH = 8 +MAX_SEQ_LEN = 8192 +MAX_NUM_TOKENS = 8192 +#: The Eagle3 drafter checkpoint's RoPE: NeoX over the whole head, theta 5e6. +DRAFT_ROPE_THETA = 5.0e6 + +#: E4M3-exact planted index levels, 1/8 apart in relative terms (0.25 .. 15): 48 distinct block scores. +LEVELS = tuple(2.0**e * (1 + m / 8) for e in range(-2, 4) for m in range(8)) +#: The index value of a block's non-peak rows; its score (8) is below every level's (>= 16). +LOW = 0.125 + + +def mapping() -> Mapping: + return Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + + +def sparse_attention_config() -> MiniMaxM3SparseAttentionConfig: + """The target's runtime: MSA, the FP8 indexer, the fused five-way projection (M3's checkpoint selection rule: + no forced initial blocks, the query's own block forced).""" + return MiniMaxM3SparseAttentionConfig( + implementation="msa", + indexer_kv_dtype="fp8", + fuse_qkv_index_projection=True, + sparse_num_index_heads=INDEX_HEADS, + num_attention_heads=HEADS, + num_key_value_heads=KV_HEADS, + ) + + +def make_manager(*, draft: bool) -> MiniMaxM3KVCacheManagerV2: + torch.cuda.init() + return MiniMaxM3KVCacheManagerV2( + KvCacheConfig( + max_tokens=4 * MAX_SEQ_LEN, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + host_cache_size=0, + event_buffer_max_size=0, + ), + CacheType.SELF, + num_layers=2, + num_kv_heads=[KV_HEADS] * (3 if draft else 2), + head_dim=HEAD_DIM, + tokens_per_block=PAGE, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + mapping=mapping(), + dtype=DataType.NVFP4, + vocab_size=16, + spec_config=( + Eagle3DecodingConfig(max_draft_len=3, speculative_model="unused") if draft else None + ), + sparse_layer_ids=[SPARSE], + sparse_attention_config=sparse_attention_config(), + ) + + +def make_msa_backend(layer: int) -> MiniMaxM3MsaSparseAttention: + """One layer's MSA backend as ``MiniMaxM3Tp8ModelingV2`` builds it (NVFP4 KV quant config on every layer).""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.NVFP4) + backend = create_attention( + "TRTLLM", + layer, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=None, + quant_config=quant, + q_scaling=1.0, + sparse_params=sparse_attention_config().to_sparse_params(layer_idx=layer), + ) + backend.update_quant_config(quant) + assert isinstance(backend, MiniMaxM3MsaSparseAttention), type(backend).__name__ + return backend + + +def make_draft_backend() -> TrtllmAttention: + """The drafter layer's attention as ``M3Eagle3Draft`` builds it: fused NeoX RoPE, an FP8 KV cache (the target + overrides the drafter's KV dtype to FP8), no other quantization.""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + backend = create_attention( + "TRTLLM", + DRAFT, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=RopeParams(dim=HEAD_DIM, theta=DRAFT_ROPE_THETA, max_positions=MAX_SEQ_LEN), + ), + quant_config=quant, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + backend.update_quant_config(quant) + assert type(backend) is TrtllmAttention, type(backend).__name__ + return backend + + +def make_msa_metadata(manager) -> MiniMaxM3MsaSparseAttentionMetadata: + metadata = MiniMaxM3MsaSparseAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + max_num_sequences=MAX_BATCH, + kv_cache_manager=manager, + mapping=mapping(), + sparse_metadata_params=sparse_attention_config().to_sparse_metadata_params(), + ) + # The target's MSA preparation choice that changes what prepare() stages (the identity tables of short steps). + metadata.msa_policy = MsaRuntimePolicy(short_sequence_identity=True) + return metadata + + +def make_draft_metadata(manager) -> TrtllmAttentionMetadata: + """Drafter metadata over the shared manager's FP8 draft view, the manager the runtime's draft swap installs.""" + view = manager.get_draft_kv_cache_view() + assert view is not None and view.dtype == DataType.FP8 + return TrtllmAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + kv_cache_manager=view, + mapping=mapping(), + ) + + +def allocate(manager, request_ids: Sequence[int], qlens: Sequence[int], cached: Sequence[int]): + manager.add_dummy_requests(list(request_ids), token_nums=[c + n for c, n in zip(cached, qlens)]) + + +def prepare(metadata, request_ids, qlens, cached, num_contexts: int) -> None: + metadata.seq_lens = torch.tensor(list(qlens), dtype=torch.int) + metadata.num_contexts = num_contexts + metadata.request_ids = list(request_ids) + metadata.prompt_lens = [ + c + n if row < num_contexts else c for row, (c, n) in enumerate(zip(cached, qlens)) + ] + metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=list(cached)) + metadata.prepare() + + +def flat_slots(manager, layer: int, request_id: int, positions: torch.Tensor) -> torch.Tensor: + """The cache slots (page slot id * PAGE + offset) of a request's positions in ``layer``.""" + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0], + dtype=torch.int64, + device="cuda", + ) + positions = positions.to(device="cuda", dtype=torch.int64) + return (pages[positions // PAGE] * PAGE + positions % PAGE).to(torch.int32) + + +def write_fp8_rows(manager, layer, request_id, positions, k, v) -> None: + """K / V rows [n, 128] (E4M3) into an FP8 layer's cache at ``positions``.""" + view = manager.get_buffers(layer, "HND") + assert view.dtype == torch.float8_e4m3fn + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches(view[:, 0], view[:, 1], None, slots, k, v, None) + + +def write_nvfp4_rows(manager, metadata, layer, request_id, positions, k, v, idx_k, quant) -> None: + """BF16 K / V rows quantized into the NVFP4 layer's cache, and index-K rows into its FP8 index cache.""" + data = manager.get_buffers(layer, "HND") + scales = manager.get_block_scale_buffers(layer, "HND") + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches_nvfp4( + data[:, 0], + data[:, 1], + scales[:, 0], + scales[:, 1], + metadata.msa_idx_k_cache(layer), + slots, + k, + v, + idx_k, + quant, + ) + + +def index_directions(gen: torch.Generator) -> tuple[torch.Tensor, torch.Tensor]: + """Two orthogonal ±1 index directions (the halves of the 128 dims), E4M3-exact.""" + signs = torch.randint(0, 2, (HEAD_DIM,), generator=gen, device="cuda").float() * 2 - 1 + u0 = torch.zeros(HEAD_DIM, device="cuda") + u1 = torch.zeros(HEAD_DIM, device="cuda") + u0[: HEAD_DIM // 2] = signs[: HEAD_DIM // 2] + u1[HEAD_DIM // 2 :] = signs[HEAD_DIM // 2 :] + return u0, u1 + + +def planted_index_keys(length: int, u0, u1, gen: torch.Generator) -> torch.Tensor: + """Index-K rows [length, 128] (bf16, E4M3-exact) whose per-block maximum score is unambiguous. + + Every block has one peak row at a random position, ``a_b u0 + c_b u1`` with ``a`` and ``c`` two independent + permutations of ``LEVELS``; its other rows are ``LOW (u0 + u1)``. A query along ``u0`` then ranks the blocks by + ``a``, one along ``u1`` by ``c``, every score 1/8 or more apart from the next.""" + blocks = (length + PAGE - 1) // PAGE + assert blocks <= len(LEVELS) + levels = torch.tensor(LEVELS, device="cuda") + a = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + c = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + rows = (LOW * (u0 + u1)).expand(length, HEAD_DIM).clone() + for b in range(blocks): + span = min(PAGE, length - b * PAGE) + peak = b * PAGE + int(torch.randint(0, span, (1,), generator=gen, device="cuda")) + rows[peak] = a[b] * u0 + c[b] * u1 + return rows.to(torch.bfloat16) + + +def nvfp4_quant(*values: torch.Tensor) -> torch.Tensor: + """The (Q, K, V) quantization scales ``fused_write_layer_caches_nvfp4`` takes: each role's amax at the top of the + E2M1 x E4M3 range.""" + quant = torch.ones(3, dtype=torch.float32, device="cuda") + for role, x in enumerate(values, start=1): + quant[role] = (448.0 * 6.0) / x.abs().max().float() + return quant diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/_m3_nvfp4_cache.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/_m3_nvfp4_cache.py new file mode 100644 index 000000000000..bfa03c09da03 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/_m3_nvfp4_cache.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A MiniMax-M3 TP8 sparse-layer decode step over a packed NVFP4 cache, for the attention entries' tests. + +The cache is written by the in-tree scatter (``fused_write_layer_caches_nvfp4``, +which quantizes with ``torch.ops.trtllm.fp4_quantize``), into one +``[pages, 2 (K, V), 1, 128, .]`` pool per role exactly as ``msa_paged_kv`` and +``get_block_scale_buffers`` hand the K and V views over. Pages are shuffled so +a kernel that ignored the block table would read the wrong ones. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches_nvfp4, +) + +#: Query heads per rank at TP8 (on one KV head). +NUM_HEADS = 8 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 + + +def token_kv_len(seq_len: int, decode_query_len: int, offset: int) -> int: + """The KV extent query token ``offset`` of a request attends, as every decode kernel derives it.""" + return max(seq_len - decode_query_len + offset + 1, 0) + + +def build_step( + seq_lens: Sequence[int], decode_query_len: int, *, identity: bool, seed: int +) -> SimpleNamespace: + """One decode step: E4M3 q [T, 8, 128], the packed NVFP4 K/V views and their scales, a [1, T, 16] + head-major top-k table (the identity table, or a random ascending selection with a -1 tail), the block + table and per-request lengths.""" + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + total_q = batch * decode_query_len + max_blocks = max(1, max((s + PAGE - 1) // PAGE for s in seq_lens)) + num_pages = batch * max_blocks + num_slots = num_pages * PAGE + block_table = ( + torch.randperm(num_pages, device="cuda", generator=g) + .to(torch.int32) + .reshape(batch, max_blocks) + ) + data_pool = torch.zeros(num_pages, 2, 1, PAGE, HEAD_DIM // 2, dtype=torch.uint8, device="cuda") + scale_pool = torch.zeros( + num_pages, 2, 1, PAGE, HEAD_DIM // 16, dtype=torch.uint8, device="cuda" + ) + quant = torch.ones(3, dtype=torch.float32, device="cuda") + values = [] + for role in (0, 1): + x = torch.randn(num_slots, HEAD_DIM, device="cuda", generator=g, dtype=torch.float32).to( + torch.bfloat16 + ) + quant[role + 1] = (448.0 * 6.0) / x.abs().max().float().item() + values.append(x) + assert fused_write_layer_caches_nvfp4( + data_pool[:, 0], + data_pool[:, 1], + scale_pool[:, 0], + scale_pool[:, 1], + None, + torch.arange(num_slots, device="cuda", dtype=torch.int32), + values[0], + values[1], + None, + quant, + ) + q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, device="cuda", generator=g).to( + torch.float8_e4m3fn + ) + topk = torch.full((1, total_q, TOPK), -1, dtype=torch.int32, device="cuda") + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + nblocks = (token_kv_len(seq_lens[req], decode_query_len, offset) + PAGE - 1) // PAGE + if identity: + assert nblocks <= TOPK, "an identity step has every request within 16 blocks" + chosen = torch.arange(nblocks, device="cuda") + else: + chosen = torch.randperm(nblocks, device="cuda", generator=g)[:TOPK].sort().values + topk[0, tok, : chosen.numel()] = chosen.to(torch.int32) + return SimpleNamespace( + q=q, + k_paged=data_pool[:, 0], + v_paged=data_pool[:, 1], + k_block_scale=scale_pool[:, 0], + v_block_scale=scale_pool[:, 1], + k_global_scale=(1.0 / quant[1:2]).contiguous(), + v_global_scale=(1.0 / quant[2:3]).contiguous(), + topk=topk, + block_table=block_table, + seq_lens=torch.tensor(list(seq_lens), dtype=torch.int32, device="cuda"), + decode_query_len=decode_query_len, + sm_scale=HEAD_DIM**-0.5, + source=values, + quant=quant, + num_slots=num_slots, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_dense_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_dense_decode.py new file mode 100644 index 000000000000..13f73c32d60d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_dense_decode.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dense decode against the fp32 reference and between valid sequence-length bounds.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_dense_decode import ( # noqa: E501 + HEAD_DIM, + HEADS, + KV_HEADS, + MODEL_MAX, + SHORT_BOUND, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_dense_decode import ( # noqa: E501 + m3_dense_decode as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_dense_decode requires a CUDA device" + +PAGE = 32 +SUBPAGES = 2 +SM_SCALE = HEAD_DIM**-0.5 +# The FP8-KV band of the in-tree dense decode test (test_minimax_m3_dense_decode.test_matches_reference). +FP8_TOL = 6e-2 + + +class _Pool: + """The one method the kernel asks of a KV cache manager.""" + + def __init__(self, pool: torch.Tensor): + self.pool = pool + + def get_kv_subpage_pool(self, layer_idx: int, kv_layout: str = "HND"): + assert kv_layout == "HND" + return self.pool, SUBPAGES + + +def _inputs(kv_lens, qlen: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(kv_lens) + pages = max((n + PAGE - 1) // PAGE for n in kv_lens) + slots = batch * pages + 3 + pool = torch.randn( + (slots * SUBPAGES, KV_HEADS, PAGE, HEAD_DIM), device="cuda", generator=g + ).clamp(-4, 4) + # Slots handed out non-contiguously, as the block manager does. + perm = torch.randperm(slots, generator=g, device="cuda")[: batch * pages] + block_table = perm.reshape(batch, pages).to(torch.int32) + # The last third of every request's values is shifted by +1: a third of each row's softmax mass sits there, far + # outside the band, so a decode that missed the tail of a context fails. + for b, n in enumerate(kv_lens): + start = 2 * n // 3 + first_full = (start + PAGE - 1) // PAGE + last = (n + PAGE - 1) // PAGE + if first_full < last: + v_idx = block_table[b, first_full:last].long() * SUBPAGES + 1 + pool[v_idx] += 1.0 + if start % PAGE: + v0 = int(block_table[b, start // PAGE]) * SUBPAGES + 1 + pool[v0, :, start % PAGE :] += 1.0 + q = torch.randn((batch * qlen, HEADS, HEAD_DIM), device="cuda", generator=g).to(torch.bfloat16) + return ( + q, + _Pool(pool.to(torch.float8_e4m3fn)), + block_table, + torch.tensor(kv_lens, device="cuda", dtype=torch.int32), + ) + + +def _args(q, manager, block_table, seq_lens, qlen, max_seq_len): + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + return ( + q, + manager, + 0, + block_table, + seq_lens, + SM_SCALE, + out, + qlen, + max_seq_len, + seq_lens.shape[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 — pytest resolves the fixture by this name + spec = cell.spec + q, manager, block_table, seq_lens = _inputs( + spec["kv_lens"], spec["qlen"], zlib.crc32(cell.why.encode()) % 2**31 + ) + args = _args(q, manager, block_table, seq_lens, spec["qlen"], spec["max_seq_len"]) + certifying(op) + ref = op.reference(*args) + out = op(*args).clone() + torch.testing.assert_close(out.float(), ref, rtol=FP8_TOL, atol=FP8_TOL) + if spec["tight_bound"]: + model_max_output = op( + *_args(q, manager, block_table, seq_lens, spec["qlen"], MODEL_MAX) + ).clone() + # The split reorders only the merge of the CTAs' fp32 partial sums before the one BF16 rounding: every element + # is within 4 of its own BF16 spacing (2^(e - 8) for |x| = m * 2^e, m in [0.5, 1)). + ref_bf16 = model_max_output.float() + _, exp = torch.frexp(ref_bf16.abs()) + spacing = torch.where( + ref_bf16 == 0, + torch.full_like(ref_bf16, 2.0**-133), + torch.ldexp(torch.ones_like(ref_bf16), exp - 8), + ) + assert bool(((out.float() - ref_bf16).abs() <= 4 * spacing).all()) + assert_within_ulp(out, model_max_output, element_ulp=4.0, rms_ulp=1.0) + + +def test_guard_refuses_more_requests_than_the_counters() -> None: + q, manager, block_table, seq_lens = _inputs((3000, 3000), 1, 3) + args = list(_args(q, manager, block_table, seq_lens, 1, SHORT_BOUND)) + args[9] = 1 + with pytest.raises(AssertionError, match="counters cover"): + with validating(op): + op(*args) + + +def test_guard_refuses_an_understated_sequence_length_bound() -> None: + q, manager, block_table, seq_lens = _inputs((20000,), 1, 7) + args = _args(q, manager, block_table, seq_lens, 1, SHORT_BOUND) + with pytest.raises(AssertionError, match="sequence length exceeds its bound"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_draft_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_draft_attention.py new file mode 100644 index 000000000000..bfa051f8c35b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_draft_attention.py @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_attention entry: the Eagle3 drafter layer's attention against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp8``): the target's hybrid NVFP4 +``MiniMaxM3KVCacheManagerV2`` with its one-model Eagle3 draft layer, the requests allocated in it, drafter +``TrtllmAttentionMetadata`` over the manager's FP8 draft view (the manager the runtime's draft swap installs) filled +by its ``prepare()``, and the drafter's ``TrtllmAttention`` built as ``M3Eagle3Draft`` builds it (fused NeoX RoPE at +the drafter checkpoint's theta, FP8 KV). Earlier rows go into the draft layer's cache through the in-tree scatter at +the slots the manager assigned; the call rotates and appends the step's own. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_draft_attention import ( # noqa: E501 + m3_draft_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp8 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_attention requires a CUDA device" + +# Band, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS. +# The reference reads the same FP8 cache rows the call attends, so what is left is the FP8 FMHA's own arithmetic: +# it takes the rotated Q in E4M3 (rounding the reference's Q to E4M3 takes the RMS gap from ~9 to ~7 ulps) and rounds +# the softmax probabilities before the PV product. Measured on GB300 over these cells (max element / RMS): 14.8 / 8.7 +# (4 x 4 rows), 13.3 / 7.4 (8 x 1), 19.1 / 7.7 (512 context rows). The band leaves about 1.5x; a reference shifted +# by one row is far outside it. +_ELEMENT_ULP = 30.0 +_RMS_ULP = 12.0 + +# Cached tokens per request of each cell; the new rows per request come from the cell. +_CACHED = { + (4, 4): (200, 129, 1000, 64), + (8, 1): (5, 127, 128, 300, 1000, 2047, 64, 900), + (1, None): (0,), +} +#: Rows past the live ones in the fused QKV, as a padded step hands them over. +_PAD = 3 + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _rope(x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """NeoX RoPE over the whole head at the drafter's theta, float64.""" + half = rt.HEAD_DIM // 2 + inv_freq = rt.DRAFT_ROPE_THETA ** ( + -torch.arange(0, rt.HEAD_DIM, 2, dtype=torch.float64, device=x.device) / rt.HEAD_DIM + ) + angles = positions.to(torch.float64)[:, None] * inv_freq[None, :] + x = x.to(torch.float64) + x1, x2 = x[..., :half], x[..., half:] + return torch.cat( + [x1 * angles.cos() - x2 * angles.sin(), x2 * angles.cos() + x1 * angles.sin()], -1 + ) + + +@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 this name + spec = cell.spec + context = spec.get("context") + cached = _CACHED[spec["reqs"], spec.get("qlen")] + qlens = (context,) if context else (spec["qlen"],) * spec["reqs"] + assert len(cached) == spec["reqs"] + contexts = 1 if context else 0 + rids = list(range(spec["reqs"])) + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=True) + try: + backend = rt.make_draft_backend() + metadata = rt.make_draft_metadata(manager) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.HEAD_DIM, generator=gen, device="cuda") + hist = hist.to(torch.float8_e4m3fn) + rt.write_fp8_rows(manager, rt.DRAFT, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + assert int(metadata.num_tokens) == total + width = (rt.HEADS + 2 * rt.KV_HEADS) * rt.HEAD_DIM + qkv = torch.randn(total + _PAD, width, generator=gen, device="cuda").to(torch.bfloat16) + + certifying(op) + out = op(backend, qkv, metadata) + ref = op.reference(backend, qkv, metadata) + assert out.shape == (total, rt.HEADS * rt.HEAD_DIM) and out.dtype == torch.bfloat16 + assert torch.isfinite(out.float()).all() + element, rms = _ulp_stats(out, ref) + print(f"\n[{cell.why}] max element {element:.2f} ulp, rms {rms:.3f} ulp") + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + # The call's append: each live row's rotated K and its V at the row's position, E4M3 (unit KV scale), to + # within one E4M3 step of the float64 rotation (the kernel rotates in FP32). + pool = manager.get_buffers(rt.DRAFT, "HND") + q_size = rt.HEADS * rt.HEAD_DIM + first = 0 + for rid, c, n in zip(rids, cached, qlens): + positions = torch.arange(c, c + n, device="cuda") + slots = rt.flat_slots(manager, rt.DRAFT, rid, positions).long() + rows = qkv[first : first + n] + want_k = _rope(rows[:, q_size : q_size + rt.HEAD_DIM].float(), positions) + want_v = rows[:, q_size + rt.HEAD_DIM :].to(torch.float8_e4m3fn) + got_k = pool[slots // rt.PAGE, 0, 0, slots % rt.PAGE].to(torch.float64) + got_v = pool[slots // rt.PAGE, 1, 0, slots % rt.PAGE] + step = want_k.abs().clamp_min(2.0**-6) * 2.0**-3 + assert bool(((got_k - want_k).abs() <= step).all()), "K is appended rotated, in E4M3" + assert torch.equal(got_v.view(torch.uint8), want_v.view(torch.uint8)) + first += n + + # The band discriminates: a reference whose rows are shifted by one token is outside it. + shifted = op.reference(backend, qkv.roll(1, 0), metadata) + with pytest.raises(AssertionError): + assert_within_ulp(out, shifted, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..241f07fd7136 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp8_indexer_qk_norm_rope entry of the MiniMax-M3 tp8 target: index-Q out and index-K in the +cache against the in-place fused norm + RoPE followed by the E4M3 cast, byte for byte up to the 0.1% budget the +in-tree test of the kernel allows.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_fp8_indexer_qk_norm_rope import ( # noqa: E501 + m3_fp8_indexer_qk_norm_rope as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp8_indexer_qk_norm_rope requires a CUDA device" + +DIM, ROTARY, EPS, THETA, PAGE = 128, 64, 1e-6, 5e6, 128 + + +def _bytes_close(a: torch.Tensor, b: torch.Tensor) -> None: + assert a.shape == b.shape and a.dtype == b.dtype == torch.float8_e4m3fn + differ = int((a.view(torch.uint8) != b.view(torch.uint8)).sum()) + assert differ <= (a.numel() - 1) // 1000, f"{differ} of {a.numel()} E4M3 bytes differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + m, heads = cell.spec["m"], cell.spec["heads"] + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + qk = torch.randn(m, (heads + 1) * DIM, generator=g, device="cuda").to(torch.bfloat16) + qw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (m,), generator=g, device="cuda", dtype=torch.int32) + pages = (m + PAGE - 1) // PAGE + 1 + cache = torch.zeros(pages, 1, PAGE, DIM, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(pages * PAGE, generator=g, device="cuda")[:m].to(torch.int32) + before = qk.clone() + args = (qk, cache, slots, heads, DIM, ROTARY, EPS, qw, kw, THETA, pos) + with validating(op): + q = op(*args) + assert torch.equal(qk, before), "the BF16 input is left untouched" + q_ref, k_ref = op.reference(*args) + _bytes_close(q.reshape(m, heads, DIM), q_ref) + flat = cache.view(-1, DIM) + _bytes_close(flat[slots.long()], k_ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fused_score_select.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fused_score_select.py new file mode 100644 index 000000000000..ea21f436a343 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fused_score_select.py @@ -0,0 +1,230 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_fused_score_select entry: MSA index scoring and top-16 block selection of a decode +span of at most two requests in one launch (``torch.ops.trtllm.minimax_m3_tp8_fused_score_select``). + +Every cell is held bit for bit to the entry's ``reference`` (``m3_select_long.reference`` over the per-KV-head maxima +of ``m3_index_score.reference``) and to the pipeline it replaces in the target: this target's ``m3_index_score`` (the +prefetching kernel at ``SPLIT_K``), the per-KV-head amax and ``m3_select_long``, on the same inputs. The selection is a +table of block ids, so the claim is exact. +""" + +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.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops import ( # noqa: E501 + m3_fused_score_select as fss, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_fused_score_select import ( # noqa: E501 + MAX_REQUESTS, + scratch_shape, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_fused_score_select import ( # noqa: E501 + m3_fused_score_select as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_index_score import ( # noqa: E501 + SPLIT_K, + m3_index_score, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, + m3_select_long, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fused_score_select requires a CUDA device" + +PAGE = 128 +HEADS = 1 # index heads per rank at TP8 +HEAD_DIM = 128 + + +def _n_valid(seq_lens, qlen: int) -> torch.Tensor: + """Each query row's causal valid-block count, as the metadata stages it (``per_token_valid_blocks``).""" + vals = [ + (max(seq - qlen + offset + 1, 0) + PAGE - 1) // PAGE + for seq in seq_lens + for offset in range(qlen) + ] + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _inputs(spec, seed: int, dtype=torch.float8_e4m3fn): + """Random FP8 index Q [T, 1, 128] and paged index-K [pages, 128, 128] behind a shuffled block table (as the index + scorer's test builds them), the request lengths and the rows' valid-block counts.""" + seq_lens, qlen = spec["seq_lens"], spec["qlen"] + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + max_blocks = (max(seq_lens) + PAGE - 1) // PAGE + pages = batch * max_blocks + block_table = ( + torch.randperm(pages, device="cuda", generator=g).to(torch.int32).view(batch, max_blocks) + ) + idx_q = torch.randn(batch * qlen, HEADS, HEAD_DIM, device="cuda", generator=g).to(dtype) + index_k = torch.randn(pages, PAGE, HEAD_DIM, device="cuda", generator=g).to(dtype) + seq = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + return idx_q, index_k, block_table, seq, _n_valid(seq_lens, qlen) + + +def _out(total_q: int, num_kv_heads: int) -> torch.Tensor: + """The [total_q, num_kv_heads, 16] view of a head-major backing, as the target hands it over.""" + return torch.full((num_kv_heads, total_q, TOPK), -7, dtype=torch.int32, device="cuda").permute( + 1, 0, 2 + ) + + +def _args(spec, inputs, out, scratch, counters): + idx_q, index_k, block_table, seq, n_valid = inputs + return ( + idx_q, + index_k, + block_table, + seq, + n_valid, + out, + scratch, + counters, + spec["qlen"], + spec["num_kv_heads"], + spec["init_blocks"], + spec["local_blocks"], + spec["max_tiles"], + False, + ) + + +def _workspace(max_tiles: int): + """Scratch filled with garbage (the op needs no initialization) and zeroed counters.""" + scratch = torch.full(scratch_shape(max_tiles), float("nan"), device="cuda") + counters = torch.zeros(MAX_REQUESTS, dtype=torch.int32, device="cuda") + return scratch, counters + + +def _pipeline(spec, inputs) -> torch.Tensor: + """What the target runs on a wider span: this target's index scorer (prefetching kernel), the per-KV-head amax over + its index heads, and the long-row selector.""" + idx_q, index_k, block_table, seq, n_valid = inputs + total_q = idx_q.shape[0] + max_tiles, kv = spec["max_tiles"], spec["num_kv_heads"] + backing = torch.full( + (HEADS, max_tiles, total_q), -float("inf"), dtype=torch.float32, device="cuda" + ) + m3_index_score( + idx_q, index_k, block_table, seq, backing.transpose(1, 2), spec["qlen"], SPLIT_K, False + ) + per_kv = backing.view(kv, HEADS // kv, max_tiles, total_q).amax(dim=1) + table = _out(total_q, kv) + return m3_select_long( + per_kv, n_valid, table, spec["init_blocks"], spec["local_blocks"], False, False + ) + + +@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 this name + spec = cell.spec + inputs = _inputs(spec, zlib.crc32(cell.why.encode()) % 2**31) + assert len(spec["seq_lens"]) == spec["batch"] + total_q = inputs[0].shape[0] + out = _out(total_q, spec["num_kv_heads"]) + scratch, counters = _workspace(spec["max_tiles"]) + args = _args(spec, inputs, out, scratch, counters) + certifying(op) + op(*args) + ref = op.reference(*args) + torch.cuda.synchronize() + + assert out.dtype is torch.int32 and out.shape == (total_q, spec["num_kv_heads"], TOPK) + assert torch.equal(out, ref), "the selected blocks differ from the reference" + assert not bool(counters.any()), "the arrival counters must be back at zero after the call" + assert torch.equal(out, _pipeline(spec, inputs)), ( + "the selected blocks differ from m3_index_score + per-KV amax + m3_select_long" + ) + + # A second call on the same scratch and counters (as the next sparse layer makes) selects the same. + again = _out(total_q, spec["num_kv_heads"]) + op(*_args(spec, inputs, again, scratch, counters)) + torch.cuda.synchronize() + assert torch.equal(again, out) and not bool(counters.any()) + + +@pytest.mark.parametrize("cell", op.CELLS[:2], ids=[c.why[:44] for c in op.CELLS[:2]]) +def test_the_target_wrapper_is_the_entry(cell) -> None: + """``ops/m3_fused_score_select``: the geometry check admits the cell, the scratch and counters are created + once per (device, engine block count), and the wrapper's table is the entry's.""" + spec = cell.spec + inputs = _inputs(spec, 3) + idx_q, index_k, block_table, seq, n_valid = inputs + max_tiles = spec["max_tiles"] + assert fss.supported(idx_q, index_k, spec["qlen"], spec["num_kv_heads"], max_tiles) is None + bufs = fss.buffers(idx_q.device, max_tiles) + assert bufs is not None and fss.buffers(idx_q.device, max_tiles) is bufs + table = fss.fused_score_select( + idx_q, + index_k, + block_table, + seq, + n_valid, + decode_query_len=spec["qlen"], + num_kv_heads=spec["num_kv_heads"], + init_blocks=spec["init_blocks"], + local_blocks=spec["local_blocks"], + max_tiles=max_tiles, + scratch=bufs[0], + counters=bufs[1], + ) + torch.cuda.synchronize() + ref = op.reference(*_args(spec, inputs, _out(idx_q.shape[0], 1), *_workspace(max_tiles))) + assert torch.equal(table, ref) + assert not bool(bufs[1].any()) + + +def test_the_geometry_check_refuses_what_the_kernel_does_not_take() -> None: + spec = op.CELLS[0].spec + idx_q, index_k, *_ = _inputs(spec, 5) + assert "not FP8" in fss.supported(idx_q.bfloat16(), index_k, 4, 1, 8192) + assert "geometry" in fss.supported(idx_q, index_k, 5, 1, 8192) + wide_q = idx_q.expand(-1, 4, -1).contiguous() + assert "geometry" in fss.supported(wide_q, index_k, 4, 1, 8192) + assert "stage ring" in fss.supported(idx_q, index_k, 4, 1, 16384) + + +def test_the_dispatch_range_is_the_kernel_limit() -> None: + assert dispatch.FUSED_SCORE_SELECT_MAX_REQUESTS == MAX_REQUESTS + assert {c.spec["batch"] for c in op.CELLS} == set(dispatch.FUSED_SCORE_SELECT.keys) + + +def test_guard_refuses_nonzero_counters() -> None: + spec = op.CELLS[0].spec + inputs = _inputs(spec, 7) + scratch, counters = _workspace(spec["max_tiles"]) + counters[0] = 1 + with pytest.raises(AssertionError, match="counters must be zero"): + with validating(op): + op(*_args(spec, inputs, _out(inputs[0].shape[0], 1), scratch, counters)) + + +def test_guard_refuses_a_request_past_its_block_table() -> None: + spec = op.CELLS[4].spec + idx_q, index_k, block_table, seq, n_valid = _inputs(spec, 9) + seq = seq.clone() + seq[1] = block_table.shape[1] * PAGE + 1 + scratch, counters = _workspace(spec["max_tiles"]) + with pytest.raises(AssertionError, match="past its block-table row"): + with validating(op): + op( + *_args( + spec, + (idx_q, index_k, block_table, seq, n_valid), + _out(idx_q.shape[0], 1), + scratch, + counters, + ) + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_index_score.py new file mode 100644 index 000000000000..090e82cb7020 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_index_score.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_index_score entry: the MSA index scorer with its pre-wait prefetch. + +Every cell is held to the fp32 reference (a few fp32 ulps: the kernel sums the 128 exact products in its k-split order) +and to the in-tree scorer (``cute_dsl_minimax_m3_index_decode_score``) bit for bit, which the pre-wait prefetch must not +change. The tp8 target runs the prefetching kernel only (spans of more than two requests; ``early`` is never +selected), so that is the schedule driven here. +""" + +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.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_index_score import ( # noqa: E501 + PAGE, + SPLIT_K, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_index_score import ( # noqa: E501 + m3_index_score as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_index_score requires a CUDA device" + +HEADS = 1 # index heads per rank at TP8 +HEAD_DIM = 128 + + +def _inputs(seq_lens, qlen: int, seed: int, dtype=torch.float8_e4m3fn): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + max_blocks = (max(seq_lens) + PAGE - 1) // PAGE + stride = ((max_blocks + 15) // 16) * 16 + pages = batch * max_blocks + block_table = ( + torch.randperm(pages, device="cuda", generator=g).to(torch.int32).view(batch, max_blocks) + ) + idx_q = torch.randn(batch * qlen, HEADS, HEAD_DIM, device="cuda", generator=g).to(dtype) + index_k = torch.randn(pages, PAGE, HEAD_DIM, device="cuda", generator=g).to(dtype) + seq = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + return idx_q, index_k, block_table, seq, stride + + +def _score(stride: int, total_q: int): + backing = torch.full( + (HEADS, stride, total_q), -float("inf"), device="cuda", dtype=torch.float32 + ) + return backing.transpose(1, 2) + + +@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 this name + spec = cell.spec + qlen = spec["qlen"] + idx_q, index_k, block_table, seq, stride = _inputs( + spec["seq_lens"], qlen, zlib.crc32(cell.why.encode()) + ) + score = _score(stride, idx_q.shape[0]) + args = (idx_q, index_k, block_table, seq, score, qlen, SPLIT_K, spec["early"]) + certifying(op) + ref = op.reference(*args) + op(*args) + stock = _score(stride, idx_q.shape[0]) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen, SPLIT_K + ) + torch.cuda.synchronize() + torch.testing.assert_close(score, ref, rtol=1e-5, atol=1e-4) + assert torch.equal(score, stock), "the scores differ from the in-tree scorer" + + +def test_bf16_index_is_the_in_tree_scorer() -> None: + idx_q, index_k, block_table, seq, stride = _inputs( + (5000, 300, 2049), 4, 11, dtype=torch.bfloat16 + ) + score, stock = _score(stride, 12), _score(stride, 12) + op(idx_q, index_k, block_table, seq, score, 4, SPLIT_K, False) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, 4, SPLIT_K + ) + assert torch.equal(score, stock) + + +def test_guard_refuses_a_request_past_its_block_table() -> None: + idx_q, index_k, block_table, seq, stride = _inputs((1000,), 1, 5) + seq = seq + 4 * PAGE + with pytest.raises(AssertionError, match="past its block-table row"): + with validating(op): + op(idx_q, index_k, block_table, seq, _score(stride + 16, 1), 1, SPLIT_K, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_msa_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_msa_attention.py new file mode 100644 index 000000000000..060daf650f9e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_msa_attention.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_msa_attention entry: one layer's MSA backend step against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp8``): a hybrid NVFP4 ``MiniMaxM3KVCacheManagerV2`` +with the requests allocated in it, ``MiniMaxM3MsaSparseAttentionMetadata`` filled by its ``prepare()``, and the layer's +``MiniMaxM3MsaSparseAttention`` built as the target builds it. Earlier rows (and, on a sparse layer, the step's own +rows, which the target's producer writes before the call) go into the cache through the in-tree scatters, at the +slots the manager assigned. The sparse layers' index-K rows are planted so every block's score is 1/8 or more apart +from every other's: the top-k selection the reference recomputes is then unambiguous. + +The cells are the target's calls (``_attention_core`` in modeling.py): a dense layer with live bf16 K / V in a +prefill chunk and live FP8 K / V in a decode step (the dense layers sit in the hybrid cache's FP8 half), a sparse +layer whose producer already wrote the caches (``k = v = None``) in a prefill chunk beyond 16 blocks and in a mixed +step, both with the NVFP4 scales the target passes. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_msa_attention import ( # noqa: E501 + m3_msa_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp8 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_msa_attention requires a CUDA device" + +# Bands, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS, +# per kernel the step reaches. The reference reads the same quantized cache rows (and casts Q to E4M3 where the FP8 +# FMHA does), so what is left is each FMHA's own arithmetic, measured on GB300 over these cells (max element / RMS): +# * fmha_sm100 over the FP8 pool (dense context rows): 10.7 / 6.0. It runs E4M3 x E4M3 and its softmax +# probabilities are rounded before the PV product. +# * trtllm-gen over the FP8 pool (dense generation rows): 7.7 / 6.1. It rounds the probabilities to E4M3 (scaled by +# 448 against the running maximum): emulating that brings a single-tile request to 0.6 RMS. +# * the NVFP4 CSR prefill (sparse context rows): 33.6 / 11.8 at 4096 rows, 24.3 / 9.1 at 512. It stages the +# reconstructed K / V in E4M3 (the in-tree A/B test, test_minimax_m3_sparse_attn_decode, notes it and holds it +# to 3e-2). +# * the Triton NVFP4 sparse decode (sparse generation rows): 1.0 / 0.19 -- the band of its own entry +# (m3_sparse_decode): exact operands, FP16 probabilities, BF16 rounding. +# Each band leaves about 1.5x the measured maximum; a shifted or wrongly selected reference is far outside them. +_FP8_CONTEXT = dict(element_ulp=16.0, rms_ulp=9.0) +_FP8_DECODE = dict(element_ulp=12.0, rms_ulp=9.0) +_NVFP4_CONTEXT = dict(element_ulp=48.0, rms_ulp=16.0) +_NVFP4_DECODE = dict(element_ulp=2.0, rms_ulp=0.5) +_BANDS = { + ("dense", "context"): (_FP8_CONTEXT, None), + ("dense", "generation"): (None, _FP8_DECODE), + ("sparse", "context"): (_NVFP4_CONTEXT, None), + ("sparse", "mixed"): (_NVFP4_CONTEXT, _NVFP4_DECODE), +} + + +def _assert_in_bands(output, ref, bands, context_rows: int) -> None: + """Gate the context rows and the generation rows each on their own kernel's band.""" + context_band, decode_band = bands + if context_rows: + assert_within_ulp(output[:context_rows], ref[:context_rows], **context_band) + if context_rows < output.shape[0]: + assert_within_ulp(output[context_rows:], ref[context_rows:], **decode_band) + + +# (request rows, cached tokens per request, new tokens per request, context rows) per cell. +_SHAPES = { + ("dense", "context"): ((384,), (1024,), 1), + ("dense", "generation"): ((127, 128, 1000, 2500), (1, 1, 1, 1), 0), + ("sparse", "context"): ((0,), (4096,), 1), + ("sparse", "mixed"): ((0, 2999, 4500), (512, 1, 1), 1), +} + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + """The two quantities ``assert_within_ulp`` gates, in BF16 ulps.""" + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _shift(x: torch.Tensor) -> torch.Tensor: + """``x``'s rows moved down by one (byte-wise, so E4M3 rows move too).""" + return x.view(torch.uint8).roll(1, 0).view(x.dtype) if x.element_size() == 1 else x.roll(1, 0) + + +def _shape(cell): + spec = cell.spec + cached, qlens, contexts = _SHAPES[spec["layer"], spec["phase"]] + if spec["phase"] == "context": + assert sum(qlens) == spec["tokens"] + elif spec["phase"] == "generation": + assert len(qlens) == spec["reqs"] + else: + assert sum(qlens[:contexts]) == spec["tokens"] and len(qlens) - contexts == spec["reqs"] + return cached, qlens, contexts + + +def _run_dense(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + rt.write_fp8_rows(manager, rt.DENSE, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + # A prefill chunk arrives in BF16 (the hybrid cache's dense layers take FP8 Q / K / V only on decode steps). + dtype = torch.bfloat16 if contexts else torch.float8_e4m3fn + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + k = torch.randn(total, rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + v = torch.randn(total, rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + # The NVFP4 target passes its layer scales on every layer; an FP8 layer's write and FMHA take none. + inv = torch.tensor([1.0, 0.5, 2.0], device="cuda") + fwd = 1.0 / inv + return (q, k, v, None, inv, fwd), dict(rids=rids, cached=cached, qlens=qlens, contexts=contexts) + + +def _run_sparse(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + u0, u1 = rt.index_directions(gen) + lengths = [c + n for c, n in zip(cached, qlens)] + keys = [torch.randn(n, rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths] + values = [torch.randn(n, rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths] + quant = rt.nvfp4_quant(torch.cat(keys), torch.cat(values)) + for rid, n, kk, vv in zip(rids, lengths, keys, values): + # Every row, the step's own included: the producer writes a sparse layer's caches before the call. + idx_k = rt.planted_index_keys(n, u0, u1, gen) + rt.write_nvfp4_rows( + manager, + metadata, + rt.SPARSE, + rid, + torch.arange(n), + kk.to(torch.bfloat16), + vv.to(torch.bfloat16), + idx_k, + quant, + ) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + # Query tokens alternate between the two planted directions, so neighbouring tokens select different blocks. + idx_q = torch.stack([u0, u1])[torch.arange(total, device="cuda") % 2].to(torch.float8_e4m3fn) + return (q, None, None, idx_q, quant, 1.0 / quant), dict( + rids=rids, cached=cached, qlens=qlens, contexts=contexts + ) + + +@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 this name + spec = cell.spec + sparse = spec["layer"] == "sparse" + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=False) + try: + backend = rt.make_msa_backend(rt.SPARSE if sparse else rt.DENSE) + metadata = rt.make_msa_metadata(manager) + run = _run_sparse if sparse else _run_dense + (q, k, v, idx_q, inv, fwd), step = run(cell, gen, manager, backend, metadata) + output = torch.full( + (q.shape[0], rt.HEADS * rt.HEAD_DIM), float("nan"), dtype=torch.bfloat16, device="cuda" + ) + certifying(op) + result = op(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert result is output + ref = op.reference(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert torch.isfinite(output.float()).all(), "every live row is written" + bands = _BANDS[spec["layer"], spec["phase"]] + context_rows = sum(step["qlens"][: step["contexts"]]) + for name, rows in ( + ("context", slice(0, context_rows)), + ("generation", slice(context_rows, None)), + ): + if output[rows].numel(): + element, rms = _ulp_stats(output[rows], ref[rows]) + print( + f"\n[{cell.why}] {name} rows: max element {element:.2f} ulp, rms {rms:.3f} ulp" + ) + _assert_in_bands(output, ref, bands, context_rows) + + if k is not None: + # The call's own cache write: the live K / V rows land, as a plain E4M3 cast, at their positions. + pool = manager.get_buffers(rt.DENSE, "HND") + first = 0 + for rid, c, n in zip(step["rids"], step["cached"], step["qlens"]): + slots = rt.flat_slots(manager, rt.DENSE, rid, torch.arange(c, c + n)).long() + for role, src in ((0, k), (1, v)): + got = pool[slots // rt.PAGE, role, 0, slots % rt.PAGE] + want = src[first : first + n].to(torch.float8_e4m3fn) + assert torch.equal(got.view(torch.uint8), want.view(torch.uint8)) + first += n + + # The band discriminates: a reference whose query rows are shifted by one token is outside it. + shifted = op.reference(backend, _shift(q), k, v, idx_q, metadata, output, sparse, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, shifted, bands, context_rows) + if sparse: + # And the selection is load-bearing: attending every causal block instead is outside it too. + dense = op.reference(backend, q, k, v, idx_q, metadata, output, False, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, dense, bands, context_rows) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_producer.py new file mode 100644 index 000000000000..a948a445f648 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_producer.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_producer entry: the pre-wait-staged NVFP4 QKV + indexer producer. + +Every cell is held to the native producer (``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert``) bit for bit on q, +index-q and every K / V / scale / index-K cache byte, and to the fp32 torch reference: the FP8 outputs within one E4M3 +step, the decoded NVFP4 K / V within one E2M1 step of their block scale. +""" + +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.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_producer import ( # noqa: E501 + FP4_BLOCK, + HEADS_Q, + PAGE, + decode_nvfp4_token, + scale_index, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_producer import ( # noqa: E501 + m3_producer as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_producer requires a CUDA device" + +PAGES = 1024 +MAX_POSITIONS = 262144 +EPS = 1e-6 +KV_QUANT = (1.0, 0.37, 0.61) + + +def _caches(): + return ( + torch.zeros(PAGES, 2, 1, PAGE, 64, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 2, 1, PAGE, 8, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 1, PAGE, 128, dtype=torch.float8_e4m3fn, device="cuda"), + ) + + +def _inputs(tokens: int, scale: float, bad_slots: bool, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + rotary = torch.randn(MAX_POSITIONS, 2, 32, device="cuda", generator=g) + weights = [(torch.randn(128, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + packed = ( + torch.randn(tokens, (HEADS_Q + 4) * 128, device="cuda", generator=g) * scale + ).bfloat16() + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + if bad_slots: + slots[1] = -1 + slots[2] = PAGES * PAGE + 5 + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + quant = torch.tensor(KV_QUANT, device="cuda") + return packed, slots, quant, weights, rotary, positions + + +def _args(inputs, caches, early: bool = True): + packed, slots, quant, weights, rotary, positions = inputs + return (packed, *caches, slots, quant, EPS, *weights, rotary, positions, HEADS_Q, True, early) + + +def _native(inputs, caches): + packed, slots, quant, weights, rotary, positions = inputs + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, *caches, slots, quant, HEADS_Q, 1, 1, 128, 64, EPS, *weights, rotary, positions + ) + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +@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 this name + spec = cell.spec + early = spec.get("early", True) + inputs = _inputs( + spec["tokens"], spec["scale"], spec.get("bad_slots", False), zlib.crc32(cell.why.encode()) + ) + stock_caches, caches = _caches(), _caches() + stock = _native(inputs, stock_caches) + args = _args(inputs, caches, early) + certifying(op) + ref = op.reference(*args) + out = op(*args) + torch.cuda.synchronize() + names = ("q", "index_q", "kv_data", "kv_scale", "index_k") + for name, s, o in zip(names, (*stock, *stock_caches), (*out, *caches)): + assert torch.equal(s.view(torch.uint8), o.view(torch.uint8)), ( + f"{name} differs from the native producer" + ) + + _within_fp8(out[0], ref["q"], "q") + _within_fp8(out[1][:, 0], ref["index_q"], "index_q") + kv_data, kv_scale, index_k = caches + slots = inputs[1] + checked = 0 + for t in range(spec["tokens"]): + slot = int(slots[t]) + if slot < 0 or slot >= PAGES * PAGE: + continue + page, within = divmod(slot, PAGE) + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") + for plane, name in ((0, "k"), (1, "v")): + scales = kv_scale[page, plane, 0].reshape(-1)[scale_index(within, plane).cuda()] + dq = decode_nvfp4_token(kv_data[page, plane, 0, within], scales, KV_QUANT[1 + plane]) + r = ref[name][t].bfloat16().float() + # Round to nearest E2M1 under the block's E4M3-rounded scale s (in value units s / quant_scale): at most + # half the widest E2M1 step (one unit), plus the clipping at 6 units when s rounded down; a scale that + # underflows E4M3 zeroes the block. A slack of one E4M3 step covers the kernel's approximate reciprocals. + qs = KV_QUANT[1 + plane] + amax = r.view(-1, FP4_BLOCK).abs().amax(dim=1) + unit = (qs * amax / 6.0).to(torch.float8_e4m3fn).float() / qs + bound = torch.where( + unit > 0, 1.15 * unit + (amax - 6.0 * unit / 1.15).clamp_min(0.0), amax + ) + bound = bound.repeat_interleave(FP4_BLOCK) + 1e-12 + assert bool(((dq - r).abs() <= bound).all()), (name, t) + checked += 1 + assert checked >= spec["tokens"] - 2 + + +def test_early_trigger_does_not_change_the_bits() -> None: + inputs = _inputs(48, 1.0, False, 7) + a, b = _caches(), _caches() + out_a = op(*_args(inputs, a, early=True)) + out_b = op(*_args(inputs, b, early=False)) + for x, y in zip((*out_a, *a), (*out_b, *b)): + assert torch.equal(x.view(torch.uint8), y.view(torch.uint8)) + + +def test_guard_refuses_a_repeated_slot() -> None: + inputs = list(_inputs(4, 1.0, False, 3)) + inputs[1] = inputs[1].clone() + inputs[1][1] = inputs[1][0] + with pytest.raises(AssertionError, match="must not repeat a slot"): + with validating(op): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_operand_off_the_device() -> None: + inputs = list(_inputs(4, 1.0, False, 9)) + inputs[5] = inputs[5].cpu() # position ids on the host + with pytest.raises(RuntimeError, match="packed's CUDA device"): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_unaligned_index_k_page_stride() -> None: + """The index-K rows are written as 32-bit words: a page stride that is not a multiple of 4 is refused.""" + kv_data, kv_scale, _ = _caches() + backing = torch.zeros(PAGES * (128 * 128 + 1), dtype=torch.float8_e4m3fn, device="cuda") + index_k = backing.as_strided((PAGES, 1, PAGE, 128), (128 * 128 + 1, 128 * 128, 128, 1)) + with pytest.raises(RuntimeError, match="index_k"): + op(*_args(_inputs(4, 1.0, False, 13), (kv_data, kv_scale, index_k))) + + +def test_op_refuses_unaligned_weights_or_rope_table() -> None: + """The norm weights are read as 8-byte words and the RoPE rows as float4: offset views are refused.""" + packed, slots, quant, weights, rotary, positions = _inputs(4, 1.0, False, 17) + offset_weight = torch.zeros(129, dtype=torch.bfloat16, device="cuda")[1:] + with pytest.raises(RuntimeError, match="norm weights"): + op( + *_args( + (packed, slots, quant, [offset_weight, *weights[1:]], rotary, positions), _caches() + ) + ) + flat = torch.zeros(rotary.numel() + 1, device="cuda") + offset_rope = flat[1:].view(rotary.shape) + with pytest.raises(RuntimeError, match="rotary_cos_sin"): + op(*_args((packed, slots, quant, weights, offset_rope, positions), _caches())) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..9d16b5bb945a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_qk_norm_rope_fp8 entry of the MiniMax-M3 tp8 target: the out-of-place FP8 form of the fused +per-head Gemma RMSNorm + partial NeoX RoPE against the in-place kernel followed by the E4M3 cast. V is a cast, so it +matches bit for bit; Q and K round once from fp32 where the composition rounds through bf16 first, so each element is +within one E4M3 step of it.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_qk_norm_rope_fp8 import ( # noqa: E501 + m3_qk_norm_rope_fp8 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_qk_norm_rope_fp8 requires a CUDA device" + +HEAD_DIM, ROTARY_DIM, EPS, THETA = 128, 64, 1e-6, 5e6 + + +def _args(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + q, kv = spec["q"], spec["kv"] + qkv = torch.randn(spec["m"], (q + 2 * kv) * HEAD_DIM, generator=g, device="cuda").to( + torch.bfloat16 + ) + qw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (spec["m"],), generator=g, device="cuda", dtype=torch.int32) + return ( + qkv, + q, + kv, + kv, + HEAD_DIM, + ROTARY_DIM, + EPS, + qw, + kw, + THETA, + True, + pos, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + args = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + before = args[0].clone() + with validating(op): + out = op(*args) + assert torch.equal(args[0], before), "the FP8 form must leave qkv untouched" + ref = op.reference(*args) + assert out.dtype == torch.float8_e4m3fn and out.shape == ref.shape + v_start = (cell.spec["q"] + cell.spec["kv"]) * HEAD_DIM + assert torch.equal(out[:, v_start:].view(torch.uint8), ref[:, v_start:].view(torch.uint8)) + assert _e4m3_steps(out[:, :v_start], ref[:, :v_start]) <= 1 + + +def _e4m3_steps(a, b) -> int: + """The largest distance between two E4M3 tensors in representable values (E4M3 codes order by magnitude).""" + + def signed(t): + c = t.view(torch.uint8).int() + # Sign-magnitude to a line: -0 and +0 both map to 0. + return torch.where(c >= 128, 128 - c, c) + + return int((signed(a) - signed(b)).abs().max()) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_qkv_insert.py new file mode 100644 index 000000000000..eecd77e17f8d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_qkv_insert.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_qkv_insert entry: the generic horizontal producer on an FP8 main cache. + +Q, index-Q and every written K / V / index-K cache row are held within one E4M3 step of the fp32 reference. +""" + +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.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_qkv_insert import ( # noqa: E501 + m3_qkv_insert as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_qkv_insert requires a CUDA device" + +PAGES, PAGE, HEADS, KV, INDEX, D, ROTARY = 256, 128, 8, 1, 1, 128, 64 +MAX_POSITIONS = 262144 +EPS = 1e-6 + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +def _args(tokens: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + packed = torch.randn( + tokens, (HEADS + 2 * KV + INDEX + 1) * D, device="cuda", generator=g + ).bfloat16() + kv = torch.zeros(PAGES, 2, KV, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + index_k = torch.zeros(PAGES, 1, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + weights = [(torch.randn(D, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + rotary = torch.randn(MAX_POSITIONS, 2, ROTARY // 2, device="cuda", generator=g) + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + return ( + packed, + kv, + None, + index_k, + slots, + None, + HEADS, + KV, + INDEX, + D, + ROTARY, + EPS, + *weights, + rotary, + positions, + ) + + +@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 this name + args = _args(cell.spec["m"], zlib.crc32(cell.why.encode())) + certifying(op) + ref = op.reference(*args) + q, index_q = op(*args) + torch.cuda.synchronize() + _within_fp8(q.view(ref["q"].shape), ref["q"], "q") + _within_fp8(index_q.view(ref["index_q"].shape), ref["index_q"], "index_q") + kv, index_k, slots = args[1], args[3], args[4] + for t in range(cell.spec["m"]): + page, within = divmod(int(slots[t]), PAGE) + _within_fp8(kv[page, 0, :, within], ref["k"][t], f"k token {t}") + _within_fp8(kv[page, 1, :, within], ref["v"][t], f"v token {t}") + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_select_blocks.py new file mode 100644 index 000000000000..c392993acf4a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_select_blocks.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_select_blocks entry of the MiniMax-M3 tp8 target: MSA top-16 block selection with the forced +init / local blocks and per-row valid-block masks.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_select_blocks import ( # noqa: E501 + m3_select_blocks as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_blocks requires a CUDA device" + +TOPK, INIT, LOCAL = 16, 0, 1 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + n_blocks, total_q = spec["n_blocks"], spec["total_q"] + # Distinct scores (a permutation) so the selection has no ties to break. + scores = ( + torch.randperm(n_blocks * total_q, generator=g, device="cuda") + .float() + .view(1, n_blocks, total_q) + ) + if spec["mixed"]: + n_valid = torch.randint( + 17, n_blocks + 1, (total_q,), generator=g, device="cuda", dtype=torch.int32 + ) + else: + n_valid = torch.full((total_q,), n_blocks, dtype=torch.int32, device="cuda") + with validating(op): + out = op(scores, TOPK, n_valid, INIT, LOCAL, True) + ref = op.reference(scores, TOPK, n_valid, INIT, LOCAL, True) + assert out.shape == ref.shape + assert torch.equal(out, ref) + assert out.permute(1, 0, 2).is_contiguous(), "head-major backing" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_select_long.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_select_long.py new file mode 100644 index 000000000000..96ad7bb7b05c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_select_long.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_select_long entry (top-16 MSA block selection for long decode rows).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_select_long import ( # noqa: E501 + m3_select_long as op, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + select_blocks_from_maxscore, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_long requires a CUDA device" + +# Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 2-, 4- and 9-key +# widths, past the row (clamped), and zero (a CUDA-graph padding row). +_MIXED_VALID = (0, 5, 16, 17, 200, 2048, 2049, 4096, 4097, 8193, 9000, 1, 300, 6000, 129, 3000) + + +def _scores(spec, g): + nb, tq = spec["num_blocks"], spec["total_q"] + s = torch.randn(spec["num_kv_heads"], nb, tq, generator=g, device="cuda", dtype=torch.float32) + kind = spec["scores"] + if kind == "ties": + # Even tokens: a flat row with ten winners, so thousands of keys tie at the 16th-largest score (the + # re-read path). Odd tokens: scores on a coarse grid, so ties are broken by block index. + s[:, :, 0::2] = 0.25 + winners = torch.randint(0, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda") + s[:, :, 0::2].scatter_(1, winners[:, :, 0::2], 3.0) + s[:, :, 1::2] = (s[:, :, 1::2] * 2).round() / 2 + elif kind == "nonfinite": + s[:, 7::50, :] = float("inf") + s[:, 11::70, :] = float("nan") + s[:, 3::40, :] = -0.0 + s[:, 4::40, :] = 0.0 + # Odd tokens: almost everything -inf, so picks run past the finite scores. + s[:, :, 1::2] = float("-inf") + s[:, 0:5, 1::2] = 1.0 + return s + + +def _valid(spec, g): + tq, nb = spec["total_q"], spec["num_blocks"] + if spec["valid"] == "mixed": + vals = [_MIXED_VALID[i % len(_MIXED_VALID)] for i in range(tq)] + else: + vals = (nb - torch.randint(0, 12, (tq,), generator=g, device="cuda")).tolist() + vals[0] = nb + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _case(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + scores = _scores(spec, g) + nvb = _valid(spec, g) + out = torch.empty( + spec["num_kv_heads"], spec["total_q"], TOPK, dtype=torch.int32, device="cuda" + ).permute(1, 0, 2) + return scores, nvb, out, spec["init_blocks"], spec["local_blocks"], False + + +def _mass_tie_rows(scores, nvb, init_blocks, local_blocks): + """[total_q, num_kv_heads] mask of rows where more than 1024 valid blocks share the 16th-largest score.""" + num_kv_heads, num_blocks, total_q = scores.shape + blocks = torch.arange(num_blocks, device=scores.device) + raw = nvb[:total_q].long() + valid = raw.clamp(0, num_blocks) + forced = scores.permute(2, 0, 1).clone() # [total_q, heads, blocks] + forced = torch.where(blocks < init_blocks, 1.0e30, forced) + forced = torch.where(blocks >= (raw - local_blocks).clamp_min(0)[:, None, None], 1.0e29, forced) + in_row = (blocks[None, :] < valid[:, None])[:, None, :] + ranked = torch.where(in_row, forced, float("-inf")).sort(dim=-1, descending=True).values + kth = ranked[..., min(TOPK, num_blocks) - 1 : min(TOPK, num_blocks)] + return ((forced == kth) & in_row).sum(dim=-1) > 1024 + + +@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 this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + + # Selection is an exact claim: the same blocks or a different attention. + assert out.dtype is torch.int32 + assert torch.equal(out, ref), "selected blocks differ from the reference" + # And it is the in-tree selector's answer, bit for bit (the op the target falls back to) -- except on rows + # where more than 1024 blocks tie at the 16th-largest score: there the in-tree op does not break the tie to + # the smallest blocks, while this op and the reference do. + scores, nvb, _, init_blocks, local_blocks, _ = args + in_tree = select_blocks_from_maxscore( + scores, + topk=TOPK, + n_valid_blocks=nvb, + init_blocks=init_blocks, + local_blocks=local_blocks, + head_major_output=True, + ) + mass_tie = _mass_tie_rows(scores, nvb, init_blocks, local_blocks) + assert torch.equal(out[~mass_tie], in_tree[~mass_tie]), ( + "selected blocks differ from the in-tree minimax_m3_select_blocks" + ) + + +def test_pdl_launch_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=1) + pdl_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, False) + op(scores, nvb, pdl_out, init_blocks, local_blocks, True) + assert torch.equal(out, pdl_out) + + +def test_early_trigger_selects_the_same() -> None: + """Releasing the dependent grid right after the select's own wait changes only when the next kernel may launch.""" + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=2) + early_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False) + op(scores, nvb, early_out, init_blocks, local_blocks, True, True) + assert torch.equal(out, early_out) + + +def test_op_refuses_a_row_past_its_width() -> None: + scores = torch.zeros(1, 9217, 1, device="cuda") + nvb = torch.ones(1, dtype=torch.int32, device="cuda") + out = torch.empty(1, 1, TOPK, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="at most 9216"): + op(scores, nvb, out, 0, 1, False) + + +def test_guard_refuses_negative_forcing() -> None: + args = list(_case(op.CELLS[0].spec, seed=2)) + args[4] = -1 + with pytest.raises(AssertionError, match="non-negative"): + with validating(op): + op(*args) + + +def test_guard_refuses_out_aliasing_scores() -> None: + scores, nvb, _, init_blocks, local_blocks, _ = _case(op.CELLS[0].spec, seed=3) + aliased = ( + scores.view(torch.int32) + .reshape(-1)[: scores.shape[2] * TOPK] + .view(scores.shape[2], 1, TOPK) + ) + with pytest.raises(AssertionError, match="alias"): + with validating(op): + op(scores, nvb, aliased, init_blocks, local_blocks, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_sparse_decode.py new file mode 100644 index 000000000000..3e6b37f90d8d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_sparse_decode.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sparse_decode entry (the target's Triton block-sparse decode and its routes).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_sparse_decode import ( # noqa: E501 + m3_sparse_decode as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_nvfp4_cache import ( # noqa: E402 — needs the path declared above + HEAD_DIM, + NUM_HEADS, + build_step, +) +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_sparse_decode requires a CUDA device" + +# Band, in BF16 ulps (2^-8) of each [8, 128] head row's largest magnitude, and of the relative RMS. Every route +# computes from exact operands (E4M3 q and E2M1 x E4M3 K/V are exact in FP16), so the gap to the float64 reference +# is the BF16 rounding on both sides (one ulp of the row at most), the FP16 probabilities before PV, the split-K +# partials kept in FP32 and merged by their log2-sum-exp, and approximate exp2/log2: 2 ulps per element, the two +# roundings (each at most half an ulp) dominating the RMS, so half an ulp. None of these depends on the head count: +# measured on GB300 over the same cells at the tp4 target's 16 heads, at most 1.80 ulps per element and 0.20 ulps +# RMS. +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + + +def _call(case, output, identity): + return op( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + identity, + ) + + +@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 this name + spec = cell.spec + case = build_step( + spec["seq_lens"], + spec["decode_query_len"], + identity=spec["identity"], + seed=zlib.crc32(cell.why.encode()) % 2**31, + ) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + certifying(op) + _call(case, output, spec["identity"]) + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + spec["identity"], + ) + assert torch.isfinite(output.float()).all(), "padding rows must be zeros, not NaN" + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_identity_flag_does_not_change_the_answer() -> None: + """On an identity table the identity routes (the producer, or the multi-query producer from 17 rows) and the + selected-table route (the long-context kernel) compute the same attention.""" + case = build_step((2048, 130, 5, 0, 1777, 900, 2000, 0), 4, identity=True, seed=17) + outs = [] + for identity in (True, False): + output = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + _call(case, output, identity) + outs.append(output) + assert_within_ulp(outs[0], outs[1], element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_guard_refuses_a_padded_selection_slot() -> None: + """The Triton producers index the block table with every one of a token's first min(16, blocks) ids, + -1 included.""" + case = build_step((1500,), 1, identity=True, seed=23) + case.topk[0, 0, 3] = -1 + output = torch.empty(1, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="blocks of the request"): + with validating(op): + _call(case, output, True) + + +def test_guard_refuses_a_false_identity_promise() -> None: + """identity_blocks lets the multi-query producer read one table for the whole request.""" + case = build_step((1500, 1600, 1700, 1800, 1500, 1600, 1700, 1800), 4, identity=True, seed=29) + case.topk[0, 5, :2] = case.topk[0, 5, :2].flip(0) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="identity_blocks promises"): + with validating(op): + _call(case, output, True) + + +def test_a_selected_step_the_lc8_kernel_declines_is_a_miss() -> None: + """A selected table at the TP8 heads is the one-launch kernel's (modeling's SPARSE_DECODE_LC8): an output it + cannot write (here a strided view) raises under require instead of silently taking the split-K chain.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 + modeling as dispatch, + ) + + case = build_step((9000, 40000), 4, identity=False, seed=31) + wide = torch.empty( + case.q.shape[0], NUM_HEADS, 2 * HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + output = wide[..., :HEAD_DIM] + assert not output.is_contiguous() + dispatch.reset_logged() + dispatch.set_require(True) + try: + with pytest.raises(dispatch.FastPathUnavailable, match="m3_sparse_decode_lc8"): + _call(case, output, False) + finally: + dispatch.set_require(False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py new file mode 100644 index 000000000000..39766f50d67f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py @@ -0,0 +1,195 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sparse_decode_lc8 entry (one-launch long-context sparse decode at the TP8 8 query heads, +``torch.ops.trtllm.minimax_m3_tp8_sparse_decode_lc8``). + +Every cell is held to the float64 reference within the derived band, its MXFP8 hand-off to the recipe bit for bit, and +the two ``prewait`` variants (staging before or after the grid wait) to each other bit for bit. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_sparse_decode_lc8 import ( # noqa: E501 + NUM_HEADS, + WIDTH, + dequant_nvfp4_pages, + mxfp8_handoff_reference, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_sparse_decode_lc8 import ( # noqa: E501 + m3_sparse_decode_lc8 as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_nvfp4_cache import NUM_HEADS as CASE_HEADS # noqa: E402 — needs the path declared above +from _m3_nvfp4_cache import build_step # noqa: E402 +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_sparse_decode_lc8 requires a CUDA device" + +#: Hand-off buffers are sized as MERGE_MXFP8 sizes them. +_HANDOFF_ROWS = 128 + +# Band, in BF16 ulps (2^-8) of each output row's largest magnitude, and of the relative RMS. +# +# The reference is float64 attention over the exactly dequantized blocks, rounded once to BF16. The kernel's +# operands are exact too (E4M3 q and E2M1 x E4M3 K/V are exact in FP16), so what separates the two is: the final +# BF16 rounding on both sides (up to one ulp of the element, at most one of the row's scale), the softmax +# probabilities rounded to FP16 before PV (2^-11 relative each, averaging down over the block), ex2.approx / +# lg2.approx (~2^-22) and FP32 accumulation order. Hence 2 ulps per element; the RMS is dominated by the two +# roundings (each at most half an ulp, about 0.3 in aggregate), so half an ulp. None of these depends on the head +# count or on the MMA orientation: the tp4 target's 16-head kernel measured at most 1.91 ulps per element and 0.25 +# ulps RMS on GB300 over the same cells. +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + + +def _args(case, out, q8, sf, linear_softmax, prewait, num_sf_k_tiles=WIDTH // 128): + return ( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + num_sf_k_tiles, + case.decode_query_len, + case.sm_scale, + linear_softmax, + False, + prewait, + ) + + +def _buffers(total_q, mxfp8): + out = torch.empty(total_q, WIDTH, dtype=torch.bfloat16, device="cuda") + if not mxfp8: + return out, None, None + q8 = torch.zeros(_HANDOFF_ROWS, WIDTH, dtype=torch.uint8, device="cuda") + sf = torch.zeros(_HANDOFF_ROWS * (WIDTH // 32), dtype=torch.uint8, device="cuda") + return out, q8, sf + + +def _case(spec, seed): + return build_step(spec["seq_lens"], spec["decode_query_len"], identity=False, seed=seed) + + +def test_the_cases_are_the_tp8_geometry() -> None: + assert CASE_HEADS == NUM_HEADS == 8 and WIDTH == 1024 + + +@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 this name + spec = cell.spec + case = _case(spec, zlib.crc32(cell.why.encode()) % 2**31) + total_q = case.q.shape[0] + out, q8, sf = _buffers(total_q, spec["mxfp8"]) + args = _args(case, out, q8, sf, spec["linear_softmax"], spec["prewait"]) + certifying(op) + op(*args) + ref = op.reference(*args) + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + # Padding rows attend nothing and must be exact zeros, not NaN. + assert torch.isfinite(out.float()).all() + if spec["mxfp8"]: + # The hand-off is a deterministic function of the op's own BF16 rows: exact. + ref_q8, ref_sf = mxfp8_handoff_reference(out, _HANDOFF_ROWS) + assert torch.equal(q8[:total_q], ref_q8), ( + "MXFP8 rows differ from the recipe applied to the BF16 rows" + ) + assert torch.equal(sf, ref_sf), "UE8M0 scales differ from the recipe / swizzle" + + # The other staging variant: the same arithmetic, so the same bits (rows and hand-off). + other_out, other_q8, other_sf = _buffers(total_q, spec["mxfp8"]) + op(*_args(case, other_out, other_q8, other_sf, spec["linear_softmax"], not spec["prewait"])) + assert torch.equal(other_out.view(torch.int16), out.view(torch.int16)), ( + "prewait and plain staging wrote different rows" + ) + if spec["mxfp8"]: + assert torch.equal(other_q8, q8) and torch.equal(other_sf, sf) + + +def test_reference_dequant_matches_the_quantizer() -> None: + """The reference reads the paged bytes through its own layout claims (K scales token-major, V scales in + 4x4 token quads). Hold them to the quantizer's own flat output: every slot was written in order, so slot + ``s`` is page ``s // 128``, token ``s % 128``, and the flat scales are token-major by construction.""" + case = build_step((3000,), 1, identity=False, seed=5) + pages = torch.arange(case.k_paged.shape[0], device="cuda") + views = ((case.k_paged, case.k_block_scale, False), (case.v_paged, case.v_block_scale, True)) + for role, (data, scale, quad) in enumerate(views): + flat_data, flat_scale = torch.ops.trtllm.fp4_quantize( + case.source[role].reshape(case.num_slots, 1, -1).contiguous(), + case.quant[role + 1 : role + 2], + 16, + False, + False, + ) + expected = dequant_nvfp4_pages( + flat_data.view(torch.uint8).reshape(-1, 1, 128, 64), + flat_scale.view(torch.uint8).reshape(-1, 1, 128, 8), + pages, + quad_scales=False, + ) + got = dequant_nvfp4_pages(data, scale, pages, quad_scales=quad) + assert torch.equal(got, expected), ( + f"role {role}: paged dequantization disagrees with the quantizer" + ) + + +def test_the_variants_agree() -> None: + """The two softmax recurrences, the hand-off switch and the staging must not change the BF16 rows beyond + rounding; the hand-off switch and the staging not at all.""" + case = build_step((5000, 12345), 4, identity=False, seed=11) + outs = {} + for linear in (False, True): + for mxfp8 in (False, True): + for prewait in (False, True): + out, q8, sf = _buffers(case.q.shape[0], mxfp8) + op(*_args(case, out, q8, sf, linear, prewait)) + outs[(linear, mxfp8, prewait)] = out + base = outs[(False, False, False)] + for key, other in outs.items(): + assert_within_ulp(other, base, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + same_recurrence = outs[(key[0], False, False)] + torch.testing.assert_close(other, same_recurrence, rtol=0, atol=0) + + +def test_guard_refuses_a_block_past_the_request() -> None: + """A selected id at or past the request's block count reads a page the token cannot attend; with every + position masked the CTA's softmax is NaN.""" + case = build_step((3000,), 1, identity=False, seed=4) + case.topk[0, 0, 0] = 40 + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="block count"): + with validating(op): + op(*_args(case, out, q8, sf, False, True)) + + +def test_guard_refuses_strided_page_rows() -> None: + """The kernel walks a page's token rows at a fixed 64-byte pitch; a padded row stride is misread.""" + case = build_step((3000,), 1, identity=False, seed=6) + wide = torch.zeros(*case.k_paged.shape[:3], 80, dtype=torch.uint8, device="cuda")[..., :64] + wide.copy_(case.k_paged) + case.k_paged = wide + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="64-byte rows"): + with validating(op): + op(*_args(case, out, q8, sf, False, True)) + + +def test_guard_refuses_a_narrow_handoff() -> None: + """The hand-off rows are 1024 wide (8 scale k-tiles); a smaller count overlaps them.""" + case = build_step((3000,), 1, identity=False, seed=7) + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="num_sf_k_tiles"): + with validating(op): + op(*_args(case, out, q8, sf, False, True, num_sf_k_tiles=4)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_msa_kv_update.py new file mode 100644 index 000000000000..458a7000d9fd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_msa_kv_update.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the msa_kv_update entry (the one-launch MSA KV-length correction).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.msa_kv_update import ( # noqa: E501 + TOKENS_PER_BLOCK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.msa_kv_update import ( # noqa: E501 + msa_kv_update as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.msa_kv_update import ( # noqa: E501 + msa_kv_update as hook, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "msa_kv_update requires a CUDA device" + +#: Spare capacity past the live rows, as the metadata's buffers have: the op must leave it alone. +_SPARE = 3 + + +def _case(spec, seed): + """The metadata's buffers for one step, with stale values in every output so a skipped write shows.""" + g = torch.Generator(device="cpu").manual_seed(seed) + batch, qlen = spec["batch"], spec["query_len"] + total_q = batch * qlen + staged = torch.randint(qlen, 5000, (batch,), generator=g, dtype=torch.int32) + accept = spec["accept"] + if accept == "full": + kv = staged.clone() + elif accept == "page_edge": + staged = torch.tensor([130, 257, 2 * TOKENS_PER_BLOCK + qlen][:batch], dtype=torch.int32) + kv = staged - torch.tensor([0, 1, qlen - 1][:batch], dtype=torch.int32) + else: + kv = staged - torch.randint(0, qlen, (batch,), generator=g, dtype=torch.int32) + if accept == "padded": + kv[batch // 2 :] = 0 + staged[batch // 2 :] = 0 + ncols = (5000 + qlen) // TOKENS_PER_BLOCK + 2 + table = ( + torch.randperm(batch * ncols + 7, generator=g)[: batch * ncols] + .to(torch.int32) + .reshape(batch, ncols) + ) + dev = dict(device="cuda", dtype=torch.int32) + stale = -12345 + return ( + torch.cat([kv, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.cat([staged, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.full((batch + _SPARE,), qlen, **dev), + torch.arange(batch, **dev).repeat_interleave(qlen), + torch.arange(qlen, **dev).repeat(batch), + table.cuda(), + TOKENS_PER_BLOCK, + batch, + total_q, + torch.full((batch + _SPARE,), stale, **dev), + torch.full((batch + 1 + _SPARE,), stale, **dev) if spec["cu_kv"] else None, + torch.full((total_q + _SPARE,), stale, **dev), + torch.full((total_q + _SPARE,), stale, **dev) if spec["n_valid"] else None, + ) + + +def _check(args, ref): + seq, cu_kv, loc, n_valid = args[9], args[10], args[11], args[12] + batch, total_q = args[7], args[8] + ref_seq, ref_cu, ref_loc, ref_nv = ref + # Integer bookkeeping: exact, and nothing past the live rows is touched. + assert torch.equal(seq[:batch], ref_seq) + assert torch.equal(loc[:total_q], ref_loc) + assert bool((seq[batch:] == -12345).all()) and bool((loc[total_q:] == -12345).all()) + assert (cu_kv is None) == (ref_cu is None) and (n_valid is None) == (ref_nv is None) + if cu_kv is not None: + assert torch.equal(cu_kv[: batch + 1], ref_cu) + assert bool((cu_kv[batch + 1 :] == -12345).all()) + if n_valid is not None: + assert torch.equal(n_valid[:total_q], ref_nv) + assert bool((n_valid[total_q:] == -12345).all()) + + +@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 this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + op(*args) + _check(args, op.reference(*args)) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_hook_takes_every_cell(cell) -> None: + """The target's hook (what the metadata calls) takes each certified step and writes what the op does.""" + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31 + 1) + assert hook(*args) is True + _check(args, op.reference(*args)) + + +def test_the_hook_declines_what_the_op_would_misread() -> None: + """A strided operand is read wrongly by the op, so the hook declines it (the metadata then runs its own + torch update) and writes nothing.""" + args = list(_case(op.CELLS[2].spec, seed=5)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + before = args[11].clone() + assert hook(*args) is False + assert torch.equal(args[11], before) + + +def test_guard_refuses_a_strided_operand() -> None: + args = list(_case(op.CELLS[2].spec, seed=6)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_request_index_past_the_batch() -> None: + args = list(_case(op.CELLS[2].spec, seed=7)) + args[3] = args[3].clone() + args[3][-1] = args[7] + with pytest.raises(AssertionError, match="q_batch_row"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_allreduce_op_matrix.py new file mode 100644 index 000000000000..7b8a360875cd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_allreduce_op_matrix.py @@ -0,0 +1,456 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the ``m3_allreduce`` entry, 4 ranks: its cells against the entry's reference, and the +fused draft loop's layer boundaries built on it (``M3Eagle3Draft.ar_norm``: the plain all-reduce, then FlashInfer's +fused add + RMSNorm) against stock Eagle3's, bit for bit. + +Stock ``Eagle3DecoderLayer`` all-reduces each projection plainly (BF16 out) and runs the ``RMSNorm`` module with the +residual: FlashInfer's ``fused_add_rmsnorm``, which adds in FP32 and normalizes the FP32 sum. The draft loop's +boundaries must be that path exactly, under every all-reduce strategy a single-node job can run (AUTO, ONESHOT, +TWOSHOT, NCCL; MNNVL, which the tp8 drafter's two-tray group may take, is multi-node only and not driven here), at the +drafter's shapes (H = 6144, eps = 1e-6, 1 to 32 rows), at the job's world size (2, 4 or 8 ranks on one node; the tp8 +drafter runs at 8 over two trays, which one node cannot form, and nothing here depends on the rank count). The inputs +are +random rows and rows built so a BF16 rounding of the residual sum would move the norm the most, including the +review's case: unit weights, sums 1 + 2^-8 on 6143 coordinates and 2 - 2^-8 on the last, both BF16 rounding +midpoints that round to even in opposite directions. + +Asserts the normalized rows and the residual sum bitwise equal to the stock path's, and reports the max ulp of each +against an fp64 reference (ulps of BF16 at the larger magnitude; a correctly rounded output is within 0.5). + +Its own launcher, one rank per device in CUDA_VISIBLE_DEVICES, as ``_m3_mc_allreduce_norm_op_matrix.py`` is; the +collected entry point is ``test_modeling_v2_m3_allreduce.py``. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_allreduce_op_matrix.py +""" + +import math +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +TOKENS = (1, 4, 8, 32) +RESULT_TAG = "[ulp]" +"""Prefix of rank 0's result lines; the pytest entry reports them.""" + +#: World sizes the matrix runs at (one node): see ``_rank_job.WORLD_SIZES``. +WORLD_SIZES = (2, 4, 8) +RANK = 0 +WORLD = 1 + +Inputs = Tuple[List[torch.Tensor], torch.Tensor, torch.Tensor] +"""(every rank's [m, H] BF16 partial rows, the [m, H] BF16 residual, the [H] BF16 norm weight).""" + + +# ---------------------------------------------------------------- inputs (every rank builds every rank's rows) + + +def _gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _spread(parts: List[torch.Tensor]) -> List[torch.Tensor]: + """Every rank's BF16 partial: ``parts`` on the first ranks, zeros on the rest.""" + assert len(parts) <= WORLD, (len(parts), WORLD) + zeros = [torch.zeros_like(parts[0]) for _ in range(WORLD - len(parts))] + return [p.bfloat16() for p in parts] + [z.bfloat16() for z in zeros] + + +def _random(seed: int, m: int, x_scale: float, r_scale: float, unit_w: bool) -> Inputs: + g = _gen(seed) + xs = [ + (torch.randn(m, HIDDEN, generator=g, device="cuda") * x_scale).bfloat16() + for _ in range(WORLD) + ] + residual = (torch.randn(m, HIDDEN, generator=g, device="cuda") * r_scale).bfloat16() + if unit_w: + w = torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16) + else: + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return xs, residual, w + + +def _on_rank0(m: int, x: torch.Tensor) -> List[torch.Tensor]: + """``x`` as rank 0's partial, zeros elsewhere (the cross-rank sum is exact).""" + return [x.bfloat16()] + [ + torch.zeros(m, HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in range(WORLD - 1) + ] + + +def _review_case(m: int) -> Inputs: + """Sums 1 + 2^-8 on the first H - 1 coordinates and 2 - 2^-8 on the last, unit weights. + + Both sums are BF16 midpoints: 1 + 2^-8 rounds (to even) down to 1 and 2 - 2^-8 up to 2, so the fused path's + norm sees a residual moved in opposite directions on the bulk and on the last coordinate. + """ + residual = torch.ones(m, HIDDEN, device="cuda") + residual[:, -1] = 2.0 - 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _review_case_flipped(m: int) -> Inputs: + """The mirror: the bulk on odd-mantissa midpoints (1 + 3 * 2^-8 rounds up), the last one rounding down.""" + residual = torch.full((m, HIDDEN), 1.0 + 2.0**-7, device="cuda") + residual[:, -1] = 2.0 - 2.0 * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _midpoints(seed: int, m: int, outliers: int) -> Inputs: + """Every sum a BF16 midpoint: residual r random in [1, 2) and x = half its ulp (+2^-8), split over the ranks. + + The bulk has an odd mantissa (rounds up); ``outliers`` coordinates per row carry an even one (rounds down) and + a larger magnitude, so they weigh in the RMS while rounding the other way. + """ + g = _gen(seed) + mant = ( + torch.randint(0, 64, (m, HIDDEN), generator=g, device="cuda") * 2 + 1 + ) # odd 7-bit mantissa + r = 1.0 + mant.float() * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + if outliers: + cols = torch.randperm(HIDDEN, generator=g, device="cuda")[:outliers] + even = torch.randint(0, 64, (m, outliers), generator=g, device="cuda") * 2 + r[:, cols] = 8.0 * (1.0 + even.float() * 2.0**-7) + x[:, cols] = 8.0 * 2.0**-8 + # Split x over the ranks in exact BF16 pieces: x/2 on rank 0, x/4 on ranks 1 and 2 (x/2 on rank 1 at two ranks), + # 0 on the rest. + parts = [x / 2, x / 2] if WORLD == 2 else [x / 2, x / 4, x / 4] + return ( + _spread(parts), + r.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _cancelling(seed: int, m: int) -> Inputs: + """Large partials that cancel across the ranks (+A, -A, +B, -B + d; +A, -A + d at two or three ranks) onto a unit + residual.""" + g = _gen(seed) + a = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + b = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + d = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 2.0**-6).bfloat16().float() + parts = [a, -a, b, (-b + d)] if WORLD >= 4 else [a, (-a + d)] + residual = (1.0 + torch.rand(m, HIDDEN, generator=g, device="cuda")).bfloat16() + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return _spread(parts), residual, w + + +def _cases() -> List[Tuple[str, Callable[[int], Inputs]]]: + cases: List[Tuple[str, Callable[[int], Inputs]]] = [ + ("review: 1+2^-8 x6143, 2-2^-8 x1", _review_case), + ("review mirrored", _review_case_flipped), + ("all midpoints, no outliers", lambda m: _midpoints(11, m, 0)), + ("all midpoints, 1 outlier", lambda m: _midpoints(12, m, 1)), + ("all midpoints, 64 outliers", lambda m: _midpoints(13, m, 64)), + ("cancelling partials", lambda m: _cancelling(14, m)), + ] + for i, (xs, rs, unit) in enumerate( + ((1.0, 4.0, False), (1.0, 1.0, True), (0.05, 8.0, False), (8.0, 0.5, False)) + ): + cases.append( + ( + f"random x{xs} r{rs}{' unit-w' if unit else ''}", + lambda m, i=i, xs=xs, rs=rs, unit=unit: _random(100 + i, m, xs, rs, unit), + ) + ) + return cases + + +# ---------------------------------------------------------------- the three paths + + +def _fp64(xs: List[torch.Tensor], residual: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + s = residual.double() + for x in xs: + s = s + x.double() + rms = torch.rsqrt(s.square().mean(-1, keepdim=True) + EPS) + return s * rms * w.double() + + +def _ulp_at(v: torch.Tensor) -> torch.Tensor: + """BF16's ulp at |v| (fp64): 2^(e - 8) for |v| in [2^(e-1), 2^e), floored at the smallest normal's.""" + _, e = torch.frexp(v.abs().double()) + return torch.ldexp(torch.ones_like(v, dtype=torch.float64), torch.clamp(e - 8, min=-133)) + + +def _ulps(out: torch.Tensor, ref: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the larger magnitude of the two.""" + o, r = out.double(), ref.double() + return ((o - r).abs() / _ulp_at(torch.maximum(o.abs(), r.abs()))).max().item() + + +def _ulps_at(out: torch.Tensor, ref: torch.Tensor, scale: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the magnitude ``scale`` (the operands'), not at the result's.""" + return ((out.double() - ref.double()).abs() / _ulp_at(scale.abs())).max().item() + + +def _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """The draft loop's boundary (``M3Eagle3Draft.ar_norm``): ``m3_allreduce`` plain, then the catalog's + ``flashinfer_fused_add_rmsnorm``.""" + from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 + flashinfer_fused_add_rmsnorm, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + + return flashinfer_fused_add_rmsnorm(m3_allreduce(x, ar, None), residual.clone(), w, EPS) + + +def _check_cells(mapping: Any) -> None: + """Every cell of ``m3_allreduce.CELLS`` (AUTO strategy, as the target builds it) against the entry's reference: + rank-order fp32 sums with one BF16 rounding of the sum and one of the norm. A strategy may round the cross-rank sum + to BF16 before adding the residual, so the residual sum is bounded at its operands' magnitude + u = ulp(max(|residual|, |sum|)): half an ulp for that rounding plus half an output ulp for each of the two final + roundings, an output ulp being at most 2u, so 2.5u (where the two cancel, the output's own ulp is far smaller). + The normed rows are the norm of the op's own residual sum, within 2 ulps of the reference norm of that sum.""" + import sys as _sys + from pathlib import Path as _Path + + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams + from tensorrt_llm.functional import AllReduceStrategy + + _sys.path.insert(0, str(_Path(__file__).resolve().parents[4])) + from _validating import validating + + ar = AllReduce(mapping=mapping, strategy=AllReduceStrategy.AUTO, dtype=torch.bfloat16) + for i, cell in enumerate(m3_allreduce.CELLS): + spec = cell.spec + xs, residual, w = _random(200 + i, spec["m"], 1.0, 4.0, False) + params = ( + None + if spec.get("plain") + else AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual.clone(), + norm_weight=w, + eps=EPS, + trigger_completion_at_end=False, + ) + ) + # The reference first: the fused path may write the updated residual into ``params.residual``. + ref = m3_allreduce.reference(xs, ar, params) + with validating(m3_allreduce): + out = m3_allreduce(xs[RANK], ar, params) + torch.cuda.synchronize() + total = sum(x.double() for x in xs) + if params is None: + err = _ulps_at(out, ref, total) + assert err <= 1.0, (cell.why, err) + else: + err = _ulps_at(out[1], ref[1], torch.maximum(total.abs(), residual.double().abs())) + assert err <= 2.5, (cell.why, "residual", err) + h = out[1].double() + own = h * torch.rsqrt(h.square().mean(-1, keepdim=True) + EPS) * w.double() + assert _ulps(out[0], own) <= 2.0, (cell.why, "normed", _ulps(out[0], own)) + dist.barrier() + if RANK == 0: + print(f"{RESULT_TAG} cell {cell.why} passed", flush=True) + + +def _stock(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """Stock Eagle3: ``o_proj``'s plain all-reduce (BF16 out), then ``RMSNorm(h, residual)``'s FlashInfer + ``fused_add_rmsnorm`` (FP32 add, norm of the FP32 sum; in place on both).""" + h = ar(x) + res = residual.clone() + torch.ops.trtllm.flashinfer_fused_add_rmsnorm(h, res, w, EPS) + return h, res + + +# ---------------------------------------------------------------- the matrix + + +def _single_node_cluster_info() -> None: + """``TorchDist`` (the communicator under ``TLLM_DISABLE_MPI``) asks Ray for each rank's node and GPU to build its + node-local group. These ranks are spawned without Ray on one node, so every rank reports the same node and its + own device.""" + from tensorrt_llm._torch.distributed.communicator import TorchDist + + def cluster_info(self): + if self.cluster_info is None: + self.cluster_info = [("127.0.0.1", r) for r in range(dist.get_world_size())] + return self.cluster_info + + TorchDist._get_cluster_info = cluster_info + + +def _run_matrix() -> None: + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers trtllm::flashinfer_fused_add_rmsnorm + from tensorrt_llm._torch.distributed import AllReduce + from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE + from tensorrt_llm.functional import AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + assert IS_FLASHINFER_AVAILABLE, "stock Eagle3's add + RMSNorm is FlashInfer's fused_add_rmsnorm" + _single_node_cluster_info() + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK, gpus_per_node=WORLD) + _check_cells(mapping) + strategies = ( + AllReduceStrategy.AUTO, + AllReduceStrategy.ONESHOT, + AllReduceStrategy.TWOSHOT, + AllReduceStrategy.NCCL, + ) + worst = 0.0 + mismatches: List[str] = [] + for strategy in strategies: + ar = AllReduce(mapping=mapping, strategy=strategy, dtype=torch.bfloat16) + assert ar.mnnvl_allreduce is None and ar.symm_mem_allreduce is None, strategy + for name, build in _cases(): + ulp, res_ulp = 0.0, 0.0 + for m in TOKENS: + xs, residual, w = build(m) + target_out, target_res = _target(ar, xs[RANK], residual, w) + stock_out, stock_res = _stock(ar, xs[RANK], residual, w) + torch.cuda.synchronize() + for what, a, b in ( + ("normed", target_out, stock_out), + ("residual", target_res, stock_res), + ): + bad = int((a.view(torch.int16) != b.view(torch.int16)).sum().item()) + if bad: + mismatches.append( + f"{strategy.name} {name} m={m}: {bad} {what} elements differ from stock" + ) + ref = _fp64(xs, residual, w) + ref_res = residual.double() + sum(x.double() for x in xs) + u, ur = _ulps(target_out, ref), _ulps(target_res, ref_res) + assert math.isfinite(u) and math.isfinite(ur), (strategy, name, m) + ulp, res_ulp = max(ulp, u), max(res_ulp, ur) + dist.barrier() + worst = max(worst, ulp) + if RANK == 0: + print( + f"{RESULT_TAG} {strategy.name:7s} | {name:34s} | target = stock vs fp64: normed {ulp:.3f}" + f" | residual {res_ulp:.3f}", + flush=True, + ) + if RANK == 0: + scope = f"all cases, m in {TOKENS}" + print( + f"{RESULT_TAG} MAX | {scope:34s} | target = stock vs fp64: normed {worst:.3f}", + flush=True, + ) + assert not mismatches, "\n".join(mismatches) + + +def _run_one_rank() -> int: + global RANK, WORLD + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD in WORLD_SIZES, f"the matrix runs at {WORLD_SIZES} ranks on one node, got {WORLD}" + try: + _run_matrix() + except BaseException: + import traceback + + print(f"[rank {RANK}] FAILED", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + os._exit(13) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker per claimed device, each in its own session; on the first nonzero exit or the deadline every + session is killed (as ``_m3_mc_allreduce_norm_op_matrix.py`` does).""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the devices this run owns (one rank each)" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size in WORLD_SIZES, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at {WORLD_SIZES}" + ) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + 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: Optional[str] = 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/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..8a77ae91b481 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,670 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``, at the job's ranks. + +Its own launcher, as the catalog's collective matrices are: run plainly it spawns +one rank process per device in CUDA_VISIBLE_DEVICES and kills them all on the +first failure or on the deadline. A rank whose peers never push polls forever, +so the deadline is what keeps a broken exchange from wedging the calling run. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_mc_allreduce_norm_op_matrix.py + +World size: one rank per visible device, 2, 4 or 8. The tp8 target runs these +entries at 8 ranks over two GB300 trays, which one node cannot form; the world +size is a runtime argument of the kernel and no cell's spec names one, so the +cells are certified at the world size of the job (``_rank_job.WORLD_SIZES``). + +Each rank builds the three multicast workspaces (8 / 16 / 48 tokens) the target +builds, over the torch.distributed group (the non-MPI path of +``_make_mnnvl_mcast_buffer``). Every rank can regenerate every rank's inputs from +the shared seed convention, which is what keeps each reference arithmetic rather +than a second collective. Producer pushes (``prepushed`` 2 / 3) are emulated by +storing the canonicalized rows through the multicast mapping at the call parity +the flags select, which is exactly what the producing kernels store. + +Not a pytest module, despite the ``check_*`` bodies: they are one fixed sequence +inside a single multi-rank job whose workspace state (call parities) carries from +one to the next. The collected entry point is +``test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py``. +""" + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Dict, List, Optional + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +CLUSTER = 8 +BUCKETS = (8, 16, 48) +EMPTY_WORD = -2147483648 + +# Bound inside the rank body: importing the entries pulls in tensorrt_llm. +allreduce_entry: Any = None +finalize_entry: Any = None +ar_module: Any = None +validating: Any = None +#: World sizes the matrix runs at (one node): see ``_rank_job.WORLD_SIZES``. +WORLD_SIZES = (2, 4, 8) +RANK = 0 +WORLD = 1 +WORKSPACES: Dict[int, Any] = {} + + +# ---------------------------------------------------------------- inputs + + +def _gen(seed: int, rank: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed * 1009 + rank + 1) + + +def _rows(seed: int, rank: int, m: int, scale: float = 1.0) -> torch.Tensor: + """One rank's [m, H] BF16 rows; ``rank`` -1 for tensors every rank holds alike.""" + g = _gen(seed, rank + 1) + return (torch.randn(m, HIDDEN, generator=g, device="cuda") * scale).to(torch.bfloat16) + + +def _residual(seed: int, m: int) -> torch.Tensor: + return _rows(seed + 7, -1, m, 4.0) + + +def _norm_w(seed: int) -> torch.Tensor: + g = _gen(seed + 11, 0) + return (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + +def _finalize_inputs(seed: int, rank: int, m: int, topk: int, dropped: int): + """One rank's unfinalized MoE outputs: fc2 [m * topk, H], perm [m, topk], fp32 scales, shared [m, H].""" + g = _gen(seed + 23, rank) + rows = m * topk + fc2 = (torch.randn(rows, HIDDEN, generator=g, device="cuda") * 0.5).to(torch.bfloat16) + perm = torch.randperm(rows, generator=g, device="cuda").to(torch.int32) + if dropped: + drop = torch.randperm(rows, generator=g, device="cuda")[:dropped] + perm[drop] = -1 + perm = perm.view(m, topk).contiguous() + scales = torch.rand(m, topk, generator=g, device="cuda", dtype=torch.float32) + shared = _rows(seed + 29, rank, m, 0.5) + return fc2, perm, scales, shared + + +# ---------------------------------------------------------------- workspace / producer emulation + + +def _make_workspace(maxt: int): + """One bucket's exchange, as ``McAllReduceWorkspace.create`` builds it, over the torch.distributed group.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + from tensorrt_llm.mapping import Mapping + + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK) + words = 4 * maxt * WORLD * (HIDDEN // 2) + group = dist.group.WORLD + handle = _make_mnnvl_mcast_buffer(group, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar_module.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=group, + max_tokens=maxt, + ) + ws.flags_for(CLUSTER) + return ws + + +def _push(ws, rows: torch.Tensor, exchange_set: int) -> None: + """Store this rank's rows into every rank's exchange, as a producer kernel does before a prepushed call. + + Slot of (token t, cluster CTA c): set * 2 * half + parity(t, c) * half + (t * world + rank) * H / 2 + c * H / 16 + words, parity read from the CTA's flag; -0.0 halves are stored as +0.0 (0x80000000 is the empty word). + """ + m = rows.shape[0] + torch.cuda.synchronize() + parity = (ws.flags_for(CLUSTER).view(ws.max_tokens, CLUSTER)[:m] & 1).long() + half = ws.max_tokens * WORLD * (HIDDEN // 2) + chunk = HIDDEN // 2 // CLUSTER + h16 = rows.contiguous().view(torch.int16) + h16 = torch.where(h16 == -32768, torch.zeros_like(h16), h16) + words = h16.view(torch.int32).view(m, CLUSTER, chunk) + t = torch.arange(m, device="cuda").view(m, 1, 1) + c = torch.arange(CLUSTER, device="cuda").view(1, CLUSTER, 1) + w = torch.arange(chunk, device="cuda").view(1, 1, chunk) + index = ( + exchange_set * 2 * half + + parity.view(m, CLUSTER, 1) * half + + (t * WORLD + RANK) * (HIDDEN // 2) + + c * chunk + + w + ) + ws.mc[index.reshape(-1)] = words.reshape(-1) + torch.cuda.synchronize() + dist.barrier() + + +# ---------------------------------------------------------------- gates + + +def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + assert out.dtype == ref.dtype and out.shape == ref.shape, ( + where, + out.dtype, + ref.dtype, + out.shape, + ) + a, b = out.contiguous().view(torch.uint8), ref.contiguous().view(torch.uint8) + bad = (a != b).sum().item() + assert bad == 0, f"{where}: {bad} of {a.numel()} bytes differ" + + +#: Share of normed elements one BF16 step off the float64 reference: measured at most 1.1e-4 per call over +#: every call of this matrix on GB300 (see ``_assert_normed``). The gate sits ~18x above that and far below what a +#: wrong statistic produces (most elements move). +NORMED_STEP_FRACTION = 0.002 +NORMED_STEP_SEEN: List[float] = [] + + +def _assert_normed(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + """At most one BF16 rounding step from the float64 reference, and rarely that. + + The kernel's RMS statistic is a shuffle / warp / cluster tree of FP32 partial sums and its ``rsqrtf`` is + approximate (~2^-22 relative): the scaled value moves by ~1e-6 relative, which changes the BF16 rounding only + for values within that distance of a rounding midpoint -- one step, never two. The fraction bound catches a + statistic that is systematically off (a missing rank, a wrong H), which moves most elements. + """ + assert out.dtype == ref.dtype == torch.bfloat16 and out.shape == ref.shape, where + o, r = out.float(), ref.float() + _, exponent = torch.frexp(torch.maximum(o.abs(), r.abs())) + step = torch.ldexp(torch.ones_like(o), exponent - 8) + diff = (o - r).abs() + assert bool((diff <= step).all()), f"{where}: max diff {diff.max().item()} beyond one BF16 step" + frac = (diff != 0).float().mean().item() + NORMED_STEP_SEEN.append(frac) + assert frac <= NORMED_STEP_FRACTION, f"{where}: {frac:.4f} of normed elements off by one step" + + +def _assert_quant(outs, ref, m: int, where: str) -> None: + """The MXFP8 copy: bit-exact against the recipe applied to the op's own normed rows (a pure function of + them), and equal to the reference's wherever the reference's normed 32-column block is the op's.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + mxfp8_rows, + sf_written, + ) + + q, sf = outs[2], outs[3] + own_q, own_sf = mxfp8_rows(outs[0]) + written = sf_written(m, HIDDEN // 32, sf.device) + _assert_bitwise(q, own_q, f"{where} e4m3 rows") + _assert_bitwise(sf[written], own_sf[written], f"{where} UE8M0 scales") + same_block = (outs[0].view(m, -1, 32) == ref[0].view(m, -1, 32)).all(-1) + qb = q.view(torch.uint8).view(m, -1, 32) + rqb = ref[2].view(torch.uint8).view(m, -1, 32) + assert bool((qb == rqb).all(-1)[same_block].all()), ( + f"{where}: e4m3 differs on a block whose normed rows agree" + ) + + +def _check_outputs(outs, ref, m: int, quant: bool, where: str) -> None: + assert len(outs) == (4 if quant else 2), (where, len(outs)) + _assert_bitwise(outs[1], ref[1], f"{where} updated residual") + _assert_normed(outs[0], ref[0], f"{where} normed") + if quant: + _assert_quant(outs, ref, m, where) + + +# ---------------------------------------------------------------- calls + + +def _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, opts): + return allreduce_entry.m3_mc_allreduce_norm( + x, + x2, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + ws.max_tokens, + quant, + prepushed, + opts, + ) + + +def _plain_case(spec: Dict[str, Any], seed: int, where: str) -> None: + m, maxt, quant, prepushed = spec["m"], spec["maxt"], spec["quant"], spec["prepushed"] + ws = WORKSPACES[maxt] + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] if spec["x2"] else None + residual, norm_w = _residual(seed, m), _norm_w(seed) + x, x2 = xs[RANK], (None if x2s is None else x2s[RANK]) + if prepushed >= 2: + _push(ws, x, 0) + if prepushed == 3: + _push(ws, x2, 1) + with validating(allreduce_entry.m3_mc_allreduce_norm): + outs = _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, spec["opts"]) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference(xs, x2s, residual, norm_w, EPS, quant) + _check_outputs(outs, ref, m, quant, where) + dist.barrier() + + +# ---------------------------------------------------------------- checks + + +def check_plain_cells() -> None: + """Every cell of ``m3_mc_allreduce_norm.CELLS`` through ``__call__`` under ``validating``, gated on its + ``reference``.""" + for i, cell in enumerate(allreduce_entry.m3_mc_allreduce_norm.CELLS): + _plain_case(cell.spec, 1000 + 17 * i, cell.why) + + +def check_finalize_cells() -> None: + """Every cell of ``m3_mc_finalize_allreduce_norm.CELLS``, gated on its ``reference``.""" + op = finalize_entry.m3_mc_finalize_allreduce_norm + for i, cell in enumerate(op.CELLS): + spec = cell.spec + m, maxt, quant, topk = spec["m"], spec["maxt"], spec["quant"], spec["topk"] + seed = 2000 + 17 * i + ws = WORKSPACES[maxt] + per_rank = [_finalize_inputs(seed, r, m, topk, spec["dropped"]) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + with validating(op): + outs = op( + fc2, + perm, + scales, + shared, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + maxt, + quant, + spec["opts"], + ) + torch.cuda.synchronize() + ref = op.reference( + [p[0] for p in per_rank], + [p[1] for p in per_rank], + [p[2] for p in per_rank], + [p[3] for p in per_rank], + residual, + norm_w, + EPS, + quant, + ) + _check_outputs(outs, ref, m, quant, cell.why) + dist.barrier() + + +def check_the_target_wrappers_select_bucket_and_options() -> None: + """The target's own wrappers (bucket by token count, option table) on top of the entries, every token count of + the 8-token bucket plus one of each larger bucket, and the finalize wrapper; the MXFP8 publication is the + prequant slot's business and is left off.""" + buckets = ar_module.McBuckets([WORKSPACES[b] for b in BUCKETS]) + for m in (*range(1, 9), 11, 40): + seed = 3000 + m + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + normed, updated = ar_module.m3_mc_allreduce_norm( + xs[RANK], residual, norm_w, EPS, buckets, x2=x2s[RANK] + ) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference(xs, x2s, residual, norm_w, EPS, False) + _check_outputs([normed, updated], ref, m, False, f"wrapper m={m}") + dist.barrier() + for m in (9, 33): + seed = 3100 + m + per_rank = [_finalize_inputs(seed, r, m, 4, 2) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + normed, updated = ar_module.m3_mc_finalize_allreduce_norm( + fc2, perm, scales, shared, residual, norm_w, EPS, buckets + ) + torch.cuda.synchronize() + ref = finalize_entry.m3_mc_finalize_allreduce_norm.reference( + *[[p[j] for p in per_rank] for j in range(4)], residual, norm_w, EPS, False + ) + _check_outputs([normed, updated], ref, m, False, f"finalize wrapper m={m}") + dist.barrier() + + +def check_cuda_graph_replays() -> None: + """Two boundaries captured in one graph (8- and 16-token buckets), replayed with fresh inputs. + + The decode graph holds every boundary call; each replay rotates the same per-CTA parities the eager calls do, + so a replay that reused a stale half or skipped a re-arm reads the previous replay's rows and fails the + bitwise residual gate (every replay gets new inputs). + """ + shapes = ((3, 8, 15), (12, 16, 15)) + norm_w = _norm_w(4000) + bufs = [] + for m, maxt, opts in shapes: + bufs.append( + dict( + m=m, + maxt=maxt, + opts=opts, + x=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + x2=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + residual=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + ) + ) + + def body(): + return [ + _call_plain( + WORKSPACES[b["maxt"]], b["x"], b["x2"], b["residual"], norm_w, False, 0, b["opts"] + ) + for b in bufs + ] + + def fill(seed: int): + refs = [] + for j, b in enumerate(bufs): + s = seed + 31 * j + xs = [_rows(s, r, b["m"]) for r in range(WORLD)] + x2s = [_rows(s + 3, r, b["m"]) for r in range(WORLD)] + b["x"].copy_(xs[RANK]) + b["x2"].copy_(x2s[RANK]) + b["residual"].copy_(_residual(s, b["m"])) + refs.append( + allreduce_entry.m3_mc_allreduce_norm.reference( + xs, x2s, b["residual"].clone(), norm_w, EPS, False + ) + ) + torch.cuda.synchronize() + dist.barrier() + return refs + + refs = fill(4100) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + warm = body() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(warm[j], refs[j], b["m"], False, f"graph warm-up m={b['m']}") + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = body() + for rep in range(4): + refs = fill(4200 + 100 * rep) + graph.replay() + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(captured[j], refs[j], b["m"], False, f"replay {rep} m={b['m']}") + dist.barrier() + del graph + dist.barrier() + + +def check_the_gate_sees_a_missing_rank() -> None: + """The residual gate is not blind to the reduction: dropping any one rank's partial from the reference + moves most elements of the updated residual (each rank's rows are distinct by construction).""" + m, seed = 4, 5000 + xs = [_rows(seed, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + full = allreduce_entry.m3_mc_allreduce_norm.reference(xs, None, residual, norm_w, EPS, False) + for r in range(WORLD): + partial = [x if i != r else torch.zeros_like(x) for i, x in enumerate(xs)] + wrong = allreduce_entry.m3_mc_allreduce_norm.reference( + partial, None, residual, norm_w, EPS, False + ) + frac = (wrong[1] != full[1]).float().mean().item() + assert frac > 0.9, f"dropping rank {r} changes only {frac:.3f} of the residual" + dist.barrier() + + +def _expect_rejected(fn, match: str, where: str) -> None: + try: + fn() + except AssertionError as exc: + assert match in str(exc), f"{where}: rejected for the wrong reason: {exc}" + else: + raise AssertionError(f"{where}: the guard admitted it") + + +def check_guards() -> None: + """``is_valid`` refuses what the op accepts and answers wrongly. Nothing here reaches the kernel: the guard + raises first, on every rank alike, so no rank is left polling.""" + m, seed = 4, 6000 + ws8, ws16 = WORKSPACES[8], WORKSPACES[16] + x, residual, norm_w = _rows(seed, RANK, m), _residual(seed, m), _norm_w(seed) + op = allreduce_entry.m3_mc_allreduce_norm + + def call(ws_uc, flags, norm): + with validating(op): + op( + x, + None, + residual, + norm, + EPS, + ws_uc, + ws8.mc.data_ptr(), + flags, + RANK, + WORLD, + CLUSTER, + False, + 8, + False, + 0, + 15, + ) + + strided = torch.empty(2 * HIDDEN, dtype=torch.bfloat16, device="cuda")[::2] + strided.copy_(norm_w) + _expect_rejected( + lambda: call(ws8.uc, ws8.flags_for(CLUSTER), strided), "contiguous", "strided norm_w" + ) + _expect_rejected( + lambda: call(ws8.uc[: ws8.uc.numel() // 2], ws8.flags_for(CLUSTER), norm_w), + "words", + "short workspace", + ) + _expect_rejected( + lambda: call(ws8.uc, ws16.flags_for(CLUSTER), norm_w), "counters", "another bucket's flags" + ) + + fin = finalize_entry.m3_mc_finalize_allreduce_norm + fc2, perm, scales, shared = _finalize_inputs(seed, RANK, 9, 4, 0) + bad = perm.clone() + bad[0, 0] = fc2.shape[0] + + def call_fin(): + with validating(fin): + fin( + fc2, + bad, + scales, + shared, + _residual(seed, 9), + norm_w, + EPS, + ws16.uc, + ws16.mc.data_ptr(), + ws16.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + 16, + False, + 15, + ) + + _expect_rejected(call_fin, "perm spans", "perm past the FC2 rows") + dist.barrier() + + +CHECKS = ( + check_plain_cells, + check_finalize_cells, + check_the_target_wrappers_select_bucket_and_options, + check_cuda_graph_replays, + check_the_gate_sees_a_missing_rank, + check_guards, +) + + +def _run_one_rank() -> int: + """Body of one spawned rank: run every check; a failing rank exits nonzero and the launcher kills its + siblings.""" + global RANK, WORLD, allreduce_entry, finalize_entry, ar_module, validating + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD in WORLD_SIZES, f"the matrix runs at {WORLD_SIZES} ranks on one node, got {WORLD}" + + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.m3_mc_allreduce as _ar # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.comm import ( # noqa: E501 + m3_mc_allreduce_norm as _plain, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.comm import ( # noqa: E501 + m3_mc_finalize_allreduce_norm as _finalize, + ) + + allreduce_entry, finalize_entry, ar_module = _plain, _finalize, _ar + # Spawned by file path, not collected by pytest: reach the sibling helper the plain way (confined to the + # rank processes). /comm -> tests/unittest/_torch/modeling_v2. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + validating = _validating_cm + for maxt in BUCKETS: + WORKSPACES[maxt] = _make_workspace(maxt) + + 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 early wedges its peers' polls; the launcher kills them on this exit. + os._exit(13) + print(f"[rank {RANK}] {check.__name__} passed", flush=True) + dist.barrier() + torch.cuda.synchronize() + print( + f"[rank {RANK}] {len(CHECKS)} checks passed; normed one-step fraction max " + f"{max(NORMED_STEP_SEEN):.5f} over {len(NORMED_STEP_SEEN)} calls", + flush=True, + ) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker process per claimed device, each in its own session; on the first nonzero exit or on 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 (one rank each)" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size in WORLD_SIZES, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at {WORLD_SIZES}" + ) + 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", + ) + # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed + # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. + 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: Optional[str] = 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/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_rank_job.py new file mode 100644 index 000000000000..72e49f20a1e6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_rank_job.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run this target's collective rank job from pytest, and report what it did. + +The same contract as ``tests/unittest/_torch/modeling_v2/comm/_rank_job.py``: the +launcher beside this file owns the per-device rank processes and the deadline +that kills a wedged collective; this only selects the devices, starts it in a +fresh interpreter (by file path: this tree is not a package) and turns its exit +code into an assertion. + +World size. The tp8 target runs 8 ranks over two GB300 trays of 4 GPUs (one +NVLink domain); a single-node unit-test job cannot form that group. The +collective entries take the world size as a runtime argument and their cells' +specs do not name one (token counts, buckets, push modes, epilogue options), so +the jobs here certify every cell at each world size in ``WORLD_SIZES`` the +node's visible devices can form: 2 and 4 on a GB300 tray, 8 only on a node with +8 devices. A world size the node cannot form is skipped with that reason, never +emulated; the 8-rank, two-tray group itself is exercised only by the target's +multi-node runs. +""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +TARGET_WORLD_SIZE = 8 +"""tp8's world size: 8 ranks over two GB300 trays, which one node's job cannot form.""" + +WORLD_SIZES = (2, 4, TARGET_WORLD_SIZE) +"""The world sizes the jobs certify the (world-independent) cells at, where the node has that many devices.""" + +_LAUNCHER_GRACE_S = 300 +"""Headroom over the launcher's own deadline, so its message wins the race.""" + + +def visible_devices() -> list: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + return [d for d in visible.split(",") if d.strip()] + import torch + + return [str(i) for i in range(torch.cuda.device_count())] + + +def require_world(world_size: int) -> str: + """The first ``world_size`` visible devices; skips (with the reason) when the node has fewer.""" + devices = visible_devices() + if len(devices) < world_size: + where = ( + "the tp8 target's own 8-rank group spans two trays and is exercised only by its multi-node runs" + if world_size == TARGET_WORLD_SIZE + else "run this job on a node with more devices" + ) + pytest.skip( + f"{world_size} ranks need {world_size} devices; {len(devices)} visible on this node ({where})" + ) + return ",".join(devices[:world_size]) + + +def _load_launcher_constants(launcher: Path): + """Read the launcher's module-level budget without importing it by name (it imports only torch at + module scope).""" + 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, world_size: int) -> str: + """Run ``__op_matrix``'s launcher over ``world_size`` devices (skipped when the node has fewer); returns + its stdout.""" + assert world_size in WORLD_SIZES, world_size + launcher = Path(__file__).resolve().with_name(f"_{entry}_op_matrix.py") + env = dict(os.environ, CUDA_VISIBLE_DEVICES=require_world(world_size)) + + # The ranks import tensorrt_llm absolutely, and a source checkout is not necessarily installed. + # tests/unittest/_torch/modeling_v2/models/minimax_m3//comm -> repo root. + repo_root = Path(__file__).resolve().parents[8] + assert (repo_root / "tensorrt_llm").is_dir(), ( + f"expected the repo root at {repo_root}, found no tensorrt_llm/ there; " + f"this file moved without its parents[] index following" + ) + env["PYTHONPATH"] = os.pathsep.join(p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p) + + constants = _load_launcher_constants(launcher) + timeout = constants.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} at world size {world_size}\n" + f"--- stdout ---\n{completed.stdout}\n" + f"--- stderr ---\n{completed.stderr}" + ) + return completed.stdout diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_allreduce.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_allreduce.py new file mode 100644 index 000000000000..413bc0eec59b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_allreduce.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the ``m3_allreduce`` certification matrix. + +The matrix is ``_m3_allreduce_op_matrix.py`` beside this file, its own launcher; see ``_rank_job``. It drives every +cell of the entry against its reference, then requires the fused draft loop's boundaries built on it (plain +all-reduce, then the fused add + RMSNorm) to equal stock Eagle3's bit for bit under every single-node strategy, and +reports both against an fp64 reference as a warning, one line per strategy and case. + +The tp8 drafter runs at 8 ranks over two trays; a single-node job cannot form that group, so the matrix runs at each +world size of ``_rank_job.WORLD_SIZES`` the node's devices can form (the claim, draft boundary = stock boundary, does +not depend on the rank count) and skips the others. MNNVL, the strategy the two-tray group may take, is multi-node +only and not driven here. +""" + +import warnings + +import _rank_job +import pytest +import torch + +assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" + +_RESULT_TAG = "[ulp]" + + +# Spawns one rank process per device of the world size; under xdist workers would fight for them. +@pytest.mark.no_xdist +@pytest.mark.parametrize("world_size", _rank_job.WORLD_SIZES, ids=lambda w: f"world{w}") +def test_m3_allreduce_op_matrix(world_size) -> None: + out = _rank_job.run("m3_allreduce", world_size) + lines = [line for line in out.splitlines() if line.startswith(_RESULT_TAG)] + assert lines, f"the launcher reported no measurements:\n{out}" + warnings.warn( + f"draft boundary (= stock Eagle3) vs fp64 at {world_size} ranks, max BF16 ulps\n" + + "\n".join(lines), + stacklevel=1, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..8bae5f075307 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the multicast all-reduce + RMSNorm entries' certification matrix. + +Covers both ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm`` (one +kernel, one exchange protocol). The matrix is ``_m3_mc_allreduce_norm_op_matrix.py`` +beside this file, its own launcher; see ``_rank_job`` for why. + +The tp8 target runs these entries at 8 ranks over two trays, which a single-node +job cannot form. The world size is a runtime argument of the kernel and no cell's +spec names one, so the matrix certifies every cell at each world size of +``_rank_job.WORLD_SIZES`` the node's devices can form (2 and 4 on a GB300 tray, 8 +on an 8-device node) and skips the others; it does not emulate the two-tray group. +""" + +import _rank_job +import pytest +import torch + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + + +# Spawns one rank process per device of the world size; under xdist workers would fight for them. +@pytest.mark.no_xdist +@pytest.mark.parametrize("world_size", _rank_job.WORLD_SIZES, ids=lambda w: f"world{w}") +def test_m3_mc_allreduce_norm_op_matrix(world_size) -> None: + _rank_job.run("m3_mc_allreduce_norm", world_size) + + +def test_the_cells_are_world_independent() -> None: + """What lets a single-node job certify the tp8 entries: no cell's spec names a world size.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.comm import ( # noqa: E501 + m3_mc_allreduce_norm, + m3_mc_finalize_allreduce_norm, + ) + + for op in ( + m3_mc_allreduce_norm.m3_mc_allreduce_norm, + m3_mc_finalize_allreduce_norm.m3_mc_finalize_allreduce_norm, + ): + for cell in op.CELLS: + assert not {"world", "world_size", "rank"} & set(cell.spec), cell.why diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..2a8a2672b3de --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests of the ``m3_mc_finalize_allreduce_norm`` entry that need no collective. + +Its cells are certified, with ``m3_mc_allreduce_norm``'s (one kernel, one exchange protocol), by the 4-rank matrix +``_m3_mc_allreduce_norm_op_matrix.py`` that ``test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py`` runs. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.comm.m3_mc_finalize_allreduce_norm import ( # noqa: E501 + finalize_partial, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.comm.m3_mc_finalize_allreduce_norm import ( # noqa: E501 + m3_mc_finalize_allreduce_norm as op, +) + + +def test_finalize_partial_rounds_the_routed_sum_before_the_shared_expert() -> None: + """As stock: the finalized routed output is BF16, and the shared expert is added to it (shared.add_(routed)). + Four routed rows of 1.0 at FP32 routing weight 0.2509765625 sum to 1.00390625, which rounds to BF16 1.0; plus the + shared -1 that is 0, where adding before the rounding would leave 0.00390625.""" + fc2 = torch.ones((4, 8), dtype=torch.bfloat16) + perm = torch.arange(4, dtype=torch.int32).view(1, 4) + scales = torch.full((1, 4), 0.2509765625, dtype=torch.float32) + shared = torch.full((1, 8), -1.0, dtype=torch.bfloat16) + out = finalize_partial(fc2, perm, scales, shared) + routed = finalize_partial(fc2, perm, scales, None) + assert torch.equal(out, shared.clone().add_(routed)) + assert torch.equal(out, torch.zeros(1, 8, dtype=torch.bfloat16)) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_each_cells_partial_is_the_stock_finalize_then_shared_add(cell) -> None: + """Per cell shape (rows, top-k, dropped slots), every rank's partial in the entry's reference is the stock MoE's: + the finalized routed rows rounded to BF16, then ``shared.add_(routed)``. The 4-rank matrix drives the cells on the + GPU against this reference.""" + spec = cell.spec + g = torch.Generator().manual_seed(spec["m"] * 31 + spec["dropped"]) + m, topk, hidden = spec["m"], spec["topk"], 64 + fc2 = torch.randn(m * topk, hidden, generator=g).to(torch.bfloat16) + perm = torch.randperm(m * topk, generator=g).view(m, topk).to(torch.int32) + perm.view(-1)[: spec["dropped"]] = -1 + scales = torch.rand(m, topk, generator=g) + shared = torch.randn(m, hidden, generator=g).to(torch.bfloat16) + routed = finalize_partial(fc2, perm, scales, None) + assert torch.equal(finalize_partial(fc2, perm, scales, shared), shared.clone().add_(routed)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/embedding/test_modeling_v2_m3_embedding.py new file mode 100644 index 000000000000..08e4f32a7761 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/embedding/test_modeling_v2_m3_embedding.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_embedding entry of the MiniMax-M3 tp8 target: an exact gather with zero rows outside the +table.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.embedding.m3_embedding import ( # noqa: E501 + m3_embedding as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_embedding requires a CUDA device" + +VOCAB, HIDDEN = 4096, 6144 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + table = torch.randn(VOCAB, HIDDEN, generator=g, device="cuda").to(torch.bfloat16) + ids = torch.randint( + 0, VOCAB, (spec["rows"],), generator=g, device="cuda", dtype=spec["id_dtype"] + ) + if spec["pad"]: + ids[1::3] = -1 + with validating(op): + out = op(ids, table) + assert torch.equal(out, op.reference(ids, table)) + if spec["pad"]: + assert not out[1::3].any(), "padded ids must embed to zero rows" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/_m3_mxfp8_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/_m3_mxfp8_operands.py new file mode 100644 index 000000000000..e44ad56d743d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/_m3_mxfp8_operands.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 operands for the MXFP8 GEMM entries' tests: E4M3 values and 128x4-swizzled UE8M0 block scales.""" + +import torch + + +def swizzle(byte: torch.Tensor) -> torch.Tensor: + """[N, K / 32] scale bytes into the 128x4-swizzled buffer of ``mxfp8_quantize(w, True)``.""" + rows, cols = byte.shape + padded_cols, padded_rows = (cols + 3) // 4 * 4, (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + + (c // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + out = torch.zeros(padded_rows * padded_cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def mxfp8(rows: int, k: int, g: torch.Generator, lo: int = 118, hi: int = 125): + """An MXFP8 operand: E4M3 values of a unit normal, UE8M0 block scales 2^(lo-127) .. 2^(hi-128).""" + v = torch.randn(rows, k, generator=g, device="cuda").to(torch.float8_e4m3fn) + sf = torch.randint(lo, hi, (rows, k // 32), generator=g, device="cuda", dtype=torch.int32) + return v, swizzle(sf.to(torch.uint8)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_bf16_linear.py new file mode 100644 index 000000000000..254bd62efb7f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_bf16_linear.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_bf16_linear entry of the MiniMax-M3 tp8 target (the drafter's projections).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm.m3_bf16_linear import ( # noqa: E501 + m3_bf16_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_bf16_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Accumulation order (cuBLAS split-K) then one bf16 rounding on both sides: at most 2 units of ULP[bf16] of the + row, RMS well under one; measured on GB300: worst 0.5 unit.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w = (torch.randn(spec["n"], spec["k"], generator=g, device="cuda") * 0.02).to(torch.bfloat16) + with validating(op): + out = op(x, w) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(x, w), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_draft_lowm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_draft_lowm.py new file mode 100644 index 000000000000..7fb3cad9c0e5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_draft_lowm.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 gemm/m3_draft_lowm entry: the Eagle3 drafter's BF16 linears on FlashInfer's fixed-tactic low-M +kernels (direct SIMT and split-K tensor-core), at the TP8 drafter shards. + +Every cell is held to the fp64 reference within the derived band. The target's drafter linear (``_draft_linear`` in +modeling.py) is held on real BF16 weights: a covered row count is the entry's call bit for bit, an uncovered one the +in-tree BF16 GEMM's (``m3_bf16_linear``). +""" + +import zlib +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm.m3_draft_lowm import ( # noqa: E501 + TACTICS, + tactic_for, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm.m3_draft_lowm import ( # noqa: E501 + m3_draft_lowm as op, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_lowm requires a CUDA device" + +# Element band, in BF16 ulps of the row scale, and the relative RMS (``assert_within_ulp``). +# +# Kernel and reference multiply the same BF16 operands (every product exact in fp32); the kernel sums in fp32 (the +# direct kernel one dot product in K order, the split-K kernel tensor-core blocks and then its split partials), the +# reference in fp64, and both round once to BF16. With w ~ N(0, 1 / K) and x ~ N(0, 1) the outputs are O(1) and the +# fp32 order difference is ~sqrt(K) * 2^-24 of the summed magnitudes, a few 1e-6 at K = 18432: far below half a BF16 +# ulp, so the two roundings disagree by at most one ulp of the element -- up to two units of the row scale for an +# element in its row maximum's binade -- and only where the sum sits that close to a rounding boundary, a few in ten +# thousand elements, hence 0.1 ulp RMS. The same derivation as the decode GEMV's band (test_modeling_v2_m3_mx_gemv), +# with no slack under it: a second ulp would mean a wrong sum. +ELEMENT_ULP = 2.0 +RMS_ULP = 0.1 + +DRAFTER_LINEARS = { + (1280, 12288): "qkv", + (6144, 1024): "o_proj", + (4608, 6144): "gate_up", + (6144, 2304): "down", + (6144, 18432): "fc", +} + + +@pytest.fixture(autouse=True) +def _policy(): + dispatch.set_require(False) + dispatch.reset_logged() + yield + dispatch.set_require(False) + dispatch.reset_logged() + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + return x, w + + +def test_the_table_is_the_tp8_drafter() -> None: + assert set(TACTICS) == set(DRAFTER_LINEARS) + assert {c.spec["nk"] for c in op.CELLS} == set(TACTICS) + + +@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 this name + (n, k), m = cell.spec["nk"], cell.spec["m"] + x, w = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + kind, tactic = tactic_for(n, k, m) + out = torch.full((m, n), float("nan"), dtype=torch.bfloat16, device="cuda") + args = ( + x, + w, + out, + kind, + tactic, + True, + ) # the target launches with PDL when the environment enables it + certifying(op) + assert op(*args) is out + ref = op.reference(*args) + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + assert_within_ulp(out, ref, element_ulp=ELEMENT_ULP, rms_ulp=RMS_ULP) + # PDL only changes the launch attribute. + plain = torch.empty_like(out) + op(x, w, plain, kind, tactic, False) + assert torch.equal(plain.view(torch.int16), out.view(torch.int16)) + + +def _drafter(weight: torch.Tensor): + """A drafter whose every linear is ``weight`` (``_prebuild_draft_lowm`` compiles per distinct shape).""" + return SimpleNamespace(w={key: weight for key in ("qkv", "o", "gu", "dn", "fc")}) + + +@pytest.mark.parametrize("nk", sorted(DRAFTER_LINEARS), ids=lambda nk: DRAFTER_LINEARS[nk]) +def test_the_drafter_linear_is_the_entry_and_falls_back_past_its_rows(nk) -> None: + n, k = nk + _, weight = build(n, k, 1, seed=n + k) + dispatch._prebuild_draft_lowm(_drafter(weight)) + use_pdl = get_env_enable_pdl() + last = TACTICS[nk][-1][0] + dispatch.set_require(True) + for m in range(1, last + 1): + assert dispatch.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = dispatch._draft_linear(x, weight) + kind, tactic = tactic_for(n, k, m) + want = op(x, weight, torch.empty(m, n, dtype=x.dtype, device="cuda"), kind, tactic, use_pdl) + assert torch.equal(got.view(torch.int16), want.view(torch.int16)), m + # Past the table: the in-tree BF16 GEMM, even under require. + for m in (last + 1, 2 * last + 3): + assert not dispatch.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = dispatch._draft_linear(x, weight) + stock = dispatch.m3_bf16_linear(x, weight) + assert torch.equal(got.view(torch.int16), stock.view(torch.int16)), m + + +def test_a_covered_call_not_prebuilt_raises_under_require(monkeypatch) -> None: + monkeypatch.setattr(dispatch, "_DRAFT_LOWM_BUILT", set()) + x, weight = build(6144, 1024, 1, seed=5) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="not compiled before warmup"): + dispatch._draft_linear(x, weight) + # Outside require it falls back to the in-tree BF16 GEMM. + dispatch.set_require(False) + got = dispatch._draft_linear(x, weight) + stock = dispatch.m3_bf16_linear(x, weight) + assert torch.equal(got.view(torch.int16), stock.view(torch.int16)) + + +def test_guard_refuses_a_strided_activation() -> None: + x, w = build(6144, 1024, 2, seed=6) + wide = torch.zeros(2, 2048, dtype=torch.bfloat16, device="cuda") + wide[:, :1024] = x + kind, tactic = tactic_for(6144, 1024, 2) + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op( + wide[:, :1024], + w, + torch.empty(2, 6144, dtype=torch.bfloat16, device="cuda"), + kind, + tactic, + False, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_fp32_linear.py new file mode 100644 index 000000000000..4daa50048d12 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_fp32_linear.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp32_linear entry of the MiniMax-M3 tp8 target (the router beyond the decode routers' rows).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm.m3_fp32_linear import ( # noqa: E501 + m3_fp32_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp32_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """torch's default fp32 matmul precision on this hardware may be TF32 (~2^-11 relative per product): the band is + stated in fp32 ulps of each row's largest logit and covers it (K = 6144 TF32 products: ~2^13 ulps worst element); + measured on GB300 with the engine's defaults: worst element ~1700 ulps, RMS ~300.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = (torch.randn(spec["m"], spec["k"], generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(spec["e"], spec["k"], generator=g, device="cuda") * 0.02 + with validating(op): + out = op(x, w) + assert out.dtype == torch.float32 and out.shape == (spec["m"], spec["e"]) + assert_within_ulp(out, op.reference(x, w), element_ulp=8192.0, rms_ulp=2048.0) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mx_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mx_gemv.py new file mode 100644 index 000000000000..960f824f877f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mx_gemv.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv entry: the decode MXFP8 GEMV on a producer-quantized activation. + +Operands as the target holds them: the weight through mxfp8_quantize with the swizzled scale layout (the MXFP8 linear +method's storage), the activation through mxfp8_quantize with the swizzled layout (what the boundary all-reduce and +the attention merge publish). +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 + O_PROJ, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 + m3_mx_gemv as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mx_gemv requires a CUDA device" + + +# Element band, in ULP units of the row scale: see compare. +ELEMENT_ULP = 1.5 + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Element ELEMENT_ULP of the row scale, relative RMS 0.1 ulp. + + Kernel and reference multiply bit-identical dequantized operands; the kernel sums in fp32 (tensor-core blocks, + then the split-K partials), the reference in fp64. Over K <= 6144 that order difference is ~1e-6 relative, so the + two bf16 roundings disagree only where the sum sits within that of a rounding boundary: one bf16 ulp of that + element, which for an element in its row maximum's binade is up to two of ``assert_within_ulp``'s units of the row + scale, and a few in ten thousand elements in aggregate. Any fp32 sum shows it: torch's IEEE fp32 matmul of the + same dequantized operands measured at most 1.22 units against the fp64 reference over every row count 1..48 on + the tp4 target's o_proj and QKV (6 seeds each, GB300; K 2048 and 6144, the tp8 projections' K is 1024..6144), the + kernel at most 1.25, both at most 0.009 ulp RMS. The band leaves the margin between that and a full extra ulp. + """ + assert_within_ulp(out, ref, element_ulp=ELEMENT_ULP, rms_ulp=0.1) + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + return x, xq.view(m, k), xsf, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + _, xq, xsf, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(xq, xsf, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + compare(out, op.reference(xq, xsf, w, wsf)) + + +def test_an_unserved_shape_returns_none() -> None: + """Shapes outside the tuned table fall back at the call site; the op says so rather than computing.""" + _, xq, xsf, w, wsf = build(256, 512, 1, seed=1) + assert op(xq, xsf, w, wsf) is None + + +def test_guard_refuses_linear_activation_scales() -> None: + """A linear [M, K / 32] scale buffer is read as the swizzled one.""" + n, k = O_PROJ + x, _, _, w, wsf = build(n, k, 4, seed=2) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, False) + with pytest.raises(AssertionError, match="swizzled"): + with validating(op): + op(xq.view(4, k), xsf, w, wsf) + + +def test_guard_refuses_a_strided_weight() -> None: + n, k = O_PROJ + _, xq, xsf, w, wsf = build(n, k, 1, seed=3) + wide = torch.zeros(n, 2 * k, dtype=w.dtype, device="cuda") + wide[:, :k] = w + with pytest.raises(AssertionError, match="contiguous e4m3"): + with validating(op): + op(xq, xsf, wide[:, :k], wsf) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..c88f84729b45 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv_bf16 entry: the decode MXFP8 GEMV quantizing a bf16 activation in its prologue. + +The weight as the MXFP8 linear method stores it (mxfp8_quantize, swizzled scales). Besides the reference band, the +prologue's quantization is held to mxfp8_quantize itself: the GEMV on a bf16 x must equal, bit for bit, the +prequantized GEMV on mxfp8_quantize(x) -- the claim the target relies on when it mixes the two paths. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 + O_PROJ, + m3_mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm.m3_mx_gemv_bf16 import ( # noqa: E501 + m3_mx_gemv_bf16 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mx_gemv_bf16 requires a CUDA device" + +# The element band of test_modeling_v2_m3_mx_gemv.compare (derived there): the same fp32 sums against fp64. +ELEMENT_ULP = 1.5 + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + return x, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + x, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(x, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + # The band of test_modeling_v2_m3_mx_gemv.compare: the quantized operands are the reference's exactly (checked + # below), the sums differ only in order. + assert_within_ulp(out, op.reference(x, w, wsf), element_ulp=ELEMENT_ULP, rms_ulp=0.1) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + pre = m3_mx_gemv(xq.view(m, k), xsf, w, wsf, use_pdl=True) + assert pre is not None and torch.equal(out.view(torch.int16), pre.view(torch.int16)), ( + "the prologue's MXFP8 quantization differs from mxfp8_quantize" + ) + + +def test_a_strided_activation_is_declined() -> None: + """The prologue reads x densely; the op declines a strided x (None) rather than computing on it (a served shape: + the same x, contiguous, is computed).""" + n, k = O_PROJ + x, w, wsf = build(n, k, 2, seed=1) + wide = torch.zeros(2, 2 * k, dtype=torch.bfloat16, device="cuda") + wide[:, :k] = x + assert op(wide[:, :k], w, wsf) is None + assert op(wide[:, :k].contiguous(), w, wsf) is not None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py new file mode 100644 index 000000000000..9ba7fa37d475 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_gemm entry of the MiniMax-M3 tp8 target (the decode-graph GEMM on a prequantized +activation).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm.m3_mxfp8_gemm import ( # noqa: E501 + m3_mxfp8_gemm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_gemm requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Same E4M3 operands and power-of-two scales on both sides; accumulation order then one bf16 rounding: at most + 2 units of ULP[bf16] of the row (the low-M GEMM's band), RMS well under one.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + act, act_sf = mxfp8(spec["m"], spec["k"], g) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + with validating(op): + out = op(act, act_sf, w, w_sf, torch.bfloat16, False) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp( + out, + op.reference(act, act_sf, w, w_sf, torch.bfloat16, False), + element_ulp=2.0, + rms_ulp=0.25, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py new file mode 100644 index 000000000000..dd7ede34f0b8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_linear entry of the MiniMax-M3 tp8 target: MXFP8 quantize + GEMM on each backend the +engine selects (native CUTLASS, FlashInfer, the per-bucket tuned quantize + GEMM).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm.m3_mxfp8_linear import ( # noqa: E501 + m3_mxfp8_linear as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_linear requires a CUDA device" + +_FLAGS = { + "native": (False, False, False, False), + "flashinfer": (False, False, True, False), + "tuned": (True, False, False, False), +} + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """The activation's MXFP8 round trip is exact on both sides (same block rule); then the accumulation order and + one bf16 rounding: at most 2 units of ULP[bf16] of the row, RMS well under one.""" + spec = cell.spec + if ( + spec["path"] == "flashinfer" + and getattr(torch.ops.trtllm, "flashinfer_mm_mxfp8", None) is None + ): + pytest.skip("this build has no trtllm::flashinfer_mm_mxfp8") + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + alpha = torch.ones(1, dtype=torch.float32, device="cuda") + args = (x, w, w_sf, alpha, torch.bfloat16, *_FLAGS[spec["path"]]) + with validating(op): + out = op(*args) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/_m3_moe_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/_m3_moe_operands.py new file mode 100644 index 000000000000..80ef9fbd7225 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/_m3_moe_operands.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One tp8 rank's MiniMax-M3 MoE layer operands for the decode-engine entries' tests. + +The routed experts are random checkpoint-format NVFP4 tensors ([w3 (up); w1 (gate)] rows, e4m3 block scales) put +through TRT-LLM's own TRTLLM-Gen NVFP4 loader shuffle (``NVFP4TRTLLMGenFusedMoEMethod._shuffle_all_experts``), which +writes the buffers the engines read in place. The per-expert scalars follow the loader's conventions (act_beta and +act_clamp divided by fc31_alpha, fc31_scale_c = fc2_input_scale x fc31_alpha) and are sized so the FC1 sums land +around the SwiGLU-OAI clamp and the outputs near unit scale. The shared expert's MXFP8 weights come from +mxfp8_quantize with the swizzled scale layout, as the MXFP8 linear method stores them. +""" + +import functools +from types import SimpleNamespace + +import torch + +# Hidden size, the routed experts' intermediate per rank at TP8 (1536 / 8 x 2: every expert on every rank), experts, +# top-k, and the shared expert's intermediate per rank. +H, I_TP, NUM_EXPERTS, TOP_K = 6144, 384, 128, 4 +SHARED_I = 384 +SWIGLU_ALPHA, SWIGLU_LIMIT = 1.702, 7.0 +ROUTED_SCALING = 2.0 +MINIMAX2_ROUTING = 5 # RoutingMethodType.MiniMax2 +# NVFP4 global scales: 448 x 6 over the amax each side is sized for. +FC31_INPUT_SCALE = 448.0 * 6.0 / 8.0 +FC2_INPUT_SCALE = 448.0 * 6.0 / 24.0 + + +def _rand_e4m3_scales(shape, gen) -> torch.Tensor: + """E4M3 block scales in [0.5, 2): biased exponents 6 and 7, any mantissa.""" + return torch.randint(0x30, 0x40, shape, dtype=torch.uint8, device="cuda", generator=gen).view( + torch.float8_e4m3fn + ) + + +@functools.lru_cache(maxsize=None) +def experts(seed: int = 20261013): + """The routed experts and the router of one MoE layer, in the layout and conventions of the TRTLLM-Gen loader. + + Returns a SimpleNamespace: ``tensors`` (the ten expert operands in ``M3MoeState.layer``'s order), ``bias``, + ``router_weight``, ``fc31_input_scale`` and ``shared`` (gate_up, its scales, down, its scales). + """ + from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4TRTLLMGenFusedMoEMethod + + gen = torch.Generator(device="cuda").manual_seed(seed) + e = NUM_EXPERTS + w31 = torch.randint( + 0, 256, (e, 2 * I_TP, H // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + w2 = torch.randint(0, 256, (e, H, I_TP // 2), dtype=torch.uint8, device="cuda", generator=gen) + w31s = _rand_e4m3_scales((e, 2 * I_TP, H // 16), gen) + w2s = _rand_e4m3_scales((e, H, I_TP // 16), gen) + module = SimpleNamespace( + w3_w1_weight=w31, + w2_weight=w2, + w3_w1_weight_scale=w31s, + w2_weight_scale=w2s, + is_gated_activation=True, + bias=False, + ) + NVFP4TRTLLMGenFusedMoEMethod()._shuffle_all_experts(module) + torch.cuda.synchronize() + + # Raw FC1 sums: rms(x) x FC31_INPUT_SCALE x rms(E2M1 value) x rms(block scale) x sqrt(H) ~ 1.0e5 at rms(x) = 1; + # fc31_alpha brings them to ~1.5 in the dequantized domain, the FC2 alpha the outputs to ~1. + base = 1.5 / (FC31_INPUT_SCALE * 2.93 * 1.3 * H**0.5) + fc31_alpha = base * (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) + fc2_alpha = (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) / ( + FC2_INPUT_SCALE * 2.93 * 1.3 * I_TP**0.5 + ) + tensors = ( + w31, + w31s, + w2, + w2s, + fc31_alpha, + FC2_INPUT_SCALE * fc31_alpha, # fc31_scale_c + torch.full((e,), SWIGLU_ALPHA, device="cuda"), # act_alpha + 1.0 / fc31_alpha, # act_beta: the (up + 1) offset over fc31_alpha + SWIGLU_LIMIT / fc31_alpha, # act_clamp over fc31_alpha + fc2_alpha, + ) + bias = (0.02 * torch.randn(e, device="cuda", generator=gen)).float() + router_weight = (torch.randn(e, H, device="cuda", generator=gen) / H**0.5).float().contiguous() + fc31_input_scale = torch.tensor([FC31_INPUT_SCALE], dtype=torch.float32, device="cuda") + + gu = (0.02 * torch.randn(2 * SHARED_I, H, device="cuda", generator=gen)).to(torch.bfloat16) + dn = (0.05 * torch.randn(H, SHARED_I, device="cuda", generator=gen)).to(torch.bfloat16) + gu_q, gu_sf = torch.ops.trtllm.mxfp8_quantize(gu, True) + dn_q, dn_sf = torch.ops.trtllm.mxfp8_quantize(dn, True) + shared = (gu_q.view(2 * SHARED_I, H), gu_sf, dn_q.view(H, SHARED_I), dn_sf) + return SimpleNamespace( + tensors=tensors, + bias=bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared=shared, + ) + + +# Routing cases: random (up to 4 M distinct experts), and every token on the same four experts (the bias outweighs +# any sigmoid difference), which leaves later slot groups empty. +CONCENTRATED = (5, 40, 77, 126) + + +def layer_for(state, case: str, shared: bool = True): + """One layer on ``state`` over :func:`experts` for a routing case. ``shared=False``: the shared expert's weights + are zero, so the engine returns the routed part alone (bf16(0 + bf16(routed))).""" + ops = experts() + bias = ops.bias.clone() + if case == "concentrated": + bias[list(CONCENTRATED)] += 1.0 + shared_weights = ops.shared + if not shared: + gu, gu_sf, dn, dn_sf = ops.shared + shared_weights = (torch.zeros_like(gu), gu_sf, torch.zeros_like(dn), dn_sf) + return state.layer( + *ops.tensors, + bias, + router_weight=ops.router_weight, + fc31_input_scale=ops.fc31_input_scale, + shared_weights=shared_weights, + ) + + +def stock_routed( + hidden: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor +) -> torch.Tensor: + """The stock TRTLLM-Gen NVFP4 MoE (fp4_quantize, then fp4_block_scale_moe_runner with SwiGLU-OAI) on the given + routing, over :func:`experts`: what the built-in MiniMaxM3MoE computes for the routed part.""" + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers the stock ops + + ops = experts() + w31, w31s, w2, w2s, fc31_alpha, scale_c, act_alpha, act_beta, act_clamp, fc2_alpha = ops.tensors + xq, xsf = torch.ops.trtllm.fp4_quantize(hidden, ops.fc31_input_scale, 16, False, False) + out = torch.ops.trtllm.fp4_block_scale_moe_runner( + None, None, xq, xsf.view(torch.float8_e4m3fn).flatten(), + w31, w31s, None, act_alpha, act_beta, act_clamp, w2, w2s, None, + scale_c, fc31_alpha, fc2_alpha, + NUM_EXPERTS, TOP_K, None, None, I_TP, 0, NUM_EXPERTS, ROUTED_SCALING, MINIMAX2_ROUTING, True, 0, + topk_weights=topk_weights.to(torch.bfloat16), topk_ids=topk_ids.to(torch.int32), + ) # fmt: skip + return out[0] if isinstance(out, (list, tuple)) else out + + +def hidden_states(m: int, seed: int) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(m, H, device="cuda", generator=gen).to(torch.bfloat16) + + +def check_routed_part(op, state, m: int, seed: int) -> None: + """Exact, independent of the reference: the engine's routing (its route_dbg record: ids, then the bf16 weights' + fp32 bits) is the reference routing bit for bit, and its routed part is the stock TRTLLM-Gen MoE's output on that + routing bit for bit -- the claim the engines are built on.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + moe_route, + ) + + layer = layer_for(state, "random", shared=False) + hidden = hidden_states(m, seed=seed) + out = op(hidden, layer) + ids = state.route_dbg[:, :4].long() + weights = state.route_dbg[:, 4:].contiguous().view(torch.float32) + ref_ids, ref_weights = moe_route(hidden.float(), layer.router_weight, layer.bias) + assert torch.equal(ids, ref_ids) and torch.equal(weights, ref_weights), ( + "routing differs from the reference" + ) + stock = stock_routed(hidden, ids, weights) + assert torch.equal(out.view(torch.int16), stock.view(torch.int16)), ( + "routed part differs from TRTLLM-Gen's" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/_m3_router.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/_m3_router.py new file mode 100644 index 000000000000..f5d91aef73a7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/_m3_router.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared operands and gates of the router entries' tests (m3_router_gemv, m3_router_gemv_fp4, m3_router_tc: one +kernel source).""" + +import zlib + +import torch + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp # noqa: E402 — needs the path declared above + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.moe.m3_router_gemv import ( # noqa: E501 + m3_router_gemv as gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.moe.m3_router_gemv import ( # noqa: E501 + near_tie, + nvfp4_quantize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.moe.m3_router_gemv_fp4 import ( # noqa: E501 + m3_router_gemv_fp4 as gemv_fp4, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.moe.m3_router_tc import ( # noqa: E501 + m3_router_tc as tc, +) + +__all__ = [ + "SF_SCALE", + "_build", + "_compare_fp4", + "_compare_logits", + "_seed", + "gemv", + "gemv_fp4", + "tc", +] + +#: The experts' NVFP4 input scale the tc cells use: 448 * 6 / 8, block scales well inside E4M3. +SF_SCALE = 336.0 + + +def _build(m: int, e: int, k: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(m, k, generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(e, k, generator=g, device="cuda") * 0.02 + return x, w + + +def _seed(cell) -> int: + return zlib.crc32(cell.why.encode()) % 2**31 + + +def _compare_logits(out: torch.Tensor, ref: torch.Tensor) -> None: + """fp32 logits that differ from a true fp32 product only in summation order. + + Over K = 6144 products a reordered fp32 sum moves by a few ulps of the sum of magnitudes, which near + cancellation is ~sqrt(K) times the row's largest logit's ulp at most. Gated in fp32 ulps of each row's largest + logit: 64 element-wise, 8 in relative RMS. Measured on GB300: GEMV 1.6 / 1.4 ulp, tensor-core router 3.8 / 2.0 + ulp at 2..48 tokens; the tensor-core logits against a true fp32 product (the TF32 rounding the band must see) + sit at 2200+ / 1700 ulp. + """ + assert out.dtype == ref.dtype == torch.float32 and out.shape == ref.shape + assert_within_ulp(out, ref, element_ulp=64.0, rms_ulp=8.0) + + +def _codes(packed: torch.Tensor) -> torch.Tensor: + """[M, K / 2] packed e2m1 bytes -> [M, K] 4-bit codes, low nibble first.""" + lo, hi = packed & 0xF, packed >> 4 + return torch.stack((lo, hi), dim=-1).reshape(packed.shape[0], -1) + + +def _compare_fp4( + q: torch.Tensor, sf: torch.Tensor, x: torch.Tensor, sf_scale: torch.Tensor +) -> None: + """Scales bit-exact; codes bit-exact except within ``TIE_WINDOW`` of an e2m1 midpoint, where the kernel's + approximate reciprocals choose the side (and then only by one code).""" + ref_q, ref_sf, scaled = nvfp4_quantize(x, sf_scale) + assert q.dtype == torch.uint8 and q.shape == ref_q.shape + assert torch.equal(sf, ref_sf), f"{(sf != ref_sf).sum().item()} block scales differ" + got, want = _codes(q), _codes(ref_q) + differ = got != want + assert not bool((differ & ~near_tie(scaled)).any()), ( + f"{(differ & ~near_tie(scaled)).sum().item()} e2m1 codes differ away from a rounding tie" + ) + magnitude = ((got & 7).int() - (want & 7).int()).abs() + assert bool((magnitude[differ] <= 1).all()), "a tie resolved more than one code away" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_moe_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_moe_engine.py new file mode 100644 index 000000000000..87655f1fc9ac --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_moe_engine.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the moe/m3_moe_engine entry: the 1..8-token MoE engine on its caller-owned state. + +One GB300 rank of the tp8 deployment: 128 experts (all local, routed and shared intermediate 384 per rank), top 4, +hidden 6144, the shared expert folded in, ``num_ctas`` = SMs - 24 as the target builds it. Operands: +``_m3_moe_operands``. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + HIDDEN, + INTER, + SHARED_INTER, + M3MoeState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_moe_operands import ( # noqa: E402 + I_TP, + SHARED_I, + H, + check_routed_part, + hidden_states, + layer_for, +) +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" + +_STATES = {} + + +def test_the_operands_are_the_tp8_engine_geometry() -> None: + assert (H, I_TP, SHARED_I) == (HIDDEN, INTER, SHARED_INTER) == (6144, 384, 384) + + +def test_the_engine_refuses_the_tp4_width() -> None: + with pytest.raises(ValueError, match="384"): + M3MoeState(torch.device("cuda"), num_tokens=1, inter=768) + + +def _state(m: int) -> M3MoeState: + if m not in _STATES: + dev = torch.device("cuda") + sms = torch.cuda.get_device_properties(dev).multi_processor_count + _STATES[m] = M3MoeState(dev, num_tokens=m, num_ctas=sms - 24, shared=True) + return _STATES[m] + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """The engine against the plain-torch MoE: BAND_ELEMENT bf16 ulp of the row's largest magnitude per element, + BAND_RMS ulp relative RMS. + + Both sides quantize the same operands with the same recipes (the activation's and the intermediate's NVFP4 codes + and block scales, the shared expert's MXFP8); they differ in summation order (tensor-core fp32 against fp64) and in + the PTX approximations (rcp / ex2 / tanh). Such a difference changes the output only where it moves a value across + a rounding boundary: one bf16 ulp of a slot output or of the sum, or, rarely, one intermediate NVFP4 code, which + reaches the output as one weight column's share. Measured over every cell, both routing cases and 3 seeds each + on the tp4 engine (GB300, intermediate 768): at most 1.48 ulp element, 0.18 ulp RMS; the tp8 engine runs the + same recipes over half the intermediate. The band leaves ~3x on both; a misread recipe (the activation's + 1 / (sf x (1 / sc)) spelling of the reciprocal on ~0.06% of its codes) measured 12 ulp element, 7 ulp RMS. + """ + assert_within_ulp(out, ref, element_ulp=BAND_ELEMENT, rms_ulp=BAND_RMS) + + +BAND_ELEMENT = 4.0 +BAND_RMS = 0.5 + + +@pytest.mark.parametrize("case", ("random", "concentrated")) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, case) -> None: + m = cell.spec["tokens"] + layer = layer_for(_state(m), case) + hidden = hidden_states(m, seed=zlib.crc32((cell.why + case).encode()) % 2**31) + with validating(op): + out = op(hidden, layer) + again = op(hidden, layer) + assert out.shape == (m, H) and out.dtype is torch.bfloat16 + # The workspace is tagged per call and never reset: a repeated call is the same call. + assert torch.equal(out.view(torch.int16), again.view(torch.int16)), ( + "a repeated call changed the output" + ) + compare(out, op.reference(hidden, layer)) + + +@pytest.mark.parametrize("m", (1, 2, 8)) +def test_routed_part_is_trtllm_gen_bit_for_bit(m) -> None: + check_routed_part(op, _state(m), m, seed=100 + m) + + +def test_layers_on_one_state_do_not_leak_into_each_other() -> None: + """Two layers share one state's workspace, in stream order: A, B, A gives A's output twice, bit for bit, and the + caller's ``out`` receives the same rows as a fresh tensor.""" + state = _state(4) + a, b = layer_for(state, "random"), layer_for(state, "concentrated") + hidden = hidden_states(4, seed=7) + first = op(hidden, a) + op(hidden, b) + out = torch.empty(4, H, dtype=torch.bfloat16, device="cuda") + second = op(hidden, a, out=out) + assert second.data_ptr() == out.data_ptr() + assert torch.equal(first.view(torch.int16), out.view(torch.int16)) + + +@pytest.mark.parametrize( + "bad", + ("misaligned hidden", "fp32 out", "short out", "long out", "int32 arp"), +) +def test_malformed_buffers_are_rejected_by_the_call(bad) -> None: + """Why the entry has no is_valid: the compiled call checks each buffer against its signature and raises.""" + state = _state(2) + layer = layer_for(state, "random") + hidden, kw = hidden_states(2, seed=3), {} + if bad == "misaligned hidden": + hidden = torch.randn(2 * H + 8, device="cuda").to(torch.bfloat16)[2 : 2 + 2 * H].view(2, H) + elif bad == "fp32 out": + kw["out"] = torch.empty(2, H, dtype=torch.float32, device="cuda") + elif bad == "short out": + kw["out"] = torch.empty(1, H, dtype=torch.bfloat16, device="cuda") + elif bad == "long out": + kw["out"] = torch.empty(3, H, dtype=torch.bfloat16, device="cuda") + else: + kw["arp"] = torch.zeros(8, dtype=torch.int32, device="cuda") + with pytest.raises(ValueError): + op(hidden, layer, **kw) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv.py new file mode 100644 index 000000000000..2e923b4ec114 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_router_gemv entry of the MiniMax-M3 tp8 target.""" + +import pytest +import torch + +__extra_import_path__ = ["../../../..", "."] +from _m3_router import ( # noqa: E402, F401 + SF_SCALE, + _build, + _compare_fp4, + _compare_logits, + _seed, + gemv, + gemv_fp4, + tc, +) +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "the router entries require a CUDA device" + + +@pytest.mark.parametrize("cell", gemv.CELLS, ids=[c.why[:48] for c in gemv.CELLS]) +def test_gemv_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], spec["e"], spec["k"], _seed(cell)) + with validating(gemv): + out = gemv(x, w, 1, False) + _compare_logits(out, gemv.reference(x, w, 1, False)) + + +def test_the_gemv_is_a_true_fp32_product() -> None: + """And the GEMV, by contrast, is no TF32 product: the TF32 reference rejects it.""" + x, w = _build(1, 128, 6144, 6) + out = gemv(x, w, 1, False) + with pytest.raises(AssertionError): + _compare_logits(out, tc.reference(x, w, None, False)[0]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv_fp4.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv_fp4.py new file mode 100644 index 000000000000..11e9982c8dcc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv_fp4.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_router_gemv_fp4 entry of the MiniMax-M3 tp8 target.""" + +import pytest +import torch + +__extra_import_path__ = ["../../../..", "."] +from _m3_router import ( # noqa: E402, F401 + SF_SCALE, + _build, + _compare_fp4, + _compare_logits, + _seed, + gemv, + gemv_fp4, + tc, +) +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "the router entries require a CUDA device" + + +@pytest.mark.parametrize("cell", gemv_fp4.CELLS, ids=[c.why[:48] for c in gemv_fp4.CELLS]) +def test_gemv_fp4_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], spec["e"], spec["k"], _seed(cell)) + sf_scale = torch.tensor([spec["sf_scale"]], dtype=torch.float32, device="cuda") + with validating(gemv_fp4): + logits, q, sf = gemv_fp4(x, w, sf_scale, False) + ref_logits, _, _ = gemv_fp4.reference(x, w, sf_scale, False) + _compare_logits(logits, ref_logits) + _compare_fp4(q, sf, x, sf_scale) + + +def test_guard_refuses_a_zero_global_scale() -> None: + x, w = _build(1, 128, 6144, 4) + with pytest.raises(AssertionError, match="positive and finite"): + with validating(gemv_fp4): + gemv_fp4(x, w, torch.zeros(1, dtype=torch.float32, device="cuda"), False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_tc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_tc.py new file mode 100644 index 000000000000..0b3638ee02eb --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_tc.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_router_tc entry of the MiniMax-M3 tp8 target.""" + +import pytest +import torch + +__extra_import_path__ = ["../../../..", "."] +from _m3_router import ( # noqa: E402, F401 + SF_SCALE, + _build, + _compare_fp4, + _compare_logits, + _seed, + gemv, + gemv_fp4, + tc, +) +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "the router entries require a CUDA device" + + +@pytest.mark.parametrize("cell", tc.CELLS, ids=[c.why[:48] for c in tc.CELLS]) +def test_tc_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], 128, spec["k"], _seed(cell)) + sf_scale = torch.tensor([SF_SCALE], dtype=torch.float32, device="cuda") if spec["fp4"] else None + with validating(tc): + outs = tc(x, w, sf_scale, False) + ref = tc.reference(x, w, sf_scale, False) + assert len(outs) == len(ref) == (3 if spec["fp4"] else 1) + _compare_logits(outs[0], ref[0]) + if spec["fp4"]: + _compare_fp4(outs[1], outs[2], x, sf_scale) + + +def test_tc_reference_is_tf32_and_the_gate_sees_it() -> None: + """The tensor-core router rounds the weight to TF32: against a true fp32 product its logits sit far outside + the band, so the TF32 rounding in ``reference`` is load-bearing rather than decoration.""" + x, w = _build(16, 128, 6144, 5) + out = tc(x, w, None, False)[0] + with pytest.raises(AssertionError): + _compare_logits(out, gemv.reference(x, w, 1, False)) + + +def test_guard_refuses_a_misaligned_activation() -> None: + """A row view starting 8 bytes into the buffer keeps a legal row stride but not a legal bulk-copy source.""" + x, w = _build(2, 128, 6144, 3) + flat = torch.empty(2 * 6144 + 4, dtype=torch.bfloat16, device="cuda") + shifted = flat[4:].view(2, 6144) + shifted.copy_(x) + with pytest.raises(AssertionError, match="16-byte aligned"): + with validating(tc): + tc(shifted, w, None, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py new file mode 100644 index 000000000000..d7a3f6fef547 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the norm/m3_embedding_firstnorm entry: the decode embedding gather + the first layer's RMSNorm. + +The table is the model's full bf16 [200064, 6144] embedding shape (random values at the checkpoint's scale), eps the +checkpoint's 1e-6. +""" + +import functools +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + VOCAB, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + m3_embedding_firstnorm as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_embedding_firstnorm requires a CUDA device" + +H = 6144 +EPS = 1e-6 + + +@functools.lru_cache(maxsize=None) +def _operands(): + g = torch.Generator(device="cuda").manual_seed(20261013) + table = (0.02 * torch.randn(VOCAB, H, generator=g, device="cuda")).to(torch.bfloat16) + weight = (1.0 + 0.1 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + return table, weight + + +def _buffers(rows): + return (torch.empty(rows, H, dtype=torch.bfloat16, device="cuda"), + torch.empty(rows, H, dtype=torch.bfloat16, device="cuda")) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + rows, id_dtype = cell.spec["rows"], cell.spec["id_dtype"] + table, weight = _operands() + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + ids = torch.randint(0, VOCAB, (rows,), generator=g, device="cuda").to(id_dtype) + raw, out = _buffers(rows) + with validating(op): + normed, residual = op(ids, table, weight, EPS, raw, out) + assert normed.data_ptr() == out.data_ptr() and residual.data_ptr() == raw.data_ptr() + ref_normed, ref_raw = op.reference(ids, table, weight, EPS, raw, out) + # The gather is a copy: exact. + assert torch.equal(raw.view(torch.int16), ref_raw.view(torch.int16)) + # The norm: the kernel's fp32 sum of squares (warp butterflies, then across 4 warps) and fast rsqrt against an + # fp64 reference, then x * rstd * w rounded once to bf16. The fp32 side is ~1e-6 relative off, so the two roundings + # disagree only next to a bf16 rounding boundary: one ulp at most per element, a small fraction in aggregate. + # Measured over these cells (3 seeds each, GB300): at most 0.51 ulp element, 0.026 ulp RMS. + assert_within_ulp(out, ref_normed, element_ulp=1.0, rms_ulp=0.1) + + +def test_ids_outside_the_table_give_zero_rows() -> None: + """Documented, not refused: an id outside [0, V) gathers a zero row whose norm is zero.""" + table, weight = _operands() + ids = torch.tensor([-1, 5, VOCAB], dtype=torch.int64, device="cuda") + raw, out = _buffers(3) + op(ids, table, weight, EPS, raw, out) + assert not raw[0].any() and not raw[2].any() and not out[0].any() and not out[2].any() + assert torch.equal(raw[1], table[5]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/_m3_spec_inputs.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/_m3_spec_inputs.py new file mode 100644 index 000000000000..42fdec65c754 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/_m3_spec_inputs.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Inputs shared by the tp8 target's spec entry tests: TP8 dimensions, single-device stand-ins for the cross-rank +exchanges, and the comparison bands.""" + +from __future__ import annotations + +import zlib + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + make_keys, +) + +#: MiniMax-M3 at TP8: hidden size, vocabulary, the target head's shard (25008 rows) and RMSNorm epsilon. +HIDDEN = 6144 +VOCAB = 200064 +WORLD = 8 +TARGET_SHARD = VOCAB // WORLD +EPS = 1e-6 + + +def seed_of(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def num_sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def randn(shape, g: torch.Generator, std: float = 1.0, dtype=torch.bfloat16) -> torch.Tensor: + return (torch.randn(shape, generator=g, device="cuda") * std).to(dtype) + + +def head_workspace(world: int, call: int = 0, slot_rows: int = SHARD_ROWS): + """(ws_uc int64 [2 * slot_rows * world], keys int64 [slot_rows], counters int32 [2]) of a fresh call ``call``.""" + uc = torch.zeros(2 * slot_rows * world, dtype=torch.int64, device="cuda") + keys = torch.zeros(slot_rows, dtype=torch.int64, device="cuda") + counters = torch.tensor([0, call], dtype=torch.int32, device="cuda") + return uc, keys, counters + + +def plant_peers( + uc: torch.Tensor, + counters: torch.Tensor, + own_best: torch.Tensor, + rank: int, + world: int, + shard: int, + g: torch.Generator, + slot_rows: int = SHARD_ROWS, + ties: bool = False, +) -> None: + """Write the other ranks' keys for each row into this rank's slots, as their multicast pushes would. + + Each peer's maximum is the own row maximum shifted by +-0.25 (well clear of the accumulation-order band), so + some rows are won by a peer and some by this rank; with ``ties`` a peer on an even row reports exactly the own + maximum, and the lower global index must win. + """ + if world == 1: + return + m = own_best.shape[0] + half = int(counters[1].item()) & 1 + for r in range(world): + if r == rank: + continue + sign = torch.where(torch.rand(m, generator=g, device="cuda") < 0.5, -1.0, 1.0) + value = own_best.float() + 0.25 * sign + if ties: + value = torch.where(torch.arange(m, device="cuda") % 2 == 0, own_best.float(), value) + index = r * shard + torch.randint(0, shard, (m,), generator=g, device="cuda") + keys = make_keys(value, index) + for t in range(m): + uc[(half * slot_rows + t) * world + r] = keys[t] + + +def assert_bf16_close(out: torch.Tensor, ref: torch.Tensor, max_frac: float = 1e-3) -> None: + """At most one bf16 ulp per element, and only on a ``max_frac`` share of the elements. + + Kernel and reference compute the same fp32 expression and round once; they differ only where a different + summation order (or rsqrtf) puts the fp32 value on the other side of a bf16 rounding boundary. + """ + assert out.shape == ref.shape and out.dtype == ref.dtype == torch.bfloat16 + o, r = out.float(), ref.float() + ulp = torch.ldexp(torch.ones_like(r), torch.frexp(r).exponent - 8) + ulp = torch.where(r == 0, torch.full_like(r, 2.0**-133), ulp) + diff = (o - r).abs() + assert bool((diff <= ulp * 1.0001).all()), f"max |diff| / ulp = {(diff / ulp).max().item():.2f}" + frac = (diff > 0).float().mean().item() + assert frac <= max_frac, f"{frac:.2e} of the elements moved by one ulp (allowed {max_frac})" + + +def craft_drafts(tokens: torch.Tensor, nreq: int, d: int) -> torch.Tensor: + """Drafts matching request b's verify tokens for its first (b + 2) % (d + 1) positions and differing from + there on, so the acceptance counts vary across requests.""" + t = tokens.view(nreq, d + 1)[:, :d].clone() + for b in range(nreq): + n = (b + 2) % (d + 1) + t[b, n:] += 1 + return t.contiguous() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_capture_norm.py new file mode 100644 index 000000000000..8720f10a3168 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_capture_norm.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_capture_norm entry (Eagle3 capture with fc_norm folded in).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_capture_norm import ( # noqa: E501 + m3_capture_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_capture_norm requires a CUDA device" + +LAYERS = 3 + + +def _case(spec, seed): + g = gen(seed) + m, slot = spec["m"], spec["slot"] + h = randn((m, HIDDEN), g, 2.0) + r = randn((m, HIDDEN), g, 2.0) if spec["residual"] else None + w = (1.0 + randn((HIDDEN,), g, 0.1).float()).to(torch.bfloat16) + buf = torch.full((m, LAYERS * HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda") + dst = buf[:, slot * HIDDEN : (slot + 1) * HIDDEN] + return (h, r, w, EPS, dst), buf + + +@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 this name + """Within one bf16 ulp per element on at most 0.1% of the elements; the other layer slots of the buffer + untouched.""" + args, buf = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + slot = cell.spec["slot"] + others = torch.cat([buf[:, : slot * HIDDEN], buf[:, (slot + 1) * HIDDEN :]], dim=1) + assert bool((others == 7.0).all()) + + +def test_guard_refuses_a_strided_weight() -> None: + (h, r, w, eps, dst), _ = _case(op.CELLS[0].spec, 1) + w2 = torch.stack([w, w], dim=1)[:, 0] + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(h, r, w2, eps, dst) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_chunk_norm.py new file mode 100644 index 000000000000..9a0fad240f3e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_chunk_norm.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_chunk_norm entry (the Eagle3 fc_norm over the captured layers).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_chunk_norm import ( # noqa: E501 + m3_chunk_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_chunk_norm requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, c, h = spec["m"], spec["c"], spec["h"] + buf = randn((m, c * h + spec["extra_cols"]), g, 2.0) + src = buf[:, : c * h + spec["extra_cols"]] + w = (1.0 + randn((c, h), g, 0.1).float()).to(torch.bfloat16) + dst = torch.full((m, c * h), 7.0, dtype=torch.bfloat16, device="cuda") + return (src, w, EPS, dst, True) + + +@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 this name + """Within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic, rsqrtf).""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + + +def test_guard_refuses_a_strided_weight() -> None: + src, w, eps, dst, pdl = _case(op.CELLS[0].spec, 1) + wt = torch.empty((w.shape[1], w.shape[0]), dtype=w.dtype, device="cuda").t() + wt.copy_(w) + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(src, wt, eps, dst, pdl) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_input.py new file mode 100644 index 000000000000..7476e6f10e5e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_input.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_input entry (embedding + both RMSNorms + concat in one launch), against its reference and +against the in-tree drafter's modules.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_draft_input import ( # noqa: E501 + m3_draft_input as op, +) +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.rms_norm import RMSNorm + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + VOCAB, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_input requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, h = spec["m"], spec["h"] + table = randn((VOCAB, h), g) + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + if spec["invalid_id"]: + ids[0] = -1 + hidden = randn((4 * m if spec["gather"] else m, h), g) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 2 if spec["gather"] else None + w_e = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + w_h = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + out = torch.full((m, 2 * h), 7.0, dtype=torch.bfloat16, device="cuda") + residual = torch.full((m, h), 7.0, dtype=torch.bfloat16, device="cuda") + return (ids, table, hidden, rows, w_e, w_h, EPS, EPS, out, residual, True) + + +@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 this name + """Normalized halves within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic in a + different summation order, rsqrtf); the gathered residual and the zero embedding of an invalid id exact.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out, residual = op(*args) + ref_out, ref_res = op.reference(*args) + torch.cuda.synchronize() + assert_bf16_close(out, ref_out) + assert torch.equal(residual, ref_res) + if cell.spec["invalid_id"]: + assert not bool(out[0, : cell.spec["h"]].float().any()) + + +def test_guard_refuses_rows_past_hidden() -> None: + gather = next(c for c in op.CELLS if c.spec["gather"]) + args = list(_case(gather.spec, 1)) + args[3] = args[3] + 100 + with pytest.raises(AssertionError, match="rows must index hidden"): + with validating(op): + op(*args) + + +@pytest.mark.parametrize("m,gather", [(4, False), (32, False), (8, True)], ids=str) +def test_the_in_tree_draft_input(m, gather) -> None: + """Against the in-tree drafter's own modules: the Embedding lookup, then the layer's input_layernorm and + hidden_norm (in-tree RMSNorm) and the concat of Eagle3DecoderLayer.forward. Within one bf16 ulp on at most 0.1% of + the elements (the RMS statistic summed in another order); the residual (the carried hidden rows) exact.""" + g = gen(seed_of(f"in-tree draft input {m} {gather}")) + h = HIDDEN + embed = Embedding(VOCAB, h, dtype=torch.bfloat16).cuda() + embed.weight.data = randn((VOCAB, h), g) + norms = [] + for _ in range(2): + norm = RMSNorm(hidden_size=h, eps=EPS, dtype=torch.bfloat16).cuda() + norm.weight.data = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + norms.append(norm) + input_layernorm, hidden_norm = norms + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + hidden = randn((4 * m if gather else m, h), g, 2.0) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 3 if gather else None + carried = hidden[rows] if gather else hidden + stock = torch.cat([input_layernorm(embed(ids)), hidden_norm(carried)], dim=-1) + out, residual = m3_spec_tail.draft_input( + ids, embed.weight, hidden, rows, input_layernorm.weight, EPS, hidden_norm.weight, EPS + ) + torch.cuda.synchronize() + assert_bf16_close(out, stock) + assert torch.equal(residual, carried) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_step.py new file mode 100644 index 000000000000..6cf9240e71a0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_step.py @@ -0,0 +1,615 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_step entry (one fused greedy draft pass: the BF16 draft head streamed as BF16 fragments, +the cross-rank argmax over the eight TP8 ranks and the draft loop's bookkeeping), and its tokens against the in-tree +drafter's head and greedy draft sampling. + +The in-tree path: each rank's LM head shard GEMM (BF16 out), ``LogitsProcessor``'s float cast, then either the argmax +of the gathered logits (``torch.argmax``) or ``greedy_sample_draft_with_tp_gather``'s per-rank ``torch.max`` and the +pick among the ranks' (index, value) pairs. The edge cases use integer data, whose fp32 sums are exact in any order, +so the kernel and the in-tree path see the same BF16 logits and must agree bit for bit. The tp8 target never builds a +coded copy of the head: the op's coded stream (``exc_off`` given) is not certified here. +""" + +import warnings +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops import ( # noqa: E501 + m3_draft_step as draft_step_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_draft_step import ( # noqa: E501 + DRAFT_LEN, + DRAFT_SHARD, + HIDDEN, + MAX_TOKENS, + bf16_fragments, + from_fragments, + head_weight, + shard_logits, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_draft_step import ( # noqa: E501 + m3_draft_step as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + WORLD, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_step requires a CUDA device" + +T = DRAFT_LEN + 1 +_LOW32 = 0xFFFFFFFF + +# ------------------------------------------------------------------ the per-row rule against the in-tree token +# +# The kernel and the in-tree head (a cuBLAS BF16 GEMM) both round an FP32 sum of the same exact BF16 products to BF16, +# in different summation orders. A row's token can differ from the in-tree one only where some candidate's FP32 sum +# can land on either side of a BF16 rounding midpoint. Each row is held to: +# +# kernel token == in-tree token, or the row is provably undecidable: with the fp64 reference logit L_i (exact BF16 +# products, fp64 sum) and the FP32-order bound E_i = gamma_K * sum_k |w_ik * x_k|, the BF16 value of any FP32 sum of +# candidate i lies in [lo_i, hi_i] = [bf16(L_i - E_i), bf16(L_i + E_i)]; the kernel's token must win for some +# choice of values within those intervals (larger value, then the lower global id), and a BF16 rounding midpoint +# must lie within E of the kernel's or the in-tree token's L. +# +# gamma_K = K * 2^-22 is la-xh2's bound (4x the round-to-nearest unit roundoff 2^-24 per addition: any summation +# order, and a truncating tensor-core accumulator); K counts the row's non-zero products (adding an exact zero does +# not round). lo / hi take one more FP32 step outward before the BF16 rounding, which covers rounding L -+ E to FP32. +# Each undecidable row is reported as a warning with its midpoint distances, and each test reports its row count; any +# other mismatch fails. +GAMMA_UNIT = 2.0**-22 + + +def _bf16_outward(v: torch.Tensor, direction: float) -> torch.Tensor: + f = v.float() + f = torch.nextafter(f, torch.full_like(f, direction)) + b = f.to(torch.bfloat16).float() + return torch.where(b == 0, torch.zeros_like(b), b) + + +def bounds(h: torch.Tensor, w: torch.Tensor): + """(fp64 logits L, bound E, lo, hi) [rows, N] of the head input rows h against the head rows w.""" + cols = (h != 0).any(dim=0) # zero columns add exact zeros + hd, wd = h[:, cols].double(), w[:, cols].double() + logits = hd @ wd.t() + bound = (h != 0).sum(dim=1, keepdim=True).double() * GAMMA_UNIT * (hd.abs() @ wd.abs().t()) + del wd + inf = float("inf") + return logits, bound, _bf16_outward(logits - bound, -inf), _bf16_outward(logits + bound, inf) + + +def midpoint_distance(v: float) -> float: + """Distance of v to the nearest BF16 rounding midpoint.""" + b = torch.tensor(v, dtype=torch.float64).float().to(torch.bfloat16).float() + if float(b) == 0.0: + return abs(v) + ulp = 2.0 ** (int(torch.frexp(b).exponent) - 8) + return min(abs(v - (float(b) + ulp / 2)), abs(v - (float(b) - ulp / 2))) + + +def _winnable(tok: int, ids: torch.Tensor, lo: torch.Tensor, hi: torch.Tensor) -> bool: + """Whether candidate ``tok`` wins the row for some values within the candidates' [lo, hi].""" + k = int((ids == tok).nonzero()[0, 0]) + others = ids != tok + beat = (hi[k] > lo[others]) | ((hi[k] == lo[others]) & (tok < ids[others])) + return bool(beat.all()) + + +def check_row(name: str, kernel_tok: int, stock_tok: int, ids, logits, bound, lo, hi) -> bool: + """The per-row rule for one row over the candidates ``ids`` [C] (logits / bound / lo / hi [C]); True when the + row is undecidable (reported as a warning), False when the tokens match; a mismatch outside the bound fails.""" + if kernel_tok == stock_tok: + return False + assert _winnable(kernel_tok, ids, lo, hi), ( + f"{name}: token {kernel_tok} differs from the in-tree {stock_tok} outside the FP32-order bound" + ) + ik = int((ids == kernel_tok).nonzero()[0, 0]) + i_stock = int((ids == stock_tok).nonzero()[0, 0]) + lk, ls, ek, es = ( + float(logits[ik]), + float(logits[i_stock]), + float(bound[ik]), + float(bound[i_stock]), + ) + dk, ds = midpoint_distance(lk), midpoint_distance(ls) + assert dk <= ek or ds <= es, f"{name}: neither logit is within its bound of a BF16 midpoint" + warnings.warn( + f"[draft-step undecidable] {name}: kernel {kernel_tok} (L {lk:.9g}, midpoint distance {dk:.3e}, bound " + f"{ek:.3e}) vs in-tree {stock_tok} (L {ls:.9g}, midpoint distance {ds:.3e}, bound {es:.3e})" + ) + return True + + +def report(name: str, rows: int, undecidable: int) -> None: + warnings.warn(f"[draft-step rows] {name}: {rows} rows, {undecidable} undecidable") + + +def key_value(keys: torch.Tensor) -> torch.Tensor: + """The fp32 value a key orders by (the inverse of the kernel's order-preserving image).""" + u = (keys >> 32) & _LOW32 + bits = torch.where(u >= 0x80000000, u - 0x80000000, _LOW32 - u) + bits = torch.where(bits >= 2**31, bits - 2**32, bits) # as a signed 32-bit pattern + return bits.to(torch.int32).view(torch.float32) + + +def key_index(keys: torch.Tensor) -> torch.Tensor: + return _LOW32 - (keys & _LOW32) + + +def _bookkeeping(nreq: int, step: int, n_x_rows: int, g: torch.Generator): + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if step == 0 + else None + ) + seq_lens = torch.full((nreq,), T, dtype=torch.int32, device="cuda") if step == 0 else None + next_new = torch.full((nreq, T), -5, dtype=torch.int32, device="cuda") + pos_src = torch.randint(100, 3000, (n_x_rows,), generator=g, device="cuda", dtype=torch.int32) + kv_lens = torch.randint(200, 3000, (nreq,), generator=g, device="cuda", dtype=torch.int32) + ids_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + pos_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + return num_acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next + + +#: The op's trailing coded-head arguments (exc_off, exc_pos, exc_w, exp_base): none, the BF16 fragments. +NO_CODING = (None, None, None, 0) + + +def _head(w: torch.Tensor): + """(wf, NO_CODING) of a BF16 shard: its BF16 fragments, as ``prepare_head`` builds them.""" + return bf16_fragments(w), NO_CODING + + +def _args( + x, + x_rows, + nreq, + wf, + rank, + world, + step, + g, + uc=None, + keys=None, + counters=None, + num_acc=None, + coded=NO_CODING, +): + """The op's argument tuple for one rank (fresh workspace and bookkeeping buffers unless given); ``coded`` = the + op's trailing coded-head arguments, always ``NO_CODING`` here.""" + acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next = _bookkeeping( + nreq, step, x.shape[0], g + ) + if num_acc is None: + num_acc = acc + if uc is None: + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + out_keys = torch.zeros(nreq, dtype=torch.int64, device="cuda") + tiles = wf.shape[0] if coded[0] is None else coded[0].numel() - 1 + return (x, x_rows, nreq, wf, rank * tiles * 16, keys, counters, uc, uc.data_ptr(), rank, world, + next_new[:, 1:], step, ids_next, pos_src, pos_next, kv_lens, num_acc, seq_lens, DRAFT_LEN, num_sms(), + True, out_keys, *coded) # fmt: skip + + +def _case(spec, seed): + g = gen(seed) + nreq, step, rank, world, vs, k = ( + spec["nreq"], + spec["step"], + spec["rank"], + spec["world"], + spec["vs"], + spec["k"], + ) + pass0 = step == 0 + x = randn((nreq * T if pass0 else nreq, k), g) + assert not spec["coded"], "the tp8 target certifies the BF16 fragments only" + wf, coded = _head(randn((vs, k), g, 0.02)) + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if pass0 + else None + ) + x_rows = (torch.arange(nreq, device="cuda") * T + num_acc.long() - 1) if pass0 else None + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + h = x[x_rows] if pass0 else x + own = shard_logits(h, wf, *coded).max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=MAX_TOKENS, ties=spec.get("ties", False) + ) + return _args(x, x_rows, nreq, wf, rank, world, step, g, uc, keys, counters, num_acc, coded) + + +@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 this name + """The bookkeeping is exact on every row (next positions, KV and sequence lengths). Each row's token is the + in-tree one (this rank's BF16 shard logits from the head's GEMM, their first maximum, against the peers' keys) or + the row is provably undecidable (check_row); the winning key's index is the token, and its value lies in the + token's [lo, hi] (or is the peer's pushed key when a peer wins).""" + args = _case(cell.spec, seed_of(cell.why)) + x, x_rows, nreq, wf, offset, _, counters, uc = args[:8] + rank, world = args[9], args[10] + half = int(counters[1].item()) & 1 + pushed = ( + uc[half * MAX_TOKENS * world : (half * MAX_TOKENS + nreq) * world].view(nreq, world).clone() + ) + h = x[x_rows] if x_rows is not None else x[:nreq] + w = head_weight(wf, x.shape[1], *args[23:27]) + best, idx = local_argmax(torch.nn.functional.linear(h, w).float(), offset) + stock = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, MAX_TOKENS) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert args[6].tolist() == [0, 1] and not bool(args[5].any()), "counters / keys not re-armed" + step = args[12] + others = [c for c in range(args[11].shape[1]) if c != step] + assert torch.equal(args[11][:, others], ref["draft_out"][:, others]) + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + assert torch.equal(args[15], ref["pos_next"]) + assert torch.equal(args[16], ref["kv_lens"]) + if args[18] is not None: + assert torch.equal(args[18], ref["seq_lens"]) + # Candidates per row: this rank's head rows (fp64 bounds) and each peer's pushed key (an exact value). + logits, bound, lo, hi = bounds(h, w) + own_ids = offset + torch.arange(w.shape[0], device="cuda") + peers = [p for p in range(world) if p != rank] + undecidable = 0 + for r in range(nreq): + pv = key_value(pushed[r, peers]).double() + ids = torch.cat([own_ids, key_index(pushed[r, peers])]) + undecidable += check_row( + f"{cell.why[:40]} row {r}", + int(args[13][r]), + int(stock[r]), + ids, + torch.cat([logits[r], pv]), + torch.cat([bound[r], torch.zeros_like(pv)]), + torch.cat([lo[r], pv.float()]), + torch.cat([hi[r], pv.float()]), + ) + report(cell.why[:40], nreq, undecidable) + out_keys = args[22] + assert torch.equal(key_index(out_keys), args[13].long()), "the winning key's index is the token" + got = key_value(out_keys) + for r in range(nreq): + c = int(args[13][r]) - offset + if 0 <= c < w.shape[0]: + assert float(lo[r, c]) <= float(got[r]) <= float(hi[r, c]), ( + f"row {r}: winning value outside [lo, hi]" + ) + else: + assert int(out_keys[r]) in pushed[r].tolist(), ( + f"row {r}: a peer won with a key it did not push" + ) + + +def test_fragments_round_trip() -> None: + g = gen(seed_of("fragments")) + w = randn((64, 96), g) + wf = bf16_fragments(w) + assert wf.shape == (4, 6, 256) + assert torch.equal(from_fragments(wf), w) + head = SimpleNamespace(weight=w, bias=None) + prepared = draft_step_kernels.prepare_head(head, 3) + assert not hasattr(prepared, "coded") + assert torch.equal(prepared.wf, wf) and prepared.vocab_offset == 3 * 64 and prepared.weight is w + + +# ------------------------------------------------------------------ against the in-tree drafter, bit for bit + + +def _stock_tokens(shards: list, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """The in-tree greedy draft tokens of the rows x over the vocab-sharded BF16 head: (argmax of the gathered float + logits, greedy_sample_draft_with_tp_gather's per-rank (index, value) pick).""" + logits = [torch.nn.functional.linear(x, w).float() for w in shards] + gathered = torch.argmax(torch.cat(logits, dim=-1), dim=-1).to(torch.int32) + combined = [] + for rank, shard_logit in enumerate(logits): + worker = SimpleNamespace(mapping=SimpleNamespace(tp_rank=rank)) + combined.append(Eagle3OneModelWorker._get_local_max_and_combined(worker, shard_logit)) + pick = Eagle3OneModelWorker._get_draft_tokens_from_gathered + pick = getattr(pick, "_torchdynamo_orig_callable", pick) + tp_gather = pick(None, torch.cat(combined, dim=-1)).to(torch.int32) + return gathered, tp_gather + + +def _fused_tokens(shards: list, x: torch.Tensor, step: int = 1) -> list: + """Each rank's fused draft tokens: every rank's own key from a one-rank run, then each rank's eight-rank pass with + its peers' keys in its slots, as their multicast pushes would put them.""" + g = gen(seed_of("fused tokens")) + nreq = x.shape[0] + heads = [_head(w) for w in shards] + own_keys = [] + for rank, (wf, cd) in enumerate(heads): + args = list(_args(x, None, nreq, wf, rank, 1, step, g, coded=cd)) + args[9] = 0 # one-rank run: this rank's slot only, global ids of `rank` + op(*args) + own_keys.append(args[22].clone()) + tokens = [] + for rank, (wf, cd) in enumerate(heads): + uc, keys, counters = head_workspace(WORLD, slot_rows=MAX_TOKENS) + for peer in range(WORLD): + if peer != rank: + for t in range(nreq): + uc[t * WORLD + peer] = own_keys[peer][t] + args = _args(x, None, nreq, wf, rank, WORLD, step, g, uc, keys, counters, coded=cd) + with validating(op): + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + tokens.append(args[13].clone()) + return tokens + + +def _integer_head(g: torch.Generator, rows: int = WORLD * DRAFT_SHARD) -> torch.Tensor: + """A [rows, 6144] head of small integers over the first 4 columns (zero elsewhere): its logits against integer + rows are exact integers in any summation order, many of which share a BF16 value.""" + w = torch.zeros((rows, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, :4] = torch.randint(-60, 61, (rows, 4), generator=g, device="cuda").to(torch.bfloat16) + return w + + +def _shards(w: torch.Tensor) -> list: + return [w[r * DRAFT_SHARD : (r + 1) * DRAFT_SHARD].contiguous() for r in range(WORLD)] + + +def _assert_stock(w: torch.Tensor, x: torch.Tensor, expect=None) -> None: + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + if expect is not None: + assert gathered.tolist() == expect, (gathered.tolist(), expect) + for rank, tokens in enumerate(_fused_tokens(shards, x)): + assert torch.equal(tokens, gathered), ( + f"rank {rank}: {tokens.tolist()} vs stock {gathered.tolist()}" + ) + + +def _ones_rows(nreq: int) -> torch.Tensor: + x = torch.zeros((nreq, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[:, :4] = 1.0 + return x + + +def test_forced_bf16_ties_take_the_lowest_global_id() -> None: + """A head row copied onto every rank that wins every row of ones (logit 244, every other at most 240): the lowest + global id must win, inside a shard and across ranks. Then integer rows in [-2, 2], whose logits (up to 480 in + magnitude, BF16 spacing 2 above 256) collide often, exactly or after BF16 rounding.""" + g = gen(seed_of("forced ties")) + w = _integer_head(g) + win = torch.tensor([61.0, 61.0, 61.0, 61.0], device="cuda").to(torch.bfloat16) + for gid in ( + 3 * DRAFT_SHARD + 5, + DRAFT_SHARD + 77, + DRAFT_SHARD + 9000, + 2 * DRAFT_SHARD + 1, + 7 * DRAFT_SHARD + 1, + ): + w[gid, :4] = win + _assert_stock(w, _ones_rows(5), expect=[DRAFT_SHARD + 77] * 5) + # Random integer rows: ties after BF16 rounding happen naturally. + x = torch.randint(-2, 3, (8, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + x[:, 4:] = 0 + _assert_stock(w, x) + + +def test_signed_zero_maxima_take_the_lowest_global_id() -> None: + """Rows whose maximum is zero: an all-zero head input (every logit +-0) and logits <= 0 with zeros by cancellation + on several ranks. -0 and +0 compare equal, so the lowest global id reaching zero wins.""" + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 0] = -1.0 + w[:, 1] = -1.0 + # Zeros by cancellation on ranks 1, 2, 3 and 6; 49999 (rank 1's 25008..50015) is the lowest of them. + for gid in (2 * DRAFT_SHARD + 5, 6 * DRAFT_SHARD + 4000, 3 * DRAFT_SHARD + 7, 49999): + w[gid, 0] = 1.0 + w[gid, 1] = -1.0 + x = torch.zeros((3, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[1, :2] = 1.0 + x[2, :2] = 2.0 + x[2, 2] = -1.0 # w[:, 2] is zero: a -0 product on every row + _assert_stock(w, x, expect=[0, 49999, 49999]) + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + (1000, 1001, "lower"), # 1001 rounds to 1000 in BF16: a tie, the lower id wins + (1001, 1002, "lower"), # 1001 and the tie-to-even midpoint 1002 both round to 1000 + (1000, 1004, "higher"), # adjacent BF16 values: the larger wins + (1003, 1000, "lower"), # 1003 rounds up to 1004 > 1000 + ], + ids=["one_fp32_step", "midpoint", "adjacent_bf16", "rounds_up"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +def test_near_equal_top2_follow_the_bf16_logits(lower, higher, winner, where) -> None: + """Top-2 logits one to four fp32 steps apart: the winner is the one the BF16 logits decide (the in-tree logits' + dtype), not the fp32 sum. The pair sits inside rank 3's shard or on ranks 1 and 6; every other logit is at most + 240 in magnitude.""" + g = gen(seed_of(f"near top2 {lower} {higher} {where}")) + w = _integer_head(g) + if where == "in_shard": + lo_id, hi_id = 3 * DRAFT_SHARD + 100, 3 * DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD + 24000, 6 * DRAFT_SHARD + 3 + for gid, value in ((lo_id, lower), (hi_id, higher)): + q, r = divmod(value, 4) # four integers <= 251, exact in BF16 + w[gid, :4] = torch.tensor([q + (i < r) for i in range(4)], device="cuda").to(torch.bfloat16) + expect = lo_id if winner == "lower" else hi_id + _assert_stock(w, _ones_rows(2), expect=[expect, expect]) + + +# ------------------------------------------------------------------ non-integer rows at a BF16 rounding midpoint + +#: A BF16 rounding midpoint (between 6.0 and 6.03125; a tie there rounds to the even 6.0) and the FP32 spacing there. +MIDPOINT = 6.0 + 2.0**-6 +FP32_ULP = 2.0**-21 + + +def _near_midpoint_head(g: torch.Generator, targets: dict): + """A full-vocabulary head [8 * 25008, 6144] and one head-input row x with 65 non-zero entries: 1, 2^-12, 2^-20 + and 2^-28 on columns 0..3 and 61 random non-integer BF16 values. Every head row carries random non-integer BF16 + values on those 61 columns (logits about 0.2 in magnitude); each row of ``targets`` (global id -> fp64 value) also + gets columns 0..3 set so its fp64 logit is the target to within 2^-30 (a coarse value, then three finer + corrections, each BF16-rounded).""" + x = torch.zeros((1, HIDDEN), dtype=torch.bfloat16, device="cuda") + scales = (1.0, 2.0**-12, 2.0**-20, 2.0**-28) + for col, scale in enumerate(scales): + x[0, col] = scale + x[0, 4:65] = randn((61,), g, 0.5) + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 4:65] = randn((WORLD * DRAFT_SHARD, 61), g, 0.05) + xd = x[0].double() + for gid, target in targets.items(): + rest = float(target - (w[gid].double() * xd).sum()) + for col, scale in enumerate(scales): + v = torch.tensor(rest / scale, dtype=torch.float64).to(torch.bfloat16) + w[gid, col] = v + rest -= float(v) * scale + assert abs(rest) <= 2.0**-30, (gid, rest) + return w, x + + +def _check_fused_against_stock(name: str, w: torch.Tensor, x: torch.Tensor, expect=None) -> int: + """Each rank's fused tokens against the in-tree tokens under the per-row rule (full-vocabulary candidates); + ``expect`` (a token per row) additionally pins rows the bound decides. Returns the undecidable rows.""" + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + fused = _fused_tokens(shards, x) + for tokens in fused[1:]: + assert torch.equal(tokens, fused[0]), "the ranks disagree on the global token" + logits, bound, lo, hi = bounds(x, w) + ids = torch.arange(w.shape[0], device="cuda") + undecidable = 0 + for r in range(x.shape[0]): + undecidable += check_row( + f"{name} row {r}", + int(fused[0][r]), + int(gathered[r]), + ids, + logits[r], + bound[r], + lo[r], + hi[r], + ) + if expect is not None: + assert int(gathered[r]) == expect[r] and int(fused[0][r]) == expect[r], ( + name, + int(gathered[r]), + int(fused[0][r]), + expect[r], + ) + others = ids != expect[r] + assert _winnable(expect[r], ids, lo[r], hi[r]), ( + f"{name}: the bound does not let {expect[r]} win" + ) + rivals = [int(i) for i in ids[others][hi[r][others] >= lo[r][expect[r]]]] + assert not any(_winnable(i, ids, lo[r], hi[r]) for i in rivals), ( + f"{name}: the row is not decided" + ) + report(name, x.shape[0], undecidable) + return undecidable + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + ( + -2, + 2, + None, + ), # straddling the midpoint by 2 FP32 ulps: either side is in-tree up to the order + (2, -2, None), + (0, -2, None), # exactly on the midpoint (ties to the even 6.0) against just below it + ( + -2048, + 2048, + "higher", + ), # ~1e-3 either side, well outside the bound: decided by BF16 rounding + ( + -4096, + -2048, + "lower", + ), # both below the midpoint, ~1e-3 apart: a BF16 tie the lower id wins + ], + ids=["straddle", "straddle_reversed", "on_midpoint", "decided_up", "decided_tie"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +def test_top2_at_a_bf16_midpoint(lower, higher, winner, where) -> None: + """Non-integer rows whose top-2 fp64 logits sit the given number of FP32 ulps from a BF16 rounding midpoint, inside + rank 3's shard or on ranks 1 and 6, every other logit far below. The per-row rule holds on every row; outside the + bound the token is the one the BF16 rounding decides, bit for bit.""" + g = gen(seed_of(f"midpoint {lower} {higher} {where}")) + if where == "in_shard": + lo_id, hi_id = 3 * DRAFT_SHARD + 100, 3 * DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD + 24000, 6 * DRAFT_SHARD + 3 + w, x = _near_midpoint_head( + g, {lo_id: MIDPOINT + lower * FP32_ULP, hi_id: MIDPOINT + higher * FP32_ULP} + ) + expect = None if winner is None else [lo_id if winner == "lower" else hi_id] + _check_fused_against_stock(f"midpoint {where} {lower:+d}/{higher:+d} ulp", w, x, expect) + + +# ------------------------------------------------------------------ the head copy prepare_head builds + + +def test_prepare_head_streams_the_bf16_fragments_at_every_k() -> None: + """prepare_head keeps the fragment-ordered BF16 copy for every shard, a K the tp4 coding took (a multiple of 64) + included; the op's head weight is the shard bit for bit, and a coded stream is refused.""" + g = gen(seed_of("prepare bf16")) + for k in (1024, 96): + w = randn((4096, k), g, 0.02) + prepared = draft_step_kernels.prepare_head(SimpleNamespace(weight=w, bias=None), 2) + assert prepared.vocab_offset == 2 * 4096 and prepared.weight is w + assert torch.equal(prepared.wf, bf16_fragments(w)) + assert torch.equal(head_weight(prepared.wf, k).view(torch.int16), w.view(torch.int16)) + with pytest.raises(NotImplementedError): + head_weight(prepared.wf, 96, torch.zeros(2, dtype=torch.int32, device="cuda")) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 + modeling, + ) + + assert not hasattr(modeling, "DRAFT_HEAD_CODED") + assert all(not c.spec["coded"] for c in op.CELLS) + + +def test_the_cells_are_the_tp8_head() -> None: + assert (WORLD, DRAFT_SHARD, HIDDEN) == (8, 25008, 6144) + for cell in op.CELLS: + assert (cell.spec["world"], cell.spec["vs"], cell.spec["k"]) == (WORLD, DRAFT_SHARD, HIDDEN) + assert 0 <= cell.spec["rank"] < WORLD + + +def test_guard_refuses_head_rows_past_x() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[1] = args[1] + 1000 + with pytest.raises(AssertionError, match="x_rows must index x"): + with validating(op): + op(*args) + + +def test_guard_refuses_too_few_ctas() -> None: + args = list(_case(op.CELLS[1].spec, 2)) + args[20] = 8 + with pytest.raises(AssertionError, match="tiles per CTA"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_exact_numerics.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_exact_numerics.py new file mode 100644 index 000000000000..af6e3c862fd7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_exact_numerics.py @@ -0,0 +1,299 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp8 target's speculative step against the in-tree (stock) Eagle3 numerics, at the real TP8 shapes. + +* Verify: the target logits are the in-tree BF16 LM head's own shard GEMM, and the fused shard argmax + acceptance + gives the in-tree greedy acceptance of the gathered logits, bit for bit, on every rank (peers' keys are the real + per-shard maxima, computed by the same kernel). +* Drafter: the worker runs the in-tree drafter's arithmetic and weights (its BF16 linears on fixed-tactic BF16 low-M + kernels at the rows ``DRAFT_LOWM`` covers, held to the fp64 product and to the entry by + test_modeling_v2_m3_draft_lowm.py); its own drafter arithmetic is the first pass's fc_norm (chunk_norm, or + capture_norm folded into the capture), held here to the in-tree ``apply_eagle3_fc`` RMSNorm modules within one bf16 + ulp (its effect on the fc output and on the full-vocabulary BF16 head's draft tokens is measured too), and the + fused draft loop's draft input and BF16 head, held to the in-tree modules and greedy draft sampling by their own + tests (test_modeling_v2_m3_draft_input.py, test_modeling_v2_m3_draft_step.py). + +Synthetic weights: hidden 6144, vocabulary 200064 (8 x 25008 shards), 3 captured layers, draft length 3. +""" + +import importlib.util +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + make_keys, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 + m3_shard_accept as shard_accept_op, +) +from tensorrt_llm._torch.models.modeling_speculative import Eagle3ForCausalLM +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 + EPS, + HIDDEN, + TARGET_SHARD, + VOCAB, + WORLD, + assert_bf16_close, + craft_drafts, + gen, + head_workspace, + num_sms, + randn, + seed_of, +) + +assert torch.cuda.is_available(), "the exact-numerics checks require a CUDA device" + +DRAFT_LEN = 3 +CAPTURED_LAYERS = 3 +#: Largest share of fc_norm outputs one bf16 ulp off the in-tree RMSNorm (fp32 statistic summed in another order). +NORM_ULP_FRAC = 1e-3 + + +def _stock_accept(logits: torch.Tensor, draft: torch.Tensor): + """The in-tree greedy strict acceptance (``SpecWorkerBase._sample_and_accept_draft_tokens_base``) of the gathered + float logits of a decode-only batch, no forcing.""" + md = SimpleNamespace(is_all_greedy_sample=True, is_ragged_verify=False, is_cuda_graph=False) + worker = SimpleNamespace(force_num_accepted_tokens=0.0) + worker._sample_tokens_for_batch = lambda *a: Eagle3OneModelWorker._sample_tokens_for_batch( + worker, *a + ) + worker._apply_force_accepted_tokens = lambda *a, **k: ( + Eagle3OneModelWorker._apply_force_accepted_tokens(worker, *a, **k) + ) + nreq = draft.shape[0] + return Eagle3OneModelWorker._sample_and_accept_draft_tokens_base( + worker, logits, draft, 0, nreq, md + ) + + +def _lm_head_shard(weight: torch.Tensor, rank: int) -> LMHead: + """Rank ``rank``'s local part of the target's vocab-sharded BF16 LM head at TP8: an LM head over that rank's + 25008 rows (the TP8 module's local GEMM is this module's; its all-gather is the caller's).""" + head = LMHead(num_embeddings=TARGET_SHARD, embedding_dim=HIDDEN, dtype=torch.bfloat16) + head.weight.data = weight[rank * TARGET_SHARD : (rank + 1) * TARGET_SHARD].contiguous() + return head + + +def _run_shard_accept(logits, rank, world, draft, keys_of_peers=None, call=0): + """The fused verify of one rank: its shard logits, the peers' keys planted in its slots.""" + nreq = draft.shape[0] + rows = logits.shape[0] + uc, keys, counters = head_workspace(world, call=call, slot_rows=SHARD_ROWS) + half = call & 1 + for peer, peer_keys in (keys_of_peers or {}).items(): + for t in range(rows): + uc[(half * SHARD_ROWS + t) * world + peer] = peer_keys[t] + acc = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + shard_accept_op( + logits, rank * TARGET_SHARD if world > 1 else 0, keys, counters, uc, uc.data_ptr(), rank, world, + draft, batch_idx, acc, num, gather, 0, next_new, num_sms(), None, + ) # fmt: skip + torch.cuda.synchronize() + return acc, num, gather, next_new + + +def _local_keys(logits: torch.Tensor, rank: int) -> torch.Tensor: + """What rank ``rank`` pushes: the key of its own shard maximum per row, found by the same kernel at world 1.""" + rows = logits.shape[0] + nreq = rows // (DRAFT_LEN + 1) + dummy = torch.full((nreq, DRAFT_LEN), -1, dtype=torch.int32, device="cuda") + acc, _, _, _ = _run_shard_accept(logits, 0, 1, dummy) + local = acc.view(-1).to(torch.int64) + value = logits.float().gather(1, local[:, None]).squeeze(1) + return make_keys(value, local + rank * TARGET_SHARD) + + +@pytest.mark.parametrize( + "nreq,ties", [(1, False), (2, False), (4, False), (8, True), (12, False)], ids=str +) +def test_fused_verify_is_the_stock_greedy_acceptance(nreq, ties) -> None: + """Target logits (in-tree BF16 head) and verify outputs bit-identical to the in-tree path, on all eight ranks.""" + g = gen(seed_of(f"verify {nreq} {ties}")) + rows = nreq * (DRAFT_LEN + 1) + x = randn((rows, HIDDEN), g, 1.0) + weight = randn((VOCAB, HIDDEN), g, 0.02) + if ties: + # Two equal head rows on ranks 1 and 6 that win every row: the lower global id must be the token. + win = (x.float().sum(0) / x.float().sum(0).norm() * 4.0).to(torch.bfloat16) + weight[TARGET_SHARD + 77] = win + weight[6 * TARGET_SHARD + 5] = win + heads = [_lm_head_shard(weight, r) for r in range(WORLD)] + # Stock: each rank's LMHead forward, all-gathered, then LogitsProcessor's float cast. + shard_logits = [head(x) for head in heads] + for head, stock in zip(heads, shard_logits): + assert torch.equal(head.apply_linear(x, None), stock), ( + "the verify's shard GEMM is not the head's" + ) + gathered = torch.cat(shard_logits, dim=1).float() + target_tokens = gathered.argmax(dim=-1).to(torch.int32) + if ties: + assert bool((target_tokens == TARGET_SHARD + 77).all()) + draft = craft_drafts(target_tokens, nreq, DRAFT_LEN) + stock_acc, stock_num = _stock_accept(gathered, draft) + peer_keys = {r: _local_keys(shard_logits[r], r) for r in range(WORLD)} + for rank in range(WORLD): + peers = {r: k for r, k in peer_keys.items() if r != rank} + acc, num, gather, next_new = _run_shard_accept( + shard_logits[rank], rank, WORLD, draft, peers + ) + assert torch.equal(acc, stock_acc.to(torch.int32)), f"rank {rank}: accepted tokens" + assert torch.equal(num, stock_num.to(torch.int32)), f"rank {rank}: accepted counts" + expect_gather = torch.arange(nreq, device="cuda") * (DRAFT_LEN + 1) + num.long() - 1 + assert torch.equal(gather, expect_gather), f"rank {rank}: first-pass gather rows" + assert torch.equal(next_new[:, 0], target_tokens[expect_gather]), f"rank {rank}: next_new" + + +class _Drafter: + """The in-tree Eagle3 first-pass pieces at the target's shapes: per-layer fc_norm RMSNorms and the fc.""" + + def __init__(self, g: torch.Generator): + self.norms = [] + for _ in range(CAPTURED_LAYERS): + norm = RMSNorm(hidden_size=HIDDEN, eps=EPS, dtype=torch.bfloat16).cuda() + norm.weight.data = (1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to( + torch.bfloat16 + ) + self.norms.append(norm) + self.fc = Linear(HIDDEN * CAPTURED_LAYERS, HIDDEN, bias=False, dtype=torch.bfloat16).cuda() + self.fc.weight.data = randn((HIDDEN, HIDDEN * CAPTURED_LAYERS), g, 0.02) + self.model = SimpleNamespace( + dtype=torch.bfloat16, + hidden_size=HIDDEN, + fc_norm=torch.nn.ModuleList(self.norms), + _norm_before_fc=False, + fc=self.fc, + ) + self.stacked = torch.stack([n.weight.detach() for n in self.norms]).contiguous() + + def stock_normed(self, hs: torch.Tensor) -> torch.Tensor: + """``apply_eagle3_fc``'s fc input: each captured layer through its own in-tree RMSNorm.""" + chunks = hs.chunk(CAPTURED_LAYERS, dim=-1) + return torch.cat([norm(c) for norm, c in zip(self.norms, chunks)], dim=-1) + + def stock_fc(self, hs: torch.Tensor) -> torch.Tensor: + return Eagle3ForCausalLM.apply_eagle3_fc(self, hs) + + +def _captures(rows: int, g: torch.Generator): + """Per captured layer (hidden, residual) and the in-tree capture buffer of bf16(hidden + residual).""" + layers = [(randn((rows, HIDDEN), g, 1.5), randn((rows, HIDDEN), g, 3.0)) for _ in range(3)] + buf = torch.cat([h + r for h, r in layers], dim=-1) + return layers, buf + + +def _row_identical(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return (a == b).all(dim=-1) + + +@pytest.mark.parametrize("rows", (4, 8, 32, 48), ids=lambda r: f"rows{r}") +def test_first_drafter_inputs_are_the_stock_fc_norm(rows) -> None: + """chunk_norm and capture_norm against ``apply_eagle3_fc``'s in-tree RMSNorms: within one bf16 ulp on at most + NORM_ULP_FRAC of the elements; the fc (the in-tree module in both paths) then differs only on rows whose input + did.""" + g = gen(seed_of(f"fc_norm {rows}")) + drafter = _Drafter(g) + layers, buf = _captures(rows, g) + stock_in = drafter.stock_normed(buf) + stock_out = drafter.stock_fc(buf) + assert torch.equal(drafter.fc(stock_in), stock_out) + + chunk_in = m3_spec_tail.chunk_norm(buf, drafter.stacked, EPS) + capture_in = torch.empty_like(buf) + for i, (h, r) in enumerate(layers): + m3_spec_tail.capture_norm( + h, r, drafter.stacked[i], EPS, capture_in[:, i * HIDDEN : (i + 1) * HIDDEN] + ) + torch.cuda.synchronize() + for name, ours in (("chunk_norm", chunk_in), ("capture_norm", capture_in)): + assert_bf16_close(ours, stock_in, max_frac=NORM_ULP_FRAC) + out = drafter.fc(ours) + same_rows = _row_identical(ours, stock_in) + assert torch.equal(out[same_rows], stock_out[same_rows]), ( + f"{name}: fc moved an unchanged row" + ) + moved = (ours != stock_in).float().mean().item() + print( + f"[exact-numerics] rows={rows} {name}: {moved:.2e} of fc inputs one ulp off, " + f"{int((~same_rows).sum())}/{rows} rows differ, fc max |diff| " + f"{(out.float() - stock_out.float()).abs().max().item():.3e}" + ) + assert torch.equal(chunk_in, capture_in), "capture_norm and chunk_norm disagree" + + +@pytest.mark.parametrize("rows", (4, 32), ids=lambda r: f"rows{r}") +def test_draft_tokens_through_the_full_vocab_head(rows) -> None: + """The fc_norm's ulp-level differences carried through the full-vocabulary BF16 head (a proxy of the draft head + input): the greedy draft token equals the stock one on every row whose stock top-2 margin exceeds the logit + perturbation.""" + g = gen(seed_of(f"draft head {rows}")) + drafter = _Drafter(g) + _, buf = _captures(rows, g) + weight = randn((VOCAB, HIDDEN), g, 0.02) + head = LMHead(num_embeddings=VOCAB, embedding_dim=HIDDEN, dtype=torch.bfloat16) + head.weight.data = weight + stock_logits = head(drafter.stock_fc(buf)).float() + ours_logits = head(drafter.fc(m3_spec_tail.chunk_norm(buf, drafter.stacked, EPS))).float() + perturb = (ours_logits - stock_logits).abs().max(dim=-1).values + top2 = stock_logits.topk(2, dim=-1).values + decidable = (top2[:, 0] - top2[:, 1]) > 2 * perturb + stock_tok, ours_tok = stock_logits.argmax(-1), ours_logits.argmax(-1) + assert torch.equal(ours_tok[decidable], stock_tok[decidable]) + print( + f"[exact-numerics] rows={rows} draft head: {int((ours_tok == stock_tok).sum())}/{rows} tokens equal, " + f"{int(decidable.sum())} decidable, max logit |diff| {perturb.max().item():.3e}" + ) + + +def test_the_drafter_is_the_in_tree_drafter() -> None: + """Outside the fused draft loop's steps the worker drafts with the in-tree loop, forward and sampling, and nothing + re-quantizes the drafter's linears or the LM heads: no hook of a quantized / pruned head or drafter projection is + left in the target, and the fused head streams the BF16 head weight itself (no coded copy).""" + assert ( + modeling.M3Eagle3Worker._forward_draft_loop is not Eagle3OneModelWorker._forward_draft_loop + ) + for name in ( + "_forward_linear_draft_loop", + "_run_draft_forward", + "sample_draft_tokens", + "sample_and_accept_draft_tokens", + ): + assert getattr(modeling.M3Eagle3Worker, name) is getattr(Eagle3OneModelWorker, name), name + for module, name in ( + (modeling, "M3DraftGemvLinearMethod"), + (modeling, "adopt_draft_linears"), + (modeling, "greedy_head_outputs"), + (modeling, "M3FP8LMHeadMethod"), + (modeling, "adopt_fp8_lm_heads"), + ): + assert not hasattr(module, name), f"{module.__name__}.{name}" + assert not hasattr(modeling.M3Eagle3Worker, "prebuild_target_head") + pkg = modeling.__name__.rsplit(".", 1)[0] + for gone in ( + "ops.m3_draft_head", + "ops.spec.m3_draft_head", + "ops.gemm.m3_draft_gemv", + "ops.m3_target_head", + "ops.gemm.m3_target_head", + "ops.m3_draft_lowm", + ): + assert importlib.util.find_spec(f"{pkg}.{gone}") is None, gone diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_force_vec.py new file mode 100644 index 000000000000..fe0d05f52d1a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_force_vec.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_force_vec entry of the MiniMax-M3 tp8 target: the fractional forced-acceptance draw.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_force_vec requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(5) + pool = torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, generator=g, device="cuda") + counter = torch.tensor([41], dtype=torch.int64, device="cuda") + out = torch.empty(spec["batch"], dtype=torch.int32, device="cuda") + args = ( + pool, + counter, + out, + spec["batch"], + 2, + 4, + spec["frac"], + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + spec["block"], + ) + ref = op.reference(*args) + with validating(op): + op(*args) + assert int(counter.item()) == 42, "the draw bumps the step counter once" + assert torch.equal(out, ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_sampled_accept.py new file mode 100644 index 000000000000..25a1db2b0943 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_sampled_accept.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 m3_sampled_accept entry (fused temperature + top-p verify) and its workspace size.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_sampled_accept import ( # noqa: E501 + m3_sampled_accept as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_sampled_accept_ws_words import ( # noqa: E501 + m3_sampled_accept_ws_words as ws_words, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_sampled_accept requires a CUDA device" + +DRAFT_LEN = 3 +#: Hot tokens of the distribution cells; every other logit sits 22 below the top (mass ~ 1e-5 of the total). +HEAD = 12 +#: Draws per distribution cell, and the total-variation band against the reference distribution. With 12 +#: categories and >= 4000 draws the expected TV of an exact sampler is ~0.02 (sd ~0.006); 0.06 is ~6 sd above it +#: while a wrong nucleus cut (one category too many or too few, >= 1 / 12 of the mass moved) lands at >= 0.08. +DRAWS = 4000 +TV_BAND = 0.06 + + +def _logits(spec, g, rows): + vs = spec["vs"] + if spec["check"] == "flat": + return (torch.randn((rows, vs), generator=g, device="cuda") * 0.3).to(torch.bfloat16) + if spec["check"] == "argmax": + return (torch.randn((rows, vs), generator=g, device="cuda") * 3).to(torch.bfloat16) + row = torch.full((vs,), -20.0, device="cuda") + hot = torch.randperm(vs, generator=g, device="cuda")[:HEAD] + row[hot] = torch.linspace(2.0, 0.2, HEAD, device="cuda") + return row.to(torch.bfloat16).expand(rows, vs).contiguous() + + +def _case(spec, seed): + g = gen(seed) + nreq = spec["nreq"] + rows = nreq * (DRAFT_LEN + 1) + logits = _logits(spec, g, rows) + words = ws_words(1) + uc = torch.zeros(words, dtype=torch.int64, device="cuda") + counters = torch.zeros(2, dtype=torch.int32, device="cuda") + tokens = torch.zeros(48, dtype=torch.int32, device="cuda") + extra = torch.zeros(1, dtype=torch.int32, device="cuda") + temperatures = torch.full((rows,), spec["temperature"], dtype=torch.float32, device="cuda") + top_ps = torch.full((rows,), spec["top_p"], dtype=torch.float32, device="cuda") + seeds = torch.randint(0, 2**62, (rows,), generator=g, device="cuda", dtype=torch.int64) + offsets = torch.zeros(rows, dtype=torch.int64, device="cuda") + # Drafts drawn from the likeliest tokens, so some prefixes match. + top = logits[0].float().topk(3).indices.to(torch.int32) + spec["vocab_offset"] + draft = top[torch.randint(0, 3, (nreq, DRAFT_LEN), generator=g, device="cuda")].contiguous() + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + force = spec.get("force", "none") + force_total = 3 if force == "total" else 0 + force_vec = (torch.randint(1, DRAFT_LEN + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if force == "vec" else None) # fmt: skip + return [logits, spec["vocab_offset"], temperatures, top_ps, seeds, offsets, uc, uc.data_ptr(), counters, tokens, + 0, 1, draft, batch_idx, acc, num, gather, force_total, next_new, force_vec, extra, True] # fmt: skip + + +@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 this name + """Statistical for the draw, exact for everything derived from it. + + Every drawn token has nonzero reference probability (lies in the nucleus); the acceptance outputs equal the + reference's strict acceptance of the drawn tokens, every call; the call counter advances once per call; and, + for the distribution cells, the draws over many calls (fresh Philox offsets) are within TV_BAND of the + reference distribution (see DRAWS / TV_BAND). + """ + spec = cell.spec + args = _case(spec, seed_of(cell.why)) + logits, offset, rows = args[0], spec["vocab_offset"], args[0].shape[0] + certifying(op) + ref = op.reference(*args) + probs = ref["probs"] + calls = 1 if spec["check"] == "argmax" else -(-DRAWS // rows) + counts = torch.zeros(logits.shape[1], dtype=torch.float64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + drawn = args[9][:rows].clone() + local = drawn.long() - offset + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + assert bool((probs.gather(1, local[:, None]) > 0).all()), ( + "a drawn token outside the nucleus" + ) + acc = ref["accept"](drawn) + assert torch.equal(args[14], acc["accepted"]) + assert torch.equal(args[15], acc["num_accepted"]) + assert torch.equal(args[16], acc["gather_ids"]) + assert torch.equal(args[18], acc["next_new"]) + assert args[8].tolist() == [0, c + 1] + counts.index_add_(0, local, torch.ones(rows, dtype=torch.float64, device="cuda")) + if spec["check"] == "distribution": + emp = counts / counts.sum() + tv = 0.5 * (emp - probs[0].double()).abs().sum().item() + assert tv <= TV_BAND, f"total variation {tv:.3f} against the reference nucleus distribution" + + +@pytest.mark.parametrize("top_p", [0.001, 0.0002]) +def test_long_rejection_chains_stay_in_the_nucleus(top_p) -> None: + """A tiny nucleus over a flat tail: rows that need more exchange rounds than the workspace has slots. + + Each round's candidate is drawn over the tokens heavier than the last one, so a row takes ~ln(1 / top_p) + 1 + rounds on average, with a tail past the 16 exchange slots (about 1% of the rows at top_p 0.001). Rounds reuse the + slots cyclically, told apart by a lap-parity bit in the wire tags, and no round accepts unconditionally: every + draw, however many rounds it took, lies in the reference nucleus. World 1 (one rank's 8 CTAs per row exchange + through the same slot ring); the cross-rank path needs the multi-device job. + """ + spec = dict(nreq=1, vs=25008, vocab_offset=25008, check="flat", top_p=top_p, temperature=1.0) + args = _case(spec, seed_of(f"flat tail top_p {top_p}")) + logits, rows = args[0], args[0].shape[0] + probs = op.reference(*args)["probs"] + calls = -(-DRAWS // rows) + outside = torch.zeros((), dtype=torch.int64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + local = args[9][:rows].long() - spec["vocab_offset"] + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + outside += (probs.gather(1, local[:, None]) == 0).sum() + assert args[8].tolist() == [0, c + 1] + draws = calls * rows + rounds = 2.0 + args[20].item() / draws + kept = int((probs > 0).sum(1).min().item()) + assert outside.item() == 0, ( + f"{outside.item()} of {draws} draws outside the nucleus (kept >= {kept} of {logits.shape[1]}, " + f"{rounds:.2f} rounds per row on average)" + ) + assert rounds > 4.0, ( + f"only {rounds:.2f} rounds per row: the rows no longer stress the slot ring" + ) + + +@pytest.mark.parametrize("cell", ws_words.CELLS, ids=[c.why[:44] for c in ws_words.CELLS]) +def test_workspace_words(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(ws_words) + world = cell.spec["world"] + assert ws_words(world) == ws_words.reference(world) + + +def test_guard_refuses_top_k_rows() -> None: + args = _case(op.CELLS[0].spec, 1) + top_ks = torch.full((args[0].shape[0],), 20, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="top-k rows"): + with validating(op): + op(*args, top_ks=top_ks) + + +def test_guard_refuses_fp64_temperatures() -> None: + args = _case(op.CELLS[0].spec, 2) + args[2] = args[2].double() + with pytest.raises(AssertionError, match="fp32 sampling parameters"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_shard_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_shard_accept.py new file mode 100644 index 000000000000..e32448e55119 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_shard_accept.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_shard_accept entry (greedy verify from this rank's logits shard).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 + m3_shard_accept as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 + craft_drafts, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_shard_accept requires a CUDA device" + + +def _case(spec, seed, edit_logits=None): + g = gen(seed) + nreq, d, vs, world, rank = ( + spec["nreq"], + spec["draft_len"], + spec["vs"], + spec["world"], + spec["rank"], + ) + m = nreq * (d + 1) + if spec["strided"]: + logits = randn((m, vs + 72), g, 3.0)[:, :vs] + else: + logits = randn((m, vs), g, 3.0) + if edit_logits is not None: + edit_logits(logits) + uc, keys, counters = head_workspace(world, slot_rows=SHARD_ROWS) + own = logits.float().max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=SHARD_ROWS, ties=spec.get("ties", False) + ) + offset = rank * vs + best, idx = local_argmax(logits.float(), offset) + tokens = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, SHARD_ROWS) + draft = craft_drafts(tokens, nreq, d) + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + force_total = 3 if spec["force"] == "total" else 0 + force_vec = (torch.randint(1, d + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if spec["force"] == "vec" else None) # fmt: skip + return (logits, offset, keys, counters, uc, uc.data_ptr(), rank, world, draft, batch_idx, acc, num, gather, + force_total, next_new, num_sms(), force_vec) # fmt: skip + + +@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 this name + """Bit-exact: the op reduces BF16 logits without arithmetic, ties (within the shard and across ranks) go to + the lowest global index on both sides.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) + assert args[3].tolist() == [0, 1] and not bool(args[2].any()) + + +def test_guard_refuses_nan_logits() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0][1, 7] = float("nan") + with pytest.raises(AssertionError, match="NaN"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_batch_index_past_the_requests() -> None: + args = list(_case(op.CELLS[0].spec, 2)) + args[9] = args[9] + 1 + with pytest.raises(AssertionError, match="batch_idx"): + with validating(op): + op(*args) + + +def _signed_zero_maxima(logits: torch.Tensor) -> None: + """Every row's maximum is zero, reached as -0.0 at a lower index and +0.0 at a higher one.""" + logits.fill_(-1.0) + logits[:, 3] = -0.0 + logits[:, 9] = 0.0 + + +@pytest.mark.parametrize("ties", [False, True], ids=["peer_rows", "peer_ties"]) +def test_signed_zero_maxima_take_the_lower_index(ties) -> None: + """-0.0 and +0.0 compare equal, so the lower index wins as with torch.argmax (also across ranks).""" + spec = dict(op.CELLS[1].spec, ties=ties) + args = _case(spec, seed_of(f"signed-zero-{ties}"), edit_logits=_signed_zero_maxima) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_slot_scatter.py new file mode 100644 index 000000000000..ba3220639b65 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_slot_scatter.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_slot_scatter entry (the spec sampler's per-slot stores in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_slot_scatter import ( # noqa: E501 + m3_slot_scatter as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_slot_scatter requires a CUDA device" + +STALE = -4 + + +def _case(spec, seed): + """Worker outputs and the store's transposed [slots, width] views of its [width, slots, 1] buffers.""" + g = gen(seed) + n, nslots = spec["n"], spec["slots"] + + def rnd(shape): + return torch.randint(0, VOCAB, shape, generator=g, device="cuda", dtype=torch.int32) + + store_new = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_next = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_lens = torch.full((nslots,), STALE, dtype=torch.int32, device="cuda") + store_draft = torch.full((nslots, spec["store_d"]), STALE, dtype=torch.int32, device="cuda") + slots = torch.randperm(nslots, generator=g, device="cuda")[:n] + lens = torch.randint(1, 5, (n,), generator=g, device="cuda", dtype=torch.int32) + return (slots, rnd((n, spec["src_w"])), store_new.squeeze(-1).T, rnd((n, spec["src_w"])), + store_next.squeeze(-1).T, lens, store_lens, rnd((n, spec["src_d"])), store_draft) # fmt: skip + + +@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 this name + """Index copies: exact, zero padding included, other slots untouched.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in (("da", args[2]), ("db", args[4]), ("dl", args[6]), ("dd", args[8])): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_repeated_slots() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0] = args[0].clone() + args[0][1] = args[0][0] + with pytest.raises(AssertionError, match="distinct"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_step_gather.py new file mode 100644 index 000000000000..e1a5aed2034b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_step_gather.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_step_gather entry (the overlap scheduler's step-input gathers in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_step_gather import ( # noqa: E501 + m3_step_gather as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_step_gather requires a CUDA device" + +#: Rows of the engine's flat buffers (the max batch), and their stale fill. +MAX_ROWS = 8 +STALE = -4 + + +def _case(spec, seed): + g = gen(seed) + p, t, d, nslots, b0 = spec["p"], spec["t"], spec["d"], spec["slots"], spec["b0"] + i32 = dict(dtype=torch.int32, device="cuda") + new_tok = torch.randint(0, VOCAB, (t, nslots, 1), generator=g, device="cuda", dtype=torch.int32) + if spec.get("two_d"): + new_tok = new_tok.squeeze(-1) + next_draft = torch.randint(0, VOCAB, (nslots, d), generator=g, device="cuda", dtype=torch.int32) + lens = torch.randint(1, t + 1, (nslots,), generator=g, device="cuda", dtype=torch.int32) + slots = torch.randperm(nslots, generator=g, device="cuda")[:p].to(torch.int32) + ppi = slots.repeat_interleave(t) + ids = torch.full((MAX_ROWS * t,), STALE, **i32) + draft = torch.full((MAX_ROWS * d,), STALE, **i32) + pos_off = torch.full((MAX_ROWS * t,), STALE, **i32) + kv_off = torch.full((MAX_ROWS,), STALE, **i32) + return ( + new_tok, + next_draft, + lens, + slots, + ppi, + p, + t, + d, + ids, + spec["ids0"], + draft, + spec["d0"], + pos_off, + kv_off, + b0, + ) + + +@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 this name + """Index copies: exact, including the zero-filled offsets and the untouched rest of ids / draft.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in ( + ("ids", args[8]), + ("draft", args[10]), + ("pos_off", args[12]), + ("kv_off", args[13]), + ): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_a_slot_past_the_store() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[3] = args[3].clone() + args[3][0] = 8 + with pytest.raises(AssertionError, match="slots must index the store"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_capture_boundary.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_capture_boundary.py new file mode 100644 index 000000000000..455c283718ad --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_capture_boundary.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The generic layer's feed-forward exit on Eagle3 capture layers, in both targets' bodies (``PrefillTarget._layer``, +``DecodeTarget._generic_layer``). + +Every layer's MoE / MLP output goes unreduced to the fused boundary (all-reduce + residual + the next layer's input +norm, the final norm after the last layer), and a capture layer captures the reduced residual sum that boundary +writes, with no residual of its own. CPU only: the attention, the feed-forward and the boundaries are stand-ins. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 + modeling, +) + +LAYERS = 4 + + +class _Spec: + def __init__(self, layers): + self.layers = set(layers) + + def is_layer_capture(self, i): + return i in self.layers + + +@pytest.fixture(params=["prefill", "decode"]) +def layer(request, monkeypatch): + """(core, the target's generic layer bound to a target over ``core``).""" + calls = [] + w = {f"l{i}_{n}": f"{n}{i}" for i in range(LAYERS) for n in ("in_norm", "post_norm")} + w["final_norm"] = "final" + c = SimpleNamespace( + w=w, num_layers=LAYERS, moe_layers=frozenset({2, 3}), rms_eps=1e-6, calls=calls + ) + monkeypatch.setattr(modeling, "_can_mc", lambda x, folded=False: False) + monkeypatch.setattr( + modeling, + "_native_allreduce_norm", + lambda core, x, r, norm: (("normed", x, norm), ("updated", x, r)), + ) + + def boundary(self, x, x2, residual, norm_w): + calls.append((x, x2, norm_w)) + return ("next", norm_w), ("sum", x, x2) + + if request.param == "prefill": + cls = modeling.PrefillTarget + monkeypatch.setattr(cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i))) + monkeypatch.setattr(cls, "_swiglu_mlp", lambda self, h, prefix: ("mlp", prefix)) + else: + cls = modeling.DecodeTarget + monkeypatch.setattr( + cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i), False, False) + ) + monkeypatch.setattr( + cls, "_swiglu_mlp", lambda self, h, prefix, push=None: (("mlp", prefix), False) + ) + monkeypatch.setattr(cls, "_attention", lambda self, i, *a, **k: (("attn", i), False)) + monkeypatch.setattr(cls, "_generic_boundary", boundary) + target = cls(c) + return c, target._layer if request.param == "prefill" else target._generic_layer + + +@pytest.mark.parametrize("i", range(LAYERS)) +def test_a_capture_layer_captures_the_boundary_residual_sum(layer, i) -> None: + core, run = layer + captured = [] + out = run( + i, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([i]), + lambda spec, layer, h, r: captured.append((layer, h, r)), + ) + norm = core.w[f"l{i + 1}_in_norm"] if i + 1 < LAYERS else "final" + assert core.calls[-1][2] == norm, ( + "the boundary applies the next layer's norm (the final norm after the last)" + ) + assert captured == [(i, out[1], None)] + if i in core.moe_layers: + assert core.calls[-1][1] == ("routed", i), "the routed part folds into the boundary" + + +def test_a_layer_without_capture_captures_nothing(layer) -> None: + _, run = layer + captured = [] + run( + 1, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([]), + lambda *a: captured.append(a), + ) + assert not captured diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_dispatch.py new file mode 100644 index 000000000000..61254b231485 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_dispatch.py @@ -0,0 +1,1054 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp8 decode target's dispatch table (``modeling.py``, imported here as ``dispatch``): its row ranges against the +catalog cells that certify them and the kernels' own limits, the selection for every decode row count, and the miss / +uncovered policy. + +CPU only: nothing here launches a kernel. +""" + +import ast +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops import ( # noqa: E501 + m3_attn_lc8, + m3_draft_step, + m3_fused_score_select, + m3_mc_allreduce, + m3_shard_accept, + triton_sparse_decode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention import ( # noqa: E501 + m3_fused_score_select as fused_score_select_entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention import ( # noqa: E501 + m3_sparse_decode_lc8 as lc8_entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm import ( # noqa: E501 + m3_draft_lowm as draft_lowm_entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.mx_gemv import ( # noqa: E501 + op as mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec import ( # noqa: E501 + m3_draft_step as draft_step_entry, +) +from tensorrt_llm._torch.modules.linear import TensorParallelMode + +_TARGET = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8" +) + +PROJECTIONS = (QKV_SPARSE, QKV, O_PROJ, DENSE_GATE_UP, DENSE_DOWN) +ROWS = range(1, dispatch.DECODE_MAX_TOKENS + 2) # every decode row count and the first one past it + +# The Eagle3 drafter's BF16 linears at TP8 (per-rank qkv, o_proj, gate_up, down shards and the replicated fc). +DRAFTER_LINEARS = ((1280, 12288), (6144, 1024), (4608, 6144), (6144, 2304), (6144, 18432)) + + +@pytest.fixture(autouse=True) +def _policy(): + """Every test starts outside require mode with nothing logged, and leaves it so.""" + dispatch.set_require(False) + dispatch.reset_logged() + yield + dispatch.set_require(False) + dispatch.reset_logged() + + +class _Log: + def __init__(self): + self.info, self.warning = [], [] + + +@pytest.fixture +def log(monkeypatch) -> _Log: + """The dispatch's log lines, by level.""" + rec = _Log() + monkeypatch.setattr(dispatch.logger, "info", rec.info.append) + monkeypatch.setattr(dispatch.logger, "warning", rec.warning.append) + return rec + + +# ------------------------------------------------------------------ the ranges against their certificates + + +@pytest.mark.parametrize("fast_path", dispatch.FAST_PATHS, ids=lambda r: r.path) +def test_every_covered_shape_is_a_certified_cell(fast_path) -> None: + uncertified = fast_path.keys - fast_path.certified() + assert not uncertified, ( + f"{fast_path.path} covers {sorted(uncertified, key=str)} but " + f"{type(fast_path.entry).__name__}.CELLS certifies none of them: add a cell or narrow the range" + ) + + +def test_the_fast_path_names_are_unique() -> None: + names = [r.path for r in dispatch.FAST_PATHS] + assert len(names) == len(set(names)) + + +def test_the_removed_tp4_paths_are_gone() -> None: + """The tp8 target carries none of the tp4 target's coded heads, low-M MXFP8 GEMM, 9..13-row MoE engine or + attention megakernels: no fast path, range constant or kernel module of theirs is left to select.""" + names = {r.path for r in dispatch.FAST_PATHS} + assert not names & { + "m3_target_head", + "mxfp8_low_m_gemm", + "m3_moe16_engine", + "attn_mega", + "attn_c1", + "fused_decode_attn", + } + for constant in ( + "TARGET_HEAD", + "TARGET_HEAD_MAX_ROWS", + "DRAFT_HEAD_CODED", + "LOW_M_GEMM", + "LOW_M_GEMM_PATH", + "LOW_M_GEMM_MAX_TOKENS", + "MOE16_ENGINE", + "MOE16_MIN_TOKENS", + "MOE16_MAX_TOKENS", + "ATTN_MEGA", + "ATTN_C1", + "ATTN_C1_QUERY_LEN", + "FUSED_DECODE_ATTN", + "FUSED_ATTN_MAX_QUERY_LEN", + "FUSED_ATTN_MAX_TOKENS", + "SHARED_GATE_UP", + "SHARED_DOWN", + ): + assert not hasattr(dispatch, constant), constant + for module in ( + "ops.m3_target_head", + "ops.mxfp8_low_m_gemm", + "ops.m3_attn_mega", + "ops.m3_attn_mega_c1", + "ops.m3_fused_attn", + "ops.m3_attn_lc", + "ops.gemm.m3_target_head", + "ops.gemm.mxfp8_low_m_gemm", + "ops.moe.m3_moe16_engine", + "ops.attention.m3_attn_mega", + "ops.attention.m3_attn_c1", + "ops.attention.m3_fused_decode_attn", + "ops.attention.m3_sparse_decode_lc", + ): + assert importlib.util.find_spec(f"{_TARGET}.{module}") is None, module + for method in ("prebuild_target_head", "_m3_coded_head"): + assert not hasattr(dispatch.M3Eagle3Worker, method), method + for method in ("moe16_engine", "_build_moe16_engine"): + assert not hasattr(dispatch.MiniMaxM3Tp8ModelingV2, method), method + + +# ------------------------------------------------------------------ the ranges against the kernels' own limits + + +@pytest.mark.parametrize("nk", PROJECTIONS, ids=str) +def test_gemv_rows_are_the_rows_the_kernel_serves(nk) -> None: + n, k = nk + served = {m for m in range(1, mx_gemv.MAX_TOKENS + 1) if mx_gemv.serves(n, k, m, False)} + assert dispatch.MX_GEMV_ROWS[nk] == served + served_bf16 = [m for m in range(1, mx_gemv.MAX_TOKENS + 1) if mx_gemv.serves(n, k, m, True)] + assert served_bf16 == list(range(1, dispatch.MX_GEMV_BF16_MAX_TOKENS[nk] + 1)) + + +def test_the_gemv_tables_are_the_tp8_projections() -> None: + assert set(dispatch.MX_GEMV_ROWS) == set(dispatch.MX_GEMV_BF16_MAX_TOKENS) == set(PROJECTIONS) + assert PROJECTIONS == ((1536, 6144), (1280, 6144), (6144, 1024), (3072, 6144), (6144, 1536)) + for n, k in PROJECTIONS: + assert mx_gemv.get_config(n, k) is not None, (n, k) + + +@pytest.mark.parametrize("nk,push", [((768, 6144), None), ((6144, 384), "moe")], ids=str) +def test_the_shared_expert_projections_take_no_gemv(nk, push) -> None: + """At TP8 the shared expert (768 x 6144 gate_up, 6144 x 384 down) runs inside the MoE engine on decode rows and on + the generic GEMM otherwise: the GEMV has no config for it and the dispatch never routes a call there.""" + assert mx_gemv.get_config(*nk) is None + assert dispatch._gemv_variants(*nk, push, 8) == [] + + +def test_the_plans_are_the_all_reduce_buckets() -> None: + assert tuple(p.max_tokens for p in dispatch.DECODE_PLANS) == m3_mc_allreduce.BUCKETS + assert dispatch.DECODE_MAX_TOKENS == m3_mc_allreduce.MAX_TOKENS == mx_gemv.MAX_TOKENS + + +def test_the_thresholds_are_the_kernels() -> None: + assert dispatch.SPARSE_DECODE_MQ_MIN_TOKENS == triton_sparse_decode.MQ_MIN_TOKENS + assert ( + dispatch.MOE_ENGINE_MAX_TOKENS < dispatch.ATTN_PUSH_MAX_TOKENS <= dispatch.DECODE_MAX_TOKENS + ) + assert dispatch.MC_FINALIZE_MIN_TOKENS == dispatch.MOE_ENGINE_MAX_TOKENS + 1 + verify_rows = dispatch.SHARD_ACCEPT_MAX_REQUESTS * (dispatch.SHARD_ACCEPT_DRAFT_LEN + 1) + assert verify_rows == dispatch.DECODE_MAX_TOKENS == m3_shard_accept.SHARD_ROWS + assert dispatch.FUSED_SCORE_SELECT_MAX_REQUESTS == fused_score_select_entry.MAX_REQUESTS == 2 + assert dispatch.INDEX_SCORE_EARLY_MAX_REQUESTS == 0 + assert dispatch.FC_EARLY_MAX_ROWS == 0 + assert dispatch.LC8_PREWAIT_MAX_ROWS == 8 + + +def test_no_fast_path_quantizes_a_bf16_head_or_the_drafter() -> None: + """The target carries the exact-numerics configuration only: the lm_head and the Eagle3 drafter run at the + checkpoint's precision, so no fast path (nor a range constant) of an FP8 / NVFP4 / pruned head or of a quantized + drafter projection is left to select.""" + names = {r.path for r in dispatch.FAST_PATHS} + assert not names & { + "fp8_lm_head", + "draft_head", + "m3_draft_gemv", + "m3_draft_in_gemv", + "m3_gemv_ar", + } + for constant in ( + "FP8_LM_HEAD", + "FP8_LM_HEAD_MAX_TOKENS", + "DRAFT_HEAD", + "DRAFT_HEAD_MAX_TOKENS", + ): + assert not hasattr(dispatch, constant), constant + + +# ------------------------------------------------------------------ selection, every decode row count + + +def _expected_moe(rows: int): + if rows <= 8: + return "m3_moe_engine" + if rows <= 48: + return "mc_finalize_allreduce_norm" + return None + + +def _expected_router(rows: int): + if rows == 1: + return "m3_router_gemv" + if rows <= 48: + return "m3_router_tc" + return None + + +@pytest.mark.parametrize("rows", ROWS) +def test_selection(rows) -> None: + assert dispatch.is_decode_step(rows) == (rows <= 48) + if rows <= 48: + assert dispatch.plan_for(rows).name == ( + "tiny" if rows <= 8 else "small" if rows <= 16 else "verify" + ) + else: + with pytest.raises(AssertionError, match="no decode plan"): + dispatch.plan_for(rows) + assert dispatch.moe_path(rows) == _expected_moe(rows) + assert dispatch.router_path(rows) == _expected_router(rows) + assert dispatch.EMBED_FIRSTNORM.covers(rows) == (rows <= 8) + assert dispatch.ATTN_PUSH.covers(rows) == (rows <= 16) + assert dispatch.MOE_PUSH.covers(rows) == (rows <= 8) + assert dispatch.MC_ALLREDUCE.covers(rows) == (rows <= 48) + assert dispatch.FUSED_SCORE_SELECT.covers(rows) == (rows <= 2) + + +@pytest.mark.parametrize("rows", ROWS) +def test_linear_selection(rows) -> None: + bf16_max = { + QKV_SPARSE: 16, + QKV: 16, + O_PROJ: 48, + DENSE_GATE_UP: 16, + DENSE_DOWN: 32, + } + for nk in PROJECTIONS: + assert dispatch.linear_path(*nk, rows, True) == ( + "m3_mx_gemv" if rows <= 48 else "flashinfer_mxfp8_gemm" + ) + assert dispatch.linear_path(*nk, rows, False) == ( + "m3_mx_gemv_bf16" if rows <= bf16_max[nk] else "mxfp8_gemm" + ) + # A projection outside the table (the shared expert's, which runs in the MoE engine) takes the generic GEMM. + for nk in ((768, 6144), (6144, 384), (4096, 4096)): + assert dispatch.linear_path(*nk, rows, True) == "flashinfer_mxfp8_gemm" + assert dispatch.linear_path(*nk, rows, False) == "mxfp8_gemm" + + +def test_draft_step_selection() -> None: + """The fused draft loop takes the configured draft length at 1..8 requests (one head row per request, the + kernel's limit); its draft input every pass's rows: the first pass's 4 per request, then one per request.""" + assert ( + dispatch.DRAFT_STEP_MAX_REQUESTS + == draft_step_entry.MAX_TOKENS + == m3_draft_step.MAX_TOKENS + == 8 + ) + assert dispatch.DRAFT_STEP_DRAFT_LEN == dispatch.SHARD_ACCEPT_DRAFT_LEN + assert draft_step_entry.WORLD == 8 and draft_step_entry.DRAFT_SHARD * 8 == 200064 + for draft_len in range(1, 6): + for nreq in range(1, 14): + covered = draft_len == 3 and nreq <= 8 + assert dispatch.DRAFT_STEP.covers((draft_len, nreq)) == covered + assert dispatch.DRAFT_INPUT.keys == frozenset(range(1, 9)) | {12, 16, 20, 24, 28, 32} + + +def test_verify_selection() -> None: + for draft_len in range(1, 6): + for nreq in range(1, 14): + covered = draft_len == 3 and nreq * (draft_len + 1) <= 48 + assert dispatch.SHARD_ACCEPT.covers((draft_len, nreq)) == covered + + +def test_the_index_scorer_runs_its_prefetching_schedule_only() -> None: + """Spans of at most two requests score in the fused score + select; the scorer covers only the prefetching + kernel (no one-request early-issue schedule is ever selected).""" + assert dispatch.INDEX_SCORE.keys == frozenset({"prefetch"}) + assert not dispatch.INDEX_SCORE.covers("early") + for cell in dispatch.INDEX_SCORE.entry.CELLS: + assert not cell.spec["early"] + assert len(cell.spec["seq_lens"]) > dispatch.FUSED_SCORE_SELECT_MAX_REQUESTS + + +# ------------------------------------------------------------------ the long-context sparse decode (lc8) + + +def test_lc8_variants_are_the_prewait_threshold() -> None: + assert dispatch.SPARSE_DECODE_LC8.keys == frozenset({"prewait", "plain"}) + for cell in lc8_entry.m3_sparse_decode_lc8.CELLS: + rows = len(cell.spec["seq_lens"]) * cell.spec["decode_query_len"] + if cell.spec["prewait"]: + assert rows <= dispatch.LC8_PREWAIT_MAX_ROWS, cell.why + + +def _lc8_operands(total_q: int): + q = torch.zeros(total_q, lc8_entry.NUM_HEADS, lc8_entry.HEAD_DIM, dtype=torch.float8_e4m3fn) + pages = torch.zeros(1, 1, lc8_entry.PAGE, lc8_entry.HEAD_DIM // 2, dtype=torch.uint8) + scales = torch.zeros(1, 1, lc8_entry.PAGE, lc8_entry.HEAD_DIM // 16, dtype=torch.uint8) + topk = torch.full((1, total_q, 16), -1, dtype=torch.int32) + out = torch.zeros(total_q, lc8_entry.NUM_HEADS, lc8_entry.HEAD_DIM, dtype=torch.bfloat16) + return q, pages, scales, topk, out + + +@pytest.mark.parametrize("total_q", [1, 4, 7, 8, 9, 16, 32, 48]) +def test_lc8_prewait_follows_the_row_threshold(monkeypatch, total_q) -> None: + """The call site stages before the grid wait exactly up to ``LC8_PREWAIT_MAX_ROWS`` query rows.""" + calls = [] + monkeypatch.setattr(m3_attn_lc8, "m3_sparse_decode_lc8", lambda *args: calls.append(args)) + q, pages, scales, topk, out = _lc8_operands(total_q) + assert m3_attn_lc8.lc8_applies(q, pages, pages, topk, scales, scales, out) + m3_attn_lc8.m3_lc8_attn( + q, + pages, + pages, + topk, + torch.zeros(1, 1, dtype=torch.int32), + torch.zeros(1, dtype=torch.int32), + sm_scale=0.1, + output=out, + decode_query_len=1, + k_block_scale=scales, + v_block_scale=scales, + k_global_scale=torch.ones(1), + v_global_scale=torch.ones(1), + ) + (args,) = calls + assert args[-1] == (total_q <= dispatch.LC8_PREWAIT_MAX_ROWS) + assert args[13] == lc8_entry.WIDTH // 128 == 8 + + +def test_lc8_takes_the_tp8_geometry_only() -> None: + q, pages, scales, topk, out = _lc8_operands(4) + assert m3_attn_lc8.lc8_applies(q, pages, pages, topk, scales, scales, out) + q16 = torch.zeros(4, 16, 128, dtype=torch.float8_e4m3fn) + out16 = torch.zeros(4, 16, 128, dtype=torch.bfloat16) + assert not m3_attn_lc8.lc8_applies(q16, pages, pages, topk, scales, scales, out16) + + +# ------------------------------------------------------------------ require mode and fallback logging + + +def test_a_miss_raises_under_require() -> None: + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_moe_engine'.* 5 rows.*not built"): + dispatch.missed("m3_moe_engine", 5, "the engine is not built") + + +def test_a_miss_logs_once_per_reason_otherwise(log) -> None: + for _ in range(3): + dispatch.missed("m3_moe_engine", 5, "the engine is not built") + dispatch.missed("m3_moe_engine", 6, "the engine failed to build") + assert len(log.warning) == 2 + assert "m3_moe_engine at 5 rows: the engine is not built" in log.warning[0] + + +# ------------------------------------------------------------------ the phase split + + +def test_the_core_selects_the_target_on_the_context_count_alone() -> None: + """Every pure generation step is decode's, whatever its width; a step with context rows (a mixed one too) is + prefill's. Nothing but ``num_contexts`` is read: the metadata stand-in carries no other field.""" + core = SimpleNamespace(_prefill="prefill", _decode="decode") + select = dispatch.MiniMaxM3Tp8ModelingV2._select_target + assert select(core, SimpleNamespace(num_contexts=0)) == "decode" + assert select(core, SimpleNamespace(num_contexts=1)) == "prefill" + assert select(core, SimpleNamespace(num_contexts=7)) == "prefill" + + +_DECODE_FIELDS = ( + "num_tokens", + "num_seqs", + "kv_cache_manager", + "msa_idx_k_cache", + "msa_out_cache_loc", + "msa_decode_span", + "msa_block_table", + "msa_seq_lens_cuda", + "_msa_short_ok", + "msa_short_table", + "msa_proxy_max_score_view", + "msa_worst_case_max_k_tiles", + "msa_n_valid_blocks", +) + + +def _decode_metadata(num_tokens: int) -> SimpleNamespace: + """The decode projection's fields and nothing else (no ``num_contexts`` / ``num_ctx_tokens``).""" + md = SimpleNamespace(**{name: object() for name in _DECODE_FIELDS}) + md.num_tokens = num_tokens + return md + + +def test_the_decode_projection_states_the_phase_and_reads_no_count_back() -> None: + md = _decode_metadata(4) + step = dispatch.DecodeTarget(SimpleNamespace()).step_args(md) + assert step["num_contexts"] == 0 and step["num_ctx_tokens"] == 0 + assert step["metadata"] is md and step["msa_short_ok"] is md._msa_short_ok + + +@pytest.mark.parametrize( + "rows,embeds", [(49, False), (64, False), (4, True)], ids=["49_rows", "64_rows", "embeds"] +) +def test_a_wide_generation_step_runs_every_layer_generically_in_decode( + monkeypatch, rows, embeds +) -> None: + """Past the decode kernels' rows (or with embeddings) the decode target runs its own generic layer for every + layer, with its own projection, and never the decode-kernel layer.""" + layers, seen = 3, [] + core = SimpleNamespace(num_layers=layers, w={"embed": None}) + monkeypatch.setattr(dispatch, "m3_embedding", lambda ids, table: ("embedded", int(ids.numel()))) + + def generic(self, i, position_ids, hidden, residual, step, spec_metadata, capture): + seen.append((i, step["num_contexts"], step["metadata"])) + return ("hidden", i), ("residual", i) + + def kernels(self, *args): + raise AssertionError("the decode-kernel layer ran on a wide step") + + monkeypatch.setattr(dispatch.DecodeTarget, "_generic_layer", generic) + monkeypatch.setattr(dispatch.DecodeTarget, "_layer", kernels) + md = _decode_metadata(rows) + ids = torch.arange(rows, dtype=torch.int32) + out = dispatch.DecodeTarget(core).forward( + md, None if embeds else ids, None, ("embeds",) if embeds else None, None, None + ) + assert out == ("hidden", layers - 1) + assert seen == [(i, 0, md) for i in range(layers)] + + +def test_uncovered_never_raises_and_logs_once(log) -> None: + dispatch.set_require(True) + for _ in range(3): + dispatch.uncovered("m3_fused_score_select", "more than 2 requests: scorer + selector") + assert len(log.info) == 1 and not log.warning + + +class _Core: + """The parts of the core ``DecodeTarget``'s embedding and push targets read: no embedding owners built.""" + + def __init__(self): + self.embed_owners = {} + self.w = {"embed": torch.arange(64 * 4, dtype=torch.float32).view(64, 4)} + self.mapping = None + self.hidden = 6144 + + +@pytest.mark.parametrize("rows", (1, 8)) +def test_embedding_without_its_owner_raises_under_require(rows) -> None: + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=rf"'embed_firstnorm'.* {rows} rows.*owner" + ): + dispatch.DecodeTarget(_Core())._embed(torch.arange(rows, dtype=torch.int32), rows) + + +def test_embedding_without_its_owner_falls_back_and_logs_once(log) -> None: + target = dispatch.DecodeTarget(_Core()) + for _ in range(2): + hidden, residual = target._embed(torch.arange(3, dtype=torch.int32), 3) + assert residual is None and hidden.shape == (3, 4) + assert len(log.warning) == 1 and "embed_firstnorm at 3 rows" in log.warning[0] + + +def test_embedding_past_its_rows_is_uncovered_in_every_mode(log) -> None: + dispatch.set_require(True) + hidden, residual = dispatch.DecodeTarget(_Core())._embed(torch.arange(9, dtype=torch.int32), 9) + assert residual is None and hidden.shape == (9, 4) + assert len(log.info) == 1 and "embed_firstnorm" in log.info[0] and not log.warning + + +def test_a_push_without_its_workspace_raises_under_require(monkeypatch) -> None: + monkeypatch.setattr(m3_mc_allreduce, "get_workspace", lambda mapping, hidden: None) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'attn_push'.* 4 rows"): + dispatch.DecodeTarget(_Core())._attn_push_target(4) + + +def test_a_push_without_its_workspace_falls_back_and_logs_once(monkeypatch, log) -> None: + monkeypatch.setattr(m3_mc_allreduce, "get_workspace", lambda mapping, hidden: None) + target = dispatch.DecodeTarget(_Core()) + assert target._attn_push_target(4) is None + assert target._moe_push_targets(4) is None + assert target._attn_push_target(4) is None + assert len(log.warning) == 2 + assert "attn_push at 4 rows" in log.warning[0] and "moe_push at 4 rows" in log.warning[1] + + +def test_the_target_has_no_dispatch_module_left() -> None: + """The dispatch tables live in ``modeling.py``: no module of the target still reaches a ``dispatch`` module.""" + target_dir = Path(dispatch.__file__).parent + assert not (target_dir / "dispatch.py").exists() + stale = [] + for path in sorted(target_dir.rglob("*.py")): + for node in ast.walk(ast.parse(path.read_text())): + if isinstance(node, ast.ImportFrom) and any(a.name == "dispatch" for a in node.names): + stale.append(f"{path.relative_to(target_dir)}:{node.lineno}") + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "dispatch" + ): + stale.append(f"{path.relative_to(target_dir)}:{node.lineno}: dispatch.{node.attr}") + assert not stale, stale + + +# ------------------------------------------------------------------ the fused index score + top-16 select + + +class _FssConfig: + num_index_heads = 1 + sparse_index_dim = 128 + num_kv_heads = 1 + init_blocks = 0 + local_blocks = 1 + + +def _fss_core(): + """What ``DecodeTarget._fused_score_select`` reads of the core: layer 5's MSA backend configuration.""" + return SimpleNamespace(attn={5: SimpleNamespace(m3_config=_FssConfig())}) + + +def _fss_step(rows: int, qlen: int, max_tiles: int = 8192, n_valid: bool = True): + """A decode span's ``step_args`` entries as the fused score + select reads them (CPU stand-ins, never + launched).""" + return dict( + msa_idx_k_cache=lambda layer_idx: torch.zeros(4, 1, 128, 128, dtype=torch.float8_e4m3fn), + msa_worst_case_max_k_tiles=max_tiles, + msa_n_valid_blocks=torch.ones(rows * qlen, dtype=torch.int32) if n_valid else None, + msa_block_table=torch.zeros(rows, 4, dtype=torch.int32), + msa_seq_lens_cuda=torch.full((rows,), 300, dtype=torch.int32), + ) + + +def _fss(rows: int, qlen: int = 4, dtype=torch.float8_e4m3fn, **md): + idx_q = torch.zeros(rows * qlen, 128, dtype=dtype) + return dispatch.DecodeTarget(_fss_core())._fused_score_select( + 5, idx_q, _fss_step(rows, qlen, **md), rows, qlen, rows * qlen + ) + + +def test_fused_score_select_covers_one_and_two_requests() -> None: + assert dispatch.FUSED_SCORE_SELECT.keys == frozenset({1, 2}) + assert {c.spec["batch"] for c in dispatch.FUSED_SCORE_SELECT.entry.CELLS} == {1, 2} + + +@pytest.mark.parametrize( + "kwargs,reason", + [ + (dict(dtype=torch.bfloat16), "not FP8"), + (dict(qlen=5), "not the kernel's geometry"), + (dict(max_tiles=16384), "do not fit the stage ring"), + (dict(n_valid=False), "not one decode span"), + ], + ids=["bf16_index", "qlen5", "ring", "no_valid_counts"], +) +@pytest.mark.parametrize("rows", [1, 2]) +def test_a_covered_span_failing_a_precondition_raises_under_require(rows, kwargs, reason) -> None: + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=rf"'m3_fused_score_select'.* {rows} rows.*{reason}" + ): + _fss(rows, **kwargs) + + +def test_a_covered_span_failing_a_precondition_falls_back_and_logs_once(log) -> None: + for _ in range(2): + assert _fss(2, dtype=torch.bfloat16) is None + assert len(log.warning) == 1 and "m3_fused_score_select at 2 rows" in log.warning[0] + + +def test_a_capture_without_the_select_scratch_raises_under_require(monkeypatch) -> None: + monkeypatch.setattr(m3_fused_score_select, "buffers", lambda device, max_tiles: None) + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=r"'m3_fused_score_select'.* 1 rows.*scratch.*captured" + ): + _fss(1) + + +@pytest.mark.parametrize("rows", [3, 8]) +def test_a_span_past_two_requests_is_uncovered_in_every_mode(rows, log) -> None: + dispatch.set_require(True) + assert _fss(rows) is None + assert len(log.info) == 1 and "m3_fused_score_select" in log.info[0] and not log.warning + + +# ------------------------------------------------------------------ the drafter's low-M linears + + +def test_draft_lowm_keys_are_the_tactic_rows() -> None: + tactics = draft_lowm_entry.TACTICS + assert set(tactics) == set(DRAFTER_LINEARS) == set(dispatch.DRAFT_LOWM_ROWS) + expected = frozenset( + (nk, m) for nk, table in tactics.items() for m in range(1, table[-1][0] + 1) + ) + assert dispatch.DRAFT_LOWM.keys == expected + for nk, table in tactics.items(): + rows = [entry[0] for entry in table] + assert rows == sorted(rows) and len(set(rows)) == len(rows), nk + for m in range(1, table[-1][0] + 1): + assert draft_lowm_entry.tactic_for(*nk, m) is not None, (nk, m) + assert draft_lowm_entry.tactic_for(*nk, table[-1][0] + 1) is None, nk + + +def _lowm_operands(nk, m: int): + """An [N, K] BF16 weight (a broadcast view: the dispatch reads its shape and dtype only) and an [m, K] input.""" + n, k = nk + weight = torch.zeros(1, 1, dtype=torch.bfloat16).expand(n, k) + return weight, torch.zeros(m, k, dtype=torch.bfloat16) + + +_IN_TREE = "the in-tree BF16 GEMM" + + +@pytest.fixture(autouse=True) +def _in_tree_gemm(monkeypatch): + """The in-tree BF16 GEMM a drafter linear falls back to, stubbed (nothing here launches a kernel).""" + monkeypatch.setattr(dispatch, "m3_bf16_linear", lambda x, weight: _IN_TREE) + + +def _draft_linear(weight, x): + return dispatch._draft_linear(x, weight) + + +@pytest.mark.parametrize("nk", DRAFTER_LINEARS, ids=str) +def test_draft_linear_without_its_prebuilt_kernel_raises_under_require(monkeypatch, nk) -> None: + monkeypatch.setattr(dispatch, "_DRAFT_LOWM_BUILT", set()) + weight, x = _lowm_operands(nk, 1) + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=r"'m3_draft_lowm'.* 1 rows.*not compiled before warmup" + ): + _draft_linear(weight, x) + + +def test_draft_linear_without_its_prebuilt_kernel_falls_back_and_logs_once( + monkeypatch, log +) -> None: + monkeypatch.setattr(dispatch, "_DRAFT_LOWM_BUILT", set()) + weight, x = _lowm_operands((6144, 1024), 2) + for _ in range(2): + assert _draft_linear(weight, x) is _IN_TREE + assert len(log.warning) == 1 and "m3_draft_lowm at 2 rows" in log.warning[0] + + +def test_draft_linear_on_a_non_bf16_activation_is_a_miss() -> None: + weight, x = _lowm_operands((6144, 1024), 1) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_draft_lowm'.*not BF16"): + _draft_linear(weight, x.float()) + + +def test_draft_linear_on_a_strided_activation_keeps_the_in_tree_gemm(log) -> None: + """The first pass's fc may read a column window of the captured hidden states: uncovered, never raised.""" + weight, x = _lowm_operands((6144, 1024), 2) + wide = x.new_zeros(2, 2048) + wide[:, :1024] = x + assert not wide[:, :1024].is_contiguous() + dispatch.set_require(True) + assert _draft_linear(weight, wide[:, :1024]) is _IN_TREE + assert len(log.info) == 1 and not log.warning + + +def test_draft_linear_past_its_rows_is_uncovered_in_every_mode(monkeypatch, log) -> None: + monkeypatch.setattr(dispatch, "_DRAFT_LOWM_BUILT", set()) + dispatch.set_require(True) + weight, x = _lowm_operands((6144, 1024), 3) # the o_proj table ends at 2 rows + assert _draft_linear(weight, x) is _IN_TREE + weight, x = _lowm_operands((1280, 12288), 9) + assert _draft_linear(weight, x) is _IN_TREE + weight, x = _lowm_operands((4096, 1024), 1) # not a drafter shape + assert _draft_linear(weight, x) is _IN_TREE + assert len(log.info) == 3 and not log.warning + + +def test_draft_linear_on_a_batched_activation_is_uncovered(monkeypatch, log) -> None: + monkeypatch.setattr(dispatch, "_DRAFT_LOWM_BUILT", set()) + dispatch.set_require(True) + weight, x = _lowm_operands((6144, 1024), 1) + assert _draft_linear(weight, x[None]) is _IN_TREE + assert len(log.info) == 1 and not log.warning + + +def test_prebuild_drafter_linears(monkeypatch) -> None: + """Covered drafter linear calls compile before warmup; without a drafter nothing is built.""" + built = [] + monkeypatch.setattr(dispatch, "_prebuild_draft_lowm", lambda dm: built.append(dm) or 17) + worker = dispatch.M3Eagle3Worker.__new__(dispatch.M3Eagle3Worker) + assert worker.prebuild_drafter_linears(None) == 0 and not built + attn = SimpleNamespace() + drafter = SimpleNamespace(attn=attn) + assert worker.prebuild_drafter_linears(drafter) == 17 + assert built == [drafter] + + +# ------------------------------------------------------------------ the fused draft loop's preconditions + +_HEAD_ROWS = 64 +_HIDDEN = 32 + + +def _drafter(head_weight=None): + """A drafter with only what M3Eagle3Worker._m3_drafter_supported reads: its vocab-sharded BF16 head at TP8 (the + drafter's one-layer shape is validated when ``M3Eagle3Draft`` is built).""" + weight = ( + head_weight + if head_weight is not None + else torch.zeros((_HEAD_ROWS, _HIDDEN), dtype=torch.bfloat16) + ) + lm_head = SimpleNamespace( + weight=weight, bias=None, tp_mode=TensorParallelMode.COLUMN, padding_size=0, tp_size=8 + ) + return SimpleNamespace(lm_head=lm_head) + + +def _worker(): + """An M3Eagle3Worker past the step's verify: the fused verify left the first-pass gather rows and the fused + first draft inputs; nothing else of the in-tree worker is read.""" + worker = dispatch.M3Eagle3Worker.__new__(dispatch.M3Eagle3Worker) + worker.mapping = SimpleNamespace(tp_size=8, tp_rank=1, enable_attention_dp=False) + worker._d2t = None + worker._m3_draft_head = None + worker._m3_first_gather_ids = torch.zeros(1, dtype=torch.int64) + worker._m3_fused_first_inputs = True + return worker + + +_ATTN_MD = SimpleNamespace(kv_lens_cuda=None) +_SPEC_MD = SimpleNamespace(runtime_draft_len=3) + + +@pytest.fixture +def no_gpu_draft_head(monkeypatch): + """The fused head's CUDA-only lookups replaced: a 148-SM device, and a built exchange workspace.""" + monkeypatch.setattr(m3_draft_step, "_num_sms", lambda: 148) + monkeypatch.setattr(m3_draft_step, "_WS", object()) + + +def test_draft_loop_without_the_prebuilt_head_raises_under_require(no_gpu_draft_head) -> None: + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_draft_step'.* 4 rows.*prebuild"): + _worker()._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, _drafter()) + + +def test_draft_loop_without_the_prebuilt_head_falls_back_and_logs_once( + no_gpu_draft_head, log +) -> None: + worker = _worker() + for _ in range(2): + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, _drafter()) + assert len(log.warning) == 1 and "m3_draft_step at 4 rows" in log.warning[0] + + +def test_draft_loop_with_a_head_built_for_another_weight_raises_under_require( + no_gpu_draft_head, +) -> None: + worker = _worker() + assert worker.prebuild_draft_head(_drafter()) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_draft_step'"): + worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 2, _drafter()) + + +def test_draft_loop_takes_a_covered_step_once_the_head_is_built(no_gpu_draft_head) -> None: + dispatch.set_require(True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + head = worker._m3_draft_head + assert head.weight is drafter.lm_head.weight and head.vocab_offset == _HEAD_ROWS + assert torch.equal(draft_step_entry.from_fragments(head.wf), drafter.lm_head.weight) + for nreq in range(1, 9): + assert worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, nreq, drafter) + + +def test_draft_loop_streams_the_bf16_fragments_only(no_gpu_draft_head) -> None: + """The tp8 draft head is the fragment-ordered BF16 copy at every K (no coded copy is built, whatever K), it + decodes to the weight bit for bit, and the op refuses a coded stream.""" + dispatch.set_require(True) + g = torch.Generator().manual_seed(11) + weight = (torch.randn((_HEAD_ROWS, 128), generator=g) * 0.02).to(torch.bfloat16) + drafter = _drafter(weight) + worker = _worker() + assert worker.prebuild_draft_head(drafter) + head = worker._m3_draft_head + assert not hasattr(head, "coded") and not hasattr(head, "exc_off") + assert head.weight is weight and head.vocab_offset == _HEAD_ROWS + decoded = draft_step_entry.head_weight(head.wf, 128) + assert torch.equal(decoded.view(torch.int16), weight.view(torch.int16)) + with pytest.raises(NotImplementedError, match="BF16 fragments only"): + draft_step_entry.head_weight(head.wf, 128, torch.zeros(5, dtype=torch.int32)) + assert all(not c.spec["coded"] for c in draft_step_entry.m3_draft_step.CELLS) + for nreq in range(1, 9): + assert worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, nreq, drafter) + + +def test_draft_loop_past_its_requests_is_uncovered_in_every_mode(no_gpu_draft_head, log) -> None: + dispatch.set_require(True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 9, drafter) + assert not worker._m3_fused_loop_ok(_ATTN_MD, SimpleNamespace(runtime_draft_len=2), 4, drafter) + assert len(log.info) == 2 and not log.warning + + +def test_draft_loop_needs_the_fused_first_inputs(no_gpu_draft_head, log) -> None: + """A step whose first draft inputs came from the in-tree path (or that the fused verify did not prepare) drafts + with the in-tree loop, silently: it is not a draft-step shape.""" + dispatch.set_require(True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + worker._m3_fused_first_inputs = False + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + worker._m3_fused_first_inputs, worker._m3_first_gather_ids = True, None + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + assert not log.info and not log.warning + + +def test_an_unsupported_drafter_is_uncovered(no_gpu_draft_head, log) -> None: + dispatch.set_require(True) + worker = _worker() + padded = _drafter() + padded.lm_head.padding_size = 16 + for drafter in (padded,): + assert not worker.prebuild_draft_head(drafter) + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + assert log.info and not log.warning + + +def test_a_capture_without_the_exchange_workspace_raises_under_require(monkeypatch) -> None: + monkeypatch.setattr(m3_draft_step, "_num_sms", lambda: 148) + monkeypatch.setattr(m3_draft_step, "_WS", None) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_draft_step'.*workspace.*captured"): + worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + + +# ------------------------------------------------------------------ the fc fork (never early at tp8) + + +class _CaptureMetadata: + def __init__(self, nreq: int): + self.layers_to_capture = (1, 29, 56) + self.num_capture_tokens = nreq * 4 + self.runtime_draft_len = 3 + self.captured = [] + + def maybe_capture_hidden_states(self, layer_id, hidden_states, residual) -> None: + self.captured.append(layer_id) + + +def _fc_early_worker(monkeypatch, eligible: bool = True): + worker = dispatch.M3Eagle3Worker.__new__(dispatch.M3Eagle3Worker) + worker._m3_capture_norm = None + worker._m3_fc_norm_w = None + worker._m3_fc_early_drafter = None + worker._m3_fc_early_stream = None + forks = [] + monkeypatch.setattr(worker, "tail_eligible", lambda *args: eligible) + + def fork(md, num_tokens, draft_model): + forks.append((num_tokens, draft_model)) + return f"stream-{len(forks)}" + + monkeypatch.setattr(worker, "_m3_fc_fork", fork) + return worker, forks + + +@pytest.mark.parametrize("nreq,eligible", [(1, True), (2, True), (3, True), (8, True), (1, False)]) +def test_no_early_fork_at_tp8(monkeypatch, nreq, eligible) -> None: + """``FC_EARLY_MAX_ROWS`` is 0: no verify (not even one request's 4 rows) forks the drafter's fc at the last + captured layer; every captured layer still captures and the fc forks beside the verify head instead.""" + worker, forks = _fc_early_worker(monkeypatch, eligible) + md = _CaptureMetadata(nreq) + worker._m3_fc_early_stream = "stale" + worker.arm_capture_norm(SimpleNamespace(num_seqs=nreq), md, object(), object()) + assert worker._m3_fc_early_drafter is None + for layer_id in md.layers_to_capture: + worker.capture_hidden_states(md, layer_id, None) + worker.capture_hidden_states(md, dispatch.FINAL_HIDDEN_LAYER_ID, None) + assert md.captured == [1, 29, 56, dispatch.FINAL_HIDDEN_LAYER_ID] + assert not forks and worker._m3_fc_early_stream is None + + +def test_no_early_fork_with_a_final_state_capture(monkeypatch) -> None: + """The final-state column is written after the model returns, so captures that include it keep the fc at its + stock place instead of forking at the last captured decoder layer.""" + worker, forks = _fc_early_worker(monkeypatch) + md = _CaptureMetadata(1) + md.layers_to_capture = (dispatch.FINAL_HIDDEN_LAYER_ID, 29, 56) + worker.arm_capture_norm(SimpleNamespace(num_seqs=1), md, object(), object()) + for layer_id in (29, 56, dispatch.FINAL_HIDDEN_LAYER_ID): + worker.capture_hidden_states(md, layer_id, None) + assert not forks and worker._m3_fc_early_stream is None + + +# ------------------------------------------------------------------ the NVFP4 producer's coverage + + +def _producer_core(nvfp4: bool = True, heads: int = 8): + """What the producer dispatch reads off the core before its first precondition.""" + return SimpleNamespace(main_kv_is_nvfp4=nvfp4, heads=heads, kv_heads=1, index_heads=1) + + +def test_the_producer_covers_the_tp8_nvfp4_layers_only() -> None: + """``PRODUCER`` covers a sparse layer on an NVFP4 cache at the TP8 heads (8 query, one KV and one index head); + elsewhere the generic producer runs without a miss.""" + covered = dispatch._producer_covered + assert covered(_producer_core()) + assert not covered(_producer_core(nvfp4=False)) + assert not covered(_producer_core(heads=16)) + assert not covered(SimpleNamespace(main_kv_is_nvfp4=True, heads=8, kv_heads=2, index_heads=2)) + assert dispatch.PRODUCER.keys == frozenset({"nvfp4_tp8"}) + assert dispatch.PRODUCER.covers("nvfp4_tp8") + assert not dispatch.PRODUCER.covers("nvfp4_tp4") + + +def _produce(core): + # The packed TP8 projection: 8 query heads, K, V, index Q and index K of 128 each. + packed = torch.zeros(4, 12 * 128, dtype=torch.bfloat16) + return dispatch._sparse_qkv_producer(core, 3, packed, None, None, None, None) + + +def test_a_covered_producer_layer_missing_a_precondition_raises_under_require() -> None: + """A TP8 NVFP4 sparse layer whose step lacks a producer precondition (here the position ids) is a miss.""" + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=r"'m3_producer' covers this step at 4 rows" + ): + _produce(_producer_core()) + + +def _runtime(fused: bool, indexer: str = "fp8", kv: str = "fp8"): + """A model config with the runtime fields ``_validate_runtime`` reads.""" + sparse = SimpleNamespace( + algorithm="minimax_m3", + implementation="msa", + fuse_qkv_index_projection=fused, + indexer_kv_dtype=indexer, + ) + mode = SimpleNamespace( + has_fp8_kv_cache=lambda: kv == "fp8", has_fp4_kv_cache=lambda: kv == "nvfp4" + ) + return SimpleNamespace( + sparse_attention_config=sparse, quant_config=SimpleNamespace(quant_mode=mode) + ) + + +@pytest.mark.parametrize("fused", [True, False], ids=["five_way", "unfused"]) +def test_the_target_takes_both_projection_layouts(fused) -> None: + """The packed five-way projection and the separate QKV + BF16 index projections + (fuse_qkv_index_projection=False) are both this target's; a BF16 indexer or a BF16 KV cache is not.""" + dispatch._validate_runtime(_runtime(fused)) + dispatch._validate_runtime(_runtime(fused, kv="nvfp4")) + with pytest.raises(ValueError, match="indexer_kv_dtype='fp8'"): + dispatch._validate_runtime(_runtime(fused, indexer="bf16")) + with pytest.raises(ValueError, match="FP8 or NVFP4 KV cache"): + dispatch._validate_runtime(_runtime(fused, kv="bf16")) + + +@pytest.mark.parametrize("target", ["PrefillTarget", "DecodeTarget"]) +def test_an_unfused_layer_the_producer_covers_is_a_miss(target) -> None: + """As the built-in target: a layer the NVFP4 producer covers needs the five-way projection, so the unfused path is + a miss there (raised under require) before anything runs, in either target.""" + dispatch.set_require(True) + attention = getattr(dispatch, target)(_producer_core())._unfused_sparse_attention + with pytest.raises(dispatch.FastPathUnavailable, match="fuse_qkv_index_projection is disabled"): + if target == "PrefillTarget": + attention(3, None, None, None) + else: + attention(3, None, None, None, None) + + +def test_a_covered_producer_miss_falls_back_and_logs_once(log) -> None: + core = _producer_core() + assert _produce(core) is None and _produce(core) is None + assert len(log.warning) == 1 and "m3_producer at 4 rows" in log.warning[0] + + +@pytest.mark.parametrize("nvfp4,heads", [(False, 8), (True, 16)], ids=["fp8_cache", "tp4_heads"]) +def test_an_uncovered_producer_layer_falls_back_silently(log, nvfp4, heads) -> None: + dispatch.set_require(True) + assert _produce(_producer_core(nvfp4, heads)) is None + assert not log.warning + + +def test_a_covered_producer_layer_with_the_wrong_pool_layout_is_not_supported() -> None: + """The pool-layout check the covered branch takes before the generic producer (then a miss).""" + pool = torch.zeros(4, 2, 1, 128, 32, dtype=torch.uint8) + scales = torch.zeros(4, 2, 1, 128, 8, dtype=torch.uint8) + assert not dispatch._producer_layout(pool, scales) + good = torch.zeros(4, 2, 1, 128, 64, dtype=torch.uint8) + assert dispatch._producer_layout(good, scales) + assert not dispatch._producer_layout(good, scales.transpose(3, 4).contiguous().transpose(3, 4)) + + +# ------------------------------------------------------------------ the shell's speculative plumbing + + +def test_the_shell_installs_guided_decoding_on_the_worker() -> None: + """With speculation the guided decoder goes to the worker (the in-tree one-engine shell's delegation); without a + worker the shell reports it cannot take one.""" + shell_cls = dispatch.ModelingV2MiniMaxM3Nvfp4Sm103Tp8 + shell = shell_cls.__new__(shell_cls) + torch.nn.Module.__init__(shell) + installed = [] + shell.spec_worker = SimpleNamespace(set_guided_decoder=lambda g: installed.append(g) or True) + decoder = object() + assert shell.set_guided_decoder(decoder) and installed == [decoder] + shell.spec_worker = None + assert shell.set_guided_decoder(decoder) is False diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py new file mode 100644 index 000000000000..d1ee4df80e69 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode kernels the core builds before the executor's warmup (``prebuild_kernels``), and that a forward after it +never compiles one: the MoE engine of every ``MOE_ENGINE`` token count, the decode GEMVs with their push variants, and +the index scorer's prefetching kernel at every decode query length (the tp8 target never selects the scorer's +early-issue schedule: ``INDEX_SCORE_EARLY_MAX_REQUESTS`` is 0). + +The engine builds are stubbed (a real one needs a loaded checkpoint's expert buffers; the engine kernels have their +own tests); the GEMVs and the scorer compile for real on the GPU. +""" + +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers the trtllm ops +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops import ( # noqa: E501 + m3_index_score as index_score_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_index_score import ( # noqa: E501 + m3_index_score, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 + DENSE_DOWN, + O_PROJ, + QKV_SPARSE, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.mx_gemv import ( # noqa: E501 + op as mx_gemv, +) + +dispatch = modeling + +assert torch.cuda.is_available(), "the prebuild checks require a CUDA device" + +INDEX_HEADS = 1 # index heads per rank at TP8 +INDEX_DIM = 128 +DRAFT_LEN = 3 +_CORE_METHODS = ( + "prebuild_kernels", + "_prebuild_engines", + "_try_build", + "engine", + "_lookup", + "_scorer_variants", + "_prebuild_index_scorer", + "decode_query_lens", +) + + +@pytest.fixture(autouse=True) +def _restore_modes(): + require, frozen = dispatch.require_mode(), mx_gemv.frozen() + yield + dispatch.set_require(require) + mx_gemv.freeze(frozen) + dispatch.reset_logged() + + +class _FakeSparseAttn: + def __init__(self, indexer_kv_dtype: str): + self.indexer_kv_dtype = indexer_kv_dtype + self.m3_config = SimpleNamespace(num_index_heads=INDEX_HEADS, sparse_index_dim=INDEX_DIM) + + +def _core(num_layers: int = 3, indexer_kv_dtype: str = "fp8", built=None, fail=False): + """The parts of ``MiniMaxM3Tp8ModelingV2`` the prebuild reads: one rank (the build agreement takes no + collective), MoE + sparse layers 1.. (layer 0 dense), no projections (the GEMV prebuild has its own test), the + speculative config; the engine builds stubbed.""" + core = SimpleNamespace( + mapping=SimpleNamespace(tp_size=1, enable_attention_dp=False), + moe_layers=frozenset(range(1, num_layers)), + sparse_layers=frozenset(range(1, num_layers)), + attn=[_FakeSparseAttn(indexer_kv_dtype) for _ in range(num_layers)], + w={"embed": torch.empty(1, device="cuda")}, + mxfp8=SimpleNamespace(quant_method=SimpleNamespace(use_cutlass=True)), + model_config=SimpleNamespace( + spec_config=SimpleNamespace(max_draft_len=DRAFT_LEN, max_total_draft_tokens=DRAFT_LEN) + ), + moe_engines={}, + projections=lambda: [], + ) + for name in _CORE_METHODS: + setattr(core, name, getattr(modeling.MiniMaxM3Tp8ModelingV2, name).__get__(core)) + built = [] if built is None else built + + def build(i, num_tokens, shared=True): + if fail: + raise ValueError("the routed experts are not TRTLLM-Gen NVFP4 SwiGLU-OAI") + built.append(("moe", i, num_tokens, shared)) + return ("engine", i, num_tokens) + + core._build_engine = build + return core + + +def _forbid_compiles(monkeypatch, core=None) -> None: + """Every compile / build entry point a forward could reach raises.""" + import cutlass.cute as cute + + def boom(*args, **kwargs): + raise AssertionError("a compile entry point ran after prebuild") + + monkeypatch.setattr(cute, "compile", boom) + monkeypatch.setattr(mx_gemv, "ensure_compiled", boom) + monkeypatch.setattr(mx_gemv, "precompile", boom) + monkeypatch.setattr(index_score_kernel.IndexScoreRunner, "_compile", boom) + if core is not None: + core._build_engine = boom + + +def test_decode_query_lens_follow_the_speculative_config() -> None: + core = _core() + assert core.decode_query_lens() == [1, 1 + DRAFT_LEN] + core.model_config.spec_config = None + assert core.decode_query_lens() == [1] + core.model_config.spec_config = SimpleNamespace( + max_draft_len=4, max_total_draft_tokens=4, draft_len_schedule={1: 4, 8: 2} + ) + assert core.decode_query_lens() == [1, 3, 5] + + +def test_prebuild_builds_every_engine_count_and_scorer_variant() -> None: + built = [] + core = _core(built=built) + times = core.prebuild_kernels() + + assert mx_gemv.frozen(), "prebuild must close the GEMV compiles" + + # Layer 0 is dense; every MoE layer gets every count, in layer order. + def per_layer(i): + return [("moe", i, n, True) for n in range(1, dispatch.MOE_ENGINE_MAX_TOKENS + 1)] + + assert built == per_layer(1) + per_layer(2) + assert set(core.moe_engines[2]) == {(n, True, True) for n in dispatch.MOE_ENGINE.keys} + assert not hasattr(core, "moe16_engines") + + want = { + (torch.float8_e4m3fn, INDEX_HEADS, qlen, INDEX_DIM, modeling.INDEX_SCORE_SPLIT_K, False) + for qlen in (1, 1 + DRAFT_LEN) + } + assert set(times["scorer_keys"]) == want + assert want <= set(index_score_kernel.IndexScoreRunner.kernel_cache) + for key in ("gemv_s", "moe_s", "scorer_s"): + assert times[key] >= 0.0 + + # A second build builds no engine again. + built.clear() + core.prebuild_kernels() + assert built == [] + + +def test_engine_lookups_after_prebuild_never_build(monkeypatch) -> None: + core = _core() + core._prebuild_engines(5) + _forbid_compiles(monkeypatch, core) + for require in (False, True): + dispatch.set_require(require) + for n in sorted(dispatch.MOE_ENGINE.keys): + assert core.engine(5, n, shared=True) == ("engine", 5, n) + + +def test_missing_engine_is_a_miss_not_a_build(monkeypatch) -> None: + core = _core() + core.moe_engines[7] = {} + _forbid_compiles(monkeypatch, core) + dispatch.set_require(False) + assert core.engine(7, 3, shared=True) is None + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="not built before warmup"): + core.engine(7, 3, shared=True) + + +def test_failed_engine_build_is_logged_then_missed() -> None: + core = _core(fail=True) + core._prebuild_engines(9) + assert all(v is None for v in core.moe_engines[9].values()) and len(core.moe_engines[9]) == 8 + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="build failed"): + core.engine(9, 1, shared=True) + + +def _projection(nk, push): + n, k = nk + w = (torch.randn(n, k, device="cuda") / k**0.5).to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + return SimpleNamespace(weight=wq.view(n, k), weight_scale=wsf, push=push) + + +@pytest.mark.parametrize( + "nk,push", + ((O_PROJ, "attn"), (QKV_SPARSE, None), (DENSE_DOWN, None)), + ids=("o_proj", "qkv_sparse", "dense_down"), +) +def test_gemvs_after_prebuild_never_compile(monkeypatch, nk, push) -> None: + use_pdl = True + module = _projection(nk, push) + variants = modeling._gemv_variants(*nk, push, 8) + pushes = [v for v in variants if v[2] is not None] + if push is None: + assert not pushes + else: + covered = dispatch.ATTN_PUSH.keys if push == "attn" else dispatch.MOE_PUSH.keys + assert {rows for rows, _, _ in pushes} == set(covered) + assert {p[1] for _, _, p in pushes} == {b for b in (8, 16) if b <= max(covered)} + mx_gemv.freeze(False) + # As prebuild_kernels compiles one projection. + mx_gemv.ensure_compiled(module.weight, module.weight_scale, use_pdl) + for rows, x_bf16, ar in variants: + mx_gemv.precompile(module.weight, module.weight_scale, rows, x_bf16, use_pdl=use_pdl, ar=ar) + mx_gemv.freeze(True) + _forbid_compiles(monkeypatch) + + n, k = nk + for rows, x_bf16, ar in variants: + if ar is not None: + # A push needs a multicast exchange to launch: its kernel is compiled (the forward finds it). + resolved = mx_gemv._resolve(rows, n, k, x_bf16, True, None, use_pdl, (0, 0, 0, *ar)) + if resolved is not None: + assert mx_gemv.is_compiled(resolved[0]), (rows, x_bf16, ar) + continue + if x_bf16: + x = torch.randn(rows, k, device="cuda").to(torch.bfloat16) + out = mx_gemv.m3_mx_gemv_bf16(x, module.weight, module.weight_scale, use_pdl=use_pdl) + else: + xq, xsf = torch.ops.trtllm.mxfp8_quantize( + torch.randn(rows, k, device="cuda").to(torch.bfloat16), True + ) + out = mx_gemv.m3_mx_gemv( + xq.view(rows, k), xsf, module.weight, module.weight_scale, use_pdl=use_pdl + ) + served = mx_gemv.serves(n, k, rows, x_bf16) + assert (out is not None) == served, (rows, x_bf16) + assert mx_gemv.last_uncompiled is None + torch.cuda.synchronize() + + +def test_frozen_gemv_declines_an_uncompiled_kernel(monkeypatch) -> None: + module = _projection(DENSE_DOWN, None) + n, k = DENSE_DOWN + mx_gemv.freeze(True) + _forbid_compiles(monkeypatch) + # A launch option nothing prebuilt (no early trigger, without a push): declined, never compiled. + xq, xsf = torch.ops.trtllm.mxfp8_quantize( + torch.randn(3, k, device="cuda").to(torch.bfloat16), True + ) + key = mx_gemv._resolve(3, n, k, False, False, None, True, None)[0] + if mx_gemv.is_compiled(key): + pytest.skip("this process already compiled the probe kernel") + out = mx_gemv.m3_mx_gemv( + xq.view(3, k), xsf, module.weight, module.weight_scale, trigger_early=False, use_pdl=True + ) + assert out is None and mx_gemv.last_uncompiled == key + + +@pytest.mark.parametrize("indexer_kv_dtype", ("fp8", "bf16")) +def test_scorer_after_prebuild_never_compiles(monkeypatch, indexer_kv_dtype) -> None: + core = _core(indexer_kv_dtype=indexer_kv_dtype) + keys = core._prebuild_index_scorer(core.decode_query_lens(), torch.device("cuda")) + # The decode query lengths 1 and 1 + DRAFT_LEN, the prefetching kernel only. + assert len(keys) == 2 and not any(key[-1] for key in keys) + _forbid_compiles(monkeypatch, core) + dtype = torch.float8_e4m3fn if indexer_kv_dtype == "fp8" else torch.bfloat16 + for _, heads, qlen, dim, split_k, early in keys: + rows = 1 if early else 2 + blocks = 3 + idx_q = torch.randn(rows * qlen, heads, dim, device="cuda").to(dtype) + idx_k = torch.randn(rows * blocks, 1, 128, dim, device="cuda").to(dtype) + block_table = torch.arange(rows * blocks, dtype=torch.int32, device="cuda").view( + rows, blocks + ) + seq_lens = torch.full((rows,), blocks * 128 - 5, dtype=torch.int32, device="cuda") + score = torch.empty(heads, blocks, rows * qlen, dtype=torch.float32, device="cuda") + m3_index_score( + idx_q, + idx_k.squeeze(1), + block_table, + seq_lens, + score.transpose(1, 2), + qlen, + split_k, + early, + ) + torch.cuda.synchronize() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py new file mode 100644 index 000000000000..81972c1c503e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode robustness without a GPU kernel: the multicast all-reduce epilogue options keep the cluster-barrier bit, +an engine build that fails on one TP rank drops that engine on every rank (two ranks, and the tp8 target's eight, +simulated with threads), and the Eagle3 fc fork's output is recorded on the stream that reads it.""" + +import threading +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops import ( # noqa: E501 + m3_mc_allreduce as mc, +) + + +@pytest.fixture(autouse=True) +def _restore_modes(): + require = modeling.require_mode() + yield + modeling.set_require(require) + modeling.reset_logged() + + +@pytest.mark.parametrize("finalize", [False, True]) +@pytest.mark.parametrize("bucket", [8, 16, 32, 48]) +def test_ar_opts_keep_cluster_barrier_bit(monkeypatch, bucket, finalize): + # Tables without bit 2 and with buckets missing: every lookup, fallback included, keeps the barrier. + monkeypatch.setattr(mc, "AR_OPTS_BY_M", (0, 1, 4, 5, 8, 9, 12, 13)) + monkeypatch.setattr(mc, "AR_OPTS_BY_BUCKET", {}) + monkeypatch.setattr(mc, "AR_FIN_OPTS_BY_BUCKET", {}) + for m in range(1, min(bucket, 8) + 1): + opts = mc._ar_opts(m, bucket, finalize=finalize) + assert opts & mc._OPT_CLUSTER_BARRIER and 0 <= opts <= 15, (m, bucket, finalize, opts) + + +def test_ar_opts_shipped_tables_unchanged(): + for m in range(1, 9): + assert mc._ar_opts(m, 8) == mc.AR_OPTS_BY_M[m - 1] + for b, o in mc.AR_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b) == o + for b, o in mc.AR_FIN_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b, finalize=True) == o + + +class _FakeComm: + """An MPI-like TP communicator for ``world`` threads: allreduce sums the ranks' ints.""" + + def __init__(self, world): + self.world = world + self.barrier = threading.Barrier(world) + self.lock = threading.Lock() + self.vals = [] + self.calls = 0 + + def allreduce(self, v): + with self.lock: + self.vals.append(v) + self.calls += 1 + self.barrier.wait() + total = sum(self.vals) + self.barrier.wait() + with self.lock: + self.vals.clear() + self.barrier.wait() + return total + + def Get_size(self): # noqa: N802 (MPI name) + return self.world + + +def _mapping(tp_size=2, adp=False): + return SimpleNamespace(tp_size=tp_size, enable_attention_dp=adp, tp_group=list(range(tp_size))) + + +def _run_ranks(fns): + out, errs = [None] * len(fns), [] + + def run(i): + try: + out[i] = fns[i]() + except BaseException as e: # noqa: BLE001 + errs.append(e) + + threads = [threading.Thread(target=run, args=(i,)) for i in range(len(fns))] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + assert not errs, errs + return out + + +def _patch_comm(monkeypatch, world: int) -> _FakeComm: + from tensorrt_llm._torch.distributed import ops + + comm = _FakeComm(world) + monkeypatch.setattr(mc, "_tp_comm", lambda mapping: comm) + monkeypatch.setattr(ops, "mpi_disabled", lambda: False) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + return comm + + +@pytest.fixture +def two_rank_comm(monkeypatch): + return _patch_comm(monkeypatch, 2) + + +def test_tp_agree_single_rank_and_adp_take_no_collective(monkeypatch): + monkeypatch.setattr(mc, "_tp_comm", mock.Mock(side_effect=AssertionError("no collective"))) + assert mc.tp_agree(_mapping(1), True) is True + assert mc.tp_agree(_mapping(1), False) is False + assert mc.tp_agree(_mapping(8, adp=True), False) is False + + +def test_tp_agree_refuses_capture(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError): + mc.tp_agree(_mapping(2), True) + + +@pytest.mark.parametrize("oks", [(True, True), (True, False), (False, True), (False, False)]) +def test_tp_agree_built_two_ranks(two_rank_comm, oks): + built = [object() if ok else None for ok in oks] + out = _run_ranks([lambda i=i: mc.tp_agree_built(_mapping(2), built[i], "t") for i in range(2)]) + assert out == (built if all(oks) else [None, None]) + + +def _core(fail: bool, world: int = 2) -> SimpleNamespace: + """The engine-building half of the core (``MiniMaxM3Tp8ModelingV2``) on one of ``world`` simulated TP ranks, + its engine builds stubbed (and failing where ``fail``).""" + core = SimpleNamespace(mapping=_mapping(world), moe_engines={}) + cls = modeling.MiniMaxM3Tp8ModelingV2 + + def build(i, num_tokens, shared=True): + if fail: + raise RuntimeError("compile failed on this rank") + return ("engine", num_tokens) + + core._build_engine = build + for name in ("_prebuild_engines", "_try_build", "engine", "_lookup"): + setattr(core, name, getattr(cls, name).__get__(core)) + return core + + +@pytest.mark.parametrize( + "world,failing", + [(2, ()), (2, (1,)), (8, ()), (8, (5,)), (8, (0, 7))], + ids=["w2_ok", "w2_rank1_fails", "w8_ok", "w8_rank5_fails", "w8_ranks0_7_fail"], +) +def test_engine_build_agreed_across_ranks(monkeypatch, world, failing): + comm = _patch_comm(monkeypatch, world) + cores = [_core(rank in failing, world) for rank in range(world)] + _run_ranks([lambda c=c: c._prebuild_engines(3) for c in cores]) + keys = sorted(cores[0].moe_engines[3]) + assert keys and comm.calls == world * len(keys) + assert not any(hasattr(c, "moe16_engines") for c in cores) + for c in cores[1:]: + assert c.moe_engines[3].keys() == cores[0].moe_engines[3].keys() + for key in keys: + for c in cores: + assert (c.moe_engines[3][key] is None) == bool(failing), key + # Every rank then takes the same path in the forward: the engine, or (after a miss) the generic MoE. + modeling.set_require(False) + for n in sorted(modeling.MOE_ENGINE.keys): + got = [c.engine(3, n, shared=True) is None for c in cores] + assert len(set(got)) == 1, (n, got) + + +@pytest.mark.parametrize("capturing", [False, True]) +def test_fc_join_records_output_on_main(monkeypatch, capturing): + M3Eagle3Worker = modeling.M3Eagle3Worker + main, fork = mock.Mock(name="main"), mock.Mock(name="fork") + hidden = mock.Mock(name="hidden", is_cuda=True) + monkeypatch.setattr(torch.cuda, "current_stream", lambda *a: main) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: capturing) + worker = SimpleNamespace(_m3_fc_hidden=hidden) + M3Eagle3Worker._m3_fc_join(worker, fork) + main.wait_stream.assert_called_once_with(fork) + if capturing: + hidden.record_stream.assert_not_called() + else: + hidden.record_stream.assert_called_once_with(main) + worker._m3_fc_hidden = None + M3Eagle3Worker._m3_fc_join(worker, fork) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py new file mode 100644 index 000000000000..6d563ff9f831 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp8 target's files that are byte copies of the tp4 target's stay so. + +The two MiniMax-M3 targets share no code by design (``test_modeling_v2_claims.test_targets_do_not_share_files``): the +tp8 target was made as a copy of the tp4 one and keeps its own copy of every file, tp8-specific or not. The files +listed here are still identical, so a fix to one copy applies to the other unchanged. This test fails on the first +edit to either side: a change to one copy must be considered for the other -- applied to both, or the file dropped +from this list with the reason the copies now differ. + +CPU only: it reads the files by path relative to the packages. +""" + +from pathlib import Path + +import pytest + +import tensorrt_llm._torch._experimental.modeling_v2 as _modeling_v2 + +# Nothing of either target is imported (no GPU, no built extensions needed): the paths are the packages', resolved +# from the imported parents rather than from this file, which lives under tests/. +pytestmark = pytest.mark.cpu_only + +_M3 = Path(_modeling_v2.__file__).resolve().parent / "models" / "minimax_m3" +_TP4 = _M3 / "minimax_m3_nvfp4__sm_103__tp4" +_TP8 = _M3 / "minimax_m3_nvfp4__sm_103__tp8" + +#: Paths relative to the target directory that are byte copies of the tp4 target's. +#: (``ops/attention/m3_msa_attention.py`` left the list: its reference dequantizes NVFP4 pages with each target's own +#: long-context sparse decode entry, ``m3_sparse_decode_lc`` at TP4 and ``m3_sparse_decode_lc8`` at TP8.) +IDENTICAL_TO_TP4 = ( + "ops/embedding_firstnorm/__init__.py", + "ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py", + "ops/embedding_firstnorm/op.py", + "ops/long_select.py", + "ops/m3_force_vec.py", + "ops/m3_index_score.py", + "ops/m3_prequant.py", + "ops/m3_router_gemv.py", + "ops/m3_sampled_accept.py", + "ops/m3_shard_accept.py", + "ops/m3_spec_tail.py", + "ops/msa_kv_update.py", + "ops/mx_gemv/m3_mx_gemv_kernel.py", + "ops/triton_sparse_decode_mq.py", + "ops/__init__.py", + "ops/activation/__init__.py", + "ops/attention/__init__.py", + "ops/attention/m3_draft_attention.py", + "ops/attention/m3_fp8_indexer_qk_norm_rope.py", + "ops/attention/m3_select_blocks.py", + "ops/attention/msa_kv_update.py", + "ops/comm/__init__.py", + "ops/embedding/__init__.py", + "ops/embedding/m3_embedding.py", + "ops/gemm/__init__.py", + "ops/gemm/m3_fp32_linear.py", + "ops/moe/__init__.py", + "ops/moe/m3_router_gemv_fp4.py", + "ops/moe/m3_router_gemv.py", + "ops/moe/m3_router_tc.py", + "ops/norm/__init__.py", + "ops/norm/m3_embedding_firstnorm.py", + "ops/spec/__init__.py", + "ops/spec/m3_capture_norm.py", + "ops/spec/m3_chunk_norm.py", + "ops/spec/m3_draft_input.py", + "ops/spec/m3_force_vec.py", + "ops/spec/m3_slot_scatter.py", + "ops/spec/m3_step_gather.py", +) + + +def test_the_targets_are_where_this_test_reads_them() -> None: + for path in (_TP4 / "modeling.py", _TP8 / "modeling.py"): + assert path.is_file(), path + + +def _differs(a: Path, b: Path) -> str: + """Empty when the two files are byte-identical, else where they first differ.""" + x, y = a.read_bytes(), b.read_bytes() + if x == y: + return "" + first = next((i for i, (p, q) in enumerate(zip(x, y)) if p != q), min(len(x), len(y))) + line = x[:first].count(b"\n") + 1 + return f"{a} and {b} differ from line {line} ({len(x)} vs {len(y)} bytes)" + + +@pytest.mark.parametrize("relative", IDENTICAL_TO_TP4) +def test_the_copy_is_the_tp4_file(relative) -> None: + ours, theirs = _TP8 / relative, _TP4 / relative + assert ours.is_file(), f"{relative}: gone from the tp8 target; drop it from IDENTICAL_TO_TP4" + assert theirs.is_file(), f"{relative}: gone from the tp4 target; drop it from IDENTICAL_TO_TP4" + diff = _differs(ours, theirs) + assert not diff, ( + f"{diff}. The targets share no code, so a change to one copy must be considered for the other: apply it to " + "both, or drop the file from IDENTICAL_TO_TP4 stating why the copies now differ." + ) + + +def test_the_list_names_every_identical_file() -> None: + """A file that became identical again (or was copied unchanged later) is listed, so it is guarded too.""" + unlisted = sorted( + str(path.relative_to(_TP8)) + for path in _TP8.rglob("*.py") + if "__pycache__" not in path.parts + and str(path.relative_to(_TP8)) not in IDENTICAL_TO_TP4 + and (_TP4 / path.relative_to(_TP8)).is_file() + and not _differs(path, _TP4 / path.relative_to(_TP8)) + ) + assert not unlisted, f"byte copies of the tp4 target not in IDENTICAL_TO_TP4: {unlisted}" diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_front_op_matrix.py b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_front_op_matrix.py index d19d722ff2f2..79c3e0fd6459 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_front_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_front_op_matrix.py @@ -9,14 +9,11 @@ import path to reach ``_rank_job``. """ -import sys -from pathlib import Path - import pytest import torch -sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "comm")) -import _rank_job # noqa: E402 +__extra_import_path__ = ["../comm"] +import _rank_job def _is_sm100() -> bool: diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m1.py b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m1.py index 9487eb8c60d2..c1e686861dd9 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m1.py +++ b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m1.py @@ -7,10 +7,9 @@ import importlib import zlib +__extra_import_path__ = ["..", "."] import _k3_moe_engines as engines import pytest - -__extra_import_path__ = [".."] from _validating import validating k3_moe_m1 = importlib.import_module( diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m2.py b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m2.py index 493eb1a5272f..f4cc27a4bad9 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m2.py +++ b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m2.py @@ -7,10 +7,9 @@ import importlib import zlib +__extra_import_path__ = ["..", "."] import _k3_moe_engines as engines import pytest - -__extra_import_path__ = [".."] from _validating import validating k3_moe_m2 = importlib.import_module( diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_chunk_kda.py b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_chunk_kda.py index 79ba6e832536..71f48ba6d271 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_chunk_kda.py +++ b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_chunk_kda.py @@ -5,7 +5,6 @@ import zlib from typing import Dict, List -import _kda_reference as kr import pytest import torch @@ -17,7 +16,8 @@ ) from tensorrt_llm._torch.modules.fla.index import prepare_chunk_indices -__extra_import_path__ = [".."] +__extra_import_path__ = ["..", "."] +import _kda_reference as kr from _validating import validating assert torch.cuda.is_available(), "chunk_kda requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_attn.py b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_attn.py index 8b7b3a728302..5680838921a5 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_attn.py +++ b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_attn.py @@ -23,7 +23,6 @@ import importlib import zlib -import _kda_cells as kc import pytest import torch @@ -32,7 +31,8 @@ K3KdaBuffers, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["..", "."] +import _kda_cells as kc from _validating import validating _attn = importlib.import_module( diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_decode_attn.py b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_decode_attn.py index 19683e0b84ea..729624da1591 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_decode_attn.py +++ b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_decode_attn.py @@ -20,7 +20,6 @@ import importlib import zlib -import _kda_cells as kc import pytest import torch import torch.nn.functional as F @@ -37,7 +36,8 @@ kda_decode, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["..", "."] +import _kda_cells as kc from _validating import validating _attn = importlib.import_module( diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_verify.py b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_verify.py index 754db075f388..01e835e6bbc9 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_verify.py +++ b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_verify.py @@ -20,11 +20,11 @@ import importlib import zlib -import _kda_cells as kc import pytest import torch -__extra_import_path__ = [".."] +__extra_import_path__ = ["..", "."] +import _kda_cells as kc from _validating import validating k3_kda_verify = importlib.import_module( diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_decode.py b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_decode.py index 1f105019ae8a..088cc6e2f607 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_decode.py +++ b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_decode.py @@ -12,7 +12,6 @@ import zlib -import _kda_cells as kc import pytest import torch @@ -20,7 +19,8 @@ kda_decode, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["..", "."] +import _kda_cells as kc from _validating import validating pytestmark = pytest.mark.skipif(not kc.sm100(), reason="certified on SM100 only") diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_prefill.py b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_prefill.py index 78f91c4a4e48..2c1b7011b16c 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_prefill.py +++ b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_prefill.py @@ -6,7 +6,6 @@ import zlib from typing import List -import _kda_reference as kr import pytest import torch @@ -14,7 +13,8 @@ kda_prefill, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["..", "."] +import _kda_reference as kr from _validating import validating assert torch.cuda.is_available(), "kda_prefill requires a CUDA device" 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 dcf8bce2550c..cb9e43bb2c6b 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -32,8 +32,11 @@ from transformers import PretrainedConfig from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + _RESOLVED_MODES, + ModelingV2Mode, NoTarget, modeling_v2_resolve, + resolved_mode, ) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp import ( routing as qwen4_exp_routing, @@ -43,6 +46,7 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_auto import AutoModelForCausalLM from tensorrt_llm.llmapi.llm_args import ( + CacheTransceiverConfig, KvCacheConfig, MTPDecodingConfig, NGramDecodingConfig, @@ -356,6 +360,14 @@ def test_r1_near_misses_do_not_match(mapping_kwargs, missed): modeling_v2_resolve(config, _llm_args("require")) +@pytest.mark.parametrize("mode", ["auto", "require"]) +def test_the_target_reads_the_mode_it_was_resolved_under(mode): + """A target raises instead of falling back under require, so it has to know which mode built it.""" + config = _model_config(_gpt_oss_config()) + target = modeling_v2_resolve(config, _llm_args(mode)) + assert resolved_mode(target) is ModelingV2Mode(mode) + + def test_resolve_class_reads_the_decision_off_the_config(): """The loader decides and writes the class on the ModelConfig; the class lookup only has to honour it.""" @@ -469,6 +481,48 @@ def test_an_unknown_mode_is_rejected_when_the_arguments_are_built(): _llm_args("yes") +def test_the_default_mode_is_off(): + assert TorchLlmArgs(model=_DUMMY_MODEL).modeling_v2 == "off" + + +def test_no_target_is_a_value_error(): + """A refusal is the ValueError ``require`` has always raised, so callers that catch ValueError still do.""" + assert issubclass(NoTarget, ValueError) + + +class _RefusingTarget: + """A target that refuses its runtime at construction.""" + + def __init__(self, config): + raise NoTarget("this runtime is not the one the target is written for") + + +class _BuiltIn: + def __init__(self, config): + self.config = config + + +@pytest.mark.parametrize("mode", ["auto", "require"]) +def test_a_refusal_at_construction_falls_back_unless_required(monkeypatch, mode): + """A target's construction-time refusal builds the built-in implementation under ``auto`` (the mode + ``default`` records) and is the error under ``require``.""" + import tensorrt_llm._torch.models.modeling_auto as modeling_auto + + monkeypatch.setattr(modeling_auto, "get_registered_model_class", lambda arch: _BuiltIn) + monkeypatch.setitem(_RESOLVED_MODES, _RefusingTarget, ModelingV2Mode(mode)) + config = replace(_model_config(_gpt_oss_config()), modeling_v2_target=_RefusingTarget) + config.extra_attrs["nvfp4_gemm_allowed_backends"] = ["cutlass"] + if mode == "require": + with pytest.raises(NoTarget, match="not the one the target is written for"): + AutoModelForCausalLM.from_config(config) + return + model = AutoModelForCausalLM.from_config(config) + assert isinstance(model, _BuiltIn) + assert model.config.modeling_v2_target is None + # The loader's settings reach the built-in model. + assert model.config.extra_attrs["nvfp4_gemm_allowed_backends"] == ["cutlass"] + + # --- speculative decoding is an identity criterion ------------------------------ @@ -564,6 +618,265 @@ def test_routes_on_the_kv_cache_algorithm_the_engine_runs( modeling_v2_resolve(config, _llm_args("require", **args)) +@pytest.mark.parametrize( + "mapping_kwargs, parallel", + [ + (dict(world_size=4, tp_size=4, moe_tp_size=4, moe_ep_size=1), "tp4"), + (dict(world_size=8, tp_size=8, moe_tp_size=8, moe_ep_size=1), "tp8"), + # experts split by expert rather than by tensor: another weight layout, no target + (dict(world_size=8, tp_size=8, moe_tp_size=1, moe_ep_size=8), None), + (dict(world_size=8, tp_size=8, moe_tp_size=8, enable_attention_dp=True), None), + (dict(world_size=2, tp_size=2, moe_tp_size=2), None), + ], +) +def test_minimax_m3_parallel_names_the_target_topology(mapping_kwargs, parallel): + """MiniMax-M3 has one target per tensor-parallel width: tp4 and tp8 load different per-rank shapes.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + assert routing._parallel(Mapping(**mapping_kwargs)) == parallel + + +# The MiniMax-M3 text model's shape fields, in ``routing._SHAPE`` order. +_M3_SHAPE_FIELDS = ( + "num_hidden_layers", + "hidden_size", + "num_local_experts", + "intermediate_size", + "num_attention_heads", + "num_key_value_heads", +) +# The MiniMax-M3 RoPE base and sparse-attention layers (dense 0-2, sparse 3-59), as its config.json reads. +_M3_LAYOUT = dict( + rope_theta=5_000_000, sparse_attention_config=dict(sparse_attention_freq=[0] * 3 + [1] * 57) +) +# Another checkpoint revision with the same architecture class and shape: another RoPE base and sparse layout. +_OTHER_LAYOUT = dict( + rope_theta=1_000_000, sparse_attention_config=dict(sparse_attention_freq=[0] * 4 + [1] * 56) +) + + +def _m3_text_config(**layout): + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + return dict(zip(_M3_SHAPE_FIELDS, routing._SHAPE), **(layout or _M3_LAYOUT)) + + +@pytest.mark.parametrize("spec", [None, "eagle3"], ids=["no_speculation", "one_model_eagle3"]) +@pytest.mark.parametrize("tp", [4, 8]) +def test_minimax_m3_route_takes_the_released_checkpoint(monkeypatch, spec, tp): + from types import SimpleNamespace + + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + monkeypatch.setattr(routing, "_eagle3_one_model_or_none", lambda spec: True) + spec_config = ( + None + if spec is None + else SimpleNamespace(spec_dec_mode=SimpleNamespace(name="EAGLE3_ONE_MODEL")) + ) + target = routing.route( + SimpleNamespace( + pretrained_config=SimpleNamespace(text_config=_m3_text_config()), + mapping=Mapping(world_size=tp, tp_size=tp, moe_tp_size=tp, moe_ep_size=1), + spec_config=spec_config, + ), + _m3_llm_args("auto", "nvfp4"), + routing._SM, + ) + assert target.__name__ == f"ModelingV2MiniMaxM3Nvfp4Sm103Tp{tp}" + + +@pytest.mark.parametrize( + "layout, reason", + [ + (_OTHER_LAYOUT, "rope_theta 1000000, not 5000000; sparse attention layers other than 3-59"), + (dict(_M3_LAYOUT, rope_theta=1_000_000), "rope_theta 1000000, not 5000000"), + ( + dict(_M3_LAYOUT, sparse_attention_config=_OTHER_LAYOUT["sparse_attention_config"]), + "sparse attention layers other than 3-59", + ), + ( + dict(sparse_attention_config=_M3_LAYOUT["sparse_attention_config"]), + "rope_theta None, not 5000000", + ), + ( + dict( + _M3_LAYOUT, + sparse_attention_config=dict( + _M3_LAYOUT["sparse_attention_config"], use_sparse_attention=False + ), + ), + "sparse attention layers other than 3-59", + ), + ], + ids=["other_revision", "rope_base", "sparse_layers", "no_rope_base", "sparse_disabled"], +) +def test_minimax_m3_route_declines_another_rope_base_or_sparse_layout(monkeypatch, layout, reason): + """The architecture name and the shape do not pin the RoPE base or which layers are sparse: a checkpoint with + the target shape but another of either (another checkpoint revision) is not routed to the MiniMax-M3 targets, + even with their quantization.""" + from types import SimpleNamespace + + from tensorrt_llm._torch._experimental.modeling_v2._router_index import NoTarget + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + text_config = _m3_text_config(**layout) + assert routing._layout_mismatch(text_config) == reason + with pytest.raises(NoTarget, match=reason.split(";")[0]): + routing.route( + SimpleNamespace( + pretrained_config=SimpleNamespace(text_config=text_config), + mapping=Mapping(world_size=4, tp_size=4, moe_tp_size=4, moe_ep_size=1), + spec_config=None, + ), + _m3_llm_args("auto", "nvfp4"), + routing._SM, + ) + + +def test_minimax_m3_layout_reads_rope_theta_from_rope_parameters(): + """A config that keeps the RoPE base under ``rope_parameters`` (newer transformers) reads the same.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + sparse = _M3_LAYOUT["sparse_attention_config"] + assert routing._layout_mismatch(_m3_text_config()) is None + for rope in (dict(rope_theta=5_000_000), dict(full_attention=dict(rope_theta=5_000_000))): + assert ( + routing._layout_mismatch(dict(sparse_attention_config=sparse, rope_parameters=rope)) + is None + ) + + +def _m3_sparse(**overrides): + """The sparse-attention runtime of the PR deployment (MSA, FP8 indexer, fused projection), with overrides.""" + from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig + + fields = dict( + implementation="msa", + indexer_kv_dtype="fp8", + fuse_qkv_index_projection=True, + sparse_disable_index_value=True, + ) + fields.update(overrides) + return MiniMaxM3SparseAttentionConfig(**fields) + + +def _m3_resolve_config(sparse, kv_checkpoint, tp=4, layout=None): + """A ModelConfig as the loader holds it when it resolves: the MiniMax-M3 VL wrapper with the target text shape, + the checkpoint's KV quantization (before kv_cache_config.dtype is applied) and the sparse-attention runtime.""" + pretrained = PretrainedConfig( + architectures=["MiniMaxM3SparseForConditionalGeneration"], + text_config=_m3_text_config(**(layout or {})), + ) + return ModelConfig( + pretrained_config=pretrained, + mapping=Mapping(world_size=tp, tp_size=tp, moe_tp_size=tp, moe_ep_size=1), + quant_config=QuantConfig(kv_cache_quant_algo=kv_checkpoint), + sparse_attention_config=sparse, + ) + + +def _m3_llm_args(mode, kv_dtype): + return _llm_args(mode, kv_cache_config=dict(dtype=kv_dtype)) + + +@pytest.fixture +def _minimax_m3_sm(_on_sm10x): + """The MiniMax-M3 targets route on sm_103 only; on sm_100 they route nowhere.""" + if _on_sm10x != (10, 3): + pytest.skip("the MiniMax-M3 targets are certified on sm_103 only") + + +@pytest.mark.usefixtures("_minimax_m3_sm") +@pytest.mark.parametrize("tp", [4, 8]) +@pytest.mark.parametrize("mode", ["auto", "require"]) +def test_minimax_m3_pr_runtime_resolves_to_the_target(monkeypatch, mode, tp): + """The PR deployment (MSA, FP8 indexer, fused projection, kv_cache_config.dtype='nvfp4') takes the target.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + config = _m3_resolve_config(_m3_sparse(), None, tp) + target = modeling_v2_resolve(config, _m3_llm_args(mode, "nvfp4")) + assert target.__name__ == f"ModelingV2MiniMaxM3Nvfp4Sm103Tp{tp}" + + +@pytest.mark.usefixtures("_minimax_m3_sm") +@pytest.mark.parametrize("tp", [4, 8]) +def test_minimax_m3_bf16_kv_cache_falls_back_under_auto(monkeypatch, tp): + """The targets run an FP8 or NVFP4 KV cache: ``auto`` with the BF16 cache resolves nothing, so the loader builds + the registered built-in MiniMax-M3 model.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + from tensorrt_llm._torch.models.modeling_minimaxm3 import MiniMaxM3VLForConditionalGeneration + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + # kv_cache_config.dtype 'auto' keeps the checkpoint's KV quantization (None: the BF16 cache). + config = _m3_resolve_config(_m3_sparse(fuse_qkv_index_projection=False), None, tp) + assert modeling_v2_resolve(config, _m3_llm_args("auto", "auto")) is None + resolved = replace(config, modeling_v2_target=None) + assert AutoModelForCausalLM._resolve_class(resolved) is MiniMaxM3VLForConditionalGeneration + + +@pytest.mark.usefixtures("_minimax_m3_sm") +def test_minimax_m3_bf16_kv_cache_raises_under_require(monkeypatch): + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + config = _m3_resolve_config(_m3_sparse(fuse_qkv_index_projection=False), None) + with pytest.raises(ValueError) as excinfo: + modeling_v2_resolve(config, _m3_llm_args("require", "auto")) + message = str(excinfo.value) + assert "MiniMaxM3SparseForConditionalGeneration" in message + assert "KV cache bf16" in message + + +@pytest.mark.usefixtures("_minimax_m3_sm") +@pytest.mark.parametrize("tp", [4, 8]) +def test_minimax_m3_other_checkpoint_revision_is_declined(monkeypatch, tp): + """Another checkpoint revision (same architecture class and shape, another RoPE base and sparse layout) with the + PR runtime: ``auto`` builds the built-in model and ``require`` names the layout.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + config = _m3_resolve_config(_m3_sparse(), None, tp, layout=_OTHER_LAYOUT) + assert modeling_v2_resolve(config, _m3_llm_args("auto", "nvfp4")) is None + with pytest.raises(ValueError, match="rope_theta 1000000, not 5000000"): + modeling_v2_resolve(config, _m3_llm_args("require", "nvfp4")) + + +@pytest.mark.usefixtures("_minimax_m3_sm") +def test_minimax_m3_kv_cache_dtype_overrides_the_checkpoint(monkeypatch): + """An explicit kv_cache_config.dtype is the KV cache the engine runs with, whatever the checkpoint declares.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + config = _m3_resolve_config(_m3_sparse(fuse_qkv_index_projection=False), None) + assert modeling_v2_resolve(config, _m3_llm_args("auto", "auto")) is None + target = modeling_v2_resolve(config, _m3_llm_args("auto", "fp8")) + assert target.__name__ == "ModelingV2MiniMaxM3Nvfp4Sm103Tp4" + + +@pytest.mark.usefixtures("_minimax_m3_sm") +@pytest.mark.parametrize("tp", [4, 8]) +@pytest.mark.parametrize("role", ["aggregated", "context", "generation"]) +def test_minimax_m3_requires_explicit_opt_in(monkeypatch, tp: int, role: str) -> None: + """Every serving role stays on the built-in model unless modeling_v2 is requested.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + monkeypatch.setenv("TRTLLM_DISAGG_ROLE", role) + config = _m3_resolve_config(_m3_sparse(), None, tp) + kwargs = dict(kv_cache_config=dict(dtype="nvfp4")) + if role != "aggregated": + kwargs["cache_transceiver_config"] = CacheTransceiverConfig() + assert modeling_v2_resolve(config, TorchLlmArgs(model=_DUMMY_MODEL, **kwargs)) is None + for mode in ("auto", "require"): + target = modeling_v2_resolve(config, _llm_args(mode, **kwargs)) + assert target.__name__ == f"ModelingV2MiniMaxM3Nvfp4Sm103Tp{tp}" + assert resolved_mode(target) is ModelingV2Mode(mode) + + _SM100 = (10, 0) _TP16_MOETP4EP4 = dict(world_size=16, tp_size=16, moe_tp_size=4, moe_ep_size=4) _TP16_MOETP16EP1 = dict(world_size=16, tp_size=16, moe_tp_size=16, moe_ep_size=1) diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 5e3ce1d1292b..975239ea79ed 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -988,6 +988,7 @@ def test_piecewise_attention_boundary_preserves_indexer_cache_contract( msa_write_idx_k=Mock(), msa_prefill_proxy_plan=None, msa_prefill_n_valid_blocks=None, + _msa_prefill_identity=None, msa_prefill_proxy_kv_indices=torch.tensor([0, 1], dtype=torch.int32), msa_qo_lens_cpu=torch.tensor([2], dtype=torch.int32), msa_kv_lens_cpu=torch.tensor([2], dtype=torch.int32), diff --git a/tests/unittest/_torch/sampler/test_greedy_tokens.py b/tests/unittest/_torch/sampler/test_greedy_tokens.py new file mode 100644 index 000000000000..e5100456aadc --- /dev/null +++ b/tests/unittest/_torch/sampler/test_greedy_tokens.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TorchSampler consuming greedy tokens a forward sampled in place of its LM head.""" + +from typing import Any, Callable + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.handle_logits import HandleLogits +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.sampler import TorchSampler +from tensorrt_llm._torch.pyexecutor.sampler.sampler import GREEDY_TOKENS_KEY, greedy_tokens_usable +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.bindings import SamplingConfig +from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams + +VOCAB_SIZE = 16 + + +def _request( + request_id: int = 0, + seq_slot: int = 0, + sampling_params: SamplingParams | None = None, + **kwargs: Any, +) -> LlmRequest: + params = sampling_params if sampling_params is not None else SamplingParams() + return LlmRequest( + request_id=request_id, + max_new_tokens=8, + input_tokens=[1, 2, 3], + sampling_config=SamplingConfig(params._get_sampling_config()), + seq_slot=seq_slot, + is_streaming=False, + **kwargs, + ) + + +def _set(name: str, value: Any) -> Callable[[LlmRequest], None]: + def apply(request: LlmRequest) -> None: + setattr(request, name, value) + + return apply + + +class TestGreedyTokensUsable: + """Whether a generation request is served by the forward's greedy token alone.""" + + pytestmark = pytest.mark.cpu_only + + def test_plain_greedy_request_is_usable(self): + assert greedy_tokens_usable(_request()) + + @pytest.mark.parametrize( + ("kwargs", "mutate"), + [ + pytest.param({}, _set("py_draft_tokens", [5, 6]), id="draft_tokens"), + pytest.param( + {}, _set("py_embedding_bias", torch.zeros(VOCAB_SIZE)), id="embedding_bias" + ), + pytest.param({}, _set("py_bad_words", [[4]]), id="bad_words"), + pytest.param({}, _set("py_no_repeat_ngram_size", 2), id="no_repeat_ngram"), + pytest.param( + {}, _set("py_logits_post_processors", [print]), id="logits_post_processor" + ), + pytest.param({"return_log_probs": True}, None, id="log_probs"), + pytest.param({"return_generation_logits": True}, None, id="generation_logits"), + pytest.param({"stop_words_list": [[13, 14]]}, None, id="multi_token_stop_word"), + ], + ) + def test_a_logits_reader_makes_the_request_unusable(self, kwargs, mutate): + request = _request(**kwargs) + if mutate is not None: + mutate(request) + + assert not greedy_tokens_usable(request) + + @pytest.mark.parametrize( + "sampling_params", + [ + pytest.param(SamplingParams(temperature=0.8, top_k=50), id="top_k"), + pytest.param(SamplingParams(temperature=0.7, top_p=0.9), id="top_p"), + pytest.param(SamplingParams(repetition_penalty=1.3), id="repetition_penalty"), + pytest.param(SamplingParams(presence_penalty=0.5), id="presence_penalty"), + pytest.param(SamplingParams(frequency_penalty=0.5), id="frequency_penalty"), + pytest.param(SamplingParams(min_tokens=2), id="min_length"), + ], + ) + def test_non_greedy_or_penalized_sampling_is_unusable(self, sampling_params): + assert not greedy_tokens_usable(_request(sampling_params=sampling_params)) + + def test_single_token_stop_words_stay_usable(self): + """Single-token stop words are checked on the host, like the end id.""" + assert greedy_tokens_usable(_request(stop_words_list=[[13], [17]])) + + def test_guided_decoding_makes_the_request_unusable(self): + """Guided decoding masks the logits after the forward, so a greedy token cannot serve it.""" + request = _request(guided_decoding_params=GuidedDecodingParams(regex="[0-9]+")) + assert request.py_guided_decoding_params is not None + + assert not greedy_tokens_usable(request) + + @pytest.mark.parametrize("dummy_flag", ["is_attention_dp_dummy", "is_cuda_graph_dummy"]) + def test_a_dummy_request_is_never_usable(self, dummy_flag): + """The stable greedy path takes no batch with a dummy, so such a batch needs logits.""" + request = _request() + setattr(request, dummy_flag, True) + + assert not greedy_tokens_usable(request) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_handle_logits_takes_greedy_only_outputs() -> None: + """The executor hands HandleLogits the outputs' logits before sampling; a greedy-only batch (no request asks for + logits) has none, and HandleLogits must return without reading them.""" + HandleLogits()([], [_request()], None, 1, [0], True) + + +class TestStableGreedyFromGreedyTokens: + """The stable greedy path takes the forward's tokens instead of an argmax over logits.""" + + @staticmethod + def _sampler(max_num_sequences: int = 3) -> TorchSampler: + return TorchSampler( + TorchSampler.Args( + max_seq_len=16, + max_draft_len=0, + max_num_sequences=max_num_sequences, + max_beam_width=1, + max_total_draft_tokens=0, + disable_overlap_scheduler=True, + ) + ) + + @staticmethod + def _batch(*requests: LlmRequest) -> ScheduledRequests: + scheduled = ScheduledRequests() + scheduled.generation_requests = list(requests) + return scheduled + + @staticmethod + def _one_hot_logits(tokens: list[int]) -> torch.Tensor: + logits = torch.zeros((len(tokens), VOCAB_SIZE), device="cuda") + logits[torch.arange(len(tokens)), torch.tensor(tokens)] = 10.0 + return logits + + def test_greedy_tokens_match_the_logits_argmax(self): + tokens = [7, 3] + slots = [2, 0] + results = {} + for name, outputs in ( + ("logits", {"logits": self._one_hot_logits(tokens)}), + ( + "greedy_tokens", + { + # No logits key, as the engine returns greedy-only outputs. + # A padded graph replays more rows than the batch has requests. + GREEDY_TOKENS_KEY: torch.tensor(tokens + [9], dtype=torch.int64, device="cuda"), + }, + ), + ): + sampler = self._sampler() + sampler.store.new_tokens.zero_() + batch = self._batch(*(_request(i, slot) for i, slot in enumerate(slots))) + + ( + requests, + seq_slots_host, + _, + _, + _, + new_tokens_host, + single_step_greedy, + ) = sampler._process_requests(batch, outputs, sampler.store.new_tokens, [0]) + torch.cuda.synchronize() + + assert single_step_greedy + assert requests == batch.generation_requests + assert seq_slots_host.tolist() == slots + results[name] = ( + new_tokens_host.reshape(-1).tolist(), + sampler.store.new_tokens[0, :, 0].tolist(), + ) + + assert results["greedy_tokens"] == results["logits"] + assert results["greedy_tokens"][0] == tokens + assert [results["greedy_tokens"][1][slot] for slot in slots] == tokens + + def test_missing_logits_for_a_batch_that_needs_them_raises(self): + sampler = self._sampler() + request = _request() + request.py_no_repeat_ngram_size = 2 + outputs = {GREEDY_TOKENS_KEY: torch.tensor([4], dtype=torch.int64, device="cuda")} + + with pytest.raises(RuntimeError, match="greedy tokens without logits"): + sampler._process_requests(self._batch(request), outputs, sampler.store.new_tokens, [0]) + + def test_a_batch_padded_with_an_attention_dp_dummy_asks_for_logits(self): + """A dummy makes the batch unusable, so the engine computes logits rather than keep greedy tokens.""" + live = _request(0, 0) + dummy = _request(1, 1) + dummy.is_attention_dp_dummy = True + batch = self._batch(live, dummy) + + assert not all(greedy_tokens_usable(request) for request in batch.generation_requests) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py index dd59b0292daa..1f7f739d23ea 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py @@ -12,12 +12,14 @@ CUDA_GRAPH_DUMMY_REQUEST_ID, ADPShapeAgreement, CUDAGraphRunner, + CUDAGraphRunnerConfig, DSparkCUDAGraphRunner, KeyType, cuda_graph_dummy_request_id, ) from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner from tensorrt_llm._torch.pyexecutor.sampler.sampler_common import SampleType +from tensorrt_llm.mapping import Mapping class _Batch: @@ -202,7 +204,7 @@ def test_decoder_selects_confidence_graph_subclass_only_when_enabled( monkeypatch, confidence_enabled ): monkeypatch.setattr(graph_module, "get_memory_buffers", lambda: None) - config = SimpleNamespace( + config = CUDAGraphRunnerConfig( use_cuda_graph=False, cuda_graph_padding_enabled=True, cuda_graph_batch_sizes=[1, 4], @@ -213,7 +215,13 @@ def test_decoder_selects_confidence_graph_subclass_only_when_enabled( is_encoder_decoder=False, enable_encoder_decoder_mixed_cuda_graph=False, enable_in_graph_sampling=False, - cuda_graph_mem_pool=None, + max_num_tokens=16, + use_mrope=False, + original_max_total_draft_tokens=0, + enable_attention_dp=False, + batch_size=4, + mapping=Mapping(), + dist=None, dynamic_draft_len_mapping=None, ) registered = [] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py index 566bdac3e270..5f16696f22a6 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py @@ -730,7 +730,7 @@ def gather(*args) -> bool: def test_graph_overlap_offset_zeroing_preserves_ragged_bucket_padding(trims: bool) -> None: decoder = object.__new__(DecoderInputPreparer) decoder._config = SimpleNamespace( - max_num_tokens=16, max_beam_width=1, disable_overlap_scheduler=False + max_num_tokens=16, max_beam_width=1, disable_overlap_scheduler=False, is_spec_decode=True ) decoder._dspark_trims_submitted_tokens = trims decoder._staged_previous_batch_indices = None diff --git a/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py b/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py new file mode 100644 index 000000000000..1ea8cdaa58f4 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py @@ -0,0 +1,254 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SpecSampler's per-slot store with and without a model's SpecStepFold.""" + +from typing import Optional + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.runtime_hooks import SpecStepFold +from tensorrt_llm._torch.pyexecutor.sampler import TorchSampler +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler +from tensorrt_llm.bindings import SamplingConfig + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + +MAX_DRAFT_LEN = 3 +MAX_NUM_SEQUENCES = 6 +SEQ_SLOTS = [4, 1] + + +class _RecordingStepFold: + """A SpecStepFold whose store decisions are fixed per test; it records what the sampler asks.""" + + def __init__(self, rows_stored: bool = False, scatter_taken: bool = False) -> None: + self._rows_stored = rows_stored + self._scatter_taken = scatter_taken + self.store: Optional[object] = None + self.rows_stored_calls: list[tuple[int, list[int]]] = [] + self.scatter_calls: list[dict[str, torch.Tensor]] = [] + + def bind_engine_buffers( + self, + input_ids: torch.Tensor, + draft_tokens: torch.Tensor, + position_id_offsets: torch.Tensor, + kv_lens_offsets: torch.Tensor, + max_batch_size: int, + ) -> None: + raise AssertionError("the sampler never binds engine buffers") + + def bind_sampler_store(self, store: object) -> None: + self.store = store + + def begin_step(self, seq_slots: Optional[list[int]], num_rows: int) -> bool: + raise AssertionError("the sampler never begins a step") + + def gather(self, *args: object) -> bool: + raise AssertionError("the sampler never gathers") + + def sampler_rows_stored(self, outputs: dict, num_skip: int, seq_slots: list[int]) -> bool: + self.rows_stored_calls.append((num_skip, list(seq_slots))) + return self._rows_stored + + def scatter( + self, + seq_slots: torch.Tensor, + new_tokens: torch.Tensor, + next_new_tokens: torch.Tensor, + new_tokens_lens: torch.Tensor, + next_draft_tokens: torch.Tensor, + ) -> bool: + self.scatter_calls.append( + dict( + seq_slots=seq_slots.clone(), + new_tokens=new_tokens.clone(), + next_new_tokens=next_new_tokens.clone(), + new_tokens_lens=new_tokens_lens.clone(), + next_draft_tokens=next_draft_tokens.clone(), + ) + ) + return self._scatter_taken + + +def _fold(**kwargs: bool) -> SpecStepFold: + return _RecordingStepFold(**kwargs) + + +def _sampler() -> SpecSampler: + sampler = SpecSampler( + TorchSampler.Args( + max_seq_len=64, + max_draft_len=MAX_DRAFT_LEN, + max_num_sequences=MAX_NUM_SEQUENCES, + max_beam_width=1, + max_total_draft_tokens=MAX_DRAFT_LEN, + disable_overlap_scheduler=True, + ) + ) + # The store is allocated uninitialized; give it a recognizable baseline. + for index, tensor in enumerate( + ( + sampler.store.new_tokens, + sampler.store.next_new_tokens, + sampler.store.next_draft_tokens, + sampler.store.new_tokens_lens, + ) + ): + tensor.fill_(-100 - index) + return sampler + + +def _store_snapshot(sampler: SpecSampler) -> dict[str, torch.Tensor]: + store = sampler.store + return dict( + new_tokens=store.new_tokens.clone(), + next_new_tokens=store.next_new_tokens.clone(), + next_draft_tokens=store.next_draft_tokens.clone(), + new_tokens_lens=store.new_tokens_lens.clone(), + ) + + +def _assert_store_equals(sampler: SpecSampler, expected: dict[str, torch.Tensor]) -> None: + torch.cuda.synchronize() + for name, tensor in _store_snapshot(sampler).items(): + assert torch.equal(tensor, expected[name]), name + + +def _step(draft_width: int): + requests = [] + for request_id, seq_slot in enumerate(SEQ_SLOTS): + request = LlmRequest( + request_id=request_id, + max_new_tokens=8, + input_tokens=[1, 2, 3], + sampling_config=SamplingConfig(), + seq_slot=seq_slot, + is_streaming=False, + ) + request.py_draft_tokens = [7] * MAX_DRAFT_LEN + requests.append(request) + scheduled = ScheduledRequests() + scheduled.generation_requests = requests + rows = len(requests) + 1 # plus a CUDA-graph padding row the sampler must not store + outputs = dict( + new_tokens=torch.arange(rows * (draft_width + 1), dtype=torch.int32, device="cuda").view( + rows, draft_width + 1 + ), + new_tokens_lens=torch.arange(1, rows + 1, dtype=torch.int32, device="cuda"), + next_draft_tokens=( + 100 + torch.arange(rows * draft_width, dtype=torch.int32, device="cuda") + ).view(rows, draft_width), + next_new_tokens=( + 200 + torch.arange(rows * (draft_width + 1), dtype=torch.int32, device="cuda") + ).view(rows, draft_width + 1), + ) + return scheduled, outputs + + +def _reference_store(before: dict[str, torch.Tensor], outputs: dict) -> dict[str, torch.Tensor]: + """What the sampler has always stored: each live row's outputs at its slot, zero-padded.""" + expected = {name: tensor.clone() for name, tensor in before.items()} + count = len(SEQ_SLOTS) + slots = torch.tensor(SEQ_SLOTS, dtype=torch.long, device="cuda") + + def padded(tensor: torch.Tensor, width: int) -> torch.Tensor: + tensor = tensor[:count, :width] + return torch.nn.functional.pad(tensor, (0, width - tensor.shape[1])) + + accepted_width = expected["new_tokens"].shape[0] + wire_width = expected["next_new_tokens"].shape[0] + draft_width = expected["next_draft_tokens"].shape[1] + expected["new_tokens"].squeeze(-1).T.index_copy_( + 0, slots, padded(outputs["new_tokens"], accepted_width) + ) + expected["next_new_tokens"].squeeze(-1).T.index_copy_( + 0, slots, padded(outputs["next_new_tokens"], wire_width) + ) + expected["new_tokens_lens"].index_copy_(0, slots, outputs["new_tokens_lens"][:count]) + expected["next_draft_tokens"].index_copy_( + 0, slots, padded(outputs["next_draft_tokens"], draft_width) + ) + return expected + + +@pytest.mark.parametrize("draft_width", [MAX_DRAFT_LEN, MAX_DRAFT_LEN - 1]) +def test_without_a_fold_the_store_is_written_as_before(draft_width): + sampler = _sampler() + scheduled, outputs = _step(draft_width) + before = _store_snapshot(sampler) + + state = sampler.sample_async(scheduled, outputs, []) + + assert sampler.spec_step_fold is None + expected = _reference_store(before, outputs) + _assert_store_equals(sampler, expected) + state.sampler_event.synchronize() + assert torch.equal(state.host.new_tokens, expected["new_tokens"].cpu()) + assert torch.equal(state.host.next_draft_tokens, expected["next_draft_tokens"].cpu()) + + +def test_binding_a_fold_hands_it_the_sampler_store(): + sampler = _sampler() + fold = _fold() + + sampler.bind_spec_step_fold(fold) + + assert sampler.spec_step_fold is fold + assert fold.store is sampler.store + + +def test_rows_the_forward_already_stored_are_left_alone(): + sampler = _sampler() + fold = _fold(rows_stored=True) + sampler.bind_spec_step_fold(fold) + scheduled, outputs = _step(MAX_DRAFT_LEN) + before = _store_snapshot(sampler) + + state = sampler.sample_async(scheduled, outputs, []) + + assert fold.rows_stored_calls == [(0, SEQ_SLOTS)] + assert fold.scatter_calls == [] + _assert_store_equals(sampler, before) + # The host copies still read the (fold-owned) store. + state.sampler_event.synchronize() + assert torch.equal(state.host.new_tokens, before["new_tokens"].cpu()) + + +def test_a_fold_scatter_replaces_the_index_copies(): + sampler = _sampler() + fold = _fold(scatter_taken=True) + sampler.bind_spec_step_fold(fold) + scheduled, outputs = _step(MAX_DRAFT_LEN) + before = _store_snapshot(sampler) + + sampler.sample_async(scheduled, outputs, []) + + _assert_store_equals(sampler, before) + assert len(fold.scatter_calls) == 1 + call = fold.scatter_calls[0] + count = len(SEQ_SLOTS) + assert call["seq_slots"].tolist() == SEQ_SLOTS + assert call["seq_slots"].dtype == torch.long and call["seq_slots"].is_cuda + assert torch.equal(call["new_tokens"], outputs["new_tokens"][:count]) + assert torch.equal(call["next_new_tokens"], outputs["next_new_tokens"][:count]) + assert torch.equal(call["new_tokens_lens"], outputs["new_tokens_lens"][:count]) + assert torch.equal(call["next_draft_tokens"], outputs["next_draft_tokens"][:count]) + + +@pytest.mark.parametrize("draft_width", [MAX_DRAFT_LEN, MAX_DRAFT_LEN - 1]) +def test_a_fold_that_declines_leaves_the_sampler_store_path(draft_width): + sampler = _sampler() + fold = _fold(rows_stored=False, scatter_taken=False) + sampler.bind_spec_step_fold(fold) + scheduled, outputs = _step(draft_width) + before = _store_snapshot(sampler) + + sampler.sample_async(scheduled, outputs, []) + + assert fold.rows_stored_calls == [(0, SEQ_SLOTS)] + assert len(fold.scatter_calls) == 1 + _assert_store_equals(sampler, _reference_store(before, outputs)) diff --git a/tests/unittest/_torch/test_connector_adp.py b/tests/unittest/_torch/test_connector_adp.py index ad1e6829630e..91ec0607af18 100644 --- a/tests/unittest/_torch/test_connector_adp.py +++ b/tests/unittest/_torch/test_connector_adp.py @@ -17,6 +17,7 @@ import pytest from pydantic import ValidationError +from tensorrt_llm._torch.distributed.communicator import Distributed from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import ( AsyncRequests, KvCacheConnectorManager, @@ -80,7 +81,21 @@ def _manager(rank: int = 1) -> KvCacheConnectorManager: def _configure_recovery_controls(executor: "PyExecutor") -> None: from tensorrt_llm._torch.pyexecutor.executor_request_queue import ExecutorRequestQueue - executor.dist = MagicMock(tp_size=1, cp_size=1, pp_size=1) + executor.dist = MagicMock( + spec_set=Distributed, + tp_size=1, + cp_size=1, + pp_size=1, + tp_rank=0, + rank=0, + cp_rank=0, + has_cp_helix=False, + has_pp=False, + cp_config={}, + ) + executor.dist.tp_cp_allgather_int64.side_effect = partial( + Distributed.tp_cp_allgather_int64, executor.dist + ) executor.enable_attention_dp = True executor.enable_early_first_token_response = False executor.kv_connector_manager = None @@ -414,6 +429,7 @@ def test_batch_start_combines_readiness_with_status(batch_sizes: tuple[int, int] _configure_recovery_controls(executor) executor.kv_connector_manager = _manager(rank) executor.dist.tp_size = 2 + executor.dist.tp_rank = executor.dist.rank = rank executor.dist.tp_allgather_int64.side_effect = partial(collective.gather_int64, rank) executors.append(executor)