diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index 0e7fc9390dc0..54a477030e66 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -78,6 +78,10 @@ list(FILTER SRC_CU EXCLUDE REGEX "fusedSampling/.*") list(FILTER SRC_CPP EXCLUDE REGEX "marlin/.*") list(FILTER SRC_CU EXCLUDE REGEX "marlin/.*") +# Scope tcgen05 attention sources to their supported architectures. +list(FILTER SRC_CU EXCLUDE REGEX + "minimaxM3/(fusedDecodeAttn|attnMega|attnC1)\\.cu$") + if(NOT ENABLE_MULTI_DEVICE) list(FILTER SRC_CU EXCLUDE REGEX "customAllReduceKernels*.*cu$") endif() @@ -111,10 +115,27 @@ if(MARLIN_SRC) set_cuda_architectures(marlin_src 89 90 120f) endif() +# Preserve the FlashInfer-matched compilation for sampling only. Do not enable +# fast math globally: attention exactness depends on its flags. +set_source_files_properties( + ${CMAKE_CURRENT_SOURCE_DIR}/minimaxM31/toppRenormDetSm103.cu + ${CMAKE_CURRENT_SOURCE_DIR}/minimaxM31/toppRenormDetV3Sm103.cu + ${CMAKE_CURRENT_SOURCE_DIR}/minimaxM31/softmaxMcSm103.cu + PROPERTIES COMPILE_OPTIONS "$<$:--use_fast_math>") + +# Preserve the historical architecture-scoped attention object library. +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/fusedSampling/fusedSamplingKernels.h b/cpp/tensorrt_llm/kernels/fusedSampling/fusedSamplingKernels.h index 8863035e5740..5a267eb9ca4b 100644 --- a/cpp/tensorrt_llm/kernels/fusedSampling/fusedSamplingKernels.h +++ b/cpp/tensorrt_llm/kernels/fusedSampling/fusedSamplingKernels.h @@ -21,6 +21,10 @@ namespace tensorrt_llm { +//! Philox subsequences from here up hold per-row streams. A shared seed/offset makes the row +//! index the subsequence, which never reaches them. +constexpr uint64_t kPerRowSubsequenceBase = 1ull << 63; + namespace kernels { diff --git a/cpp/tensorrt_llm/kernels/fusedSampling/fusedSamplingKernelsCommon.cuh b/cpp/tensorrt_llm/kernels/fusedSampling/fusedSamplingKernelsCommon.cuh index f8b6c31de06b..62b669ed4a84 100644 --- a/cpp/tensorrt_llm/kernels/fusedSampling/fusedSamplingKernelsCommon.cuh +++ b/cpp/tensorrt_llm/kernels/fusedSampling/fusedSamplingKernelsCommon.cuh @@ -129,10 +129,6 @@ struct RowParams bool needMinP; }; -//! Philox subsequences from here up hold per-row streams. A shared seed/offset makes the row -//! index the subsequence, which never reaches them. -constexpr uint64_t kPerRowSubsequenceBase = 1ull << 63; - //! Seeds one row's Philox stream (see FusedSamplingParams::perRowRng). __device__ inline void initRowRng(FusedSamplingParams const& p, int row, curandStatePhilox4_32_10_t* state) { 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/minimaxM31/chainSpecSamplingMcSm103.cu b/cpp/tensorrt_llm/kernels/minimaxM31/chainSpecSamplingMcSm103.cu new file mode 100644 index 000000000000..fa62f8357e4d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/chainSpecSamplingMcSm103.cu @@ -0,0 +1,418 @@ +/* + * 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. + */ + +// ST-01-REJ: the stock chainSpeculativeSamplingKernel<1024, 4> on ctasPerRow CTAs +// per request, bitwise. v2 (bitwise v1 and stock, 24.8 vs 32.3 us at one request) +// moves loads only, every add and compare in its v1 order: the acceptance loop's probabilities are loaded in parallel +// first, the chain loop is unrolled, and the leader reads the per-tile values from shared memory. See +// chainSpecSamplingMcSm103.h for the contract. Compiled with the repo's default flags, as +// kernels/speculativeDecoding/chainSpeculativeSamplingKernels.cu. + +#include "tensorrt_llm/kernels/fusedSampling/fusedSamplingKernels.h" +#include "tensorrt_llm/kernels/minimaxM31/chainSpecSamplingMcSm103.h" + +#include +#include +#include + +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM31 +{ + +namespace +{ + +using speculative_decoding::ChainSpeculativeSamplingParams; + +constexpr int32_t kBlockSize = 1024; +constexpr int32_t kVecSize = 4; +constexpr int32_t kTileSize = kBlockSize * kVecSize; +constexpr int32_t kMaxTiles = kChainSpecMcMaxVocab / kTileSize; + +//! One request's cross-CTA state. +struct RowWorkspace +{ + //! The stock kernel's per-thread (target, residual) mass chains. + float2 chain[kBlockSize]; + //! Per tile: the stock per-tile residual (or target) sum, and the last index with mass (-1 if none). + float tileSum[kMaxTiles]; + int32_t tileLast[kMaxTiles]; + float total; + int32_t targetOnly; + //! The row barrier: CTAs arrived in this generation, and the generation (only increases). + unsigned arrive; + unsigned gen; + int32_t pad[3]; +}; + +// --------------------------------------------------------------------------------------------------------------------- +// Identical to kernels/speculativeDecoding/chainSpeculativeSamplingKernels.cu (Float2SumOp, loadProbs, loadResidual). +// --------------------------------------------------------------------------------------------------------------------- +struct Float2SumOp +{ + __device__ __forceinline__ float2 operator()(float2 a, float2 b) const + { + return make_float2(a.x + b.x, a.y + b.y); + } +}; + +template +__device__ __forceinline__ void loadProbs(float const* row, uint32_t base, uint32_t vocabSize, float (&out)[VEC_SIZE]) +{ + if constexpr (VEC_SIZE == 4) + { + if (base + VEC_SIZE <= vocabSize) + { + float4 const v = *reinterpret_cast(row + base); + out[0] = v.x; + out[1] = v.y; + out[2] = v.z; + out[3] = v.w; + return; + } + } +#pragma unroll + for (int32_t j = 0; j < VEC_SIZE; ++j) + { + out[j] = base + j < vocabSize ? row[base + j] : 0.0f; + } +} + +//! The distribution a request resamples from: max(0, target - draft) at the first rejected +//! position, or the target itself when every draft token was accepted (draft == nullptr) or +//! when the residual has no mass. +template +__device__ __forceinline__ void loadResidual( + float const* target, float const* draft, bool targetOnly, uint32_t base, uint32_t vocabSize, float (&out)[VEC_SIZE]) +{ + loadProbs(target, base, vocabSize, out); + if (targetOnly) + { + return; + } + float d[VEC_SIZE]; + loadProbs(draft, base, vocabSize, d); +#pragma unroll + for (int32_t j = 0; j < VEC_SIZE; ++j) + { + out[j] = fmaxf(out[j] - d[j], 0.0f); + } +} + +// --------------------------------------------------------------------------------------------------------------------- + +__device__ __forceinline__ unsigned loadAcquire(unsigned* p) +{ + return cuda::atomic_ref(*p).load(cuda::memory_order_acquire); +} + +//! A barrier of the request's G CTAs; the last CTA to arrive runs leader() (all its threads) before releasing the +//! others. +template +__device__ __forceinline__ void rowSync(RowWorkspace* w, int32_t G, F leader) +{ + __shared__ unsigned sGen; + __shared__ int32_t sLast; + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + sGen = loadAcquire(&w->gen); + sLast = atomicAdd(&w->arrive, 1u) == static_cast(G - 1); + } + __syncthreads(); + if (sLast) + { + __threadfence(); + leader(); + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + atomicExch(&w->arrive, 0u); + cuda::atomic_ref(w->gen).store(sGen + 1, cuda::memory_order_release); + } + } + else if (threadIdx.x == 0) + { + while (loadAcquire(&w->gen) == sGen) + { + __nanosleep(20); + } + } + __syncthreads(); +} + +__global__ void __launch_bounds__(kBlockSize, 1) + chainSpecSamplingMcKernel(ChainSpeculativeSamplingParams params, RowWorkspace* ws) +{ + using BlockReduce = cub::BlockReduce; + using BlockReducePair = cub::BlockReduce; + using BlockScan = cub::BlockScan; + + __shared__ union + { + typename BlockReduce::TempStorage reduce; + typename BlockReducePair::TempStorage reducePair; + typename BlockScan::TempStorage scan; + } tempStorage; + + __shared__ int32_t sPos; + __shared__ float sU; + __shared__ int32_t sWinner; + __shared__ float sPrefix; + __shared__ int32_t sTile; + __shared__ int32_t sTileLast; + + int64_t const row = blockIdx.y; + int32_t const G = static_cast(gridDim.x); + int32_t const cta = static_cast(blockIdx.x); + int32_t const tid = static_cast(threadIdx.x); + int32_t const numDraft = params.numDraftTokens; + uint32_t const vocabSize = static_cast(params.vocabSize); + RowWorkspace* w = ws + row; + int32_t* outTokens = params.outputTokenIds + row * (numDraft + 1); + + // 1. Acceptance: the stock tid-0 loop, identical in every CTA of the request. The probabilities it may read are + // loaded by numDraft threads in parallel first (the loop's values and order unchanged). + __shared__ float sPTarget[kChainSpecMcMaxDraft]; + __shared__ float sPDraft[kChainSpecMcMaxDraft]; + if (tid < numDraft) + { + int32_t const token = params.draftTokenIds[row * numDraft + tid]; + sPTarget[tid] = params.targetProbs[(row * (numDraft + 1) + tid) * vocabSize + token]; + sPDraft[tid] = params.draftProbs[(row * numDraft + tid) * vocabSize + token]; + } + __syncthreads(); + if (tid == 0) + { + curandStatePhilox4_32_10_t state; + curand_init(static_cast(params.seed[row]), + kPerRowSubsequenceBase | static_cast(params.offset[row]), 0, &state); + int32_t pos = numDraft; + for (int32_t i = 0; i < numDraft; ++i) + { + float const pTarget = sPTarget[i]; + float const pDraft = sPDraft[i]; + if (!(curand_uniform(&state) * pDraft < pTarget)) + { + pos = i; + break; + } + } + sPos = pos; + // curand_uniform is in (0, 1]; the inverse-CDF search below wants [0, 1). + sU = 1.0f - curand_uniform(&state); + } + __syncthreads(); + + int32_t const pos = sPos; + float const* target = params.targetProbs + (row * (numDraft + 1) + pos) * vocabSize; + float const* draft = pos < numDraft ? params.draftProbs + (row * numDraft + pos) * vocabSize : nullptr; + + // 2. The stock per-thread (target, residual) chains, spread over the request's CTAs; one CTA reduces them. + int32_t const chainsPerCta = kBlockSize / G; + if (tid < chainsPerCta) + { + int32_t const t = cta * chainsPerCta + tid; + // x: target mass, y: residual mass. + float2 sums = make_float2(0.0f, 0.0f); +#pragma unroll 8 + for (uint32_t base = t * kVecSize; base < vocabSize; base += kTileSize) + { + float tv[kVecSize]; + loadProbs(target, base, vocabSize, tv); + float dv[kVecSize] = {}; + if (draft != nullptr) + { + loadProbs(draft, base, vocabSize, dv); + } +#pragma unroll + for (int32_t j = 0; j < kVecSize; ++j) + { + sums.x += tv[j]; + sums.y += fmaxf(tv[j] - dv[j], 0.0f); + } + } + w->chain[t] = sums; + } + rowSync(w, G, + [&] + { + float2 const blockSums + = BlockReducePair(tempStorage.reducePair).Reduce(__ldcg(&w->chain[tid]), Float2SumOp{}); + if (tid == 0) + { + int32_t const targetOnly = draft == nullptr || !(blockSums.y > 0.0f); + w->targetOnly = targetOnly; + w->total = targetOnly ? blockSums.x : blockSums.y; + } + }); + bool const targetOnly = __ldcg(&w->targetOnly) != 0; + + // 3. Tile sums (the stock per-tile cub reduce) and each tile's last index with mass (the stock fallback takes the + // max over the tiles it visited, i.e. up to and including the tile where the threshold falls). + int32_t const numTiles = static_cast((vocabSize + kTileSize - 1) / kTileSize); + for (int32_t tile = cta; tile < numTiles; tile += G) + { + if (tid == 0) + { + sTileLast = -1; + } + __syncthreads(); + uint32_t const base = tile * kTileSize + tid * kVecSize; + float value[kVecSize]; + loadResidual(target, draft, targetOnly, base, vocabSize, value); + int32_t localLast = -1; +#pragma unroll + for (int32_t j = 0; j < kVecSize; ++j) + { + if (value[j] > 0.0f) + { + localLast = static_cast(base + j); + } + } + if (localLast >= 0) + { + atomicMax(&sTileLast, localLast); + } + float const tileSum = BlockReduce(tempStorage.reduce).Sum(value); + __syncthreads(); + if (tid == 0) + { + w->tileSum[tile] = tileSum; + w->tileLast[tile] = sTileLast; + } + __syncthreads(); + } + + // 4. One CTA: the tile prefix in tile order, then the stock in-tile scan; it writes the outputs. + rowSync(w, G, + [&] + { + float const threshold = sU * __ldcg(&w->total); + // The per-tile values in shared memory, read by numTiles threads in parallel. + __shared__ float sTileSum[kMaxTiles]; + __shared__ int32_t sTileLastOf[kMaxTiles]; + if (tid < numTiles) + { + sTileSum[tid] = __ldcg(&w->tileSum[tid]); + sTileLastOf[tid] = __ldcg(&w->tileLast[tid]); + } + __syncthreads(); + if (tid == 0) + { + float prefix = 0.0f; + int32_t found = -1; + for (int32_t tile = 0; tile < numTiles; ++tile) + { + float const tileSum = sTileSum[tile]; + if (prefix + tileSum > threshold) + { + found = tile; + break; + } + prefix += tileSum; + } + sTile = found; + sPrefix = prefix; + sWinner = static_cast(vocabSize); + } + __syncthreads(); + if (sTile >= 0) + { + uint32_t const base = sTile * kTileSize + tid * kVecSize; + float value[kVecSize]; + loadResidual(target, draft, targetOnly, base, vocabSize, value); + float inclusive[kVecSize]; + BlockScan(tempStorage.scan).InclusiveSum(value, inclusive); + float const prefix = sPrefix; +#pragma unroll + for (int32_t j = 0; j < kVecSize; ++j) + { + if (value[j] > 0.0f && prefix + inclusive[j] > threshold) + { + atomicMin(&sWinner, static_cast(base + j)); + break; + } + } + } + __syncthreads(); + if (tid == 0) + { + // Float rounding can leave the summed prefix just short of the threshold; the last token with mass + // (over the tiles the stock loop visits) is then the one the threshold falls on. + int32_t const upto = sTile >= 0 ? sTile : numTiles - 1; + int32_t lastValid = -1; + for (int32_t tile = 0; tile <= upto; ++tile) + { + lastValid = max(lastValid, sTileLastOf[tile]); + } + int32_t const token = sWinner < static_cast(vocabSize) ? sWinner : max(lastValid, 0); + for (int32_t i = 0; i < pos; ++i) + { + outTokens[i] = params.draftTokenIds[row * numDraft + i]; + } + outTokens[pos] = token; + for (int32_t i = pos + 1; i <= numDraft; ++i) + { + outTokens[i] = -1; + } + params.outputNumTokens[row] = pos + 1; + } + }); +} + +} // namespace + +cudaError_t invokeChainSpecSamplingMc( + ChainSpeculativeSamplingParams const& params, int ctasPerRow, void* workspace, cudaStream_t stream) +{ + if (params.batchSize <= 0) + { + return cudaSuccess; + } + if (ctasPerRow < 1 || ctasPerRow > kBlockSize || kBlockSize % ctasPerRow != 0 || params.vocabSize <= 0 + || params.vocabSize % kVecSize != 0 || params.vocabSize > kChainSpecMcMaxVocab || params.numDraftTokens < 0 + || params.numDraftTokens > kChainSpecMcMaxDraft || params.batchSize > 65535) + { + return cudaErrorInvalidValue; + } + cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(static_cast(ctasPerRow), static_cast(params.batchSize)); + cfg.blockDim = dim3(kBlockSize); + cfg.dynamicSmemBytes = 0; + cfg.stream = stream; + return cudaLaunchKernelEx(&cfg, chainSpecSamplingMcKernel, params, static_cast(workspace)); +} + +size_t chainSpecSamplingMcWorkspaceBytesPerRow() +{ + return sizeof(RowWorkspace); +} + +} // namespace minimaxM31 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM31/chainSpecSamplingMcSm103.h b/cpp/tensorrt_llm/kernels/minimaxM31/chainSpecSamplingMcSm103.h new file mode 100644 index 000000000000..33a37dc071c7 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/chainSpecSamplingMcSm103.h @@ -0,0 +1,84 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/speculativeDecoding/chainSpeculativeSamplingKernels.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM31 +{ + +//! The widest vocab invokeChainSpecSamplingMc takes (64 tiles of 4096). +constexpr int kChainSpecMcMaxVocab = 64 * 4096; + +//! The most draft tokens per request invokeChainSpecSamplingMc takes (its acceptance probabilities sit in shared +//! memory). +constexpr int kChainSpecMcMaxDraft = 32; + +//! ST-01-REJ: chain speculative sampling (kernels/speculativeDecoding/chainSpeculativeSamplingKernels.cu, the +//! VEC_SIZE 4 kernel behind trtllm::chain_speculative_sampling) as ONE launch of batch x ctasPerRow CTAs (1024 threads +//! each) instead of one CTA per request. +//! +//! Semantics: the stock kernel's (see ChainSpeculativeSamplingParams): per request, the Philox acceptance loop, then +//! the corrected (or bonus) token by an inverse-CDF search over max(0, target - draft) (or the target), the outputs +//! [batch, numDraft + 1] tokens and [batch] counts. +//! +//! Bitwise contract vs the stock kernel (tokens and counts, every input of its vectorized path): +//! 1. every CTA of a request runs the stock tid-0 acceptance loop unchanged (same seed, subsequence +//! kPerRowSubsequenceBase | offset, same draws), so all agree on the rejected position and the uniform; +//! 2. the (target, residual) mass: the stock kernel's 1024 per-thread chains (thread t adds the elements at +//! 4t + 4096k in k order) are spread over the request's CTAs unchanged, and one CTA reduces them with the same +//! cub BlockReduce and Float2SumOp; +//! 3. each 4096-element tile's sum is the stock cub BlockReduce::Sum over the same per-thread vec4, and +//! each tile's last index with mass is a max (order-free); +//! 4. one CTA walks the tile sums in tile order with the stock fp32 prefix adds, and in the tile where the threshold +//! falls runs the stock cub BlockScan::InclusiveSum and atomicMin search; the fallback (no winner) is +//! the last index with mass over the tiles the stock loop visits. +//! Validated bitwise (tokens + counts) vs trtllm::chain_speculative_sampling on GB300, K7, vocab 200064, batch +//! 1 / 2 / 4 / 8 x 5 cases (realistic, same, far, ties, onehot) x 64 seeds at ctasPerRow 4 / 8 / 16, eager and graph +//! (microbenchmark): 31.2 vs 41.2 us per call at batch 1 with ctasPerRow 16. +//! +//! Compiled with the repo's default flags, as the stock kernel (only adds, max and compares on the probs; no +//! --use_fast_math). +//! +//! Contract: vocabSize % 4 == 0, vocabSize <= kChainSpecMcMaxVocab, draftProbs / targetProbs 16-byte aligned (the +//! stock kernel's VEC_SIZE 4 path); 1024 % ctasPerRow == 0. The request's CTAs synchronize through global-memory +//! barriers, so all batch x ctasPerRow CTAs must be co-resident (one 1024-thread CTA per SM): pick ctasPerRow with +//! batch x ctasPerRow <= the SM count (ctasPerRow 1 needs no co-residency). +//! +//! Workspace: chainSpecSamplingMcWorkspaceBytesPerRow() x batch bytes, 16-byte aligned, ZERO-initialized once before +//! the first launch and then reused as is: every launch leaves it ready for the next (arrival counters reset, the +//! barrier generation only increases), so it is CUDA-graph safe without an init kernel. Launches that share one +//! workspace must be stream-ordered. +cudaError_t invokeChainSpecSamplingMc(speculative_decoding::ChainSpeculativeSamplingParams const& params, + int ctasPerRow, void* workspace, cudaStream_t stream); + +//! Workspace bytes per request of invokeChainSpecSamplingMc. +size_t chainSpecSamplingMcWorkspaceBytesPerRow(); + +} // namespace minimaxM31 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM31/ctxResidualRmsNormMxfp8Sm103.cu b/cpp/tensorrt_llm/kernels/minimaxM31/ctxResidualRmsNormMxfp8Sm103.cu new file mode 100644 index 000000000000..9627b1ce8cd0 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/ctxResidualRmsNormMxfp8Sm103.cu @@ -0,0 +1,232 @@ +/* + * 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/minimaxM31/ctxResidualRmsNormMxfp8Sm103.h" +#include "tensorrt_llm/kernels/quantization.cuh" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::minimaxM31Sm103 +{ +namespace +{ + +constexpr int kPackedSize = 8; // BF16 elements per 16-byte access, = CVT_ELTS_PER_THREAD +constexpr int kWarpSize = 32; + +union PackedBf16 +{ + int4 packed; + __nv_bfloat162 unpacked[4]; +}; + +// The helpers below are customAllReduceKernels.cu's (reduce_fusion::warp_reduce_sum / block_reduce_sum, +// accumulate, rms_norm, add128b): the same arithmetic in the same order, so the norm is bitwise +// rms_norm_rows_kernel's. +__device__ __forceinline__ float warpReduceSum(float val) +{ + val += __shfl_xor_sync(~0, val, 16); + val += __shfl_xor_sync(~0, val, 8); + val += __shfl_xor_sync(~0, val, 4); + val += __shfl_xor_sync(~0, val, 2); + val += __shfl_xor_sync(~0, val, 1); + return val; +} + +__device__ __forceinline__ float blockReduceSum(float val) +{ + __shared__ float smem[kWarpSize]; + int const laneId = threadIdx.x % kWarpSize; + int const warpId = threadIdx.x / kWarpSize; + int const warpNum = blockDim.x / kWarpSize; + val = warpReduceSum(val); + if (laneId == 0) + { + smem[warpId] = val; + } + __syncthreads(); + val = laneId < warpNum ? smem[laneId] : 0.f; + val = warpReduceSum(val); + return val; +} + +__device__ __forceinline__ int4 add128b(PackedBf16& a, PackedBf16& b) +{ + PackedBf16 c; + c.unpacked[0] = a.unpacked[0] + b.unpacked[0]; + c.unpacked[1] = a.unpacked[1] + b.unpacked[1]; + c.unpacked[2] = a.unpacked[2] + b.unpacked[2]; + c.unpacked[3] = a.unpacked[3] + b.unpacked[3]; + return c.packed; +} + +__device__ __forceinline__ float accumulate(float acc, PackedBf16& vec) +{ +#pragma unroll + for (int i = 0; i < kPackedSize; ++i) + { + float v = static_cast(reinterpret_cast<__nv_bfloat16*>(vec.unpacked)[i]); + acc += v * v; + } + return acc; +} + +__device__ __forceinline__ int4 rmsNorm(float denom, PackedBf16& vec, PackedBf16& weight) +{ + PackedBf16 ret; +#pragma unroll + for (int i = 0; i < kPackedSize; ++i) + { + float v1 = static_cast(reinterpret_cast<__nv_bfloat16*>(vec.unpacked)[i]); + float v2 = static_cast(reinterpret_cast<__nv_bfloat16*>(weight.unpacked)[i]); + reinterpret_cast<__nv_bfloat16*>(ret.unpacked)[i] = static_cast<__nv_bfloat16>(v1 * denom * v2); + } + return ret.packed; +} + +//! rms_norm_rows_kernel (persistent rows, next row prefetched, +//! one 8-element chunk per thread) whose per-row output also goes through quantize_mxfp8_rows' per-element math +//! (cvt_warp_fp16_to_mxfp8, swizzled K32 scales) while it is still in registers. The scale rows that pad +//! the row count to a multiple of 128 are zero-filled, as quantize_mxfp8_rows does. +template +__global__ void ctxResidualRmsNormMxfp8Kernel(__nv_bfloat16* inter, __nv_bfloat16 const* residual, + __nv_bfloat16 const* weight, __nv_bfloat16* normOut, uint64_t* qOut, uint32_t* sfOut, uint32_t* fp4Out, + uint32_t* fp4SfOut, float const* fp4GlobalScale, int rows, int hidden, float eps) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + static constexpr int kThreadsPerSf = 32 / CVT_ELTS_PER_THREAD; + int const colIdx = threadIdx.x; + int const offset = colIdx * kPackedSize; + int const colThreads = hidden / kPackedSize; + std::optional const batchIdx = 0; + std::optional const numRows = rows; + + cudaGridDependencySynchronize(); + + // quantize_with_block_size's per-tensor scale (SFScaleVal) for the NVFP4 copy. + float const fp4SfScale = kFp4 ? fp4GlobalScale[0] : 1.0f; + PackedBf16 weightVec; + weightVec.packed = *reinterpret_cast(weight + offset); + PackedBf16 interVec, residualVec; + int row = blockIdx.x; + if (row < rows) + { + interVec.packed = *reinterpret_cast(inter + static_cast(row) * hidden + offset); + residualVec.packed = *reinterpret_cast(residual + static_cast(row) * hidden + offset); + } + for (; row < rows; row += gridDim.x) + { + int64_t const rowOffset = static_cast(row) * hidden + offset; + int const next = row + gridDim.x; + PackedBf16 nextInterVec, nextResidualVec; + if (next < rows) + { + nextInterVec.packed = *reinterpret_cast(inter + static_cast(next) * hidden + offset); + nextResidualVec.packed + = *reinterpret_cast(residual + static_cast(next) * hidden + offset); + } + interVec.packed = add128b(interVec, residualVec); + *reinterpret_cast(inter + rowOffset) = interVec.packed; + float acc = accumulate(0.f, interVec); + acc = blockReduceSum(acc); + float const denom = rsqrtf(acc / hidden + eps); + PackedBf16 outVec; + outVec.packed = rmsNorm(denom, interVec, weightVec); + if constexpr (kWriteNorm) + { + *reinterpret_cast(normOut + rowOffset) = outVec.packed; + } + PackedVec<__nv_bfloat16> quantIn; + *reinterpret_cast(&quantIn) = outVec.packed; + auto sf = cvt_quant_get_sf_out_offset( + batchIdx, row, colIdx, numRows, hidden / 32, sfOut, QuantizationSFLayout::SWIZZLED); + qOut[static_cast(row) * colThreads + colIdx] = cvt_warp_fp16_to_mxfp8<__nv_bfloat16, 32>(quantIn, sf); + if constexpr (kFp4) + { + // fp4_quantize(norm, globalScale, 16, ue8m0 = false, swizzled = false): E2M1 x 8 per thread, E4M3 scale per + // 16 elements in the linear layout (no row padding). + PackedVec<__nv_bfloat16> fp4In; + *reinterpret_cast(&fp4In) = outVec.packed; + auto sf4 = cvt_quant_get_sf_out_offset( + batchIdx, row, colIdx, numRows, hidden / 16, fp4SfOut, QuantizationSFLayout::LINEAR); + fp4Out[static_cast(row) * colThreads + colIdx] + = cvt_warp_fp16_to_fp4<__nv_bfloat16, 16, false>(fp4In, fp4SfScale, sf4); + } + // blockReduceSum's shared memory is written again by the next row. + __syncthreads(); + interVec.packed = nextInterVec.packed; + residualVec.packed = nextResidualVec.packed; + } + int const paddedRows = (rows + 127) / 128 * 128; + for (int padRow = rows + blockIdx.x; padRow < paddedRows; padRow += gridDim.x) + { + auto sf = cvt_quant_get_sf_out_offset( + batchIdx, padRow, colIdx, numRows, hidden / 32, sfOut, QuantizationSFLayout::SWIZZLED); + cvt_store_replicated_sf<1>(sf, 0x00); + } + // Every thread's stores are visible before any thread reports PDL completion (as in quantize_mxfp8_rows). + __threadfence(); + __syncthreads(); + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +} // namespace + +void invokeCtxResidualRmsNormMxfp8(__nv_bfloat16* inter, __nv_bfloat16 const* residual, __nv_bfloat16 const* weight, + __nv_bfloat16* normOut, uint64_t* qOut, uint32_t* sfOut, uint32_t* fp4Out, uint32_t* fp4SfOut, + float const* fp4GlobalScale, int rows, int hidden, float eps, bool enablePdl, cudaStream_t stream) +{ + static int const multiProcessorCount = [] + { + int device = 0; + int count = 0; + cudaGetDevice(&device); + cudaDeviceGetAttribute(&count, cudaDevAttrMultiProcessorCount, device); + return count; + }(); + cudaLaunchConfig_t config = {0}; + config.gridDim = std::max(1, std::min(rows, multiProcessorCount)); + config.blockDim = hidden / kPackedSize; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = enablePdl; + config.attrs = attribute; + config.numAttrs = 1; + auto launch = [&](auto kernel) + { + cudaLaunchKernelEx(&config, kernel, inter, residual, weight, normOut, qOut, sfOut, fp4Out, fp4SfOut, + fp4GlobalScale, rows, hidden, eps); + }; + bool const fp4 = fp4Out != nullptr; + if (normOut != nullptr) + { + fp4 ? launch(ctxResidualRmsNormMxfp8Kernel) : launch(ctxResidualRmsNormMxfp8Kernel); + } + else + { + fp4 ? launch(ctxResidualRmsNormMxfp8Kernel) : launch(ctxResidualRmsNormMxfp8Kernel); + } +} + +} // namespace kernels::minimaxM31Sm103 + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM31/ctxResidualRmsNormMxfp8Sm103.h b/cpp/tensorrt_llm/kernels/minimaxM31/ctxResidualRmsNormMxfp8Sm103.h new file mode 100644 index 000000000000..9cba469d348a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/ctxResidualRmsNormMxfp8Sm103.h @@ -0,0 +1,43 @@ +/* + * 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::minimaxM31Sm103 +{ + +//! MiniMax-M3 CTX prefill: the residual add + RMSNorm of the all-reduce's RESIDUAL_RMS_NORM epilogue +//! (rms_norm_rows_kernel, bitwise the same) fused with the MXFP8 quantization of its output that the next MXFP8 +//! linear performs (quantize_mxfp8_rows: K32 UE8M0 scales, swizzled 128x4 layout, bitwise the same). +//! In place: `inter` (the all-reduce output) becomes `inter + residual`, as in the stock epilogue. +//! normOut may be null (the BF16 norm output is then not written). With fp4Out (and fp4SfOut / fp4GlobalScale) set, +//! it also writes fp4_quantize(norm, globalScale, 16, ue8m0 = false, swizzled = false) of the same rows. +//! hidden % 256 == 0, hidden / 8 <= 1024. +void invokeCtxResidualRmsNormMxfp8(__nv_bfloat16* inter, __nv_bfloat16 const* residual, __nv_bfloat16 const* weight, + __nv_bfloat16* normOut, uint64_t* qOut, uint32_t* sfOut, uint32_t* fp4Out, uint32_t* fp4SfOut, + float const* fp4GlobalScale, int rows, int hidden, float eps, bool enablePdl, cudaStream_t stream); + +} // namespace kernels::minimaxM31Sm103 + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM31/selectPrefillTileSm103.cu b/cpp/tensorrt_llm/kernels/minimaxM31/selectPrefillTileSm103.cu new file mode 100644 index 000000000000..2a40178ca22c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/selectPrefillTileSm103.cu @@ -0,0 +1,545 @@ +/* + * 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.1 GB300 (sm_103) target-local copy of the #20629 prefill tile select +// (cpp/tensorrt_llm/kernels/minimaxM3/selectPrefillTile.* @ cf8a77e816); namespace and op name are M3.1's. + +// Prefill top-16 block select over head-major scores [num_kv_heads, num_blocks, total_q] with unit query stride. +// +// One CTA per (KV head, 8 consecutive queries). The 8 queries of one block are one 32-byte sector, so the CTA reads +// each sector it needs exactly once, as two 16-byte loads. The tile's rows are staged in shared memory as ordered keys +// of the forced scores. A long tile's key rows fill most of shared memory, so one CTA runs per SM; its loads then come +// from extra staging warps (kWideThreads from kWideMinBlocks blocks), which exit before the selection. Each of the +// first 8 warps then selects one row: +// - the 32 lane-stride group minima are sorted in the warp and the 16th of them bounds the row's 16th smallest key, +// so every key at or below that bound is a candidate; the candidates (compacted in ascending block order) are +// ranked exactly by (key, block); +// - a row with more than kMaxCand candidates (heavy ties) falls back to an exact 4 x 8-bit radix select. +// Ordering is the descending ordered key of the float bits used by minimax_m3_select_blocks, ties to the smaller +// block, so the result is the same top-16 set. + +#include "tensorrt_llm/kernels/minimaxM31/selectPrefillTileSm103.h" + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/logger.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::minimaxM31Sm103 +{ + +namespace +{ +constexpr int kTopK = 16; +constexpr int kTileQ = 8; +constexpr int kThreads = kTileQ * 32; +// From this many blocks on, the tile is staged by kWideThreads threads with kWideUnroll blocks in flight each. +constexpr int kWideMinBlocks = 2600; +constexpr int kWideThreads = 1024; +constexpr int kWideUnroll = 4; +constexpr uint32_t kFull = 0xFFFFFFFFu; +constexpr float kInitScore = 1.0e30f; +constexpr float kLocalScore = 1.0e29f; +constexpr uint32_t kNegInfKey = 0xFF800000u; +// Per-warp scratch: the radix histogram (256 bins) or the candidate list (keys then blocks). +constexpr int kMaxCand = 256; +constexpr int kWarpScratch = 2 * kMaxCand; +// Blocks each thread has in flight while staging a narrow tile. +constexpr int kStageUnroll = 8; + +__device__ __forceinline__ uint32_t orderedKey(float x) +{ + uint32_t const bits = __float_as_uint(x); + return (bits & 0x80000000u) ? bits : (~bits & 0x7fffffffu); +} + +// Row pitch in keys: >= num_blocks and == 4 (mod 32), so the staging writes of +// 8 rows x 4 consecutive blocks per warp hit 32 distinct banks. +__host__ __device__ __forceinline__ int rowPitch(int numBlocks) +{ + int p = (numBlocks + 31) / 32 * 32 + 4; + return p; +} + +// kNt staging threads with kUnroll blocks in flight each; warps 0 .. kTileQ - 1 select. kIlp builds each row's 32 +// stride-class minima while staging (a staging thread only ever sees one class) and issues the candidate scan's +// shared-memory loads four at a time; the minima, the bound and every decision are the same as the scalar scans'. +template +__global__ void __launch_bounds__(kNt) selectPrefillTileKernel(float const* __restrict__ scores, int64_t headStride, + int64_t blockStride, int32_t const* __restrict__ nValid, int32_t* __restrict__ out, int64_t outQueryStride, + int64_t outHeadStride, int numBlocks, int totalQ, int initBlocks, int localBlocks) +{ + extern __shared__ uint32_t smem[]; + __shared__ uint32_t sClassMin[kTileQ * 32]; // [row][b % 32] (kIlp) + int const pitch = rowPitch(numBlocks); + uint32_t* keys = smem; // [kTileQ][pitch] + uint32_t* hist = smem + kTileQ * pitch; // [kTileQ][kWarpScratch] + + int const head = blockIdx.y; + int const q0 = blockIdx.x * kTileQ; + int const tid = threadIdx.x; + int const warp = tid / 32; + int const lane = tid % 32; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); +#endif + + // ---- stage: thread loads queries 4 * jh .. 4 * jh + 3 of block b = tid / 2 + (kNt / 2) * i as one float4 ---- + { + int const jh = tid % 2; + int rowEnd[4]; + int64_t localStart[4]; + int tileEnd = 0; +#pragma unroll + for (int c = 0; c < 4; ++c) + { + int const q = q0 + 4 * jh + c; + rowEnd[c] = 0; + localStart[c] = 0; + if (q < totalQ) + { + int const raw = nValid[q]; + rowEnd[c] = max(0, min(raw, numBlocks)); + localStart[c] = max(static_cast(raw) - localBlocks, static_cast(0)); + } + tileEnd = max(tileEnd, rowEnd[c]); + } + tileEnd = max(tileEnd, __shfl_xor_sync(kFull, tileEnd, 1)); + if constexpr (kIlp) + { + if (tid < kTileQ * 32) + { + sClassMin[tid] = 0xFFFFFFFFu; + } + } + __shared__ int sTileEnd[kNt / 32]; + if (lane == 0) + { + sTileEnd[warp] = tileEnd; + } + __syncthreads(); + tileEnd = 0; +#pragma unroll + for (int w = 0; w < kNt / 32; ++w) + { + tileEnd = max(tileEnd, sTileEnd[w]); + } + // Two threads read a tile's 8 queries of one block, one aligned 32-byte sector, when the base and the + // strides keep 16-byte alignment; otherwise, and for a partial last tile, they read one float at a time. + bool const fullVec = (reinterpret_cast(scores) & 15) == 0 && ((headStride | blockStride) & 3) == 0 + && (q0 + 4 * jh + 3) < totalQ; + float const* base = scores + static_cast(head) * headStride + q0 + 4 * jh; + constexpr int kBlocksPerIter = kNt / 2; + static_assert(kBlocksPerIter % 32 == 0, "a staging thread stays in one stride class b % 32"); + uint32_t classMin[4] = {0xFFFFFFFFu, 0xFFFFFFFFu, 0xFFFFFFFFu, 0xFFFFFFFFu}; + for (int b0 = tid / 2; b0 < tileEnd; b0 += kBlocksPerIter * kUnroll) + { + float4 v[kUnroll]; +#pragma unroll + for (int u = 0; u < kUnroll; ++u) + { + int const b = b0 + u * kBlocksPerIter; + float4 x = make_float4(0.f, 0.f, 0.f, 0.f); + if (b < tileEnd) + { + float const* p = base + static_cast(b) * blockStride; + if (fullVec) + { + x = __ldg(reinterpret_cast(p)); + } + else + { + x.x = b < rowEnd[0] ? __ldg(p) : 0.f; + x.y = b < rowEnd[1] ? __ldg(p + 1) : 0.f; + x.z = b < rowEnd[2] ? __ldg(p + 2) : 0.f; + x.w = b < rowEnd[3] ? __ldg(p + 3) : 0.f; + } + } + v[u] = x; + } +#pragma unroll + for (int u = 0; u < kUnroll; ++u) + { + int const b = b0 + u * kBlocksPerIter; + if (b < tileEnd) + { + float const vals[4] = {v[u].x, v[u].y, v[u].z, v[u].w}; +#pragma unroll + for (int c = 0; c < 4; ++c) + { + uint32_t key = kNegInfKey; + if (b < rowEnd[c]) + { + float s = vals[c]; + if (b < initBlocks) + { + s = kInitScore; + } + if (b >= localStart[c]) + { + s = kLocalScore; + } + key = orderedKey(s); + if constexpr (kIlp) + { + classMin[c] = min(classMin[c], key); + } + } + keys[(4 * jh + c) * pitch + b] = key; + } + } + } + } + if constexpr (kIlp) + { + int const cls = (tid / 2) % 32; +#pragma unroll + for (int c = 0; c < 4; ++c) + { + if (classMin[c] != 0xFFFFFFFFu) + { + atomicMin(&sClassMin[(4 * jh + c) * 32 + cls], classMin[c]); + } + } + } + } + __syncthreads(); + if (warp >= kTileQ) + { + return; + } + + // ---- select: warp w owns query q0 + w ---- + int const q = q0 + warp; + if (q >= totalQ) + { + return; + } + int const raw = nValid[q]; + int const len = max(0, min(raw, numBlocks)); + uint32_t const* row = keys + warp * pitch; + uint32_t* h = hist + warp * kWarpScratch; + int32_t* dst = out + static_cast(q) * outQueryStride + static_cast(head) * outHeadStride; + + // Threshold key K and how many keys equal to K to take (tieTake). Keys < K are all in. + if (len > kTopK) + { + // Fast path. Group g = lanes' stride classes b % 32: each lane's best (smallest) key. + // T = the 16th smallest of the 32 group bests, so at least 16 keys are <= T and the + // row's 16th smallest key is <= T: every top-16 key (and every key tied with one) is a + // candidate. Candidates are compacted in ascending block order and ranked exactly. + uint32_t best = 0xFFFFFFFFu; + bool const any = lane < len; + if constexpr (kIlp) + { + best = sClassMin[warp * 32 + lane]; + } + else + { + for (int b = lane; b < len; b += 32) + { + best = min(best, row[b]); + } + } + // Ascending bitonic sort of the 32 group bests; empty groups sort last. + uint32_t v = any ? best : 0xFFFFFFFFu; + bool vEmpty = !any; +#pragma unroll + for (int size = 2; size <= 32; size *= 2) + { +#pragma unroll + for (int stride = size / 2; stride > 0; stride /= 2) + { + uint32_t const pv = __shfl_xor_sync(kFull, v, stride); + bool const pe = __shfl_xor_sync(kFull, vEmpty, stride); + bool const up = ((lane & size) == 0) == ((lane & stride) == 0); + // order: (empty, value); smaller first when 'up' + bool const pLess = (!pe && vEmpty) || (pe == vEmpty && pv < v); + bool const pGreater = (pe && !vEmpty) || (pe == vEmpty && pv > v); + if ((up && pLess) || (!up && pGreater)) + { + v = pv; + vEmpty = pe; + } + } + } + uint32_t const bound = __shfl_sync(kFull, v, kTopK - 1); + uint32_t* candKey = h; + uint32_t* candBlk = h + kMaxCand; + int nCand = 0; + int cb0 = 0; + if constexpr (kIlp) + { + for (; cb0 + 128 <= len; cb0 += 128) + { + uint32_t k[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + k[j] = row[cb0 + 32 * j + lane]; + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + bool const c = k[j] <= bound; + uint32_t const m = __ballot_sync(kFull, c); + int const slot = nCand + __popc(m & ((1u << lane) - 1u)); + if (c && slot < kMaxCand) + { + candKey[slot] = k[j]; + candBlk[slot] = static_cast(cb0 + 32 * j + lane); + } + nCand += __popc(m); + } + } + } + for (int b0 = cb0; b0 < len; b0 += 32) + { + int const b = b0 + lane; + bool const c = b < len && row[b] <= bound; + uint32_t const m = __ballot_sync(kFull, c); + int const slot = nCand + __popc(m & ((1u << lane) - 1u)); + if (c && slot < kMaxCand) + { + candKey[slot] = row[b]; + candBlk[slot] = static_cast(b); + } + nCand += __popc(m); + } + __syncwarp(); + if (nCand <= kMaxCand) + { + // Exact rank by (key, block); candidates are already in ascending block order, so a + // second ballot compaction of the rank < 16 ones writes the output ascending. + int written = 0; + for (int i0 = 0; i0 < nCand; i0 += 32) + { + int const i = i0 + lane; + bool sel = false; + uint32_t k = 0xFFFFFFFFu; + uint32_t blk = 0; + if (i < nCand) + { + k = candKey[i]; + blk = candBlk[i]; + int rank = 0; + for (int jdx = 0; jdx < nCand; ++jdx) + { + uint32_t const kj = candKey[jdx]; + rank += (kj < k) || (kj == k && jdx < i); + } + sel = rank < kTopK; + } + uint32_t const m = __ballot_sync(kFull, sel && k != kNegInfKey); + int const slot = written + __popc(m & ((1u << lane) - 1u)); + if (sel && k != kNegInfKey) + { + dst[slot] = static_cast(blk); + } + written += __popc(m); + } + for (int r = written + lane; r < kTopK; r += 32) + { + dst[r] = -1; + } + return; + } + __syncwarp(); + } + + uint32_t thresh = 0xFFFFFFFFu; + int tieTake = kTopK; + if (len > kTopK) + { + uint32_t prefix = 0; + int remaining = kTopK; // rank (1-based) of the wanted key among keys matching prefix +#pragma unroll 1 + for (int pass = 0; pass < 4; ++pass) + { + int const shift = 24 - 8 * pass; + uint32_t const hiMask = pass == 0 ? 0u : (0xFFFFFFFFu << (shift + 8)); + for (int i = lane; i < 256; i += 32) + { + h[i] = 0; + } + __syncwarp(); + for (int b0 = 0; b0 < len; b0 += 32) + { + int const b = b0 + lane; + bool const active = b < len; + uint32_t const k = active ? row[b] : 0u; + bool const match = active && ((k & hiMask) == prefix); + uint32_t const digit = (k >> shift) & 0xFFu; + uint32_t const m = __ballot_sync(kFull, match); + if (match) + { + uint32_t const peers = __match_any_sync(m, digit); + if ((__ffs(peers) - 1) == lane) + { + atomicAdd(&h[digit], static_cast(__popc(peers))); + } + } + } + __syncwarp(); + // Find the digit holding the remaining-th smallest key: each lane owns 8 bins. + uint32_t c[8]; + uint32_t local = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + c[i] = h[lane * 8 + i]; + local += c[i]; + } + uint32_t incl = local; +#pragma unroll + for (int o = 1; o < 32; o *= 2) + { + uint32_t const v = __shfl_up_sync(kFull, incl, o); + if (lane >= o) + { + incl += v; + } + } + uint32_t const excl = incl - local; + bool const mine = excl < static_cast(remaining) && incl >= static_cast(remaining); + uint32_t const owner = __ballot_sync(kFull, mine); + int const ownerLane = __ffs(owner) - 1; + int digitSel = 0; + uint32_t below = 0; + if (lane == ownerLane) + { + uint32_t acc = excl; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + if (acc + c[i] >= static_cast(remaining)) + { + digitSel = lane * 8 + i; + below = acc; + break; + } + acc += c[i]; + } + } + digitSel = __shfl_sync(kFull, digitSel, ownerLane); + below = __shfl_sync(kFull, below, ownerLane); + prefix |= static_cast(digitSel) << shift; + remaining -= static_cast(below); + __syncwarp(); + } + thresh = prefix; + tieTake = remaining; // keys equal to thresh to take, smallest blocks first + } + + // ---- collect in ascending block order ---- + // len <= 16: every valid block is a pick (thresh = max, tieTake = 16 covers all). + int count = 0; + int ties = 0; + for (int b0 = 0; b0 < len && count < kTopK; b0 += 32) + { + int const b = b0 + lane; + bool const active = b < len; + uint32_t const k = active ? row[b] : 0xFFFFFFFFu; + bool const strict = active && (len <= kTopK || k < thresh); + bool const tie = active && len > kTopK && k == thresh; + uint32_t const tieMask = __ballot_sync(kFull, tie); + int const tieRank = ties + __popc(tieMask & ((1u << lane) - 1u)); + bool const take = strict || (tie && tieRank < tieTake); + uint32_t const takeMask = __ballot_sync(kFull, take); + int const slot = count + __popc(takeMask & ((1u << lane) - 1u)); + if (take && slot < kTopK) + { + dst[slot] = (k == kNegInfKey) ? -1 : b; + } + count += __popc(takeMask); + ties += __popc(tieMask); + } + __syncwarp(); + // Picks forced to -inf were written as -1 in place; compact the valid ids to the front, -1 tail. + if (lane == 0) + { + int32_t ids[kTopK]; + int n = 0; + int const picks = min(count, kTopK); + for (int i = 0; i < picks; ++i) + { + int32_t const v = dst[i]; + if (v >= 0) + { + ids[n++] = v; + } + } + for (int i = 0; i < kTopK; ++i) + { + dst[i] = i < n ? ids[i] : -1; + } + } +} +} // namespace + +cudaError_t invokeMinimaxM3SelectPrefillTile(float const* scores, int64_t headStride, int64_t blockStride, + 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, cudaStream_t stream) +{ + if (totalQ == 0 || numKvHeads == 0) + { + return cudaSuccess; + } + if (numBlocks > kMinimaxM3SelectPrefillTileMaxBlocks) + { + return cudaErrorInvalidValue; + } + // TRTLLM_M3_SELECT_TILE_V2=1 keeps the previous launch: 256 threads for every tile, one load at a time per scan. + static bool const v2Only = common::getBoolEnv("TRTLLM_M3_SELECT_TILE_V2"); + static std::once_flag logged; + std::call_once(logged, + [] + { + TLLM_LOG_INFO(v2Only + ? "minimaxM3SelectPrefillTile: staging reads each block's sector as two 16-byte loads" + : "minimaxM3SelectPrefillTile: v4, stride-class minima while staging, 4-wide candidate " + "scan and 1024 staging threads from 2600 blocks"); + }); + size_t const smemBytes + = (static_cast(kTileQ) * rowPitch(numBlocks) + kTileQ * kWarpScratch) * sizeof(uint32_t); + bool const wide = !v2Only && numBlocks >= kWideMinBlocks; + auto kernel = v2Only ? selectPrefillTileKernel + : wide ? selectPrefillTileKernel + : selectPrefillTileKernel; + cudaError_t err + = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smemBytes)); + if (err != cudaSuccess) + { + return err; + } + dim3 const grid((totalQ + kTileQ - 1) / kTileQ, numKvHeads); + common::launchWithPdlWhenEnabled("minimaxM3SelectPrefillTile", kernel, grid, dim3(wide ? kWideThreads : kThreads), + smemBytes, stream, scores, headStride, blockStride, nValidBlocks, out, outQueryStride, outHeadStride, numBlocks, + totalQ, initBlocks, localBlocks); + return cudaGetLastError(); +} + +} // namespace kernels::minimaxM31Sm103 + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM31/selectPrefillTileSm103.h b/cpp/tensorrt_llm/kernels/minimaxM31/selectPrefillTileSm103.h new file mode 100644 index 000000000000..0e5d5cc030af --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/selectPrefillTileSm103.h @@ -0,0 +1,51 @@ +/* + * 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.1 GB300 (sm_103) target-local copy of the #20629 prefill tile select +// (cpp/tensorrt_llm/kernels/minimaxM3/selectPrefillTile.* @ cf8a77e816); namespace and op name are M3.1's. + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::minimaxM31Sm103 +{ + +//! Largest row (blocks per query) the prefill tile select stages in shared memory: 8 query rows of keys plus the +//! per-warp scratch stay under 200 KB. +constexpr int32_t kMinimaxM3SelectPrefillTileMaxBlocks = 5856; + +//! MiniMax-M3 top-16 block select for prefill: per-KV-head scores [numKvHeads, numBlocks, totalQ] with unit query +//! stride, one CTA per (KV head, 8 consecutive queries) so every 32-byte sector of the scores is read once. +//! +//! Same contract as minimax_m3_select_blocks: per (query, head) row the clamp(nValidBlocks[query], 0, numBlocks) +//! leading blocks; the last localBlocks of them (counted from the unclamped count) forced to 1e29 and the first +//! initBlocks to 1e30, local winning; ranked by the descending ordered key of the float bits (NaN above +inf, +0 above +//! -0, negative NaN below -inf), ties to the smaller block; picks whose forced score is -inf emitted as -1; ascending +//! ids with -1 tail padding. Writes out[query * outQueryStride + head * outHeadStride + 0..15]. +cudaError_t invokeMinimaxM3SelectPrefillTile(float const* scores, int64_t headStride, int64_t blockStride, + 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, cudaStream_t stream); + +} // namespace kernels::minimaxM31Sm103 + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM31/softmaxMcSm103.cu b/cpp/tensorrt_llm/kernels/minimaxM31/softmaxMcSm103.cu new file mode 100644 index 000000000000..6fc16f1d1147 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/softmaxMcSm103.cu @@ -0,0 +1,295 @@ +/* + * 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. + */ + +// ST-01-SM: FlashInfer softmax(logits, temperature) Path A with the Reduce kernel's +// write pass spread over many CTAs per row. See softmaxMcSm103.h for the contract. Compiled with --use_fast_math (set +// on this source in kernels/CMakeLists.txt), as FlashInfer's JIT compiles it. + +#include "tensorrt_llm/kernels/minimaxM31/softmaxMcSm103.h" + +#include +#include + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM31 +{ + +namespace +{ + +// --------------------------------------------------------------------------------------------------------------------- +// Vendored from FlashInfer (v0.6.15; identical in v0.6.18): include/flashinfer/utils.cuh (ceil_div, round_up), +// include/flashinfer/sampling.cuh (MaxReduceOp, Float2SoftmaxReduceOp, PartialSoftmaxResult) and +// include/flashinfer/vec_dtypes.cuh (vec_t: only the members the kernels below use; cast_load / +// cast_store are the same-type branch of FlashInfer's cast_load_impl / cast_store_impl, i.e. load / store). +// Copyright (c) 2023-2026 by FlashInfer team, licensed under the Apache License, Version 2.0. Copied verbatim +// (formatting kept) so that the element order, the float4 accesses and the reduction operators are FlashInfer's own; +// FlashInfer's headers are not on the TensorRT-LLM C++ include path. +// --------------------------------------------------------------------------------------------------------------------- +namespace fi +{ +// clang-format off +#define FLASHINFER_INLINE inline __attribute__((always_inline)) __device__ + +template +__forceinline__ __device__ __host__ constexpr T1 ceil_div(const T1 x, const T2 y) noexcept { + return (x + y - 1) / y; +} + +template +__forceinline__ __device__ __host__ constexpr T1 round_up(const T1 x, const T2 y) noexcept { + return ceil_div(x, y) * y; +} + +using MaxReduceOp = cuda::maximum<>; + +struct Float2SoftmaxReduceOp { + __device__ __forceinline__ float2 operator()(const float2& a, const float2& b) const { + if (isinf(a.x)) return b; + if (isinf(b.x)) return a; + + float new_max = max(a.x, b.x); + float new_denom = a.y * __expf(a.x - new_max) + b.y * __expf(b.x - new_max); + return make_float2(new_max, new_denom); + } +}; + +struct PartialSoftmaxResult { + float max_val; + float denominator; +}; + +template +struct vec_t; + +// float x 4 or more +template +struct vec_t { + static_assert(vec_size % 4 == 0, "Invalid vector size"); + float4 data[vec_size / 4]; + + FLASHINFER_INLINE float& operator[](size_t i) { return ((float*)(data))[i]; } + FLASHINFER_INLINE const float& operator[](size_t i) const { return ((const float*)(data))[i]; } + FLASHINFER_INLINE void fill(float val) { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + data[i] = make_float4(val, val, val, val); + } + } + FLASHINFER_INLINE void load(const float* ptr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + data[i] = ((float4*)ptr)[i]; + } + } + FLASHINFER_INLINE void store(float* ptr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + ((float4*)ptr)[i] = data[i]; + } + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + static_assert(std::is_same_v, "only the same-type branch of cast_load_impl is vendored"); + load(ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + static_assert(std::is_same_v, "only the same-type branch of cast_store_impl is vendored"); + store(ptr); + } +}; + +#undef FLASHINFER_INLINE +// clang-format on +} // namespace fi + +// FlashInfer's OnlineSoftmax on sm_80+ (DISPATCH_COMPUTE_CAP_NUM_THREADS) and vec_size = gcd(16 / sizeof(float), d) +// = 4 for d % 4 == 0. +constexpr uint32_t kThreads = 1024; +constexpr uint32_t kVec = 4; + +// FlashInfer OnlineSoftmaxMapKernel<1024, 4, float>: the same body (the shared temp storage is static instead of +// FlashInfer's dynamic OnlineSoftmaxTempStorage union; no PDL griddepcontrol, as the launches here are plain). +__global__ void __launch_bounds__(kThreads) softmaxMcMapKernel(float const* logits, fi::PartialSoftmaxResult* partials, + float const* temperature_arr, float temperature_val, uint32_t d, uint32_t num_slices) +{ + // clang-format off + const uint32_t bx = blockIdx.x; + const uint32_t by = blockIdx.y; // slice index + const uint32_t tx = threadIdx.x; + float temperature = temperature_arr == nullptr ? temperature_val : temperature_arr[bx]; + const float inv_temp = (temperature == 0.f) ? 0.f : 1.f / temperature; + const uint32_t vec_alignment_elems = alignof(fi::vec_t) / sizeof(float); + const uint32_t slice_stride = fi::round_up(fi::ceil_div(d, num_slices), vec_alignment_elems); + const uint32_t slice_start = by * slice_stride; + const uint32_t slice_size = min((by + 1) * slice_stride, d) - slice_start; + if (slice_start >= d) return; + __shared__ typename cub::BlockReduce::TempStorage reduce; + __shared__ float s_max, s_den; + fi::vec_t logits_vec; + float running_max = -cuda::std::numeric_limits::infinity(); + float running_denominator = 0.0f; + float threadlocal_running_denominator = 0.0f; +#pragma unroll 2 + for (uint32_t i = 0; i < fi::ceil_div(slice_size, kThreads * kVec); ++i) { + logits_vec.fill(-cuda::std::numeric_limits::infinity()); + if ((i * kThreads + tx) * kVec < slice_size) { + logits_vec.cast_load(logits + bx * d + slice_start + (i * kThreads + tx) * kVec); + } + float thread_max = -cuda::std::numeric_limits::infinity(); +#pragma unroll + for (uint32_t j = 0; j < kVec; ++j) { + logits_vec[j] *= inv_temp; + thread_max = max(thread_max, logits_vec[j]); + } + float block_max = cub::BlockReduce(reduce).Reduce(thread_max, fi::MaxReduceOp{}); + if (tx == 0) s_max = block_max; + __syncthreads(); + block_max = s_max; + // if block_max is -inf, then this block contains all -inf values, so we can skip updating + if (!isinf(block_max)) { + float threadlocal_sum = 0.0f; +#pragma unroll + for (uint32_t j = 0; j < kVec; ++j) { + threadlocal_sum += __expf(logits_vec[j] - block_max); + } + float new_max = max(running_max, block_max); + threadlocal_running_denominator = threadlocal_running_denominator * __expf(running_max - new_max) + + threadlocal_sum * __expf(block_max - new_max); + running_max = new_max; + } + } + running_denominator = cub::BlockReduce(reduce).Sum(threadlocal_running_denominator); + if (tx == 0) s_den = running_denominator; + __syncthreads(); + running_denominator = s_den; + if (tx == 0) { + partials[bx * num_slices + by] = {running_max, running_denominator}; + } + // clang-format on +} + +// FlashInfer OnlineSoftmaxReduceKernel<1024, 4, float> with its write loop split over gridDim.y CTAs per row: the +// prologue (the partials' reduction) is FlashInfer's, unchanged, in every CTA of the row. +__global__ void __launch_bounds__(kThreads) + softmaxMcReduceWriteKernel(float const* logits, float* output, fi::PartialSoftmaxResult const* partials, + float const* temperature_arr, float temperature_val, uint32_t d, uint32_t num_slices) +{ + // clang-format off + const uint32_t bx = blockIdx.x; + const uint32_t tx = threadIdx.x; + float temperature = temperature_arr == nullptr ? temperature_val : temperature_arr[bx]; + const float inv_temp = (temperature == 0.f) ? 0.f : 1.f / temperature; + __shared__ typename cub::BlockReduce::TempStorage reduce_pair; + __shared__ float s_max, s_den; + const fi::Float2SoftmaxReduceOp reduce_op; + float2 thread_aggregate = make_float2(-cuda::std::numeric_limits::infinity(), 0.0f); + for (uint32_t i = tx; i < num_slices; i += kThreads) { + fi::PartialSoftmaxResult partial = partials[bx * num_slices + i]; + float2 partial_pair = make_float2(partial.max_val, partial.denominator); + thread_aggregate = reduce_op(thread_aggregate, partial_pair); + } + float2 block_result = cub::BlockReduce(reduce_pair).Reduce(thread_aggregate, reduce_op); + if (tx == 0) { + s_max = block_result.x; + s_den = block_result.y; + } + __syncthreads(); + const float final_max = s_max; + const float inv_denominator = 1.0f / s_den; + // FlashInfer's write loop iterates i over ceil_div(d, 1024 * 4) tiles of 4096 elements; this CTA takes tiles + // i = blockIdx.y, blockIdx.y + gridDim.y, ... (elementwise, so the split does not change any value). + fi::vec_t logits_vec; + fi::vec_t prob_vec; + for (uint32_t i = blockIdx.y; i < fi::ceil_div(d, kThreads * kVec); i += gridDim.y) { + logits_vec.fill(-cuda::std::numeric_limits::infinity()); + if ((i * kThreads + tx) * kVec < d) { + logits_vec.cast_load(logits + bx * d + (i * kThreads + tx) * kVec); + } +#pragma unroll + for (uint32_t j = 0; j < kVec; ++j) { + logits_vec[j] *= inv_temp; + float p = __expf(static_cast(logits_vec[j]) - final_max) * inv_denominator; + prob_vec[j] = static_cast(p); + } + if ((i * kThreads + tx) * kVec < d) { + prob_vec.cast_store(output + bx * d + (i * kThreads + tx) * kVec); + } + } + // clang-format on +} + +uint32_t numSlices(int d) +{ + return fi::ceil_div(static_cast(d), static_cast(kSoftmaxMcSliceSize)); +} + +} // namespace + +cudaError_t invokeSoftmaxMc(float const* logits, float* out, float const* temperatureArr, float temperatureVal, + int rows, int d, int writeCtas, void* partials, cudaStream_t stream) +{ + if (rows <= 0 || d <= 0) + { + return cudaSuccess; + } + if (d % static_cast(kVec) != 0 || writeCtas < 1 || writeCtas > 65535 || rows > 65535) + { + return cudaErrorInvalidValue; + } + uint32_t const slices = numSlices(d); + auto* part = static_cast(partials); + cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(static_cast(rows), slices); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = 0; + cfg.stream = stream; + cudaError_t const err = cudaLaunchKernelEx( + &cfg, softmaxMcMapKernel, logits, part, temperatureArr, temperatureVal, static_cast(d), slices); + if (err != cudaSuccess) + { + return err; + } + cfg.gridDim = dim3(static_cast(rows), static_cast(writeCtas)); + return cudaLaunchKernelEx(&cfg, softmaxMcReduceWriteKernel, logits, out, + static_cast(part), temperatureArr, temperatureVal, static_cast(d), + slices); +} + +size_t softmaxMcPartialsBytes(int rows, int d) +{ + if (rows <= 0 || d <= 0) + { + return 0; + } + return static_cast(rows) * numSlices(d) * sizeof(fi::PartialSoftmaxResult); +} + +} // namespace minimaxM31 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM31/softmaxMcSm103.h b/cpp/tensorrt_llm/kernels/minimaxM31/softmaxMcSm103.h new file mode 100644 index 000000000000..37326130f669 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/softmaxMcSm103.h @@ -0,0 +1,71 @@ +/* + * 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 minimaxM31 +{ + +//! FlashInfer's Path A slice size (OnlineSoftmax DEFAULT_SLICE_SIZE): the Map kernel's grid is rows x +//! ceil(d / kSoftmaxMcSliceSize) CTAs, one partial (max, denominator) float2 each. +constexpr int kSoftmaxMcSliceSize = 8192; + +//! ST-01-SM: FlashInfer softmax(logits, temperature) Path A (OnlineSoftmaxMapKernel + OnlineSoftmaxReduceKernel, +//! batch <= 128 and vocab >= 24576) with the Reduce kernel's write pass spread over writeCtas CTAs per row instead of +//! one (FlashInfer: one 1024-thread CTA per row, ~19 us at 7-8 rows of a 200064 vocab on GB300). +//! +//! Semantics: logits fp32 [rows, d] (contiguous, d % 4 == 0, 16-byte aligned) -> out fp32 [rows, d] (same layout), +//! out[r, i] = softmax(logits[r,:] / T_r)[i] with T_r = temperatureArr[r] when temperatureArr is not null, else +//! temperatureVal (inverse temperature 0 when T_r == 0, as FlashInfer). +//! +//! Bitwise contract vs flashinfer.sampling.softmax(logits, temperature) (FlashInfer 0.6.15 == 0.6.18, Path A): +//! - the Map kernel is FlashInfer's OnlineSoftmaxMapKernel<1024, 4, float> (same grid, same cub BlockReduce calls, +//! same per-element source form), writing the same per-slice partials; +//! - every one of a row's writeCtas CTAs runs FlashInfer's Reduce prologue unchanged (the same cub +//! BlockReduce with Float2SoftmaxReduceOp over the same partials in the same order), so each gets +//! FlashInfer's final max and 1 / denominator bit for bit, and writes only its share of the row's 4096-element +//! tiles (tiles writeCtas apart) with the Reduce kernel's per-element arithmetic. The write is elementwise, so the +//! split changes no value. +//! Validated bitwise at writeCtas 4 / 8 / 16 on GB300 (microbenchmark): 9.2 vs FlashInfer 21.8 us per call +//! at 7-8 rows with writeCtas 16. +//! +//! MUST be compiled with --use_fast_math (FlashInfer's JIT flags: __expf, the approximate 1 / x, flush-to-zero); +//! cpp/tensorrt_llm/kernels/CMakeLists.txt sets the flag on this one source. +//! +//! partials: rows x ceil(d / kSoftmaxMcSliceSize) float2 (softmaxMcPartialsBytes), 8-byte aligned; fully overwritten +//! by every launch (no initialization, CUDA-graph safe). Launches that share it must be stream-ordered. Two plain +//! launches (no PDL), no co-residency requirement. The caller keeps the contract rows <= 128 and d >= 24576 (the +//! shapes where FlashInfer itself takes Path A); 1 <= writeCtas <= 65535. +cudaError_t invokeSoftmaxMc(float const* logits, float* out, float const* temperatureArr, float temperatureVal, + int rows, int d, int writeCtas, void* partials, cudaStream_t stream); + +//! Bytes of invokeSoftmaxMc's partials workspace for rows x d. +size_t softmaxMcPartialsBytes(int rows, int d); + +} // namespace minimaxM31 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM31/sparseDecodeLc8Tp8.cu b/cpp/tensorrt_llm/kernels/minimaxM31/sparseDecodeLc8Tp8.cu new file mode 100644 index 000000000000..567986ea2f1c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/sparseDecodeLc8Tp8.cu @@ -0,0 +1,621 @@ +/* + * 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.1 TP8 long-context sparse decode attention in one launch (8 query heads on one KV head). The M3.1 +// target's copy of kernels/minimaxM3Tp8/sparseDecodeLc8.cu with one change (SYNC3): every thread reads its two +// heads' row max and row sum from the four per-warp partials itself, in the same order as the single-thread +// rowm / rowl steps it replaces, so two of the softmax's five __syncthreads go and the rows are bit-identical; and one +// Reserved shared-memory layout for the kernel. +// +// 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/minimaxM31/sparseDecodeLc8Tp8.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 = MinimaxM31Tp8SparseDecodeLc8Params; + +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]; // Reserved shared-memory layout for the kernel. + float recv_l[kTopk]; // Reserved shared-memory layout for the kernel. + float part[kHalfDim]; + float red[4][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) m31_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(); +#pragma unroll + for (int e = 0; e < 2; ++e) + { + int const h = 2 * qd + e; + m_row[e] = fmaxf(fmaxf(s.red[0][h], s.red[1][h]), fmaxf(s.red[2][h], s.red[3][h])); + } + 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(); // row-max reads of red done before it 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(); +#pragma unroll + for (int e = 0; e < 2; ++e) + { + int const h = 2 * qd + e; + l_row[e] = (s.red[0][h] + s.red[1][h]) + (s.red[2][h] + s.red[3][h]); + } + } + + // Normalize; push (head, dim) to CTA 2 * head + dim / 64, slot recv[rank][dim % 64], plus each head's LSE (log2) + // Reserved shared-memory layout for the kernel. + // 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 o = 0.F; + { + 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; +#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 = m31_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 = m31_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 invokeMinimaxM31Tp8SparseDecodeLc8(MinimaxM31Tp8SparseDecodeLc8Params 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/minimaxM31/sparseDecodeLc8Tp8.h b/cpp/tensorrt_llm/kernels/minimaxM31/sparseDecodeLc8Tp8.h new file mode 100644 index 000000000000..84689bcf778a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/sparseDecodeLc8Tp8.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.1 TP8 one-launch long-context sparse decode (the M3.1 target's copy of +//! kernels/minimaxM3Tp8/sparseDecodeLc8; the M3 kernel is unchanged). 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 MinimaxM31Tp8SparseDecodeLc8Params +{ + 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. +//! +//! +//! SM100-family only; on other architectures the kernel compiles to an empty body, so callers check the device first. +cudaError_t invokeMinimaxM31Tp8SparseDecodeLc8(MinimaxM31Tp8SparseDecodeLc8Params 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/kernels/minimaxM31/sparseDecodeLcTp4.cu b/cpp/tensorrt_llm/kernels/minimaxM31/sparseDecodeLcTp4.cu new file mode 100644 index 000000000000..a9277a3fd2ea --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/sparseDecodeLcTp4.cu @@ -0,0 +1,635 @@ +/* + * 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.1 TP4 long-context sparse decode attention in one launch: the M3.1 target's copy of +// kernels/minimaxM3/sparseDecodeLc.cu with two exact options (BTSMEM: the block-table row staged in shared memory +// before the grid wait; SYNC3: the softmax's row max / sum read by every thread from the per-warp partials) and one +// Reserved shared-memory layout for the kernel. +// +// +// 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/minimaxM31/sparseDecodeLcTp4.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 = MinimaxM31Tp4SparseDecodeLcParams; +constexpr int kBtCap = 8192; // block-table entries staged in shared memory (BTSMEM; 8192 x 128 = 1M tokens) + +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]; // Reserved shared-memory layout for the kernel. + float recv_l[kTopk]; // Reserved shared-memory layout for the kernel. + float wgt[kTopk]; // Reserved shared-memory layout for the kernel. + float wl[kTopk]; // Reserved shared-memory layout for the kernel. + 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) +}; + +template +constexpr size_t smemBytes() +{ + return BTSMEM ? sizeof(Smem) + kBtCap * sizeof(int) : sizeof(Smem); +} + +#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) m31_tp4_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) + { + // Reserved shared-memory layout for the kernel. + // 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)); + } + } + + // BTSMEM: the request's block-table row (written before the step's graph runs) is staged before the wait too. + int* const s_bt = reinterpret_cast(smem_raw + sizeof(Smem)); + bool bt_smem = false; + if constexpr (BTSMEM) + { + 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); + } + } + + 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; + if constexpr (BTSMEM) + { + __syncthreads(); // s_bt complete + } + int64_t const page_pw = !active ? 0 : (bt_smem ? s_bt[blk] : 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 constexpr (SYNC3) + { +#pragma unroll + for (int r = 0; r < 2; ++r) + { + int const h = g + r * 8; + m_row[r] = fmaxf(fmaxf(s.red[0][h], s.red[1][h]), fmaxf(s.red[2][h], s.red[3][h])); + } + } + else + { + 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 constexpr (SYNC3) + { +#pragma unroll + for (int r = 0; r < 2; ++r) + { + int const h = g + r * 8; + l_row[r] = (s.red[0][h] + s.red[1][h]) + (s.red[2][h] + s.red[3][h]); + } + } + else + { + 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 o = 0.F; + { + 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; +#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 = m31_tp4_lc_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 = m31_tp4_lc_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 launchOptions(Params const& prm, int num_tokens, bool btsmem, bool sync3, bool use_pdl, cudaStream_t stream) +{ + if (btsmem) + { + return sync3 ? launch(prm, num_tokens, use_pdl, stream) + : launch(prm, num_tokens, use_pdl, stream); + } + return sync3 ? launch(prm, num_tokens, use_pdl, stream) + : launch(prm, num_tokens, use_pdl, stream); +} + +} // namespace + +cudaError_t invokeMinimaxM31Tp4SparseDecodeLc(MinimaxM31Tp4SparseDecodeLcParams const& params, int32_t numTokens, + bool linearSoftmax, bool mxfp8, bool earlyTrigger, bool btsmem, bool sync3, bool usePdl, cudaStream_t stream) +{ + if (!earlyTrigger) + { + return cudaErrorNotSupported; + } + if (linearSoftmax) + { + return mxfp8 ? launchOptions(params, numTokens, btsmem, sync3, usePdl, stream) + : launchOptions(params, numTokens, btsmem, sync3, usePdl, stream); + } + return mxfp8 ? launchOptions(params, numTokens, btsmem, sync3, usePdl, stream) + : launchOptions(params, numTokens, btsmem, sync3, usePdl, stream); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM31/sparseDecodeLcTp4.h b/cpp/tensorrt_llm/kernels/minimaxM31/sparseDecodeLcTp4.h new file mode 100644 index 000000000000..b228ff3d7cd0 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/sparseDecodeLcTp4.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 +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +//! Operands of the MiniMax-M3.1 TP4 one-launch long-context sparse decode (the M3.1 target's copy of +//! kernels/minimaxM3/sparseDecodeLc; the M3 kernel is unchanged). 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 MinimaxM31Tp4SparseDecodeLcParams +{ + 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). +//! +//! Two exact options over the M3 kernel, each bit-identical to it: +//! - btsmem: the request's block-table row (up to 8192 entries) is staged in dynamic shared memory before the grid +//! dependency wait, so the selected block's page id is read from shared memory after it (longer rows: global). +//! - sync3: every thread reads its heads' row max and row sum from the four per-warp partials itself, in the same +//! order as the single-thread rowm / rowl steps it replaces (two of the softmax's five __syncthreads go). +//! +//! Early trigger releases 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 invokeMinimaxM31Tp4SparseDecodeLc(MinimaxM31Tp4SparseDecodeLcParams const& params, int32_t numTokens, + bool linearSoftmax, bool mxfp8, bool earlyTrigger, bool btsmem, bool sync3, bool usePdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM31/toppRenormDetSm103.cu b/cpp/tensorrt_llm/kernels/minimaxM31/toppRenormDetSm103.cu new file mode 100644 index 000000000000..19b548860e6f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/toppRenormDetSm103.cu @@ -0,0 +1,535 @@ +/* + * 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. + */ + +// ST-01 v2: FlashInfer AirTopPRenormProb in one +// persistent launch, with the per-row keep-all bypass for top_p >= 1. See toppRenormDetSm103.h for the contract. +// Compiled with --use_fast_math (set on this source in kernels/CMakeLists.txt), as FlashInfer's JIT compiles it. + +#include "tensorrt_llm/kernels/minimaxM31/toppRenormDetSm103.h" + +#include +#include +#include +#include + +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM31 +{ + +namespace +{ + +// --------------------------------------------------------------------------------------------------------------------- +// Vendored from FlashInfer, include/flashinfer/air_top_p.cuh (FlashInfer v0.6.15; identical in v0.6.18), namespace +// flashinfer::sampling::air_top_p. Copyright (c) 2026 by FlashInfer team, licensed under the Apache License, +// Version 2.0. Copied verbatim (only the pieces this kernel uses; formatting kept) so that the bucket keys, masks +// and the deterministic histogram -> float conversion are FlashInfer's own; FlashInfer's headers are not on the +// TensorRT-LLM C++ include path. +// --------------------------------------------------------------------------------------------------------------------- +namespace fi_air_top_p +{ +// clang-format off +static constexpr int BITS_PER_PASS = 11; +static constexpr int NUM_BUCKETS = 1 << BITS_PER_PASS; + +template +static constexpr int NUM_PASSES = (sizeof(T) * 8 + BITS_PER_PASS - 1) / BITS_PER_PASS; + +template +__device__ int constexpr calcStartBit(int pass) { + int startBit = static_cast(sizeof(T) * 8) - (pass + 1) * BITS_PER_PASS; + return startBit < 0 ? 0 : startBit; +} + +template +__device__ uint32_t constexpr calcMask(int pass) { + int numBits = calcStartBit(pass - 1) - calcStartBit(pass); + return (1 << numBits) - 1; +} + +template +__device__ typename cub::Traits::UnsignedBits twiddleIn(T key, bool selectMin) { + auto bits = reinterpret_cast::UnsignedBits&>(key); + bits = cub::Traits::TwiddleIn(bits); + if (!selectMin) bits = ~bits; + return bits; +} + +template +__device__ T twiddleOut(typename cub::Traits::UnsignedBits bits, bool selectMin) { + if (!selectMin) bits = ~bits; + bits = cub::Traits::TwiddleOut(bits); + return reinterpret_cast(bits); +} + +__device__ inline uint32_t calcMantissa(float value) { + union { + uint32_t bits; + float value; + } input; + input.value = value; + constexpr uint32_t numMantissa = 23; + return input.bits & ((1u << numMantissa) - 1); +} + +__device__ inline uint32_t calcExponent(float value) { + union { + uint32_t bits; + float value; + } input; + input.value = value; + constexpr uint32_t numMantissa = 23; + return input.bits & ~((1u << numMantissa) - 1); +} + +__device__ inline float calcFloatValue(uint32_t count, uint32_t exponent, uint64_t bitSum) { + constexpr uint32_t numTotalBits = 64; + constexpr uint32_t numMantissa = 23; + uint64_t extraInMantissa = (bitSum >> numMantissa); + extraInMantissa = (exponent == 0) ? extraInMantissa : extraInMantissa + count; + uint32_t numExtra = numTotalBits - __clzll(extraInMantissa); + int numNorm = (exponent == 0) ? 0 : -1; + exponent = exponent + ((numExtra + numNorm) << numMantissa); + uint32_t mantissa; + if (extraInMantissa != 0) { + int numMove = numMantissa - (numExtra - 1); + uint32_t mask = (1u << (numExtra - 1)) - 1; + extraInMantissa = extraInMantissa & mask; + if (numMove > 0) { + extraInMantissa = extraInMantissa << numMove; + mask = (1u << numMantissa) - 1; + mantissa = ((bitSum & mask) >> (numExtra - 1)) | extraInMantissa; + } else { + mantissa = extraInMantissa >> (-1 * numMove); + } + } else { + mantissa = bitSum; + } + uint32_t bitFloat = exponent | mantissa; + return reinterpret_cast(bitFloat); +} + +// clang-format on +} // namespace fi_air_top_p + +// --------------------------------------------------------------------------------------------------------------------- +// End of the vendored FlashInfer code. +// --------------------------------------------------------------------------------------------------------------------- + +namespace fa = fi_air_top_p; +namespace cg = cooperative_groups; + +constexpr int kThreads = 1024; +constexpr int kBuckets = fa::NUM_BUCKETS; // 2048 +constexpr int kSmemBytes = 32768; +constexpr int kCopies = 4; // privatized pass histograms +constexpr size_t kDynSmemBytes = kCopies * kBuckets * (sizeof(unsigned long long) + sizeof(int)); // 96 KB + +// One row's persistent state. Zero once; every launch leaves it zeroed except `gen` (monotonic) and the per-launch +// scalars it rewrites before reading. +struct RowWorkspace +{ + unsigned long long hist[kBuckets]; // the pass's det mantissa-sum histogram (FlashInfer's uint64 HisT) + int cnt[kBuckets]; // the pass's count histogram + float chain[kThreads]; // the 1024 per-thread chain sums of the kept values + unsigned arrive; // group barrier: CTAs arrived + unsigned gen; // group barrier: generation + unsigned kbits; // the selected key prefix (FlashInfer's kthValueBits) + float sum; // the remaining target sum (FlashInfer's counter->sum) + float norm; + int kept; + int pad[2]; +}; + +__device__ __forceinline__ unsigned ldAcquire(unsigned* p) +{ + return cuda::atomic_ref(*p).load(cuda::memory_order_acquire); +} + +__device__ __forceinline__ void stRelease(unsigned* p, unsigned v) +{ + cuda::atomic_ref(*p).store(v, cuda::memory_order_release); +} + +// Group barrier of the row's G CTAs; the last CTA to arrive runs `leader` (whole CTA) before releasing the others. +template +__device__ __forceinline__ void groupSync(RowWorkspace* w, int G, F leader) +{ + __shared__ unsigned sGen; + __shared__ int sLast; + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + unsigned gen = ldAcquire(&w->gen); + sGen = gen; + unsigned old = atomicAdd(&w->arrive, 1u); + sLast = (old == static_cast(G - 1)); + } + __syncthreads(); + if (sLast) + { + __threadfence(); + leader(); + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + atomicExch(&w->arrive, 0u); + stRelease(&w->gen, sGen + 1); + } + } + else if (threadIdx.x == 0) + { + while (ldAcquire(&w->gen) == sGen) + { + __nanosleep(20); + } + } + __syncthreads(); +} + +// FlashInfer's last-block selection for `pass` (air_top_p.cuh AirTopPRenormRadixKernel, IsDeterministic branch), on +// the row's global integer histograms. histVal and cntS (the bucket counts, cached for the sequential scan) reuse the +// CTA's scratch. +__device__ __forceinline__ void selectPass(RowWorkspace* w, int pass, float p, float* histVal, int* cntS) +{ + constexpr int WARP_SIZE = 32; + constexpr int WARP_COUNT = kBuckets / WARP_SIZE; + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); + __shared__ float warpSum[WARP_COUNT]; + __shared__ cuda::atomic blockSum; + for (int i = threadIdx.x; i < kBuckets; i += blockDim.x) + { + unsigned long long value = __ldcg(&w->hist[i]); + int count = __ldcg(&w->cnt[i]); + cntS[i] = count; + if (count != 0) + { + uint32_t sb = fa::calcStartBit(pass); + uint32_t bv = __ldcg(&w->kbits); + if (pass == 0) + { + bv = i << sb; + } + histVal[i] = fa::calcFloatValue( + static_cast(count), fa::calcExponent(fa::twiddleOut(bv, false)), value); + } + else + { + histVal[i] = 0.0f; + } + } + __syncthreads(); + for (int i = threadIdx.x; i < kBuckets; i += blockDim.x) + { + cg::reduce_store_async(warp, warpSum + i / WARP_SIZE, histVal[i], cg::plus{}); + } + __syncthreads(); + if (threadIdx.x < WARP_SIZE) + { + cg::reduce_store_async(warp, blockSum, warpSum[threadIdx.x], cg::plus{}); + cg::reduce_update_async(warp, blockSum, warpSum[threadIdx.x + WARP_SIZE], cg::plus{}); + } + __syncthreads(); + if (threadIdx.x == 0) + { + float currentSum = pass == 0 ? static_cast(blockSum) * p : __ldcg(&w->sum); + float prev = 0; + int targetStep = 0; + for (int i = 0; i < WARP_COUNT; i++) + { + if (warpSum[i]) + { + targetStep = i; + if ((prev + warpSum[i]) >= currentSum) + { + break; + } + prev += warpSum[i]; + } + } + int targetIdx = 0; + for (int i = targetStep * WARP_SIZE; i < kBuckets; i++) + { + if (cntS[i]) + { + targetIdx = i; + if ((prev + histVal[i]) >= currentSum) + { + break; + } + prev += histVal[i]; + } + } + w->sum = currentSum - prev; + unsigned kb = pass == 0 ? 0u : __ldcg(&w->kbits); + w->kbits = kb | (static_cast(targetIdx) << fa::calcStartBit(pass)); + } + __syncthreads(); + for (int i = threadIdx.x; i < kBuckets; i += blockDim.x) + { + w->hist[i] = 0; + w->cnt[i] = 0; + } +} + +__global__ void __launch_bounds__(kThreads, 1) toppRenormDetKernel( + float const* __restrict__ probs, float* __restrict__ out, float const* topP, int d, RowWorkspace* ws, int* keptOut) +{ + __shared__ __align__(16) unsigned char smem[kSmemBytes]; + // Privatized pass histograms: kCopies copies, one per group of warps (dynamic shared memory). + extern __shared__ __align__(16) unsigned char dyn[]; + unsigned long long* hsCopies = reinterpret_cast(dyn); + int* csCopies = reinterpret_cast(dyn + kCopies * kBuckets * sizeof(unsigned long long)); + int const row = blockIdx.y; + int const G = gridDim.x; + int const cta = blockIdx.x; + int const tid = threadIdx.x; + int const lane = tid & 31; + RowWorkspace* w = ws + row; + float const* in = probs + static_cast(row) * d; + float* o = out + static_cast(row) * d; + float const p = topP[row]; + float4 const* in4 = reinterpret_cast(in); + int const n4 = d / 4; + int const stride4 = G * kThreads; + int const start4 = cta * kThreads + tid; + + unsigned long long* hs = reinterpret_cast(smem); // 16 KB + int* cs = reinterpret_cast(smem + kBuckets * sizeof(unsigned long long)); // 8 KB + unsigned kbits = 0; + + // top_p >= 1 (the keep-all exception): no radix passes, the threshold is -inf. + bool const keepAll = p >= 1.0f; + for (int pass = 0; pass < (keepAll ? 0 : fa::NUM_PASSES); ++pass) + { + for (int i = tid; i < kCopies * kBuckets; i += kThreads) + { + hsCopies[i] = 0; + csCopies[i] = 0; + } + __syncthreads(); + int const copy = (tid >> 5) % kCopies; + unsigned long long* myHs = hsCopies + copy * kBuckets; + int* myCs = csCopies + copy * kBuckets; + int const sb = fa::calcStartBit(pass); + uint32_t const mask = fa::calcMask(pass); + int const psb = pass == 0 ? 0 : fa::calcStartBit(pass - 1); + // Warp-uniform trip count: every lane runs the same iterations (masked), so the warp collectives are + // full-warp. + for (int j0 = start4 - lane; j0 < n4; j0 += stride4) + { + int const j = j0 + lane; + bool const valid = j < n4; + float4 v = valid ? in4[j] : make_float4(0.f, 0.f, 0.f, 0.f); + float const e[4] = {v.x, v.y, v.z, v.w}; +#pragma unroll + for (int k = 0; k < 4; ++k) + { + float const x = e[k]; + auto const bits = fa::twiddleIn(x, false); + bool const ok = valid && (pass == 0 || ((bits >> psb) << psb) == kbits); + int const b = ok ? static_cast((bits >> sb) & mask) : -1; + unsigned const m = ok ? fa::calcMantissa(x) : 0u; + // Adaptive per warp (integer sums, so order-free and exact): a warp whose 32 lanes share one bucket + // (near-uniform rows) adds it with ONE aggregated atomic; otherwise every eligible lane adds into its + // warp group's private copy (few lanes collide on rows spread over many buckets). + int const b0 = __shfl_sync(0xffffffffu, b, 0); + if (__all_sync(0xffffffffu, b == b0)) + { + if (b0 >= 0) + { + unsigned const s = __reduce_add_sync(0xffffffffu, m); + if (lane == 0) + { + atomicAdd(myHs + b0, static_cast(s)); + atomicAdd(myCs + b0, 32); + } + } + } + else if (ok) + { + atomicAdd(myHs + b, static_cast(m)); + atomicAdd(myCs + b, 1); + } + } + } + __syncthreads(); + for (int i = tid; i < kBuckets; i += kThreads) + { + unsigned long long h = 0; + int c = 0; +#pragma unroll + for (int k = 0; k < kCopies; ++k) + { + h += hsCopies[k * kBuckets + i]; + c += csCopies[k * kBuckets + i]; + } + hs[i] = h; + cs[i] = c; + } + __syncthreads(); + for (int i = tid; i < kBuckets; i += kThreads) + { + if (hs[i] != 0) + { + atomicAdd(&w->hist[i], hs[i]); + } + if (cs[i] != 0) + { + atomicAdd(&w->cnt[i], cs[i]); + } + } + groupSync(w, G, + [&] { + selectPass(w, pass, p, reinterpret_cast(smem), + reinterpret_cast(smem + kBuckets * sizeof(float))); + }); + kbits = __ldcg(&w->kbits); + } + + // Apply: the threshold, FlashInfer's per-thread chain sums (spread over the G CTAs), the cub block sum, then the + // renorm write. + float const thr = keepAll ? -INFINITY : fa::twiddleOut(kbits, false); + int const cpc = kThreads / G; // chains per CTA + int const c0 = cta * cpc; + int const kTile = (kSmemBytes / 4) / cpc; + int const kTotal = (d + kThreads - 1) / kThreads; + float* stage = reinterpret_cast(smem); + float acc = 0.0f; + int kept = 0; + for (int k0 = 0; k0 < kTotal; k0 += kTile) + { + int const kn = min(kTile, kTotal - k0); + for (int q = tid; q < kn * cpc; q += kThreads) + { + int const kk = q / cpc, c = q - kk * cpc; + int const i = c0 + c + kThreads * (k0 + kk); + stage[q] = i < d ? in[i] : 0.0f; + } + __syncthreads(); + if (tid < cpc) + { + for (int kk = 0; kk < kn; ++kk) + { + int const i = c0 + tid + kThreads * (k0 + kk); + float const val = stage[kk * cpc + tid]; + if (i < d && val >= thr) + { + acc += val; + ++kept; + } + } + } + __syncthreads(); + } + if (tid < cpc) + { + w->chain[c0 + tid] = acc; + } + typedef cub::BlockReduce BlockReduce; + __shared__ typename BlockReduce::TempStorage temp; + __shared__ int ctaKept; + if (tid == 0) + { + ctaKept = 0; + } + __syncthreads(); + if (tid < cpc && kept) + { + atomicAdd(&ctaKept, kept); + } + __syncthreads(); + if (tid == 0 && ctaKept) + { + atomicAdd(&w->kept, ctaKept); + } + groupSync(w, G, + [&] + { + float totalSum = BlockReduce(temp).Sum(__ldcg(&w->chain[tid])); + if (tid == 0) + { + w->norm = (totalSum > 1e-8f) ? (1.0f / totalSum) : 1.0f; + if (keptOut) + { + keptOut[row] = __ldcg(&w->kept); + } + w->kept = 0; + } + }); + float const norm = __ldcg(&w->norm); + float4* o4 = reinterpret_cast(o); + for (int j = start4; j < n4; j += stride4) + { + float4 v = in4[j]; + float4 r; + r.x = (v.x >= thr) ? v.x * norm : 0.0f; + r.y = (v.y >= thr) ? v.y * norm : 0.0f; + r.z = (v.z >= thr) ? v.z * norm : 0.0f; + r.w = (v.w >= thr) ? v.w * norm : 0.0f; + o4[j] = r; + } +} + +} // namespace + +cudaError_t invokeToppRenormDet(float const* probs, float* out, float const* topP, int rows, int d, int ctasPerRow, + void* workspace, int* keptOut, cudaStream_t stream) +{ + if (rows <= 0 || d <= 0) + { + return cudaSuccess; + } + if (ctasPerRow < 1 || ctasPerRow > kThreads || kThreads % ctasPerRow != 0 || d % 4 != 0) + { + return cudaErrorInvalidValue; + } + static cudaError_t const attrErr = cudaFuncSetAttribute( + toppRenormDetKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(kDynSmemBytes)); + if (attrErr != cudaSuccess) + { + return attrErr; + } + cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(static_cast(ctasPerRow), static_cast(rows)); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = kDynSmemBytes; + cfg.stream = stream; + return cudaLaunchKernelEx( + &cfg, toppRenormDetKernel, probs, out, topP, d, static_cast(workspace), keptOut); +} + +size_t toppRenormDetWorkspaceBytesPerRow() +{ + return sizeof(RowWorkspace); +} + +} // namespace minimaxM31 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM31/toppRenormDetSm103.h b/cpp/tensorrt_llm/kernels/minimaxM31/toppRenormDetSm103.h new file mode 100644 index 000000000000..0e5e1eb7912e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/toppRenormDetSm103.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 + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM31 +{ + +//! ST-01: the spec-decoding sampler's top-p renorm (FlashInfer top_p_renorm_probs) as ONE persistent launch of +//! rows x ctasPerRow CTAs (1024 threads each), flat in the kept count. +//! +//! Semantics: probs fp32 [rows, d] (contiguous, d % 4 == 0, 16-byte aligned), topP fp32 [rows]. Per row, the +//! threshold thr is the value FlashInfer's AirTopPRenormProb finds with its three +//! radix passes (11 / 11 / 10 bits of the descending twiddled key); out[i] = probs[i] * norm where probs[i] >= thr, +//! else 0, with norm = 1 / sum (1 when sum <= 1e-8) and sum the kept values' total in FlashInfer's order. +//! +//! Bitwise contract vs FlashInfer top_p_renorm_probs(probs, top_p, is_deterministic=True) (air_top_p.cuh, FlashInfer +//! 0.6.15 == 0.6.18), for every row with top_p < 1: +//! - the det histograms are INTEGER (uint64 mantissa sums + int counts), so the warp-aggregated updates here give +//! FlashInfer's exact histograms in any order; +//! - each pass's bucket selection is FlashInfer's last-block code (calcFloatValue, cooperative_groups warp / block +//! sums, the sequential scan) run on integer-identical histograms with the same lane <-> bucket mapping; +//! - the kept sum: thread chain t adds the kept values at i = t (mod 1024) in order (FlashInfer's per-thread input +//! to cub BlockReduce); the chains are spread over the row's CTAs and one CTA runs the same +//! cub BlockReduce::Sum on them. +//! Validated bitwise (renormed probs and the sampled token) on GB300 over R 7 / 8 / 56 / 64 x 8 row classes x +//! 4 top_p modes x 16 seeds, eager and graph (microbenchmark); v1.2's binning (each pass's histogram adds: +//! one aggregated atomic for a warp whose 32 lanes share a bucket, else per-lane atomics into one of 4 privatized +//! shared-memory copies; v1.1's cached bucket counts) is bitwise v1 on every row of the MB grid at 54.3 vs v1's 85.7 us +//! per 7-row call. +//! +//! Exception: a row with top_p >= 1 skips the radix search and keeps EVERY element +//! (threshold -inf), then the same sum and write. FlashInfer itself drops tail buckets at top_p 1 in ~43% of rows, +//! in BOTH its default and its deterministic mode; keep-all is the intended law. The choice is row-local (all CTAs of +//! a row take the same branch), so the other rows of a batch stay bitwise. +//! +//! MUST be compiled with --use_fast_math (FlashInfer's JIT flags: flush-to-zero and the approximate 1 / x of the norm); +//! without it the output is not bitwise FlashInfer's. cpp/tensorrt_llm/kernels/CMakeLists.txt sets the flag on this +//! one source. +//! +//! Workspace: toppRenormDetWorkspaceBytesPerRow() x rows bytes, ZERO-initialized once before the first launch and +//! then reused as is: every launch leaves it ready for the next (histograms cleared, arrival counters reset, the +//! barrier generation only increases), so it is CUDA-graph safe without an init kernel. Launches that share one +//! workspace must be stream-ordered. The row's CTAs synchronize through global-memory barriers, so all +//! rows x ctasPerRow CTAs must be co-resident (one 1024-thread CTA per SM): pick ctasPerRow with +//! rows x ctasPerRow <= the SM count (ctasPerRow 1 needs no co-residency). 1024 % ctasPerRow == 0. +//! keptOut (optional, int [rows]) receives the per-row count of elements >= the threshold. +cudaError_t invokeToppRenormDet(float const* probs, float* out, float const* topP, int rows, int d, int ctasPerRow, + void* workspace, int* keptOut, cudaStream_t stream); + +//! Workspace bytes per row of invokeToppRenormDet. +size_t toppRenormDetWorkspaceBytesPerRow(); + +} // namespace minimaxM31 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM31/toppRenormDetV3Sm103.cu b/cpp/tensorrt_llm/kernels/minimaxM31/toppRenormDetV3Sm103.cu new file mode 100644 index 000000000000..3ba14d32938d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/toppRenormDetV3Sm103.cu @@ -0,0 +1,650 @@ +/* + * 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. + */ + +// TOPP_DET v3: toppRenormDetSm103.cu (v1.2) with the redundant selection and an optional programmatic dependent +// launch. See toppRenormDetV3Sm103.h for the contract (bitwise v1.2's output). Compiled with --use_fast_math (set on +// this source in kernels/CMakeLists.txt), as FlashInfer's JIT compiles it. + +#include "tensorrt_llm/kernels/minimaxM31/toppRenormDetV3Sm103.h" + +#include +#include +#include +#include + +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM31 +{ + +namespace toppv3 +{ + +// --------------------------------------------------------------------------------------------------------------------- +// Vendored from FlashInfer, include/flashinfer/air_top_p.cuh (FlashInfer v0.6.15; identical in v0.6.18), namespace +// flashinfer::sampling::air_top_p. Copyright (c) 2026 by FlashInfer team, licensed under the Apache License, +// Version 2.0. Copied verbatim (only the pieces this kernel uses; formatting kept) so that the bucket keys, masks +// and the deterministic histogram -> float conversion are FlashInfer's own; FlashInfer's headers are not on the +// TensorRT-LLM C++ include path. +// --------------------------------------------------------------------------------------------------------------------- +namespace fi_air_top_p +{ +// clang-format off +static constexpr int BITS_PER_PASS = 11; +static constexpr int NUM_BUCKETS = 1 << BITS_PER_PASS; + +template +static constexpr int NUM_PASSES = (sizeof(T) * 8 + BITS_PER_PASS - 1) / BITS_PER_PASS; + +template +__device__ int constexpr calcStartBit(int pass) { + int startBit = static_cast(sizeof(T) * 8) - (pass + 1) * BITS_PER_PASS; + return startBit < 0 ? 0 : startBit; +} + +template +__device__ uint32_t constexpr calcMask(int pass) { + int numBits = calcStartBit(pass - 1) - calcStartBit(pass); + return (1 << numBits) - 1; +} + +template +__device__ typename cub::Traits::UnsignedBits twiddleIn(T key, bool selectMin) { + auto bits = reinterpret_cast::UnsignedBits&>(key); + bits = cub::Traits::TwiddleIn(bits); + if (!selectMin) bits = ~bits; + return bits; +} + +template +__device__ T twiddleOut(typename cub::Traits::UnsignedBits bits, bool selectMin) { + if (!selectMin) bits = ~bits; + bits = cub::Traits::TwiddleOut(bits); + return reinterpret_cast(bits); +} + +__device__ inline uint32_t calcMantissa(float value) { + union { + uint32_t bits; + float value; + } input; + input.value = value; + constexpr uint32_t numMantissa = 23; + return input.bits & ((1u << numMantissa) - 1); +} + +__device__ inline uint32_t calcExponent(float value) { + union { + uint32_t bits; + float value; + } input; + input.value = value; + constexpr uint32_t numMantissa = 23; + return input.bits & ~((1u << numMantissa) - 1); +} + +__device__ inline float calcFloatValue(uint32_t count, uint32_t exponent, uint64_t bitSum) { + constexpr uint32_t numTotalBits = 64; + constexpr uint32_t numMantissa = 23; + uint64_t extraInMantissa = (bitSum >> numMantissa); + extraInMantissa = (exponent == 0) ? extraInMantissa : extraInMantissa + count; + uint32_t numExtra = numTotalBits - __clzll(extraInMantissa); + int numNorm = (exponent == 0) ? 0 : -1; + exponent = exponent + ((numExtra + numNorm) << numMantissa); + uint32_t mantissa; + if (extraInMantissa != 0) { + int numMove = numMantissa - (numExtra - 1); + uint32_t mask = (1u << (numExtra - 1)) - 1; + extraInMantissa = extraInMantissa & mask; + if (numMove > 0) { + extraInMantissa = extraInMantissa << numMove; + mask = (1u << numMantissa) - 1; + mantissa = ((bitSum & mask) >> (numExtra - 1)) | extraInMantissa; + } else { + mantissa = extraInMantissa >> (-1 * numMove); + } + } else { + mantissa = bitSum; + } + uint32_t bitFloat = exponent | mantissa; + return reinterpret_cast(bitFloat); +} + +// clang-format on +} // namespace fi_air_top_p + +// --------------------------------------------------------------------------------------------------------------------- +// End of the vendored FlashInfer code. +// --------------------------------------------------------------------------------------------------------------------- + +namespace fa = fi_air_top_p; +namespace cg = cooperative_groups; + +constexpr int kThreads = 1024; +constexpr int kBuckets = fa::NUM_BUCKETS; // 2048 +constexpr int kSmemBytes = 32768; +constexpr int kCopies = 4; // privatized pass histograms +constexpr size_t kDynSmemBytes = kCopies * kBuckets * (sizeof(unsigned long long) + sizeof(int)); // 96 KB + +constexpr int kSyncs = fa::NUM_PASSES + 1; // one per radix pass, then the norm +// Candidate list: passes 1 .. NUM_PASSES - 1 run locally on the members of pass 0's threshold bucket when it holds at +// most this many elements (the list lives in the workspace, then in the dynamic shared memory of each CTA). +constexpr int kCandMax = 4096; + +// One row's persistent state. Zero once; every launch leaves it zeroed except the arrival counters, which only grow. +struct RowWorkspace +{ + unsigned long long hist[fa::NUM_PASSES][kBuckets]; // per pass: det mantissa-sum histogram (uint64 HisT) + int cnt[fa::NUM_PASSES][kBuckets]; // per pass: count histogram + float chain[kThreads]; // the 1024 per-thread chain sums of the kept values + unsigned long long arrive[kSyncs]; // per sync: CTAs arrived, monotone across launches + float cand[kCandMax]; // the candidate list (pass 0's bucket members) + unsigned nCand; // its length, zero between launches + int kept; + int pad[2]; +}; + +// Arrive at sync `si` and wait until all G CTAs of the row have arrived. The counter is never reset: each CTA adds +// kThreads / G (G divides kThreads), so every launch adds exactly kThreads per sync it runs whatever its G (a row's +// CTAs all take the same branch) and the counter is a multiple of kThreads between launches; the target is the next +// one. +__device__ __forceinline__ void allArrive(RowWorkspace* w, int G, int si) +{ + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + unsigned long long const step = static_cast(kThreads / G); + unsigned long long const old = atomicAdd(&w->arrive[si], step); + unsigned long long const target = (old / kThreads + 1) * kThreads; + cuda::atomic_ref ref(w->arrive[si]); + while (ref.load(cuda::memory_order_acquire) < target) + { + __nanosleep(20); + } + } + __syncthreads(); +} + +// FlashInfer's last-block selection for `pass` (air_top_p.cuh AirTopPRenormRadixKernel, IsDeterministic branch), on +// the row's global integer histograms of this pass, run by EVERY CTA of the row (identical inputs, identical code, so +// identical bits). The running sum and key prefix are CTA-local (shared memory) instead of the workspace. histVal and +// cntS (the bucket counts, cached for the sequential scan) reuse the CTA's scratch. +// +// kGlobal: the histograms are the row's global ones (read through L2); otherwise this CTA's shared-memory histograms of +// the candidate list, which hold the same integers. +template +__device__ __forceinline__ void selectPassOn(unsigned long long const* hist, int const* cnt, int pass, float p, + float* histVal, int* cntS, float* sSum, unsigned* sKbits) +{ + constexpr int WARP_SIZE = 32; + constexpr int WARP_COUNT = kBuckets / WARP_SIZE; + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); + __shared__ float warpSum[WARP_COUNT]; + __shared__ cuda::atomic blockSum; + for (int i = threadIdx.x; i < kBuckets; i += blockDim.x) + { + unsigned long long value = kGlobal ? __ldcg(&hist[i]) : hist[i]; + int count = kGlobal ? __ldcg(&cnt[i]) : cnt[i]; + cntS[i] = count; + if (count != 0) + { + uint32_t sb = fa::calcStartBit(pass); + uint32_t bv = *sKbits; + if (pass == 0) + { + bv = i << sb; + } + histVal[i] = fa::calcFloatValue( + static_cast(count), fa::calcExponent(fa::twiddleOut(bv, false)), value); + } + else + { + histVal[i] = 0.0f; + } + } + __syncthreads(); + for (int i = threadIdx.x; i < kBuckets; i += blockDim.x) + { + cg::reduce_store_async(warp, warpSum + i / WARP_SIZE, histVal[i], cg::plus{}); + } + __syncthreads(); + if (threadIdx.x < WARP_SIZE) + { + cg::reduce_store_async(warp, blockSum, warpSum[threadIdx.x], cg::plus{}); + cg::reduce_update_async(warp, blockSum, warpSum[threadIdx.x + WARP_SIZE], cg::plus{}); + } + __syncthreads(); + if (threadIdx.x == 0) + { + float currentSum = pass == 0 ? static_cast(blockSum) * p : *sSum; + float prev = 0; + int targetStep = 0; + for (int i = 0; i < WARP_COUNT; i++) + { + if (warpSum[i]) + { + targetStep = i; + if ((prev + warpSum[i]) >= currentSum) + { + break; + } + prev += warpSum[i]; + } + } + int targetIdx = 0; + for (int i = targetStep * WARP_SIZE; i < kBuckets; i++) + { + if (cntS[i]) + { + targetIdx = i; + if ((prev + histVal[i]) >= currentSum) + { + break; + } + prev += histVal[i]; + } + } + *sSum = currentSum - prev; + unsigned kb = pass == 0 ? 0u : *sKbits; + *sKbits = kb | (static_cast(targetIdx) << fa::calcStartBit(pass)); + } + __syncthreads(); +} + +__device__ __forceinline__ void selectPass( + RowWorkspace* w, int pass, float p, float* histVal, int* cntS, float* sSum, unsigned* sKbits) +{ + selectPassOn(w->hist[pass], w->cnt[pass], pass, p, histVal, cntS, sSum, sKbits); +} + +// CTA 0 clears pass `pass`'s global histograms once every CTA of the row has finished reading them. +__device__ __forceinline__ void clearPass(RowWorkspace* w, int pass) +{ + if (blockIdx.x == 0) + { + for (int i = threadIdx.x; i < kBuckets; i += blockDim.x) + { + w->hist[pass][i] = 0; + w->cnt[pass][i] = 0; + } + } +} + +template +__global__ void __launch_bounds__(kThreads, 1) toppRenormDetV3Kernel( + float const* __restrict__ probs, float* __restrict__ out, float const* topP, int d, RowWorkspace* ws, int* keptOut) +{ + __shared__ __align__(16) unsigned char smem[kSmemBytes]; + // Privatized pass histograms: kCopies copies, one per group of warps (dynamic shared memory). + extern __shared__ __align__(16) unsigned char dyn[]; + unsigned long long* hsCopies = reinterpret_cast(dyn); + int* csCopies = reinterpret_cast(dyn + kCopies * kBuckets * sizeof(unsigned long long)); + int const row = blockIdx.y; + int const G = gridDim.x; + int const cta = blockIdx.x; + int const tid = threadIdx.x; + int const lane = tid & 31; + RowWorkspace* w = ws + row; + float const* in = probs + static_cast(row) * d; + float* o = out + static_cast(row) * d; + __shared__ float sSum; + __shared__ unsigned sKbits; + if (tid == 0) + { + sSum = 0.0f; + sKbits = 0u; + } + // Everything above is CTA-private; probs, top_p and the shared workspace are only touched after the wait. + if constexpr (kPdl) + { + cudaGridDependencySynchronize(); + } + __syncthreads(); + float const p = topP[row]; + float4 const* in4 = reinterpret_cast(in); + int const n4 = d / 4; + int const stride4 = G * kThreads; + int const start4 = cta * kThreads + tid; + + unsigned long long* hs = reinterpret_cast(smem); // 16 KB + int* cs = reinterpret_cast(smem + kBuckets * sizeof(unsigned long long)); // 8 KB + unsigned kbits = 0; + + // top_p >= 1 (the keep-all exception): no radix passes, the threshold is -inf. + bool const keepAll = p >= 1.0f; + for (int pass = 0; pass < (keepAll ? 0 : fa::NUM_PASSES); ++pass) + { + for (int i = tid; i < kCopies * kBuckets; i += kThreads) + { + hsCopies[i] = 0; + csCopies[i] = 0; + } + __syncthreads(); + int const copy = (tid >> 5) % kCopies; + unsigned long long* myHs = hsCopies + copy * kBuckets; + int* myCs = csCopies + copy * kBuckets; + int const sb = fa::calcStartBit(pass); + uint32_t const mask = fa::calcMask(pass); + int const psb = pass == 0 ? 0 : fa::calcStartBit(pass - 1); + // Warp-uniform trip count: every lane runs the same iterations (masked), so the warp collectives are + // full-warp. + for (int j0 = start4 - lane; j0 < n4; j0 += stride4) + { + int const j = j0 + lane; + bool const valid = j < n4; + float4 v = valid ? in4[j] : make_float4(0.f, 0.f, 0.f, 0.f); + float const e[4] = {v.x, v.y, v.z, v.w}; +#pragma unroll + for (int k = 0; k < 4; ++k) + { + float const x = e[k]; + auto const bits = fa::twiddleIn(x, false); + bool const ok = valid && (pass == 0 || ((bits >> psb) << psb) == kbits); + int const b = ok ? static_cast((bits >> sb) & mask) : -1; + unsigned const m = ok ? fa::calcMantissa(x) : 0u; + // Adaptive per warp (integer sums, so order-free and exact): a warp whose 32 lanes share one bucket + // (near-uniform rows) adds it with ONE aggregated atomic; otherwise every eligible lane adds into its + // warp group's private copy (few lanes collide on rows spread over many buckets). + int const b0 = __shfl_sync(0xffffffffu, b, 0); + if (__all_sync(0xffffffffu, b == b0)) + { + if (b0 >= 0) + { + unsigned const s = __reduce_add_sync(0xffffffffu, m); + if (lane == 0) + { + atomicAdd(myHs + b0, static_cast(s)); + atomicAdd(myCs + b0, 32); + } + } + } + else if (ok) + { + atomicAdd(myHs + b, static_cast(m)); + atomicAdd(myCs + b, 1); + } + } + } + __syncthreads(); + for (int i = tid; i < kBuckets; i += kThreads) + { + unsigned long long h = 0; + int c = 0; +#pragma unroll + for (int k = 0; k < kCopies; ++k) + { + h += hsCopies[k * kBuckets + i]; + c += csCopies[k * kBuckets + i]; + } + hs[i] = h; + cs[i] = c; + } + __syncthreads(); + for (int i = tid; i < kBuckets; i += kThreads) + { + if (hs[i] != 0) + { + atomicAdd(&w->hist[pass][i], hs[i]); + } + if (cs[i] != 0) + { + atomicAdd(&w->cnt[pass][i], cs[i]); + } + } + allArrive(w, G, pass); + if (pass > 0) + { + clearPass(w, pass - 1); // every CTA finished selectPass(pass - 1) before arriving here + } + selectPass(w, pass, p, reinterpret_cast(smem), reinterpret_cast(smem + kBuckets * sizeof(float)), + &sSum, &sKbits); + kbits = sKbits; + if constexpr (kCandList) + { + // Pass 0's threshold bucket count (cached by selectPass) is the same in every CTA of the row, so all of + // them take the same branch. A small bucket: list its members once, then finish the radix search locally. + int const sb0 = fa::calcStartBit(0); + int const nCand = reinterpret_cast(smem + kBuckets * sizeof(float))[kbits >> sb0]; + if (pass == 0 && nCand <= kCandMax) + { + for (int j0 = start4 - lane; j0 < n4; j0 += stride4) + { + int const j = j0 + lane; + if (j < n4) + { + float4 const v = in4[j]; + float const e[4] = {v.x, v.y, v.z, v.w}; +#pragma unroll + for (int k = 0; k < 4; ++k) + { + if (((fa::twiddleIn(e[k], false) >> sb0) << sb0) == kbits) + { + w->cand[atomicAdd(&w->nCand, 1u)] = e[k]; + } + } + } + } + allArrive(w, G, 1); + clearPass(w, 0); // every CTA finished selectPass(0) before arriving here + // The list in this CTA's dynamic shared memory (privatized copy 1's space), the local histograms in + // copy 0's: the same integer histograms the global passes 1 .. NUM_PASSES - 1 would build. + float* candS = reinterpret_cast(hsCopies + kBuckets); + for (int i = tid; i < nCand; i += kThreads) + { + candS[i] = __ldcg(&w->cand[i]); + } + __syncthreads(); + for (int lp = 1; lp < fa::NUM_PASSES; ++lp) + { + for (int i = tid; i < kBuckets; i += kThreads) + { + hsCopies[i] = 0; + csCopies[i] = 0; + } + __syncthreads(); + int const lsb = fa::calcStartBit(lp); + uint32_t const lmask = fa::calcMask(lp); + int const lpsb = fa::calcStartBit(lp - 1); + for (int i = tid; i < nCand; i += kThreads) + { + float const x = candS[i]; + auto const bits = fa::twiddleIn(x, false); + if (((bits >> lpsb) << lpsb) == kbits) + { + int const b = static_cast((bits >> lsb) & lmask); + atomicAdd(hsCopies + b, static_cast(fa::calcMantissa(x))); + atomicAdd(csCopies + b, 1); + } + } + __syncthreads(); + selectPassOn(hsCopies, csCopies, lp, p, reinterpret_cast(smem), + reinterpret_cast(smem + kBuckets * sizeof(float)), &sSum, &sKbits); + kbits = sKbits; + } + break; + } + } + } + + // Apply: the threshold, FlashInfer's per-thread chain sums (spread over the G CTAs), the cub block sum, then the + // renorm write. + float const thr = keepAll ? -INFINITY : fa::twiddleOut(kbits, false); + int const cpc = kThreads / G; // chains per CTA + int const c0 = cta * cpc; + int const kTile = (kSmemBytes / 4) / cpc; + int const kTotal = (d + kThreads - 1) / kThreads; + float* stage = reinterpret_cast(smem); + float acc = 0.0f; + int kept = 0; + for (int k0 = 0; k0 < kTotal; k0 += kTile) + { + int const kn = min(kTile, kTotal - k0); + for (int q = tid; q < kn * cpc; q += kThreads) + { + int const kk = q / cpc, c = q - kk * cpc; + int const i = c0 + c + kThreads * (k0 + kk); + stage[q] = i < d ? in[i] : 0.0f; + } + __syncthreads(); + if (tid < cpc) + { + for (int kk = 0; kk < kn; ++kk) + { + int const i = c0 + tid + kThreads * (k0 + kk); + float const val = stage[kk * cpc + tid]; + if (i < d && val >= thr) + { + acc += val; + ++kept; + } + } + } + __syncthreads(); + } + if (tid < cpc) + { + w->chain[c0 + tid] = acc; + } + typedef cub::BlockReduce BlockReduce; + __shared__ typename BlockReduce::TempStorage temp; + __shared__ int ctaKept; + if (tid == 0) + { + ctaKept = 0; + } + __syncthreads(); + if (tid < cpc && kept) + { + atomicAdd(&ctaKept, kept); + } + __syncthreads(); + if (tid == 0 && ctaKept) + { + atomicAdd(&w->kept, ctaKept); + } + allArrive(w, G, fa::NUM_PASSES); + if (!keepAll) + { + clearPass(w, fa::NUM_PASSES - 1); // every CTA finished the last selectPass before arriving here + } + __shared__ float sNorm; + { + float totalSum = BlockReduce(temp).Sum(__ldcg(&w->chain[tid])); + if (tid == 0) + { + sNorm = (totalSum > 1e-8f) ? (1.0f / totalSum) : 1.0f; + } + } + __syncthreads(); + if (cta == 0 && tid == 0) + { + if (keptOut) + { + keptOut[row] = __ldcg(&w->kept); + } + w->kept = 0; // only CTA 0 reads it; the next launch adds to it after its own grid-dependency wait + if constexpr (kCandList) + { + w->nCand = 0; // every CTA copied the list before arriving at the norm barrier + } + } + // The workspace is left ready for the next launch here (the last clear above), so dependents may start; their + // grid-dependency wait still covers this launch's renorm write. + if constexpr (kPdl) + { + cudaTriggerProgrammaticLaunchCompletion(); + } + float const norm = sNorm; + float4* o4 = reinterpret_cast(o); + for (int j = start4; j < n4; j += stride4) + { + float4 v = in4[j]; + float4 r; + r.x = (v.x >= thr) ? v.x * norm : 0.0f; + r.y = (v.y >= thr) ? v.y * norm : 0.0f; + r.z = (v.z >= thr) ? v.z * norm : 0.0f; + r.w = (v.w >= thr) ? v.w * norm : 0.0f; + o4[j] = r; + } +} + +} // namespace toppv3 + +cudaError_t invokeToppRenormDetV3(float const* probs, float* out, float const* topP, int rows, int d, int ctasPerRow, + void* workspace, int* keptOut, bool pdl, bool candidateList, cudaStream_t stream) +{ + using namespace toppv3; + if (rows <= 0 || d <= 0) + { + return cudaSuccess; + } + if (ctasPerRow < 1 || ctasPerRow > kThreads || kThreads % ctasPerRow != 0 || d % 4 != 0) + { + return cudaErrorInvalidValue; + } + static cudaError_t const attrErr = [] + { + for (auto* k : {toppRenormDetV3Kernel, toppRenormDetV3Kernel, + toppRenormDetV3Kernel, toppRenormDetV3Kernel}) + { + cudaError_t const e + = cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(kDynSmemBytes)); + if (e != cudaSuccess) + { + return e; + } + } + return cudaSuccess; + }(); + if (attrErr != cudaSuccess) + { + return attrErr; + } + cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(static_cast(ctasPerRow), static_cast(rows)); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = kDynSmemBytes; + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + auto* ws = static_cast(workspace); + auto* kernel = pdl ? (candidateList ? toppRenormDetV3Kernel : toppRenormDetV3Kernel) + : (candidateList ? toppRenormDetV3Kernel : toppRenormDetV3Kernel); + return cudaLaunchKernelEx(&cfg, kernel, probs, out, topP, d, ws, keptOut); +} + +size_t toppRenormDetV3WorkspaceBytesPerRow() +{ + return sizeof(toppv3::RowWorkspace); +} + +} // namespace minimaxM31 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM31/toppRenormDetV3Sm103.h b/cpp/tensorrt_llm/kernels/minimaxM31/toppRenormDetV3Sm103.h new file mode 100644 index 000000000000..f7a08d3fb619 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM31/toppRenormDetV3Sm103.h @@ -0,0 +1,58 @@ +/* + * 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 minimaxM31 +{ + +//! TOPP_DET v3: the same renorm as invokeToppRenormDet (toppRenormDetSm103.h: semantics, the bitwise contract vs +//! FlashInfer's deterministic top_p_renorm_probs, the top_p >= 1 keep-all rule, --use_fast_math, the co-residency +//! rule for ctasPerRow), with the output bitwise equal to v1.2's on every row. What changes is the synchronization: +//! - no leader CTA: after each radix pass every CTA of the row waits for all of them, then runs FlashInfer's +//! selection itself on the row's global integer histograms of that pass (identical inputs and code, so identical +//! bits in every CTA); the norm is likewise reduced by every CTA from the same 1024 chain sums; +//! - one global histogram buffer per pass, cleared by CTA 0 one barrier later (once every CTA has read it); +//! - per-barrier arrival counters that only grow (each launch adds 1024 per barrier it runs, whatever its +//! ctasPerRow), so a workspace needs no reset between launches and replays, and launches with different +//! ctasPerRow may share it. +//! candidateList: when pass 0's threshold bucket holds at most 4096 elements, its members are listed once (one barrier) +//! and the remaining radix passes run locally in every CTA on the list (the same integer histograms, so the same +//! bits); larger buckets take the global passes. +//! pdl: launched as a programmatic dependent; everything that reads probs, top_p or the workspace waits for the +//! preceding grid, and dependents are released once the workspace is ready for the next launch. +//! Workspace: toppRenormDetV3WorkspaceBytesPerRow() x rows bytes, ZERO-initialized once before the first launch, +//! then reused as is (CUDA-graph safe). Launches sharing one workspace must be stream-ordered. +cudaError_t invokeToppRenormDetV3(float const* probs, float* out, float const* topP, int rows, int d, int ctasPerRow, + void* workspace, int* keptOut, bool pdl, bool candidateList, cudaStream_t stream); + +//! Workspace bytes per row of invokeToppRenormDetV3. +size_t toppRenormDetV3WorkspaceBytesPerRow(); + +} // namespace minimaxM31 +} // 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/kernels/speculativeDecoding/chainSpeculativeSamplingKernels.cu b/cpp/tensorrt_llm/kernels/speculativeDecoding/chainSpeculativeSamplingKernels.cu new file mode 100644 index 000000000000..f919dc72bf68 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/chainSpeculativeSamplingKernels.cu @@ -0,0 +1,283 @@ +/* + * 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 "chainSpeculativeSamplingKernels.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/fusedSampling/fusedSamplingKernels.h" + +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::speculative_decoding +{ + +namespace +{ + +constexpr int32_t kBlockSize = 1024; + +struct Float2SumOp +{ + __device__ __forceinline__ float2 operator()(float2 a, float2 b) const + { + return make_float2(a.x + b.x, a.y + b.y); + } +}; + +struct MaxOp +{ + __device__ __forceinline__ int32_t operator()(int32_t a, int32_t b) const + { + return a > b ? a : b; + } +}; + +template +__device__ __forceinline__ void loadProbs(float const* row, uint32_t base, uint32_t vocabSize, float (&out)[VEC_SIZE]) +{ + if constexpr (VEC_SIZE == 4) + { + if (base + VEC_SIZE <= vocabSize) + { + float4 const v = *reinterpret_cast(row + base); + out[0] = v.x; + out[1] = v.y; + out[2] = v.z; + out[3] = v.w; + return; + } + } +#pragma unroll + for (int32_t j = 0; j < VEC_SIZE; ++j) + { + out[j] = base + j < vocabSize ? row[base + j] : 0.0f; + } +} + +//! The distribution a request resamples from: max(0, target - draft) at the first rejected +//! position, or the target itself when every draft token was accepted (draft == nullptr) or +//! when the residual has no mass. +template +__device__ __forceinline__ void loadResidual( + float const* target, float const* draft, bool targetOnly, uint32_t base, uint32_t vocabSize, float (&out)[VEC_SIZE]) +{ + loadProbs(target, base, vocabSize, out); + if (targetOnly) + { + return; + } + float d[VEC_SIZE]; + loadProbs(draft, base, vocabSize, d); +#pragma unroll + for (int32_t j = 0; j < VEC_SIZE; ++j) + { + out[j] = fmaxf(out[j] - d[j], 0.0f); + } +} + +template +__global__ void __launch_bounds__(BLOCK_SIZE) chainSpeculativeSamplingKernel(ChainSpeculativeSamplingParams params) +{ + using BlockReduce = cub::BlockReduce; + using BlockReducePair = cub::BlockReduce; + using BlockReduceInt = cub::BlockReduce; + using BlockScan = cub::BlockScan; + + __shared__ union + { + typename BlockReduce::TempStorage reduce; + typename BlockReducePair::TempStorage reducePair; + typename BlockReduceInt::TempStorage reduceInt; + typename BlockScan::TempStorage scan; + } tempStorage; + + __shared__ int32_t sPos; + __shared__ float sU; + __shared__ bool sTargetOnly; + __shared__ float sTotal; + __shared__ float sTileSum[2]; + __shared__ int32_t sWinner; + __shared__ int32_t sLastValid; + + int64_t const row = blockIdx.x; + int32_t const tid = static_cast(threadIdx.x); + int32_t const numDraft = params.numDraftTokens; + uint32_t const vocabSize = static_cast(params.vocabSize); + int32_t* outTokens = params.outputTokenIds + row * (numDraft + 1); + + if (tid == 0) + { + curandStatePhilox4_32_10_t state; + curand_init(static_cast(params.seed[row]), + kPerRowSubsequenceBase | static_cast(params.offset[row]), 0, &state); + int32_t pos = numDraft; + for (int32_t i = 0; i < numDraft; ++i) + { + int32_t const token = params.draftTokenIds[row * numDraft + i]; + float const pTarget = params.targetProbs[(row * (numDraft + 1) + i) * vocabSize + token]; + float const pDraft = params.draftProbs[(row * numDraft + i) * vocabSize + token]; + if (curand_uniform(&state) * pDraft < pTarget) + { + outTokens[i] = token; + } + else + { + pos = i; + break; + } + } + sPos = pos; + // curand_uniform is in (0, 1]; the inverse-CDF search below wants [0, 1). + sU = 1.0f - curand_uniform(&state); + sWinner = static_cast(vocabSize); + } + __syncthreads(); + + int32_t const pos = sPos; + float const* target = params.targetProbs + (row * (numDraft + 1) + pos) * vocabSize; + float const* draft = pos < numDraft ? params.draftProbs + (row * numDraft + pos) * vocabSize : nullptr; + uint32_t const tileSize = BLOCK_SIZE * VEC_SIZE; + + // x: target mass, y: residual mass. + float2 sums = make_float2(0.0f, 0.0f); + for (uint32_t base = tid * VEC_SIZE; base < vocabSize; base += tileSize) + { + float t[VEC_SIZE]; + loadProbs(target, base, vocabSize, t); + float d[VEC_SIZE] = {}; + if (draft != nullptr) + { + loadProbs(draft, base, vocabSize, d); + } +#pragma unroll + for (int32_t j = 0; j < VEC_SIZE; ++j) + { + sums.x += t[j]; + sums.y += fmaxf(t[j] - d[j], 0.0f); + } + } + float2 const blockSums = BlockReducePair(tempStorage.reducePair).Reduce(sums, Float2SumOp{}); + if (tid == 0) + { + sTargetOnly = draft == nullptr || !(blockSums.y > 0.0f); + sTotal = sTargetOnly ? blockSums.x : blockSums.y; + } + __syncthreads(); + + // Inverse-CDF search: every tile is only summed, and the one tile the threshold falls in is + // scanned to find the token. + bool const targetOnly = sTargetOnly; + float const threshold = sU * sTotal; + float prefix = 0.0f; + int32_t localLastValid = -1; + int32_t parity = 0; + for (uint32_t tileBase = 0; tileBase < vocabSize; tileBase += tileSize, parity ^= 1) + { + uint32_t const base = tileBase + tid * VEC_SIZE; + float value[VEC_SIZE]; + loadResidual(target, draft, targetOnly, base, vocabSize, value); +#pragma unroll + for (int32_t j = 0; j < VEC_SIZE; ++j) + { + if (value[j] > 0.0f) + { + localLastValid = static_cast(base + j); + } + } + float const tileSum = BlockReduce(tempStorage.reduce).Sum(value); + if (tid == 0) + { + sTileSum[parity] = tileSum; + } + __syncthreads(); + if (prefix + sTileSum[parity] > threshold) + { + float inclusive[VEC_SIZE]; + BlockScan(tempStorage.scan).InclusiveSum(value, inclusive); +#pragma unroll + for (int32_t j = 0; j < VEC_SIZE; ++j) + { + if (value[j] > 0.0f && prefix + inclusive[j] > threshold) + { + atomicMin(&sWinner, static_cast(base + j)); + break; + } + } + __syncthreads(); + break; + } + prefix += sTileSum[parity]; + } + + if (sWinner == static_cast(vocabSize)) + { + // Float rounding can leave the summed prefix just short of the threshold; the last + // token with mass is then the one the threshold falls on. + int32_t const lastValid = BlockReduceInt(tempStorage.reduceInt).Reduce(localLastValid, MaxOp{}); + if (tid == 0) + { + sLastValid = lastValid; + } + } + + if (tid == 0) + { + int32_t const token = sWinner < static_cast(vocabSize) ? sWinner : max(sLastValid, 0); + outTokens[pos] = token; + for (int32_t i = pos + 1; i <= numDraft; ++i) + { + outTokens[i] = -1; + } + params.outputNumTokens[row] = pos + 1; + } +} + +bool isAligned(void const* ptr, uintptr_t alignment) +{ + return reinterpret_cast(ptr) % alignment == 0; +} + +} // namespace + +void invokeChainSpeculativeSampling(ChainSpeculativeSamplingParams const& params, cudaStream_t stream) +{ + if (params.batchSize == 0) + { + return; + } + dim3 const grid(params.batchSize); + dim3 const block(kBlockSize); + bool const vectorized = params.vocabSize % 4 == 0 && isAligned(params.targetProbs, sizeof(float4)) + && isAligned(params.draftProbs, sizeof(float4)); + if (vectorized) + { + chainSpeculativeSamplingKernel<<>>(params); + } + else + { + chainSpeculativeSamplingKernel<<>>(params); + } + sync_check_cuda_error(stream); +} + +} // namespace kernels::speculative_decoding + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/chainSpeculativeSamplingKernels.h b/cpp/tensorrt_llm/kernels/speculativeDecoding/chainSpeculativeSamplingKernels.h new file mode 100644 index 000000000000..e3da5ec5b78f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/chainSpeculativeSamplingKernels.h @@ -0,0 +1,65 @@ +/* + * 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::speculative_decoding +{ + +//! Chain speculative sampling: verifies each request's linear draft against the target. +//! +//! Draft token i is accepted with probability min(1, p_target / p_draft); the first rejected +//! position is replaced by a sample from norm(max(0, p_target - p_draft)), and when every +//! draft token is accepted a bonus token is sampled from the last target distribution. The +//! output distribution is the target's. +//! +//! Each request draws from the Philox subsequence its own offset names, so its result does not +//! depend on its row or on the other requests. +struct ChainSpeculativeSamplingParams +{ + //! [batchSize, numDraftTokens, vocabSize] + float const* draftProbs{nullptr}; + //! [batchSize, numDraftTokens] + int32_t const* draftTokenIds{nullptr}; + //! [batchSize, numDraftTokens + 1, vocabSize] + float const* targetProbs{nullptr}; + //! [batchSize] + int64_t const* seed{nullptr}; + //! [batchSize] + int64_t const* offset{nullptr}; + //! Output [batchSize, numDraftTokens + 1]: the accepted draft tokens, then the corrected or + //! bonus token, then -1. + int32_t* outputTokenIds{nullptr}; + //! Output [batchSize]: accepted draft tokens + 1. + int32_t* outputNumTokens{nullptr}; + int32_t batchSize{0}; + int32_t numDraftTokens{0}; + int32_t vocabSize{0}; +}; + +void invokeChainSpeculativeSampling(ChainSpeculativeSamplingParams const& params, cudaStream_t stream); + +} // namespace kernels::speculative_decoding + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index ab424318c862..5b23408cf16c 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -76,6 +76,38 @@ add_th_common_family( nvfp4MlaKvCacheGatherOp.cpp trtllmGenQKVProcessOp.cpp) +# MiniMax-M3.1 GB300 native dependencies; explicit to avoid duplicate +# registrations. +target_sources( + th_common_attention + PRIVATE minimaxM31/minimaxM31ChainSpecSamplingMcOp.cpp + minimaxM31/minimaxM31Sm103CtxNormMxfp8Op.cpp + minimaxM31/minimaxM31Sm103SelectPrefillTileOp.cpp + minimaxM31/minimaxM31SoftmaxMcOp.cpp + minimaxM31/minimaxM31ToppRenormDetOp.cpp + minimaxM31/minimaxM31ToppRenormDetV3Op.cpp + minimaxM31/minimaxM31Tp4SparseDecodeLcOp.cpp + minimaxM31/minimaxM31Tp8SparseDecodeLc8Op.cpp + chainSpeculativeSamplingOp.cpp + minimaxM3/minimaxM3AttnC1Op.cpp + minimaxM3/minimaxM3AttnMegaOp.cpp + minimaxM3/minimaxM3DraftStepOp.cpp + minimaxM3/minimaxM3FusedDecodeAttnOp.cpp + minimaxM3/minimaxM3McAllreduceNormOp.cpp + minimaxM3/minimaxM3MsaKvUpdateOp.cpp + minimaxM3/minimaxM3Mxfp8LowMGemmOp.cpp + minimaxM3/minimaxM3ProducerPwOp.cpp + minimaxM3/minimaxM3RouterGemvOp.cpp + minimaxM3/minimaxM3SampledAcceptOp.cpp + minimaxM3/minimaxM3SelectLongOp.cpp + minimaxM3/minimaxM3ShardAcceptOp.cpp + minimaxM3/minimaxM3SparseDecodeLcOp.cpp + minimaxM3/minimaxM3SpecTailOp.cpp + minimaxM3/minimaxM3TargetHeadOp.cpp + minimaxM3/minimaxM3Tp8FusedScoreSelectOp.cpp + minimaxM3/minimaxM3Tp8SparseDecodeLc8Op.cpp + minimaxM3/stepInputsScatterOp.cpp) + add_th_common_family( th_common_gemm_quant cutlassScaledMM.cpp diff --git a/cpp/tensorrt_llm/thop/chainSpeculativeSamplingOp.cpp b/cpp/tensorrt_llm/thop/chainSpeculativeSamplingOp.cpp new file mode 100644 index 000000000000..ae9b5cb80a2b --- /dev/null +++ b/cpp/tensorrt_llm/thop/chainSpeculativeSamplingOp.cpp @@ -0,0 +1,91 @@ +/* + * 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/speculativeDecoding/chainSpeculativeSamplingKernels.h" + +#include +#include + +#include + +namespace torch_ext +{ + +namespace tk = tensorrt_llm::kernels::speculative_decoding; + +//! Returns (output_token_ids [batch, numDraft + 1], num_tokens [batch]); see +//! ChainSpeculativeSamplingParams for the semantics. +std::tuple chain_speculative_sampling(torch::Tensor const& draftProbs, + torch::Tensor const& draftTokenIds, torch::Tensor const& targetProbs, torch::Tensor const& seed, + torch::Tensor const& offset) +{ + TORCH_CHECK(draftProbs.dim() == 3, "draft_probs must be [batch, numDraft, vocab]"); + TORCH_CHECK(draftTokenIds.dim() == 2, "draft_token_ids must be [batch, numDraft]"); + TORCH_CHECK(targetProbs.dim() == 3, "target_probs must be [batch, numDraft + 1, vocab]"); + + auto const batchSize = draftProbs.size(0); + auto const numDraft = draftProbs.size(1); + auto const vocabSize = draftProbs.size(2); + TORCH_CHECK(draftTokenIds.size(0) == batchSize && draftTokenIds.size(1) == numDraft, + "draft_token_ids must be [batch, numDraft]"); + TORCH_CHECK( + targetProbs.size(0) == batchSize && targetProbs.size(1) == numDraft + 1 && targetProbs.size(2) == vocabSize, + "target_probs must be [batch, numDraft + 1, vocab]"); + + for (auto const* t : {&draftProbs, &draftTokenIds, &targetProbs, &seed, &offset}) + { + TORCH_CHECK(t->is_cuda() && t->is_contiguous(), "chain_speculative_sampling inputs must be contiguous CUDA"); + } + TORCH_CHECK(draftProbs.scalar_type() == torch::kFloat32 && targetProbs.scalar_type() == torch::kFloat32, + "draft_probs and target_probs must be float32"); + TORCH_CHECK(draftTokenIds.scalar_type() == torch::kInt32, "draft_token_ids must be int32"); + TORCH_CHECK(seed.scalar_type() == torch::kInt64 && offset.scalar_type() == torch::kInt64, + "seed and offset must be int64 tensors"); + TORCH_CHECK( + seed.numel() == batchSize && offset.numel() == batchSize, "seed and offset must hold one entry per request"); + + auto outputTokenIds = torch::empty({batchSize, numDraft + 1}, draftTokenIds.options()); + auto outputNumTokens = torch::empty({batchSize}, draftTokenIds.options()); + + tk::ChainSpeculativeSamplingParams params; + params.draftProbs = draftProbs.const_data_ptr(); + params.draftTokenIds = draftTokenIds.const_data_ptr(); + params.targetProbs = targetProbs.const_data_ptr(); + params.seed = seed.const_data_ptr(); + params.offset = offset.const_data_ptr(); + params.outputTokenIds = outputTokenIds.data_ptr(); + params.outputNumTokens = outputNumTokens.data_ptr(); + params.batchSize = static_cast(batchSize); + params.numDraftTokens = static_cast(numDraft); + params.vocabSize = static_cast(vocabSize); + tk::invokeChainSpeculativeSampling(params, at::cuda::getCurrentCUDAStream()); + return {outputTokenIds, outputNumTokens}; +} + +} // namespace torch_ext + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "chain_speculative_sampling(Tensor draft_probs, Tensor draft_token_ids, Tensor target_probs, " + "Tensor seed, Tensor offset) -> (Tensor, Tensor)"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("chain_speculative_sampling", &torch_ext::chain_speculative_sampling); +} 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/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/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/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31ChainSpecSamplingMcOp.cpp b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31ChainSpecSamplingMcOp.cpp new file mode 100644 index 000000000000..aa2e7e5fdd68 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31ChainSpecSamplingMcOp.cpp @@ -0,0 +1,140 @@ +/* + * 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/minimaxM31/chainSpecSamplingMcSm103.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! ST-01-REJ's multi-CTA chain speculative sampling (see kernels/minimaxM31/chainSpecSamplingMcSm103.h): the inputs +//! and outputs of trtllm::chain_speculative_sampling (output_token_ids [batch, numDraft + 1], num_tokens [batch]), +//! bitwise its results, on batch x ctas_per_row CTAs. workspace: a zero-initialized byte buffer of at least +//! batch x minimax_m31_chain_speculative_sampling_mc_ws_bytes(), reused across calls; ctas_per_row divides 1024 and +//! batch x ctas_per_row must not exceed the SM count when ctas_per_row > 1. +std::tuple minimaxM31ChainSpecSamplingMc(torch::Tensor const& draftProbs, + torch::Tensor const& draftTokenIds, torch::Tensor const& targetProbs, torch::Tensor const& seed, + torch::Tensor const& offset, torch::Tensor& workspace, int64_t ctasPerRow) +{ + TORCH_CHECK(draftProbs.dim() == 3, + "minimax_m31_chain_speculative_sampling_mc: draft_probs must be [batch, " + "numDraft, vocab]"); + TORCH_CHECK(draftTokenIds.dim() == 2, + "minimax_m31_chain_speculative_sampling_mc: draft_token_ids must be [batch, " + "numDraft]"); + TORCH_CHECK(targetProbs.dim() == 3, + "minimax_m31_chain_speculative_sampling_mc: target_probs must be [batch, numDraft + 1, vocab]"); + auto const batchSize = draftProbs.size(0); + auto const numDraft = draftProbs.size(1); + auto const vocabSize = draftProbs.size(2); + TORCH_CHECK(draftTokenIds.size(0) == batchSize && draftTokenIds.size(1) == numDraft, + "minimax_m31_chain_speculative_sampling_mc: draft_token_ids must be [batch, numDraft]"); + TORCH_CHECK( + targetProbs.size(0) == batchSize && targetProbs.size(1) == numDraft + 1 && targetProbs.size(2) == vocabSize, + "minimax_m31_chain_speculative_sampling_mc: target_probs must be [batch, numDraft + 1, vocab]"); + for (auto const* t : {&draftProbs, &draftTokenIds, &targetProbs, &seed, &offset}) + { + TORCH_CHECK(t->is_cuda() && t->is_contiguous() && t->get_device() == draftProbs.get_device(), + "minimax_m31_chain_speculative_sampling_mc: inputs must be contiguous CUDA tensors on one device"); + } + TORCH_CHECK(draftProbs.scalar_type() == torch::kFloat32 && targetProbs.scalar_type() == torch::kFloat32, + "minimax_m31_chain_speculative_sampling_mc: draft_probs and target_probs must be float32"); + TORCH_CHECK(draftTokenIds.scalar_type() == torch::kInt32, + "minimax_m31_chain_speculative_sampling_mc: draft_token_ids must be int32"); + TORCH_CHECK(seed.scalar_type() == torch::kInt64 && offset.scalar_type() == torch::kInt64, + "minimax_m31_chain_speculative_sampling_mc: seed and offset must be int64 tensors"); + TORCH_CHECK(seed.numel() == batchSize && offset.numel() == batchSize, + "minimax_m31_chain_speculative_sampling_mc: seed and offset must hold one entry per request"); + TORCH_CHECK(vocabSize % 4 == 0 && vocabSize <= tensorrt_llm::kernels::minimaxM31::kChainSpecMcMaxVocab + && reinterpret_cast(draftProbs.data_ptr()) % 16 == 0 + && reinterpret_cast(targetProbs.data_ptr()) % 16 == 0, + "minimax_m31_chain_speculative_sampling_mc: the probs must be 16-byte aligned with a vocab divisible by 4 and " + "at most ", + tensorrt_llm::kernels::minimaxM31::kChainSpecMcMaxVocab); + TORCH_CHECK(batchSize <= 65535, "minimax_m31_chain_speculative_sampling_mc: batch must be <= 65535"); + TORCH_CHECK(ctasPerRow >= 1 && ctasPerRow <= 1024 && 1024 % ctasPerRow == 0, + "minimax_m31_chain_speculative_sampling_mc: ctas_per_row must divide 1024"); + auto const wsBytesPerRow + = static_cast(tensorrt_llm::kernels::minimaxM31::chainSpecSamplingMcWorkspaceBytesPerRow()); + TORCH_CHECK(workspace.is_cuda() && workspace.is_contiguous() && workspace.get_device() == draftProbs.get_device() + && reinterpret_cast(workspace.data_ptr()) % 16 == 0 + && workspace.numel() * workspace.element_size() >= wsBytesPerRow * batchSize, + "minimax_m31_chain_speculative_sampling_mc: workspace must be a contiguous, 16-byte aligned CUDA buffer of at " + "least ", + wsBytesPerRow * batchSize, " bytes on draft_probs' device"); + + c10::cuda::CUDAGuard const deviceGuard{draftProbs.device()}; + TORCH_CHECK( + ctasPerRow == 1 || batchSize * ctasPerRow <= at::cuda::getCurrentDeviceProperties()->multiProcessorCount, + "minimax_m31_chain_speculative_sampling_mc: batch x ctas_per_row must not exceed the SM count"); + auto outputTokenIds = torch::empty({batchSize, numDraft + 1}, draftTokenIds.options()); + auto outputNumTokens = torch::empty({batchSize}, draftTokenIds.options()); + if (batchSize == 0) + { + return {outputTokenIds, outputNumTokens}; + } + + tensorrt_llm::kernels::speculative_decoding::ChainSpeculativeSamplingParams params; + params.draftProbs = draftProbs.const_data_ptr(); + params.draftTokenIds = draftTokenIds.const_data_ptr(); + params.targetProbs = targetProbs.const_data_ptr(); + params.seed = seed.const_data_ptr(); + params.offset = offset.const_data_ptr(); + params.outputTokenIds = outputTokenIds.data_ptr(); + params.outputNumTokens = outputNumTokens.data_ptr(); + params.batchSize = static_cast(batchSize); + params.numDraftTokens = static_cast(numDraft); + params.vocabSize = static_cast(vocabSize); + auto const err = tensorrt_llm::kernels::minimaxM31::invokeChainSpecSamplingMc(params, static_cast(ctasPerRow), + workspace.data_ptr(), at::cuda::getCurrentCUDAStream(draftProbs.get_device())); + TORCH_CHECK( + err == cudaSuccess, "minimax_m31_chain_speculative_sampling_mc launch failed: ", cudaGetErrorString(err)); + return {outputTokenIds, outputNumTokens}; +} + +//! Workspace bytes per request of minimax_m31_chain_speculative_sampling_mc. +int64_t minimaxM31ChainSpecSamplingMcWsBytes() +{ + return static_cast(tensorrt_llm::kernels::minimaxM31::chainSpecSamplingMcWorkspaceBytesPerRow()); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m31_chain_speculative_sampling_mc(Tensor draft_probs, Tensor draft_token_ids, Tensor target_probs, " + "Tensor seed, Tensor offset, Tensor(a!) workspace, int ctas_per_row) -> (Tensor, Tensor)"); + // No tensor arguments: a catch-all kernel (as moe_a2a_get_timeout), callable without a device. + m.def("minimax_m31_chain_speculative_sampling_mc_ws_bytes() -> int", + &tensorrt_llm::torch_ext::minimaxM31ChainSpecSamplingMcWsBytes); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m31_chain_speculative_sampling_mc", &tensorrt_llm::torch_ext::minimaxM31ChainSpecSamplingMc); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31Sm103CtxNormMxfp8Op.cpp b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31Sm103CtxNormMxfp8Op.cpp new file mode 100644 index 000000000000..580623a89393 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31Sm103CtxNormMxfp8Op.cpp @@ -0,0 +1,100 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM31/ctxResidualRmsNormMxfp8Sm103.h" + +#include +#include +#include + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! MiniMax-M3.1 sm_103 targets (target-local copy of the MiniMax-M3 SM107 CTX op; PDL is a caller argument). +//! Residual add + RMSNorm of an all-reduce output (in place: inter becomes inter + residual), with the MXFP8 +//! quantization of the normalized rows in the same pass. Returns (norm_out, mxfp8, sf, fp4, fp4_sf) where mxfp8 / sf +//! match mxfp8_quantize(norm_out, swizzled=True) bit for bit (mxfp8 [rows, hidden] Float8_e4m3fn, a 1-D UINT8 scale +//! tensor of the swizzled 128x4 layout). With fp4GlobalScale, fp4 / fp4_sf match +//! fp4_quantize(norm_out, fp4GlobalScale, 16, sf_use_ue8m0=False, is_sf_swizzled=False) bit for bit (fp4 [rows, +//! hidden / 2] UINT8, a 1-D UINT8 linear scale tensor of rows * hidden / 16); otherwise both are empty. norm_out is +//! empty when writeNorm is false. +std::tuple minimaxM31CtxResidualRmsNormMxfp8( + torch::Tensor& inter, torch::Tensor const& residual, torch::Tensor const& weight, double eps, bool writeNorm, + std::optional const& fp4GlobalScale, bool enablePdl) +{ + TORCH_CHECK(inter.is_cuda() && inter.is_contiguous() && inter.scalar_type() == at::kBFloat16 && inter.dim() >= 2); + TORCH_CHECK( + residual.is_contiguous() && residual.scalar_type() == at::kBFloat16 && residual.sizes() == inter.sizes()); + int64_t const hidden = inter.size(-1); + int64_t const rows = inter.numel() / hidden; + TORCH_CHECK(weight.is_contiguous() && weight.scalar_type() == at::kBFloat16 && weight.numel() == hidden); + TORCH_CHECK(hidden % 256 == 0 && hidden / 8 <= 1024, "hidden must be a multiple of 256 and at most 8192"); + at::cuda::CUDAGuard const guard(inter.device()); + auto const* deviceProperties = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(deviceProperties->major == 10 && deviceProperties->minor == 3, + "minimax_m31_ctx_residual_rmsnorm_mxfp8 requires SM103"); + auto const opts = inter.options(); + torch::Tensor normOut = writeNorm ? torch::empty_like(inter) : torch::empty({0}, opts); + torch::Tensor qOut = torch::empty(inter.sizes(), opts.dtype(at::ScalarType::Float8_e4m3fn)); + int64_t const sfSize = (rows + 127) / 128 * 128 * ((hidden / 32 + 3) / 4 * 4); + torch::Tensor sfOut = torch::empty({sfSize}, opts.dtype(at::ScalarType::Byte)); + bool const fp4 = fp4GlobalScale.has_value(); + if (fp4) + { + TORCH_CHECK( + fp4GlobalScale->is_cuda() && fp4GlobalScale->scalar_type() == at::kFloat && fp4GlobalScale->numel() == 1); + } + std::vector fp4Shape(inter.sizes().begin(), inter.sizes().end()); + fp4Shape.back() = hidden / 2; + torch::Tensor fp4Out = fp4 ? torch::empty(fp4Shape, opts.dtype(at::ScalarType::Byte)) : torch::empty({0}, opts); + torch::Tensor fp4SfOut + = fp4 ? torch::empty({rows * (hidden / 16)}, opts.dtype(at::ScalarType::Byte)) : torch::empty({0}, opts); + tensorrt_llm::kernels::minimaxM31Sm103::invokeCtxResidualRmsNormMxfp8( + reinterpret_cast<__nv_bfloat16*>(inter.data_ptr()), reinterpret_cast<__nv_bfloat16 const*>(residual.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(weight.data_ptr()), + writeNorm ? reinterpret_cast<__nv_bfloat16*>(normOut.data_ptr()) : nullptr, + reinterpret_cast(qOut.data_ptr()), reinterpret_cast(sfOut.data_ptr()), + fp4 ? reinterpret_cast(fp4Out.data_ptr()) : nullptr, + fp4 ? reinterpret_cast(fp4SfOut.data_ptr()) : nullptr, + fp4 ? fp4GlobalScale->data_ptr() : nullptr, static_cast(rows), static_cast(hidden), + static_cast(eps), enablePdl, at::cuda::getCurrentCUDAStream(inter.get_device())); + return {normOut, qOut, sfOut, fp4Out, fp4SfOut}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m31_ctx_residual_rmsnorm_mxfp8(Tensor(a!) inter, Tensor residual, Tensor weight, float eps, " + "bool write_norm, Tensor? fp4_global_scale=None, bool enable_pdl=False) -> (Tensor, Tensor, Tensor, Tensor, " + "Tensor)"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m31_ctx_residual_rmsnorm_mxfp8", &tensorrt_llm::torch_ext::minimaxM31CtxResidualRmsNormMxfp8); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31Sm103SelectPrefillTileOp.cpp b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31Sm103SelectPrefillTileOp.cpp new file mode 100644 index 000000000000..6242fe6ce896 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31Sm103SelectPrefillTileOp.cpp @@ -0,0 +1,86 @@ +/* + * 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.1 GB300 (sm_103) target-local copy of the #20629 prefill tile select +// (cpp/tensorrt_llm/kernels/minimaxM3/selectPrefillTile.* @ cf8a77e816); namespace and op name are M3.1's. + +#include "tensorrt_llm/kernels/minimaxM31/selectPrefillTileSm103.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! Prefill top-16 block select over per-KV-head scores with unit query stride, into a caller-owned table. +//! +//! @param scores Float32 CUDA tensor [num_kv_heads, num_blocks, total_q] with stride(2) == 1. +//! @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 (a +//! head-major backing is a permuted view). Every entry is written. +void minimaxM31SelectPrefillTile(torch::Tensor const& scores, torch::Tensor const& nValidBlocks, torch::Tensor& out, + int64_t initBlocks, int64_t localBlocks) +{ + TORCH_CHECK(scores.is_cuda() && scores.scalar_type() == torch::kFloat32 && scores.dim() == 3, + "minimax_m31_select_prefill_tile: scores must be a 3-D float32 CUDA tensor [num_kv_heads, num_blocks, " + "total_q]"); + TORCH_CHECK(scores.stride(2) == 1, "minimax_m31_select_prefill_tile: scores must have unit query stride"); + TORCH_CHECK(nValidBlocks.is_cuda() && nValidBlocks.scalar_type() == torch::kInt32 && nValidBlocks.is_contiguous() + && nValidBlocks.numel() >= scores.size(2), + "minimax_m31_select_prefill_tile: 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::minimaxM31Sm103::kMinimaxM3SelectPrefillTileMaxBlocks, + "minimax_m31_select_prefill_tile supports at most ", + tensorrt_llm::kernels::minimaxM31Sm103::kMinimaxM3SelectPrefillTileMaxBlocks, " 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_m31_select_prefill_tile: 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::minimaxM31Sm103::invokeMinimaxM3SelectPrefillTile(scores.data_ptr(), + scores.stride(0), scores.stride(1), 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), stream); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "minimax_m31_select_prefill_tile launch failed: ", cudaGetErrorString(err)); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m31_select_prefill_tile(Tensor scores, Tensor n_valid_blocks, Tensor(a!) out, int init_blocks, " + "int local_blocks) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m31_select_prefill_tile", &tensorrt_llm::torch_ext::minimaxM31SelectPrefillTile); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31SoftmaxMcOp.cpp b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31SoftmaxMcOp.cpp new file mode 100644 index 000000000000..6d64fbadc8a4 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31SoftmaxMcOp.cpp @@ -0,0 +1,95 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM31/softmaxMcSm103.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! ST-01-SM's multi-CTA softmax (see kernels/minimaxM31/softmaxMcSm103.h): logits fp32 [rows, d] -> out fp32 +//! [rows, d], bitwise flashinfer.sampling.softmax(logits, temperature) on FlashInfer's Path A shapes (rows <= 128, +//! d >= 24576). temperature: an fp32 [rows] tensor, or None for the scalar temperature_val. partials: a CUDA buffer +//! of at least rows x ceil(d / 8192) float2, overwritten by every call; write_ctas: the CTAs per row of the write pass. +void minimaxM31SoftmaxMc(torch::Tensor const& logits, std::optional const& temperature, + double temperatureVal, torch::Tensor& partials, torch::Tensor& out, int64_t writeCtas) +{ + TORCH_CHECK( + logits.is_cuda() && logits.scalar_type() == torch::kFloat32 && logits.dim() == 2 && logits.is_contiguous(), + "minimax_m31_softmax_mc: logits must be a contiguous 2-D float32 CUDA tensor"); + int64_t const rows = logits.size(0); + int64_t const d = logits.size(1); + TORCH_CHECK(d % 4 == 0 && reinterpret_cast(logits.data_ptr()) % 16 == 0, + "minimax_m31_softmax_mc: logits must be 16-byte aligned with a width divisible by 4"); + TORCH_CHECK(rows <= 128 && d >= 24576, + "minimax_m31_softmax_mc: FlashInfer's Path A shapes only (rows <= 128, width >= 24576)"); + TORCH_CHECK(rows * d <= std::numeric_limits::max(), "minimax_m31_softmax_mc: logits is too large"); + float const* temperatureArr = nullptr; + if (temperature.has_value()) + { + auto const& t = temperature.value(); + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kFloat32 && t.dim() == 1 && t.is_contiguous() + && t.size(0) == rows && t.get_device() == logits.get_device(), + "minimax_m31_softmax_mc: temperature must be a contiguous float32 [rows] tensor on logits' device"); + temperatureArr = t.data_ptr(); + } + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kFloat32 && out.dim() == 2 && out.is_contiguous() + && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 16 == 0 + && out.get_device() == logits.get_device(), + "minimax_m31_softmax_mc: out must be a contiguous, 16-byte aligned float32 [rows, d] on logits' device"); + TORCH_CHECK(writeCtas >= 1 && writeCtas <= 65535, "minimax_m31_softmax_mc: write_ctas must be in [1, 65535]"); + auto const partialsBytes = static_cast( + tensorrt_llm::kernels::minimaxM31::softmaxMcPartialsBytes(static_cast(rows), static_cast(d))); + TORCH_CHECK(partials.is_cuda() && partials.is_contiguous() && partials.get_device() == logits.get_device() + && reinterpret_cast(partials.data_ptr()) % 8 == 0 + && partials.numel() * partials.element_size() >= partialsBytes, + "minimax_m31_softmax_mc: partials must be a contiguous, 8-byte aligned CUDA buffer of at least ", partialsBytes, + " bytes on logits' device"); + if (rows == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{logits.device()}; + auto const err = tensorrt_llm::kernels::minimaxM31::invokeSoftmaxMc(logits.data_ptr(), out.data_ptr(), + temperatureArr, static_cast(temperatureVal), static_cast(rows), static_cast(d), + static_cast(writeCtas), partials.data_ptr(), at::cuda::getCurrentCUDAStream(logits.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m31_softmax_mc launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m31_softmax_mc(Tensor logits, Tensor? temperature, float temperature_val, Tensor(a!) partials, " + "Tensor(b!) out, int write_ctas) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m31_softmax_mc", &tensorrt_llm::torch_ext::minimaxM31SoftmaxMc); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31ToppRenormDetOp.cpp b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31ToppRenormDetOp.cpp new file mode 100644 index 000000000000..bea1b0bbc95f --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31ToppRenormDetOp.cpp @@ -0,0 +1,98 @@ +/* + * 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/minimaxM31/toppRenormDetSm103.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! ST-01's det top-p renorm (see kernels/minimaxM31/toppRenormDetSm103.h): probs fp32 [rows, d] and top_p fp32 +//! [rows] -> out fp32 [rows, d], bitwise FlashInfer's top_p_renorm_probs(is_deterministic=True) on rows with +//! top_p < 1, keep-all on rows with top_p >= 1. workspace: a zero-initialized byte buffer of at least +//! rows x minimax_m31_topp_renorm_det_ws_bytes(), reused across calls; ctas_per_row divides 1024 and +//! rows x ctas_per_row must not exceed the SM count when ctas_per_row > 1. +void minimaxM31ToppRenormDet(torch::Tensor const& probs, torch::Tensor const& topP, torch::Tensor& workspace, + torch::Tensor& out, int64_t ctasPerRow) +{ + TORCH_CHECK(probs.is_cuda() && probs.scalar_type() == torch::kFloat32 && probs.dim() == 2 && probs.is_contiguous(), + "minimax_m31_topp_renorm_det: probs must be a contiguous 2-D float32 CUDA tensor"); + int64_t const rows = probs.size(0); + int64_t const d = probs.size(1); + TORCH_CHECK(d % 4 == 0 && reinterpret_cast(probs.data_ptr()) % 16 == 0, + "minimax_m31_topp_renorm_det: probs must be 16-byte aligned with a width divisible by 4"); + TORCH_CHECK(d <= std::numeric_limits::max() && rows <= 65535, + "minimax_m31_topp_renorm_det: probs is too large (rows <= 65535, width < 2^31)"); + TORCH_CHECK(topP.is_cuda() && topP.scalar_type() == torch::kFloat32 && topP.dim() == 1 && topP.is_contiguous() + && topP.size(0) == rows && topP.get_device() == probs.get_device(), + "minimax_m31_topp_renorm_det: top_p must be a contiguous float32 [rows] tensor on probs' device"); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kFloat32 && out.dim() == 2 && out.is_contiguous() + && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 16 == 0 + && out.get_device() == probs.get_device(), + "minimax_m31_topp_renorm_det: out must be a contiguous, 16-byte aligned float32 [rows, d] on probs' device"); + TORCH_CHECK(ctasPerRow >= 1 && ctasPerRow <= 1024 && 1024 % ctasPerRow == 0, + "minimax_m31_topp_renorm_det: ctas_per_row must divide 1024"); + auto const wsBytesPerRow + = static_cast(tensorrt_llm::kernels::minimaxM31::toppRenormDetWorkspaceBytesPerRow()); + TORCH_CHECK(workspace.is_cuda() && workspace.is_contiguous() && workspace.get_device() == probs.get_device() + && reinterpret_cast(workspace.data_ptr()) % 16 == 0 + && workspace.numel() * workspace.element_size() >= wsBytesPerRow * rows, + "minimax_m31_topp_renorm_det: workspace must be a contiguous, 16-byte aligned CUDA buffer of at least ", + wsBytesPerRow * rows, " bytes on probs' device"); + if (rows == 0 || d == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{probs.device()}; + auto const err + = tensorrt_llm::kernels::minimaxM31::invokeToppRenormDet(probs.data_ptr(), out.data_ptr(), + topP.data_ptr(), static_cast(rows), static_cast(d), static_cast(ctasPerRow), + workspace.data_ptr(), nullptr, at::cuda::getCurrentCUDAStream(probs.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m31_topp_renorm_det launch failed: ", cudaGetErrorString(err)); +} + +//! Workspace bytes per row of minimax_m31_topp_renorm_det. +int64_t minimaxM31ToppRenormDetWsBytes() +{ + return static_cast(tensorrt_llm::kernels::minimaxM31::toppRenormDetWorkspaceBytesPerRow()); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m31_topp_renorm_det(Tensor probs, Tensor top_p, Tensor(a!) workspace, Tensor(b!) out, " + "int ctas_per_row) -> ()"); + // No tensor arguments: a catch-all kernel (as moe_a2a_get_timeout), callable without a device. + m.def("minimax_m31_topp_renorm_det_ws_bytes() -> int", &tensorrt_llm::torch_ext::minimaxM31ToppRenormDetWsBytes); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m31_topp_renorm_det", &tensorrt_llm::torch_ext::minimaxM31ToppRenormDet); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31ToppRenormDetV3Op.cpp b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31ToppRenormDetV3Op.cpp new file mode 100644 index 000000000000..d125c4b79c0e --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31ToppRenormDetV3Op.cpp @@ -0,0 +1,98 @@ +/* + * 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/minimaxM31/toppRenormDetV3Sm103.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! TOPP_DET v3 (see kernels/minimaxM31/toppRenormDetV3Sm103.h): minimax_m31_topp_renorm_det's renorm, bitwise its +//! output, with the redundant per-CTA selection and an optional programmatic dependent launch. workspace: a +//! zero-initialized byte buffer of at least rows x minimax_m31_topp_renorm_det_v3_ws_bytes(), reused across calls; +//! ctas_per_row divides 1024 and rows x ctas_per_row must not exceed the SM count when ctas_per_row > 1. +void minimaxM31ToppRenormDetV3(torch::Tensor const& probs, torch::Tensor const& topP, torch::Tensor& workspace, + torch::Tensor& out, int64_t ctasPerRow, bool pdl, bool candidateList) +{ + TORCH_CHECK(probs.is_cuda() && probs.scalar_type() == torch::kFloat32 && probs.dim() == 2 && probs.is_contiguous(), + "minimax_m31_topp_renorm_det_v3: probs must be a contiguous 2-D float32 CUDA tensor"); + int64_t const rows = probs.size(0); + int64_t const d = probs.size(1); + TORCH_CHECK(d % 4 == 0 && reinterpret_cast(probs.data_ptr()) % 16 == 0, + "minimax_m31_topp_renorm_det_v3: probs must be 16-byte aligned with a width divisible by 4"); + TORCH_CHECK(d <= std::numeric_limits::max() && rows <= 65535, + "minimax_m31_topp_renorm_det_v3: probs is too large (rows <= 65535, width < 2^31)"); + TORCH_CHECK(topP.is_cuda() && topP.scalar_type() == torch::kFloat32 && topP.dim() == 1 && topP.is_contiguous() + && topP.size(0) == rows && topP.get_device() == probs.get_device(), + "minimax_m31_topp_renorm_det_v3: top_p must be a contiguous float32 [rows] tensor on probs' device"); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kFloat32 && out.dim() == 2 && out.is_contiguous() + && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 16 == 0 + && out.get_device() == probs.get_device(), + "minimax_m31_topp_renorm_det_v3: out must be a contiguous, 16-byte aligned float32 [rows, d] on probs' device"); + TORCH_CHECK(ctasPerRow >= 1 && ctasPerRow <= 1024 && 1024 % ctasPerRow == 0, + "minimax_m31_topp_renorm_det_v3: ctas_per_row must divide 1024"); + auto const wsBytesPerRow + = static_cast(tensorrt_llm::kernels::minimaxM31::toppRenormDetV3WorkspaceBytesPerRow()); + TORCH_CHECK(workspace.is_cuda() && workspace.is_contiguous() && workspace.get_device() == probs.get_device() + && reinterpret_cast(workspace.data_ptr()) % 16 == 0 + && workspace.numel() * workspace.element_size() >= wsBytesPerRow * rows, + "minimax_m31_topp_renorm_det_v3: workspace must be a contiguous, 16-byte aligned CUDA buffer of at least ", + wsBytesPerRow * rows, " bytes on probs' device"); + if (rows == 0 || d == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{probs.device()}; + auto const err + = tensorrt_llm::kernels::minimaxM31::invokeToppRenormDetV3(probs.data_ptr(), out.data_ptr(), + topP.data_ptr(), static_cast(rows), static_cast(d), static_cast(ctasPerRow), + workspace.data_ptr(), nullptr, pdl, candidateList, at::cuda::getCurrentCUDAStream(probs.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m31_topp_renorm_det_v3 launch failed: ", cudaGetErrorString(err)); +} + +//! Workspace bytes per row of minimax_m31_topp_renorm_det_v3. +int64_t minimaxM31ToppRenormDetV3WsBytes() +{ + return static_cast(tensorrt_llm::kernels::minimaxM31::toppRenormDetV3WorkspaceBytesPerRow()); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m31_topp_renorm_det_v3(Tensor probs, Tensor top_p, Tensor(a!) workspace, Tensor(b!) out, " + "int ctas_per_row, bool pdl, bool candidate_list) -> ()"); + // No tensor arguments: a catch-all kernel (as moe_a2a_get_timeout), callable without a device. + m.def( + "minimax_m31_topp_renorm_det_v3_ws_bytes() -> int", &tensorrt_llm::torch_ext::minimaxM31ToppRenormDetV3WsBytes); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m31_topp_renorm_det_v3", &tensorrt_llm::torch_ext::minimaxM31ToppRenormDetV3); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31Tp4SparseDecodeLcOp.cpp b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31Tp4SparseDecodeLcOp.cpp new file mode 100644 index 000000000000..5d188c43c702 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31Tp4SparseDecodeLcOp.cpp @@ -0,0 +1,144 @@ +/* + * 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/minimaxM31/sparseDecodeLcTp4.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_m31_tp4_sparse_decode_lc: ", name, + " must be a one-byte CUDA view (uint8 / e4m3 / packed e2m1)"); +} + +} // namespace + +//! One-launch MiniMax-M3.1 TP4 long-context sparse decode over the selected NVFP4 blocks (16 query heads on one KV +//! head, head_dim 128, page 128, top-k 16): the M3.1 target's copy of minimax_m3_sparse_decode_lc with two exact +//! options, bit-identical to it: btsmem (block-table row staged in shared memory before the grid wait) and sync3 +//! (softmax row max / sum read by every thread from the per-warp partials). +//! The compatibility one_div argument must be false. +//! +//! @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 minimaxM31Tp4SparseDecodeLc(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, + bool btsmem, bool sync3, bool oneDiv) +{ + TORCH_CHECK(!oneDiv, "one_div is excluded from this exact handoff"); + TORCH_CHECK(earlyTrigger, "minimax_m31_tp4_sparse_decode_lc is instantiated with early_trigger=True only"); + TORCH_CHECK(q.is_cuda(), "minimax_m31_tp4_sparse_decode_lc: q must be a CUDA tensor"); + c10::cuda::CUDAGuard const deviceGuard{q.device()}; + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 10, + "minimax_m31_tp4_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_m31_tp4_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_m31_tp4_sparse_decode_lc: global scales must be float32"); + TORCH_CHECK(topk.scalar_type() == torch::kInt32 && topk.dim() == 2 && topk.size(1) == kTopK, + "minimax_m31_tp4_sparse_decode_lc: topk must be int32 [T, 16]"); + TORCH_CHECK(blockTable.scalar_type() == torch::kInt32 && seqLens.scalar_type() == torch::kInt32, + "minimax_m31_tp4_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_m31_tp4_sparse_decode_lc: out must be bfloat16 [T, 2048] with unit last-dim stride"); + TORCH_CHECK(q8.has_value() == sf.has_value(), "minimax_m31_tp4_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::MinimaxM31Tp4SparseDecodeLcParams 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::invokeMinimaxM31Tp4SparseDecodeLc( + prm, static_cast(q.size(0)), linearSoftmax, mxfp8, earlyTrigger, btsmem, sync3, usePdl, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m31_tp4_sparse_decode_lc launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m31_tp4_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, bool btsmem, bool sync3, bool one_div=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m31_tp4_sparse_decode_lc", &tensorrt_llm::torch_ext::minimaxM31Tp4SparseDecodeLc); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31Tp8SparseDecodeLc8Op.cpp b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31Tp8SparseDecodeLc8Op.cpp new file mode 100644 index 000000000000..0290b536fa29 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM31/minimaxM31Tp8SparseDecodeLc8Op.cpp @@ -0,0 +1,143 @@ +/* + * 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/minimaxM31/sparseDecodeLc8Tp8.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_m31_tp8_sparse_decode_lc8: ", name, + " must be a one-byte CUDA view (uint8 / e4m3 / packed e2m1)"); +} + +} // namespace + +//! One-launch MiniMax-M3.1 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): the M3.1 target's copy of the M3 LC8 op, with the SYNC3 softmax (rows +//! bit-identical to minimax_m3_tp8_sparse_decode_lc8). 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). +//! @param oneDiv Compatibility argument; must be false. +void minimaxM31Tp8SparseDecodeLc8(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, + bool oneDiv) +{ + TORCH_CHECK(!oneDiv, "one_div is excluded from this exact handoff"); + TORCH_CHECK(q.is_cuda(), "minimax_m31_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_m31_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_m31_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_m31_tp8_sparse_decode_lc8: global scales must be float32"); + TORCH_CHECK(topk.scalar_type() == torch::kInt32 && topk.dim() == 2 && topk.size(1) == kTopK, + "minimax_m31_tp8_sparse_decode_lc8: topk must be int32 [T, 16]"); + TORCH_CHECK(blockTable.scalar_type() == torch::kInt32 && seqLens.scalar_type() == torch::kInt32, + "minimax_m31_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_m31_tp8_sparse_decode_lc8: out must be bfloat16 [T, 1024] with unit last-dim stride"); + TORCH_CHECK(q8.has_value() == sf.has_value(), "minimax_m31_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::MinimaxM31Tp8SparseDecodeLc8Params 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::invokeMinimaxM31Tp8SparseDecodeLc8( + prm, static_cast(q.size(0)), linearSoftmax, mxfp8, prewait, usePdl, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m31_tp8_sparse_decode_lc8 launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m31_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, bool one_div=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m31_tp8_sparse_decode_lc8", &tensorrt_llm::torch_ext::minimaxM31Tp8SparseDecodeLc8); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index cdb288f7220f..d54443b4ed74 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -37,7 +37,7 @@ import enum import os from importlib import import_module -from typing import TYPE_CHECKING, Optional, Tuple +from typing import TYPE_CHECKING, Dict, Optional, Tuple from tensorrt_llm.logger import logger @@ -49,6 +49,7 @@ # architectures[0] -> routing module, relative to this package. MODELING_V2_ROUTERS = { + "MiniMaxM3SparseForConditionalGeneration": "models.minimax_m3.routing", "GptOssForCausalLM": "models.gpt_oss.routing", "DeepseekV3ForCausalLM": "models.deepseek_v3.routing", } @@ -101,6 +102,62 @@ class NoTarget(Exception): """ +_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. + + Inside one node the IPC workspace (``comm/allreduce``) serves it. Across + nodes the targets take the MNNVL workspace (``comm/mnnvl_allreduce``), which + needs one NVLink domain over those nodes, as on GB200 and GB300 trays. This + is the check the built-in ``AllReduce`` makes before it takes MNNVL. Where it + fails, the built-in model falls back to NCCL, which the targets do not + implement, so a routing module refuses and ``auto`` keeps the built-in model. + Attention DP reduces nothing in its norm slots, and its allgather and + reduce-scatter run over NCCL. + """ + if mapping.enable_attention_dp or mapping.tp_size == 1 or not mapping.is_multi_node(): + return True + import torch + + from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce + + # bf16 is the activation dtype of both default targets. + return MNNVLAllReduce.is_mnnvl(mapping, torch.bfloat16) + + +def _kv_cache_quant_algo(config: "ModelConfig", llm_args: "TorchLlmArgs"): + """The KV-cache algorithm the engine will run with, which routing runs too early to read. + + The loader resolves the target before `validate_and_set_kv_cache_quant` + applies `kv_cache_config.dtype` to the config, so this applies the same + rule here: ``auto`` keeps the checkpoint's algorithm, ``fp8`` and + ``nvfp4`` replace it. Any other dtype comes back as the string given, so a + route refuses it by name: ``fp8_ds_mla`` is FP8 to the loader but also + selects a packed sparse-MLA cache layout, and an unknown dtype is one the + loader raises on later. + """ + from tensorrt_llm.quantization.mode import QuantAlgo + + dtype = llm_args.kv_cache_config.dtype + if dtype == "auto": + algo = config.quant_config.kv_cache_quant_algo + return None if algo is None else QuantAlgo(algo) + return {"fp8": QuantAlgo.FP8, "nvfp4": QuantAlgo.NVFP4}.get(dtype, dtype) + + def routing_module(arch: str): """Import and return the routing module for ``arch``, or None.""" name = MODELING_V2_ROUTERS.get(arch) @@ -143,6 +200,17 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti raise NoTarget( f"no routing module for this architecture; routed architectures: {known}" ) + if arch == "MiniMaxM3SparseForConditionalGeneration": + target = routing.route( + pretrained_config, + config.mapping, + _device_compute_capability(), + config.spec_config, + _kv_cache_quant_algo(config, llm_args), + ) + if target is not None: + _RESOLVED_MODES[target] = mode + return target return routing.route( pretrained_config, config.mapping, _device_compute_capability(), config.spec_config ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py new file mode 100644 index 000000000000..97bfeb2cf715 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py @@ -0,0 +1,360 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-head QK RMS norm followed by RoPE, in place on a packed qkv buffer.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* + +from .._op import Arch, Cell, OpWrapper + + +class FusedQkNormRope(OpWrapper): + """In place on `qkv`: RMS-norm the q and k heads, then rotate them. + Returns `qkv` -- the same buffer, mutated and handed back so a call site + can show what it touched without allocating. + + The v heads are not touched. `qkv` is one `[num_tokens, (Hq + Hk + Hv) * + head_dim]` buffer, which is the layout the fused projection already + produces, so nothing is split or copied. The norm is `rmsnorm(x) * w`, or + the Gemma form `rmsnorm(x) * (1 + w)` with `use_gemma`; it covers the whole + head, the rotation only its first `rotary_dim` dims. With `use_mrope`, + `position_ids` is `[3, T]` and rotary pair `j` takes its position from row 1 + when `j % 3 == 1 and j < 3 * mrope_section1`, from row 2 when + `j % 3 == 2 and j < 3 * mrope_section2`, from row 0 otherwise -- the + interleaved mRoPE layout. The caller owns the QKV projection before the + call and the KV-cache append and attention after it. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # gpt-oss-120b's default target asks for the RoPE half only: + # `is_qk_norm=False`, neox pairing, and YaRN-blended frequencies. 64 query + # heads, 8 key and 8 value, head_dim 64, all of which the cells carry + # because the op reads them as scalars rather than deriving them. Under + # attention TP each rank sees its own share of the heads. + _GPT_OSS = dict( + num_heads_q=64, + num_heads_k=8, + num_heads_v=8, + head_dim=64, + rotary_dim=64, + base=150000.0, + is_neox=True, + is_qk_norm=False, + # The checkpoint's YaRN correction range, from beta_fast=32 / + # beta_slow=1 over its original 4096-token context. + factor=32.0, + low=8.0, + high=18.0, + attention_factor=1.0, + ) + + # qwen4_exp (Qwen3.8-Flash-Next) runs Gemma QK norm and an unscaled rope + # on the first 64 dims (32 pairs), base 1e7, at two geometries. Its + # mrope_section is [11, 11, 10], passed as sections 11 / 11. Positions are + # drawn uniformly below `max_position` (a test axis, not an op argument): + # the kernel's fast sincos loses phase as `position * freq` grows, so the + # position range is part of what a cell certifies. + _QWEN_ROPE = dict( + rotary_dim=64, + base=1e7, + is_neox=True, + is_qk_norm=True, + use_gemma=True, + max_position=32768, + ) + _QWEN_MROPE = dict(use_mrope=True, mrope_section1=11, mrope_section2=11, **_QWEN_ROPE) + # Main attention (trunk and MTP): 24 query heads, 2 kv, head_dim 256. + _QWEN_MAIN = dict(num_heads_q=24, num_heads_k=2, num_heads_v=2, head_dim=256, **_QWEN_MROPE) + # The QSA index plane's query rope: 4 index heads plus the one k and one v + # slot of its scratch layout, over index_head_dim 128 (== 2 * rotary_dim, + # so the partial factor is not reapplied). + _QWEN_INDEX_HEADS = dict(num_heads_q=4, num_heads_k=1, num_heads_v=1, head_dim=128) + _QWEN_INDEX_MROPE = dict(**_QWEN_INDEX_HEADS, **_QWEN_MROPE) + _QWEN_INDEX_PLAIN = dict(**_QWEN_INDEX_HEADS, **_QWEN_ROPE) + + CELLS: tuple[Cell, ...] = ( + Cell(why="decode: one token", spec=dict(tokens=1, **_GPT_OSS)), + Cell(why="prefill-shaped batch", spec=dict(tokens=2048, **_GPT_OSS)), + Cell( + why="a token count that is not a multiple of any tile", + spec=dict(tokens=37, **_GPT_OSS), + ), + Cell( + why="default target at attention TP8: 8q/1k/1v heads per rank", + spec=dict( + tokens=64, **{**_GPT_OSS, "num_heads_q": 8, "num_heads_k": 1, "num_heads_v": 1} + ), + ), + Cell( + why="qwen4_exp main rope, decode: one token, Gemma norm and interleaved mRoPE " + "with three distinct position rows", + spec=dict(tokens=1, **_QWEN_MAIN), + ), + Cell( + why="qwen4_exp main rope, 19 tokens: the size its wiring and saturation tests run at", + spec=dict(tokens=19, **_QWEN_MAIN), + ), + Cell( + why="qwen4_exp main rope, 8192-token prefill: the Gemma one-ulp tail, under " + "GEMMA_PREFILL_ATOL", + spec=dict(tokens=8192, gemma_prefill_band=True, **_QWEN_MAIN), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 1 token: one k and one v head, the " + "narrowest layout; 5 warps of an 8-warp block", + spec=dict(tokens=1, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 2 tokens: the second token's heads " + "straddle a block boundary", + spec=dict(tokens=2, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 5 tokens: the last block runs one warp", + spec=dict(tokens=5, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 17 tokens: one past a power of two", + spec=dict(tokens=17, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 19 tokens: the size its wiring and " + "saturation tests run at", + spec=dict(tokens=19, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 33 tokens: one past a power of two", + spec=dict(tokens=33, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 65 tokens: one past a power of two", + spec=dict(tokens=65, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 8192-token prefill: the index site sees the " + "whole packed batch; under GEMMA_PREFILL_ATOL", + spec=dict(tokens=8192, gemma_prefill_band=True, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (plain), decode: the form the index plane runs " + "today, row 0 alone as a flat [T] position_ids", + spec=dict(tokens=1, **_QWEN_INDEX_PLAIN), + ), + Cell( + why="qwen4_exp index rope (plain), 19 tokens: an odd count across blocks", + spec=dict(tokens=19, **_QWEN_INDEX_PLAIN), + ), + Cell( + why="qwen4_exp index rope (plain), 8192-token prefill: under GEMMA_PREFILL_ATOL", + spec=dict(tokens=8192, gemma_prefill_band=True, **_QWEN_INDEX_PLAIN), + ), + ) + + note = """ + Certified by caller: gpt-oss runs `is_qk_norm=False` with YaRN; qwen4_exp + runs `is_qk_norm=True, use_gemma=True`, unscaled, with and without + interleaved mRoPE. Uncertified because no shipped target reaches them: + `is_qk_norm=True` with `use_gemma=False` (the plain-weight norm the op's + name advertises), `is_neox=False`, YaRN together with the norm or with + mRoPE. A target that wants one adds the cell rather than assuming the path + works because the argument exists. + + mRoPE preconditions. `position_ids` is a contiguous int32 `[3, T]`; the op + rejects a strided view loudly, so an engine's `[3, 1, max_num_tokens]` + position buffer has to be copied, not only reshaped. A HuggingFace + `mrope_section` `[s0, s1, s2]` summing to `rotary_dim / 2` is passed as + `mrope_section1 = s1`, `mrope_section2 = s2`; `s0` is implied. Only the + interleaved layout exists here -- the contiguous-section variant is a + different operator, measured ~500x out when wired in its place. Sections + saturate: at 32 pairs `mrope_section2` 10, 11 and 12 are bit-identical (the + largest `j % 3 == 2` below 32 is 29), while `mrope_section1` 10 moves pair + 31 back to row 0. Three equal rows reduce bitwise to the plain + `use_mrope=False` call on that row, so a text batch cannot see a wrong + split; an image span is where one shows. + + Metadata consumed: none. + """ + + def __call__( + self, + qkv: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + is_neox: bool, + position_ids: torch.Tensor, + factor: float = 1.0, + low: float = 0.0, + high: float = 0.0, + attention_factor: float = 1.0, + is_qk_norm: bool = True, + use_gemma: bool = False, + use_mrope: bool = False, + mrope_section1: int = 0, + mrope_section2: int = 0, + ) -> torch.Tensor: + """Writes `qkv` in place and returns it -- the op destroys the q/k + heads' prior contents, so a caller that wants them has to have kept + its own copy beforehand.""" + torch.ops.trtllm.fused_qk_norm_rope( + 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, + ) + return qkv + + def reference( + self, + qkv: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + is_neox: bool, + position_ids: torch.Tensor, + factor: float = 1.0, + low: float = 0.0, + high: float = 0.0, + attention_factor: float = 1.0, + is_qk_norm: bool = True, + use_gemma: bool = False, + use_mrope: bool = False, + mrope_section1: int = 0, + mrope_section2: int = 0, + ) -> torch.Tensor: + """The buffer the op writes back, fp32 throughout. + + Returns rather than mutates so a cell can drive both sides from the + same input; the in-place write is what `__call__` does. + """ + num_heads = num_heads_q + num_heads_k + num_heads_v + x = qkv.float().view(-1, num_heads, head_dim).clone() + half = rotary_dim // 2 + inv_freq = self._inv_freq(rotary_dim, base, factor, low, high, qkv.device) + + if use_mrope: + angle = position_ids.float()[:, :, None] * inv_freq # [3, T, half] + cos, sin = angle.cos(), angle.sin() + + def pick(c: torch.Tensor) -> torch.Tensor: + out = c[0].clone() + out[:, 1 : mrope_section1 * 3 : 3] = c[1][:, 1 : mrope_section1 * 3 : 3] + out[:, 2 : mrope_section2 * 3 : 3] = c[2][:, 2 : mrope_section2 * 3 : 3] + return out + + cos, sin = pick(cos), pick(sin) + else: + angle = position_ids.float()[:, None] * inv_freq # [T, half] + cos, sin = angle.cos(), angle.sin() + cos = cos[:, None, :] * attention_factor + sin = sin[:, None, :] * attention_factor + + for start, count, weight in ( + (0, num_heads_q, q_weight.float()), + (num_heads_q, num_heads_k, k_weight.float()), + ): + h = x[:, start : start + count, :] + if is_qk_norm: + h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + eps) + h = h * ((1.0 + weight) if use_gemma else weight) + r = h[..., :rotary_dim] + if is_neox: + x1, x2 = r[..., :half], r[..., half:] + rotated = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) + else: + x1, x2 = r[..., ::2], r[..., 1::2] + rotated = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten( + -2 + ) + x[:, start : start + count, :] = torch.cat([rotated, h[..., rotary_dim:]], -1) + return x.view(qkv.shape[0], -1).to(qkv.dtype) + + @staticmethod + def _inv_freq( + rotary_dim: int, base: float, factor: float, low: float, high: float, device + ) -> torch.Tensor: + """YaRN-blended inverse frequencies; plain RoPE when factor == 1.""" + j = torch.arange(rotary_dim // 2, dtype=torch.float32, device=device) + pos_freqs = base ** (2.0 * j / rotary_dim) + inv_extrapolation = 1.0 / pos_freqs + inv_interpolation = 1.0 / (factor * pos_freqs) + if high == low: + high = high + 0.001 # the kernel guards the ramp singularity the same way + ramp = ((j - low) / (high - low)).clamp(0.0, 1.0) + extrapolation = 1.0 - ramp + return inv_interpolation * (1.0 - extrapolation) + inv_extrapolation * extrapolation + + def is_valid( + self, + qkv: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + is_neox: bool, + position_ids: torch.Tensor, + *unused_args: object, + **unused_kwargs: object, + ) -> None: + # `eps`, `q_weight`, `k_weight`, `base` and `is_neox` are named though + # nothing here reads them: they mirror `__call__`, and gpt-oss passes + # them by keyword. + # + # The op reads the head counts as scalars and addresses qkv from + # data_ptr(); a width that disagrees with them is read as whatever the + # counts say and the heads silently come out shifted. + expected = (num_heads_q + num_heads_k + num_heads_v) * head_dim + assert qkv.shape[-1] == expected, ( + f"qkv's last dim must be (Hq + Hk + Hv) * head_dim = {expected}; got " + f"{qkv.shape[-1]}. The op trusts the counts and reads the buffer flat" + ) + assert qkv.is_contiguous(), "qkv must be contiguous; the op writes it through data_ptr()" + assert rotary_dim <= head_dim, ( + f"rotary_dim ({rotary_dim}) exceeds head_dim ({head_dim}); the rotation " + "would read past the end of each head" + ) + assert rotary_dim % 2 == 0, f"rotary_dim must be even to pair; got {rotary_dim}" + assert position_ids.shape[-1] == qkv.shape[0], ( + f"one position per token: {position_ids.shape[-1]} ids for {qkv.shape[0]} tokens" + ) + + +fused_qk_norm_rope = FusedQkNormRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather_cast.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather_cast.py new file mode 100644 index 000000000000..24d10fae21bc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather_cast.py @@ -0,0 +1,540 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Vocab-dim all-gather of bf16 logit shards into fp32 full-vocabulary rows, in one launch. + +``out = torch.cat(all_gather(shard), -1).float()``: a column-parallel LM head's verify logits as the sampler takes +them, without the NCCL all-gather, the concat copy and the cast. Every rank pushes its shard through a multicast +NVLink workspace and raises per-CTA ready flags (the epoch of the call), so no bit pattern is reserved as a +sentinel and the result is the stock path's bit for bit, -0.0 and NaN payloads included. + +``allgather_f32`` (MC-ALLGATHER): the same gather for fp32 shards, copied as is, ``out = torch.cat(all_gather(shard), +-1)``: the drafter's sharded draft logits as the rejection sampler takes them (fp32 [B, K, S] -> [B, K, world * S]), +without the NCCL all-gather and the concat copy. It shares the workspace (a slot holds MAX_ROWS / 2 fp32 rows). + +Compiled on the first call per shape, which must run outside CUDA-graph capture (``precompile``). The workspace is +collective on its first use: every rank of the TP group calls ``workspace_args`` at the same point. +""" + +from __future__ import annotations + +import importlib +import os +import threading +from typing import Callable, Dict, Optional, Sequence, Tuple + +import torch + +from .._op import Arch, Cell, OpWrapper + +MAX_ROWS = 64 # capacity: fixed: the workspace's rows per call (K3 verify rows at batch 16) +MAX_GRID = 152 # capacity: fixed: the workspace's flag columns per rank +# capacity: fixed: fp32 rows per call (an fp32 row is twice a bf16 row's bytes) +MAX_ROWS_FP32 = MAX_ROWS // 2 +# Launched without programmatic dependent launch: a plain launch waits for all prior stream work, so the previous +# call has completed whatever kernels run between two calls (the kernel's ordering condition holds by construction). +# It still releases its own dependents early. ``pdl=1`` on a call launches it with PDL (measurement only). +LAUNCH_PDL = False + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} +_workspaces: Dict[object, dict] = {} +_modules: Dict[str, object] = {} + + +def _kernel_module(): + """The CuTe DSL kernel beside this file, imported on first use (it pulls in cutlass).""" + mod = _modules.get("kernel") + if mod is None: + mod = _modules["kernel"] = importlib.import_module(".allgather_cast_kernel", __package__) + return mod + + +# CTAs by world size and rows: the fastest bitwise grid measured at each K3 verify row count on GB300 (graph-us at +# 96 / 128 / 152 CTAs). TP8 over two trays (RESULT_M31_N3A1K_MICROBENCH + its addendum-4 sweep): 4 rows 12.54 / 12.83 / +# 12.88, 8 rows 18.79 / 16.55 / 16.56, 16 rows 25.27 / 23.25 / 24.10, the other row counts fastest at 152. TP4 on one +# tray (RESULT_M31_N3A1K4_MICROBENCH): 4 rows 12.24 / 11.90 / 11.96, 8 rows 15.22 / 14.29 / 14.61, 16 rows 22.02 / +# 22.43 / 20.47, 32 rows 36.66 / 35.25 / 34.48 (96 is kept at 4 rows: within 0.5 us). Row counts in between take the +# next measured one's grid, rows past the table MAX_GRID; other world sizes take the TP8 table. A function of the +# world size, the rows and the device only, so every rank picks the same grid. +GRID_BY_ROWS = ((4, 96), (8, 128), (12, 152), (16, 128), (32, 152)) +GRID_BY_WORLD_ROWS = {4: ((4, 96), (8, 128), (32, 152)), 8: GRID_BY_ROWS} + + +def pick_grid(rows: int, device=None, world: int = 8) -> int: + """CTAs for a call of ``rows`` rows over ``world`` ranks (capped by the SM count and MAX_GRID).""" + sms = torch.cuda.get_device_properties( + device if device is not None else torch.cuda.current_device() + ).multi_processor_count + table = GRID_BY_WORLD_ROWS.get(world, GRID_BY_ROWS) + grid = next((g for bound, g in table if rows <= bound), MAX_GRID) + return min(MAX_GRID, sms, grid) + + +def workspace_args(mapping, shard_cols: int, group=None) -> dict: + """This TP group's workspace as the call's ``ws_*``, ``rank`` and ``world`` arguments. + + Three rotating buffers of [world][MAX_ROWS, shard_cols] bf16 and the [world][MAX_GRID] ready flags behind one + multicast mapping, plus the rank-local epoch words; built once per (mapping, shard_cols, group) and reused. + Collective on first use, outside CUDA-graph capture: every rank of the TP group calls it at the same point. + ``group`` overrides the communicator (a torch.distributed group when MPI is disabled). + """ + key = (mapping, shard_cols, id(group) if group is not None else None) + ws = _workspaces.get(key) + if ws is not None: + return ws["args"] + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "allgather_cast: the workspace must be allocated outside CUDA-graph capture" + ) + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + + kern = _kernel_module() + world = mapping.tp_size + comm = group if group is not None else _get_mnnvl_workspace_comm(mapping) + use_fabric_handle = ( + os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" or mapping.is_multi_node() + ) + error: Optional[Exception] = None + try: + words = kern.workspace_words(MAX_ROWS, shard_cols, world, MAX_GRID) + handle = _make_mnnvl_mcast_buffer(comm, words * 4, mapping, use_fabric_handle) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + with torch.inference_mode(): + uc.zero_() + flags = torch.zeros(kern.FLAG_WORDS, dtype=torch.int32, device=uc.device) + torch.cuda.synchronize() + args = dict(ws_uc=uc, ws_mc=mc, ws_flags=flags, rank=mapping.tp_rank, world=world) + ws = dict(handle=handle, comm=comm, args=args) + except Exception as exc: # noqa: BLE001 -- reported to every rank below, then re-raised + error = exc + # Also the barrier that keeps any rank from raising a peer's flag before the peer has zeroed it. + if not _mnnvl_workspace_all_succeeded(comm, error is None): + raise RuntimeError("allgather_cast workspace failed on at least one rank") from error + _workspaces[key] = ws + return ws["args"] + + +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=0) + + +def supports(shard: torch.Tensor, world: int) -> bool: + """Whether the kernel takes this shard: bf16 [rows <= MAX_ROWS, S] contiguous, S a multiple of 8, on CUDA.""" + return ( + shard.is_cuda + and shard.dim() == 2 + and shard.dtype == torch.bfloat16 + and shard.is_contiguous() + and 0 < shard.shape[0] <= MAX_ROWS + and shard.shape[1] % 8 == 0 + and shard.data_ptr() % 16 == 0 + and 1 < world <= 8 + ) + + +def supports_f32(shard: torch.Tensor, world: int) -> bool: + """Whether ``allgather_f32`` takes this shard: fp32 [rows <= MAX_ROWS_FP32, S] contiguous, S a multiple of 4.""" + return ( + shard.is_cuda + and shard.dim() == 2 + and shard.dtype == torch.float32 + and shard.is_contiguous() + and 0 < shard.shape[0] <= MAX_ROWS_FP32 + and shard.shape[1] % 4 == 0 + and shard.data_ptr() % 16 == 0 + and 1 < world <= 8 + ) + + +def _consts(rows: int, cols: int, world: int, grid: int, fp32: bool = False) -> tuple: + kern = _kernel_module() + return ( + grid, + rows, + cols, + world, + kern.slot_words(MAX_ROWS, cols), + kern.flag_base(MAX_ROWS, cols, world), + MAX_GRID, + bool(fp32), + ) + + +def _default_grid(rows: int, device, world: int, fp32: bool) -> int: + """An fp32 row moves a bf16 row's bytes twice: it takes the grid of twice the rows.""" + return pick_grid(2 * rows if fp32 else rows, device, world) + + +def _kernel_for( + shard, + out, + ws_uc, + ws_mc, + ws_flags, + rank: int, + world: int, + grid: int, + pdl: int, + fp32: bool = False, +): + """The compiled launcher of this call's shape (compiled here on first use, outside capture), its tensor + arguments and the stream.""" + import cuda.bindings.driver as cuda_driver + import cutlass.cute as cute + + kern = _kernel_module() + if not (supports_f32(shard, world) if fp32 else supports(shard, world)): + raise ValueError(f"allgather_cast: unsupported shard {tuple(shard.shape)} {shard.dtype}") + rows, cols = shard.shape + if grid <= 0: + grid = _default_grid(rows, shard.device, world, fp32) + if not 0 < grid <= MAX_GRID or ws_uc.numel() < kern.workspace_words( + MAX_ROWS, cols, world, MAX_GRID + ): + raise ValueError("allgather_cast: the call does not fit the workspace") + args = ( + _arg(shard.view(-1).view(torch.int32)), + _arg(out.view(-1).view(torch.int32)), + _arg(ws_uc.view(-1)), + _arg(ws_mc.view(-1)), + _arg(ws_flags.view(-1), align=4), + ) + use_pdl = LAUNCH_PDL if pdl < 0 else bool(pdl) + stream = cuda_driver.CUstream(torch.cuda.current_stream(shard.device).cuda_stream) + consts = _consts(rows, cols, world, grid, fp32) + key = consts + (use_pdl,) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "trtllm::allgather_cast must run once per shape outside CUDA-graph capture first" + ) + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kern.allgather_cast, *args, int(rank), *consts, use_pdl, stream + ) + return fn, args, stream + + +@torch.library.custom_op("trtllm::allgather_cast", mutates_args=()) +def _allgather_cast( + shard: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc: torch.Tensor, + ws_flags: torch.Tensor, + rank: int, + world: int, + grid: int, + pdl: int, +) -> torch.Tensor: + out = torch.empty( + shard.shape[0], world * shard.shape[1], dtype=torch.float32, device=shard.device + ) + fn, args, stream = _kernel_for(shard, out, ws_uc, ws_mc, ws_flags, rank, world, grid, pdl) + fn(*args, int(rank), stream) + return out + + +@_allgather_cast.register_fake +def _(shard, ws_uc, ws_mc, ws_flags, rank, world, grid, pdl): + return shard.new_empty((shard.shape[0], world * shard.shape[1]), dtype=torch.float32) + + +@torch.library.custom_op("trtllm::allgather_f32", mutates_args=()) +def _allgather_f32( + shard: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc: torch.Tensor, + ws_flags: torch.Tensor, + rank: int, + world: int, + grid: int, + pdl: int, +) -> torch.Tensor: + out = torch.empty( + shard.shape[0], world * shard.shape[1], dtype=torch.float32, device=shard.device + ) + fn, args, stream = _kernel_for( + shard, out, ws_uc, ws_mc, ws_flags, rank, world, grid, pdl, fp32=True + ) + fn(*args, int(rank), stream) + return out + + +@_allgather_f32.register_fake +def _(shard, ws_uc, ws_mc, ws_flags, rank, world, grid, pdl): + return shard.new_empty((shard.shape[0], world * shard.shape[1]), dtype=torch.float32) + + +def gather_last_dim_f32(logits: torch.Tensor, ws_uc, ws_mc, ws_flags, rank, world, grid=0, pdl=-1): + """``torch.cat(all_gather(logits), -1)`` for an fp32 [..., S] shard (the leading dims flattened into rows), or + None when ``allgather_f32`` does not take it (the caller then runs the stock all-gather; shapes decide, so every + rank declines alike).""" + lead = tuple(logits.shape[:-1]) + flat = logits.reshape(-1, logits.shape[-1]) + if not supports_f32(flat, world): + return None + out = torch.ops.trtllm.allgather_f32( + flat, ws_uc, ws_mc, ws_flags, int(rank), int(world), int(grid), int(pdl) + ) + return out.view(*lead, world * logits.shape[-1]) + + +def is_compiled( + rows: int, cols: int, world: int, grid: int = 0, device=None, fp32: bool = False +) -> bool: + """Whether the (rows, cols, world, grid) shape is compiled (for the caller's capture-time decision).""" + grid = grid if grid > 0 else _default_grid(rows, device, world, fp32) + return _consts(rows, cols, world, grid, fp32) + (LAUNCH_PDL,) in _compiled + + +def precompile( + mapping, + cols: int, + rows: Tuple[int, ...], + device, + group=None, + agree: Optional[Callable[[bool], bool]] = None, +) -> bool: + """Compile every row count in ``rows`` outside capture, then run each shape once (collective: every rank). + + The compiles launch nothing. ``agree`` (every rank's verdict, e.g. a TP all-reduce of success) is consulted + between the compiles and the first launch, so a compile that fails on one rank turns the item off on every rank + instead of leaving its peers polling for flags it never raises. Returns whether the shapes are ready.""" + ws = workspace_args(mapping, cols, group) + ok = True + try: + for r in rows: + shard = torch.zeros(r, cols, dtype=torch.bfloat16, device=device) + out = torch.empty(r, ws["world"] * cols, dtype=torch.float32, device=device) + _kernel_for( + shard, out, ws["ws_uc"], ws["ws_mc"], ws["ws_flags"], ws["rank"], ws["world"], 0, -1 + ) + except Exception as exc: # noqa: BLE001 -- reported to every rank through ``agree`` + from tensorrt_llm.logger import logger + + logger.warning(f"allgather_cast: compile failed on rank {ws['rank']}: {exc}") + ok = False + if agree is not None: + ok = agree(ok) + if not ok: + return False + for r in rows: + allgather_cast(torch.zeros(r, cols, dtype=torch.bfloat16, device=device), **ws) + torch.cuda.synchronize() + return True + + +def precompile_f32( + mapping, + cols: int, + rows: Tuple[int, ...], + device, + group=None, + agree: Optional[Callable[[bool], bool]] = None, +) -> bool: + """``precompile`` for ``allgather_f32`` (fp32 shards of ``cols`` columns, each row count in ``rows``).""" + ws = workspace_args(mapping, cols, group) + ok = True + try: + for r in rows: + shard = torch.zeros(r, cols, dtype=torch.float32, device=device) + out = torch.empty(r, ws["world"] * cols, dtype=torch.float32, device=device) + _kernel_for( + shard, out, ws["ws_uc"], ws["ws_mc"], ws["ws_flags"], ws["rank"], ws["world"], 0, -1, + fp32=True, + ) # fmt: skip + except Exception as exc: # noqa: BLE001 -- reported to every rank through ``agree`` + from tensorrt_llm.logger import logger + + logger.warning(f"allgather_f32: compile failed on rank {ws['rank']}: {exc}") + ok = False + if agree is not None: + ok = agree(ok) + if not ok: + return False + for r in rows: + allgather_f32(torch.zeros(r, cols, dtype=torch.float32, device=device), **ws) + torch.cuda.synchronize() + return True + + +class AllgatherCast(OpWrapper): + """fp32 [rows, world * S]: every rank's bf16 [rows, S] shard, in rank order along the last dim, widened to fp32. + + ``shard`` bf16 [rows, S] contiguous (rows <= MAX_ROWS, S a multiple of 8): this rank's vocabulary columns. + ``ws_*``, ``rank``, ``world``: ``workspace_args``'s. ``grid``: CTAs (0 picks). + ``pdl``: -1 launches as LAUNCH_PDL (off); 0 / 1 force it (measurement only). + + Every rank returns the same tensor. The call is a collective: every rank of the group calls it with the same + rows, S and grid, in the same order. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # MiniMax-M3.1's target LM head at TP8 over two GB300 trays: vocabulary 200064 (shard 25008), K3 verify rows + # B * (K + 1): K3 at batches 1..16, K7 at batches 1..8. + CELLS: tuple[Cell, ...] = ( + Cell( + why="M3.1 TP8, K3 verify at batch 1: 4 rows of the 25008-column shard", + spec=dict(world=8, rows=4, cols=25008), + ), + Cell( + why="M3.1 TP8, K7 verify at batch 1 (the headline c1): 8 rows", + spec=dict(world=8, rows=8, cols=25008), + ), + Cell( + why="M3.1 TP8, K7 verify at batch 2: 16 rows", + spec=dict(world=8, rows=16, cols=25008), + ), + Cell( + why="M3.1 TP8, K3 verify at batch 8: 32 rows", + spec=dict(world=8, rows=32, cols=25008), + ), + Cell( + why="M3.1 TP8, K3 verify at batch 16: 64 rows, the workspace's capacity", + spec=dict(world=8, rows=64, cols=25008), + ), + Cell( + why="a row count between the K3 batch sizes (odd vectors per CTA)", + spec=dict(world=8, rows=7, cols=25008), + ), + Cell( + why="one tray (TP4): the 50016-column shard", + spec=dict(world=4, rows=4, cols=50016), + ), + Cell( + why="one tray (TP4), K3 verify at batch 12: 48 rows", + spec=dict(world=4, rows=48, cols=50016), + ), + ) + + note = """ + Rows past MAX_ROWS (64) and shards that are not whole 16-byte vectors are declined (``supports``); the caller + takes the stock all-gather. Every rank must make the same decision, and does: it depends only on shapes. + """ + + def __call__( + self, + shard: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc: torch.Tensor, + ws_flags: torch.Tensor, + rank: int, + world: int, + grid: int = 0, + pdl: int = -1, + ) -> torch.Tensor: + return torch.ops.trtllm.allgather_cast( + shard, ws_uc, ws_mc, ws_flags, int(rank), int(world), int(grid), int(pdl) + ) + + def reference(self, shards_in_rank_order: Sequence[torch.Tensor]) -> torch.Tensor: + """The stock path: the column-parallel head's all-gather, the vocab-dim concat and ``.float()``. + + The documented collective exception to mirroring ``__call__``: it takes every rank's shard in rank order. + """ + return torch.cat(list(shards_in_rank_order), dim=-1).float() + + def is_valid( + self, + shard: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc: torch.Tensor, + ws_flags: torch.Tensor, + rank: int, + world: int, + grid: int = 0, + pdl: int = -1, + ) -> None: + assert supports(shard, world), (tuple(shard.shape), shard.dtype, shard.is_contiguous()) + assert 0 <= rank < world, (rank, world) + + +allgather_cast = AllgatherCast() + + +class AllgatherF32(OpWrapper): + """fp32 [rows, world * S]: every rank's fp32 [rows, S] shard, in rank order along the last dim, copied as is. + + ``shard`` fp32 [rows, S] contiguous (rows <= MAX_ROWS_FP32, S a multiple of 4): this rank's vocabulary columns. + ``ws_*``, ``rank``, ``world``: ``workspace_args``'s (the allgather_cast workspace of the same S). ``grid``: CTAs + (0 picks). ``pdl``: as ``allgather_cast``'s. A collective, as ``allgather_cast``. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # MiniMax-M3.1's DSpark drafter at TP8 over two GB300 trays: the rejection sampler's draft logits, B * K rows of + # the 25008-column vocabulary shard. + CELLS: tuple[Cell, ...] = ( + Cell( + why="M3.1 TP8, K7 draft logits at batch 1 (the headline c1): 7 rows", + spec=dict(world=8, rows=7, cols=25008), + ), + Cell( + why="M3.1 TP8, K7 draft logits at batch 2: 14 rows", + spec=dict(world=8, rows=14, cols=25008), + ), + Cell( + why="M3.1 TP8, K3 draft logits at batch 8: 24 rows", + spec=dict(world=8, rows=24, cols=25008), + ), + Cell( + why="the fp32 capacity: 32 rows", + spec=dict(world=8, rows=32, cols=25008), + ), + Cell( + why="one tray (TP4): the 50016-column shard, K7 at batch 1", + spec=dict(world=4, rows=7, cols=50016), + ), + ) + + note = """ + Rows past MAX_ROWS_FP32 (32) and shards that are not whole 16-byte vectors are declined (``supports_f32``); the + caller takes the stock all-gather. Every rank declines alike: it depends only on shapes. + """ + + def __call__( + self, + shard: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc: torch.Tensor, + ws_flags: torch.Tensor, + rank: int, + world: int, + grid: int = 0, + pdl: int = -1, + ) -> torch.Tensor: + return torch.ops.trtllm.allgather_f32( + shard, ws_uc, ws_mc, ws_flags, int(rank), int(world), int(grid), int(pdl) + ) + + def reference(self, shards_in_rank_order: Sequence[torch.Tensor]) -> torch.Tensor: + """The stock path: all_gather + the vocab-dim concat (the documented collective exception).""" + return torch.cat(list(shards_in_rank_order), dim=-1) + + def is_valid( + self, + shard: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc: torch.Tensor, + ws_flags: torch.Tensor, + rank: int, + world: int, + grid: int = 0, + pdl: int = -1, + ) -> None: + assert supports_f32(shard, world), (tuple(shard.shape), shard.dtype, shard.is_contiguous()) + assert 0 <= rank < world, (rank, world) + + +allgather_f32 = AllgatherF32() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather_cast_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather_cast_kernel.py new file mode 100644 index 000000000000..a77ad777283f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather_cast_kernel.py @@ -0,0 +1,327 @@ +# 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. +# ============================================================================= +# Vocab-dim all-gather of bf16 logit shards into fp32 full-vocabulary rows -- CTM (prims/cute) kernel +# ============================================================================= +# +# Every rank r holds shard_r, bf16 [rows, S] (its vocabulary columns [r S, r S + S)). Every rank writes +# out[i, r S + j] = float(shard_r[i, j]) for every rank r, row i < rows, column j < S +# i.e. torch.cat(all_gather(shard), -1).float(): the gather, the concat and the cast in one launch. The cast is a bit +# shift (bf16 is the high half of fp32), so out holds every rank's bits exactly: -0.0, NaN payloads and infinities +# included. +# +# Geometry: G CTAs of 256 threads. The shard is N = rows S / 8 16-byte vectors (S a multiple of 8); CTA c owns +# vectors [c N / G, (c + 1) N / G) of every rank's shard. CTA c pushes its vectors of this rank's shard into slot +# [buffer][rank] of every rank's workspace (multimem stores through the multicast mapping), then raises its ready +# flag [rank][c] in every rank's workspace (one multimem max, release, system scope). It converts its vectors of this +# rank's own shard straight from the input meanwhile, then polls the flags [r][c] of the other ranks in its own copy +# (acquire, system scope) and converts their vectors from its own buffer. +# +# fp32 mode (MC-ALLGATHER, the draft-logits gather): the shard is fp32 [rows, S] (S a multiple of 4), the 16-byte +# vectors hold 4 fp32, and every vector is stored as is (no widening): out = torch.cat(all_gather(shard), -1), bit for +# bit. A slot then holds max_rows / 2 rows. +# +# Flags carry the call's epoch, not the data: there is no empty-word sentinel, so no bit pattern of the logits is +# reserved. A flag only grows (multimem max of the epoch), so it needs no re-arming, and calls with different grids +# share it: a reader waits until flag >= its epoch. +# +# ORDERING: every call must start after the previous call on its rank completed. The op launches the kernel without +# programmatic dependent launch, so a launch waits for all prior work on the stream, whatever kernels run between two +# calls (a PDL launch would rely on every kernel in between executing griddepcontrol.wait). The kernel still releases +# its own dependents early. +# +# Buffers: 3 rotating buffers, call n using buffer n % 3. Rank r's call n + 3 reuses the buffer that rank q reads in +# call n only after r's call n + 2 completed, which waited for q's pushes of call n + 2, which q issues after its call +# n + 1 (and so n) completed. +# +# flags (rank-local int32): [0] epoch of this call, [1] CTAs arrived, [4] debug code, [5] debug detail. The last CTA to +# arrive advances the epoch. A flag poll that waits longer than POLL_TIMEOUT_NS (a peer that never raises the column, +# e.g. a grid that differs between ranks) writes the debug words and traps instead of hanging. +# ============================================================================= +"""CTM vocab-dim all-gather + concat + fp32 cast of bf16 logit shards over a multicast NVLink workspace, in one +launch.""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +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 + +THREADS = 256 +BUFFERS = 3 +FLAG_WORDS = 8 +POLL_TIMEOUT_NS = 30_000_000_000 # 30 s: far past any straggler, short of a hung job +DEBUG_POLL = 2 +VEC_ELEMS = 8 # bf16 per 16-byte vector +VEC_ELEMS_FP32 = 4 # fp32 per 16-byte vector (fp32 mode) +HI_MASK = -65536 # 0xFFFF0000 + + +def vectors(rows: int, shard: int) -> int: + """16-byte vectors of one rank's [rows, shard] bf16 slice.""" + return rows * shard // VEC_ELEMS + + +def slot_words(max_rows: int, shard: int) -> int: + """Int32 words of one rank's slot in one buffer.""" + return max_rows * shard // 2 + + +def workspace_words(max_rows: int, shard: int, world: int, max_grid: int) -> int: + """Int32 words of the multicast workspace: the buffers, then the [world][max_grid] ready flags.""" + return BUFFERS * world * slot_words(max_rows, shard) + world * max_grid + + +def flag_base(max_rows: int, shard: int, world: int) -> int: + return BUFFERS * world * slot_words(max_rows, shard) + + +@dsl_user_op +def _atomic_add_acq_rel(addr_i64, val, *, loc=None, ip=None): + """atom.acq_rel.gpu.global.add.u32, returning the old value.""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [addr_i64.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + "atom.acq_rel.gpu.global.add.u32 $0, [$1], $2;", "=r,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 _multimem_st_v4(addr_i64, a, b, c, d, *, loc=None, ip=None): + """multimem.st.global.v4.f32: one 16-byte store into every rank's copy (the bits, as stored).""" + _llvm.inline_asm( + None, + [addr_i64.ir_value(loc=loc, ip=ip), a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip), + c.ir_value(loc=loc, ip=ip), d.ir_value(loc=loc, ip=ip)], + "multimem.st.global.v4.f32 [$0], {$1, $2, $3, $4};", "l,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 _multimem_max_release_sys(addr_i64, val, *, loc=None, ip=None): + """multimem.red.release.sys.global.max.u32: raise every rank's copy of a flag to val, ordered after the CTA's + stores (the caller's barrier makes them the releasing thread's).""" + _llvm.inline_asm( + None, + [addr_i64.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + "multimem.red.release.sys.global.max.u32 [$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 _ld_acquire_sys(addr_i64, *, loc=None, ip=None): + """ld.acquire.sys.global.u32.""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [addr_i64.ir_value(loc=loc, ip=ip)], + "ld.acquire.sys.global.u32 $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 _globaltimer(*, loc=None, ip=None): + """%globaltimer (ns).""" + 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 _trap(*, loc=None, ip=None): + """fence.sc.sys (the debug words reach memory), then trap.""" + _llvm.inline_asm( + None, [], "fence.sc.sys;\n\ttrap;", "", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +def _store_fp32x8(out, word, got): + """The 8 bf16 of a 16-byte vector, widened to fp32 bit for bit, as two 16-byte stores at out word `word`.""" + lo = [] + for q in range(4): + w = cutlass.Int32(got[q]) + lo.append(cutlass.Int32(w << cutlass.Int32(16))) + lo.append(cutlass.Int32(w & cutlass.Int32(HI_MASK))) + out.store((lo[0], lo[1], lo[2], lo[3]), idx=word, alignment=16) + out.store((lo[4], lo[5], lo[6], lo[7]), idx=word + cutlass.Int32(4), alignment=16) + + +def _store_out(out, word, got, fp32: bool): + """One 16-byte input vector at out word `word`: 4 fp32 copied as is (fp32 mode), else 8 bf16 widened.""" + if cutlass.const_expr(fp32): + out.store( + ( + cutlass.Int32(got[0]), + cutlass.Int32(got[1]), + cutlass.Int32(got[2]), + cutlass.Int32(got[3]), + ), + idx=word, + alignment=16, + ) + else: + _store_fp32x8(out, word, got) + + +@cute.kernel +def allgather_cast_kernel( + shard: cutlass.Array, # int32 words of bf16 [rows, S]: this rank's shard + out: cutlass.Array, # int32 bits of fp32 [rows, world * S], out + buf_uc: cutlass.Array, # int32 words: this rank's workspace + buf_mc: cutlass.Array, # int32 words: its multicast mapping + flags: cutlass.Array, # int32 [FLAG_WORDS], rank-local + rank: cutlass.Int32, + grid: cutlass.Constexpr[int], + rows: cutlass.Constexpr[int], + shard_cols: cutlass.Constexpr[int], # S + world: cutlass.Constexpr[int], + slot: cutlass.Constexpr[int], # int32 words per (buffer, rank) slot + flag_off: cutlass.Constexpr[int], # first flag word of the workspace + max_grid: cutlass.Constexpr[int], # flag columns per rank + fp32: cutlass.Constexpr[bool], # fp32 shard copied as is (else bf16 widened to fp32) +): + ve = VEC_ELEMS_FP32 if fp32 else VEC_ELEMS + n_vec = rows * shard_cols // ve + row_vecs = shard_cols // ve + out_row = world * shard_cols + + tx, _, _ = cute.arch.thread_idx() + cta, _, _ = cute.arch.block_idx() + s_epoch = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + + # The shard is the previous kernel's output and the epoch the previous call's: both after the grid wait. + prims.griddepcontrol(prims.GridDepAction.WAIT) + if tx == 0: + s_epoch.store(flags.load(idx=0, is_volatile=True), idx=0) + prims.barrier_cta_sync(0) + epoch = s_epoch.load(idx=0) + buf_base = (epoch % cutlass.Int32(BUFFERS)) * cutlass.Int32(world * slot) + + lo = cta * cutlass.Int32(n_vec) // cutlass.Int32(grid) + hi = (cta + cutlass.Int32(1)) * cutlass.Int32(n_vec) // cutlass.Int32(grid) + + # Push this CTA's vectors of this rank's shard into every rank's slot [buffer][rank]. + mine = buf_base + rank * cutlass.Int32(slot) + v = lo + tx + while v < hi: + got = shard.load(idx=v * cutlass.Int32(4), vector_size=4, alignment=16) + _multimem_st_v4( + buf_mc.data_ptr(mine + v * cutlass.Int32(4)).toint(), + cutlass.Int32(got[0]), cutlass.Int32(got[1]), cutlass.Int32(got[2]), cutlass.Int32(got[3]), + ) # fmt: skip + v += cutlass.Int32(THREADS) + prims.barrier_cta_sync(0) + target = epoch + cutlass.Int32(1) + if tx == 0: + _multimem_max_release_sys( + buf_mc.data_ptr(cutlass.Int32(flag_off) + rank * cutlass.Int32(max_grid) + cta).toint(), + target, + ) + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + + # This rank's own columns, straight from the input while the peers' pushes are in flight. + v = lo + tx + while v < hi: + got = shard.load(idx=v * cutlass.Int32(4), vector_size=4, alignment=16) + row = v // cutlass.Int32(row_vecs) + col = v - row * cutlass.Int32(row_vecs) + _store_out( + out, + row * cutlass.Int32(out_row) + + rank * cutlass.Int32(shard_cols) + + col * cutlass.Int32(ve), + got, + fp32, + ) + v += cutlass.Int32(THREADS) + + # Thread r waits for rank r's flag [r][cta]; the barrier hands its acquire to the whole CTA. + if tx < cutlass.Int32(world): + if tx != rank: + addr = buf_uc.data_ptr( + cutlass.Int32(flag_off) + tx * cutlass.Int32(max_grid) + cta + ).toint() + start = _globaltimer() + while _ld_acquire_sys(addr) - target < cutlass.Int32(0): + if _globaltimer() - start > cutlass.Int64(POLL_TIMEOUT_NS): + flags.store(cutlass.Int32(DEBUG_POLL), idx=4) + flags.store(tx * cutlass.Int32(65536) + cta, idx=5) + _trap() + prims.barrier_cta_sync(0) + + for r in cutlass.range_constexpr(world): + if cutlass.Int32(r) != rank: + src = buf_base + cutlass.Int32(r * slot) + v = lo + tx + while v < hi: + got = buf_uc.load( + idx=src + v * cutlass.Int32(4), vector_size=4, alignment=16, is_volatile=True + ) + row = v // cutlass.Int32(row_vecs) + col = v - row * cutlass.Int32(row_vecs) + _store_out( + out, row * cutlass.Int32(out_row) + cutlass.Int32(r * shard_cols) + col * cutlass.Int32(ve), + got, fp32, + ) # fmt: skip + v += cutlass.Int32(THREADS) + + # Every CTA read the epoch before arriving; the last one advances it for the next call. + if tx == 0: + arrived = _atomic_add_acq_rel(flags.data_ptr(1).toint(), cutlass.Int32(1)) + if arrived == cutlass.Int32(grid - 1): + flags.store(cutlass.Int32(0), idx=1) + flags.store(target, idx=0) + + +@cute.jit +def allgather_cast( + shard: cute.Tensor, + out: cute.Tensor, + buf_uc: cute.Tensor, + buf_mc: cute.Tensor, + flags: cute.Tensor, + rank: cutlass.Int32, + grid: cutlass.Constexpr[int], + rows: cutlass.Constexpr[int], + shard_cols: cutlass.Constexpr[int], + world: cutlass.Constexpr[int], + slot: cutlass.Constexpr[int], + flag_off: cutlass.Constexpr[int], + max_grid: cutlass.Constexpr[int], + fp32: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + allgather_cast_kernel( + shard, out, buf_uc, buf_mc, flags, rank, grid, rows, shard_cols, world, slot, flag_off, max_grid, fp32, + ).launch( + grid=[grid, 1, 1], + block=[THREADS, 1, 1], + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/fakes/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/fakes/__init__.py new file mode 100644 index 000000000000..4cadfaf48c99 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/fakes/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fake (meta) impls of model-specific C++ torch ops, registered for torch.compile / fake-tensor tracing.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/fakes/m31_fake_ops.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/fakes/m31_fake_ops.py new file mode 100644 index 000000000000..96c98d1b1afd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/fakes/m31_fake_ops.py @@ -0,0 +1,161 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fake (meta) impls of the MiniMax-M3.1 C++ torch ops (cpp/tensorrt_llm/thop/minimaxM31/). + +The ops exist only on builds that carry them, either in the main library or in an overlay library the M3.1 targets load +on first use, so each fake is registered only once its op exists: ``register()`` runs at import and again after a +target loads an overlay library. It is idempotent. +""" + +import torch + +_REGISTERED = set() + + +def _ctx_residual_rmsnorm_mxfp8_fake( + inter, residual, weight, eps, write_norm, fp4_global_scale=None, enable_pdl=False +): + """minimaxM31Sm103CtxNormMxfp8Op.cpp: (norm_out, MXFP8 rows, swizzled scale bytes, NVFP4 rows, linear NVFP4 scale + bytes); inter is updated in place. Outputs that are off are empty 1-D tensors of inter's dtype.""" + hidden = inter.shape[-1] + rows = inter.numel() // hidden + empty = inter.new_empty((0,)) + norm_out = torch.empty_like(inter) if write_norm else empty + mxfp8 = torch.empty_like(inter, dtype=torch.float8_e4m3fn) + sf = inter.new_empty( + ((rows + 127) // 128 * 128 * ((hidden // 32 + 3) // 4 * 4),), dtype=torch.uint8 + ) + if fp4_global_scale is None: + return (norm_out, mxfp8, sf, empty, empty) + fp4 = inter.new_empty((*inter.shape[:-1], hidden // 2), dtype=torch.uint8) + fp4_sf = inter.new_empty((rows * (hidden // 16),), dtype=torch.uint8) + return (norm_out, mxfp8, sf, fp4, fp4_sf) + + +def _select_prefill_tile_fake(scores, n_valid_blocks, out, init_blocks, local_blocks): + """minimaxM31Sm103SelectPrefillTileOp.cpp: writes the block table into out; returns nothing.""" + return None + + +def _sparse_decode_lc_fake( + 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, + btsmem, + sync3, + one_div=False, +): + """minimaxM31Tp4SparseDecodeLcOp.cpp (LC4; LongAttn a7154ae3fa): writes out (and q8 / sf when given) in place; + returns nothing.""" + return None + + +def _sparse_decode_lc8_fake( + 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, + one_div=False, +): + """minimaxM31Tp8SparseDecodeLc8Op.cpp (LC8; LongAttn a7154ae3fa): writes out (and q8 / sf when given) in place; + returns nothing.""" + return None + + +def _topp_renorm_det_v3_fake(probs, top_p, workspace, out, ctas_per_row, pdl, candidate_list): + """m31_topp_det v3: as minimax_m31_topp_renorm_det (in place into out, workspace updated); returns nothing. + minimax_m31_topp_renorm_det_v3_ws_bytes takes no tensor and needs no fake.""" + return None + + +def _topp_renorm_det_fake(probs, top_p, workspace, out, ctas_per_row): + """m31_topp_det (ST-01): writes the renormed [rows, d] fp32 probs into out in place (and updates the persistent + workspace); returns nothing. minimax_m31_topp_renorm_det_ws_bytes takes no tensor and needs no fake.""" + return None + + +def _softmax_mc_fake(logits, temperature, temperature_val, partials, out, write_ctas): + """m31_softmax_mc (ST-01-SM): writes the [rows, d] fp32 probs into out in place (and overwrites the partials + workspace); returns nothing.""" + return None + + +def _chain_spec_sampling_mc_fake( + draft_probs, draft_token_ids, target_probs, seed, offset, workspace, ctas_per_row +): + """m31_chain_spec_mc (ST-01-REJ): (output_token_ids [batch, numDraft + 1], num_tokens [batch]), int32 as + trtllm::chain_speculative_sampling; updates the persistent workspace in place. + minimax_m31_chain_speculative_sampling_mc_ws_bytes takes no tensor and needs no fake.""" + batch_size, num_draft = draft_token_ids.shape + return ( + draft_token_ids.new_empty((batch_size, num_draft + 1)), + draft_token_ids.new_empty((batch_size,)), + ) + + +# op name -> fake impl (schemas as of img9, ONE_DIV included). +FAKES = { + "minimax_m31_ctx_residual_rmsnorm_mxfp8": _ctx_residual_rmsnorm_mxfp8_fake, + "minimax_m31_select_prefill_tile": _select_prefill_tile_fake, + "minimax_m31_tp4_sparse_decode_lc": _sparse_decode_lc_fake, + "minimax_m31_tp8_sparse_decode_lc8": _sparse_decode_lc8_fake, + "minimax_m31_topp_renorm_det": _topp_renorm_det_fake, + "minimax_m31_topp_renorm_det_v3": _topp_renorm_det_v3_fake, + "minimax_m31_softmax_mc": _softmax_mc_fake, + "minimax_m31_chain_speculative_sampling_mc": _chain_spec_sampling_mc_fake, +} + + +def _has_fake(qualname): + """True when the op already has a fake: a C++ Meta kernel or a torch.library fake (register_fake raises on both). + simple_registry is private torch API (checked on the image's torch 2.14); torch.library has no public query.""" + if torch._C._dispatch_has_kernel_for_dispatch_key(qualname, "Meta"): + return True + from torch._library.simple_registry import singleton + + return singleton.find(qualname).fake_impl.kernel is not None + + +def register(): + """Register the fake of every M3.1 op this process has loaded and that has none yet (another module may have + registered one first; that one is kept).""" + for name, fake in FAKES.items(): + if name in _REGISTERED or not hasattr(torch.ops.trtllm, name): + continue + if not _has_fake(f"trtllm::{name}"): + torch.library.register_fake(f"trtllm::{name}")(fake) + _REGISTERED.add(name) + + +register() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/__init__.py new file mode 100644 index 000000000000..5e898b4e1168 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/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.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.py new file mode 100644 index 000000000000..978decf58a8a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.py @@ -0,0 +1,450 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4-weight / NVFP4-activation mixture-of-experts layer: caller-supplied +top-k + grouped FC1 GEMM + gated activation + NVFP4 requantization + grouped +FC2 GEMM + routing-weighted combine, in one trtllm-gen call.""" + +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 + +#: NVFP4 block size: one e4m3 scale per 16 elements along K. +SF_VEC_SIZE = 16 + +#: Largest e2m1 magnitude, and the largest finite e4m3 magnitude. Their product +#: is the global scale a checkpoint's `input_scale` encodes. +E2M1_MAX = 6.0 +E4M3_MAX = 448.0 + +#: `act_type` for SwiGlu, the only activation this entry certifies. +ACT_TYPE_SWIGLU = 0 + +#: The widest unchunked call a stock engine issues: trtllm's default +#: `max_num_tokens`. A caller whose own token count exceeds it must chunk. +MAX_CERTIFIED_TOKENS = 8192 + +#: The 16 e2m1 code points in code order: sign bit 3, exponent bits 2:1, +#: mantissa bit 0. +_E2M1_VALUES = ( + 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, +) + +#: One expert stack's `(codes, block scales)`, unshuffled and unswizzled. +Operand = Tuple[torch.Tensor, torch.Tensor] + + +def _e2m1_rne(a: torch.Tensor) -> torch.Tensor: + """Round to nearest even on the e2m1 grid, saturating at +-6. + + The grid step is 0.5 below 2, 1 below 4 and 2 above -- the binade step of + e2m1's three exponents, the subnormal binade sharing the first normal one's. + `torch.round` is banker's rounding, which is the format's ties-to-even-code + rule: 0.25 -> 0, 0.75 -> 1, 1.25 -> 1, 1.75 -> 2, 2.5 -> 2, 3.5 -> 4, + 5.0 -> 4. + """ + step = torch.where(a.abs() < 2.0, 0.5, torch.where(a.abs() < 4.0, 1.0, 2.0)) + return torch.sign(a) * torch.clamp(torch.round(a.abs() / step) * step, max=E2M1_MAX) + + +def quantize_nvfp4(x: torch.Tensor, global_scale: float) -> torch.Tensor: + """NVFP4 quantize-dequantize of `x` under `global_scale`. + + Per 16 consecutive columns: `sf = e4m3(g * blockmax / 6)` and + `data = e2m1(x * g / sf)`, returning `data * sf` -- the value a downstream + block-scaled MMA multiplies, which is `g` times the reconstruction of `x`. + This models the FC1 epilogue, which the entry's test pins bit-exactly. + """ + rows, cols = x.shape + blocks = x.reshape(rows, cols // SF_VEC_SIZE, SF_VEC_SIZE) + sf = ( + (global_scale * blocks.abs().amax(dim=-1, keepdim=True) / E2M1_MAX) + .clamp(max=E4M3_MAX) + .to(torch.float8_e4m3fn) + .float() + ) + out_scale = torch.where(sf == 0, torch.zeros_like(sf), global_scale / sf) + return (_e2m1_rne(blocks * out_scale) * sf).reshape(rows, cols) + + +def dequantize_nvfp4(codes: torch.Tensor, block_scales: torch.Tensor) -> torch.Tensor: + """One expert's `[N, K]` fp32 weight from its codes and e4m3 block scales.""" + table = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=codes.device) + return table[codes.long()] * block_scales.float().repeat_interleave(SF_VEC_SIZE, dim=1) + + +class Fp4BlockScaleMoeRunner(OpWrapper): + """One NVFP4-weight MoE layer over NVFP4 (e2m1 + e4m3 block-scale) + activations, entered pre-routed. + + The caller owns top-k: `topk_ids` carries global expert ids and + `topk_weights` the combine weights, neither renormalized here. This rank + answers for `[local_expert_offset, local_expert_offset + local_num_experts)` + and contributes zero for every other id, so an expert-parallel split is N + calls whose outputs sum -- which is what the deepseek default target does + under moe_ep, one call per rank over its window; qwen4_exp makes one call + over its whole expert stack. + + With `do_finalize=True` the result is a one-element list holding the + combined `[T, H]` bf16 rows, or an empty tensor when `output` was given. + With `do_finalize=False` it is three tensors: per-slot expert rows, an + unwritten scale buffer, and the expanded-index -> permuted-row map. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # DeepSeek-R1-0528's routed geometry. The deepseek default target makes one + # call per rank over its moe_ep window at its moe_tp slice of the + # intermediate (dep4 with moe ep4: a 64-wide window at the full 2048; tp8: + # all 256 experts at 256), so the window and intermediate-slice cells are + # the shapes it runs and the full stack is the single-rank one. + _R1 = dict(hidden=7168, intermediate=2048, num_experts=256, top_k=8) + + # qwen4_exp's (Qwen3.8-Flash-Next) routed geometry: many small experts, a + # top_k past 8, an intermediate width of 640, all 512 experts local. Its + # router renormalizes the top-10 softmax weights to sum to 1 before the + # call (`routing="renormalized"`, a test axis rather than an op argument; + # the R1 cells draw uniform [0, 1) weights). Nothing is padded: + # H % 256, 2I % 128, H/16 % 4 and I/16 % 4 are all zero. + _QWEN4_EXP = dict( + hidden=2560, + intermediate=640, + num_experts=512, + top_k=10, + local_expert_offset=0, + local_num_experts=512, + routing="renormalized", + ) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="dep4's own call: the 64-wide expert window at offset 0, R1 geometry", + spec=dict(tokens=64, local_expert_offset=0, local_num_experts=64, **_R1), + ), + Cell( + why="a window at a non-zero offset -- the ids this rank must answer for move with it", + spec=dict(tokens=64, local_expert_offset=192, local_num_experts=64, **_R1), + ), + Cell( + why="the single-rank shape: all 256 experts local, and one token, the decode row count", + spec=dict(tokens=1, local_expert_offset=0, local_num_experts=256, **_R1), + ), + Cell( + why="deepseek: all 256 experts, intermediate split 2 ways (1024)", + spec=dict( + tokens=64, + local_expert_offset=0, + local_num_experts=256, + **dict(_R1, intermediate=1024), + ), + ), + Cell( + why="deepseek default at tp4: all 256 experts, intermediate split 4 ways (512)", + spec=dict( + tokens=64, + local_expert_offset=0, + local_num_experts=256, + **dict(_R1, intermediate=512), + ), + ), + Cell( + why="deepseek default at tp8: all 256 experts, intermediate split 8 ways (256)", + spec=dict( + tokens=64, + local_expert_offset=0, + local_num_experts=256, + **dict(_R1, intermediate=256), + ), + ), + Cell( + why="deepseek default at ep8 (tp8 + ep8 and dep8): the last 32-wide window", + spec=dict(tokens=64, local_expert_offset=224, local_num_experts=32, **_R1), + ), + Cell( + why="deepseek default at dep16: a 16-wide window", + spec=dict(tokens=64, local_expert_offset=112, local_num_experts=16, **_R1), + ), + Cell( + why="qwen4_exp decode: one token over 512 experts, top-10 renormalized, H 2560 / I 640", + spec=dict(tokens=1, **_QWEN4_EXP), + ), + Cell( + why="qwen4_exp at 4 rows: the last count the pre-routed path's static-block " + "permutation kernel takes", + spec=dict(tokens=4, **_QWEN4_EXP), + ), + Cell( + why="qwen4_exp at 5 rows: the first count on the single-cluster permutation " + "kernel (the dynamic-block one stops at 256 experts)", + spec=dict(tokens=5, **_QWEN4_EXP), + ), + Cell( + why="qwen4_exp prefill: 8192 rows, the single-cluster kernel's cap and the top " + "of the certified column", + spec=dict(tokens=8192, **_QWEN4_EXP), + ), + ) + + note = """ + Certified pre-routed only: `routing_logits` is `None` at every cell, so the + kernel's own routing is not exercised and `routing_method_type`, `n_group`, + `topk_group` and `routed_scaling_factor` are inert -- they are passed + because the signature takes them, not because they do anything here. + + That holds at 512 experts / top-10 too, where it would matter if it were + read: which `(num_experts, top_k)` pairs the routing kernels are compiled + for is a per-policy tier list, and a pair with no tier is a logged error, + not a raise. The pre-routed path substitutes its own preprocess/postprocess + pair before dispatch, and `routing_method_type` 0, 2, 4, 5, 6 were measured + bitwise identical to 1 there. It picks its permutation kernel by row count + instead: a static-block kernel up to 4 rows, a dynamic-block one up to 16 + rows only at <= 256 experts, a single-cluster one up to 8192 rows. + + `act_type` 0 (SwiGlu) only. `gemm1_bias`, `gemm1_alpha`, `gemm1_beta`, + `gemm1_clamp_limit` and `gemm2_bias` are `None`. `hidden_states_scale` is + the linear (unswizzled) e4m3 block-scale layout, one scale per 16 elements + of each row. + + At 512 experts only the whole stack is certified: no expert-parallel window + of it, `do_finalize=False` or a caller-supplied `output` was measured there. + + `tune_max_num_tokens` is an autotuner bucket hint and was measured inert; + it is not a token count. `T` is `hidden_states.shape[0]`, and 8192 is the + top of the certified column. A warm autotuner cache was measured bitwise + inert against a cold one at the R1 stack, its windows and the 512-expert + stack; `local_num_experts` separates the tuner's cache entries, so each + layout warms its own. + + The caller owns the weight layout entirely: FC1 and FC2 weights are + pre-shuffled and their scales additionally 128x4 swizzled. Nothing here + checks that, and getting it wrong is 300-3400 bf16 ulp with no error -- the + entry's test builds the layout twice, once in torch and once through the + trtllm preprocessing ops, and compares them before either reaches the op. + """ + + def __call__( + self, + routing_logits: Optional[torch.Tensor], + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_bias: Optional[torch.Tensor], + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm2_bias: Optional[torch.Tensor], + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int, + do_finalize: bool, + act_type: int = 0, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + tune_max_num_tokens: int = 8192, + use_dp: bool = False, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.fp4_block_scale_moe_runner( + routing_logits, + routing_bias, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + gemm1_bias, + gemm1_alpha, + gemm1_beta, + gemm1_clamp_limit, + gemm2_weights, + gemm2_weights_scale, + gemm2_bias, + output1_scale_scalar, + output1_scale_gate_scalar, + output2_scale_scalar, + num_experts, + top_k, + n_group, + topk_group, + intermediate_size, + local_expert_offset, + local_num_experts, + routed_scaling_factor, + routing_method_type, + do_finalize, + act_type, + topk_weights=topk_weights, + topk_ids=topk_ids, + output=output, + tune_max_num_tokens=tune_max_num_tokens, + use_dp=use_dp, + ) + + def reference( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + up: Operand, + gate: Operand, + down: Operand, + fc2_global_scale: float, + local_expert_offset: int = 0, + local_num_experts: Optional[int] = None, + ) -> torch.Tensor: + """The layer, in plain torch over dequantized weights, fp32 throughout. + + The documented exception to mirroring: `__call__` takes the *prepared* + operands -- packed e2m1 codes, shuffled rows, swizzled scales -- and a + reference given those would be re-deriving the layout rather than the + computation. So this takes the per-expert `(codes, scales)` pairs the + layout was built from, which is where a target's checkpoint starts too, + plus the routing this call was entered with. + + `up`, `gate` and `down` are each `(codes[E, M, K/2], scales[E, M, K/16])` + in their unshuffled, unswizzled form. `topk_ids` carries *global* ids; + this rank contributes only for the ids its window owns, and every other + slot contributes zero -- which is what makes an expert-parallel split + summable. + + The FC1 activation is requantized to NVFP4 under `fc2_global_scale` + before FC2, because the kernel does: the FC1 epilogue emits `g2 * act` + as NVFP4 and FC2's alpha divides it back out. Skipping that step is a + reference that disagrees with a correct kernel by ~23 ulp RMS, so it is + not an optional refinement. + """ + num_tokens, hidden = hidden_states.shape + up_c, up_s = up + gt_c, gt_s = gate + dn_c, dn_s = down + local_num_experts = up_c.shape[0] if local_num_experts is None else local_num_experts + out = torch.zeros(num_tokens, hidden, dtype=torch.float32, device=hidden_states.device) + for local_e in range(local_num_experts): + tok, slot = (topk_ids == local_expert_offset + local_e).nonzero(as_tuple=True) + if tok.numel() == 0: + continue + xe = hidden_states[tok] + u = xe @ dequantize_nvfp4(up_c[local_e], up_s[local_e]).t() + g = xe @ dequantize_nvfp4(gt_c[local_e], gt_s[local_e]).t() + act = u * g * torch.sigmoid(g) + act = quantize_nvfp4(act, fc2_global_scale) / fc2_global_scale + y = act @ dequantize_nvfp4(dn_c[local_e], dn_s[local_e]).t() + out.index_add_(0, tok, y * topk_weights[tok, slot].float().unsqueeze(1)) + return out + + def is_valid( + self, + routing_logits: Optional[torch.Tensor], + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_bias: Optional[torch.Tensor], + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm2_bias: Optional[torch.Tensor], + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int, + do_finalize: bool, + act_type: int = 0, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + tune_max_num_tokens: int = 8192, + use_dp: bool = False, + ) -> None: + """The two ways this op takes an input and answers it wrongly. + + Everything else a caller can get wrong here the op rejects loudly: a + mismatched `intermediate_size`, a `top_k` outside `(0, num_experts)`, a + hidden width that disagrees between the two weight operands. Those are + not restated -- swapping their errors for ours would only make them + worse. + """ + # The kernel takes raw data pointers and assumes a dense row-major + # layout for every tensor. A strided view is accepted without complaint + # and silently reads -- or writes -- the wrong elements. Observed on + # this machine for hidden_states, topk_weights and the weight operands. + for name, tensor in ( + ("routing_logits", routing_logits), + ("routing_bias", routing_bias), + ("hidden_states", hidden_states), + ("hidden_states_scale", hidden_states_scale), + ("gemm1_weights", gemm1_weights), + ("gemm1_weights_scale", gemm1_weights_scale), + ("gemm1_bias", gemm1_bias), + ("gemm1_alpha", gemm1_alpha), + ("gemm1_beta", gemm1_beta), + ("gemm1_clamp_limit", gemm1_clamp_limit), + ("gemm2_weights", gemm2_weights), + ("gemm2_weights_scale", gemm2_weights_scale), + ("gemm2_bias", gemm2_bias), + ("output1_scale_scalar", output1_scale_scalar), + ("output1_scale_gate_scalar", output1_scale_gate_scalar), + ("output2_scale_scalar", output2_scale_scalar), + ("topk_weights", topk_weights), + ("topk_ids", topk_ids), + ("output", output), + ): + if tensor is not None: + assert tensor.is_contiguous(), ( + f"{name} must be contiguous; a strided view is read as if dense " + "and silently produces wrong results" + ) + + # With the block-scale layout this entry contracts -- row shuffle plus + # 128x4 scale swizzle -- a hidden size that is only a multiple of 128 + # makes the kernel read the FC1 weight scales for the wrong blocks. + # Measured wrong at hidden 384 / 640 / 896: no error, 300-3400 bf16 ulp. + assert gemm1_weights.shape[-1] * 2 % 256 == 0, ( + "hidden size (gemm1_weights.shape[-1] * 2) must be a multiple of 256; " + "the prepared block-scale layout is silently misread otherwise" + ) + + +fp4_block_scale_moe_runner = Fp4BlockScaleMoeRunner() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/__init__.py new file mode 100644 index 000000000000..7d0c28eefd77 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Quantization entries.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py new file mode 100644 index 000000000000..58e9ee52f2df --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 quantization: packed e2m1 data plus one e4m3 scale per 16 elements.""" + +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 + +_E4M3_MAX = 448.0 +_E2M1_MAX = 6.0 + +#: Midpoints between consecutive e2m1 magnitudes. A value landing exactly on one +#: is a rounding tie; `_TIE_UP[i]` is True when rounding *up* yields the even +#: code, which is the direction round-to-nearest-even takes. +_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 its per-block output scale through two `rcp.approx.ftz.f32` +#: reciprocals, each ~2^-23 relative. A value within a few fp32 ulps of a +#: midpoint can therefore round to either neighbour, and no reference can say +#: which. Everything further out than this must match bit for bit. Largest +#: deviation observed on this machine: 7.9e-8 relative, ~12x inside the window. +TIE_WINDOW = 2.0**-20 + + +def _pad_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +class Fp4Quantize(OpWrapper): + """Quantize to FP4: `(packed e2m1 data [..., K/2] uint8, scale bytes, 1-D)`. + + One scale per `sf_vec_size` contiguous elements of the last dim. NVFP4 -- + the only form any shipped target uses -- is `sf_vec_size=16` with + `sf_use_ue8m0=False`, which makes the scale an e4m3 byte derived from the + block maximum and the caller's `global_scale`. + + `is_sf_swizzled_layout` selects the 128x4 swizzled scale order (True) or the + row-major linear order (False); the buffer is 1-D either way, so the layout + is something a consumer has to be told, not something its shape shows. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # deepseek-r1 reaches this op from both layouts: the swizzled one feeding + # nvfp4_gemm in the dense and shared-expert MLPs, the linear one feeding the + # routed-expert path, at its hidden size and its intermediate. qwen3.8 + # (qwen4_exp) calls it only for its routed experts: linear scales at hidden + # 2560. A linear buffer has no row padding, so at a row count that is not a + # multiple of 128 it is shorter than the swizzled one; at a multiple of 128 + # the two are the same length and only the cells' byte order tells them + # apart. + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode, hidden 7168, swizzled scales -- the dense MLP's input", + spec=dict(shape=(1, 7168), dtype=torch.bfloat16, vec=16, swizzled=True), + ), + Cell( + why="prefill chunk, hidden 7168, swizzled", + spec=dict(shape=(2048, 7168), dtype=torch.bfloat16, vec=16, swizzled=True), + ), + Cell( + why="linear scales -- the routed-expert path asks for this order", + spec=dict(shape=(2048, 7168), dtype=torch.bfloat16, vec=16, swizzled=False), + ), + Cell( + why="the activation width, 2 x moe_intermediate, between the two GEMMs", + spec=dict(shape=(512, 4096), dtype=torch.bfloat16, vec=16, swizzled=True), + ), + Cell( + why="a row count that is not a multiple of the swizzle's 128-row tile", + spec=dict(shape=(37, 7168), dtype=torch.bfloat16, vec=16, swizzled=True), + ), + Cell( + why="deepseek default at attention TP8: the shared expert's down input, 2048/8 wide", + spec=dict(shape=(16, 256), dtype=torch.bfloat16, vec=16, swizzled=True), + ), + Cell( + why="qwen3.8 decode: 1 token, hidden 2560, linear -- its routed-expert input", + spec=dict(shape=(1, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), + ), + Cell( + why="qwen3.8, 7 rows at hidden 2560, linear: several rows, no row padding", + spec=dict(shape=(7, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), + ), + Cell( + why="qwen3.8 prefill at max_num_tokens 8192, hidden 2560, linear", + spec=dict(shape=(8192, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), + ), + ) + + note = """ + `sf_use_ue8m0=True` selects power-of-two scales instead of e4m3. No shipped + target uses it and this entry does not certify it. + """ + + def __call__( + self, + input: torch.Tensor, + global_scale: Optional[torch.Tensor], + sf_vec_size: int, + sf_use_ue8m0: bool = False, + is_sf_swizzled_layout: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.trtllm.fp4_quantize( + input, global_scale, sf_vec_size, sf_use_ue8m0, is_sf_swizzled_layout + ) + + def reference( + self, + input: torch.Tensor, + global_scale: Optional[torch.Tensor], + sf_vec_size: int, + sf_use_ue8m0: bool = False, + is_sf_swizzled_layout: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Native-torch NVFP4, with the scales in the requested layout. + + Per block: `vecmax = max|x|`, the scale is `e4m3(global_scale * vecmax / + 6)`, and the data is the input times `global_scale / scale` rounded to + e2m1. torch's fp32 -> e4m3 cast emits NaN above 464 where the kernel's + saturates, so the scale is clamped first -- between 448 and 464 both + round to 448 anyway. + """ + packed, sf, _ = self._quantize(input, global_scale, sf_vec_size) + if not is_sf_swizzled_layout: + return packed, sf.reshape(-1) + return packed, self._swizzle(sf) + + def scaled( + self, input: torch.Tensor, global_scale: Optional[torch.Tensor], sf_vec_size: int + ) -> torch.Tensor: + """The pre-rounding product, which is what `near_tie` reads. + + Exposed because a caller cannot reconstruct it from the op's outputs: + rounding is the step that loses it. + """ + return self._quantize(input, global_scale, sf_vec_size)[2] + + @staticmethod + def near_tie(scaled: torch.Tensor) -> torch.Tensor: + """Elements whose rounding this entry refuses to predict. + + Inside `TIE_WINDOW` of an e2m1 midpoint the kernel's approximate + reciprocals decide the direction, so a code that differs there is not a + defect. Outside it, a difference is. + """ + 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 is_valid( + self, + input: torch.Tensor, + global_scale: Optional[torch.Tensor], + sf_vec_size: int, + sf_use_ue8m0: bool = False, + is_sf_swizzled_layout: bool = True, + ) -> None: + # The kernel loads one scalar from global_scale and ignores every + # element past the first, so a per-token [num_tokens] tensor is silently + # applied as global_scale[0] to every row -- observed to return a + # valid-looking result, never to raise. + assert global_scale is None or global_scale.numel() == 1, ( + "global_scale must hold exactly one element; extra elements are " + "silently ignored and global_scale[0] is applied to every row" + ) + assert input.shape[-1] % sf_vec_size == 0, ( + f"the quantized dim must be a whole number of {sf_vec_size}-element " + f"blocks; got {input.shape[-1]}" + ) + assert input.is_contiguous(), "input must be contiguous; the kernel reads it flat" + + # ── the arithmetic, shared by reference() and scaled() ────────────────── + + def _quantize( + self, input: torch.Tensor, global_scale: Optional[torch.Tensor], vec: int + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + gs = 1.0 if global_scale is None else global_scale.float().reshape(()) + k = input.shape[-1] + rows = input.numel() // k + flat = input.reshape(rows, k) + vecmax = flat.abs().reshape(rows, k // vec, vec).amax(-1).float() + 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 = (flat.float().reshape(rows, k // vec, vec) * out_scale.unsqueeze(-1)).reshape( + rows, k + ) + codes = self._e2m1_codes(scaled) + packed = codes[:, 0::2] | (codes[:, 1::2] << 4) + return packed, sf.view(torch.uint8), scaled + + @staticmethod + def _e2m1_codes(v: torch.Tensor) -> torch.Tensor: + """fp32 -> 4-bit e2m1 codes (bit 3 = sign), round-to-nearest-even.""" + mid = torch.tensor(_E2M1_MIDPOINTS, device=v.device, dtype=torch.float32) + tie_up = torch.tensor(_TIE_UP, device=v.device, dtype=torch.bool) + a = v.abs().unsqueeze(-1) + greater = (a > mid).sum(-1) # strict: ties provisionally round down + up = ((a == mid) & tie_up).any(-1) # ties whose even side is up + code = (greater + up).to(torch.uint8) + negative = (v < 0) | ((v == 0) & torch.signbit(v)) + return code | (negative.to(torch.uint8) << 3) + + @staticmethod + def _swizzle(sf: torch.Tensor) -> torch.Tensor: + """Scatter [rows, cols] scale bytes into the 128x4-swizzled buffer.""" + rows, cols = sf.shape + padded_cols, padded_rows = _pad_up(cols, 4), _pad_up(rows, 128) + r = torch.arange(rows, device=sf.device).view(-1, 1) + c = torch.arange(cols, device=sf.device).view(1, -1) + index = ( + (c % 4) + + (c // 4) * (4 * 128) + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + out = torch.zeros(padded_rows * padded_cols, dtype=torch.uint8, device=sf.device) + out[index.reshape(-1)] = sf.reshape(-1) + return out + + +fp4_quantize = Fp4Quantize() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/__init__.py new file mode 100644 index 000000000000..4e871f35b7f0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Speculative-decoding entries: the drafter's own ops.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_block_tables.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_block_tables.py new file mode 100644 index 000000000000..f7c89b37486d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_block_tables.py @@ -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. +# ============================================================================= +# DSpark / DFlash drafter block-table refresh: one launch per step +# ============================================================================= +# +# It replaces, bit for bit, what DFlashWorker._refresh_ctx_block_tables does with the draft pool's encoded block +# offsets of this step (field 0 of the K/V pair, pool ``pool_idx``, the first ``num_seqs`` rows): +# counts[r] = number of entries >= 0 in row r (placeholders are negative) +# tables[r, j] = max(offset[r, j], 0) // divisor (int32) +# Integer arithmetic only, so the result is exact; the stock path spends one launch per torch op on it. + +from typing import Optional + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _refresh_kernel( + src_ptr, # encoded offsets, row r at src_ptr + r * src_stride_row, entry j at + j * src_stride_col + src_stride_row, + src_stride_col, + tables_ptr, # int32 [rows, width], row-major + tables_stride, + counts_ptr, # int64 [rows] + width, + divisor, + BLOCK: tl.constexpr, +): + r = tl.program_id(0).to(tl.int64) + count = tl.zeros((), dtype=tl.int64) + for start in range(0, width, BLOCK): + o = start + tl.arange(0, BLOCK) + m = o < width + v = tl.load(src_ptr + r * src_stride_row + o * src_stride_col, mask=m, other=-1).to( + tl.int64 + ) + count += tl.sum((v >= 0).to(tl.int64), axis=0) + # Non-negative after the clamp, so integer division is the stock floor division. + tl.store( + tables_ptr + r * tables_stride + o, (tl.maximum(v, 0) // divisor).to(tl.int32), mask=m + ) + tl.store(counts_ptr + r, count) + + +@triton.jit +def _refresh_kernel_2d( + src_ptr, + src_stride_row, + src_stride_col, + tables_ptr, + tables_stride, + counts_ptr, + width, + divisor, + BLOCK: tl.constexpr, + COUNT_BLOCK: tl.constexpr, +): + """S10c-B2: grid (rows, column chunks). Each program decodes one chunk in int32 (the offsets fit int32 and are + non-negative after the clamp, so int32 floor division is the stock floor division); chunk 0 of a row also counts + the row's non-placeholder entries over the whole row. No atomics: every output has exactly one writer.""" + r = tl.program_id(0).to(tl.int64) + c = tl.program_id(1) + o = c * BLOCK + tl.arange(0, BLOCK) + m = o < width + v = tl.load(src_ptr + r * src_stride_row + o * src_stride_col, mask=m, other=-1).to(tl.int32) + tl.store(tables_ptr + r * tables_stride + o, tl.maximum(v, 0) // divisor, mask=m) + if c == 0: + count = tl.zeros((), dtype=tl.int64) + for start in range(0, width, COUNT_BLOCK): + oc = start + tl.arange(0, COUNT_BLOCK) + vc = tl.load( + src_ptr + r * src_stride_row + oc * src_stride_col, mask=oc < width, other=-1 + ) + count += tl.sum((vc >= 0).to(tl.int64), axis=0) + tl.store(counts_ptr + r, count) + + +def refresh_block_tables( + offsets: torch.Tensor, + pool_idx: int, + num_seqs: int, + divisor: int, + tables: torch.Tensor, + counts: torch.Tensor, +) -> bool: + """Decode this step's draft block offsets ``offsets[pool_idx, :num_seqs, 0]`` into ``tables[:num_seqs]`` and their + non-placeholder counts into ``counts[:num_seqs]``, in one launch; returns False (nothing written) when the layout + is not one the kernel reproduces, so the caller keeps the stock path.""" + if not supported(offsets, tables, counts, divisor): + return False + if num_seqs <= 0: + return True + src = offsets[pool_idx, :num_seqs, 0] + width = src.shape[-1] + if _parallel[0] and divisor <= 2**31 - 1 and (offsets.dtype == torch.int32 or _fits_int32(src)): + block = min(1024, triton.next_power_of_2(width)) + _refresh_kernel_2d[(num_seqs, triton.cdiv(width, block))]( + src, + src.stride(0), + src.stride(1), + tables, + tables.stride(0), + counts, + width, + divisor, + BLOCK=block, + COUNT_BLOCK=min(4096, triton.next_power_of_2(width)), + num_warps=4, + ) + return True + _refresh_kernel[(num_seqs,)]( + src, + src.stride(0), + src.stride(1), + tables, + tables.stride(0), + counts, + width, + divisor, + BLOCK=min(1024, triton.next_power_of_2(width)), + num_warps=1, + ) + return True + + +_parallel = [False] # S10c-B2: the 2-D int32 decode (see enable_parallel_refresh) + + +def enable_parallel_refresh(enabled: bool = True) -> bool: + """S10c-B2: decode with the (row, chunk) grid in int32 instead of one int64 program per row (the same values: the + offsets fit int32 and are non-negative after the clamp). Off by default.""" + _parallel[0] = bool(enabled) + return _parallel[0] + + +def _fits_int32(src: torch.Tensor) -> bool: + """int64 offsets are decoded in int32 only when the caller's layout guarantees they fit; without a cheap + capture-safe check, int64 offsets keep the int64 kernel.""" + return False + + +def supported( + offsets: torch.Tensor, tables: torch.Tensor, counts: torch.Tensor, divisor: Optional[int] +) -> bool: + """[pools, rows, 2, width] int32 / int64 offsets on the tables' device; int32 row-major tables of the same width; + int64 counts; a positive integer divisor.""" + return ( + offsets.dim() == 4 + and offsets.dtype in (torch.int32, torch.int64) + and tables.dtype == torch.int32 + and tables.dim() == 2 + and tables.stride(1) == 1 + and tables.shape[1] == offsets.shape[-1] + and counts.dtype == torch.int64 + and offsets.device == tables.device == counts.device + and isinstance(divisor, int) + and divisor >= 1 + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_ctx_advance.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_ctx_advance.py new file mode 100644 index 000000000000..5d20f9f90217 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_ctx_advance.py @@ -0,0 +1,137 @@ +# 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. +# ============================================================================= +# DSpark / DFlash drafter context advance + read view: one launch per step +# ============================================================================= +# +# It replaces, bit for bit, what DFlashWorker's generation branch does after the context store on a bound draft pool +# (_advance_ctx_len, then the advertised context length and the rebased read view of dflash_rebased_ctx_view): +# ctx_len[slot] = min(ctx_len[slot] + num_accepted, max_ctx); ctx_len[dummy_slot] = 0 +# num_ctx = min(ctx_len[slot], max(block_counts[row] * page - block_size, 0)) +# rebase = min(base_pages[slot], num_ctx // page) +# num_ctx_out = max(num_ctx - base_pages[slot] * page, 0) (int64) +# read[row, j] = tables[row, min(j + rebase, width - 1)] (int32) +# One program per generation request, so each request's advance is read back by the same program; rows on the dummy +# slot (padding) skip the add and read 0, which the stock path's zeroing after the add makes them read too. + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _advance_kernel( + slots_ptr, # [n] int64 + rows_ptr, # [n] int64: block-table rows of the requests + nacc_ptr, # [n] int + ctx_len_ptr, # [slots] int64, in / out + counts_ptr, # [rows] int64 + base_ptr, # [slots] int64 + tables_ptr, # int32 [rows, width] + read_ptr, # int32 [rows, width] + out_ptr, # [n] int64 advertised context lengths + width, + tables_stride, + read_stride, + max_ctx, + dummy_slot, + page, + block_size, + BLOCK: tl.constexpr, +): + r = tl.program_id(0).to(tl.int64) + slot = tl.load(slots_ptr + r) + row = tl.load(rows_ptr + r) + is_dummy = slot == dummy_slot + old = tl.load(ctx_len_ptr + slot) + new = tl.minimum(old + tl.load(nacc_ptr + r).to(tl.int64), max_ctx) + if not is_dummy: + tl.store(ctx_len_ptr + slot, new) + if r == 0: + tl.store(ctx_len_ptr + dummy_slot, tl.zeros((), dtype=tl.int64)) + new = tl.where(is_dummy, 0, new) + allocated = tl.maximum(tl.load(counts_ptr + row) * page - block_size, 0) + num_ctx = tl.minimum(new, allocated) + base = tl.load(base_ptr + slot) + rebase = tl.minimum(base, num_ctx // page) + tl.store(out_ptr + r, tl.maximum(num_ctx - base * page, 0)) + for start in range(0, width, BLOCK): + j = start + tl.arange(0, BLOCK) + m = j < width + src = tl.minimum(j + rebase, width - 1) + v = tl.load(tables_ptr + row * tables_stride + src, mask=m) + tl.store(read_ptr + row * read_stride + j, v, mask=m) + + +def advance_ctx_and_read_view( + slots: torch.Tensor, + rows: torch.Tensor, + num_accepted: torch.Tensor, + ctx_len: torch.Tensor, + max_ctx: int, + dummy_slot: int, + block_counts: torch.Tensor, + base_pages: torch.Tensor, + tables: torch.Tensor, + tables_read: torch.Tensor, + page_size: int, + block_size: int, +): + """Advance the generation requests' context lengths and build their read view in one launch: returns the + advertised context lengths [n] (int64) and writes ``tables_read[rows]``; None (nothing written) for a layout the + kernel does not reproduce, so the caller keeps the stock path.""" + n = slots.shape[0] + if not ( + slots.dim() == 1 + and rows.shape == (n,) + and num_accepted.shape == (n,) + and num_accepted.dtype in (torch.int32, torch.int64) + and ctx_len.dtype == torch.int64 + and block_counts.dtype == torch.int64 + and base_pages.dtype == torch.int64 + and tables.dtype == torch.int32 + and tables_read.dtype == torch.int32 + and tables.shape == tables_read.shape + and tables.stride(1) == 1 + and tables_read.stride(1) == 1 + and page_size >= 1 + ): + return None + out = torch.empty(n, dtype=torch.int64, device=ctx_len.device) + if n == 0: + ctx_len[dummy_slot] = 0 + return out + width = tables.shape[1] + _advance_kernel[(n,)]( + slots.to(torch.int64).contiguous(), + rows.to(torch.int64).contiguous(), + num_accepted.contiguous(), + ctx_len, + block_counts, + base_pages, + tables, + tables_read, + out, + width, + tables.stride(0), + tables_read.stride(0), + int(max_ctx), + int(dummy_slot), + int(page_size), + int(block_size), + BLOCK=min(1024, triton.next_power_of_2(width)), + num_warps=4, + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_ctx_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_ctx_store.py new file mode 100644 index 000000000000..26b67405c20c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_ctx_store.py @@ -0,0 +1,309 @@ +# 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. +# ============================================================================= +# DSpark / DFlash drafter context K / V store after the fused K / V GEMM: one launch for every drafter layer +# ============================================================================= +# +# Adapted from SGLang's DFlash KV materialization (_fused_norm_rope_kernel_stacked in +# python/sglang/kernels/ops/speculative/fused_kv_materialize.py, https://github.com/sgl-project/sglang, commit +# 6ed9843b57aaecbf0566f24a4eaf78c7c92ef25d, Apache-2.0): one program per (row, layer, kv head) reads the stacked +# [K | V] GEMM output, applies the layer's k_norm and NeoX RoPE to K. Changed here: the 0/1 accept mask on K and V, the +# paged HND store through the request's block table (instead of a dense output), and the stock path's rounding points. +# +# It replaces, bit for bit, what DFlashWorker does after DFlashForCausalLM.precompute_context_kv's GEMM. The k_norm +# stays torch's own F.rms_norm (as precompute_context_kv runs it: an in-kernel row sum did not reproduce torch's +# reduction order on every input); the bf16 weight multiply runs in the kernel (``fold_k_norm_weight``, TAILBK_SPEC F: +# fp32 product rounded to bf16, as torch's bf16 multiply) or as torch's own; this one launch does the rest: +# k = rms_norm(k) (fp32, rounded to bf16) * k_norm weight (bf16) -> flashinfer cos/sin-cache NeoX RoPE over the +# first ROT dims, ROT = the cache's width (the rotary dim; a partial rotary leaves dims ROT..HD as they are) (fp32 +# from the bf16 k, x*cos + (-/+)p*sin as fma(x, cos, (-/+)p*sin), rounded to bf16) -> k, v *= accept mask (bf16) -> +# flashinfer append_paged_kv_cache per layer (HND) at (block_table[row, col // page], col % page). +# Row r belongs to request r // (K+1) at block offset r % (K+1); it is accepted when that offset < num_accepted. +# Masked rows still write (zeros) at their clamped column, as the stock fixed-size write does. + +from typing import Dict, Optional, Sequence, Tuple, Union + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + + +@triton.jit +def _ctx_store_kernel( + kv_ptr, # [N, L * 2 * NKV * HD] the fused K / V GEMM output, per layer [K heads | V heads] (V is read here) + k_ptr, # [N, L, NKV, HD] K after torch's k_norm and (without FOLD_W) the weight multiply + w_ptr, # [L, HD] bf16 k_norm weight (FOLD_W) + cs_ptr, # [max_pos, ROT] fp32 cos | sin (flashinfer's cos/sin cache; ROT = the rotary dim) + pos_ptr, # [N] RoPE positions + nacc_ptr, # [N // KP1] accepted tokens per request + rows_ptr, # [N // KP1] block-table row per request + col_ptr, # [N] cache column per row (unused with COL_FROM_POS) + counts_ptr, # [rows] int64 pages per block-table row (COL_FROM_POS on a bound pool) + max_col, # the private arena's last column (COL_FROM_POS without a bound pool) + table_ptr, # [rows, pages_per_row] page ids + cache_ptr, # stacked [L, pages, 2, NKV, PS, HD] cache, or an int64 array of L per-layer cache addresses + kv_stride_row, + table_stride, + cache_stride_layer, + KP1: tl.constexpr, + L: tl.constexpr, + NKV: tl.constexpr, + HD: tl.constexpr, + ROT: tl.constexpr, + PS: tl.constexpr, + PTR_LIST: tl.constexpr, + COL_FROM_POS: tl.constexpr, # derive the column from the position as the stock clamp does + BOUND: tl.constexpr, # COL_FROM_POS on a bound pool (clamp to the row's pages) vs the private arena + FOLD_W: tl.constexpr, # k_ptr holds rms_norm(k) (bf16); the bf16 weight multiply runs here +): + r = tl.program_id(0) + layer = tl.program_id(1) + h = tl.program_id(2) + half: tl.constexpr = HD // 2 + rhalf: tl.constexpr = ROT // 2 + o = tl.arange(0, half) + ro = tl.arange(0, rhalf) + req = r // KP1 + m = tl.where((r % KP1) < tl.load(nacc_ptr + req), 1.0, 0.0) + vbase = kv_ptr + r * kv_stride_row + (layer * 2 * NKV + NKV + h) * HD + kbase = k_ptr + ((r * L + layer) * NKV + h) * HD + a = tl.load(kbase + ro) + b = tl.load(kbase + rhalf + ro) + if FOLD_W: # torch's bf16 * bf16: the fp32 product rounded to bf16 + wbase = w_ptr + layer * HD + a = (a.to(tl.float32) * tl.load(wbase + ro).to(tl.float32)).to(tl.bfloat16) + b = (b.to(tl.float32) * tl.load(wbase + rhalf + ro).to(tl.float32)).to(tl.bfloat16) + a = a.to(tl.float32) + b = b.to(tl.float32) + p = tl.load(pos_ptr + r) + c = tl.load(cs_ptr + p * ROT + ro) + s = tl.load(cs_ptr + p * ROT + rhalf + ro) + ra = (tl.fma(a, c, -b * s).to(tl.bfloat16).to(tl.float32) * m).to(tl.bfloat16) + rb = (tl.fma(b, c, a * s).to(tl.bfloat16).to(tl.float32) * m).to(tl.bfloat16) + v0 = (tl.load(vbase + o).to(tl.float32) * m).to(tl.bfloat16) + v1 = (tl.load(vbase + half + o).to(tl.float32) * m).to(tl.bfloat16) + row = tl.load(rows_ptr + req) + last = (r % KP1) == KP1 - 1 + if COL_FROM_POS: + # The stock column: ctx_len + offset (= the position), clamped to the request's pages (bound pool: + # min(col, pages * PS - 1).clamp(min=0)) or to the arena (min(col, max_ctx - 1)). + p_nxt = tl.load(pos_ptr + tl.where(last, r, r + 1)) + if BOUND: + cap = tl.load(counts_ptr + row) * PS - 1 + col = tl.maximum(tl.minimum(p, cap), 0) + nxt = tl.maximum(tl.minimum(p_nxt, cap), 0) + else: + col = tl.minimum(p, max_col) + nxt = tl.minimum(p_nxt, max_col) + else: + col = tl.load(col_ptr + r) + nxt = tl.load( + col_ptr + tl.where(last, r, r + 1) + ) # the next row of the same request (itself for the last) + page = tl.load(table_ptr + row * table_stride + col // PS).to(tl.int64) + if PTR_LIST: + cache = tl.load(cache_ptr + layer).to(tl.pointer_type(tl.bfloat16)) + else: + cache = cache_ptr + layer.to(tl.int64) * cache_stride_layer + kdst = cache + ((page * 2 * NKV + h) * PS + col % PS) * HD + vdst = cache + ((page * 2 * NKV + NKV + h) * PS + col % PS) * HD + # Clamped rows of one request share a column; the last of them writes. Stock's order there is a race between the + # append's thread blocks; that no reader reads such a cell is shown in the S10c-A in-process gate. + write = last | (nxt != col) + tl.store(kdst + ro, ra, mask=write) + tl.store(kdst + rhalf + ro, rb, mask=write) + if ROT < HD: # partial rotary: dims ROT..HD are not rotated, only masked + po = ROT + tl.arange(0, HD - ROT) + kp = tl.load(kbase + po) + if FOLD_W: + kp = (kp.to(tl.float32) * tl.load(w_ptr + layer * HD + po).to(tl.float32)).to( + tl.bfloat16 + ) + kp = (kp.to(tl.float32) * m).to(tl.bfloat16) + tl.store(kdst + po, kp, mask=write) + tl.store(vdst + o, v0, mask=write) + tl.store(vdst + half + o, v1, mask=write) + + +_fold_w = [ + False +] # TAILBK_SPEC F: the k_norm weight multiply in the kernel (see fold_k_norm_weight) + + +def fold_k_norm_weight(enabled: bool = True) -> bool: + """Run the stacked k_norm weight multiply inside the store kernel (one launch fewer per call) when the weight is a + contiguous bf16 [L, HD] and rms_norm's output is bf16: the fp32 product rounded to bf16, bitwise torch's bf16 + multiply. Off by default; process-wide.""" + _fold_w[0] = bool(enabled) + return _fold_w[0] + + +# Per-layer cache lists (the bound draft pool) -> their int64 address array, built once outside capture. +_PTR_ARRAYS: Dict[Tuple[int, ...], torch.Tensor] = {} + + +def _cache_arg( + kv_buf: Union[torch.Tensor, Sequence[torch.Tensor]], +) -> Optional[Tuple[torch.Tensor, bool, int]]: + if isinstance(kv_buf, torch.Tensor): + return kv_buf, False, kv_buf.stride(0) + key = tuple(t.data_ptr() for t in kv_buf) + arr = _PTR_ARRAYS.get(key) + if arr is None: + if torch.cuda.is_current_stream_capturing(): + return None # never seen eagerly: the caller keeps the stock store for this capture + arr = torch.tensor(key, dtype=torch.int64, device=kv_buf[0].device) + _PTR_ARRAYS[key] = arr + return arr, True, 0 + + +def _pow2(n: int) -> bool: + return n >= 1 and n & (n - 1) == 0 + + +def _rotary_dim(draft_model, hd: int) -> int: + """The rotary dim the drafter's RoPE uses (the width of its fp32 cos/sin cache: flashinfer rotates that many leading + dims of each head), or 0 when the kernel cannot reproduce it.""" + cos_sin = draft_model._get_cos_sin_cache() + if cos_sin.dim() != 2 or cos_sin.dtype != torch.float32 or not cos_sin.is_contiguous(): + return 0 + rot = cos_sin.shape[1] + if rot % 2 or rot > hd or not _pow2(rot // 2) or (rot < hd and not _pow2(hd - rot)): + return 0 + return rot + + +def ctx_store_supported(draft_model, kv_buf) -> bool: + """Whether this drafter's context store is one the kernel reproduces: GQA K and V into a bf16 HND paged cache, + stacked per-layer k_norm, NeoX RoPE over the leading rotary dims (whole or partial head) from the fp32 cos/sin + cache, head_dim a power of two.""" + first = kv_buf if isinstance(kv_buf, torch.Tensor) else kv_buf[0] + layer_shape = first.shape[1:] if isinstance(kv_buf, torch.Tensor) else first.shape + hd = getattr(draft_model, "_head_dim", 0) + return ( + getattr(draft_model, "_k_norm_stacked", None) is not None + and getattr(draft_model, "_k_norm_eps", None) is not None + and getattr(draft_model, "_is_neox", False) + and getattr(draft_model, "_num_kv_heads", 0) >= 1 + and hd >= 16 + and hd & (hd - 1) == 0 + and first.dtype == torch.bfloat16 + and len(layer_shape) == 5 + and layer_shape[1] == 2 + and layer_shape[2] == draft_model._num_kv_heads + and layer_shape[4] == hd + and all( + t.is_contiguous() for t in ([kv_buf] if isinstance(kv_buf, torch.Tensor) else kv_buf) + ) + and _rotary_dim(draft_model, hd) > 0 + ) + + +def fused_ctx_store( + draft_model, + projected: torch.Tensor, + positions: torch.Tensor, + num_accepted: torch.Tensor, + request_rows: torch.Tensor, + cols: Optional[torch.Tensor], + block_table: torch.Tensor, + kv_buf: Union[torch.Tensor, Sequence[torch.Tensor]], + tokens_per_request: int, + *, + capacity_counts: Optional[torch.Tensor] = None, + max_col: Optional[int] = None, +) -> bool: + """Store the context K / V of ``projected`` (fc + hidden_norm'd rows, request-major, ``tokens_per_request`` rows + each) into the paged draft cache; returns False (nothing written) when the stock path must run instead. + + ``request_rows`` / ``num_accepted`` are per request, ``positions`` / ``cols`` per row. With ``cols`` None the kernel + derives each row's column from its position as the stock path does: clamped to the request's pages on a bound pool + (``capacity_counts``: pages per block-table row, int64) or to ``max_col`` on the private arena.""" + if not ctx_store_supported(draft_model, kv_buf): + return False + if cols is None and ( + (capacity_counts is None) == (max_col is None) + or (capacity_counts is not None and capacity_counts.dtype != torch.int64) + ): + return False + cache = _cache_arg(kv_buf) + if cache is None: + return False + cache_arg, ptr_list, cache_stride_layer = cache + if draft_model._fused_kv_weight is None: + draft_model._build_fused_kv_buffers() + weight = draft_model._fused_kv_weight + # precompute_context_kv's GEMM, unchanged. + if getattr(draft_model, "_input_ln_eps", None) is not None: + ph = projected.float() + ph = ph * torch.rsqrt(ph.pow(2).mean(-1, keepdim=True) + draft_model._input_ln_eps) + projected = ph.to(weight.dtype) + elif projected.dtype != weight.dtype: + projected = projected.to(weight.dtype) + kv_flat = F.linear(projected, weight, draft_model._fused_kv_bias) + n = kv_flat.shape[0] + if n == 0: + return True + num_layers, nkv, hd = ( + draft_model._num_attn_layers, + draft_model._num_kv_heads, + draft_model._head_dim, + ) + # precompute_context_kv's k_norm, unchanged: contiguous K, F.rms_norm, the stacked weight multiply. + k = kv_flat.view(n, num_layers, 2, nkv, hd)[:, :, 0].contiguous() + k = F.rms_norm(k, (hd,), eps=draft_model._k_norm_eps) + w = draft_model._k_norm_stacked + fold = ( + _fold_w[0] + and k.dtype == torch.bfloat16 + and w.dtype == torch.bfloat16 + and w.is_contiguous() + and w.numel() == num_layers * hd + and w.device == k.device + ) + if not fold: + k = k * w.view(1, num_layers, 1, hd) + first = kv_buf if isinstance(kv_buf, torch.Tensor) else kv_buf[0] + cos_sin = draft_model._get_cos_sin_cache() + _ctx_store_kernel[(n, num_layers, nkv)]( + kv_flat, + k, + w, + cos_sin, + positions.contiguous(), + num_accepted.contiguous(), + request_rows.contiguous(), + positions if cols is None else cols.contiguous(), + capacity_counts if capacity_counts is not None else positions, + int(max_col) if max_col is not None else 0, + block_table, + cache_arg, + kv_flat.stride(0), + block_table.stride(0), + cache_stride_layer, + KP1=tokens_per_request, + L=num_layers, + NKV=nkv, + HD=hd, + ROT=cos_sin.shape[1], + PS=first.shape[-2], + PTR_LIST=ptr_list, + COL_FROM_POS=cols is None, + BOUND=capacity_counts is not None, + FOLD_W=fold, + num_warps=1, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_drafter_inputs.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_drafter_inputs.py new file mode 100644 index 000000000000..56d9cef29894 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_drafter_inputs.py @@ -0,0 +1,143 @@ +# 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. +# ============================================================================= +# DSpark / DFlash drafter step inputs: one launch per step +# ============================================================================= +# +# It replaces, bit for bit, the start of DFlashWorker.prepare_1st_drafter_inputs' generation branch, for a drafter whose +# token embedding is a plain row gather of a replicated BF16 table: +# bonus[r] = accepted[r, max(num_accepted[r] - 1, 0)] (int64) +# query_pos[r, j] = min(ctx_len[slot[r]] + num_accepted[r], max_ctx) + j, j < block (int64) +# ctx_pos[r, j] = ctx_len[slot[r]] + j, j < K + 1 (int64) +# noise[r, 0, :] = table[bonus[r]], or +0 for an id outside [0, vocab) (m3_embedding's masked_fill) +# noise[r, j>0, :] = mask_row (the drafter's trained mask embedding, else table[mask_token_id]) +# Integer arithmetic and row copies only, so the result is exact; the stock path spends a launch per torch op on it. + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _inputs_kernel( + acc_ptr, # accepted tokens [n, acc_stride] (int32 / int64) + acc_stride, + nacc_ptr, # accepted counts [n] (int) + slots_ptr, # [n] int64: the request's ctx slot + ctx_len_ptr, # [slots] int64 + max_ctx, + table_ptr, # [vocab, hidden] bf16 embedding table + mask_ptr, # [hidden] bf16 mask row + bonus_ptr, # out [n] int64 + qpos_ptr, # out [n, BLOCK_Q] int64 + cpos_ptr, # out [n, KP1] int64 + noise_ptr, # out [n, BLOCK_Q, hidden] bf16 + hidden, + vocab, + BLOCK_Q: tl.constexpr, + KP1: tl.constexpr, + BLOCK_QP: tl.constexpr, # next powers of two of BLOCK_Q / KP1 (tl.arange sizes) + KP1P: tl.constexpr, + BLOCK_H: tl.constexpr, +): + r = tl.program_id(0).to(tl.int64) + j = tl.program_id(1) # the block row this program copies + hb = tl.program_id(2) + nacc = tl.load(nacc_ptr + r).to(tl.int64) + bonus = tl.load(acc_ptr + r * acc_stride + tl.maximum(nacc - 1, 0)).to(tl.int64) + o = hb * BLOCK_H + tl.arange(0, BLOCK_H) + m = o < hidden + if j == 0: + valid = (bonus >= 0) & (bonus < vocab) + row = tl.load(table_ptr + tl.where(valid, bonus, 0) * hidden + o, mask=m & valid, other=0.0) + else: + row = tl.load(mask_ptr + o, mask=m) + tl.store(noise_ptr + (r * BLOCK_Q + j) * hidden + o, row, mask=m) + if (j == 0) & (hb == 0): + ctx = tl.load(ctx_len_ptr + tl.load(slots_ptr + r)) + tl.store(bonus_ptr + r, bonus) + now = tl.minimum(ctx + nacc, max_ctx) + q = tl.arange(0, BLOCK_QP) + tl.store(qpos_ptr + r * BLOCK_Q + q, now + q, mask=q < BLOCK_Q) + c = tl.arange(0, KP1P) + tl.store(cpos_ptr + r * KP1 + c, ctx + c, mask=c < KP1) + + +def drafter_inputs( + accepted: torch.Tensor, + num_accepted: torch.Tensor, + slots: torch.Tensor, + ctx_len: torch.Tensor, + max_ctx: int, + table: torch.Tensor, + mask_row: torch.Tensor, + block_size: int, + tokens_per_request: int, +) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]: + """(bonus [n] int64, query positions [n, block_size] int64, ctx positions [n, tokens_per_request] int64, + noise block [n, block_size, hidden]) in one launch, or None when the inputs are not a layout the kernel reproduces + (the caller keeps the stock path). ``mask_row`` [hidden] is the row every slot after the first takes.""" + n = num_accepted.shape[0] + if not ( + accepted.dim() == 2 + and accepted.shape[0] == n + and accepted.stride(1) == 1 + and accepted.dtype in (torch.int32, torch.int64) + and num_accepted.dtype in (torch.int32, torch.int64) + and slots.shape == (n,) + and ctx_len.dtype == torch.int64 + and table.dim() == 2 + and table.is_contiguous() + and mask_row.shape == (table.shape[1],) + and mask_row.dtype == table.dtype + and mask_row.is_contiguous() + and block_size >= 1 + and tokens_per_request >= 1 + ): + return None + hidden = table.shape[1] + dev = table.device + bonus = torch.empty(n, dtype=torch.int64, device=dev) + qpos = torch.empty(n, block_size, dtype=torch.int64, device=dev) + cpos = torch.empty(n, tokens_per_request, dtype=torch.int64, device=dev) + noise = torch.empty(n, block_size, hidden, dtype=table.dtype, device=dev) + if n == 0: + return bonus, qpos, cpos, noise + block_h = 1024 + _inputs_kernel[(n, block_size, triton.cdiv(hidden, block_h))]( + accepted, + accepted.stride(0), + num_accepted.contiguous(), + slots.to(torch.int64).contiguous(), + ctx_len, + int(max_ctx), + table, + mask_row, + bonus, + qpos, + cpos, + noise, + hidden, + table.shape[0], + BLOCK_Q=block_size, + KP1=tokens_per_request, + BLOCK_QP=triton.next_power_of_2(block_size), + KP1P=triton.next_power_of_2(tokens_per_request), + BLOCK_H=block_h, + num_warps=4, + ) + return bonus, qpos, cpos, noise diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_forward_indices.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_forward_indices.py new file mode 100644 index 000000000000..42448af2b144 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_forward_indices.py @@ -0,0 +1,134 @@ +# 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. +# ============================================================================= +# DSpark / DFlash drafter forward: the paged block decode's index prelude in one launch +# ============================================================================= +# +# It replaces, bit for bit, the integer lines at the top of DFlashForCausalLM.dflash_forward on a paged backend: +# cache_seqlens = num_ctx[:B].to(int32) +# cache_batch_idx = ctx_cache_batch_idx.to(int32) +# seq_lens_after = cache_seqlens + block_size +# kv_indptr = arange(0, (B + 1) * pages_per_slot, pages_per_slot) (int32) +# kv_last_page_len = ((seq_lens_after - 1) % page_size) + 1 (floored modulo, as torch) +# append_positions = (cache_seqlens.view(-1, 1) + arange(block_size)).reshape(-1) (int32) +# query_pos_i32 = query_positions.reshape(-1).to(int32) +# The int64 -> int32 casts keep the low 32 bits, as torch's do. One program writes every output. + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _indices_kernel( + num_ctx_ptr, # [B] int32 / int64 + batch_idx_ptr, # [B] int32 / int64 + qpos_ptr, # [B, block] int32 / int64, strided + qpos_s0, + qpos_s1, + seqlens_ptr, # out [B] int32 + batch_out_ptr, # out [B] int32 + after_ptr, # out [B] int32 + indptr_ptr, # out [B + 1] int32 + last_ptr, # out [B] int32 + append_ptr, # out [B * block] int32 + qflat_ptr, # out [B * block] int32 + B, + block, + pages_per_slot, + page_size, + BLOCK_B: tl.constexpr, + BLOCK_T: tl.constexpr, +): + b = tl.arange(0, BLOCK_B) + mb = b < B + seqlens = tl.load(num_ctx_ptr + b, mask=mb, other=0).to(tl.int32) + tl.store(seqlens_ptr + b, seqlens, mask=mb) + tl.store(batch_out_ptr + b, tl.load(batch_idx_ptr + b, mask=mb, other=0).to(tl.int32), mask=mb) + after = seqlens + block + tl.store(after_ptr + b, after, mask=mb) + r = (after - 1) % page_size + r = tl.where(r < 0, r + page_size, r) # torch's floored modulo for a positive divisor + tl.store(last_ptr + b, r + 1, mask=mb) + mi = b <= B + tl.store(indptr_ptr + b, b * pages_per_slot, mask=mi) + t = tl.arange(0, BLOCK_T) + mt = t < B * block + rb = t // block + j = t - rb * block + base = tl.load(num_ctx_ptr + rb, mask=mt, other=0).to(tl.int32) + tl.store(append_ptr + t, base + j, mask=mt) + q = tl.load(qpos_ptr + rb * qpos_s0 + j * qpos_s1, mask=mt, other=0).to(tl.int32) + tl.store(qflat_ptr + t, q, mask=mt) + + +def forward_indices( + num_ctx: torch.Tensor, + batch_idx: torch.Tensor, + query_positions: torch.Tensor, + block_size: int, + pages_per_slot: int, + page_size: int, +): + """The paged dflash_forward's index prelude in one launch: (cache_seqlens, cache_batch_idx, seq_lens_after, + kv_indptr, kv_last_page_len, append_positions, query_positions_flat), all int32 and bitwise the stock lines; or + None (nothing computed) for a layout the kernel does not reproduce, so the caller keeps the stock lines.""" + ok_dtypes = (torch.int32, torch.int64) + B = num_ctx.shape[0] if num_ctx.dim() == 1 else -1 + if not ( + B >= 1 + and num_ctx.dtype in ok_dtypes + and batch_idx.dtype in ok_dtypes + and query_positions.dtype in ok_dtypes + and batch_idx.shape == (B,) + and query_positions.dim() == 2 + and tuple(query_positions.shape) == (B, block_size) + and block_size >= 1 + and pages_per_slot >= 1 + and page_size >= 1 + and num_ctx.is_cuda + ): + return None + dev = num_ctx.device + i32 = dict(dtype=torch.int32, device=dev) + seqlens = torch.empty(B, **i32) + batch_out = torch.empty(B, **i32) + after = torch.empty(B, **i32) + indptr = torch.empty(B + 1, **i32) + last = torch.empty(B, **i32) + append = torch.empty(B * block_size, **i32) + qflat = torch.empty(B * block_size, **i32) + _indices_kernel[(1,)]( + num_ctx, + batch_idx, + query_positions, + query_positions.stride(0), + query_positions.stride(1), + seqlens, + batch_out, + after, + indptr, + last, + append, + qflat, + B, + block_size, + pages_per_slot, + page_size, + BLOCK_B=triton.next_power_of_2(B + 1), + BLOCK_T=triton.next_power_of_2(B * block_size), + num_warps=1, + ) + return seqlens, batch_out, after, indptr, last, append, qflat diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_markov_bias.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_markov_bias.py new file mode 100644 index 000000000000..5c2cc449e3cf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_markov_bias.py @@ -0,0 +1,128 @@ +# 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. +# ============================================================================= +# DSpark vanilla-Markov head: the biased block logits' glue around the bias GEMM, in two launches +# ============================================================================= +# +# It replaces, bit for bit, the elementwise lines of the batched biased-logits form (the tp targets' S4-logits path) +# prevs = cat([first_prev.view(B, 1), tokens[:, :K - 1].long()], 1).reshape(-1) +# emb = markov_prev_embeddings(prevs, w1) (get_masked_input_and_mask, F.embedding, masked_fill: an anchor +# outside [0, V) gives a zero row) +# bias = F.linear(emb, w2_shard) (unchanged: the caller's GEMM) +# out = shard_logits.float() + bias.view(B, K, N).to(float32) +# with ``prev_embeddings`` (one launch for the 11 index / mask / gather / fill launches) and ``widen_add`` (one launch +# for the two widening copies and the add). The gather copies w1's rows bit for bit (zero rows are +0.0, as +# masked_fill(0)); the add is one fp32 IEEE add of two exactly widened bf16 values, as torch's. + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _prev_embeddings_kernel( + first_prev_ptr, # [B] int32 / int64 + tokens_ptr, # [B, K] int32 / int64, contiguous + w1_ptr, # [V, R] bf16, contiguous + out_ptr, # out [B * K, R] bf16 + K, + V, + R: tl.constexpr, +): + row = tl.program_id(0) + b = row // K + k = row - b * K + if k == 0: + prev = tl.load(first_prev_ptr + b).to(tl.int64) + else: + prev = tl.load(tokens_ptr + b * K + k - 1).to(tl.int64) + valid = (prev >= 0) & (prev < V) + safe = tl.where(valid, prev, 0) + c = tl.arange(0, R) + vals = tl.load(w1_ptr + safe * R + c) + vals = tl.where(valid, vals, tl.zeros_like(vals)) + tl.store(out_ptr + row.to(tl.int64) * R + c, vals) + + +@triton.jit +def _widen_add_kernel( + logits_ptr, # [M, N] bf16, contiguous + bias_ptr, # [M, N] bf16, contiguous + out_ptr, # out [M, N] fp32 + N, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + c = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + m = c < N + a = tl.load(logits_ptr + row * N + c, mask=m, other=0.0).to(tl.float32) + bb = tl.load(bias_ptr + row * N + c, mask=m, other=0.0).to(tl.float32) + tl.store(out_ptr + row * N + c, a + bb, mask=m) + + +_INT = (torch.int32, torch.int64) +WIDEN_BLOCK = 4096 + + +def supports( + shard_logits: torch.Tensor, first_prev: torch.Tensor, tokens: torch.Tensor, w1: torch.Tensor +) -> bool: + """Whether both launches take these operands; a caller that gets False runs the in-tree lines.""" + if shard_logits.dim() != 3 or tokens.dim() != 2 or first_prev.dim() != 1 or w1.dim() != 2: + return False + B, K, _ = shard_logits.shape + R = w1.shape[1] + return ( + shard_logits.is_cuda + and shard_logits.dtype == torch.bfloat16 + and shard_logits.is_contiguous() + and w1.dtype == torch.bfloat16 + and w1.is_contiguous() + and w1.device == shard_logits.device + and R >= 16 + and (R & (R - 1)) == 0 + and R <= 1024 + and first_prev.dtype in _INT + and tokens.dtype in _INT + and first_prev.is_contiguous() + and tokens.is_contiguous() + and first_prev.numel() == B + and tuple(tokens.shape) == (B, K) + and B > 0 + and K > 0 + ) + + +def prev_embeddings( + first_prev: torch.Tensor, tokens: torch.Tensor, w1: torch.Tensor +) -> torch.Tensor: + """[B * K, R] bf16: row b * K + k is w1[prev] with prev = first_prev[b] (k = 0) or tokens[b, k - 1], and a zero + row for a prev outside [0, V) -- bitwise markov_prev_embeddings of the in-tree prevs.""" + B, K = tokens.shape + V, R = w1.shape + out = torch.empty(B * K, R, dtype=w1.dtype, device=w1.device) + _prev_embeddings_kernel[(B * K,)](first_prev, tokens, w1, out, K, V, R=R, num_warps=1) + return out + + +def widen_add(shard_logits: torch.Tensor, bias: torch.Tensor) -> torch.Tensor: + """fp32 ``shard_logits.float() + bias.to(float32)`` ([B, K, N] bf16 each, contiguous), bitwise torch's.""" + B, K, N = shard_logits.shape + if bias.dtype != torch.bfloat16 or not bias.is_contiguous() or bias.numel() != B * K * N: + raise ValueError("widen_add: bias must be contiguous bf16 with the logits' element count") + out = torch.empty(B, K, N, dtype=torch.float32, device=shard_logits.device) + grid = (B * K, triton.cdiv(N, WIDEN_BLOCK)) + _widen_add_kernel[grid](shard_logits, bias, out, N, BLOCK=WIDEN_BLOCK, num_warps=8) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_markov_chain.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_markov_chain.py new file mode 100644 index 000000000000..d32a58d50c15 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_markov_chain.py @@ -0,0 +1,651 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The DSpark draft head's vanilla-Markov chain over vocab-sharded draft logits, in one launch per pass. + +For every request b and block position k (greedy): ``bias = bf16(w2_shard . w1[prev])``, ``value = bf16(base[b, k] ++ bias)``, ``token[b, k]`` = the first maximum of ``value`` over the whole vocabulary (every rank's shard), ``prev = +token[b, k]``; ``prev`` for k = 0 is the anchor token. The in-tree chain gathers the full-vocabulary logits and runs +K launches of GEMV + add + argmax on every rank; this is one launch per rank, exchanging one (value, index) entry per +CTA and position over a multicast NVLink buffer. + +Shapes: the kernel splits a shard over an even number of co-resident CTAs, each owning a multiple of 64 rows (at most +``SMEM_ROW_BUDGET``). A shard that does not split so is padded: ``padded_rows`` gives the padded row count; the +caller allocates the draft logits with that last dimension (``pad_logits_buffer``, tail filled with -inf, never a +maximum) and pads ``w2_shard`` once at load (``pad_w2_shard``, zero rows). + +Compiled on the first call per shape, which must run outside CUDA-graph capture (``precompile``). The exchange +workspace is collective on its first use: every rank of the TP group calls ``workspace_args`` at the same point. + +Variants: batch-1 calls can take ``dspark_markov_chain_kernel_b1`` (``enable_batch1_variant``) and, on top of it, its +LL flag-in-payload exchange ``dspark_markov_chain_kernel_b1_ll`` (``enable_ll_exchange``); all are bitwise the same +tokens. Every rank must run the same variant for a call (the switches are process-wide; set them the same way on +every rank). +""" + +from __future__ import annotations + +import importlib +import os +import threading +from typing import Dict, Optional, Sequence, Tuple + +import torch +import torch.nn.functional as F + +from .._op import Arch, Cell, OpWrapper + +MAX_BLOCK = 8 # capacity: fixed: the workspace's positions per call +MAX_GRID = 152 # capacity: fixed: the workspace's CTAs per rank (all co-resident) +CERTIFIED_BATCH = ( + 8 # the largest batch the cells drive (the kernel takes up to the kernel module's MAX_BATCH) +) + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} +_workspaces: Dict[object, dict] = {} +_modules: Dict[str, object] = {} +_dummy: Dict[torch.device, torch.Tensor] = {} +_const_inputs: Dict[tuple, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {} + + +_batch1_variant = [ + False +] # S4B8: the batch-1 calls take dspark_markov_chain_kernel_b1 (see enable_batch1_variant) + + +_ll_exchange = [ + False +] # MARKOV_LL: the batch-1 calls take dspark_markov_chain_kernel_b1_ll (see enable_ll_exchange) + + +_cached_inputs = [ + False +] # TAILBK_SPEC: the kernel's read-only accepted inputs come from a per-shape cache (see enable_cached_inputs) +_cached_taken = [False] # a call took its accepted inputs from the cache (see cached_inputs_taken) + + +def _kernel_module(batch1: bool = False, ll: bool = False): + """The CuTe DSL kernel beside this file (or, for ``batch1``, its batch-1 variant; with ``ll`` too, the batch-1 + variant with the LL exchange), imported on first use (it pulls in cutlass).""" + name = ("kernel_b1_ll" if ll else "kernel_b1") if batch1 else "kernel" + mod = _modules.get(name) + if mod is None: + mod = _modules[name] = importlib.import_module(f".dspark_markov_chain_{name}", __package__) + return mod + + +_fp32_sum = [ + True +] # the in-tree chain's arithmetic: argmax on the unrounded fp32 sum (see enable_fp32_sum) + + +def enable_fp32_sum(enabled: bool = True) -> bool: + """Compare base + bias as the unrounded fp32 sum, as the in-tree chain does: its draft logits are fp32 + (LogitsProcessor's .float()) and it widens the bf16 bias, so its argmax sees fp32 sums. Off: round the sum to + bf16 first (a bf16-logits reference). On by default; set before the first call of a shape.""" + _fp32_sum[0] = bool(enabled) + return _fp32_sum[0] + + +def enable_batch1_variant(enabled: bool = True) -> bool: + """Run batch-1 calls on dspark_markov_chain_kernel_b1: the same chain with the position loop's base loads issued + before the W1 wait and the exchange polled in one first pass (bitwise the same tokens; faster at one request, + slower at more, so only batch 1 takes it). Off by default; set before the first call of a shape.""" + _batch1_variant[0] = bool(enabled) + return _batch1_variant[0] + + +def enable_ll_exchange(enabled: bool = True) -> bool: + """Run the batch-1 variant's calls on dspark_markov_chain_kernel_b1_ll: the same chain with the per-position + exchange as one 16-byte {value, tag, index, tag} multicast store per CTA (no fence, no flag word) polled and + reduced by one warp (no CTA barrier after the poll); bitwise the same tokens. Only batch-1 calls with the batch-1 + variant on take it. Off by default; set it the same way on every rank, before the first call of a shape.""" + _ll_exchange[0] = bool(enabled) + return _ll_exchange[0] + + +def enable_cached_inputs(enabled: bool = True) -> bool: + """Take the kernel's read-only accepted inputs (``num_accepted`` ones, ``accepted_rows`` arange, ``accepted`` + zeros; they only feed ``next_new``, which the op discards) from a cache filled on the first call per (device, + batch, block) outside CUDA-graph capture, instead of three fill/arange launches per call. The kernels never write + them, so the tokens are bitwise the same. Off by default; set it before the first call of a shape.""" + _cached_inputs[0] = bool(enabled) + return _cached_inputs[0] + + +def cached_inputs_taken() -> bool: + """Whether a call since process start took its accepted inputs from the cache (the engagement of + ``enable_cached_inputs``; a cache fill alone does not count).""" + return _cached_taken[0] + + +def _accepted_inputs( + batch: int, block: int, device +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """(num_accepted ones [B], accepted_rows arange [B], accepted zeros [B, block + 1]), int32: cached under + ``enable_cached_inputs`` (filled outside capture only; a miss under capture allocates afresh, uncached).""" + key = (device, batch, block) + inputs = _const_inputs.get(key) if _cached_inputs[0] else None + if inputs is not None: + _cached_taken[0] = True + else: + inputs = ( + torch.ones(batch, dtype=torch.int32, device=device), + torch.arange(batch, dtype=torch.int32, device=device), + torch.zeros(batch, block + 1, dtype=torch.int32, device=device), + ) + if _cached_inputs[0] and not torch.cuda.is_current_stream_capturing(): + _const_inputs[key] = inputs + return inputs + + +def ll_exchange_engaged(batch: int) -> bool: + """Whether a call of ``batch`` requests runs the LL exchange kernel under the current switches.""" + return _batch1_variant[0] and _ll_exchange[0] and int(batch) == 1 + + +def _sm_count(device) -> int: + return torch.cuda.get_device_properties(device).multi_processor_count + + +def pick_grid(shard: int, block: int, batch: int, device=None) -> int: + """CTAs for a (padded) shard: 128 when supported, else the largest supported count up to the SM count; 0 if + none.""" + kern = _kernel_module() + limit = min(MAX_GRID, _sm_count(device if device is not None else torch.cuda.current_device())) + if 128 <= limit and kern.supports(shard, 128, block, batch): + return 128 + for grid in range(limit, 0, -1): + if kern.supports(shard, grid, block, batch): + return grid + return 0 + + +def padded_rows(shard: int, max_batch: int, device=None) -> int: + """The smallest row count >= ``shard`` that the kernel splits for every batch <= ``max_batch`` and block <= + MAX_BLOCK; 0 if none fits the shared-memory budget.""" + kern = _kernel_module() + step = kern.WARPS * kern.BLOCK_ROWS + rows = -(-shard // step) * step + limit = rows + MAX_GRID * step + while rows <= limit: + if all(pick_grid(rows, MAX_BLOCK, b, device) > 0 for b in (1, max_batch)): + return rows + rows += step + return 0 + + +def pad_w2_shard(w2_shard: torch.Tensor, rows: int) -> torch.Tensor: + """``w2_shard`` [S, R] bf16 with zero rows appended up to ``rows`` (done once at load).""" + out = w2_shard.new_zeros(rows, w2_shard.shape[1]) + out[: w2_shard.shape[0]].copy_(w2_shard) + return out + + +def pad_logits_buffer( + batch: int, block: int, rows: int, device, dtype=torch.bfloat16 +) -> torch.Tensor: + """A [batch, block, rows] draft-logits buffer whose padding columns hold -inf (the draft head writes [..., :S]).""" + return torch.full((batch, block, rows), float("-inf"), dtype=dtype, device=device) + + +def workspace_args(mapping, group=None) -> dict: + """This TP group's exchange workspace as the call's ``ws_*``, ``rank``, ``slots`` and ``buf_words`` arguments. + + Three rotating Lamport buffers (every word 0x80000000 when armed) and, after them, the LL variant's two regions + of the same size (zero-filled: tag 0 is never a call's tag) behind one multicast mapping, and their flag words + (the Lamport variants' [0..3], the LL variant's [4..7]); built once per (mapping, group) and reused. Collective + on first use, outside CUDA-graph capture: every rank of the TP group calls it at the same point. ``group`` + overrides the communicator (a torch.distributed group when MPI is disabled). + """ + key = (mapping, id(group) if group is not None else None) + ws = _workspaces.get(key) + if ws is not None: + return ws["args"] + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "dspark_markov_chain: the workspace must be allocated outside CUDA-graph capture" + ) + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + + kern = _kernel_module() + ll = _kernel_module(batch1=True, ll=True) + world = mapping.tp_size + buf_words = kern.buffer_words(MAX_BLOCK, kern.MAX_BATCH, world, MAX_GRID) + comm = group if group is not None else _get_mnnvl_workspace_comm(mapping) + use_fabric_handle = ( + os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" or mapping.is_multi_node() + ) + error: Optional[Exception] = None + try: + lamport_words = kern.BUFFERS * buf_words + words = ll.ll_workspace_words(buf_words) + handle = _make_mnnvl_mcast_buffer(comm, words * 4, mapping, use_fabric_handle) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + with torch.inference_mode(): + uc[:lamport_words].fill_(kern.EMPTY_WORD) + uc[lamport_words:].zero_() + flags = torch.zeros( + max(kern.FLAG_WORDS, ll.LL_FLAG_WORDS), dtype=torch.int32, device=uc.device + ) + torch.cuda.synchronize() + args = dict( + ws_uc=uc, + ws_mc=mc, + ws_flags=flags, + rank=mapping.tp_rank, + slots=world, + buf_words=buf_words, + ) + ws = dict(handle=handle, comm=comm, args=args) + except Exception as exc: # noqa: BLE001 -- reported to every rank below, then re-raised + error = exc + # Also the barrier that keeps any rank from pushing into a peer's buffer before the peer has armed it. + if not _mnnvl_workspace_all_succeeded(comm, error is None): + raise RuntimeError( + "dspark_markov_chain Lamport buffers failed on at least one rank" + ) from error + _workspaces[key] = ws + return ws["args"] + + +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=0) + + +def supports(base: torch.Tensor, w1: torch.Tensor, w2_shard: torch.Tensor) -> bool: + """Whether the kernel takes these operands (dtypes, ranks, the padded shard's split); a caller that gets False + runs the in-tree chain.""" + kern = _kernel_module() + if base.dim() != 3 or base.dtype != torch.bfloat16 or not base.is_cuda: + return False + batch, block, rows = base.shape + return ( + w1.dtype == torch.bfloat16 + and w2_shard.dtype == torch.bfloat16 + and w1.dim() == 2 + and w1.shape[1] == kern.MARKOV_RANK + and tuple(w2_shard.shape) == (rows, kern.MARKOV_RANK) + and 0 < block <= MAX_BLOCK + and 0 < batch <= kern.MAX_BATCH + and pick_grid(rows, block, batch, base.device) > 0 + ) + + +def _launch( + base: torch.Tensor, + first_prev: torch.Tensor, + w1: torch.Tensor, + w2_shard: torch.Tensor, + shard_offset: int, + ws_uc: torch.Tensor, + ws_mc: torch.Tensor, + ws_flags: torch.Tensor, + rank: int, + slots: int, + buf_words: int, + kv_lens: Optional[torch.Tensor] = None, + rewind: Optional[torch.Tensor] = None, + rewind_first: int = 0, +) -> torch.Tensor: + """One chain launch; with ``kv_lens`` and ``rewind``, also ``kv_lens[rewind_first + i] = max(kv_lens[...] - + rewind[i], 0)`` for every i (after the kernel's grid-dependency wait).""" + import cuda.bindings.driver as cuda_driver + import cutlass.cute as cute + + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + kern = _kernel_module() + if not supports(base, w1, w2_shard): + raise ValueError(f"dspark_markov_chain: unsupported operands base {tuple(base.shape)} {base.dtype}, " + f"w2_shard {tuple(w2_shard.shape)}") # fmt: skip + if first_prev.dtype != torch.int64 or first_prev.numel() != base.shape[0]: + raise ValueError("dspark_markov_chain: first_prev must be int64 [B]") + batch, block, rows = base.shape + batch1 = _batch1_variant[0] and batch == 1 + ll = batch1 and _ll_exchange[0] + grid = pick_grid(rows, block, batch, base.device) + if kern.buffer_words(block, batch, slots, grid) > buf_words: + raise ValueError("dspark_markov_chain: the exchange does not fit the workspace buffers") + if ll: + ll_kern = _kernel_module(batch1=True, ll=True) + if ( + ll_kern.ll_region_words(block, slots, grid) > buf_words + or ws_uc.numel() < ll_kern.ll_workspace_words(buf_words) + or ws_flags.numel() < ll_kern.LL_FLAG_WORDS + ): + raise ValueError("dspark_markov_chain: the LL exchange does not fit the workspace") + device = base.device + tokens = torch.empty(batch, block, dtype=torch.int32, device=device) + next_new = torch.empty(batch, block + 1, dtype=torch.int32, device=device) + dummy = _dummy.get(device) + if dummy is None: + # The kernel's optional outputs and inputs (corrected logits, accepted rows, KV-length rewind) are unused + # here: one zero word stands for each, and an empty rewind leaves it untouched. + dummy = _dummy[device] = torch.zeros(4, dtype=torch.int32, device=device) + rewind_count = 0 if rewind is None else int(rewind.numel()) + if rewind_count: + if not supports_rewind(kv_lens, rewind, rewind_first): + raise ValueError( + "dspark_markov_chain: unsupported KV-length rewind (see supports_rewind())" + ) + else: + kv_lens, rewind = dummy, dummy + nacc, rows_idx, acc = _accepted_inputs(batch, block, device) + args = ( + _arg(base.contiguous().view(-1).view(torch.int16)), + _arg(first_prev.contiguous().view(-1), align=8), + _arg(w1.contiguous().view(-1).view(torch.int32)), + _arg(w2_shard.contiguous().view(-1).view(torch.int32)), + _arg(dummy.view(torch.float32)), + _arg(tokens.view(-1)), + _arg(next_new.view(-1)), + _arg(acc.view(-1), align=4), + _arg(nacc, align=4), + _arg(rows_idx, align=4), + _arg(kv_lens, align=4), + _arg(rewind, align=4), + _arg(ws_uc.view(-1)), + _arg(ws_mc.view(-1)), + _arg(ws_flags.view(-1)), + ) + scalars = ( + int(w1.shape[0]), + int(shard_offset), + int(rank), + int(block + 1), + int(rewind_first), + rewind_count, + ) + use_pdl = get_env_enable_pdl() + stream = cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) + consts = (grid, rows // grid, rows, block, batch, slots, 1, buf_words, True, _fp32_sum[0]) + key = consts + (use_pdl, batch1, ll) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "trtllm::dspark_markov_chain must run once per shape outside CUDA-graph capture first" + ) + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + _kernel_module(batch1, ll).k3_markov, *args, *scalars, *consts, use_pdl, stream + ) + fn(*args, *scalars, stream) + return tokens + + +@torch.library.custom_op("trtllm::dspark_markov_chain", mutates_args=()) +def _dspark_markov_chain( + base: torch.Tensor, + first_prev: torch.Tensor, + w1: torch.Tensor, + w2_shard: torch.Tensor, + shard_offset: int, + ws_uc: torch.Tensor, + ws_mc: torch.Tensor, + ws_flags: torch.Tensor, + rank: int, + slots: int, + buf_words: int, +) -> torch.Tensor: + return _launch( + base, first_prev, w1, w2_shard, shard_offset, ws_uc, ws_mc, ws_flags, rank, slots, buf_words + ) + + +@_dspark_markov_chain.register_fake +def _(base, first_prev, w1, w2_shard, shard_offset, ws_uc, ws_mc, ws_flags, rank, slots, buf_words): + return base.new_empty((base.shape[0], base.shape[1]), dtype=torch.int32) + + +@torch.library.custom_op("trtllm::dspark_markov_chain_rewind", mutates_args=("kv_lens",)) +def _dspark_markov_chain_rewind( + base: torch.Tensor, + first_prev: torch.Tensor, + w1: torch.Tensor, + w2_shard: torch.Tensor, + shard_offset: int, + ws_uc: torch.Tensor, + ws_mc: torch.Tensor, + ws_flags: torch.Tensor, + rank: int, + slots: int, + buf_words: int, + kv_lens: torch.Tensor, + rewind: torch.Tensor, + rewind_first: int, +) -> torch.Tensor: + return _launch( + base, first_prev, w1, w2_shard, shard_offset, ws_uc, ws_mc, ws_flags, rank, slots, buf_words, + kv_lens, rewind, rewind_first, + ) # fmt: skip + + +@_dspark_markov_chain_rewind.register_fake +def _(base, first_prev, w1, w2_shard, shard_offset, ws_uc, ws_mc, ws_flags, rank, slots, buf_words, kv_lens, + rewind, rewind_first): # fmt: skip + return base.new_empty((base.shape[0], base.shape[1]), dtype=torch.int32) + + +def supports_rewind( + kv_lens: Optional[torch.Tensor], rewind: Optional[torch.Tensor], rewind_first: int +) -> bool: + """Whether a chain call can also apply the KV-length rewind ``kv_lens[first + i] = max(kv_lens[first + i] - + rewind[i], 0)``: int32 contiguous CUDA operands on one device, 1..``kernel MAX_BATCH`` amounts, in range. A caller + that gets False applies the rewind itself.""" + if kv_lens is None or rewind is None: + return False + n = rewind.numel() + return ( + kv_lens.is_cuda + and rewind.is_cuda + and kv_lens.device == rewind.device + and kv_lens.dtype == torch.int32 + and rewind.dtype == torch.int32 + and kv_lens.is_contiguous() + and rewind.is_contiguous() + and 0 < n <= _kernel_module().MAX_BATCH + and 0 <= int(rewind_first) + and int(rewind_first) + n <= kv_lens.numel() + ) + + +def dspark_markov_chain_rewind( + base: torch.Tensor, + first_prev: torch.Tensor, + w1: torch.Tensor, + w2_shard: torch.Tensor, + shard_offset: int, + kv_lens: torch.Tensor, + rewind: torch.Tensor, + rewind_first: int, + *, + ws_uc: torch.Tensor, + ws_mc: torch.Tensor, + ws_flags: torch.Tensor, + rank: int, + slots: int, + buf_words: int, +) -> torch.Tensor: + """``dspark_markov_chain``'s tokens (the same launch, bitwise the same tokens) plus the target's KV-length rewind + after the draft forward in the same launch (TAILBK_SPEC site H): ``kv_lens[rewind_first + i] = + max(kv_lens[rewind_first + i] - rewind[i], 0)``, applied after the kernel's grid-dependency wait (every reader of + ``kv_lens`` in the draft forward finished). Check ``supports_rewind`` first; the call is a collective as + ``dspark_markov_chain``'s, and each rank rewinds its own ``kv_lens``. The kernel module's compile is shared with + ``dspark_markov_chain`` (the rewind's operands are runtime arguments).""" + return torch.ops.trtllm.dspark_markov_chain_rewind( + base, first_prev, w1, w2_shard, int(shard_offset), ws_uc, ws_mc, ws_flags, int(rank), int(slots), + int(buf_words), kv_lens, rewind, int(rewind_first), + ) # fmt: skip + + +def precompile( + mapping, + rows: int, + w1: torch.Tensor, + w2_shard: torch.Tensor, + shard_offset: int, + blocks: Tuple[int, ...], + max_batch: int, + group=None, +) -> None: + """Compile every (batch <= ``max_batch``, block in ``blocks``) shape outside capture (collective: every rank).""" + ws = workspace_args(mapping, group) + for block in blocks: + for batch in range(1, max_batch + 1): + base = pad_logits_buffer(batch, block, rows, w2_shard.device) + prev = torch.zeros(batch, dtype=torch.long, device=w2_shard.device) + dspark_markov_chain(base, prev, w1, w2_shard, shard_offset, **ws) + torch.cuda.synchronize() + + +class DsparkMarkovChain(OpWrapper): + """Greedy draft tokens [B, K] (int32) of the DSpark vanilla-Markov head, from this rank's padded vocab shard + of the block logits. + + ``base`` [B, K, rows] bf16: this rank's shard of the draft head's block logits (columns ``[0, shard)``), the + padding columns -inf. ``first_prev`` [B] int64: the anchor tokens. ``w1`` [V, 256] bf16: ``markov_w1``, whole. + ``w2_shard`` [rows, 256] bf16: this rank's rows of ``markov_w2``, zero-padded. ``shard_offset``: the vocabulary + index of this rank's first column. ``ws_*``, ``rank``, ``slots``, ``buf_words``: ``workspace_args``'s. + + Every rank returns the same tokens. The call is a collective: every rank of the group calls it with the same + (B, K), in the same order. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # MiniMax-M3.1's DSpark drafter: vocabulary 200064, Markov rank 256, K = 3 and 7 draft positions, the decode + # batches the M3.1 targets serve with it (1..8). TP4 is one GB300 tray, TP8 two. + _M31 = dict(vocab=200064, batches=tuple(range(1, CERTIFIED_BATCH + 1))) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="M3.1 TP4, K = 3: the 50016-row shard, padded to 50688", + spec=dict(world=4, block=3, **_M31), + ), + Cell( + why="M3.1 TP4, K = 7: seven chained exchanges per call", + spec=dict(world=4, block=7, **_M31), + ), + Cell( + why="a shard that splits without padding (32768 / 4 = 8192 rows)", + spec=dict(world=4, block=3, vocab=32768, batches=(1, 8)), + ), + Cell( + why="M3.1 TP8 over two trays, K = 3: the 25008-row shard, padded to 25088; fabric-handle multicast", + spec=dict(world=8, block=3, **_M31), + ), + Cell(why="M3.1 TP8 over two trays, K = 7", spec=dict(world=8, block=7, **_M31)), + ) + + note = """ + Greedy only, and only the vanilla Markov head: no d2t map, no candidate selector, no sampling. A caller with + any of those runs the in-tree chain. + + The kernel takes batches up to its module's MAX_BATCH (16) and blocks up to MAX_BLOCK (8); the cells certify + batches 1..8. Measured: at 16 requests the one-launch chain is 0.93x the in-tree chain's time, so the M3.1 + targets serve batches up to 8 only. + + The kernel also computes next_new_tokens, the accepted-row bookkeeping and the target's KV-length rewind of + the K3 fused verify; this entry passes none of them (one dummy word each) and returns only the tokens. + """ + + def __call__( + self, + base: torch.Tensor, + first_prev: torch.Tensor, + w1: torch.Tensor, + w2_shard: torch.Tensor, + shard_offset: int, + ws_uc: torch.Tensor, + ws_mc: torch.Tensor, + ws_flags: torch.Tensor, + rank: int, + slots: int, + buf_words: int, + ) -> torch.Tensor: + return torch.ops.trtllm.dspark_markov_chain( + base, + first_prev, + w1, + w2_shard, + int(shard_offset), + ws_uc, + ws_mc, + ws_flags, + int(rank), + int(slots), + int(buf_words), + ) + + def reference( + self, + base_in_rank_order: Sequence[torch.Tensor], + first_prev: torch.Tensor, + w1: torch.Tensor, + w2_shard_in_rank_order: Sequence[torch.Tensor], + shard: int, + ) -> torch.Tensor: + """The DSpark worker's in-tree chain on the gathered full-vocabulary logits: per position ``bias = + F.linear(w1[prev], w2)`` in bf16, ``logits + bias`` in the logits' dtype, ``torch.argmax`` (first maximum). + With fp32_sum (default) the logits are widened to fp32 first, as LogitsProcessor's ``.float()`` hands them + to the in-tree chain, so the bias is widened and the sum is not rounded; without it, bf16 logits + bf16 bias. + + Not more accurate than the in-tree chain: the tokens must equal its tokens bit for bit, so this is that + chain. The documented collective exception to mirroring ``__call__``: it takes every rank's ``base`` and + ``w2_shard`` in rank order, and ``shard``, the unpadded columns of each. + """ + full = torch.cat([b[..., :shard] for b in base_in_rank_order], dim=-1) + if _fp32_sum[0]: + full = full.float() + w2 = torch.cat([w[:shard] for w in w2_shard_in_rank_order], dim=0) + prev = first_prev + tokens = [] + for k in range(full.shape[1]): + bias = F.linear(w1[prev], w2).to(full.dtype) + prev = torch.argmax(full[:, k] + bias, dim=-1) + tokens.append(prev) + return torch.stack(tokens, 1).to(torch.int32) + + def is_valid( + self, + base: torch.Tensor, + first_prev: torch.Tensor, + w1: torch.Tensor, + w2_shard: torch.Tensor, + shard_offset: int, + ws_uc: torch.Tensor, + ws_mc: torch.Tensor, + ws_flags: torch.Tensor, + rank: int, + slots: int, + buf_words: int, + ) -> None: + # The kernel takes every column of `base` as a candidate and offsets its index by `shard_offset`: a padding + # column that is not -inf can win and return an index past this rank's shard, and a wrong offset returns + # another rank's token ids. Neither is checked on the way to the kernel. + vocab = w1.shape[0] + assert vocab % slots == 0, (vocab, slots) + shard = vocab // slots + assert shard_offset == rank * shard, (shard_offset, rank, shard) + assert base.shape[-1] >= shard, (base.shape, shard) + assert bool(torch.isneginf(base[..., shard:]).all()), ( + "the padding columns of base must hold -inf" + ) + # An anchor outside the vocabulary indexes `w1` past its end inside the kernel. + assert bool(((first_prev >= 0) & (first_prev < vocab)).all()), ( + "first_prev outside [0, vocab)" + ) + + +dspark_markov_chain = DsparkMarkovChain() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_markov_chain_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_markov_chain_kernel.py new file mode 100644 index 000000000000..a6bdf3861a22 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_markov_chain_kernel.py @@ -0,0 +1,583 @@ +# 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. +# ============================================================================= +# DSpark vanilla-Markov draft chain for a vocab-sharded draft head -- CTM (prims/cute) kernel +# ============================================================================= +# +# For every request b and block position k = 0 .. K-1 (greedy, temperature 0): +# bias[v] = bf16(markov_w2[v] . markov_w1[prev]) v in this rank's vocab shard +# value[v] = bf16(base[b, k, v] + bias[v]) (the reference's bf16 add: fp32 sum, one rounding) +# token[b, k] = the first maximum of value over the whole vocabulary (all ranks) +# prev = token[b, k] (prev for k = 0 is the anchor, first_prev[b]) +# so the K positions are a chain of dependent global argmaxes across the tensor-parallel ranks. +# +# Numerics are dsparkMarkovChainKernel's (the C++ kernel this replaces): lane l of a warp sums the 8 columns +# 8l .. 8l+7 of a row by FMAs from 0, the 32 lane sums are added by the xor tree 16, 8, 4, 2, 1, the sum is rounded +# to bf16 (cvt.rn), added to the base in fp32 and the sum rounded to bf16 (cvt.rn), as the reference chain's bf16 +# logits + bf16 bias does. Argmax order: larger value, then lower vocabulary index; NaN never +# wins. The order is total, so every split of the reduction gives the same token. +# +# Geometry: G CTAs of 256 threads, CTA c owning the RC = S / G rows [c RC, c RC + RC) of the shard (RC a multiple +# of 64). Warp w owns the RC / 8 contiguous rows [w RC / 8, ...), in blocks of 8 (row j's sum ends on lane 4j after +# the reduce-scatter form of the xor tree). The CTA's rows of markov_w2 are bulk-copied into shared memory before the +# grid-dependency wait (a weight) and stay there for all K positions. markov_w1[prev] (512 bytes per request) is one +# bulk copy per CTA and position into shared memory, issued by the thread that reduces request b's global argmax, on +# one phase of an mbarrier per position. +# +# Exchange (one NVLink hop per position): every CTA reduces its rows to one (value, index) entry per request and +# multicast-stores it into slot [k][b][rank][c] of every rank's Lamport buffer; every CTA then polls all +# W x G entries of [k][b] in its own rank's copy (the data is the ready signal: 0x80000000 = empty; values carry +# -0.0 as +0.0 and indices are non-negative, so no entry is the sentinel) and reduces them with the same order. +# Every CTA of every rank computes the same global argmax: no leader, no grid barrier, no publish word. +# +# Buffers: 3 rotating Lamport buffers. Call n uses buffer n % 3 and re-arms buffer (n - 1) % 3 (its readers, all +# CTAs of call n - 1, finished before this call's grid wait returned; its next writer is a remote rank's call n + 2, +# which needs this rank's call n + 1 pushes). flags: [0] buffer of this call, [1] CTAs arrived, [2] words the +# previous call used (to re-arm), [3] unused. The last CTA to arrive advances them. +# +# Roles: warp 0's elected lane issues the weight bulk copy (and the early dependent trigger); thread b < B pushes +# request b's CTA entry, reduces its global argmax and requests the next markov_w1 row; all 8 warps compute and poll +# (the chain inside a position is serial, and the poll is spread over all 256 threads: one load wave). Per position: +# rows → barrier → push → poll → barrier → (thread b) argmax, next row; no barrier between the argmax and the next +# position (its readers wait for the row's phase). +# ============================================================================= +"""CTM DSpark Markov chain: the draft head's greedy intra-block Markov bias, the per-position global argmax +across the tensor-parallel vocab shards (the draft tokens), next_new_tokens and the target's KV-length rewind after +the draft forward, in one launch.""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +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 + +THREADS = 256 +WARPS = THREADS // 32 +MARKOV_RANK = 256 # 32 lanes x 8 columns +ROW_WORDS = MARKOV_RANK // 2 # int32 words (bf16 pairs) per markov row +ROW_BYTES = MARKOV_RANK * 2 +BLOCK_ROWS = 8 # rows per warp block: one per epilogue lane +MAX_BATCH = 16 +MAX_BLOCK = 16 +BUFFERS = 3 +FLAG_WORDS = 4 +EMPTY_WORD = -(2**31) # 0x80000000: fp32 -0.0, never a pushed word +INT_MAX = 2**31 - 1 +NEG_INF_BITS = -8388608 # 0xFF800000 +COPY_ROWS = 32 # rows per weight bulk copy (16 KB) +SMEM_ROW_BUDGET = 400 # 200 KB of markov_w2 rows per CTA +SMEM_BYTES = 227 * 1024 + + +def rows_per_cta(shard: int, grid: int) -> int: + return shard // grid + + +def supports(shard: int, grid: int, block: int, batch: int) -> bool: + """Shapes the kernel runs: whole 8-row blocks for every warp (RC a multiple of 64), the CTA's markov_w2 rows + within the shared-memory budget, block <= 16, batch <= 16, an even entry count per exchange.""" + if grid <= 0 or shard % grid != 0: + return False + rc = rows_per_cta(shard, grid) + return ( + rc % (WARPS * BLOCK_ROWS) == 0 + and rc <= SMEM_ROW_BUDGET + and smem_bytes(rc, batch) <= SMEM_BYTES + and grid % 2 == 0 + and 0 < block <= MAX_BLOCK + and 0 < batch <= MAX_BATCH + ) + + +def smem_bytes(rows: int, batch: int) -> int: + """Shared memory of a CTA: its markov_w2 rows, the markov_w1[prev] rows, the two [b][warp] reductions, the + barriers, flags and prev (with the allocator's alignment padding).""" + return rows * ROW_BYTES + batch * ROW_BYTES + 2 * batch * WARPS * 8 + 1024 + + +def buffer_words(block: int, batch: int, slots: int, grid: int) -> int: + """Int32 words of one Lamport buffer: one (value, index) entry per position, request, rank slot and CTA.""" + return block * batch * slots * grid * 2 + + +@dsl_user_op +def _atomic_add_acq_rel(addr_i64, val, *, loc=None, ip=None): + """atom.acq_rel.gpu.global.add.u32, returning the old value.""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [addr_i64.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + "atom.acq_rel.gpu.global.add.u32 $0, [$1], $2;", "=r,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 _fma_rn(a, b, c, *, loc=None, ip=None): + """fma.rn.f32 (fmaf).""" + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip), c.ir_value(loc=loc, ip=ip)], + "fma.rn.f32 $0, $1, $2, $3;", "=f,f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + """add.rn.f32: an fp32 add the compiler cannot contract with a neighbouring multiply.""" + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "add.rn.f32 $0, $1, $2;", "=f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _pack_bf16x2(hi, lo, *, loc=None, ip=None): + """(bf16(hi) << 16) | bf16(lo), round to nearest even (cvt.rn, as __float2bfloat16_rn).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [hi.ir_value(loc=loc, ip=ip), lo.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", "=r,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _lo_f32(word): + """The bf16 in the low half of a word, as fp32 (exact).""" + return cutlass.Int32(word << cutlass.Int32(16)).bitcast(cutlass.Float32) + + +def _hi_f32(word): + """The bf16 in the high half of a word, as fp32 (exact).""" + return cutlass.Int32(word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + + +def _better(value, index, best_value, best_index): + """dsparkMarkovChainKernel's keepBetter order: larger value, or an equal value at a lower index. NaN never + compares better, so it never wins.""" + return (value > best_value) | ((value == best_value) & (index < best_index)) + + +def _keep_better(best_value, best_index, value, index): + take = _better(value, index, best_value, best_index) + return cutlass.Float32(cutlass.select_(take, value, best_value)), cutlass.Int32( + cutlass.select_(take, index, best_index) + ) + + +def _warp_best(value, index): + """Every lane gets the warp's best (value, index) (xor tree: the order is total, so all lanes agree).""" + for offset in [16, 8, 4, 2, 1]: + other_value = cute.arch.shuffle_sync_bfly(value, offset=offset, mask=-1, mask_and_clamp=31) + other_index = cute.arch.shuffle_sync_bfly(index, offset=offset, mask=-1, mask_and_clamp=31) + value, index = _keep_better(value, index, other_value, other_index) + return value, index + + +def _block_dot(s_w2, first_row, lane, w1f): + """markov_w2[first_row + j] . markov_w1[prev] for the block's 8 rows, in the C++ kernel's order: lane l sums the + 8 columns 8l .. 8l+7 by an FMA chain from 0, then the 32 lane sums are added in pairs at lane distances 16, 8, 4, 2, + 1 (the xor tree). Done as a reduce-scatter: at distances 16, 8 and 4 a lane keeps half of its rows and adds its + partner's partial of each kept row (the pairs of the full tree, so the same sums: 9 shuffles instead of 40), then + the full tree at 2 and 1. Returns row (lane >> 2)'s sum.""" + parts = [] + for r in range(BLOCK_ROWS): + w2v = s_w2.load(idx=(first_row + cutlass.Int32(r)) * cutlass.Int32(ROW_WORDS) + lane * cutlass.Int32(4), + vector_size=4, alignment=16) # fmt: skip + acc = cutlass.Float32(0.0) + for q in range(4): + word = cutlass.Int32(w2v[q]) + acc = _fma_rn(w1f[2 * q], _lo_f32(word), acc) + acc = _fma_rn(w1f[2 * q + 1], _hi_f32(word), acc) + parts.append(acc) + for offset in [16, 8, 4]: + upper = (lane & cutlass.Int32(offset)) != cutlass.Int32(0) + half = len(parts) // 2 + kept = [] + for i in range(half): + send = cutlass.Float32(cutlass.select_(upper, parts[i], parts[half + i])) + keep = cutlass.Float32(cutlass.select_(upper, parts[half + i], parts[i])) + kept.append( + _add_rn( + keep, + cute.arch.shuffle_sync_bfly(send, offset=offset, mask=-1, mask_and_clamp=31), + ) + ) + parts = kept + dot = parts[0] + for offset in [2, 1]: + dot = _add_rn( + dot, cute.arch.shuffle_sync_bfly(dot, offset=offset, mask=-1, mask_and_clamp=31) + ) + return dot + + +def _fetch_w1_row(s_w1, s_valid, w1, mbar, b, prev, vocab): + """Request b's markov_w1[prev] row into shared memory: one 512-byte bulk copy on the position's barrier phase (a + prev outside the vocabulary copies row 0 and clears s_valid[b]: the readers zero its bias). s_valid[b] is stored + before the arrive, whose release the readers' phase wait acquires.""" + valid = (prev >= cutlass.Int32(0)) & (prev < vocab) + row = cutlass.Int32(cutlass.select_(valid, prev, cutlass.Int32(0))) + s_valid.store(cutlass.Int32(cutlass.select_(valid, cutlass.Int32(1), cutlass.Int32(0))), idx=b) + prims.mbarrier_arrive_expect_tx(mbar, ROW_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_w1.subview(b * cutlass.Int32(ROW_WORDS)), + w1.subview(row * cutlass.Int32(ROW_WORDS)), + mbar, + ROW_BYTES, + ) + + +def _reduce_slots(slots_arr, first_word): + """The best of 8 (value bits, index) pairs at first_word .. in warp order.""" + value = cutlass.Int32(NEG_INF_BITS).bitcast(cutlass.Float32) + index = cutlass.Int32(INT_MAX) + for w in range(WARPS): + pair = slots_arr.load(idx=first_word + cutlass.Int32(w * 2), vector_size=2, alignment=8) + value, index = _keep_better( + value, index, cutlass.Int32(pair[0]).bitcast(cutlass.Float32), cutlass.Int32(pair[1]) + ) + return value, index + + +@cute.kernel +def k3_markov_kernel( + base: cutlass.Array, # fp32 [B * K * S] (or its int16 bf16 view with base_bf16): the draft head's shard logits + first_prev: cutlass.Array, # int64 [B]: the anchor token of each request + w1: cutlass.Array, # int32 words of bf16 markov_w1 [V, 256] + w2: cutlass.Array, # int32 words of bf16 markov_w2[shard rows] [S, 256] + corrected: cutlass.Array, # unused (kept for the launch signature) + tokens: cutlass.Array, # int32 [B * K], out + next_new: cutlass.Array, # int32 [B * (K + 1)], out + accepted: cutlass.Array, # int32 [*, accepted_stride]: accepted tokens + num_accepted: cutlass.Array, # int32 [B] + accepted_rows: cutlass.Array, # int32 [B]: the row of accepted of each request + kv_lens: cutlass.Array, # int32: the target's KV lengths, rewound here after the draft forward + rewind: cutlass.Array, # int32 [rewind_count]: subtracted from kv_lens[rewind_first ...], clamped at 0 + buf_uc: cutlass.Array, # int32 words: this rank's Lamport buffers + buf_mc: cutlass.Array, # int32 words: their multicast mapping + flags: cutlass.Array, # int32 [4] + vocab: cutlass.Int32, # rows of markov_w1: a prev outside [0, vocab) gets a zero bias + vocab_offset: cutlass.Int32, # first vocabulary index of this rank's shard + rank: cutlass.Int32, + accepted_stride: cutlass.Int32, + rewind_first: cutlass.Int32, + rewind_count: cutlass.Int32, # 0: no rewind + grid: cutlass.Constexpr[int], + rows: cutlass.Constexpr[int], # RC + shard: cutlass.Constexpr[int], # S + block: cutlass.Constexpr[int], # K + batch: cutlass.Constexpr[int], # B + slots: cutlass.Constexpr[int], # rank slots per (k, b): world x push_copies + push_copies: cutlass.Constexpr[ + int + ], # test only: each rank's entry goes into push_copies consecutive slots + buf_words: cutlass.Constexpr[int], + base_bf16: cutlass.Constexpr[bool], + fp32_sum: cutlass.Constexpr[bool], +): + rows_per_warp = rows // WARPS + blocks_per_warp = rows_per_warp // BLOCK_ROWS + entries = slots * grid # (value, index) entries per (k, b) + vectors = entries // 2 # 16-byte loads per (k, b) + vec_per_thread = (vectors + THREADS - 1) // THREADS + used_words = block * batch * entries * 2 + + tx, _, _ = cute.arch.thread_idx() + cta, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tx % cutlass.Int32(32) + + s_w2 = cutlass.Array( + cutlass.Int32, rows * ROW_WORDS, space=cutlass.AddressSpace.smem, alignment=128 + ) + # [b]: markov_w1[prev] of the position (one bulk copy per request, issued by the thread that learns prev). + s_w1 = cutlass.Array( + cutlass.Int32, batch * ROW_WORDS, space=cutlass.AddressSpace.smem, alignment=128 + ) + s_mbar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + s_mbar_w1 = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + # [b][warp] (value bits, index): the CTA reduction of the rows, then of the polled entries. + s_row_best = cutlass.Array( + cutlass.Int32, batch * WARPS * 2, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_poll_best = cutlass.Array( + cutlass.Int32, batch * WARPS * 2, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_flags = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_valid = cutlass.Array(cutlass.Int32, MAX_BATCH, space=cutlass.AddressSpace.smem, alignment=16) + + if warp == 0: + if prims.elect_sync(): + prims.mbarrier_init(s_mbar, 1) + prims.mbarrier_init(s_mbar_w1, batch) + prims.fence_mbarrier_init() + prims.barrier_cta_sync(0) + + # The CTA's markov_w2 rows: a weight, loaded before the grid dependency. + if warp == 0: + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(s_mbar, rows * ROW_BYTES) + for c in cutlass.range_constexpr((rows + COPY_ROWS - 1) // COPY_ROWS): + n = min(COPY_ROWS, rows - c * COPY_ROWS) + prims.cp_async_bulk_shared_cluster_global( + s_w2.subview(c * COPY_ROWS * ROW_WORDS), + w2.subview( + cta * cutlass.Int32(rows * ROW_WORDS) + + cutlass.Int32(c * COPY_ROWS * ROW_WORDS) + ), + s_mbar, + n * ROW_BYTES, + ) + # Dependents wait for this whole grid before reading its outputs. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + + prims.griddepcontrol(prims.GridDepAction.WAIT) + if tx == 0: + s_flags.store(flags.load(idx=0, is_volatile=True), idx=0) + s_flags.store(flags.load(idx=2, is_volatile=True), idx=2) + if tx < cutlass.Int32(batch): + _fetch_w1_row( + s_w1, s_valid, w1, s_mbar_w1, tx, cutlass.Int32(first_prev.load(idx=tx)), vocab + ) + prims.barrier_cta_sync(0) + cur = s_flags.load(idx=0) + cur_base = cur * cutlass.Int32(buf_words) + + # Re-arm the previous call's buffer (this CTA's 16-byte chunks of the words it used). + dirty_base = ((cur + cutlass.Int32(2)) % cutlass.Int32(BUFFERS)) * cutlass.Int32(buf_words) + dirty_words = s_flags.load(idx=2) + empty = cutlass.Int32(EMPTY_WORD) + w_clear = (cta * cutlass.Int32(THREADS) + tx) * cutlass.Int32(4) + while w_clear < dirty_words: + buf_uc.store((empty, empty, empty, empty), idx=dirty_base + w_clear, alignment=16) + w_clear += cutlass.Int32(grid * THREADS * 4) + + # next_new_tokens[b] = [accepted[row_b, num_accepted[b] - 1], tokens[b, :]] (-1 wraps, as torch indexing). + if cta == 0: + if tx < cutlass.Int32(batch): + col = num_accepted.load(idx=tx) - cutlass.Int32(1) + col = cutlass.Int32(cutlass.select_(col < cutlass.Int32(0), col + accepted_stride, col)) + row = accepted_rows.load(idx=tx) + next_new.store( + accepted.load(idx=row * accepted_stride + col), idx=tx * cutlass.Int32(block + 1) + ) + # The KV-length rewind after the draft forward (whose kernels, the readers of kv_lens, all finished before + # this grid's wait): kv_lens[first + b] = max(kv_lens[first + b] - rewind[b], 0). + if tx < rewind_count: + left = kv_lens.load(idx=rewind_first + tx) - rewind.load(idx=tx) + kv_lens.store(cutlass.Int32(cutlass.select_(left < cutlass.Int32(0), cutlass.Int32(0), left)), + idx=rewind_first + tx) # fmt: skip + + # The weights have landed (issued before the wait). + while not cute.arch.mbarrier_test_wait(s_mbar.data_ptr(), 0): + pass + + row0 = warp * cutlass.Int32(rows_per_warp) # first row of the warp within the CTA + shard_row0 = cta * cutlass.Int32(rows) + row0 # ... within the shard + for k in cutlass.range(block, unroll=1): + for b in cutlass.range(batch, unroll=1): + base_row = (b * cutlass.Int32(block) + k) * cutlass.Int32(shard) + # Lane 4j finishes row j of each 8-row block (the reduce-scatter leaves row j's sum on lanes 4j .. 4j+3): + # the warp's base logits, in flight during the wait below. + has_row = (lane & cutlass.Int32(3)) == cutlass.Int32(0) + lane_row = shard_row0 + (lane >> cutlass.Int32(2)) + base_vals = [] + for g in cutlass.range_constexpr(blocks_per_warp): + if cutlass.const_expr(base_bf16): + base_vals.append( + cutlass.Int32( + cutlass.Int32( + base.load(idx=base_row + lane_row + cutlass.Int32(g * BLOCK_ROWS)) + ) + << cutlass.Int32(16) + ).bitcast(cutlass.Float32) + ) + else: + base_vals.append( + base.load(idx=base_row + lane_row + cutlass.Int32(g * BLOCK_ROWS)) + ) + # markov_w1[prev] landed (the position's phase of the W1 barrier). + while not cute.arch.mbarrier_test_wait(s_mbar_w1.data_ptr(), k & cutlass.Int32(1)): + pass + valid = s_valid.load(idx=b) != cutlass.Int32(0) + # markov_w1[prev], columns 8 lane .. 8 lane + 7, as fp32 (zeros for an invalid prev). + w1v = s_w1.load( + idx=b * cutlass.Int32(ROW_WORDS) + lane * 4, vector_size=4, alignment=16 + ) + w1f = [] + for q in cutlass.range_constexpr(4): + word = cutlass.Int32( + cutlass.select_(valid, cutlass.Int32(w1v[q]), cutlass.Int32(0)) + ) + w1f.append(_lo_f32(word)) + w1f.append(_hi_f32(word)) + best_value = cutlass.Int32(NEG_INF_BITS).bitcast(cutlass.Float32) + best_index = cutlass.Int32(INT_MAX) + for g in cutlass.range_constexpr(blocks_per_warp): + dot = _block_dot(s_w2, row0 + cutlass.Int32(g * BLOCK_ROWS), lane, w1f) + bias = _lo_f32(_pack_bf16x2(dot, dot)) + # fp32_sum (default): the in-tree chain's arithmetic. Its logits are fp32 (LogitsProcessor's + # .float()) and the bf16 bias is widened, so the argmax compares the unrounded fp32 sum. Without + # fp32_sum: a bf16 reference (bf16 logits + bf16 bias in bf16) rounds that sum once. The corrected + # logits are not stored (the chain's consumers read only the tokens). + if cutlass.const_expr(fp32_sum): + value = _add_rn(base_vals[g], bias) + else: + value = _lo_f32( + _pack_bf16x2(_add_rn(base_vals[g], bias), _add_rn(base_vals[g], bias)) + ) + if has_row: + my_row = lane_row + cutlass.Int32(g * BLOCK_ROWS) + best_value, best_index = _keep_better( + best_value, best_index, value, vocab_offset + my_row + ) + best_value, best_index = _warp_best(best_value, best_index) + if lane == 0: + s_row_best.store((best_value.bitcast(cutlass.Int32), best_index), + idx=(b * cutlass.Int32(WARPS) + warp) * cutlass.Int32(2), alignment=8) # fmt: skip + prims.barrier_cta_sync(0) + # Thread b pushes request b's CTA maximum into every rank's buffer. + if tx < cutlass.Int32(batch): + cta_value, cta_index = _reduce_slots(s_row_best, tx * cutlass.Int32(WARPS * 2)) + bits = cta_value.bitcast(cutlass.Int32) + bits = cutlass.Int32( + cutlass.select_(bits == cutlass.Int32(EMPTY_WORD), cutlass.Int32(0), bits) + ) + for c in cutlass.range_constexpr(push_copies): + s_idx = rank * cutlass.Int32(push_copies) + cutlass.Int32(c) + word = cur_base + (((k * cutlass.Int32(batch) + tx) * cutlass.Int32(slots) + s_idx) + * cutlass.Int32(grid) + cta) * cutlass.Int32(2) # fmt: skip + buf_mc.store((bits, cta_index), idx=word, alignment=8) + # Drain the posted multicast store now (the polls below issue no release that would). + prims.fence_acq_rel(prims.MemScope.CLUSTER) + # Poll every rank's entries of [k][b] (this thread's 16-byte vectors, all loads in flight together) until + # none is empty; the pass that finds them all complete also reduces them. + for b in cutlass.range(batch, unroll=1): + region = cur_base + (k * cutlass.Int32(batch) + b) * cutlass.Int32(entries * 2) + best_value = cutlass.Int32(NEG_INF_BITS).bitcast(cutlass.Float32) + best_index = cutlass.Int32(INT_MAX) + pending = cutlass.Boolean(True) + while pending: + pass_value = cutlass.Int32(NEG_INF_BITS).bitcast(cutlass.Float32) + pass_index = cutlass.Int32(INT_MAX) + pass_pending = cutlass.Boolean(False) + for i in cutlass.range_constexpr(vec_per_thread): + v_idx = tx + cutlass.Int32(i * THREADS) + live = v_idx < cutlass.Int32(vectors) + got = buf_uc.load( + idx=region + + cutlass.Int32(cutlass.select_(live, v_idx, cutlass.Int32(0))) + * cutlass.Int32(4), + vector_size=4, + alignment=16, + is_volatile=True, + ) + for e in cutlass.range_constexpr(2): + vw = cutlass.Int32(got[2 * e]) + iw = cutlass.Int32(got[2 * e + 1]) + pass_pending = pass_pending | (live & ((vw == empty) | (iw == empty))) + take = live & _better( + vw.bitcast(cutlass.Float32), iw, pass_value, pass_index + ) + pass_value = cutlass.Float32( + cutlass.select_(take, vw.bitcast(cutlass.Float32), pass_value) + ) + pass_index = cutlass.Int32(cutlass.select_(take, iw, pass_index)) + best_value = pass_value + best_index = pass_index + pending = pass_pending + best_value, best_index = _warp_best(best_value, best_index) + if lane == 0: + s_poll_best.store((best_value.bitcast(cutlass.Int32), best_index), + idx=(b * cutlass.Int32(WARPS) + warp) * cutlass.Int32(2), alignment=8) # fmt: skip + prims.barrier_cta_sync(0) + if tx < cutlass.Int32(batch): + g_value, g_index = _reduce_slots(s_poll_best, tx * cutlass.Int32(WARPS * 2)) + # The next position's markov_w1 row, requested as soon as prev is known: every reader of s_w1 and + # s_valid for this position passed the barrier before the push. No barrier follows: the other threads + # go on to the next position's base logits and wait for the row's phase. s_poll_best is next written + # after the next push barrier, which this thread reaches only after these reads. + if k + cutlass.Int32(1) < cutlass.Int32(block): + _fetch_w1_row(s_w1, s_valid, w1, s_mbar_w1, tx, g_index, vocab) + if cta == 0: + tokens.store(g_index, idx=tx * cutlass.Int32(block) + k) + next_new.store(g_index, idx=tx * cutlass.Int32(block + 1) + k + cutlass.Int32(1)) + + # Every CTA read the flags before arriving; the last one advances them for the next call. + if tx == 0: + arrived = _atomic_add_acq_rel(flags.data_ptr(1).toint(), cutlass.Int32(1)) + if arrived == cutlass.Int32(grid - 1): + flags.store(cutlass.Int32(0), idx=1) + flags.store(cutlass.Int32(used_words), idx=2) + flags.store((cur + cutlass.Int32(1)) % cutlass.Int32(BUFFERS), idx=0) + + +@cute.jit +def k3_markov( + base: cute.Tensor, + first_prev: cute.Tensor, + w1: cute.Tensor, + w2: cute.Tensor, + corrected: cute.Tensor, + tokens: cute.Tensor, + next_new: cute.Tensor, + accepted: cute.Tensor, + num_accepted: cute.Tensor, + accepted_rows: cute.Tensor, + kv_lens: cute.Tensor, + rewind: cute.Tensor, + buf_uc: cute.Tensor, + buf_mc: cute.Tensor, + flags: cute.Tensor, + vocab: cutlass.Int32, + vocab_offset: cutlass.Int32, + rank: cutlass.Int32, + accepted_stride: cutlass.Int32, + rewind_first: cutlass.Int32, + rewind_count: cutlass.Int32, + grid: cutlass.Constexpr[int], + rows: cutlass.Constexpr[int], + shard: cutlass.Constexpr[int], + block: cutlass.Constexpr[int], + batch: cutlass.Constexpr[int], + slots: cutlass.Constexpr[int], + push_copies: cutlass.Constexpr[int], + buf_words: cutlass.Constexpr[int], + base_bf16: cutlass.Constexpr[bool], + fp32_sum: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + k3_markov_kernel( + base, first_prev, w1, w2, corrected, tokens, next_new, accepted, num_accepted, accepted_rows, kv_lens, + rewind, buf_uc, buf_mc, flags, vocab, vocab_offset, rank, accepted_stride, rewind_first, rewind_count, + grid, rows, shard, block, batch, slots, push_copies, buf_words, base_bf16, fp32_sum, + ).launch( + grid=[grid, 1, 1], + block=[THREADS, 1, 1], + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_markov_chain_kernel_b1.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_markov_chain_kernel_b1.py new file mode 100644 index 000000000000..97c83e4c01be --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_markov_chain_kernel_b1.py @@ -0,0 +1,616 @@ +# 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. +# ============================================================================= +# DSpark vanilla-Markov draft chain for a vocab-sharded draft head -- CTM (prims/cute) kernel, BATCH-1 VARIANT (S4B8) +# Identical to dspark_markov_chain_kernel.py except the position loop: every request's base loads are issued before the +# position's W1 wait, and the exchange is polled in one first pass for every request before the per-request fallback. +# Bitwise the same tokens (same per-row FMA order and roundings, total argmax order). Measured (TP8, S4B8): +# faster at one request (K3 19.7 -> 18.9 us, K7 38.5 -> 35.2), slower at more (register pressure), so the op runs it +# only for batch 1. +# ============================================================================= +# +# For every request b and block position k = 0 .. K-1 (greedy, temperature 0): +# bias[v] = bf16(markov_w2[v] . markov_w1[prev]) v in this rank's vocab shard +# value[v] = bf16(base[b, k, v] + bias[v]) (the reference's bf16 add: fp32 sum, one rounding) +# token[b, k] = the first maximum of value over the whole vocabulary (all ranks) +# prev = token[b, k] (prev for k = 0 is the anchor, first_prev[b]) +# so the K positions are a chain of dependent global argmaxes across the tensor-parallel ranks. +# +# Numerics are dsparkMarkovChainKernel's (the C++ kernel this replaces): lane l of a warp sums the 8 columns +# 8l .. 8l+7 of a row by FMAs from 0, the 32 lane sums are added by the xor tree 16, 8, 4, 2, 1, the sum is rounded +# to bf16 (cvt.rn), added to the base in fp32 and the sum rounded to bf16 (cvt.rn), as the reference chain's bf16 +# logits + bf16 bias does. Argmax order: larger value, then lower vocabulary index; NaN never +# wins. The order is total, so every split of the reduction gives the same token. +# +# Geometry: G CTAs of 256 threads, CTA c owning the RC = S / G rows [c RC, c RC + RC) of the shard (RC a multiple +# of 64). Warp w owns the RC / 8 contiguous rows [w RC / 8, ...), in blocks of 8 (row j's sum ends on lane 4j after +# the reduce-scatter form of the xor tree). The CTA's rows of markov_w2 are bulk-copied into shared memory before the +# grid-dependency wait (a weight) and stay there for all K positions. markov_w1[prev] (512 bytes per request) is one +# bulk copy per CTA and position into shared memory, issued by the thread that reduces request b's global argmax, on +# one phase of an mbarrier per position. +# +# Exchange (one NVLink hop per position): every CTA reduces its rows to one (value, index) entry per request and +# multicast-stores it into slot [k][b][rank][c] of every rank's Lamport buffer; every CTA then polls all +# W x G entries of [k][b] in its own rank's copy (the data is the ready signal: 0x80000000 = empty; values carry +# -0.0 as +0.0 and indices are non-negative, so no entry is the sentinel) and reduces them with the same order. +# Every CTA of every rank computes the same global argmax: no leader, no grid barrier, no publish word. +# +# Buffers: 3 rotating Lamport buffers. Call n uses buffer n % 3 and re-arms buffer (n - 1) % 3 (its readers, all +# CTAs of call n - 1, finished before this call's grid wait returned; its next writer is a remote rank's call n + 2, +# which needs this rank's call n + 1 pushes). flags: [0] buffer of this call, [1] CTAs arrived, [2] words the +# previous call used (to re-arm), [3] unused. The last CTA to arrive advances them. +# +# Roles: warp 0's elected lane issues the weight bulk copy (and the early dependent trigger); thread b < B pushes +# request b's CTA entry, reduces its global argmax and requests the next markov_w1 row; all 8 warps compute and poll +# (the chain inside a position is serial, and the poll is spread over all 256 threads: one load wave). Per position: +# rows → barrier → push → poll → barrier → (thread b) argmax, next row; no barrier between the argmax and the next +# position (its readers wait for the row's phase). +# ============================================================================= +"""CTM DSpark Markov chain: the draft head's greedy intra-block Markov bias, the per-position global argmax +across the tensor-parallel vocab shards (the draft tokens), next_new_tokens and the target's KV-length rewind after +the draft forward, in one launch.""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +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 + +THREADS = 256 +WARPS = THREADS // 32 +MARKOV_RANK = 256 # 32 lanes x 8 columns +ROW_WORDS = MARKOV_RANK // 2 # int32 words (bf16 pairs) per markov row +ROW_BYTES = MARKOV_RANK * 2 +BLOCK_ROWS = 8 # rows per warp block: one per epilogue lane +MAX_BATCH = 16 +MAX_BLOCK = 16 +BUFFERS = 3 +FLAG_WORDS = 4 +EMPTY_WORD = -(2**31) # 0x80000000: fp32 -0.0, never a pushed word +INT_MAX = 2**31 - 1 +NEG_INF_BITS = -8388608 # 0xFF800000 +COPY_ROWS = 32 # rows per weight bulk copy (16 KB) +SMEM_ROW_BUDGET = 400 # 200 KB of markov_w2 rows per CTA +SMEM_BYTES = 227 * 1024 + + +def rows_per_cta(shard: int, grid: int) -> int: + return shard // grid + + +def supports(shard: int, grid: int, block: int, batch: int) -> bool: + """Shapes the kernel runs: whole 8-row blocks for every warp (RC a multiple of 64), the CTA's markov_w2 rows + within the shared-memory budget, block <= 16, batch <= 16, an even entry count per exchange.""" + if grid <= 0 or shard % grid != 0: + return False + rc = rows_per_cta(shard, grid) + return ( + rc % (WARPS * BLOCK_ROWS) == 0 + and rc <= SMEM_ROW_BUDGET + and smem_bytes(rc, batch) <= SMEM_BYTES + and grid % 2 == 0 + and 0 < block <= MAX_BLOCK + and 0 < batch <= MAX_BATCH + ) + + +def smem_bytes(rows: int, batch: int) -> int: + """Shared memory of a CTA: its markov_w2 rows, the markov_w1[prev] rows, the two [b][warp] reductions, the + barriers, flags and prev (with the allocator's alignment padding).""" + return rows * ROW_BYTES + batch * ROW_BYTES + 2 * batch * WARPS * 8 + 1024 + + +def buffer_words(block: int, batch: int, slots: int, grid: int) -> int: + """Int32 words of one Lamport buffer: one (value, index) entry per position, request, rank slot and CTA.""" + return block * batch * slots * grid * 2 + + +@dsl_user_op +def _atomic_add_acq_rel(addr_i64, val, *, loc=None, ip=None): + """atom.acq_rel.gpu.global.add.u32, returning the old value.""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [addr_i64.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + "atom.acq_rel.gpu.global.add.u32 $0, [$1], $2;", "=r,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 _fma_rn(a, b, c, *, loc=None, ip=None): + """fma.rn.f32 (fmaf).""" + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip), c.ir_value(loc=loc, ip=ip)], + "fma.rn.f32 $0, $1, $2, $3;", "=f,f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + """add.rn.f32: an fp32 add the compiler cannot contract with a neighbouring multiply.""" + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "add.rn.f32 $0, $1, $2;", "=f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _pack_bf16x2(hi, lo, *, loc=None, ip=None): + """(bf16(hi) << 16) | bf16(lo), round to nearest even (cvt.rn, as __float2bfloat16_rn).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [hi.ir_value(loc=loc, ip=ip), lo.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", "=r,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _lo_f32(word): + """The bf16 in the low half of a word, as fp32 (exact).""" + return cutlass.Int32(word << cutlass.Int32(16)).bitcast(cutlass.Float32) + + +def _hi_f32(word): + """The bf16 in the high half of a word, as fp32 (exact).""" + return cutlass.Int32(word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + + +def _better(value, index, best_value, best_index): + """dsparkMarkovChainKernel's keepBetter order: larger value, or an equal value at a lower index. NaN never + compares better, so it never wins.""" + return (value > best_value) | ((value == best_value) & (index < best_index)) + + +def _keep_better(best_value, best_index, value, index): + take = _better(value, index, best_value, best_index) + return cutlass.Float32(cutlass.select_(take, value, best_value)), cutlass.Int32( + cutlass.select_(take, index, best_index) + ) + + +def _warp_best(value, index): + """Every lane gets the warp's best (value, index) (xor tree: the order is total, so all lanes agree).""" + for offset in [16, 8, 4, 2, 1]: + other_value = cute.arch.shuffle_sync_bfly(value, offset=offset, mask=-1, mask_and_clamp=31) + other_index = cute.arch.shuffle_sync_bfly(index, offset=offset, mask=-1, mask_and_clamp=31) + value, index = _keep_better(value, index, other_value, other_index) + return value, index + + +def _block_dot(s_w2, first_row, lane, w1f): + """markov_w2[first_row + j] . markov_w1[prev] for the block's 8 rows, in the C++ kernel's order: lane l sums the + 8 columns 8l .. 8l+7 by an FMA chain from 0, then the 32 lane sums are added in pairs at lane distances 16, 8, 4, 2, + 1 (the xor tree). Done as a reduce-scatter: at distances 16, 8 and 4 a lane keeps half of its rows and adds its + partner's partial of each kept row (the pairs of the full tree, so the same sums: 9 shuffles instead of 40), then + the full tree at 2 and 1. Returns row (lane >> 2)'s sum.""" + parts = [] + for r in range(BLOCK_ROWS): + w2v = s_w2.load(idx=(first_row + cutlass.Int32(r)) * cutlass.Int32(ROW_WORDS) + lane * cutlass.Int32(4), + vector_size=4, alignment=16) # fmt: skip + acc = cutlass.Float32(0.0) + for q in range(4): + word = cutlass.Int32(w2v[q]) + acc = _fma_rn(w1f[2 * q], _lo_f32(word), acc) + acc = _fma_rn(w1f[2 * q + 1], _hi_f32(word), acc) + parts.append(acc) + for offset in [16, 8, 4]: + upper = (lane & cutlass.Int32(offset)) != cutlass.Int32(0) + half = len(parts) // 2 + kept = [] + for i in range(half): + send = cutlass.Float32(cutlass.select_(upper, parts[i], parts[half + i])) + keep = cutlass.Float32(cutlass.select_(upper, parts[half + i], parts[i])) + kept.append( + _add_rn( + keep, + cute.arch.shuffle_sync_bfly(send, offset=offset, mask=-1, mask_and_clamp=31), + ) + ) + parts = kept + dot = parts[0] + for offset in [2, 1]: + dot = _add_rn( + dot, cute.arch.shuffle_sync_bfly(dot, offset=offset, mask=-1, mask_and_clamp=31) + ) + return dot + + +def _fetch_w1_row(s_w1, s_valid, w1, mbar, b, prev, vocab): + """Request b's markov_w1[prev] row into shared memory: one 512-byte bulk copy on the position's barrier phase (a + prev outside the vocabulary copies row 0 and clears s_valid[b]: the readers zero its bias). s_valid[b] is stored + before the arrive, whose release the readers' phase wait acquires.""" + valid = (prev >= cutlass.Int32(0)) & (prev < vocab) + row = cutlass.Int32(cutlass.select_(valid, prev, cutlass.Int32(0))) + s_valid.store(cutlass.Int32(cutlass.select_(valid, cutlass.Int32(1), cutlass.Int32(0))), idx=b) + prims.mbarrier_arrive_expect_tx(mbar, ROW_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_w1.subview(b * cutlass.Int32(ROW_WORDS)), + w1.subview(row * cutlass.Int32(ROW_WORDS)), + mbar, + ROW_BYTES, + ) + + +def _reduce_slots(slots_arr, first_word): + """The best of 8 (value bits, index) pairs at first_word .. in warp order.""" + value = cutlass.Int32(NEG_INF_BITS).bitcast(cutlass.Float32) + index = cutlass.Int32(INT_MAX) + for w in range(WARPS): + pair = slots_arr.load(idx=first_word + cutlass.Int32(w * 2), vector_size=2, alignment=8) + value, index = _keep_better( + value, index, cutlass.Int32(pair[0]).bitcast(cutlass.Float32), cutlass.Int32(pair[1]) + ) + return value, index + + +@cute.kernel +def k3_markov_kernel( + base: cutlass.Array, # fp32 [B * K * S] (or its int16 bf16 view with base_bf16): the draft head's shard logits + first_prev: cutlass.Array, # int64 [B]: the anchor token of each request + w1: cutlass.Array, # int32 words of bf16 markov_w1 [V, 256] + w2: cutlass.Array, # int32 words of bf16 markov_w2[shard rows] [S, 256] + corrected: cutlass.Array, # unused (kept for the launch signature) + tokens: cutlass.Array, # int32 [B * K], out + next_new: cutlass.Array, # int32 [B * (K + 1)], out + accepted: cutlass.Array, # int32 [*, accepted_stride]: accepted tokens + num_accepted: cutlass.Array, # int32 [B] + accepted_rows: cutlass.Array, # int32 [B]: the row of accepted of each request + kv_lens: cutlass.Array, # int32: the target's KV lengths, rewound here after the draft forward + rewind: cutlass.Array, # int32 [rewind_count]: subtracted from kv_lens[rewind_first ...], clamped at 0 + buf_uc: cutlass.Array, # int32 words: this rank's Lamport buffers + buf_mc: cutlass.Array, # int32 words: their multicast mapping + flags: cutlass.Array, # int32 [4] + vocab: cutlass.Int32, # rows of markov_w1: a prev outside [0, vocab) gets a zero bias + vocab_offset: cutlass.Int32, # first vocabulary index of this rank's shard + rank: cutlass.Int32, + accepted_stride: cutlass.Int32, + rewind_first: cutlass.Int32, + rewind_count: cutlass.Int32, # 0: no rewind + grid: cutlass.Constexpr[int], + rows: cutlass.Constexpr[int], # RC + shard: cutlass.Constexpr[int], # S + block: cutlass.Constexpr[int], # K + batch: cutlass.Constexpr[int], # B + slots: cutlass.Constexpr[int], # rank slots per (k, b): world x push_copies + push_copies: cutlass.Constexpr[ + int + ], # test only: each rank's entry goes into push_copies consecutive slots + buf_words: cutlass.Constexpr[int], + base_bf16: cutlass.Constexpr[bool], + fp32_sum: cutlass.Constexpr[bool], +): + rows_per_warp = rows // WARPS + blocks_per_warp = rows_per_warp // BLOCK_ROWS + entries = slots * grid # (value, index) entries per (k, b) + vectors = entries // 2 # 16-byte loads per (k, b) + vec_per_thread = (vectors + THREADS - 1) // THREADS + used_words = block * batch * entries * 2 + + tx, _, _ = cute.arch.thread_idx() + cta, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tx % cutlass.Int32(32) + + s_w2 = cutlass.Array( + cutlass.Int32, rows * ROW_WORDS, space=cutlass.AddressSpace.smem, alignment=128 + ) + # [b]: markov_w1[prev] of the position (one bulk copy per request, issued by the thread that learns prev). + s_w1 = cutlass.Array( + cutlass.Int32, batch * ROW_WORDS, space=cutlass.AddressSpace.smem, alignment=128 + ) + s_mbar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + s_mbar_w1 = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + # [b][warp] (value bits, index): the CTA reduction of the rows, then of the polled entries. + s_row_best = cutlass.Array( + cutlass.Int32, batch * WARPS * 2, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_poll_best = cutlass.Array( + cutlass.Int32, batch * WARPS * 2, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_flags = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_valid = cutlass.Array(cutlass.Int32, MAX_BATCH, space=cutlass.AddressSpace.smem, alignment=16) + + if warp == 0: + if prims.elect_sync(): + prims.mbarrier_init(s_mbar, 1) + prims.mbarrier_init(s_mbar_w1, batch) + prims.fence_mbarrier_init() + prims.barrier_cta_sync(0) + + # The CTA's markov_w2 rows: a weight, loaded before the grid dependency. + if warp == 0: + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(s_mbar, rows * ROW_BYTES) + for c in cutlass.range_constexpr((rows + COPY_ROWS - 1) // COPY_ROWS): + n = min(COPY_ROWS, rows - c * COPY_ROWS) + prims.cp_async_bulk_shared_cluster_global( + s_w2.subview(c * COPY_ROWS * ROW_WORDS), + w2.subview( + cta * cutlass.Int32(rows * ROW_WORDS) + + cutlass.Int32(c * COPY_ROWS * ROW_WORDS) + ), + s_mbar, + n * ROW_BYTES, + ) + # Dependents wait for this whole grid before reading its outputs. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + + prims.griddepcontrol(prims.GridDepAction.WAIT) + if tx == 0: + s_flags.store(flags.load(idx=0, is_volatile=True), idx=0) + s_flags.store(flags.load(idx=2, is_volatile=True), idx=2) + if tx < cutlass.Int32(batch): + _fetch_w1_row( + s_w1, s_valid, w1, s_mbar_w1, tx, cutlass.Int32(first_prev.load(idx=tx)), vocab + ) + prims.barrier_cta_sync(0) + cur = s_flags.load(idx=0) + cur_base = cur * cutlass.Int32(buf_words) + + # Re-arm the previous call's buffer (this CTA's 16-byte chunks of the words it used). + dirty_base = ((cur + cutlass.Int32(2)) % cutlass.Int32(BUFFERS)) * cutlass.Int32(buf_words) + dirty_words = s_flags.load(idx=2) + empty = cutlass.Int32(EMPTY_WORD) + w_clear = (cta * cutlass.Int32(THREADS) + tx) * cutlass.Int32(4) + while w_clear < dirty_words: + buf_uc.store((empty, empty, empty, empty), idx=dirty_base + w_clear, alignment=16) + w_clear += cutlass.Int32(grid * THREADS * 4) + + # next_new_tokens[b] = [accepted[row_b, num_accepted[b] - 1], tokens[b,:]] (-1 wraps, as torch indexing). + if cta == 0: + if tx < cutlass.Int32(batch): + col = num_accepted.load(idx=tx) - cutlass.Int32(1) + col = cutlass.Int32(cutlass.select_(col < cutlass.Int32(0), col + accepted_stride, col)) + row = accepted_rows.load(idx=tx) + next_new.store( + accepted.load(idx=row * accepted_stride + col), idx=tx * cutlass.Int32(block + 1) + ) + # The KV-length rewind after the draft forward (whose kernels, the readers of kv_lens, all finished before + # this grid's wait): kv_lens[first + b] = max(kv_lens[first + b] - rewind[b], 0). + if tx < rewind_count: + left = kv_lens.load(idx=rewind_first + tx) - rewind.load(idx=tx) + kv_lens.store(cutlass.Int32(cutlass.select_(left < cutlass.Int32(0), cutlass.Int32(0), left)), + idx=rewind_first + tx) # fmt: skip + + # The weights have landed (issued before the wait). + while not cute.arch.mbarrier_test_wait(s_mbar.data_ptr(), 0): + pass + + row0 = warp * cutlass.Int32(rows_per_warp) # first row of the warp within the CTA + shard_row0 = cta * cutlass.Int32(rows) + row0 # ... within the shard + has_row = (lane & cutlass.Int32(3)) == cutlass.Int32(0) + lane_row = shard_row0 + (lane >> cutlass.Int32(2)) + for k in cutlass.range(block, unroll=1): + # S4B8: every request's base logits are issued together before the position's W1 wait (one exposed + # global-load latency per position instead of one per request). Values, FMA order and argmax order unchanged. + base_all = [] + for b in cutlass.range_constexpr(batch): + base_row = (cutlass.Int32(b * block) + k) * cutlass.Int32(shard) + # Lane 4j finishes row j of each 8-row block (the reduce-scatter leaves row j's sum on lanes 4j .. 4j+3). + base_vals = [] + for g in cutlass.range_constexpr(blocks_per_warp): + if cutlass.const_expr(base_bf16): + base_vals.append( + cutlass.Int32( + cutlass.Int32( + base.load(idx=base_row + lane_row + cutlass.Int32(g * BLOCK_ROWS)) + ) + << cutlass.Int32(16) + ).bitcast(cutlass.Float32) + ) + else: + base_vals.append( + base.load(idx=base_row + lane_row + cutlass.Int32(g * BLOCK_ROWS)) + ) + base_all.append(base_vals) + # markov_w1[prev] landed for every request (the position's phase of the W1 barrier). + while not cute.arch.mbarrier_test_wait(s_mbar_w1.data_ptr(), k & cutlass.Int32(1)): + pass + for b in cutlass.range_constexpr(batch): + base_vals = base_all[b] + valid = s_valid.load(idx=b) != cutlass.Int32(0) + # markov_w1[prev], columns 8 lane .. 8 lane + 7, as fp32 (zeros for an invalid prev). + w1v = s_w1.load( + idx=cutlass.Int32(b * ROW_WORDS) + lane * 4, vector_size=4, alignment=16 + ) + w1f = [] + for q in cutlass.range_constexpr(4): + word = cutlass.Int32( + cutlass.select_(valid, cutlass.Int32(w1v[q]), cutlass.Int32(0)) + ) + w1f.append(_lo_f32(word)) + w1f.append(_hi_f32(word)) + best_value = cutlass.Int32(NEG_INF_BITS).bitcast(cutlass.Float32) + best_index = cutlass.Int32(INT_MAX) + for g in cutlass.range_constexpr(blocks_per_warp): + dot = _block_dot(s_w2, row0 + cutlass.Int32(g * BLOCK_ROWS), lane, w1f) + bias = _lo_f32(_pack_bf16x2(dot, dot)) + # fp32_sum (default): the in-tree chain's arithmetic (fp32 logits from LogitsProcessor's .float() + # plus the widened bf16 bias): compare the unrounded fp32 sum. Without fp32_sum: a bf16 reference + # rounds that sum once. + if cutlass.const_expr(fp32_sum): + value = _add_rn(base_vals[g], bias) + else: + value = _lo_f32( + _pack_bf16x2(_add_rn(base_vals[g], bias), _add_rn(base_vals[g], bias)) + ) + if has_row: + my_row = lane_row + cutlass.Int32(g * BLOCK_ROWS) + best_value, best_index = _keep_better( + best_value, best_index, value, vocab_offset + my_row + ) + best_value, best_index = _warp_best(best_value, best_index) + if lane == 0: + s_row_best.store((best_value.bitcast(cutlass.Int32), best_index), + idx=cutlass.Int32(b * WARPS * 2) + warp * cutlass.Int32(2), alignment=8) # fmt: skip + prims.barrier_cta_sync(0) + # Thread b pushes request b's CTA maximum into every rank's buffer. + if tx < cutlass.Int32(batch): + cta_value, cta_index = _reduce_slots(s_row_best, tx * cutlass.Int32(WARPS * 2)) + bits = cta_value.bitcast(cutlass.Int32) + bits = cutlass.Int32( + cutlass.select_(bits == cutlass.Int32(EMPTY_WORD), cutlass.Int32(0), bits) + ) + for c in cutlass.range_constexpr(push_copies): + s_idx = rank * cutlass.Int32(push_copies) + cutlass.Int32(c) + word = cur_base + (((k * cutlass.Int32(batch) + tx) * cutlass.Int32(slots) + s_idx) + * cutlass.Int32(grid) + cta) * cutlass.Int32(2) # fmt: skip + buf_mc.store((bits, cta_index), idx=word, alignment=8) + # Drain the posted multicast store now (the polls below issue no release that would). + prims.fence_acq_rel(prims.MemScope.CLUSTER) + # S4B8: one first pass loads every request's [k][b] entries together (B load waves in flight at once + # instead of one after another); a request still incomplete after it falls back to the v1 poll loop. The + # argmax order is total, so the result is the same whichever pass completes a request. + firsts = [] + for b in cutlass.range_constexpr(batch): + region = cur_base + (k * cutlass.Int32(batch) + cutlass.Int32(b)) * cutlass.Int32( + entries * 2 + ) + gots = [] + for i in cutlass.range_constexpr(vec_per_thread): + v_idx = tx + cutlass.Int32(i * THREADS) + live = v_idx < cutlass.Int32(vectors) + gots.append((live, buf_uc.load( + idx=region + cutlass.Int32(cutlass.select_(live, v_idx, cutlass.Int32(0))) * cutlass.Int32(4), + vector_size=4, alignment=16, is_volatile=True))) # fmt: skip + firsts.append((region, gots)) + for b in cutlass.range_constexpr(batch): + region, gots = firsts[b] + best_value = cutlass.Int32(NEG_INF_BITS).bitcast(cutlass.Float32) + best_index = cutlass.Int32(INT_MAX) + pending = cutlass.Boolean(False) + for live, got in gots: + for e in cutlass.range_constexpr(2): + vw = cutlass.Int32(got[2 * e]) + iw = cutlass.Int32(got[2 * e + 1]) + pending = pending | (live & ((vw == empty) | (iw == empty))) + take = live & _better(vw.bitcast(cutlass.Float32), iw, best_value, best_index) + best_value = cutlass.Float32( + cutlass.select_(take, vw.bitcast(cutlass.Float32), best_value) + ) + best_index = cutlass.Int32(cutlass.select_(take, iw, best_index)) + while pending: + pass_value = cutlass.Int32(NEG_INF_BITS).bitcast(cutlass.Float32) + pass_index = cutlass.Int32(INT_MAX) + pass_pending = cutlass.Boolean(False) + for i in cutlass.range_constexpr(vec_per_thread): + v_idx = tx + cutlass.Int32(i * THREADS) + live = v_idx < cutlass.Int32(vectors) + got = buf_uc.load( + idx=region + + cutlass.Int32(cutlass.select_(live, v_idx, cutlass.Int32(0))) + * cutlass.Int32(4), + vector_size=4, + alignment=16, + is_volatile=True, + ) + for e in cutlass.range_constexpr(2): + vw = cutlass.Int32(got[2 * e]) + iw = cutlass.Int32(got[2 * e + 1]) + pass_pending = pass_pending | (live & ((vw == empty) | (iw == empty))) + take = live & _better( + vw.bitcast(cutlass.Float32), iw, pass_value, pass_index + ) + pass_value = cutlass.Float32( + cutlass.select_(take, vw.bitcast(cutlass.Float32), pass_value) + ) + pass_index = cutlass.Int32(cutlass.select_(take, iw, pass_index)) + best_value = pass_value + best_index = pass_index + pending = pass_pending + best_value, best_index = _warp_best(best_value, best_index) + if lane == 0: + s_poll_best.store((best_value.bitcast(cutlass.Int32), best_index), + idx=cutlass.Int32(b * WARPS * 2) + warp * cutlass.Int32(2), alignment=8) # fmt: skip + prims.barrier_cta_sync(0) + if tx < cutlass.Int32(batch): + g_value, g_index = _reduce_slots(s_poll_best, tx * cutlass.Int32(WARPS * 2)) + # The next position's markov_w1 row, requested as soon as prev is known: every reader of s_w1 and + # s_valid for this position passed the barrier before the push. No barrier follows: the other threads + # go on to the next position's base logits and wait for the row's phase. s_poll_best is next written + # after the next push barrier, which this thread reaches only after these reads. + if k + cutlass.Int32(1) < cutlass.Int32(block): + _fetch_w1_row(s_w1, s_valid, w1, s_mbar_w1, tx, g_index, vocab) + if cta == 0: + tokens.store(g_index, idx=tx * cutlass.Int32(block) + k) + next_new.store(g_index, idx=tx * cutlass.Int32(block + 1) + k + cutlass.Int32(1)) + + # Every CTA read the flags before arriving; the last one advances them for the next call. + if tx == 0: + arrived = _atomic_add_acq_rel(flags.data_ptr(1).toint(), cutlass.Int32(1)) + if arrived == cutlass.Int32(grid - 1): + flags.store(cutlass.Int32(0), idx=1) + flags.store(cutlass.Int32(used_words), idx=2) + flags.store((cur + cutlass.Int32(1)) % cutlass.Int32(BUFFERS), idx=0) + + +@cute.jit +def k3_markov( + base: cute.Tensor, + first_prev: cute.Tensor, + w1: cute.Tensor, + w2: cute.Tensor, + corrected: cute.Tensor, + tokens: cute.Tensor, + next_new: cute.Tensor, + accepted: cute.Tensor, + num_accepted: cute.Tensor, + accepted_rows: cute.Tensor, + kv_lens: cute.Tensor, + rewind: cute.Tensor, + buf_uc: cute.Tensor, + buf_mc: cute.Tensor, + flags: cute.Tensor, + vocab: cutlass.Int32, + vocab_offset: cutlass.Int32, + rank: cutlass.Int32, + accepted_stride: cutlass.Int32, + rewind_first: cutlass.Int32, + rewind_count: cutlass.Int32, + grid: cutlass.Constexpr[int], + rows: cutlass.Constexpr[int], + shard: cutlass.Constexpr[int], + block: cutlass.Constexpr[int], + batch: cutlass.Constexpr[int], + slots: cutlass.Constexpr[int], + push_copies: cutlass.Constexpr[int], + buf_words: cutlass.Constexpr[int], + base_bf16: cutlass.Constexpr[bool], + fp32_sum: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + k3_markov_kernel( + base, first_prev, w1, w2, corrected, tokens, next_new, accepted, num_accepted, accepted_rows, kv_lens, + rewind, buf_uc, buf_mc, flags, vocab, vocab_offset, rank, accepted_stride, rewind_first, rewind_count, + grid, rows, shard, block, batch, slots, push_copies, buf_words, base_bf16, fp32_sum, + ).launch( + grid=[grid, 1, 1], + block=[THREADS, 1, 1], + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_markov_chain_kernel_b1_ll.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_markov_chain_kernel_b1_ll.py new file mode 100644 index 000000000000..af191f8cbc4b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_markov_chain_kernel_b1_ll.py @@ -0,0 +1,624 @@ +# 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. +# ============================================================================= +# DSpark vanilla-Markov draft chain for a vocab-sharded draft head -- CTM (prims/cute) kernel, BATCH-1 VARIANT WITH +# THE LL (FLAG-IN-PAYLOAD) EXCHANGE (MARKOV_LL) +# A copy of dspark_markov_chain_kernel_b1.py (S4B8) for batch 1 only, with the per-position exchange replaced: +# - post: one 16-byte multicast store per CTA and position, {value bits, tag, index, tag}; no fence and no separate +# flag word (the tag is in the payload); +# - poll + combine: warp 0 alone polls every entry of the position (16-byte loads) until both tag words of every +# entry hold this launch's tag, and reduces them with warp shuffles; no CTA barrier after the poll (the other warps +# go on to the next position's base loads and wait for the markov_w1 row, which only warp 0 can request); +# - the CTA's row reduction before the post: warp 0 waits on a named barrier the other warps only arrive at. +# Everything else is b1's: smem-resident markov_w2, the GEMV FMA order, the bf16 roundings, the fp32-sum option, the +# PDL wait placement, the argmax order (larger value, then lower vocabulary index; NaN never wins). Bitwise the same +# tokens as the b1 kernel: every CTA entry is computed as b1 computes it, and the global reduction folds the same set +# of entries under the same total order (any fold order gives the same winner). +# ============================================================================= +# +# For request 0 (batch 1) and block position k = 0 .. K-1 (greedy, temperature 0): +# bias[v] = bf16(markov_w2[v] . markov_w1[prev]) v in this rank's vocab shard +# value[v] = bf16(base[0, k, v] + bias[v]) (the reference's bf16 add: fp32 sum, one rounding) +# token[0, k] = the first maximum of value over the whole vocabulary (all ranks) +# prev = token[0, k] (prev for k = 0 is the anchor, first_prev[0]) +# +# Numerics: dspark_markov_chain_kernel_b1.py's, unchanged (lane l sums the 8 columns 8l .. 8l+7 of a row by FMAs from +# 0, xor-tree lane sums, cvt.rn to bf16, fp32 add, cvt.rn; with fp32_sum the unrounded fp32 sum is compared). +# +# Geometry: b1's (G CTAs of 256 threads, RC rows per CTA, warp w owns RC / 8 rows in blocks of 8, markov_w2 rows in +# shared memory before the grid-dependency wait, markov_w1[prev] one bulk copy per position on an mbarrier phase). +# +# LL exchange. The workspace (dspark_markov_chain.workspace_args) holds the 3 Lamport buffers of the other variants +# (untouched here) followed by LL_PARITIES = 2 LL regions of buf_words words each (zero-filled at allocation): +# region p = buf_uc[(BUFFERS + p) * buf_words ...], entry e = (k * slots + rank slot) * G + CTA, 4 words at 4 e: +# {value bits, tag, index, tag} +# Each 8-byte half carries the tag, so a reader that sees the tag in both halves has seen both halves of this write +# (8-byte single-copy granularity of a 16-byte store, the assumption of NCCL's LL protocol). +# Tag: seq, the launch sequence number of this workspace's LL calls (flags[LL_SEQ] holds the last one; this call +# uses seq = that + 1, skipping 0 on wrap; every rank runs the same calls in the same order, so every rank computes +# the same seq). Launch seq uses region seq & 1. The position is in the address, not the tag. +# Why a stale entry is never accepted: an entry is accepted only if both tag words equal this call's seq. Every LL +# call writes every entry it reads (k < K, every rank slot, every CTA) with its own seq, so an entry left by another +# call carries a different seq unless 2^32 LL calls passed since it was written (> 40 h of back-to-back calls of +# only a smaller shape), and the zero fill (tag 0) is never a seq. +# Why no entry is overwritten while a reader still needs it: call n + 1 on rank q uses the other region; call n + 2 +# (this region again) on rank q starts after q's call n + 1 completed, which waited for every rank's call n + 1 +# entries, which a rank posts only after its call n completed (griddepcontrol.wait before the first post). So no +# rank still reads call n's region when it is next written. (One region would do for K >= 2; two make K = 1 safe and +# keep calls of different shapes apart: the region stride is fixed, the layout inside a region follows the shape.) +# Memory ordering: the post is multimem.st.relaxed.sys (no release, no fence: the payload is the data, nothing else +# is published); the poll is ld.volatile (relaxed, system scope) 16-byte loads of this rank's copy (no acquire: no +# other data is read through the tag). flags: read after griddepcontrol.wait, advanced by the last CTA to arrive +# (atom.acq_rel.gpu), as in b1. +# +# Roles per position: all 8 warps compute their rows; lane 0 of each warp writes the warp's best to shared memory; +# warps 1..7 arrive at named barrier LL_BAR and go on to the next position's base loads and the markov_w1 wait; warp +# 0 waits at LL_BAR, thread 0 reduces the 8 warp bests and posts the CTA entry, warp 0 polls the W x G entries +# (ceil(W G / 32) 16-byte loads per lane, all issued before any is used), folds them per lane and across the warp +# (xor tree), and thread 0 requests the next markov_w1 row (and, in CTA 0, stores the token). The next position's +# readers of s_row_best / s_w1 / s_valid write or read them only after the row's mbarrier phase, which thread 0 +# completes after its reads of s_row_best (mbarrier arrive: release; the phase wait: acquire). +# ============================================================================= +"""CTM DSpark Markov chain, batch-1 variant with the LL (flag-in-payload) exchange: the draft head's greedy +intra-block Markov bias, the per-position global argmax across the tensor-parallel vocab shards (the draft tokens), +next_new_tokens and the target's KV-length rewind after the draft forward, in one launch.""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +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 + +THREADS = 256 +WARPS = THREADS // 32 +MARKOV_RANK = 256 # 32 lanes x 8 columns +ROW_WORDS = MARKOV_RANK // 2 # int32 words (bf16 pairs) per markov row +ROW_BYTES = MARKOV_RANK * 2 +BLOCK_ROWS = 8 # rows per warp block: one per epilogue lane +MAX_BATCH = 16 +MAX_BLOCK = 16 +BUFFERS = 3 +FLAG_WORDS = 4 +EMPTY_WORD = -(2**31) # 0x80000000: fp32 -0.0, never a pushed word +INT_MAX = 2**31 - 1 +NEG_INF_BITS = -8388608 # 0xFF800000 +COPY_ROWS = 32 # rows per weight bulk copy (16 KB) +SMEM_ROW_BUDGET = 400 # 200 KB of markov_w2 rows per CTA +SMEM_BYTES = 227 * 1024 +LL_PARITIES = 2 # LL regions after the Lamport buffers, buf_words words each +LL_FLAG_WORDS = 8 # flags words this variant needs: [4] last LL seq, [5] LL CTAs arrived ([0..3]: the other variants) +LL_SEQ = 4 +LL_ARRIVED = 5 +LL_ENTRY_WORDS = 4 # {value bits, tag, index, tag} +LL_BAR = 1 # named barrier: warp 0 waits for the 8 warp bests (barrier 0 is __syncthreads) + + +def rows_per_cta(shard: int, grid: int) -> int: + return shard // grid + + +def supports(shard: int, grid: int, block: int, batch: int) -> bool: + """Shapes the kernel runs: whole 8-row blocks for every warp (RC a multiple of 64), the CTA's markov_w2 rows + within the shared-memory budget, block <= 16, batch <= 16, an even entry count per exchange.""" + if grid <= 0 or shard % grid != 0: + return False + rc = rows_per_cta(shard, grid) + return ( + rc % (WARPS * BLOCK_ROWS) == 0 + and rc <= SMEM_ROW_BUDGET + and smem_bytes(rc, batch) <= SMEM_BYTES + and grid % 2 == 0 + and 0 < block <= MAX_BLOCK + and 0 < batch <= MAX_BATCH + ) + + +def smem_bytes(rows: int, batch: int) -> int: + """Shared memory of a CTA: its markov_w2 rows, the markov_w1[prev] rows, the two [b][warp] reductions, the + barriers, flags and prev (with the allocator's alignment padding).""" + return rows * ROW_BYTES + batch * ROW_BYTES + 2 * batch * WARPS * 8 + 1024 + + +def buffer_words(block: int, batch: int, slots: int, grid: int) -> int: + """Int32 words of one Lamport buffer: one (value, index) entry per position, request, rank slot and CTA.""" + return block * batch * slots * grid * 2 + + +@dsl_user_op +def _atomic_add_acq_rel(addr_i64, val, *, loc=None, ip=None): + """atom.acq_rel.gpu.global.add.u32, returning the old value.""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [addr_i64.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + "atom.acq_rel.gpu.global.add.u32 $0, [$1], $2;", "=r,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 _fma_rn(a, b, c, *, loc=None, ip=None): + """fma.rn.f32 (fmaf).""" + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip), c.ir_value(loc=loc, ip=ip)], + "fma.rn.f32 $0, $1, $2, $3;", "=f,f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + """add.rn.f32: an fp32 add the compiler cannot contract with a neighbouring multiply.""" + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "add.rn.f32 $0, $1, $2;", "=f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _pack_bf16x2(hi, lo, *, loc=None, ip=None): + """(bf16(hi) << 16) | bf16(lo), round to nearest even (cvt.rn, as __float2bfloat16_rn).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [hi.ir_value(loc=loc, ip=ip), lo.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", "=r,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _lo_f32(word): + """The bf16 in the low half of a word, as fp32 (exact).""" + return cutlass.Int32(word << cutlass.Int32(16)).bitcast(cutlass.Float32) + + +def _hi_f32(word): + """The bf16 in the high half of a word, as fp32 (exact).""" + return cutlass.Int32(word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + + +def _better(value, index, best_value, best_index): + """dsparkMarkovChainKernel's keepBetter order: larger value, or an equal value at a lower index. NaN never + compares better, so it never wins.""" + return (value > best_value) | ((value == best_value) & (index < best_index)) + + +def _keep_better(best_value, best_index, value, index): + take = _better(value, index, best_value, best_index) + return cutlass.Float32(cutlass.select_(take, value, best_value)), cutlass.Int32( + cutlass.select_(take, index, best_index) + ) + + +def _warp_best(value, index): + """Every lane gets the warp's best (value, index) (xor tree: the order is total, so all lanes agree).""" + for offset in [16, 8, 4, 2, 1]: + other_value = cute.arch.shuffle_sync_bfly(value, offset=offset, mask=-1, mask_and_clamp=31) + other_index = cute.arch.shuffle_sync_bfly(index, offset=offset, mask=-1, mask_and_clamp=31) + value, index = _keep_better(value, index, other_value, other_index) + return value, index + + +def _block_dot(s_w2, first_row, lane, w1f): + """markov_w2[first_row + j] . markov_w1[prev] for the block's 8 rows, in the C++ kernel's order: lane l sums the + 8 columns 8l .. 8l+7 by an FMA chain from 0, then the 32 lane sums are added in pairs at lane distances 16, 8, 4, 2, + 1 (the xor tree). Done as a reduce-scatter: at distances 16, 8 and 4 a lane keeps half of its rows and adds its + partner's partial of each kept row (the pairs of the full tree, so the same sums: 9 shuffles instead of 40), then + the full tree at 2 and 1. Returns row (lane >> 2)'s sum.""" + parts = [] + for r in range(BLOCK_ROWS): + w2v = s_w2.load(idx=(first_row + cutlass.Int32(r)) * cutlass.Int32(ROW_WORDS) + lane * cutlass.Int32(4), + vector_size=4, alignment=16) # fmt: skip + acc = cutlass.Float32(0.0) + for q in range(4): + word = cutlass.Int32(w2v[q]) + acc = _fma_rn(w1f[2 * q], _lo_f32(word), acc) + acc = _fma_rn(w1f[2 * q + 1], _hi_f32(word), acc) + parts.append(acc) + for offset in [16, 8, 4]: + upper = (lane & cutlass.Int32(offset)) != cutlass.Int32(0) + half = len(parts) // 2 + kept = [] + for i in range(half): + send = cutlass.Float32(cutlass.select_(upper, parts[i], parts[half + i])) + keep = cutlass.Float32(cutlass.select_(upper, parts[half + i], parts[i])) + kept.append( + _add_rn( + keep, + cute.arch.shuffle_sync_bfly(send, offset=offset, mask=-1, mask_and_clamp=31), + ) + ) + parts = kept + dot = parts[0] + for offset in [2, 1]: + dot = _add_rn( + dot, cute.arch.shuffle_sync_bfly(dot, offset=offset, mask=-1, mask_and_clamp=31) + ) + return dot + + +def _fetch_w1_row(s_w1, s_valid, w1, mbar, b, prev, vocab): + """Request b's markov_w1[prev] row into shared memory: one 512-byte bulk copy on the position's barrier phase (a + prev outside the vocabulary copies row 0 and clears s_valid[b]: the readers zero its bias). s_valid[b] is stored + before the arrive, whose release the readers' phase wait acquires.""" + valid = (prev >= cutlass.Int32(0)) & (prev < vocab) + row = cutlass.Int32(cutlass.select_(valid, prev, cutlass.Int32(0))) + s_valid.store(cutlass.Int32(cutlass.select_(valid, cutlass.Int32(1), cutlass.Int32(0))), idx=b) + prims.mbarrier_arrive_expect_tx(mbar, ROW_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_w1.subview(b * cutlass.Int32(ROW_WORDS)), + w1.subview(row * cutlass.Int32(ROW_WORDS)), + mbar, + ROW_BYTES, + ) + + +def _reduce_slots(slots_arr, first_word): + """The best of 8 (value bits, index) pairs at first_word .. in warp order.""" + value = cutlass.Int32(NEG_INF_BITS).bitcast(cutlass.Float32) + index = cutlass.Int32(INT_MAX) + for w in range(WARPS): + pair = slots_arr.load(idx=first_word + cutlass.Int32(w * 2), vector_size=2, alignment=8) + value, index = _keep_better( + value, index, cutlass.Int32(pair[0]).bitcast(cutlass.Float32), cutlass.Int32(pair[1]) + ) + return value, index + + +def ll_region_words(block: int, slots: int, grid: int) -> int: + """Int32 words of one LL region a call uses: one 4-word entry per position, rank slot and CTA (batch 1).""" + return block * slots * grid * LL_ENTRY_WORDS + + +def ll_workspace_words(buf_words: int) -> int: + """Int32 words of the whole workspace with the LL regions: the Lamport buffers, then LL_PARITIES regions.""" + return (BUFFERS + LL_PARITIES) * buf_words + + +@dsl_user_op +def _post_ll(addr_i64, a, b, c, d, *, loc=None, ip=None): + """multimem.st.relaxed.sys.global.v4.f32: one 16-byte store of {a, b, c, d} into every rank's copy (the bits, as + stored); relaxed at system scope so the peers' volatile polls observe it, with no release (nothing else is + published through it).""" + _llvm.inline_asm( + None, + [addr_i64.ir_value(loc=loc, ip=ip), a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip), + c.ir_value(loc=loc, ip=ip), d.ir_value(loc=loc, ip=ip)], + "multimem.st.relaxed.sys.global.v4.f32 [$0], {$1, $2, $3, $4};", "l,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 _bar_sync(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.sync on a named barrier: wait until `threads` threads arrived (or synced); memory ordering as bar.sync.""" + _llvm.inline_asm( + None, [], f"bar.sync {bar_id}, {threads};", "", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive on a named barrier: count this warp's arrival without waiting (the barrier's completion orders the + arriving threads' prior memory accesses before the waiting threads' later ones, as bar.sync).""" + _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, + ) # fmt: skip + + +@cute.kernel +def k3_markov_kernel( + base: cutlass.Array, # fp32 [K * S] (or its int16 bf16 view with base_bf16): the draft head's shard logits + first_prev: cutlass.Array, # int64 [1]: the anchor token + w1: cutlass.Array, # int32 words of bf16 markov_w1 [V, 256] + w2: cutlass.Array, # int32 words of bf16 markov_w2[shard rows] [S, 256] + corrected: cutlass.Array, # unused (kept for the launch signature) + tokens: cutlass.Array, # int32 [K], out + next_new: cutlass.Array, # int32 [K + 1], out + accepted: cutlass.Array, # int32 [*, accepted_stride]: accepted tokens + num_accepted: cutlass.Array, # int32 [1] + accepted_rows: cutlass.Array, # int32 [1]: the row of accepted of the request + kv_lens: cutlass.Array, # int32: the target's KV lengths, rewound here after the draft forward + rewind: cutlass.Array, # int32 [rewind_count]: subtracted from kv_lens[rewind_first ...], clamped at 0 + buf_uc: cutlass.Array, # int32 words: this rank's workspace (Lamport buffers, then the LL regions) + buf_mc: cutlass.Array, # int32 words: its multicast mapping + flags: cutlass.Array, # int32 [LL_FLAG_WORDS] + vocab: cutlass.Int32, # rows of markov_w1: a prev outside [0, vocab) gets a zero bias + vocab_offset: cutlass.Int32, # first vocabulary index of this rank's shard + rank: cutlass.Int32, + accepted_stride: cutlass.Int32, + rewind_first: cutlass.Int32, + rewind_count: cutlass.Int32, # 0: no rewind + grid: cutlass.Constexpr[int], + rows: cutlass.Constexpr[int], # RC + shard: cutlass.Constexpr[int], # S + block: cutlass.Constexpr[int], # K + batch: cutlass.Constexpr[int], # B: 1 (the op takes this variant only for batch 1) + slots: cutlass.Constexpr[int], # rank slots per k: world x push_copies + push_copies: cutlass.Constexpr[ + int + ], # test only: each rank's entry goes into push_copies consecutive slots + buf_words: cutlass.Constexpr[int], + base_bf16: cutlass.Constexpr[bool], + fp32_sum: cutlass.Constexpr[bool], +): + rows_per_warp = rows // WARPS + blocks_per_warp = rows_per_warp // BLOCK_ROWS + entries = slots * grid # LL entries per position + vec_per_lane = (entries + 31) // 32 # 16-byte entry loads per lane of warp 0 + ll_base = BUFFERS * buf_words # the first LL region + + tx, _, _ = cute.arch.thread_idx() + cta, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tx % cutlass.Int32(32) + + s_w2 = cutlass.Array( + cutlass.Int32, rows * ROW_WORDS, space=cutlass.AddressSpace.smem, alignment=128 + ) + # markov_w1[prev] of the position (one bulk copy, issued by thread 0 once it learns prev). + s_w1 = cutlass.Array(cutlass.Int32, ROW_WORDS, space=cutlass.AddressSpace.smem, alignment=128) + s_mbar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + s_mbar_w1 = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + # [warp] (value bits, index): the CTA reduction of the rows. + s_row_best = cutlass.Array( + cutlass.Int32, WARPS * 2, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_flags = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_valid = cutlass.Array(cutlass.Int32, MAX_BATCH, space=cutlass.AddressSpace.smem, alignment=16) + + if warp == 0: + if prims.elect_sync(): + prims.mbarrier_init(s_mbar, 1) + prims.mbarrier_init(s_mbar_w1, 1) + prims.fence_mbarrier_init() + prims.barrier_cta_sync(0) + + # The CTA's markov_w2 rows: a weight, loaded before the grid dependency. + if warp == 0: + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(s_mbar, rows * ROW_BYTES) + for c in cutlass.range_constexpr((rows + COPY_ROWS - 1) // COPY_ROWS): + n = min(COPY_ROWS, rows - c * COPY_ROWS) + prims.cp_async_bulk_shared_cluster_global( + s_w2.subview(c * COPY_ROWS * ROW_WORDS), + w2.subview( + cta * cutlass.Int32(rows * ROW_WORDS) + + cutlass.Int32(c * COPY_ROWS * ROW_WORDS) + ), + s_mbar, + n * ROW_BYTES, + ) + # Dependents wait for this whole grid before reading its outputs. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + + prims.griddepcontrol(prims.GridDepAction.WAIT) + if tx == 0: + s_flags.store(flags.load(idx=LL_SEQ, is_volatile=True), idx=0) + _fetch_w1_row( + s_w1, s_valid, w1, s_mbar_w1, tx, cutlass.Int32(first_prev.load(idx=0)), vocab + ) + prims.barrier_cta_sync(0) + # This call's tag: the last LL call's seq + 1, skipping 0 (the zero fill) on wrap; its region: seq & 1. + seq = s_flags.load(idx=0) + cutlass.Int32(1) + seq = cutlass.Int32(cutlass.select_(seq == cutlass.Int32(0), cutlass.Int32(2), seq)) + region0 = cutlass.Int32(ll_base) + (seq & cutlass.Int32(1)) * cutlass.Int32(buf_words) + + # next_new_tokens[0] = [accepted[row_0, num_accepted[0] - 1], tokens[0, :]] (-1 wraps, as torch indexing). + if cta == 0: + if tx == 0: + col = num_accepted.load(idx=0) - cutlass.Int32(1) + col = cutlass.Int32(cutlass.select_(col < cutlass.Int32(0), col + accepted_stride, col)) + row = accepted_rows.load(idx=0) + next_new.store(accepted.load(idx=row * accepted_stride + col), idx=0) + # The KV-length rewind after the draft forward (whose kernels, the readers of kv_lens, all finished before + # this grid's wait): kv_lens[first + b] = max(kv_lens[first + b] - rewind[b], 0). + if tx < rewind_count: + left = kv_lens.load(idx=rewind_first + tx) - rewind.load(idx=tx) + kv_lens.store(cutlass.Int32(cutlass.select_(left < cutlass.Int32(0), cutlass.Int32(0), left)), + idx=rewind_first + tx) # fmt: skip + + # The weights have landed (issued before the wait). + while not cute.arch.mbarrier_test_wait(s_mbar.data_ptr(), 0): + pass + + row0 = warp * cutlass.Int32(rows_per_warp) # first row of the warp within the CTA + shard_row0 = cta * cutlass.Int32(rows) + row0 # ... within the shard + has_row = (lane & cutlass.Int32(3)) == cutlass.Int32(0) + lane_row = shard_row0 + (lane >> cutlass.Int32(2)) + empty = cutlass.Int32(EMPTY_WORD) + for k in cutlass.range(block, unroll=1): + # The position's base logits, issued before the W1 wait (b1's order). + base_row = k * cutlass.Int32(shard) + # Lane 4j finishes row j of each 8-row block (the reduce-scatter leaves row j's sum on lanes 4j .. 4j+3). + base_vals = [] + for g in cutlass.range_constexpr(blocks_per_warp): + if cutlass.const_expr(base_bf16): + base_vals.append( + cutlass.Int32( + cutlass.Int32( + base.load(idx=base_row + lane_row + cutlass.Int32(g * BLOCK_ROWS)) + ) + << cutlass.Int32(16) + ).bitcast(cutlass.Float32) + ) + else: + base_vals.append(base.load(idx=base_row + lane_row + cutlass.Int32(g * BLOCK_ROWS))) + # markov_w1[prev] landed (the position's phase of the W1 barrier). + while not cute.arch.mbarrier_test_wait(s_mbar_w1.data_ptr(), k & cutlass.Int32(1)): + pass + valid = s_valid.load(idx=0) != cutlass.Int32(0) + # markov_w1[prev], columns 8 lane .. 8 lane + 7, as fp32 (zeros for an invalid prev). + w1v = s_w1.load(idx=lane * 4, vector_size=4, alignment=16) + w1f = [] + for q in cutlass.range_constexpr(4): + word = cutlass.Int32(cutlass.select_(valid, cutlass.Int32(w1v[q]), cutlass.Int32(0))) + w1f.append(_lo_f32(word)) + w1f.append(_hi_f32(word)) + best_value = cutlass.Int32(NEG_INF_BITS).bitcast(cutlass.Float32) + best_index = cutlass.Int32(INT_MAX) + for g in cutlass.range_constexpr(blocks_per_warp): + dot = _block_dot(s_w2, row0 + cutlass.Int32(g * BLOCK_ROWS), lane, w1f) + bias = _lo_f32(_pack_bf16x2(dot, dot)) + # fp32_sum (default): the in-tree chain's arithmetic (fp32 logits from LogitsProcessor's .float() plus + # the widened bf16 bias): compare the unrounded fp32 sum. Without fp32_sum: a bf16 reference rounds that + # sum once. + if cutlass.const_expr(fp32_sum): + value = _add_rn(base_vals[g], bias) + else: + value = _lo_f32( + _pack_bf16x2(_add_rn(base_vals[g], bias), _add_rn(base_vals[g], bias)) + ) + if has_row: + my_row = lane_row + cutlass.Int32(g * BLOCK_ROWS) + best_value, best_index = _keep_better( + best_value, best_index, value, vocab_offset + my_row + ) + best_value, best_index = _warp_best(best_value, best_index) + if lane == 0: + s_row_best.store((best_value.bitcast(cutlass.Int32), best_index), idx=warp * cutlass.Int32(2), + alignment=8) # fmt: skip + if warp == 0: + # Every warp's best of this position is in shared memory. + _bar_sync(LL_BAR, THREADS) + region = region0 + k * cutlass.Int32(entries * LL_ENTRY_WORDS) + # Thread 0 posts the CTA entry {value bits, tag, index, tag} into every rank's region: one 16-byte + # multicast store, no fence. + if tx == 0: + cta_value, cta_index = _reduce_slots(s_row_best, cutlass.Int32(0)) + bits = cta_value.bitcast(cutlass.Int32) + # -0.0 is posted as +0.0, as b1 posts it (equal under the order either way). + bits = cutlass.Int32(cutlass.select_(bits == empty, cutlass.Int32(0), bits)) + for c in cutlass.range_constexpr(push_copies): + s_idx = rank * cutlass.Int32(push_copies) + cutlass.Int32(c) + word = region + (s_idx * cutlass.Int32(grid) + cta) * cutlass.Int32( + LL_ENTRY_WORDS + ) + _post_ll(buf_mc.data_ptr(word).toint(), bits, seq, cta_index, seq) + # Warp 0 polls all W x G entries of the position in this rank's copy: lane l reads entries l, l + 32, .. + # (every load of a pass issued before any is used); a pass in which any entry of the lane does not + # carry this call's tag in both halves is repeated in full. The order is total, so the lane best and + # the xor tree give b1's winner. + pending = cutlass.Boolean(True) + g_value = cutlass.Int32(NEG_INF_BITS).bitcast(cutlass.Float32) + g_index = cutlass.Int32(INT_MAX) + while pending: + gots = [] + for i in cutlass.range_constexpr(vec_per_lane): + e_idx = lane + cutlass.Int32(i * 32) + live = e_idx < cutlass.Int32(entries) + gots.append((live, buf_uc.load( + idx=region + cutlass.Int32(cutlass.select_(live, e_idx, cutlass.Int32(0))) + * cutlass.Int32(LL_ENTRY_WORDS), + vector_size=4, alignment=16, is_volatile=True))) # fmt: skip + pass_value = cutlass.Int32(NEG_INF_BITS).bitcast(cutlass.Float32) + pass_index = cutlass.Int32(INT_MAX) + pass_pending = cutlass.Boolean(False) + for live, got in gots: + vw = cutlass.Int32(got[0]) + iw = cutlass.Int32(got[2]) + tagged = (cutlass.Int32(got[1]) == seq) & (cutlass.Int32(got[3]) == seq) + stale = (cutlass.Int32(got[1]) != seq) | (cutlass.Int32(got[3]) != seq) + pass_pending = pass_pending | (live & stale) + take = ( + live + & tagged + & _better(vw.bitcast(cutlass.Float32), iw, pass_value, pass_index) + ) + pass_value = cutlass.Float32( + cutlass.select_(take, vw.bitcast(cutlass.Float32), pass_value) + ) + pass_index = cutlass.Int32(cutlass.select_(take, iw, pass_index)) + g_value = pass_value + g_index = pass_index + pending = pass_pending + g_value, g_index = _warp_best(g_value, g_index) + if tx == 0: + # The next position's markov_w1 row, requested as soon as prev is known: every reader of s_w1 and + # s_valid for this position arrived at LL_BAR after its reads. The other warps wait for the row's + # phase before writing s_row_best again, and this thread read s_row_best before the arrive. + if k + cutlass.Int32(1) < cutlass.Int32(block): + _fetch_w1_row(s_w1, s_valid, w1, s_mbar_w1, tx, g_index, vocab) + if cta == 0: + tokens.store(g_index, idx=k) + next_new.store(g_index, idx=k + cutlass.Int32(1)) + else: + _bar_arrive(LL_BAR, THREADS) + + # Every CTA read the flags before arriving; the last one advances the LL seq for the next call. + if tx == 0: + arrived = _atomic_add_acq_rel(flags.data_ptr(LL_ARRIVED).toint(), cutlass.Int32(1)) + if arrived == cutlass.Int32(grid - 1): + flags.store(cutlass.Int32(0), idx=LL_ARRIVED) + flags.store(seq, idx=LL_SEQ) + + +@cute.jit +def k3_markov( + base: cute.Tensor, + first_prev: cute.Tensor, + w1: cute.Tensor, + w2: cute.Tensor, + corrected: cute.Tensor, + tokens: cute.Tensor, + next_new: cute.Tensor, + accepted: cute.Tensor, + num_accepted: cute.Tensor, + accepted_rows: cute.Tensor, + kv_lens: cute.Tensor, + rewind: cute.Tensor, + buf_uc: cute.Tensor, + buf_mc: cute.Tensor, + flags: cute.Tensor, + vocab: cutlass.Int32, + vocab_offset: cutlass.Int32, + rank: cutlass.Int32, + accepted_stride: cutlass.Int32, + rewind_first: cutlass.Int32, + rewind_count: cutlass.Int32, + grid: cutlass.Constexpr[int], + rows: cutlass.Constexpr[int], + shard: cutlass.Constexpr[int], + block: cutlass.Constexpr[int], + batch: cutlass.Constexpr[int], + slots: cutlass.Constexpr[int], + push_copies: cutlass.Constexpr[int], + buf_words: cutlass.Constexpr[int], + base_bf16: cutlass.Constexpr[bool], + fp32_sum: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + k3_markov_kernel( + base, first_prev, w1, w2, corrected, tokens, next_new, accepted, num_accepted, accepted_rows, kv_lens, + rewind, buf_uc, buf_mc, flags, vocab, vocab_offset, rank, accepted_stride, rewind_first, rewind_count, + grid, rows, shard, block, batch, slots, push_copies, buf_words, base_bf16, fp32_sum, + ).launch( + grid=[grid, 1, 1], + block=[THREADS, 1, 1], + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_qk_norm_rope.py new file mode 100644 index 000000000000..08b6630400ab --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/dspark_qk_norm_rope.py @@ -0,0 +1,309 @@ +# 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. +# ============================================================================= +# DSpark / DFlash drafter block: q / k RMSNorm + partial NeoX RoPE in one launch +# ============================================================================= +# +# It replaces, bit for bit, the three launches of DFlashForCausalLM.dflash_forward's flashinfer-RoPE branch at head_dim +# 128: flashinfer_rmsnorm on the q heads, flashinfer_rmsnorm on the k heads (both flashinfer's CuTe RMSNormKernel), +# then flashinfer_apply_rope_with_cos_sin_cache_inplace (NeoX, the cache's leading ROT dims rotated, the rest kept). +# The arithmetic is flashinfer's, made explicit (every op inline PTX, so nothing is contracted or reordered): +# RMSNorm: 16 threads x 8 contiguous elements per row; each thread sums its x^2 sequentially from 0; the 16 sums are +# added by the butterfly at offsets 1, 2, 4, 8 (adjacent pairs, then pairs of pairs); mean = sum * (1 / 128); +# rstd = rsqrt.approx.ftz(mean + eps); y = bf16((x * rstd) * w). +# RoPE (on the bf16 norm output, cos / sin fp32 from the cache): out_i = fma(x_i, cos_i, -x_{i+R/2} * sin_i) and +# out_{i+R/2} = fma(x_{i+R/2}, cos_i, x_i * sin_i) for i < R / 2, rounded to bf16. +# Measured against the stock launches: 0 mismatching elements at 4 and 32 rows. +# +# DRAFT_APPEND_FUSE (``qk_norm_rope_append``): the same launch also writes each row's K (the bits it returns) and V +# (copied from qkv) into the layer's paged HND cache [pages, 2, NKV, PS, HD] where flashinfer.page.append_paged_kv_cache +# would: page kv_indices[kv_indptr[b] + pos // PS], slot pos % PS, with b = batch_indices[row] and pos = positions[row]. +# The caller then skips that append launch; the cache bits are the same. + +import torch +import triton +import triton.language as tl + +HEAD_DIM = 128 + + +@triton.jit +def _add(a, b): + return tl.inline_asm_elementwise( + "add.rn.f32 $0, $1, $2;", "=r,r,r", [a, b], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _mul(a, b): + return tl.inline_asm_elementwise( + "mul.rn.f32 $0, $1, $2;", "=r,r,r", [a, b], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _fma(a, b, c): + return tl.inline_asm_elementwise( + "fma.rn.f32 $0, $1, $2, $3;", "=r,r,r,r", [a, b, c], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _rsqrt(a): + return tl.inline_asm_elementwise( + "rsqrt.approx.ftz.f32 $0, $1;", "=r,r", [a], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _pick(v, idx, j): + """Column j of a [R, C] tile as [R] (exact: one nonzero term).""" + return tl.sum(tl.where(idx == j, v, 0.0), axis=1) + + +@triton.jit +def _pair_sum(v, n: tl.constexpr): + """[n] -> [n / 2]: element k = v[2k] + v[2k + 1].""" + v2 = tl.reshape(v, (n // 2, 2)) + idx = tl.arange(0, 2)[None, :] + tl.arange(0, n // 2)[:, None] * 0 + return _add(_pick(v2, idx, 0), _pick(v2, idx, 1)) + + +@triton.jit +def _qk_norm_rope_kernel( + qkv_ptr, + qkv_stride, + pos_ptr, + cs_ptr, + cs_stride, + wq_ptr, + wk_ptr, + q_out, + k_out, + eps_q, + eps_k, + cache_ptr, + cache_page_stride, + bidx_ptr, + apos_ptr, + kvi_ptr, + kvp_ptr, + v_col, + HQ: tl.constexpr, + HKV: tl.constexpr, + HD: tl.constexpr, + HALF: tl.constexpr, + PS: tl.constexpr, + STORE: tl.constexpr, +): + pid = tl.program_id(0) + t = pid // (HQ + HKV) + h = pid % (HQ + HKV) + is_q = h < HQ + base = qkv_ptr + t * qkv_stride + h * HD + r16 = tl.arange(0, 16) + c8 = tl.arange(0, 8) + x2 = tl.load(base + r16[:, None] * 8 + c8[None, :]).to(tl.float32) + sq = _mul(x2, x2) + col = c8[None, :] + r16[:, None] * 0 + s = tl.zeros([16], dtype=tl.float32) + for j in tl.static_range(8): + s = _add(s, _pick(sq, col, j)) + s8 = _pair_sum(s, 16) + s4 = _pair_sum(s8, 8) + s2 = _pair_sum(s4, 4) + two = tl.arange(0, 2) + total = _add(tl.sum(tl.where(two == 0, s2, 0.0)), tl.sum(tl.where(two == 1, s2, 0.0))) + eps = tl.where(is_q, eps_q, eps_k) + rstd = _rsqrt(_add(_mul(total, 1.0 / HD), eps)) + d = tl.arange(0, HD) + x = tl.load(base + d).to(tl.float32) + w = tl.where(is_q, tl.load(wq_ptr + d), tl.load(wk_ptr + d)).to(tl.float32) + y = _mul(_mul(x, rstd), w).to(tl.bfloat16) + pos = tl.load(pos_ptr + t) + i = tl.arange(0, HALF) + cos = tl.load(cs_ptr + pos * cs_stride + i) + sin = tl.load(cs_ptr + pos * cs_stride + HALF + i) + yf = y.to(tl.float32) + x1 = tl.sum(tl.where(d[None, :] == i[:, None], yf[None, :], 0.0), axis=1) + x2r = tl.sum(tl.where(d[None, :] == (i[:, None] + HALF), yf[None, :], 0.0), axis=1) + o1 = _fma(x1, cos, _mul(-x2r, sin)) + o2 = _fma(x2r, cos, _mul(x1, sin)) + obase = tl.where(is_q, q_out + t * (HQ * HD) + h * HD, k_out + t * (HKV * HD) + (h - HQ) * HD) + tl.store(obase + d, y, mask=d >= 2 * HALF) + tl.store(obase + i, o1.to(tl.bfloat16)) + tl.store(obase + HALF + i, o2.to(tl.bfloat16)) + if STORE: + # K / V heads only: K's bits as stored above, V copied from qkv, at the append's page and slot. + is_k = h >= HQ + hk = tl.where(is_k, h - HQ, 0) + b = tl.load(bidx_ptr + t).to(tl.int64) + p = tl.load(apos_ptr + t).to(tl.int64) + page = tl.load(kvi_ptr + tl.load(kvp_ptr + b).to(tl.int64) + p // PS).to(tl.int64) + slot = p % PS + kdst = cache_ptr + page * cache_page_stride + (hk * PS + slot) * HD + vdst = cache_ptr + page * cache_page_stride + ((HKV + hk) * PS + slot) * HD + tl.store(kdst + d, y, mask=(d >= 2 * HALF) & is_k) + tl.store(kdst + i, o1.to(tl.bfloat16), mask=is_k) + tl.store(kdst + HALF + i, o2.to(tl.bfloat16), mask=is_k) + v = tl.load(qkv_ptr + t * qkv_stride + v_col + hk * HD + d, mask=is_k & (d >= 0), other=0.0) + tl.store(vdst + d, v, mask=is_k & (d >= 0)) + + +def _plain_rmsnorm(norm) -> bool: + """The stock q_norm / k_norm takes flashinfer's plain (non-Gemma) RMSNorm on a bf16 input.""" + return ( + getattr(norm, "use_gemma", True) is False + and not getattr(norm, "use_cuda_tile", False) + and not getattr(norm, "is_nvfp4", False) + and not getattr(norm, "return_hp_output", False) + and getattr(norm, "weight", None) is not None + and norm.weight.dtype == torch.bfloat16 + and norm.weight.shape == (HEAD_DIM,) + ) + + +def qk_norm_rope(qkv, q_size, kv_size, head_dim, q_norm, k_norm, positions, cos_sin_cache, is_neox): + """q / k RMSNorm + RoPE of the drafter block in one launch: returns new contiguous (q [T, q_size], k [T, kv_size]) + bitwise equal to the stock flashinfer path, or None (nothing computed) for anything it does not reproduce.""" + if not qk_norm_rope_supported( + qkv, q_size, kv_size, head_dim, q_norm, k_norm, positions, cos_sin_cache, is_neox + ): + return None + return _launch(qkv, q_size, kv_size, q_norm, k_norm, positions, cos_sin_cache) + + +def qk_norm_rope_supported( + qkv, q_size, kv_size, head_dim, q_norm, k_norm, positions, cos_sin_cache, is_neox +) -> bool: + """Whether the one-launch q / k RMSNorm + RoPE reproduces the stock path for these operands.""" + from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE + + return bool( + IS_FLASHINFER_AVAILABLE + and is_neox + and head_dim == HEAD_DIM + and qkv.dtype == torch.bfloat16 + and qkv.dim() == 2 + and qkv.stride(1) == 1 + and qkv.shape[1] >= q_size + kv_size + and q_size % head_dim == 0 + and kv_size % head_dim == 0 + and positions.dtype == torch.int32 + and positions.dim() == 1 + and positions.shape[0] == qkv.shape[0] + and cos_sin_cache.dtype == torch.float32 + and cos_sin_cache.dim() == 2 + and cos_sin_cache.stride(1) == 1 + and cos_sin_cache.shape[1] in (16, 32, 64, 128) + and _plain_rmsnorm(q_norm) + and _plain_rmsnorm(k_norm) + ) + + +def _launch(qkv, q_size, kv_size, q_norm, k_norm, positions, cos_sin_cache, append=None): + T = qkv.shape[0] + hq, hkv = q_size // HEAD_DIM, kv_size // HEAD_DIM + q = torch.empty(T, q_size, dtype=torch.bfloat16, device=qkv.device) + k = torch.empty(T, kv_size, dtype=torch.bfloat16, device=qkv.device) + if T == 0: + return q, k + if append is None: + cache, bidx, apos, kvi, kvp, ps, store = ( + q, + positions, + positions, + positions, + positions, + 1, + False, + ) + else: + cache, bidx, apos, kvi, kvp = append + ps, store = cache.shape[-2], True + _qk_norm_rope_kernel[(T * (hq + hkv),)]( + qkv, + qkv.stride(0), + positions, + cos_sin_cache, + cos_sin_cache.stride(0), + q_norm.weight, + k_norm.weight, + q, + k, + float(q_norm.variance_epsilon), + float(k_norm.variance_epsilon), + cache, + cache.stride(0) if store else 0, + bidx, + apos, + kvi, + kvp, + q_size + kv_size, + HQ=hq, + HKV=hkv, + HD=HEAD_DIM, + HALF=cos_sin_cache.shape[1] // 2, + PS=ps, + STORE=store, + num_warps=1, + ) + return q, k + + +_INT = (torch.int32, torch.int64) + + +def append_supported( + qkv, kv_size, paged_kv_cache, batch_indices, positions, kv_indices, kv_indptr +) -> bool: + """Whether ``qk_norm_rope_append`` reproduces flashinfer.page.append_paged_kv_cache for these operands: a + contiguous bf16 HND layer cache [pages, 2, NKV, PS, HEAD_DIM] with NKV = kv_size / HEAD_DIM, one append row per qkv + row, integer index tensors on the same device.""" + c = paged_kv_cache + T = qkv.shape[0] + return ( + isinstance(c, torch.Tensor) + and c.dim() == 5 + and c.dtype == torch.bfloat16 + and c.is_contiguous() + and c.device == qkv.device + and c.shape[1] == 2 + and c.shape[2] == kv_size // HEAD_DIM + and c.shape[4] == HEAD_DIM + and c.shape[3] >= 1 + and all( + x.dtype in _INT and x.device == qkv.device and x.dim() == 1 and x.is_contiguous() + for x in (batch_indices, positions, kv_indices, kv_indptr) + ) + and batch_indices.shape[0] == T + and positions.shape[0] == T + ) + + +def qk_norm_rope_append(qkv, q_size, kv_size, head_dim, q_norm, k_norm, positions, cos_sin_cache, is_neox, + paged_kv_cache, append_batch_indices, append_positions, kv_indices, kv_indptr): # fmt: skip + """``qk_norm_rope`` that also appends the rows' K / V into ``paged_kv_cache`` exactly as + flashinfer.page.append_paged_kv_cache(k, v, append_batch_indices, append_positions, paged_kv_cache, kv_indices, + kv_indptr, ..., kv_layout="HND") would (same bits); None (nothing computed or written) when either part is not + reproduced. The caller skips its append on a non-None return.""" + if not qk_norm_rope_supported(qkv, q_size, kv_size, head_dim, q_norm, k_norm, positions, cos_sin_cache, + is_neox): # fmt: skip + return None + if not append_supported(qkv, kv_size, paged_kv_cache, append_batch_indices, append_positions, kv_indices, + kv_indptr): # fmt: skip + return None + return _launch(qkv, q_size, kv_size, q_norm, k_norm, positions, cos_sin_cache, + (paged_kv_cache, append_batch_indices, append_positions, kv_indices, kv_indptr)) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/strict_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/strict_accept.py new file mode 100644 index 000000000000..eda460592842 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/strict_accept.py @@ -0,0 +1,482 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The strict draft acceptance and the next step's input tokens of a decode-only speculative step, in one launch each. + +The in-tree strict path verifies ``D`` draft tokens per request against the target's sampled tokens with about ten +small torch launches (empty / ones / reshape / == / int / cumprod / sum / += / slice copy), and builds the next +step's input tokens with three more (index, unsqueeze, concat). Both are integer copies and compares, so these +kernels are bitwise to that path by construction; the forced-acceptance override stays the caller's (it advances +the in-tree RNG pool and counters). + +Compiled on the first call per (draft length, dtypes), which must run outside CUDA-graph capture (``precompile``). +""" + +from __future__ import annotations + +import importlib +import threading +from typing import Dict, Sequence + +import torch + +from .._op import Arch, Cell, OpWrapper + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} +_modules: Dict[str, object] = {} + + +def _kernel_module(): + """The CuTe DSL kernels beside this file, imported on first use (they pull in cutlass).""" + mod = _modules.get("kernel") + if mod is None: + mod = _modules["kernel"] = importlib.import_module(".strict_accept_kernel", __package__) + return mod + + +def _arg(t: torch.Tensor, align: int): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=0) + + +def _int_align(t: torch.Tensor) -> int: + return t.element_size() + + +def max_requests() -> int: + return _kernel_module().MAX_REQUESTS + + +def _compiled_fn(name: str, key: tuple, args: tuple, scalars: tuple, consts: tuple, stream): + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"trtllm::{name} must run once per shape outside CUDA-graph capture first" + ) + import cutlass.cute as cute + + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + getattr(_kernel_module(), name), *args, *scalars, *consts, stream + ) + return fn + + +@torch.library.custom_op("trtllm::spec_strict_accept", mutates_args=()) +def _spec_strict_accept(target: torch.Tensor, draft: torch.Tensor) -> list[torch.Tensor]: + import cuda.bindings.driver as cuda_driver + + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + n_req, draft_len = draft.shape + if target.numel() != n_req * (draft_len + 1): + raise ValueError( + f"spec_strict_accept: target {tuple(target.shape)} vs draft {tuple(draft.shape)}" + ) + if n_req > max_requests() or not target.is_contiguous() or not draft.is_contiguous(): + raise ValueError("spec_strict_accept: unsupported operands (see supports())") + device = target.device + accepted = torch.empty(n_req, draft_len + 1, dtype=torch.int32, device=device) + num_accepted = torch.empty(n_req, dtype=torch.int32, device=device) + args = ( + _arg(target.view(-1), _int_align(target)), + _arg(draft.view(-1), _int_align(draft)), + _arg(accepted.view(-1), 4), + _arg(num_accepted, 4), + ) + use_pdl = get_env_enable_pdl() + stream = cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) + key = ("strict_accept", target.dtype, draft.dtype, draft_len, use_pdl) + fn = _compiled_fn("strict_accept", key, args, (n_req,), (draft_len, use_pdl), stream) + fn(*args, n_req, stream) + return [accepted, num_accepted] + + +@_spec_strict_accept.register_fake +def _(target, draft): + n_req, draft_len = draft.shape + return [ + target.new_empty((n_req, draft_len + 1), dtype=torch.int32), + target.new_empty((n_req,), dtype=torch.int32), + ] + + +@torch.library.custom_op("trtllm::spec_next_new_tokens", mutates_args=()) +def _spec_next_new_tokens( + accepted: torch.Tensor, + num_accepted: torch.Tensor, + batch_idx: torch.Tensor, + next_draft: torch.Tensor, +) -> torch.Tensor: + import cuda.bindings.driver as cuda_driver + + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + n_req, draft_len = next_draft.shape + if n_req > max_requests() or accepted.stride(1) != 1 or next_draft.stride(1) != 1: + raise ValueError("spec_next_new_tokens: unsupported operands (see supports())") + device = accepted.device + out = torch.empty( + n_req, + draft_len + 1, + dtype=torch.promote_types(accepted.dtype, next_draft.dtype), + device=device, + ) + # Flat views over the storage each row index reads (strided rows: the row strides are passed). + acc_flat = accepted.as_strided( + ((accepted.shape[0] - 1) * accepted.stride(0) + accepted.shape[1],), (1,) + ) + nd_flat = next_draft.as_strided(((n_req - 1) * next_draft.stride(0) + draft_len,), (1,)) + args = ( + _arg(acc_flat, _int_align(accepted)), + _arg(num_accepted, _int_align(num_accepted)), + _arg(batch_idx, _int_align(batch_idx)), + _arg(nd_flat, _int_align(next_draft)), + _arg(out.view(-1), _int_align(out)), + ) + scalars = (int(accepted.stride(0)), int(next_draft.stride(0)), n_req) + use_pdl = get_env_enable_pdl() + stream = cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) + key = ( + "next_new_tokens", + accepted.dtype, + num_accepted.dtype, + batch_idx.dtype, + next_draft.dtype, + draft_len, + use_pdl, + ) + fn = _compiled_fn("next_new_tokens", key, args, scalars, (draft_len, use_pdl), stream) + fn(*args, *scalars, stream) + return out + + +@_spec_next_new_tokens.register_fake +def _(accepted, num_accepted, batch_idx, next_draft): + n_req, draft_len = next_draft.shape + return accepted.new_empty( + (n_req, draft_len + 1), dtype=torch.promote_types(accepted.dtype, next_draft.dtype) + ) + + +@torch.library.custom_op("trtllm::spec_kv_rewind", mutates_args=("kv_lens", "rewind")) +def _spec_kv_rewind( + num_accepted: torch.Tensor, kv_lens: torch.Tensor, rewind: torch.Tensor, nc: int, n: int +) -> None: + import cuda.bindings.driver as cuda_driver + + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + if not 0 < n <= max_requests() or rewind.numel() < n or rewind.dtype != num_accepted.dtype: + raise ValueError("spec_kv_rewind: unsupported operands (see supports_kv_rewind())") + device = num_accepted.device + args = ( + _arg(num_accepted, _int_align(num_accepted)), + _arg(kv_lens, _int_align(kv_lens)), + _arg(rewind, _int_align(rewind)), + ) + use_pdl = get_env_enable_pdl() + stream = cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) + key = ("kv_rewind", num_accepted.dtype, kv_lens.dtype, use_pdl) + fn = _compiled_fn("kv_rewind", key, args, (int(nc), int(n)), (use_pdl,), stream) + fn(*args, int(nc), int(n), stream) + + +@_spec_kv_rewind.register_fake +def _(num_accepted, kv_lens, rewind, nc, n): + return None + + +_INT_DTYPES = (torch.int32, torch.int64) + + +def supports_strict_accept(target: torch.Tensor, draft: torch.Tensor) -> bool: + """Whether the kernel takes these operands; a caller that gets False runs the in-tree path.""" + return ( + target.is_cuda + and draft.is_cuda + and target.dtype in _INT_DTYPES + and draft.dtype in _INT_DTYPES + and draft.dim() == 2 + and 0 < draft.shape[0] <= max_requests() + and draft.shape[1] > 0 + and target.numel() == draft.shape[0] * (draft.shape[1] + 1) + and target.is_contiguous() + and draft.is_contiguous() + ) + + +def supports_next_new_tokens( + accepted: torch.Tensor, + num_accepted: torch.Tensor, + batch_idx: torch.Tensor, + next_draft: torch.Tensor, +) -> bool: + """Whether the kernel takes these operands; a caller that gets False runs the in-tree path.""" + return ( + all( + t.is_cuda and t.dtype in _INT_DTYPES + for t in (accepted, num_accepted, batch_idx, next_draft) + ) + and accepted.dim() == 2 + and next_draft.dim() == 2 + and accepted.stride(1) == 1 + and next_draft.stride(1) == 1 + and 0 < next_draft.shape[0] <= max_requests() + and num_accepted.dim() == 1 + and batch_idx.dim() == 1 + and num_accepted.is_contiguous() + and batch_idx.is_contiguous() + and num_accepted.numel() >= next_draft.shape[0] + and batch_idx.numel() >= next_draft.shape[0] + ) + + +def supports_kv_rewind( + num_accepted: torch.Tensor, kv_lens: torch.Tensor, rewind: torch.Tensor, nc: int, n: int +) -> bool: + """Whether the kernel takes these operands; a caller that gets False runs the in-tree lines.""" + return ( + all(t.is_cuda and t.dtype in _INT_DTYPES for t in (num_accepted, kv_lens, rewind)) + and all(t.dim() == 1 and t.is_contiguous() for t in (num_accepted, kv_lens, rewind)) + and rewind.dtype == num_accepted.dtype + and 0 < n <= max_requests() + and nc >= 0 + and num_accepted.numel() >= nc + n + and kv_lens.numel() >= nc + n + and rewind.numel() >= n + ) + + +def is_compiled_kv_rewind(num_accepted: torch.Tensor, kv_lens: torch.Tensor) -> bool: + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + return ("kv_rewind", num_accepted.dtype, kv_lens.dtype, get_env_enable_pdl()) in _compiled + + +def precompile_kv_rewind(device, dtypes: Sequence[torch.dtype] = (torch.int32,)) -> None: + """Compile the kv_rewind variants (num_accepted dtype x kv_lens dtype) outside capture.""" + for a_dtype in dtypes: + for k_dtype in dtypes: + num = torch.ones(1, dtype=a_dtype, device=device) + torch.ops.trtllm.spec_kv_rewind( + num, torch.zeros(1, dtype=k_dtype, device=device), torch.empty_like(num), 0, 1 + ) + torch.cuda.synchronize() + + +def precompile( + draft_lens: Sequence[int], device, dtypes: Sequence[torch.dtype] = _INT_DTYPES +) -> None: + """Compile every (draft length, operand dtype) variant a one-model worker passes, outside capture: target + tokens and drafts of either int dtype; int32 accepted / num_accepted (this op's) and batch indices (the spec + metadata's); next drafts of either int dtype.""" + for draft_len in draft_lens: + batch_idx = torch.zeros(1, dtype=torch.int32, device=device) + for t_dtype in dtypes: + for d_dtype in dtypes: + target = torch.zeros(draft_len + 1, dtype=t_dtype, device=device) + draft = torch.zeros(1, draft_len, dtype=d_dtype, device=device) + accepted, num = torch.ops.trtllm.spec_strict_accept(target, draft) + for d_dtype in dtypes: + draft = torch.zeros(1, draft_len, dtype=d_dtype, device=device) + torch.ops.trtllm.spec_next_new_tokens(accepted, num, batch_idx, draft) + torch.cuda.synchronize() + + +def is_compiled_strict_accept(target: torch.Tensor, draft: torch.Tensor) -> bool: + """Whether this call's variant is compiled (a captured step may only call compiled variants).""" + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + key = ("strict_accept", target.dtype, draft.dtype, draft.shape[1], get_env_enable_pdl()) + return key in _compiled + + +def is_compiled_next_new_tokens(accepted, num_accepted, batch_idx, next_draft) -> bool: + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + key = ( + "next_new_tokens", + accepted.dtype, + num_accepted.dtype, + batch_idx.dtype, + next_draft.dtype, + next_draft.shape[1], + get_env_enable_pdl(), + ) + return key in _compiled + + +class SpecStrictAccept(OpWrapper): + """Strict acceptance of a decode-only batch: ``target`` [n_req * (D + 1)] (the target's sampled token per verify + row, request-major) and ``draft`` [n_req, D] (the drafts being verified), int32 or int64. + + Returns ``(accepted [n_req, D + 1] int32, num_accepted [n_req] int32)``: the target tokens, and 1 + the length of + the prefix of drafts equal to the target's tokens. No forced-acceptance override (the caller's). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="K3 c1: one request, int32 tokens (the fused sampler's) vs int32 drafts", + spec=dict(n_req=1, draft_len=3), + ), + Cell( + why="K3 c8: the 32-row step, mixed prefix lengths", + spec=dict(n_req=8, draft_len=3), + ), + Cell( + why="K3 at 12 requests: the 48-row bound of the decode targets", + spec=dict(n_req=12, draft_len=3), + ), + Cell( + why="K7 c1 and c8: 8 verify rows per request", + spec=dict(n_req=8, draft_len=7), + ), + Cell( + why="int64 target tokens (torch.argmax, the all-greedy path)", + spec=dict(n_req=8, draft_len=3, target_dtype=torch.int64), + ), + Cell( + why="padded / dropped drafts (-1): never equal to a token", + spec=dict(n_req=8, draft_len=3, pad_drafts=True), + ), + Cell( + why="every draft accepted, and none accepted", + spec=dict(n_req=8, draft_len=3, extremes=True), + ), + Cell( + why="the kernel's request bound (one CTA)", + spec=dict(n_req=64, draft_len=3), + ), + ) + + note = """ + Integer compares and copies only: the reference is exact, and the op must equal it bit for bit. Decode-only + batches (no context rows). The ragged verify layout and the forced-acceptance override are not this op's. + """ + + def __call__(self, target: torch.Tensor, draft: torch.Tensor): + return torch.ops.trtllm.spec_strict_accept(target, draft) + + def reference(self, target: torch.Tensor, draft: torch.Tensor): + """The in-tree strict path (``SpecWorkerBase._sample_and_accept_draft_tokens_base``, no contexts).""" + n_req, draft_len = draft.shape + accepted = torch.empty((n_req, draft_len + 1), dtype=torch.int, device=target.device) + num_accepted = torch.ones(n_req, dtype=torch.int, device=target.device) + gen_target = target.reshape(n_req, draft_len + 1) + num_accepted += torch.cumprod((draft == gen_target[:, :draft_len]).int(), dim=-1).sum(1) + accepted[:, : draft_len + 1] = gen_target + return [accepted, num_accepted] + + +class SpecNextNewTokens(OpWrapper): + """The next step's input tokens: ``out[b] = [accepted[batch_idx[b], num_accepted[b] - 1], *next_draft[b]]``. + + ``accepted`` [rows, >= D + 1] (row stride free), ``num_accepted`` [>= n_req], ``batch_idx`` [>= n_req], + ``next_draft`` [n_req, D] (row stride free); int32 or int64 each. The output dtype is torch's promotion of + ``accepted`` and ``next_draft`` (as ``torch.concat``). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="K3 c1, int32 throughout", + spec=dict(n_req=1, draft_len=3), + ), + Cell( + why="K3 c8 with shuffled batch indices", + spec=dict(n_req=8, draft_len=3, shuffle=True), + ), + Cell( + why="K7 c8", + spec=dict(n_req=8, draft_len=7), + ), + Cell( + why="int64 next drafts: the promoted int64 output", + spec=dict(n_req=8, draft_len=3, draft_dtype=torch.int64), + ), + Cell( + why="a wider accepted buffer (max draft length > runtime) and strided drafts", + spec=dict(n_req=8, draft_len=3, wide=True), + ), + ) + + note = """ + Index copies only: the reference is exact. The op does not bound batch_idx or num_accepted; is_valid states them. + """ + + def __call__(self, accepted, num_accepted, batch_idx, next_draft): + return torch.ops.trtllm.spec_next_new_tokens(accepted, num_accepted, batch_idx, next_draft) + + def reference(self, accepted, num_accepted, batch_idx, next_draft): + """The in-tree ``SpecWorkerBase._prepare_next_new_tokens``.""" + n_req = next_draft.shape[0] + first = accepted[batch_idx[:n_req], num_accepted[:n_req] - 1].unsqueeze(1) + return torch.concat([first, next_draft], dim=1) + + def is_valid(self, accepted, num_accepted, batch_idx, next_draft) -> None: + n_req = next_draft.shape[0] + assert bool(((batch_idx[:n_req] >= 0) & (batch_idx[:n_req] < accepted.shape[0])).all()) + assert bool( + ((num_accepted[:n_req] >= 1) & (num_accepted[:n_req] <= accepted.shape[1])).all() + ) + + +class SpecKvRewind(OpWrapper): + """DFlash's draft-forward KV bookkeeping for the generation rows ``nc .. nc + n`` after the forced-acceptance + override: ``rewind[:n] = 1 - num_accepted[nc : nc + n]`` and ``kv_lens[nc : nc + n] += 1`` (in place). + + ``num_accepted`` and ``kv_lens`` int32 or int64 (1-D, contiguous); ``rewind`` [>= n] of num_accepted's dtype. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="K3 c1: one generation row, int32 counts and lengths", + spec=dict(nc=0, n=1), + ), + Cell( + why="K3 c8: eight generation rows", + spec=dict(nc=0, n=8), + ), + Cell( + why="context rows first (nc > 0): only the generation rows change", + spec=dict(nc=3, n=5), + ), + Cell( + why="int64 kv_lens", + spec=dict(nc=0, n=8, kv_dtype=torch.int64), + ), + Cell( + why="the kernel's request bound (one CTA)", + spec=dict(nc=0, n=64), + ), + ) + + note = """ + Integer arithmetic only: the reference is exact. kv_lens is updated in place; rows outside nc .. nc + n are + untouched. The caller decides the step (it runs where the in-tree lines run, after the forced override). + """ + + def __call__(self, num_accepted, kv_lens, rewind, nc: int, n: int): + torch.ops.trtllm.spec_kv_rewind(num_accepted, kv_lens, rewind, int(nc), int(n)) + return rewind[:n] + + def reference(self, num_accepted, kv_lens, rewind, nc: int, n: int): + """DFlashWorker._prepare_kv_for_draft_forward's two lines.""" + out = 1 - num_accepted[nc : nc + n] + kv_lens[nc : nc + n] += 1 + return out + + +spec_strict_accept = SpecStrictAccept() +spec_next_new_tokens = SpecNextNewTokens() +spec_kv_rewind = SpecKvRewind() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/strict_accept_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/strict_accept_kernel.py new file mode 100644 index 000000000000..bfec9e460a7d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/spec_decode/strict_accept_kernel.py @@ -0,0 +1,174 @@ +# 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. +# ============================================================================= +# Strict draft acceptance and the next step's input tokens -- CuTe DSL kernels (one CTA, one thread per request) +# ============================================================================= +# +# strict_accept, for request b (draft length D, K1 = D + 1 verify rows per request): +# accepted[b, j] = target[b K1 + j] j = 0 .. D (int32) +# num_accepted[b] = 1 + the length of the prefix k = 0, 1, .. where draft[b, k] == target[b K1 + k] +# which is the strict path's `1 + cumprod(draft == target[:, :D]).sum(1)`. Tokens are compared as int64, so an +# int32 / int64 mix compares the values torch's promotion compares. +# +# next_new_tokens, for request b (D next draft tokens): +# out[b, 0] = accepted[batch_idx[b], num_accepted[b] - 1] +# out[b, 1 + j] = next_draft[b, j] j = 0 .. D - 1 +# which is `concat([accepted[batch_idx, num_accepted - 1].unsqueeze(1), next_draft], 1)`. +# +# kv_rewind, for generation request b (rows nc .. nc + n of the batch), after the forced-acceptance override: +# rewind[b] = 1 - num_accepted[nc + b] +# kv_lens[nc + b] += 1 +# which is DFlash's `_kv_rewind_amount = 1 - num_accepted[nc:bs]` and `kv_lens_cuda[nc:bs] += 1`. +# +# Integer copies and compares only: bitwise to the torch path by construction. Thread b handles request b; the +# grid is one CTA of MAX_REQUESTS threads, so a launch serves up to MAX_REQUESTS requests. +# ============================================================================= +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +from cutlass.experimental import primitives as prims + +MAX_REQUESTS = 64 # capacity: fixed: one CTA, one thread per request + + +@cute.kernel +def strict_accept_kernel( + target: cute.Tensor, # int32 or int64 [n_req * (D + 1)] + draft: cute.Tensor, # int32 or int64 [n_req * D] + accepted: cute.Tensor, # int32 [n_req * (D + 1)], out + num_accepted: cute.Tensor, # int32 [n_req], out + n_req: cutlass.Int32, + draft_len: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + prims.griddepcontrol(prims.GridDepAction.WAIT) + if tidx < n_req: + k1 = draft_len + 1 + run = cutlass.Int32(1) + count = cutlass.Int32(0) + for k in cutlass.range_constexpr(draft_len): + tok = cutlass.Int64(target[tidx * k1 + k]) + drf = cutlass.Int64(draft[tidx * draft_len + k]) + accepted[tidx * k1 + k] = cutlass.Int32(tok) + run = cutlass.select_(tok == drf, run, cutlass.Int32(0)) + count = count + run + accepted[tidx * k1 + draft_len] = cutlass.Int32( + cutlass.Int64(target[tidx * k1 + draft_len]) + ) + num_accepted[tidx] = count + cutlass.Int32(1) + + +@cute.jit +def strict_accept( + target: cute.Tensor, + draft: cute.Tensor, + accepted: cute.Tensor, + num_accepted: cute.Tensor, + n_req: cutlass.Int32, + draft_len: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + strict_accept_kernel(target, draft, accepted, num_accepted, n_req, draft_len).launch( + grid=[1, 1, 1], block=[MAX_REQUESTS, 1, 1], stream=stream, use_pdl=use_pdl + ) + + +@cute.kernel +def next_new_tokens_kernel( + accepted: cute.Tensor, # [rows, accepted_stride] (flat), any int dtype + num_accepted: cute.Tensor, # [n_req], any int dtype + batch_idx: cute.Tensor, # [n_req], any int dtype + next_draft: cute.Tensor, # [n_req, draft_stride] (flat), any int dtype + out: cute.Tensor, # [n_req * (D + 1)], out + accepted_stride: cutlass.Int64, + draft_stride: cutlass.Int64, + n_req: cutlass.Int32, + draft_len: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + prims.griddepcontrol(prims.GridDepAction.WAIT) + if tidx < n_req: + k1 = draft_len + 1 + row = cutlass.Int64(batch_idx[tidx]) + col = cutlass.Int64(num_accepted[tidx]) - cutlass.Int64(1) + out[tidx * k1] = out.element_type(accepted[row * accepted_stride + col]) + for j in cutlass.range_constexpr(draft_len): + out[tidx * k1 + 1 + j] = out.element_type( + next_draft[cutlass.Int64(tidx) * draft_stride + j] + ) + + +@cute.jit +def next_new_tokens( + accepted: cute.Tensor, + num_accepted: cute.Tensor, + batch_idx: cute.Tensor, + next_draft: cute.Tensor, + out: cute.Tensor, + accepted_stride: cutlass.Int64, + draft_stride: cutlass.Int64, + n_req: cutlass.Int32, + draft_len: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + next_new_tokens_kernel( + accepted, + num_accepted, + batch_idx, + next_draft, + out, + accepted_stride, + draft_stride, + n_req, + draft_len, + ).launch(grid=[1, 1, 1], block=[MAX_REQUESTS, 1, 1], stream=stream, use_pdl=use_pdl) + + +@cute.kernel +def kv_rewind_kernel( + num_accepted: cute.Tensor, # [>= nc + n], any int dtype + kv_lens: cute.Tensor, # [>= nc + n], any int dtype, updated in place + rewind: cute.Tensor, # [n], num_accepted's dtype, out + nc: cutlass.Int32, + n: cutlass.Int32, +): + tidx, _, _ = cute.arch.thread_idx() + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + prims.griddepcontrol(prims.GridDepAction.WAIT) + if tidx < n: + row = nc + tidx + rewind[tidx] = rewind.element_type(cutlass.Int64(1) - cutlass.Int64(num_accepted[row])) + kv_lens[row] = kv_lens.element_type(cutlass.Int64(kv_lens[row]) + cutlass.Int64(1)) + + +@cute.jit +def kv_rewind( + num_accepted: cute.Tensor, + kv_lens: cute.Tensor, + rewind: cute.Tensor, + nc: cutlass.Int32, + n: cutlass.Int32, + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + kv_rewind_kernel(num_accepted, kv_lens, rewind, nc, n).launch( + grid=[1, 1, 1], block=[MAX_REQUESTS, 1, 1], stream=stream, use_pdl=use_pdl + ) 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..106f61daffe6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/fake_ops.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fake (meta) implementations of the torch ops that the MiniMax-M3 decode targets add (``trtllm::step_inputs_scatter`` +and the ``trtllm::minimax_m3_*`` decode ops in ``cpp/tensorrt_llm/thop/minimaxM3/``), for torch.compile and fake-tensor +tracing. Only loaded native ops without an existing fake are registered. Importing the family must remain safe +when the native extension is absent, so unsupported routes can fall back to the built-in model. Call register() +again after loading an overlay; registration is idempotent.""" + +import torch + +from ...catalog.fakes.m31_fake_ops import _has_fake + +_REGISTERED = set() + + +def _step_inputs_scatter_fake(stage): + # Writes the staged rows into their destinations; returns nothing. + return None + + +def _minimax_m3_mxfp8_low_m_gemm_fake(x, w, wsf, warps, cluster_k, row_tiles, pdl): + return x.new_empty((x.shape[0], w.shape[0])) + + +def _minimax_m3_router_gemv_fake(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), + ] + + +def _minimax_m3_router_gemv_fp4_fake(x, w, sf_scale, pdl): + return _minimax_m3_router_fp4_outputs(x, w) + + +def _minimax_m3_router_tc_fake(x, w, sf_scale, pdl): + outputs = _minimax_m3_router_fp4_outputs(x, w) + return outputs if sf_scale is not None else outputs[:1] + + +def _minimax_m3_producer_pw_fake( + 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), + ] + + +def _minimax_m3_mc_allreduce_norm_fake( + 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 + + +FAKES = { + "step_inputs_scatter": _step_inputs_scatter_fake, + "minimax_m3_mxfp8_low_m_gemm": _minimax_m3_mxfp8_low_m_gemm_fake, + "minimax_m3_router_gemv": _minimax_m3_router_gemv_fake, + "minimax_m3_router_gemv_fp4": _minimax_m3_router_gemv_fp4_fake, + "minimax_m3_router_tc": _minimax_m3_router_tc_fake, + "minimax_m3_producer_pw": _minimax_m3_producer_pw_fake, + "minimax_m3_mc_allreduce_norm": _minimax_m3_mc_allreduce_norm_fake, +} + + +def register(): + """Register available native ops, preserving existing Meta kernels or Python fakes.""" + for name, fake in FAKES.items(): + if name in _REGISTERED or not hasattr(torch.ops.trtllm, name): + continue + qualname = f"trtllm::{name}" + if not _has_fake(qualname): + torch.library.register_fake(qualname)(fake) + _REGISTERED.add(name) + + +register() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/__init__.py new file mode 100644 index 000000000000..a5aeb0abbdbb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 NVFP4 / sm_103 / tp4.""" + +# Routing has selected this target before its native fake registrations are needed. +from .. import fake_ops + +fake_ops.register() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/dspark.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/dspark.py new file mode 100644 index 000000000000..ed54c31205a7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/dspark.py @@ -0,0 +1,631 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-model DSpark on this target (phase W0): the in-tree DSpark drafter and worker, unchanged. + +The target keeps its own forward; the shell hands the step to ``M3DSparkWorker``, which is the in-tree +``DSparkWorker`` plus the five calls the shell makes on any worker (``capture_hidden_states``, +``arm_capture_norm``, ``tail_eligible``, ``prebuild``, ``__call__``). In W0 none of them takes a fused path: the +captures are the in-tree DFlash captures, the target logits go through the logits processor and the in-tree +strict accept verifies them, so the speculation is the in-tree one by construction. + +The drafter is the in-tree standalone DSpark drafter for the checkpoint's architecture, built through the draft +model registry. Its context K/V lives in the separate draft KV cache, which must keep the drafter's projection +dtype (BF16): a pool of another dtype is not bound and the worker falls back to a private arena sized by the +served sequence length. So the drafter never takes the target's KV quantization: neither the FP8 draft-KV +override (an Eagle3 need) nor the target's own KV algo (NVFP4 under ``kv_cache_config.dtype="nvfp4"``). +""" + +from __future__ import annotations + +import os +import time +from typing import Tuple + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm import ( + allgather_cast as _allgather_cast, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode import ( + strict_accept as _strict_accept, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_block_tables import ( + enable_parallel_refresh, + refresh_block_tables, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_ctx_advance import ( + advance_ctx_and_read_view, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_ctx_store import ( + fused_ctx_store, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_drafter_inputs import ( + drafter_inputs as _drafter_inputs, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_forward_indices import ( + forward_indices, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_qk_norm_rope import ( + qk_norm_rope, +) +from tensorrt_llm._torch.distributed import allgather +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_speculative import ( + external_drafter_config_kwargs, + get_draft_model, +) +from tensorrt_llm._torch.speculative.dspark import DSparkWorker +from tensorrt_llm._torch.speculative.interface import should_use_separate_draft_kv_cache +from tensorrt_llm.logger import logger + +from .dspark_s4 import MAX_FAST_BATCH, S4WorkerMixin +from .m31_switches import m31_disabled +from .mxdraft import route_mx_linears +from .ops.embedding.m3_embedding import m3_embedding + +# MXDRAFT (MXFP8 drafter GEMV): ON by default. With the K3 bundle off (the ship config), its powered K7 gate c +# (n 1,776) passes: mean d +0.009 AL, CI [-0.019, +0.037]. Kill switch TRTLLM_M31_DISABLE_MXDRAFT=1, read at import. +MXDRAFT = not m31_disabled("MXDRAFT") + + +def is_standalone_dspark(spec_config) -> bool: + """Whether the run speculates with a standalone DSpark drafter (its own checkpoint, ``decoding_type: DSpark``).""" + return bool( + spec_config is not None + and spec_config.spec_dec_mode.use_one_engine() + and spec_config.spec_dec_mode.is_dspark() + and not getattr(spec_config, "draft_is_embedded_in_target", False) + ) + + +class M3EmbedTokens(nn.Module): + """The target's token embedding as the module the in-tree drafter borrows (``load_weights_from_target_model``). + + The drafter embeds the bonus and mask tokens through ``embed_tokens(ids)``; the target keeps its replicated + table as a raw tensor and gathers it with ``m3_embedding``, so this module calls the same gather on the same + table (an exact gather, so the drafter sees the rows the target embeds). + """ + + def __init__(self, core: nn.Module) -> None: + super().__init__() + self._core = [core] # not a submodule: the table belongs to the target + + @property + def weight(self) -> torch.Tensor: + return self._core[0].w["embed"] + + def forward(self, ids: torch.Tensor) -> torch.Tensor: + return m3_embedding(ids, self.weight) + + +def s10c_c_drafter_inputs( + embed_tokens, + accepted, + num_accepted, + slots, + ctx_len, + max_ctx, + mask_token_id, + trained_mask_embedding, + block_size, + tokens_per_request, +): + """S10c-C: the drafter's step inputs in one launch (catalog/spec_decode/dspark_drafter_inputs), for this target's + embedding only: ``M3EmbedTokens`` is an exact gather of the replicated table, so the noise block is table rows (the + trained mask row if the drafter keeps one). None for any other embedding: the stock path runs.""" + if not isinstance(embed_tokens, M3EmbedTokens): + return None + table = embed_tokens.weight + mask_row = ( + trained_mask_embedding if trained_mask_embedding is not None else table[int(mask_token_id)] + ) + return _drafter_inputs( + accepted, + num_accepted, + slots, + ctx_len, + max_ctx, + table, + mask_row, + block_size, + tokens_per_request, + ) + + +# N3 A2 (switch TRTLLM_M31_DISABLE_N3A): the target's strict acceptance and next-step input tokens in one launch each +# (catalog/spec_decode/strict_accept), bitwise to the in-tree torch path. Read once, at import. +N3A_DISABLED = m31_disabled("N3A") +# N3A-KV (switch TRTLLM_M31_DISABLE_N3A_KV, and off with N3A): DFlash's draft-forward KV bookkeeping (the rewind amount +# 1 - num_accepted and kv_lens += 1) in one launch where the in-tree lines run, after the forced-acceptance override, +# bitwise (integer arithmetic). Read once, at import. +N3A_KV_DISABLED = N3A_DISABLED or m31_disabled("N3A_KV") +# Gate-only cross-check (TRTLLM_M31_CHECK_N3A_KV=1; never in a perf arm): every N3A-KV call is compared with the stock +# lines on copies (graph-safe device counters; reported on eager steps). +N3A_KV_CHECK = os.environ.get("TRTLLM_M31_CHECK_N3A_KV", "").strip() not in ("", "0") +# N3 P-gid + P-cat (switch TRTLLM_M31_DISABLE_N3P): the verify-row target logits without the identity gather_ids copy +# and with the vocab concat and the fp32 cast in one strided copy, bitwise to the logits processor. Read at import. +N3P_DISABLED = m31_disabled("N3P") +# N3 A1 (switch TRTLLM_M31_DISABLE_A1): the verify logits' vocab all-gather, concat and fp32 cast in one launch over a +# multicast NVLink workspace (catalog/comm/allgather_cast), bitwise to the NCCL all-gather + strided copy below. K3 +# verify rows (B x (K + 1), up to 64 rows: K3 batches 1..16, K7 batches 1..8). Read once, at import. +A1_DISABLED = m31_disabled("A1") +A1_MAX_ROWS = 64 # capacity: fixed: allgather_cast.MAX_ROWS (K3 at batch 16, K7 at batch 8) +# A1 at draft lengths other than 3 (switch TRTLLM_M31_DISABLE_A1_K7, K3 keeps A1). +A1_K7_DISABLED = m31_disabled("A1_K7") +_n3a_logged = False + + +def _n3a_log_once() -> None: + global _n3a_logged + if not _n3a_logged: + _n3a_logged = True + logger.info( + "[M3 dispatch] n3_strict_accept taken: fused strict acceptance + next-step input tokens" + ) + + +_n3p_logged = False +_n3a_kv_logged = False + + +def _n3a_kv_log_once() -> None: + global _n3a_kv_logged + if not _n3a_kv_logged: + _n3a_kv_logged = True + logger.info( + "[M3 dispatch] n3a_kv_rewind taken: draft-forward KV rewind + kv_lens bump in one launch" + ) + + +_a1_logged = set() + + +def _a1_log_once(key, message: str) -> None: + if key not in _a1_logged: + _a1_logged.add(key) + logger.info(message) + + +def _n3p_log_once() -> None: + global _n3p_logged + if not _n3p_logged: + _n3p_logged = True + logger.info( + "[M3 dispatch] n3p_target_logits taken: verify-row logits without the gather_ids copy and the concat" + ) + + +class M3DSparkWorker(S4WorkerMixin, DSparkWorker): + """The in-tree ``DSparkWorker`` with the calls this target's shell makes on its worker. W1: the one-launch Markov + chain (``S4WorkerMixin``, switch ``TRTLLM_M31_DRAFT_DISABLE_S4``). N3 A2: the strict acceptance and the next-step + input tokens in one launch each (switch ``TRTLLM_M31_DISABLE_N3A``). Everything else is the in-tree path.""" + + def capture_hidden_states( + self, spec_metadata, layer_id: int, hidden_states, residual=None + ) -> None: + """The in-tree DFlash capture: layer output (hidden + residual) into the static capture buffer.""" + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + + def arm_capture_norm( + self, attn_metadata, spec_metadata, lm_head: nn.Module, draft_model + ) -> None: + """No capture norm: DSpark's fc reads the raw captured layers.""" + return None + + def tail_eligible(self, attn_metadata, spec_metadata, lm_head: nn.Module) -> bool: + """Never: the target logits always go through the logits processor and the in-tree strict accept.""" + return False + + _a1 = None # the N3 A1 workspace (allgather_cast's arguments) once set up + + def prebuild(self, lm_head: nn.Module, draft_model) -> None: + """The one-launch Markov chain's shards, buffers, exchange workspace and compiles, then the N3 A2 compiles + (rank-local) and the N3 A1 workspace and compiles (collective), outside capture, every rank, in this order; + no fused verify, no coded target head, no fused draft loop.""" + rank = getattr(getattr(self, "mapping", None), "tp_rank", None) + # S10: the drafter's strided-view glue (bitwise; kill switch TRTLLM_M31_DISABLE_S10, read at import). + lean = not m31_disabled("S10") + self.enable_lean_glue(draft_model, lean) + k = getattr(getattr(self, "spec_config", None), "max_draft_len", 0) or 0 + logger.info(f"M3.1 DSpark S10: lean drafter glue {'on' if lean else 'off'} on rank {rank}") + # S10c-A: the drafter's context K / V epilogue + accept mask + paged store in one launch after the K / V GEMM + # (bitwise). Kill switches TRTLLM_M31_DISABLE_S10C (all of S10c) / TRTLLM_M31_DISABLE_S10C_A, read at import. + s10c_a = self.enable_fused_ctx_store( + None if m31_disabled("S10C") or m31_disabled("S10C_A") else fused_ctx_store + ) + logger.info( + f"M3.1 DSpark S10c-A: fused drafter ctx store {'on' if s10c_a else 'off'} on rank {rank}" + ) + # S10c-C: the drafter's step inputs (bonus, positions, noise block) in one launch (bitwise: integers and row + # copies). Kill switches TRTLLM_M31_DISABLE_S10C (all of S10c) / TRTLLM_M31_DISABLE_S10C_C, read at import. + s10c_c = self.enable_fused_drafter_inputs( + None if m31_disabled("S10C") or m31_disabled("S10C_C") else s10c_c_drafter_inputs + ) + logger.info( + f"M3.1 DSpark S10c-C: fused drafter step inputs {'on' if s10c_c else 'off'} on rank {rank}" + ) + # S10c-D: the drafter's context advance + read view in one launch (bitwise, integer). Kill switches + # TRTLLM_M31_DISABLE_S10C (all of S10c) / TRTLLM_M31_DISABLE_S10C_D, read at import. + s10c_d = self.enable_fused_ctx_advance( + None if m31_disabled("S10C") or m31_disabled("S10C_D") else advance_ctx_and_read_view + ) + logger.info( + f"M3.1 DSpark S10c-D: fused drafter context advance {'on' if s10c_d else 'off'} on rank {rank}" + ) + # S10c-B: the drafter's block-table refresh (decode + placeholder count) in one launch (bitwise, integer). + # Kill switches TRTLLM_M31_DISABLE_S10C (all of S10c) / TRTLLM_M31_DISABLE_S10C_B, read at import. + s10c_b = self.enable_fused_block_refresh( + None if m31_disabled("S10C") or m31_disabled("S10C_B") else refresh_block_tables + ) + # S10c-B2: that refresh on a (row, column-chunk) grid (bitwise; flat at short tables, removes the serial + # kernel's growth past ~11K blocks per sequence). Kill switch TRTLLM_M31_DISABLE_S10C_B2, read at import. + s10c_b2 = enable_parallel_refresh(s10c_b and not m31_disabled("S10C_B2")) + logger.info( + f"M3.1 DSpark S10c-B: fused drafter block-table refresh {'on' if s10c_b else 'off'} on rank {rank}" + f" (S10c-B2 2-D grid {'on' if s10c_b2 else 'off'})" + ) + # S10c-E: the draft model's paged block-decode index prelude in one launch (bitwise, integer). Kill switches + # TRTLLM_M31_DISABLE_S10C (all of S10c) / TRTLLM_M31_DISABLE_S10C_E, read at import. + s10c_e = self.enable_fused_forward_indices( + draft_model, None if m31_disabled("S10C") or m31_disabled("S10C_E") else forward_indices + ) + logger.info( + f"M3.1 DSpark S10c-E: fused drafter forward indices {'on' if s10c_e else 'off'} on rank {rank}" + ) + # EX-11: the drafter block's q / k RMSNorm + partial RoPE in one launch (bitwise to flashinfer's three). Kill + # switch TRTLLM_M31_DISABLE_EX11, read at import. + ex11 = self.enable_fused_qk_norm_rope( + draft_model, None if m31_disabled("EX11") else qk_norm_rope + ) + logger.info( + f"M3.1 DSpark EX-11: fused drafter q/k norm + RoPE {'on' if ex11 else 'off'} on rank {rank}" + ) + # MXDRAFT: the drafter's qkv / o / gate_up / down GEMMs on this target's MXFP8 GEMV for the decode passes' + # rows (B x K, B <= MAX_FAST_BATCH), from the exact MXFP8 form of the shipped weights (drafter-only precision). + # ON by default (see MXDRAFT); kill switch TRTLLM_M31_DISABLE_MXDRAFT. + mx = ( + 0 + if not MXDRAFT or k <= 0 + else route_mx_linears(draft_model, [b * k for b in range(1, MAX_FAST_BATCH + 1)]) + ) + logger.info(f"M3.1 DSpark MXDRAFT: MXFP8 drafter GEMV on {mx} projections on rank {rank}") + self.prebuild_s4(lm_head, draft_model) + if not N3A_DISABLED and lm_head is not None: + _strict_accept.precompile((self.max_draft_len,), lm_head.weight.device) + self._a1 = self._prebuild_a1(lm_head) if lm_head is not None else None + n3a_kv = False + if not N3A_KV_DISABLED and lm_head is not None: + _strict_accept.precompile_kv_rewind(lm_head.weight.device) + self._n3a_kv_buf = torch.empty( + _strict_accept.max_requests(), dtype=torch.int32, device=lm_head.weight.device + ) + n3a_kv = self.enable_fused_kv_rewind(self._n3a_kv_rewind) + logger.info( + f"M3.1 N3A-KV: one-launch draft-forward KV bookkeeping {'on' if n3a_kv else 'off'} on rank " + f"{getattr(getattr(self, 'mapping', None), 'tp_rank', None)}" + ) + + _n3a_kv_buf = None # the persistent rewind amounts of N3A-KV (int32, MAX_REQUESTS) + _n3a_kv_counts = None # gate-only cross-check counters (calls, mismatches) + + def _n3a_kv_rewind(self, num_accepted, kv_lens, nc: int, bs: int): + """N3A-KV for the DFlash hook: ``1 - num_accepted[nc:bs]`` (a view of the persistent buffer) with + ``kv_lens[nc:bs] += 1`` in one launch, or None (the stock lines) for operands or a variant it does not take.""" + n = bs - nc + buf = self._n3a_kv_buf + if buf is None or not _strict_accept.supports_kv_rewind(num_accepted, kv_lens, buf, nc, n): + return None + if not _strict_accept.is_compiled_kv_rewind(num_accepted, kv_lens): + return None + _n3a_kv_log_once() + if not N3A_KV_CHECK: + return _strict_accept.spec_kv_rewind(num_accepted, kv_lens, buf, nc, n) + want_rewind = 1 - num_accepted[nc:bs] + want_kv = kv_lens[nc:bs] + 1 + out = _strict_accept.spec_kv_rewind(num_accepted, kv_lens, buf, nc, n) + if self._n3a_kv_counts is None: + self._n3a_kv_counts = torch.zeros(2, dtype=torch.int64, device=kv_lens.device) + bad = (out != want_rewind).any() | (kv_lens[nc:bs] != want_kv).any() + self._n3a_kv_counts[0] += 1 + self._n3a_kv_counts[1] += bad.to(torch.int64) + if not torch.cuda.is_current_stream_capturing(): + calls, mism = self._n3a_kv_counts.tolist() + logger.info(f"M3.1 N3A-KV check: {calls} calls, {mism} mismatches") + return out + + def _a1_all_ok(self, local: bool) -> bool: + """Whether ``local`` holds on every rank of the TP group: A1 is set up on every rank or on none (a rank that + is not eligible, or whose compile failed, turns it off everywhere).""" + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _mnnvl_workspace_all_succeeded, + ) + + return _mnnvl_workspace_all_succeeded(_get_mnnvl_workspace_comm(self.mapping), local) + + def _prebuild_a1(self, lm_head: nn.Module): + """The N3 A1 workspace and its K3 verify-row compiles, or None (every rank alike): off by its switch, at + another draft length, without a column-parallel bf16 head over a TP group, or when a compile fails on any + rank.""" + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + mapping = self.mapping + if mapping.tp_size <= 1 or mapping.enable_attention_dp: + return None + cols = int(getattr(lm_head, "out_features", 0)) + sms = torch.cuda.get_device_properties(lm_head.weight.device).multi_processor_count + # Why this rank would not set A1 up (empty: it would). The grid is a function of rows and the SM count: + # requiring every SM up to the workspace's MAX_GRID makes it a function of rows alone, the same on every rank. + off = [ + why + for why, bad in ( + ("TRTLLM_M31_DISABLE_A1 set", A1_DISABLED), + ( + f"draft length {int(self.max_draft_len)} with TRTLLM_M31_DISABLE_A1_K7 set", + int(self.max_draft_len) != 3 and A1_K7_DISABLED, + ), + ( + f"draft length {int(self.max_draft_len)}: more than {A1_MAX_ROWS} verify rows at batch 1", + int(self.max_draft_len) + 1 > A1_MAX_ROWS, + ), + ( + "not a column-parallel, gathered, unpadded, bias-free head", + getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN + or not getattr(lm_head, "gather_output", False) + or getattr(lm_head, "padding_size", 0) != 0 + or getattr(lm_head, "bias", None) is not None, + ), + (f"head dtype {lm_head.weight.dtype}", lm_head.weight.dtype != torch.bfloat16), + ( + f"{cols} shard columns (whole 16-byte vectors needed)", + cols <= 0 or cols % 8 != 0, + ), + ( + f"{sms} SMs (fewer than {_allgather_cast.MAX_GRID})", + sms < _allgather_cast.MAX_GRID, + ), + ) + if bad + ] + if not self._a1_all_ok(not off): + logger.info( + f"M3.1 N3 A1: off on rank {mapping.tp_rank} ({'; '.join(off) or 'off on another rank'}); " + "the verify logits use the NCCL all-gather" + ) + return None + k1 = int(self.max_draft_len) + 1 + rows = tuple(b * k1 for b in range(1, A1_MAX_ROWS // k1 + 1)) + start = time.perf_counter() + if not _allgather_cast.precompile( + lm_head.mapping, cols, rows, lm_head.weight.device, agree=self._a1_all_ok + ): + logger.warning( + "M3.1 N3 A1: a compile failed on some rank; every rank keeps the NCCL all-gather" + ) + return None + ws = _allgather_cast.workspace_args(lm_head.mapping, cols) + # The engagement line: one per rank, at server start. + logger.info( + f"M3.1 N3 A1: one-launch verify all-gather set up on rank {mapping.tp_rank}/{mapping.tp_size} " + f"(shard {cols} columns, rows {rows}, workspace {ws['ws_uc'].numel() * 4} bytes/rank, " + f"set-up {time.perf_counter() - start:.2f} s)" + ) + return ws + + def _n3a_eligible( + self, logits, draft_tokens, num_contexts: int, batch_size: int, spec_metadata + ) -> bool: + """Whether this step's strict acceptance takes the fused op: a decode-only, non-ragged step whose logits hold + exactly batch_size x (D + 1) verify rows and whose drafts the op takes. Host state and shapes only, so it is + fixed per captured graph and the same on every rank. Decided before the target tokens are sampled, so the + sampler runs exactly once either way.""" + if N3A_DISABLED or spec_metadata is None or spec_metadata.is_ragged_verify: + return False + if num_contexts != 0 or batch_size <= 0 or draft_tokens.dim() != 2: + return False + n_req, draft_len = draft_tokens.shape + if n_req != batch_size or draft_len <= 0 or n_req > _strict_accept.max_requests(): + return False + rows = logits.shape[0] if logits.dim() > 1 else 1 + return rows == batch_size * (draft_len + 1) and draft_tokens.is_contiguous() + + def _sample_and_accept_draft_tokens_base( + self, logits, draft_tokens, num_contexts, batch_size, spec_metadata + ): + """The in-tree strict acceptance with its epilogue in one launch (N3 A2): the same target tokens (the in-tree + sampler, called once), the same accepted tokens and counts, then the in-tree forced-acceptance override (its + RNG pool and counters advance as on an in-tree step).""" + if not self._n3a_eligible(logits, draft_tokens, num_contexts, batch_size, spec_metadata): + return super()._sample_and_accept_draft_tokens_base( + logits, draft_tokens, num_contexts, batch_size, spec_metadata + ) + if logits.dim() == 1: + logits = logits.unsqueeze(0) + target_tokens = self._sample_tokens_for_batch( + logits, spec_metadata, num_contexts, batch_size + ) + if _strict_accept.supports_strict_accept( + target_tokens, draft_tokens + ) and _strict_accept.is_compiled_strict_accept(target_tokens, draft_tokens): + accepted, num_accepted = _strict_accept.spec_strict_accept(target_tokens, draft_tokens) + _n3a_log_once() + else: + # The sampler already ran: finish with the in-tree epilogue (the op's reference is that code). + accepted, num_accepted = _strict_accept.spec_strict_accept.reference( + target_tokens, draft_tokens + ) + num_accepted = self._apply_force_accepted_tokens( + num_accepted, num_contexts, draft_tokens.shape[-1], spec_metadata=spec_metadata + ) + return accepted, num_accepted + + def target_logits(self, hidden_states, attn_metadata, spec_metadata, lm_head: nn.Module): + """The verify rows' gathered fp32 target logits (N3 P-gid + P-cat), or None for the logits processor. + + On a decode-only, non-ragged step every row is a verify row and ``gather_ids`` is ``arange(num_tokens)`` (each + generation request extends it with the next ``1 + D`` row indices), so the rows are a slice, not a gather. + The processor's head is the column-parallel BF16 GEMM, an all-gather along the vocab (NCCL, then a concat) + and ``.float()``; here the same all-gather runs along dim 0 (no concat) and one strided copy writes the + rank-major shards into the fp32 output: the same bytes, an exact widening. Shapes and host state only, so + the decision is fixed per captured graph and the same on every rank.""" + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + if ( + N3P_DISABLED + or spec_metadata is None + or getattr(spec_metadata, "is_ragged_verify", False) + ): + return None + if attn_metadata.num_contexts != 0 or spec_metadata.gather_ids is None: + return None + rows = hidden_states.shape[0] + if spec_metadata.gather_ids.shape[0] != rows or rows == 0: + return None + if ( + getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN + or not getattr(lm_head, "gather_output", False) + or getattr(lm_head, "padding_size", 0) != 0 + or getattr(lm_head, "bias", None) is not None + ): + return None + shard = lm_head.apply_linear(hidden_states, None) + _n3p_log_once() + tp = lm_head.tp_size + if tp == 1: + return shard.float() + # N3 A1: one launch for the all-gather, the concat and the cast. The decision depends on the rows and the + # set-up only (the same on every rank, fixed per captured graph). + if self._a1 is not None: + if rows <= A1_MAX_ROWS and _allgather_cast.is_compiled( + rows, shard.shape[-1], tp, 0, shard.device + ): + _a1_log_once( + "taken", + "[M3 dispatch] n3_a1 allgather_cast taken: verify logits gathered in one launch", + ) + return _allgather_cast.allgather_cast(shard, **self._a1) + _a1_log_once( + rows, + f"[M3 dispatch] n3_a1 allgather_cast not taken: {rows} verify rows (compiled: verify rows up to " + f"{A1_MAX_ROWS}); NCCL all-gather", + ) + gathered = allgather(shard, lm_head.mapping, dim=0) + vocab_shard = shard.shape[-1] + out = torch.empty((rows, tp * vocab_shard), dtype=torch.float32, device=shard.device) + out.view(rows, tp, vocab_shard).copy_(gathered.view(tp, rows, vocab_shard).transpose(0, 1)) + return out + + def _prepare_next_new_tokens( + self, + accepted_tokens, + next_draft_tokens, + batch_indices_cuda, + batch_size, + num_accepted_tokens, + ): + """The next step's input tokens in one launch (N3 A2), else the in-tree index + concat.""" + batch_idx = batch_indices_cuda[:batch_size] + if ( + not N3A_DISABLED + and next_draft_tokens.dim() == 2 + and next_draft_tokens.shape[0] == batch_size + and _strict_accept.supports_next_new_tokens( + accepted_tokens, num_accepted_tokens, batch_idx, next_draft_tokens + ) + and _strict_accept.is_compiled_next_new_tokens( + accepted_tokens, num_accepted_tokens, batch_idx, next_draft_tokens + ) + ): + return _strict_accept.spec_next_new_tokens( + accepted_tokens, num_accepted_tokens, batch_idx, next_draft_tokens + ) + return super()._prepare_next_new_tokens( + accepted_tokens, next_draft_tokens, batch_indices_cuda, batch_size, num_accepted_tokens + ) + + +def build_dspark( + shell, model_config: ModelConfig +) -> Tuple[ModelConfig, nn.Module, M3DSparkWorker, bool]: + """The drafter config, the drafter and the worker of one-model standalone DSpark (mirrors + ``SpecDecOneEngineForCausalLM`` for an external drafter).""" + spec = model_config.spec_config + if not is_standalone_dspark(spec): + raise ValueError("build_dspark serves the standalone DSpark drafter only") + draft_config = ModelConfig.from_pretrained( + spec.speculative_model, **external_drafter_config_kwargs(model_config, spec) + ) + # The drop is safe only because draft_config owns its QuantConfig: from_pretrained reads the drafter checkpoint + # and takes no quant_config from the target. Never pass the target's quant_config in here. + keep_drafter_kv_dtype(draft_config) + draft_config.extra_attrs = model_config.extra_attrs + use_separate = should_use_separate_draft_kv_cache(spec) + draft = get_draft_model(model_config, draft_config, shell.lm_head, shell.model) + # A parallel drafter (DSpark) projects its block rows through the target's logits processor, as the one-engine + # construction sets it. + draft.logits_processor = shell.logits_processor + worker = M3DSparkWorker(spec, model_config.mapping, use_separate_draft_kv_cache=use_separate) + worker.set_draft_model(draft) + shell.epilogue.append(draft) + shell.epilogue.append(worker) + logger.info( + f"M3.1 target DSpark: {type(worker).__name__} over {type(draft).__name__}, K={spec.max_draft_len}, " + f"separate draft KV cache={use_separate}" + ) + return draft_config, draft, worker, use_separate + + +def load_dspark_draft_weights(shell, weights, weight_mapper=None) -> None: + """The drafter's own checkpoint, then the target's embedding and head (mirrors + ``SpecDecOneEngineForCausalLM.load_draft_weights`` for DSpark). + + The drafter borrows the target's embedding and head, and ``M3EmbedTokens.weight`` IS the target's table, so a + drafter checkpoint that ships either would load into a module the target owns or be shadowed by it: refused. + """ + borrowed = sorted( + k for k in weights.keys() if "embed_tokens." in k or k.split(".")[-2:-1] == ["lm_head"] + ) + if borrowed: + raise ValueError( + f"the DSpark drafter checkpoint ships tensors this target lends it (embedding / head): {borrowed[:4]}" + ) + shell.draft_model.load_weights(weights=weights, weight_mapper=weight_mapper) + shell.draft_model.load_weights_from_target_model(shell) + if shell.draft_model.lm_head is not shell.lm_head or ( + shell.draft_model.draft_model_full.model.embed_tokens is not shell.model.embed_tokens + ): + raise RuntimeError("the DSpark drafter does not borrow this target's embedding and head") + own = sorted( + n for n, _ in shell.named_parameters() if not n.startswith(("model.w.", "draft_model.")) + ) + logger.info(f"M3.1 target DSpark: shell parameters outside the core and the drafter: {own}") + + +def keep_drafter_kv_dtype(draft_config: ModelConfig) -> None: + """Drop any KV quantization from the drafter's config, so its context pool is allocated in its projection dtype. + + The one-engine construction copies the target's KV algo onto the drafter; on this target that is NVFP4, which + allocates a packed int8 pool the worker refuses to bind (it falls back to its private arena, and a reused + prefix's drafter K/V is lost). The cached quant modes are dropped with it, as the draft KV budget does. + """ + for quant in [draft_config.quant_config, *(draft_config.quant_config_dict or {}).values()]: + quant.kv_cache_quant_algo = None + quant.__dict__.pop("quant_mode", None) + quant.__dict__.pop("layer_quant_mode", None) + + +def draft_kv_override_applies(spec_config) -> bool: + """The target's FP8 draft-KV override is an Eagle3 need; a DSpark drafter keeps its BF16 context pool.""" + return not is_standalone_dspark(spec_config) + + +__all__ = [ + "M3DSparkWorker", + "M3EmbedTokens", + "build_dspark", + "draft_kv_override_applies", + "is_standalone_dspark", + "keep_drafter_kv_dtype", + "load_dspark_draft_weights", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/dspark_s4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/dspark_s4.py new file mode 100644 index 000000000000..5131aa51b300 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/dspark_s4.py @@ -0,0 +1,491 @@ +# 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. +"""S4: the catalog's one-launch Markov chain (``catalog/spec_decode/dspark_markov_chain``) in this target's DSpark +worker (each M3.1 target keeps its own copy of this glue: targets share only the catalog). + +``S4WorkerMixin`` goes before ``DSparkWorker`` in the target's worker. For a generation step it serves (switch not +set, 1 <= generation requests <= ``MAX_FAST_BATCH``, greedy drafting, no d2t / guided decoder / DFlash-2 path, a +vocab-sharded borrowed head): + +- the draft head computes this rank's vocab shard of the block logits (the head's local GEMM, no gather) and copies + it into a padded buffer whose tail holds -inf (the step's logits processor, swapped in for the step); +- ``_refine_block_logits`` runs the catalog op over the shards in one launch and keeps the tokens; +- ``sample_draft_tokens`` returns them. + +The tokens are bitwise the in-tree chain's on the gathered logits: the gathered logits are the concatenation of the +same per-rank GEMMs, and the op reproduces the chain's bf16 bias, bf16 add and first-maximum argmax. When the op +does not support a step after the logits were sharded, the in-tree sharded path (``vocab_slice`` + the TP-aware +argmax) runs instead, with the same tokens. + +Setup (``prebuild_s4``, from the worker's ``prebuild``, outside capture, on every rank): the padded ``markov_w2`` +shard, one padded logits buffer per (batch, block), the exchange workspace (collective) and the compiles. + +Every routing decision is Python, so it is fixed into each decode graph when the graph is captured: +``TRTLLM_M31_DRAFT_DISABLE_S4`` is read once, at import, not a switch that can be toggled at run time. The set-up +decision must be the same on every rank (the workspace is collective); it follows from the config and the +environment, and ``prebuild_s4`` checks that the ranks agree before the collective. + +Kill switches (``m31_switches``, ``TRTLLM_M31_DISABLE_=1``, read once at import; all bitwise the same tokens): +``S4B8`` -- batch-1 calls on the chain's batch-1 kernel (back to the general kernel). ``MARKOV_LL`` -- the batch-1 +kernel's LL flag-in-payload exchange (``dspark_markov_chain_kernel_b1_ll``; back to the S4B8 kernel's Lamport +exchange; no effect with S4B8 off). ``S4_FP32SUM`` -- the fp32 sum compare (back to bf16-rounded sums). ``S4_LOGITS`` +-- S4 under advanced draft sampling (back to the in-tree chain there). + +""" + +from __future__ import annotations + +import atexit +import os +from typing import Dict, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode import ( + dspark_markov_bias as markov_bias, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode import ( + dspark_markov_chain as chain, +) +from tensorrt_llm._torch.models.modeling_speculative import dspark_markov_step_bias +from tensorrt_llm.logger import logger + +from .m31_switches import m31_disabled + +MARKOV_RANK = 256 +MAX_FAST_BATCH = ( + 8 # capacity: measured: above 8 requests the chain is within 10% of the reference (16: 0.93x) +) +DISABLE_ENV = "TRTLLM_M31_DRAFT_DISABLE_S4" +S4_COUNT_LOG_EVERY = 1000 # generation steps between the per-rank "served x/y" engagement lines + +# MARKOV_LL ( GB TP8 K7 b1 exchange 2.25 us + post 0.5 us per position against a +# ~1.25 us single-hop floor): the batch-1 kernel's per-position exchange as one 16-byte {value, tag, index, tag} +# multicast store per CTA (no fence, no flag word), polled and reduced by warp 0 alone (no CTA barrier after the +# poll). Bitwise the S4B8 kernel's tokens (same entries, same total argmax order). Default ON; kill switch +# TRTLLM_M31_DISABLE_MARKOV_LL=1, read at import. +MARKOV_LL = not m31_disabled("MARKOV_LL") + +# TAILBK_SPEC (DSpark-worker sites): bookkeeping launches +# of the draft tail removed exactly. Here, site C: the chain op's three read-only accepted inputs (ones, arange, +# zeros; 3 launches per step) come from a per-shape cache filled at precompile. Site H: the target's KV-length rewind +# after the draft forward (kv_lens[nc:bs] -= rewind; clamp_(min=0): 2 launches per step) runs in the chain's launch +# (its rewind path, after the grid-dependency wait; nothing between the chain and the stock rewind reads kv_lens). +# Site S4L (S4-logits, the advanced-sampling contract path): the biased block logits' glue around the bias GEMM -- the +# prevs cat / casts, get_masked_input_and_mask, F.embedding, masked_fill, the two widening copies and the add (14 +# launches per step) -- in two launches (catalog dspark_markov_bias; the GEMM unchanged). +# Bitwise the same tokens, kv_lens and biased logits. Default ON; kill switch TRTLLM_M31_DISABLE_TAILBK_SPEC=1, read +# at import. +TAILBK_SPEC = not m31_disabled("TAILBK_SPEC") + +# Read once, at import: every rank and every step of a process sees the same setting. Any value but "" or "0" turns +# the item off, as the targets' m31_switches reads its switches. +_DISABLED = os.environ.get(DISABLE_ENV, "").strip() not in ("", "0") + + +def disabled() -> bool: + return _DISABLED + + +class _S4State: + def __init__(self, mapping, rows: int, shard: int, offset: int, w1: torch.Tensor, w2s: torch.Tensor, + block: int, ws: dict) -> None: # fmt: skip + self.mapping = mapping + self.ws = ws + self.rows = rows + self.shard = shard + self.offset = offset + self.w1 = w1 + self.w2s = w2s + self.block = block + self.bufs: Dict[Tuple[int, int], torch.Tensor] = {} + + def buffer(self, batch: int) -> torch.Tensor: + return self.bufs[(batch, self.block)] + + +# S4-logits bias form: True = all K positions' bias in one GEMM over B*K rows (S4LC: bitwise the per-step GEMMs); +# False = the in-tree per-step loop over the chain's tokens (exact by construction). +S4_LOGITS_BATCHED = True + + +def _biased_shard_logits( + shard_logits: torch.Tensor, + first_prev: torch.Tensor, + tokens: torch.Tensor, + w1: torch.Tensor, + w2_shard: torch.Tensor, + fp32: bool = True, +) -> torch.Tensor: + """The in-tree chain's biased block logits on this rank's vocab shard, given the chain's tokens: step k adds + bf16(w2 @ w1[prev_k]) with prev_0 = the anchor and prev_k = token k-1 (dspark_markov_step_bias). With ``fp32`` + (the in-tree arithmetic: LogitsProcessor hands the chain fp32 logits) the bf16 shard logits and the bias are + widened and added in fp32; without it, one bf16 add.""" + if ( + fp32 + and S4_LOGITS_BATCHED + and TAILBK_SPEC + and markov_bias.supports(shard_logits, first_prev, tokens, w1) + ): + # TAILBK_SPEC S4L: the same prevs, embedding rows, GEMM and fp32 sum, in two launches around the GEMM. + B, K, N = shard_logits.shape + bias = F.linear(markov_bias.prev_embeddings(first_prev, tokens, w1), w2_shard) + return markov_bias.widen_add(shard_logits, bias.view(B, K, N)) + if fp32: + shard_logits = shard_logits.float() + B, K, N = shard_logits.shape + prevs = torch.cat([first_prev.view(B, 1), tokens[:, : K - 1].long()], dim=1) + if S4_LOGITS_BATCHED: + bias = dspark_markov_step_bias(prevs.reshape(-1), w1, w2_shard).view(B, K, N) + return shard_logits + bias.to(shard_logits.dtype) + out = torch.empty_like(shard_logits) + for k in range(K): + out[:, k] = shard_logits[:, k] + dspark_markov_step_bias(prevs[:, k], w1, w2_shard).to( + shard_logits.dtype + ) + return out + + +class _ShardedDraftLogits(nn.Module): + """The step's draft logits processor: this rank's vocab shard of ``lm_head`` over the gathered block rows, + copied into the padded buffer; any other call goes to the drafter's own processor.""" + + def __init__(self, inner: nn.Module, state: _S4State, batch: int) -> None: + super().__init__() + self.inner = inner + self._state = [state] + self.batch = batch + + def forward(self, hidden_states, lm_head, attn_metadata, return_context_logits: bool = False): + state = self._state[0] + if hidden_states.dim() != 2 or hidden_states.shape[0] != self.batch * state.block: + return self.inner(hidden_states, lm_head, attn_metadata, return_context_logits) + logits = lm_head.apply_linear(hidden_states, None) + buf = state.buffer(self.batch) + buf[:, :, : state.shard].copy_(logits.view(self.batch, state.block, state.shard)) + return logits + + +def _head_is_sharded(lm_head: nn.Module, vocab: int, tp: int) -> bool: + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + return ( + tp > 1 + and getattr(lm_head, "tp_mode", None) == TensorParallelMode.COLUMN + and getattr(lm_head, "padding_size", 0) == 0 + and getattr(lm_head, "bias", None) is None + and lm_head.out_features * tp == vocab + and lm_head.weight.dtype == torch.bfloat16 + ) + + +class S4WorkerMixin: + """Mixed into the target's ``DSparkWorker`` subclass (before it in the bases).""" + + _s4: Optional[_S4State] = None + _s4_step: Optional[_S4State] = None + _s4_tokens: Optional[torch.Tensor] = None + _s4_step_meta: Optional[Tuple[int, int, object]] = None + _s4_draft: Optional[nn.Module] = None + _s4_served: int = 0 + _s4_gen_steps: int = 0 + _s4_ll_logged: bool = False + _s4_step_attn = ( + None # the step's attention metadata (the KV-length rewind's kv_lens), while S4 serves it + ) + _tailbk_rewound: bool = ( + False # TAILBK_SPEC H: this step's chain launch applied the KV-length rewind + ) + _tailbk_h_logged: bool = False + _tailbk_s4l_logged: bool = False + _tailbk_c_logged: bool = False + + def prebuild_s4(self, lm_head: nn.Module, draft_model) -> bool: + """Set up the one-launch chain for this worker; False (and nothing set up) when it does not apply.""" + self._s4 = None + if disabled(): + logger.info(f"M3.1 DSpark S4: off ({DISABLE_ENV} set); the in-tree Markov chain runs") + return False + if not getattr(draft_model, "has_markov_head", False): + return False + w1 = draft_model.markov_w1 + w2 = draft_model.markov_w2 + mapping = self.mapping + tp = mapping.tp_size + vocab = w1.shape[0] + if ( + w1.dtype != torch.bfloat16 + or w2.dtype != torch.bfloat16 + or w1.shape[1] != MARKOV_RANK + or tuple(w2.shape) != tuple(w1.shape) + or mapping.enable_attention_dp + or not _head_is_sharded(lm_head, vocab, tp) + ): + return False + block = int(self.max_draft_len) + shard = vocab // tp + rows = chain.padded_rows(shard, MAX_FAST_BATCH, w1.device) + if not self._s4_ranks_agree(rows != 0 and block <= chain.MAX_BLOCK): + return False + offset = mapping.tp_rank * shard + w2s = chain.pad_w2_shard(w2[offset : offset + shard].contiguous(), rows) + state = _S4State( + mapping, rows, shard, offset, w1, w2s, block, chain.workspace_args(mapping) + ) + for batch in range(1, MAX_FAST_BATCH + 1): + state.bufs[(batch, block)] = chain.pad_logits_buffer(batch, block, rows, w1.device) + # S4B8: batch-1 calls take the chain's batch-1 kernel (bitwise the same tokens; faster at one request only). + # Kill switch TRTLLM_M31_DISABLE_S4B8, read at import. + s4b8 = chain.enable_batch1_variant(not m31_disabled("S4B8")) + # MARKOV_LL: those batch-1 calls exchange through the LL flag-in-payload kernel (bitwise; see MARKOV_LL). + markov_ll = chain.enable_ll_exchange(MARKOV_LL) + # TAILBK_SPEC site C: the chain's read-only accepted inputs from a per-shape cache (bitwise; see TAILBK_SPEC). + tailbk_c = chain.enable_cached_inputs(TAILBK_SPEC) + # S4 fp32 sum: the chain compares base + bias as the in-tree chain does, the unrounded fp32 sum (its logits are + # fp32 from LogitsProcessor's .float()). Kill switch TRTLLM_M31_DISABLE_S4_FP32SUM (bf16-rounded sums), read + # at import. + self._s4_fp32 = chain.enable_fp32_sum(not m31_disabled("S4_FP32SUM")) + # S4-logits: under advanced (rejection) draft sampling the chain still computes the greedy chain's tokens; the + # biased block logits are rebuilt from them (_biased_shard_logits) and the stock sampler gathers and samples + # them as it would the in-tree chain's. Kill switch TRTLLM_M31_DISABLE_S4_LOGITS, read at import. + self._s4_logits_on = not m31_disabled("S4_LOGITS") + chain.precompile(mapping, rows, w1, w2s, offset, blocks=(block,), max_batch=MAX_FAST_BATCH) + self._s4 = state + self._s4_draft = draft_model + atexit.register(self._s4_log_count) + # The engagement line: one per rank, at server start. + logger.info( + f"M3.1 DSpark S4: one-launch Markov chain set up on rank {mapping.tp_rank}/{tp} " + f"(shard {shard} rows padded to {rows}, K={block}, batches 1..{MAX_FAST_BATCH}; S4B8 batch-1 kernel " + f"{'on' if s4b8 else 'off'}; MARKOV_LL {'on' if markov_ll and s4b8 else 'off'}; " + f"TAILBK_SPEC cached inputs {'on' if tailbk_c else 'off'}; " + f"S4-logits {'on' if self._s4_logits_on else 'off'}; fp32 sum " + f"{'on' if self._s4_fp32 else 'off'})" + ) + return True + + def _s4_ranks_agree(self, local: bool) -> bool: + """``local`` if every rank of the TP group decided the same; raises when they differ (a rank going on into + the collective workspace alone would hang its peers).""" + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _mnnvl_workspace_all_succeeded, + ) + + comm = _get_mnnvl_workspace_comm(self.mapping) + all_on = _mnnvl_workspace_all_succeeded(comm, local) + all_off = _mnnvl_workspace_all_succeeded(comm, not local) + if not (all_on or all_off): + raise RuntimeError( + "m31 DSpark S4: the TP ranks disagree on setting up the one-launch Markov chain" + ) + return all_on + + def _s4_serves(self, num_gens: int, spec_metadata, draft_model) -> bool: + return self._s4_mode(num_gens, spec_metadata, draft_model) is not None + + def _s4_mode(self, num_gens: int, spec_metadata, draft_model): + """ "tokens" (greedy drafting: the chain's tokens are the draft), "logits" (advanced draft sampling with + S4-logits on: the chain's corrected logits go to the stock sampler), or None (the in-tree chain runs).""" + if not self._s4_common(num_gens, draft_model): + return None + if not getattr(spec_metadata, "wants_advanced_draft_sampling", False): + return "tokens" + return "logits" if getattr(self, "_s4_logits_on", False) else None + + def _s4_common(self, num_gens: int, draft_model) -> bool: + return ( + self._s4 is not None + and 1 <= num_gens <= MAX_FAST_BATCH + and getattr(self, "guided_decoder", None) is None + and getattr(self, "_d2t", None) is None + and getattr(draft_model, "lilicorr", None) is None + and not getattr(draft_model, "has_candidate_selector", False) + ) + + def _host_step_hook(self, spec_metadata) -> None: + """The engagement counter (host only, no CUDA): called by the spec metadata's prepare() on every iteration, + outside any captured graph, it evaluates the same eligibility the captured step was routed by and counts + the generation steps the chain serves. One line per rank every S4_COUNT_LOG_EVERY generation steps and at + exit: "M3.1 DSpark S4: served / generation steps". The counts start at server start, so they include + warmup and padding steps: read engagement as the difference of two lines logged after the server is ready.""" + if self._s4 is None: + return + num_gens = int(getattr(spec_metadata, "num_generations", 0) or 0) + if num_gens <= 0: + return + self._s4_gen_steps += 1 + self._s4_served += int(self._s4_serves(num_gens, spec_metadata, self._s4_draft)) + if self._s4_gen_steps % S4_COUNT_LOG_EVERY == 0: + self._s4_log_count() + + def _s4_log_count(self) -> None: + if self._s4 is not None and self._s4_gen_steps: + logger.info( + f"M3.1 DSpark S4: served {self._s4_served}/{self._s4_gen_steps} generation steps on rank " + f"{self.mapping.tp_rank}" + ) + + def _forward_impl(self, input_ids, position_ids, hidden_states, logits, attn_metadata, spec_metadata, draft_model, + resource_manager=None): # fmt: skip + num_gens = attn_metadata.num_seqs - attn_metadata.num_contexts + mode = self._s4_mode(num_gens, spec_metadata, draft_model) + served = mode is not None + if not served: + return super()._forward_impl(input_ids, position_ids, hidden_states, logits, attn_metadata, + spec_metadata, draft_model, resource_manager) # fmt: skip + inner = draft_model.logits_processor + draft_model.logits_processor = _ShardedDraftLogits(inner, self._s4, num_gens) + self._s4_step = self._s4 + self._s4_step_mode = mode + self._s4_step_meta = (attn_metadata.num_contexts, attn_metadata.num_seqs, spec_metadata) + self._s4_step_attn = attn_metadata + try: + return super()._forward_impl(input_ids, position_ids, hidden_states, logits, attn_metadata, + spec_metadata, draft_model, resource_manager) # fmt: skip + finally: + draft_model.logits_processor = inner + self._s4_step = None + self._s4_step_mode = None + self._s4_tokens = None + self._s4_step_meta = None + self._s4_step_attn = None + self._tailbk_rewound = False + + def _refine_block_logits(self, draft_model, gen_logits, inputs, spec_metadata): + state = self._s4_step + if state is not None and gen_logits.dim() == 3 and gen_logits.shape[-1] == state.shard: + base = state.buffer(gen_logits.shape[0]) + if ( + chain.supports(base, state.w1, state.w2s) + and getattr(self, "_s4_step_mode", None) == "logits" + ): + # S4-logits: the chain's tokens are the in-tree greedy chain's (bitwise); the biased shard logits are + # rebuilt from them with the in-tree bias arithmetic and handed to the stock advanced sampler. + first_prev = inputs["first_prev_tokens"].long() + tokens = self._s4_chain(state, base, first_prev, spec_metadata) + self._s4_log_ll(base.shape[0]) + biased = _biased_shard_logits( + gen_logits, + first_prev, + tokens, + state.w1, + state.w2s[: state.shard], + self._s4_fp32, + ) + self._s4l_log(gen_logits, first_prev, tokens, state.w1) + return biased + if chain.supports(base, state.w1, state.w2s): + tokens = self._s4_chain( + state, base, inputs["first_prev_tokens"].long(), spec_metadata + ) + self._s4_log_ll(base.shape[0]) + self._s4_tokens = tokens + return gen_logits + return super()._refine_block_logits(draft_model, gen_logits, inputs, spec_metadata) + + def _s4_chain( + self, state: _S4State, base: torch.Tensor, first_prev: torch.Tensor, spec_metadata + ): + """The chain launch; under TAILBK_SPEC (site H) it also applies the step's deferred KV-length rewind when + the stock _apply_kv_rewind_after_draft would (a real step: not the eager warmup of a graph batch, whose + kv_lens is restored wholesale), and marks it done for that method.""" + rewind = self._tailbk_rewind_args(base.shape[0], spec_metadata) + if rewind is None: + tokens = chain.dspark_markov_chain( + base, first_prev, state.w1, state.w2s, state.offset, **state.ws + ) + self._tailbk_c_log() + return tokens + kv_lens, amount, first = rewind + tokens = chain.dspark_markov_chain_rewind(base, first_prev, state.w1, state.w2s, state.offset, kv_lens, + amount, first, **state.ws) # fmt: skip + self._tailbk_rewound = True + self._tailbk_c_log() + if not self._tailbk_h_logged: + self._tailbk_h_logged = True + logger.info( + "[M3 dispatch] TAILBK_SPEC H: KV-length rewind in the Markov chain launch " + f"on rank {self.mapping.tp_rank}" + ) + return tokens + + def _tailbk_c_log(self) -> None: + """The TAILBK_SPEC C engagement line: once per rank, after the first chain call that took its accepted inputs + from the precompile cache (host only; a captured step logs it while it is captured).""" + if self._tailbk_c_logged or not getattr(chain, "cached_inputs_taken", lambda: False)(): + return + self._tailbk_c_logged = True + logger.info( + "[M3 dispatch] TAILBK_SPEC C: chain accepted inputs from precompile cache (taken) " + f"on rank {self.mapping.tp_rank}" + ) + + def _tailbk_rewind_args(self, batch: int, spec_metadata): + """(kv_lens, rewind amounts, first row) for the chain to apply, or None (the stock rewind runs).""" + if not TAILBK_SPEC or not getattr(self, "_kv_rewind_pending", False): + return None + amount = getattr(self, "_kv_rewind_amount", None) + attn = self._s4_step_attn + if amount is None or attn is None or not hasattr(attn, "kv_lens_cuda"): + return None + if spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing(): + return None # the warmup: the stock method skips the rewind (kv_lens is restored wholesale) + nc, bs = self._kv_rewind_nc, self._kv_rewind_bs + if bs - nc != batch or amount.numel() != batch: + return None + kv_lens = attn.kv_lens_cuda + if not chain.supports_rewind(kv_lens, amount, nc): + return None + return kv_lens, amount, nc + + def _apply_kv_rewind_after_draft(self, attn_metadata, spec_metadata): + if self._tailbk_rewound: + # TAILBK_SPEC H: the chain launch applied it (see _s4_chain). + self._tailbk_rewound = False + self._kv_rewind_pending = False + return + super()._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) + + def _s4l_log(self, shard_logits, first_prev, tokens, w1) -> None: + """The TAILBK_SPEC S4L engagement line: once per rank, on the first step whose biased logits took the + two-launch glue (host only; a captured step logs it while it is captured).""" + if self._tailbk_s4l_logged or not (self._s4_fp32 and S4_LOGITS_BATCHED and TAILBK_SPEC): + return + if markov_bias.supports(shard_logits, first_prev, tokens, w1): + self._tailbk_s4l_logged = True + logger.info( + "[M3 dispatch] TAILBK_SPEC S4L: biased block logits glue in two launches " + f"on rank {self.mapping.tp_rank}" + ) + + def _s4_log_ll(self, batch: int) -> None: + """The MARKOV_LL engagement line: once per rank, on the first served step whose chain call took the LL + exchange kernel (host only; a captured step logs it while it is captured).""" + if not self._s4_ll_logged and chain.ll_exchange_engaged(batch): + self._s4_ll_logged = True + logger.info( + "[M3 dispatch] MARKOV_LL: k3_markov LL flag-in-payload exchange (batch 1) " + f"on rank {self.mapping.tp_rank}" + ) + + def sample_draft_tokens(self, logits, spec_metadata, batch_size, **kwargs): + tokens = self._s4_tokens + if tokens is not None and logits.dim() == 3 and tokens.shape == logits.shape[:2]: + self._s4_tokens = None + return tokens if tokens.dtype == torch.int32 else tokens.to(torch.int32) + return super().sample_draft_tokens(logits, spec_metadata, batch_size, **kwargs) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/m31_switches.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/m31_switches.py new file mode 100644 index 000000000000..51f966c595fd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/m31_switches.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kill switches of this MiniMax-M3.1 target's default-on items (each target keeps its own copy: targets share +only the catalog). + +Every item is on by default; ``TRTLLM_M31_DISABLE_=1`` turns item ```` off and restores the path it replaced. +The environment is read once, at import, so every rank and every step of a process sees the same setting (a step must +not change its graph mid-run). Each item lists its id and what the switch restores where it reads it. + +Values (case-insensitive, surrounding blanks ignored): ``1`` / ``true`` / ``on`` / ``yes`` set a switch; empty / ``0`` / +``false`` / ``off`` / ``no`` leave it unset; anything else raises at import, naming the variable, so a typo fails loudly +instead of silently turning an item on or off. +""" + +import os +from typing import FrozenSet + +_PREFIX = "TRTLLM_M31_DISABLE_" +_SET = frozenset(("1", "true", "on", "yes")) +_UNSET = frozenset(("", "0", "false", "off", "no")) + + +def _is_set(name: str, value: str) -> bool: + """Whether the switch variable ``name`` is set to ``value``; raises on a value of neither kind.""" + word = value.strip().lower() + if word in _SET: + return True + if word in _UNSET: + return False + raise ValueError( + f"{name}={value!r}: expected 1 / true / on / yes (set) or 0 / false / off / no / empty (unset)" + ) + + +# The item ids switched off in this process, read once at import. +_DISABLED: FrozenSet[str] = frozenset( + name[len(_PREFIX) :] + for name, value in os.environ.items() + if name.startswith(_PREFIX) and _is_set(name, value) +) + + +def m31_disabled(item_id: str) -> bool: + """Whether ``TRTLLM_M31_DISABLE_`` turned the item off for this process.""" + return item_id in _DISABLED + + +# Items that are OFF by default (rule 1b without a passing gate c, or a failed gate c): ``TRTLLM_M31_ENABLE_=1`` +# opts in; the item's ``TRTLLM_M31_DISABLE_`` still wins. +_ENABLE_PREFIX = "TRTLLM_M31_ENABLE_" + +_ENABLED: FrozenSet[str] = frozenset( + name[len(_ENABLE_PREFIX) :] + for name, value in os.environ.items() + if name.startswith(_ENABLE_PREFIX) and _is_set(name, value) +) + + +def m31_enabled(item_id: str) -> bool: + """Whether ``TRTLLM_M31_ENABLE_`` opted a default-off item in for this process.""" + return item_id in _ENABLED diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/modeling.py new file mode 100644 index 000000000000..afe87215a221 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/modeling.py @@ -0,0 +1,5647 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ModelingV2 target: MiniMax-M3.1 (MXFP8 + packed NVFP4 experts) / sm_103 / tp4 -- self-contained modeling code. + +MiniMax-M3.1's language model has MiniMax-M3's shape (60 layers, hidden 6144, 64 query heads and 4 KV heads of dim +128, 128 routed experts top 4 plus one shared expert) with a 1e7 RoPE base and MSA block-sparse attention on EVERY +layer (an index branch scores 128-token blocks and the top 16 are attended); layers 0-2 keep the dense SwiGLU-OAI MLP, +layers 3-59 the MoE. The checkpoint is an MXFP8 export whose routed experts are packed NVFP4 (compressed-tensors +names, global scales stored as multipliers; ``weights`` takes their FP32 reciprocals); the KV cache is NVFP4 with an +FP8 index-K on every layer. The layer kinds are read from the config (``sparse_attention_freq``, ``moe_layer_freq``), +so this code is the MiniMax-M3 target's with no dense-attention layer in use. + +``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 up to 1 + 7 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. + +**Kill switches** (``m31_switches``, ``TRTLLM_M31_DISABLE_=1``, read once at import): ``R1`` -- the 64-row +decode step bound (back to 48: steps of 49..64 rows run every layer's generic form). ``PROXYPB`` -- the MSA indexer +proxy's fmha_sm100 variants built before warmup (back to compiling each on its first call, which can be after ready). +``LC4_BTSMEM`` / ``LC4_SYNC3`` -- the long-context sparse decode's block-table staging / softmax sync trim (back to the +M3 op once all LC4 options are off; ``ops/attention/m3_sparse_decode_lc``). ``TOPP_DET`` -- the spec sampler's +top-p renorm as the ST-01 C++ op, exact vs FlashInfer's deterministic renorm (back to the stock FlashInfer call; +``ops/sampling/m31_topp_det``). ``SOFTMAX_MC`` -- the spec sampler's softmax with its write pass on 16 CTAs per row +(ST-01-SM), exact vs FlashInfer's softmax (back to the stock FlashInfer call; ``ops/sampling/m31_softmax_mc``). +``REJ_MC`` -- the spec sampler's chain speculative sampling on 16 CTAs per request (ST-01-REJ), exact vs the stock +op (back to ``trtllm::chain_speculative_sampling``; ``ops/sampling/m31_chain_spec_mc``). +``HOSTTAIL`` -- all six between-replay +host items (each also alone: ``HOST_D2H_SIDE``, ``SMALL_H2D``, ``BATCH_CHANGED_ONLY``, ``SKIP_ACCEPTED_UPLOAD``, +``PHILOX_PACK``, ``KV_LENS_ONCE``; back to the stock copies, see ``HOSTTAIL_*``). ``PACK_IN_GRAPH`` / +``PACK_EARLY_COPY`` / ``H2D_CHANGED_SPAN`` (VR2GB-01/-03) -- the step-input pack issued by the decode graph's prologue, +its window uploaded ahead of the replay, changed-span page-table uploads (back to the eager pack and whole-view +uploads; see ``PACK_IN_GRAPH``). ``BOUNDARY_PACK`` (BP-1, needs ``DSFOLD``) -- under DSpark, a batch-1 decode graph's +tail does the spec sampler's four slot stores and the next step's step gather (back to both between the graph replays; +see ``BOUNDARY_PACK``). + +**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 types import MappingProxyType +from typing import ( + Any, + Callable, + Dict, + FrozenSet, + Hashable, + Iterable, + List, + Mapping, + Optional, + Tuple, + Union, +) + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + ModelingV2Mode, + 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.fakes import ( + m31_fake_ops as _m31_fake_ops, +) +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, + require_msa_module, +) +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.custom_ops.torch_custom_ops import MXFP8QuantizeRunner +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.pyexecutor.sampler.ops.flashinfer import ( + set_softmax_override, + set_top_p_renorm_override, +) +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, + set_chain_speculative_sampling_override, + should_use_separate_draft_kv_cache, +) +from tensorrt_llm._torch.utils import ActType_TrtllmGen +from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm.functional import AllReduceStrategy, PositionEmbeddingType +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantAlgo + +from . import dspark as _dspark +from . import weights as _weights +from .m31_switches import m31_disabled, m31_enabled +from .ops import gqamix_sparse_decode as _gqamix_kernel +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_bp_tail, m3_mc_allreduce, m3_prequant, m3_spec_tail +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_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.activation.m3_swq import m3_swq +from .ops.attention.m3_attn_c1 import m3_attn_c1 +from .ops.attention.m3_attn_mega import m3_attn_mega +from .ops.attention.m3_dense_decode import m3_dense_decode +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, m3_mxfp8_gemm_cute_dsl +from .ops.gemm.m3_mxfp8_linear import m3_mxfp8_linear +from .ops.gemm.m3_mxfp8_pinned_quantize import m3_mxfp8_pinned_quantize +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.sampling import m31_chain_spec_mc as _chain_spec_mc +from .ops.sampling import m31_softmax_mc as _softmax_mc +from .ops.sampling import m31_topp_det as _topp_det +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 +# -------------------------------------------------------------------------------------------------------------------- + +# Kill switch ``TRTLLM_M31_DISABLE_PROXYPB``: compile the MSA indexer proxy's fmha_sm100 variants before warmup +# (``_prebuild_msa_proxy``). Compile only; off, each compiles on its first call. +MSA_PROXY_PREBUILD = not m31_disabled("PROXYPB") +# CTXNORMQ (CTX lane): a prefill layer boundary of at least ``CTXNORMQ_MIN_ROWS`` rows runs the plain all-reduce, then +# one kernel for the residual add + RMSNorm that also writes the next MXFP8 projection's activation (and, at a MoE +# boundary, the routed experts' NVFP4 input), instead of the fused all-reduce + norm followed by the consumers' own +# quantize passes. The quantized outputs equal mxfp8_quantize / fp4_quantize of the norm bit for bit; the norm itself +# is rule 1b against the fused all-reduce's (reduction order). Kill switch: TRTLLM_M31_DISABLE_CTXNORMQ=1. +CTXNORMQ = not m31_disabled("CTXNORMQ") +CTXNORMQ_MIN_ROWS = ( + 1024 # capacity: NORMQ microbenchmark: the fused pass wins at 1160 and 16384 rows +) +# AR64AUTO: a generation step past the multicast all-reduce's rows (K7 c8: 64 +# verify rows) runs the built-in fused all-reduce + residual + RMSNorm with the fixed strategy MNNVL instead of the +# configured one; the serving configs pin NCCL, 29.1 vs 11.6 us per boundary at 64 rows on GB300 TP8 (with the next +# projection's MXFP8 quantize). Pinned, not AUTO: AUTO's choice is timing-autotuned per process; MNNVL is bitwise +# identical across processes at 49..64 rows. Rule 1b against the configured strategy (another reduction order +# across ranks); every other boundary keeps the configured instance. Kill switch: TRTLLM_M31_DISABLE_AR64AUTO=1. +AR64AUTO = not m31_disabled("AR64AUTO") +# TSEL (CTX lane): a PrefillTarget sparse layer's top-16 block select runs the #20629 prefill tile select (one CTA per +# KV head and 8 queries, each 32-byte score sector read once) on sm_103, registered on the MSA indexer's selector hook. +# The table equals minimax_m3_select_blocks' bit for bit; steps outside its rows / blocks range keep the built-in +# select. Kill switch: TRTLLM_M31_DISABLE_TSEL=1. +TSEL = not m31_disabled("TSEL") +TSEL_MIN_ROWS = 256 # microbenchmark: x1.4 (128 blocks) to x4.4 (1024), x1.6 at 1160 rows +TSEL_MAX_BLOCKS = 5856 # the op's shared-memory bound (kMinimaxM3SelectPrefillTileMaxBlocks) +# Gate a in-model (correctness runs only): TRTLLM_M31_CHECK_TSEL=1 also runs minimax_m3_select_blocks on every +# tile-selected step and raises on any table difference. +TSEL_CHECK = os.environ.get("TRTLLM_M31_CHECK_TSEL") == "1" +# GQAMIX (MSA): a mixed ctx+gen step's sparse generation rows (MsaDecodeFmha) run the Triton block-sparse decode +# through ops/gqamix_sparse_decode, a verbatim copy of the in-tree kernel module whose two decode kernels do not +# specialize on the step's row count or on the [num_ctx:] slices' alignment, registered on MsaDecodeFmha's +# sparse_decode_hook. Those specializations compiled new variants after the server was ready (microbenchmark: +# bitwise in every cell, 0 new variants vs 1, at most +0.07 us). Kill switch: TRTLLM_M31_DISABLE_GQAMIX=1. +GQAMIX = not m31_disabled("GQAMIX") +# DSFOLD: under the standalone DSpark worker, the engine's previous-batch input gathers (input ids, draft tokens, +# position / KV-length offsets) run as one m3_spec_tail.step_gather launch instead of ~a dozen eager index and copy +# launches at the step boundary. Integer copies of the same values into the same buffers; the sampler keeps its +# in-tree slot stores. Kill switch: TRTLLM_M31_DISABLE_DSFOLD=1. +DSFOLD = not m31_disabled("DSFOLD") +# BOUNDARY_PACK BP-1 (with DSFOLD): a batch-1 decode graph ends with the spec sampler's four +# slot stores (the same values at the same slot of the same store) and the engine's next-step gathers, in one launch +# (ops/m3_bp_tail) gated by a device (slot, live) word the fold writes before the replay, so neither runs eagerly +# between two graph replays: the sampler skips its stores (``sampler_rows_stored``) and the engine its gathers +# (``begin_step``) on the steps the graph did them, and a replay with no identified row (a dummy-only batch) writes +# nothing. Only batch-1 graphs (never padded: the row is the step's one request); every other step keeps today's path. +# Kill switch: TRTLLM_M31_DISABLE_BOUNDARY_PACK=1. +BOUNDARY_PACK = DSFOLD and not m31_disabled("BOUNDARY_PACK") +BOUNDARY_PACK_MAX_ROWS = 1 +# The key of the forward outputs that tells the sampler the step's graph stored its rows (BOUNDARY_PACK). +_BOUNDARY_PACK_MARK = "m31_boundary_pack" +# Prebuild/workspace capacity retains the historical R1 shapes. Fused-layer dispatch +# is limited separately below; each operation also retains its own supported row range. +# Wider pure-generation steps use the target generic layers. +DECODE_MAX_TOKENS = ( + 48 if m31_disabled("R1") else 64 +) # prebuild/workspace capacity; dispatch is bounded separately +# Fused-layer dispatch scope; keep the larger prebuild/workspace capacity above. +FUSED_DECODE_MAX_TOKENS = 32 + + +@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); +# rows of the verify plan past the all-reduce's rows (m3_mc_allreduce.max_rows()) take the built-in all-reduce. +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 dense layers' decode attention (trtllm-gen, msa_decode.dense_decode_max_seq_len): the max_seq_len it is handed +# per request count. It only sizes the multi-CTA KV split (every CTA still covers its share of the actual length; the +# result changes only by the split-merge order), and the model maximum over-splits AgentX-length caches at one +# request: GB300 TP4, 16 heads, FP8 KV, 64K context: 15.1 us at 16K vs 28.6 us at 1M. Two requests are neutral (the +# grid is unchanged) and keep the step's bound. +DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({1: 16384}) + +# 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 = min(DECODE_MAX_TOKENS, m3_shard_accept_kernel.SHARD_ROWS) // ( + SHARD_ACCEPT_DRAFT_LEN + 1 +) # capacity: fixed: the kernel's verify 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 its widest MMA +# (mx_gemv.MAX_TOKENS) but the unsplit K=768 shared down projection at 17..32 rows, where FlashInfer's GEMM is faster. +_ALL_ROWS = frozenset(range(1, min(DECODE_MAX_TOKENS, mx_gemv.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, +} +# SHDOWN pin: the shared down projection on the rows the GEMV does not serve runs FlashInfer's CuTe DSL MXFP8 GEMM, +# fixed. The AutoTuner tunes this shape live per start and can pick CUTLASS (cold decode weights, 57 per step, GB300: +# 6.5 vs 3.1 us per call; microbenchmark); the pin removes that per-start choice. TRTLLM_M31_DISABLE_SHDOWN_PIN=1 +# restores the tuned backend. +SHDOWN_PIN_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = ( + {} if m31_disabled("SHDOWN_PIN") else {SHARED_DOWN: frozenset(range(17, 33))} +) +_SHDOWN_LOGGED: set = set() +# SWQ: at these rows the shared expert's SwiGLU writes the down projection's MXFP8 activation directly (ops/activation +# /m3_swq.py, bitwise to silu_and_mul + mxfp8_quantize), which then takes the prequant branch of ``_linear`` (the SHDOWN +# pin at 17..32). GB300 TP4, 32 rows: 5.41 vs 6.38 us per call with its GEMM; at <= 16 rows the +# bf16 GEMV (in-kernel quantize) stays faster or is a different quantize. TRTLLM_M31_DISABLE_SWQ=1 restores the bf16 +# SwiGLU + quantize. +SWQ_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = ( + {} if m31_disabled("SWQ") else {SHARED_DOWN: frozenset(range(17, 33))} +) +_SWQ_LOGGED: set = set() +# TOPP_DET: the spec sampler's top-p renorm (FlashInfer top_p_renorm_probs in the no-top-k rejection +# path, the draft's and the verify's call each step) runs as one persistent C++ launch +# (ops/sampling/m31_topp_det, kernels/minimaxM31/toppRenormDetSm103.cu) through the sampler's top-p renorm override, +# instead of FlashInfer's five launches with a one-CTA-per-row apply. Bitwise FlashInfer's deterministic renorm on +# rows with top_p < 1 (microbenchmark: probs and sampled token, 128 cells, eager and graph); rows with top_p >= 1 +# keep every token (FlashInfer drops tail buckets there in ~43% of rows). GB300 per call at +# 7 / 8 peaked rows: 54.3 / 54.8 vs 104.2 / 105.1 us (v1.2 MB; -100 us/step at K7 c1), flat on high-entropy +# rows (near-uniform 59.0 vs 248.5; v1.2 bitwise v1 on every MB row). Taken up to +# ``_topp_det.MAX_ROWS`` rows; larger calls, inputs outside the contract and builds without the op keep the stock +# call. Kill switch: TRTLLM_M31_DISABLE_TOPP_DET=1 (stock FlashInfer top_p_renorm_probs). +TOPP_DET = not m31_disabled("TOPP_DET") +_TOPP_DET_STATE: dict = {} + + +def _install_topp_det() -> None: + """Once per process: install ``_topp_det.top_p_renorm_det`` as the sampler's top-p renorm override when TOPP_DET + is on and the op is built, and log the engagement (or why it is not taken) once per rank.""" + if _TOPP_DET_STATE: + return + if not TOPP_DET: + state = "off (TRTLLM_M31_DISABLE_TOPP_DET=1); stock FlashInfer top_p_renorm_probs" + elif _topp_det.topp_det_op() is None: + state = "not taken: the op is not built; stock FlashInfer top_p_renorm_probs" + else: + set_top_p_renorm_override(_topp_det.top_p_renorm_det) + ctas = _topp_det.ctas_per_row(_topp_det.MAX_ROWS, _topp_det.sm_count(torch.device("cuda"))) + state = f"engaged (ctas_per_row={ctas}, rows <= {_topp_det.MAX_ROWS})" + _TOPP_DET_STATE["installed"] = state.startswith("engaged") + logger.info(f"[M3 dispatch] ST-01 TOPP_DET det top-p renorm: {state}") + + +def _prebuild_topp_det() -> None: + """The override's persistent workspace on this rank's device, allocated before warmup and any graph capture.""" + if _TOPP_DET_STATE.get("installed"): + _topp_det.prebuild() + + +# SOFTMAX_MC: the spec sampler's softmax (FlashInfer softmax(logits, temperature) in softmax_op and +# compute_probs_from_logits_op, the draft's and the verify's call each step) keeps FlashInfer's Map kernel and runs +# its Reduce kernel's write pass on ``_softmax_mc.WRITE_CTAS`` CTAs per row (ops/sampling/m31_softmax_mc, +# kernels/minimaxM31/softmaxMcSm103.cu) through the sampler's softmax override, instead of one 1024-thread CTA per +# row. Bitwise FlashInfer's softmax on its Path A shapes (rows <= 128, vocab >= 24576; microbenchmark at write +# CTAs 4 / 8 / 16). GB300 per call at 7-8 rows: 9.2 vs 21.8 us. Taken up to ``_softmax_mc.MAX_ROWS`` rows; larger +# calls, inputs outside the contract and builds without the op keep the stock call. Kill switch: +# TRTLLM_M31_DISABLE_SOFTMAX_MC=1 (stock FlashInfer softmax). +SOFTMAX_MC = not m31_disabled("SOFTMAX_MC") +_SOFTMAX_MC_STATE: dict = {} + + +def _install_softmax_mc() -> None: + """Once per process: install ``_softmax_mc.softmax_mc`` as the sampler's softmax override when SOFTMAX_MC is on + and the op is built, and log the engagement (or why it is not taken) once per rank.""" + if _SOFTMAX_MC_STATE: + return + if not SOFTMAX_MC: + state = "off (TRTLLM_M31_DISABLE_SOFTMAX_MC=1); stock FlashInfer softmax" + elif _softmax_mc.softmax_mc_op() is None: + state = "not taken: the op is not built; stock FlashInfer softmax" + else: + set_softmax_override(_softmax_mc.softmax_mc) + state = f"engaged (write_ctas={_softmax_mc.WRITE_CTAS}, rows <= {_softmax_mc.MAX_ROWS})" + _SOFTMAX_MC_STATE["installed"] = state.startswith("engaged") + logger.info(f"[M3 dispatch] ST-01-SM softmax multi-CTA write: {state}") + + +def _prebuild_softmax_mc() -> None: + """The override's persistent partials workspace on this rank's device, allocated before warmup and any graph + capture.""" + if _SOFTMAX_MC_STATE.get("installed"): + _softmax_mc.prebuild() + + +# REJ_MC: the spec sampler's rejection step (trtllm::chain_speculative_sampling in +# speculative/interface.rejection_sampling_one_model, once per verify step) runs on ctas_per_row CTAs per request +# (ops/sampling/m31_chain_spec_mc, kernels/minimaxM31/chainSpecSamplingMcSm103.cu) through the rejection sampler's +# override, instead of one 1024-thread CTA per request. Bitwise the stock op's tokens and counts (microbenchmark: +# K7, batch 1 / 2 / 4 / 8 x 5 cases x 64 seeds at ctas_per_row 4 / 8 / 16, eager and graph). GB300 per call at batch 1: +# 31.2 vs 41.2 us. Taken up to ``_chain_spec_mc.MAX_BATCH`` requests; larger calls, inputs outside the contract and +# builds without the op keep the stock op. Kill switch: TRTLLM_M31_DISABLE_REJ_MC=1 (stock +# trtllm::chain_speculative_sampling). +REJ_MC = not m31_disabled("REJ_MC") +_REJ_MC_STATE: dict = {} + + +def _install_rej_mc() -> None: + """Once per process: install ``_chain_spec_mc.chain_spec_mc`` as the rejection sampler's override when REJ_MC is + on and the op is built, and log the engagement (or why it is not taken) once per rank.""" + if _REJ_MC_STATE: + return + if not REJ_MC: + state = "off (TRTLLM_M31_DISABLE_REJ_MC=1); stock trtllm::chain_speculative_sampling" + elif _chain_spec_mc.chain_spec_mc_op() is None: + state = "not taken: the op is not built; stock trtllm::chain_speculative_sampling" + else: + set_chain_speculative_sampling_override(_chain_spec_mc.chain_spec_mc) + ctas = _chain_spec_mc.ctas_per_row( + _chain_spec_mc.MAX_BATCH, _chain_spec_mc.sm_count(torch.device("cuda")) + ) + state = f"engaged (ctas_per_row={ctas}, batch <= {_chain_spec_mc.MAX_BATCH})" + _REJ_MC_STATE["installed"] = state.startswith("engaged") + logger.info(f"[M3 dispatch] ST-01-REJ chain speculative sampling multi-CTA: {state}") + + +def _prebuild_rej_mc() -> None: + """The override's persistent workspace on this rank's device, allocated before warmup and any graph capture.""" + if _REJ_MC_STATE.get("installed"): + _chain_spec_mc.prebuild() + + +# QPIN: the decode-graph MXFP8 quantize ahead of the tuned GEMM runs a fixed tactic instead of the AutoTuner's pick, +# which PERSIST freezes from one learn run's profiling (stack4: (16, 6144) persisted TRT-LLM at a noisy 10.5 us; +# measured 2.85 vs CuTe DSL 1.78 us, microbenchmark). Both tactics are bitwise equal; CuTe DSL is faster at every +# measured shape but (32, 768), where TRT-LLM is (1.88 vs 2.03). The GEMM keeps its tuned tactic (same cache key). +# TRTLLM_M31_DISABLE_QPIN=1 restores the tuned quantize + GEMM. +QPIN = not m31_disabled("QPIN") +QPIN_TACTIC_EXCEPTIONS: Dict[Tuple[int, int], int] = {(32, 768): MXFP8QuantizeRunner.TRTLLM} +_QPIN_LOGGED: set = set() +# HOSTTAIL (LongAttn; VR lineage c972d6f162): six exact host-side items between two decode graph replays, each a +# RuntimeHooks field below (``_runtime_hooks``). Same bytes reach the device; only where / whether a copy is issued +# changes. ``TRTLLM_M31_DISABLE_HOSTTAIL=1`` turns all six off; each also has its own kill switch: +# * ``HOST_D2H_SIDE`` (KIMI-HO-04, sampler_d2h_side_stream): the spec sampler's store D2H copies go on its D2H side +# stream, so they no longer sit between graph N's stores and graph N+1 on the main stream. +# * ``SMALL_H2D`` (KIMI-HO-12a, spec_small_h2d): the sampler's slot ids stay in a device buffer re-uploaded only on +# change, and the accepted-draft counts ride the step-input pack. +# * ``BATCH_CHANGED_ONLY`` (KIMI-HO-03 / ST-18(d), spec_batch_map_changed_only): the DFlash/DSpark drafter's +# batch_indices and batch-to-slot map are uploaded only when they change (constant in steady decode). +# * ``SKIP_ACCEPTED_UPLOAD`` (KIMI-HO-09, skip_accepted_draft_upload): the per-step accepted-draft-count upload is +# skipped; nothing on this target reads spec_metadata.num_accepted_draft_tokens (it is None). +# * ``PHILOX_PACK`` (KIMI-ST-18(e), pack_spec_rng_uploads): a sampled step's four Philox seed / offset uploads ride the +# step-input pack instead of four copies between the graphs. +# * ``KV_LENS_ONCE`` (VR2GB-04, kv_lens_update_once): a speculative overlap step re-derives the MSA KV-length metadata +# once, after the KV-length correction; the first pass at the head of the decode graph is rewritten in full. +_HOSTTAIL_OFF = m31_disabled("HOSTTAIL") +HOSTTAIL_D2H_SIDE = not (_HOSTTAIL_OFF or m31_disabled("HOST_D2H_SIDE")) +HOSTTAIL_SMALL_H2D = not (_HOSTTAIL_OFF or m31_disabled("SMALL_H2D")) +HOSTTAIL_BATCH_CHANGED_ONLY = not (_HOSTTAIL_OFF or m31_disabled("BATCH_CHANGED_ONLY")) +HOSTTAIL_SKIP_ACCEPTED_UPLOAD = not (_HOSTTAIL_OFF or m31_disabled("SKIP_ACCEPTED_UPLOAD")) +HOSTTAIL_PHILOX_PACK = not (_HOSTTAIL_OFF or m31_disabled("PHILOX_PACK")) +HOSTTAIL_KV_LENS_ONCE = not (_HOSTTAIL_OFF or m31_disabled("KV_LENS_ONCE")) +# PACK_PSCATTER (VR2GB-02, pack_parallel_scatter): the step-input pack issues its non-overlapping copies with one +# program per copy instead of the serial one-CTA scatter (same bytes); TRTLLM_M31_DISABLE_PACK_PSCATTER=1 restores it. +PACK_PSCATTER = not m31_disabled("PACK_PSCATTER") +# VR2GB-01 / VR2GB-03 (between-replay step-input hooks; each switch also turns off the items that need it): +# * ``PACK_IN_GRAPH`` (pack_in_graph; needs PACK_PSCATTER): the decode graphs are captured with the step-input pack's +# prologue, so the pack's upload and scatter run inside the replay and no operation is left between the graphs. +# TRTLLM_M31_DISABLE_PACK_IN_GRAPH=1 (or a disabled PACK_PSCATTER) restores the eager upload and scatter. +# * ``PACK_EARLY_COPY`` (pack_early_copy; needs PACK_IN_GRAPH): the prologue's window is uploaded on a side copy stream +# while the previous graph still runs, so the replay does not start with a host copy. +# TRTLLM_M31_DISABLE_PACK_EARLY_COPY=1 restores the prologue's own copy of the pinned window. +# * ``H2D_CHANGED_SPAN`` (h2d_changed_span): a changed packed page-table upload stages only each row's changed column +# span, so a page-crossing step keeps fitting the in-graph window. TRTLLM_M31_DISABLE_H2D_CHANGED_SPAN=1 restores +# the whole-view upload. +PACK_IN_GRAPH = PACK_PSCATTER and not m31_disabled("PACK_IN_GRAPH") +PACK_EARLY_COPY = PACK_IN_GRAPH and not m31_disabled("PACK_EARLY_COPY") +H2D_CHANGED_SPAN = not m31_disabled("H2D_CHANGED_SPAN") +# 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, min(DECODE_MAX_TOKENS, m3_mc_allreduce.max_rows())), + 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, min(DECODE_MAX_TOKENS, m3_mc_allreduce.max_rows())), + 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"]),), +) + +# Requests of a decode step whose dense-layer attention runs at a capped KV split. +DENSE_DECODE_SPLIT = RowRange( + "m3_dense_decode", + frozenset(DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH), + m3_dense_decode, + lambda spec: (len(spec["kv_lens"]),) if spec["capped"] else (), +) + +# 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, + DENSE_DECODE_SPLIT, + 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 <= min(FUSED_DECODE_MAX_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_rows()). +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 ValueError( + "this MiniMax-M3 target requires sparse_attention_config.algorithm='minimax_m3' with implementation='msa'" + ) + if getattr(sparse, "indexer_kv_dtype", None) != "fp8": + raise ValueError( + "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 ValueError("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: + prequant = _CTXNORMQ_MX.take(x) if CTXNORMQ and not mode.tuned_quantize_gemm else None + if prequant is not None: + return _ctxnormq_mxfp8_gemm(core, prequant, w, s, mode), False + 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") + if rows in SHDOWN_PIN_ROWS.get((n, k), ()): + return _shdown_pinned(act, act_scale, w, s, "mxfp8"), False + 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", + ) + if x.dtype == torch.bfloat16 and rows in SHDOWN_PIN_ROWS.get((n, k), ()): + act, act_scale = torch.ops.trtllm.mxfp8_quantize(x, True) + return _shdown_pinned(act.view(rows, k), act_scale, w, s, "bf16"), False + if QPIN and mode.tuned_quantize_gemm and x.dtype == torch.bfloat16: + tactic = QPIN_TACTIC_EXCEPTIONS.get((rows, k), MXFP8QuantizeRunner.CUTE_DSL) + act, act_scale = m3_mxfp8_pinned_quantize(x, tactic) + if (rows, n, k) not in _QPIN_LOGGED: + _QPIN_LOGGED.add((rows, n, k)) + logger.info( + f"[M3 dispatch] QPIN: mxfp8 quantize tactic {tactic} rows={rows} [{n}, {k}]" + ) + return m3_mxfp8_gemm( + act.view(rows, k), act_scale, w, s, torch.bfloat16, mode.autotune + ), False + return _mxfp8_linear(core, x, w, s, mode), False + + +def _shdown_pinned(act, act_scale, w, s, source: str) -> torch.Tensor: + """The SHDOWN pin (SHDOWN_PIN_ROWS): FlashInfer's CuTe DSL MXFP8 GEMM, whatever the AutoTuner (or a persisted + tactic cache) holds for the shape. The shared expert's down projection reaches it with a bf16 activation (the + SwiGLU output), quantized here with the native MXFP8 quantize.""" + rows, (n, k) = act.shape[0], w.shape + if (rows, source) not in _SHDOWN_LOGGED: + _SHDOWN_LOGGED.add((rows, source)) + logger.info( + f"[M3 dispatch] SHDOWN pin: cute_dsl rows={rows} [{n}, {k}] ({source} activation)" + ) + return m3_mxfp8_gemm_cute_dsl(act, act_scale, w, s, torch.bfloat16) + + +# -------------------------------------------------------------------------------------------------------------------- +# 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_rows() 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, + ) + + +# The CTXNORMQ hand-offs: the boundary's MXFP8 activation for the first prefill MXFP8 projection reading the normed +# rows, and its NVFP4 form for the routed experts. Separate from ``m3_prequant.SLOT`` so only a CTXNORMQ boundary +# feeds the prefill path. +_CTXNORMQ_MX = m3_prequant.PrequantSlot() +_CTXNORMQ_FP4 = m3_prequant.PrequantSlot() +_CTXNORMQ_STATE = {} + + +def _ctxnormq_op(): + """``torch.ops.trtllm.minimax_m31_ctx_residual_rmsnorm_mxfp8``, or None. A build without the op in its library + loads it from ``tensorrt_llm/libs/libm31_ctxnormq.so`` when that file exists (an overlay); otherwise the boundary + keeps the fused all-reduce. Decided once per process.""" + if "op" not in _CTXNORMQ_STATE: + op = getattr(torch.ops.trtllm, "minimax_m31_ctx_residual_rmsnorm_mxfp8", None) + if op is None: + import tensorrt_llm + + lib = os.path.join(os.path.dirname(tensorrt_llm.__file__), "libs", "libm31_ctxnormq.so") + if os.path.exists(lib): + torch.ops.load_library(lib) + _m31_fake_ops.register() + op = getattr(torch.ops.trtllm, "minimax_m31_ctx_residual_rmsnorm_mxfp8", None) + _CTXNORMQ_STATE["op"] = op + logger.info( + f"[M3 dispatch] CTXNORMQ prefill boundary norm + quantize: " + f"{'engaged' if op is not None else 'not taken: the op is not built'} (rows >= {CTXNORMQ_MIN_ROWS})" + ) + return _CTXNORMQ_STATE["op"] + + +def _ctxnormq_allreduce_norm(core, x, residual, norm_w, fp4_scale): + """A prefill boundary as the plain all-reduce + the fused residual add, RMSNorm and quantize, or None (switched + off, fewer than ``CTXNORMQ_MIN_ROWS`` rows, op not built). Returns (normed, updated) like the fused all-reduce; + the MXFP8 activation (and with ``fp4_scale`` the NVFP4 one) wait in the CTXNORMQ slots for their consumers.""" + if not CTXNORMQ or x.dim() != 2 or x.shape[0] < CTXNORMQ_MIN_ROWS: + return None + op = _ctxnormq_op() + if op is None: + return None + reduced = m3_allreduce(x, core.allreduce, None) + _record("m31_ctxnormq", x.shape[0], fp4_scale is not None) + normed, xq, xsf, xfp4, xfp4_sf = op( + reduced, residual.contiguous(), norm_w, core.rms_eps, True, fp4_scale, get_env_enable_pdl() + ) + _CTXNORMQ_MX.publish(normed, xq, xsf) + if fp4_scale is not None: + _CTXNORMQ_FP4.publish(normed, xfp4, xfp4_sf) + else: + _CTXNORMQ_FP4.clear() + return normed, reduced + + +def _ctxnormq_mxfp8_gemm(core, prequant, w, s, mode): + """The GEMM ``m3_mxfp8_linear`` runs after its own quantize, on the CTXNORMQ activation (bit-identical operands).""" + act, act_scale = prequant + _record("m31_ctxnormq_gemm", tuple(w.shape), act.shape[0], mode.flashinfer, mode.native_tuning) + if mode.flashinfer: + return torch.ops.trtllm.flashinfer_mm_mxfp8(act, act_scale, w, s, torch.bfloat16) + gemm = ( + torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned + if mode.native_tuning + else torch.ops.trtllm.mxfp8_mxfp8_gemm + ) + return gemm(act, act_scale, w, s, core.mx_alpha, torch.bfloat16) + + +_TSEL_STATE = {} + + +def _tsel_op(): + """``torch.ops.trtllm.minimax_m31_select_prefill_tile`` on sm_103, or None. A build without the op in its library + loads it from ``tensorrt_llm/libs/libm31_tsel.so`` when that file exists (an overlay). Decided once per process.""" + if "op" not in _TSEL_STATE: + op = None + if torch.cuda.get_device_capability() == (10, 3): + op = getattr(torch.ops.trtllm, "minimax_m31_select_prefill_tile", None) + if op is None: + import tensorrt_llm + + lib = os.path.join(os.path.dirname(tensorrt_llm.__file__), "libs", "libm31_tsel.so") + if os.path.exists(lib): + torch.ops.load_library(lib) + _m31_fake_ops.register() + op = getattr(torch.ops.trtllm, "minimax_m31_select_prefill_tile", None) + _TSEL_STATE["op"] = op + logger.info( + f"[M3 dispatch] TSEL prefill tile select: " + f"{'engaged' if op is not None else 'not taken: not sm_103 or the op is not built'} " + f"(rows >= {TSEL_MIN_ROWS}, blocks <= {TSEL_MAX_BLOCKS})" + ) + return _TSEL_STATE["op"] + + +def _tsel_select(max_score_kv, n_valid_blocks, init_blocks, local_blocks): + """The MSA indexer's selector hook: the head-major [total_q, num_kv_heads, 16] table of + ``select_blocks_from_maxscore``, or None (outside the tile select's rows / blocks range, or no op) to keep it.""" + if ( + torch.cuda.is_current_stream_capturing() # the host-side rows / blocks gate is not graph-stable + or max_score_kv.dtype != torch.float32 + or max_score_kv.dim() != 3 + or max_score_kv.stride(2) != 1 + or max_score_kv.shape[2] < TSEL_MIN_ROWS + or max_score_kv.shape[1] > TSEL_MAX_BLOCKS + ): + return None + op = _tsel_op() + if op is None: + return None + num_kv_heads, num_blocks, total_q = max_score_kv.shape + _record("m31_tsel", num_blocks, total_q >= 1024) + out = torch.empty( + (num_kv_heads, total_q, MSA_REQUIRED_TOPK), dtype=torch.int32, device=max_score_kv.device + ).permute(1, 0, 2) + nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous() + op(max_score_kv, nvb, out, init_blocks, local_blocks) + if TSEL_CHECK: + ref = torch.ops.trtllm.minimax_m3_select_blocks( + max_score_kv, nvb, MSA_REQUIRED_TOPK, init_blocks, local_blocks, True + ) + if not torch.equal(out, ref): + raise AssertionError( + f"TSEL check: tile table mismatch against minimax_m3_select_blocks " + f"({int((out != ref).any(dim=-1).sum())} rows; {num_blocks} blocks, {total_q} rows)" + ) + _TSEL_STATE["checked"] = _TSEL_STATE.get("checked", 0) + 1 + if _TSEL_STATE["checked"] in (1, 100, 1000): + logger.info( + f"[M3 dispatch] TSEL check: {_TSEL_STATE['checked']} tile-selected steps equal" + ) + return out + + +def _tsel_register(backend) -> None: + """Register ``_tsel_select`` on the layer backend's MSA indexer (an indexer without the hook is left alone).""" + indexer = getattr(backend, "indexer", None) + if indexer is not None and getattr(indexer, "selector_hook", False) is None: + indexer.selector_hook = _tsel_select + + +_GQAMIX_STATE: dict = {} + + +def _gqamix_decode(*args, **kwargs) -> None: + """``minimax_m3_sparse_attn_decode`` for MsaDecodeFmha's sparse generation rows, through the target-local copy.""" + if not _GQAMIX_STATE.get("logged"): + _GQAMIX_STATE["logged"] = True + logger.info( + "[M3 dispatch] GQAMIX: mixed-step sparse decode through ops/gqamix_sparse_decode" + ) + _gqamix_kernel.minimax_m3_sparse_attn_decode(*args, **kwargs) + + +def _gqamix_register(backend) -> None: + """Register ``_gqamix_decode`` on the layer backend's MsaDecodeFmha (a library without the hook is left alone).""" + manager = getattr(backend, "_fmha_manager", None) + for fmha in getattr(manager, "fmha_libs", ()): + if getattr(fmha, "sparse_decode_hook", False) is None: + fmha.sparse_decode_hook = _gqamix_decode + + +def _native_allreduce_norm(core, x, residual, norm_w): + """The built-in fused all-reduce + residual + RMSNorm: with AR64AUTO, the strategy-MNNVL instance for the rows past + the multicast all-reduce's, up to the decode bound; the configured instance otherwise.""" + params = AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=norm_w, + eps=core.rms_eps, + trigger_completion_at_end=False, + ) + wide = getattr(core, "allreduce_wide", None) + rows = x.shape[0] + if wide is not None and m3_mc_allreduce.max_rows() < rows <= DECODE_MAX_TOKENS: + # The MNNVL module itself (not AllReduce(MNNVL)), so a declined shape returns None here instead of falling + # through to AllReduce's AUTO / autotuned path: the configured strategy runs then (fixed order either way). + out = m3_allreduce(x.contiguous(), wide, params) + if out is not None: + _record("ar64auto", rows) + _log_once( + logger.info, + "ar64auto", + "taken", + f"[M3.1] AR64AUTO taken: MNNVL all-reduce at {rows} rows", + ) + return out + _log_once( + logger.warning, + "ar64auto", + f"declined {rows}", + f"[M3.1] AR64AUTO declined: MNNVL refused {rows} rows; the configured all-reduce runs", + ) + _record("m3_allreduce", rows) + return m3_allreduce(x, core.allreduce, params) + + +# -------------------------------------------------------------------------------------------------------------------- +# 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] = {} + + +# AR-DET: AllReduce AUTO's per-process timing autotuner made +# VR runs differ across server starts (0 / 888 same-node and 0 / 296 across node +# pairs with it off, bitwise to the modal autotuned output). The target turns it off for the whole process at init (AUTO +# then takes the C++ size / world heuristic, a pure function of the shape), unless TRTLLM_M31_ENABLE_AR_AUTOTUNE=1. An +# explicit TLLM_DISABLE_ALLREDUCE_AUTOTUNE is left as given. Read once, at import. +AR_AUTOTUNE_ENABLED = m31_enabled("AR_AUTOTUNE") # the m31_switches value contract (1/true/on/yes) + + +def _pin_allreduce_heuristic(rank: int) -> bool: + """Default the AllReduce autotune off for this process (see AR-DET); returns whether it is off. Logs one line.""" + if not AR_AUTOTUNE_ENABLED: + os.environ.setdefault("TLLM_DISABLE_ALLREDUCE_AUTOTUNE", "1") + off = os.environ.get("TLLM_DISABLE_ALLREDUCE_AUTOTUNE", "0") == "1" + logger.info( + f"M3.1 AR-DET: AllReduce autotune {'off (C++ size / world heuristic)' if off else 'on'} on rank {rank}" + ) + return off + + +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 + _pin_allreduce_heuristic(mapping.tp_rank) + 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 MXFP8 index projection (fuse_qkv_index_projection=False) run on two streams. Packed unless the + # deployment sets the flag explicitly (``weights.fused_projection``). + self.fused_projection = _weights.fused_projection(model_config.sparse_attention_config) + + 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 + ) + # AR64AUTO: the strategy-MNNVL instance of the generation boundaries past the multicast all-reduce's rows. + # The MNNVL module of an AllReduce(MNNVL) (kept alive by self._allreduce_wide_owner); None leaves every + # boundary on the configured instance. + self.allreduce_wide = None + self._allreduce_wide_owner = None + if AR64AUTO and DECODE_MAX_TOKENS > m3_mc_allreduce.max_rows(): + self._allreduce_wide_owner = AllReduce( + mapping=mapping, strategy=AllReduceStrategy.MNNVL, dtype=cfg.torch_dtype + ) + self.allreduce_wide = getattr(self._allreduce_wide_owner, "mnnvl_allreduce", None) + if self.allreduce_wide is not None: + logger.info( + f"[M3.1] AR64AUTO: boundaries of {m3_mc_allreduce.max_rows() + 1}..{DECODE_MAX_TOKENS} rows take " + f"the strategy-MNNVL all-reduce (configured: {model_config.allreduce_strategy})" + ) + else: + logger.warning( + "[M3.1] AR64AUTO: MNNVL unavailable on this mapping; every boundary keeps the configured all-reduce" + ) + + # 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__}" + ) + if i not in self.sparse_layers: + # The KV split trtllm-gen dense decode runs at, per request count. + backend.dense_decode_max_seq_len_by_batch = DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH + 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: + # The verify width the config itself declares, whatever it names its draft count. + per_step = getattr(spec, "tokens_per_gen_step", None) + if isinstance(per_step, int) and per_step > 1: + lens.add(per_step) + 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() + self._prebuild_shdown_pin() + torch.cuda.synchronize() + times["shdown_pin_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 + start = time.perf_counter() + times["proxy_variants"] = self._prebuild_msa_proxy() if MSA_PROXY_PREBUILD else [] + times["proxy_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), MSA indexer proxy %.1f s (%d variants)" + % ( + times["gemv_s"], + times["moe_s"], + times["scorer_s"], + len(times["scorer_keys"]), + times["proxy_s"], + len(times["proxy_variants"]), + ) + ) + return times + + def _prebuild_shdown_pin(self) -> None: + """Compile and load the pinned CuTe DSL GEMM (SHDOWN_PIN_ROWS) at every pinned row count before warmup and + capture: FlashInfer compiles it on first use per configuration, which must not happen inside a graph capture.""" + for wk, sk, _ in self.projections(): + weight, scale = self.w[wk], self.w[sk] + pinned = SHDOWN_PIN_ROWS.get(tuple(weight.shape), ()) + if not pinned: + continue + k = weight.shape[1] + for rows in sorted(pinned): + x = torch.zeros(rows, k, dtype=torch.bfloat16, device=weight.device) + act, act_scale = torch.ops.trtllm.mxfp8_quantize(x, True) + m3_mxfp8_gemm_cute_dsl(act.view(rows, k), act_scale, weight, scale, torch.bfloat16) + logger.info( + f"[M3 prebuild] SHDOWN pin: cute_dsl compiled for {list(weight.shape)} " + f"rows {min(pinned)}..{max(pinned)}" + ) + # SWQ (Triton) compiles on first use too; one shape covers every row count. + if SWQ_ROWS.get(tuple(weight.shape)): + m3_swq( + torch.zeros(max(pinned), 2 * k, dtype=torch.bfloat16, device=weight.device), + self.swiglu_limit, + self.swiglu_alpha, + self.swiglu_beta, + ) + logger.info(f"[M3 prebuild] SWQ: compiled for intermediate {k}") + return + + # ------------------------------------------------------------------------------------------ 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_msa_proxy(self) -> list: + """Compile and load every fmha_sm100 variant the MSA indexer's proxy pass can pick (OnlyScore over the paged + index-K: one KV head, page 128, no split-KV), so no context chunk compiles one after the server is ready (a + 65..128-token chunk's variant did, stalling the steps behind it by 0.8-1.6 ms). The variant follows only the + longest query row: fmha_sm100 packs a row of up to 32 tokens by its pack factor, then takes the 128- or 256-row + query tile and the single-warpgroup kernel from the packed length, so query lengths 1..129 reach every variant. + Compile and load only, nothing runs: no numerics change. fmha_sm100 keys its cache by its installed sources, so + a source overlay builds its own variants. Returns the variant names.""" + require_msa_module() + from fmha_sm100 import api as fmha_api + from fmha_sm100 import jit as fmha_jit + + only_score, page = 2, 128 + names = [] + for i in sorted(self.sparse_layers): + backend = self.attn[i] + dtype = torch.float8_e4m3fn if backend.indexer_kv_dtype == "fp8" else torch.bfloat16 + heads = int(backend.m3_config.num_index_heads) + for qlen in range(1, 130): + pack = fmha_api._compute_pack_factor(qlen, heads, 1) + packed = qlen * pack + key = ( + fmha_jit._dlpack_dtype_code(dtype), + 128 if packed <= 128 else 256, + packed <= 64, + only_score, + page, + False, + pack, + ) + name, _ = fmha_jit._variant_key_from_runtime(*key) + if name not in names: + fmha_jit.get_fmha_variant(*key) + names.append(name) + return names + + 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 +# -------------------------------------------------------------------------------------------------------------------- + + +def _swiglu_act(core, gu: torch.Tensor) -> torch.Tensor: + """The measured default SwiGLU-OAI activation path.""" + return m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + + +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 = _swiglu_act(core, gu) + 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 + # CTXNORMQ: the boundary already wrote this activation's NVFP4 form. + fp4 = _CTXNORMQ_FP4.take(x) if CTXNORMQ else None + if num_chunks <= 1: + if fp4 is not None: + return _routed_experts(core, i, fp4[0], fp4[1], router_logits, True) + 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) + sfs = [None] * num_chunks + if fp4 is not None: + xs = fp4[0].split(sizes) + sfs = [t.flatten() for t in fp4[1].view(rows, -1).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], sfs[c], logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], sfs[c], 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) + if norm_w is not core.w["final_norm"]: + fused = _ctxnormq_allreduce_norm(core, x, residual, norm_w, None) + if fused is not None: + return fused + 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: + fused = _ctxnormq_allreduce_norm( + core, + hidden, + residual, + core.w[f"l{i}_post_norm"], + core.moe[i].fc31_input_scale if i in core.moe_layers else None, + ) + hidden, residual = ( + fused + if fused is not None + else _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 MXFP8 index projection + (the checkpoint's block scales applied) and the FP8 indexer kernel on the attention aux stream; then the MSA + backend (it writes the main KV cache) and the o_proj. The NVFP4 producer does not apply: it packs the + five-way projection's output.""" + core = self.core + if _producer_covered(core): + # Not a miss: the producer packs the five-way projection's output, which this configuration lacks. + uncovered( + PRODUCER.path, + "fuse_qkv_index_projection is disabled: QKV + separate MXFP8 index projection", + ) + + 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, _ = _linear(core, hidden, f"l{i}_idx_w", f"l{i}_idx_s", True) + 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 + if TSEL and sparse: + _tsel_register(core.attn[i]) + if GQAMIX and sparse: + _gqamix_register(core.attn[i]) + _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 32 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) + # Fused-layer rows are limited independently of prebuild/workspace capacity. + 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: + uncovered( + "decode_kernels", + f"outside {FUSED_DECODE_MAX_TOKENS}-row fused scope, embeddings or 2-D ids: generic layers", + ) + 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): + captured = spec_metadata is not None and spec_metadata.is_layer_capture(i) + if not kernels or residual is None or captured: + 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 and not MC_ALLREDUCE.covers(x.shape[0]): + uncovered( + MC_ALLREDUCE.path, + f"more than {m3_mc_allreduce.max_rows()} rows: the built-in all-reduce runs", + ) + use_mc = False + 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=False) + + quant = plan.qkv_prequant and i + 1 < core.num_layers + path = moe_path(num_tokens) + if path is None: + uncovered( + MC_FINALIZE.path, + f"more than {m3_mc_allreduce.max_rows()} rows: the finalized MoE and the boundary all-reduce run", + ) + 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) + n, k = core.w[f"{prefix}_dn_w"].shape + rows = gu.shape[0] + if rows in SWQ_ROWS.get((n, k), ()) and _mx_mode(core.mxfp8).decode_graph: + # The MXFP8 activation is published keyed by ``gu`` and taken by the down projection's prequant branch. + act, act_scale = m3_swq(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + m3_prequant.SLOT.publish(gu, act, act_scale) + if rows not in _SWQ_LOGGED: + _SWQ_LOGGED.add(rows) + logger.info(f"[M3 dispatch] SWQ: SwiGLU + MXFP8 producer rows={rows} [{n}, {k}]") + return _linear(core, gu, f"{prefix}_dn_w", f"{prefix}_dn_s", low_m_gemm, push) + act = _swiglu_act(core, gu) + 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 MXFP8 index projection + (the checkpoint's block scales applied) and the FP8 indexer kernel on the attention aux stream; then the MSA + backend (it writes the main KV cache) and the o_proj. The NVFP4 producer does not apply: it packs the + five-way projection's output. ``push`` is the attention + boundary exchange the o_proj GEMV writes its rows into, as in the five-way path.""" + core = self.core + if _producer_covered(core): + # Not a miss: the producer packs the five-way projection's output, which this configuration lacks. + uncovered( + PRODUCER.path, + "fuse_qkv_index_projection is disabled: QKV + separate MXFP8 index projection", + ) + + 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, _ = _linear(core, hidden, f"l{i}_idx_w", f"l{i}_idx_s", True) + 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 " + f"{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 M3DSparkStepFold(M3SpecStepFold): + """The step fold under the standalone DSpark worker (DSFOLD): the engine's previous-batch gathers in one launch + (``step_gather``). ``taken`` / ``declined`` count the steps the gather ran / fell back to the engine's own copies. + + With ``boundary_pack`` (BOUNDARY_PACK BP-1) a batch-1 decode graph also ends with the sampler's slot stores and the + next step's gather (``graph_tail``: one ``ops/m3_bp_tail`` launch), so the boundary between two replays runs + neither: + + * ``begin_step`` (before every replay) writes the device control word (slot, live) when it changed: live 1 with + the step's slot when the step is one row the fold identifies (a slot inside the store), else live 0, so a replay + the fold did not take (a dummy-only replay, an unidentified step) writes nothing at all; it reports the inputs + staged when the previous graph gathered them for exactly these rows; + * the graph's tail waits for the sampler's previous side-stream store copies (``store_copies_done``), then with live + 1 stores the row as the sampler's ``_store_outputs`` (same output row, same store, zero padding / truncation to + the store widths, at the slot) and, with a static draft length, writes the engine's next-step gathers of that row + (P = 1, ids0 = d0 = b0 = 0, both offset buffers zeroed elsewhere), and marks its outputs; + * ``sampler_rows_stored`` (after the replay) accepts a marked step whose rows are the live ones, and arms the next + ``begin_step``. + + Every other step (context or mixed steps, batches over one row, eager forwards, a draft length that can change, + first steps, slots the fold could not identify) keeps the sampler's stores and the engine's gathers. + ``bp_stores`` / ``bp_gathers`` count the steps whose stores / gathers the graph took.""" + + store_copies_done: Optional[torch.cuda.Event] = None + + def __init__(self, boundary_pack: bool = False, tokens_per_step: Optional[int] = None) -> None: + super().__init__() + self.taken = 0 + self.declined = 0 + self.boundary_pack = boundary_pack + # T of every decode step (the draft length cannot change at runtime), else None: the graph then stores only. + self.tokens_per_step = tokens_per_step + self.bp_stores = 0 + self.bp_gathers = 0 + self._bp_ctl: Optional[torch.Tensor] = None # int64 [2] on the device: (slot, live) + self._bp_ctl_host: Optional[list] = None # what it holds (the last write) + self._bp_block = 0 # the tail's row block (m3_bp_tail.block_for the store) + self._bp_rows: Optional[list] = None # the step's rows while live, else None + self._bp_staged: Optional[list] = None # the rows the last stored step's graph gathered for + self._bp_token = id(self) + + def bind_engine_buffers( + self, input_ids, draft_tokens, position_id_offsets, kv_lens_offsets, max_batch_size + ): + super().bind_engine_buffers( + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets, max_batch_size + ) + self._bp_reset() + + def bind_sampler_store(self, store: object) -> None: + super().bind_sampler_store(store) + self._bp_reset() + + def _bp_reset(self) -> None: + """A new binding (an engine or a sampler built again): BOUNDARY_PACK's buffers are made again for it.""" + self._bp_ctl = self._bp_ctl_host = None + self._bp_rows = self._bp_staged = None + self.store_copies_done = None + + def gather(self, *args, **kwargs) -> bool: + if not super().gather(*args, **kwargs): + self.declined += 1 + uncovered( + "DSFOLD", + "step_gather declined the inputs (dtype / layout); the engine's copies run", + ) + return False + if self.taken == 0: + logger.info( + "[M3 dispatch] DSFOLD DSpark step-input gather: engaged (one step_gather launch)" + ) + self.taken += 1 + return True + + def scatter( + self, seq_slots, new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens + ) -> bool: + return False + + def _bp_ready(self) -> bool: + """BOUNDARY_PACK: the control word (live 0), the tail kernel's two variants and the store-copies event, made + once outside capture when the sampler store and the engine buffers are bound and the tail can take them.""" + if not self.boundary_pack: + return False + if self._bp_ctl is not None: + return True + store = self.store + names = ("new_tokens", "next_new_tokens", "new_tokens_lens", "next_draft_tokens") + if ( + self.step_inputs is None + or not all(isinstance(getattr(store, name, None), torch.Tensor) for name in names) + or torch.cuda.is_current_stream_capturing() + ): + return False + tensors = [getattr(store, name) for name in names] + list(self.step_inputs) + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in tensors) + or not all(t.is_contiguous() for t in self.step_inputs) + or not m3_bp_tail.aligned(*tensors) + ): + uncovered( + "BOUNDARY_PACK", + "the store or the engine buffers are not aligned int32 CUDA tensors", + ) + return False + ctl = torch.zeros(2, dtype=torch.long, device=store.new_tokens_lens.device) + block = m3_bp_tail.block_for(store) + m3_bp_tail.warm_up(ctl, store, self.step_inputs, block) + # Recorded by the sampler after its side-stream store copies; the graph's tail waits on it. A CUDA event is + # created on its first record and a wait on one never created is skipped (a capture would drop the wait). + done = torch.cuda.Event(external=True) + done.record() + self.store_copies_done = done + self._bp_ctl, self._bp_ctl_host, self._bp_block = ctl, [0, 0], block + return True + + def begin_step(self, seq_slots: Optional[list], num_rows: int) -> bool: + staged, self._bp_staged = self._bp_staged, None + self._bp_rows = None + if not self._bp_ready(): + return False + num_slots = self.store.new_tokens_lens.shape[0] + rows = list(seq_slots) if seq_slots is not None else None + live = ( + rows is not None + and len(rows) == num_rows + and 0 < num_rows <= BOUNDARY_PACK_MAX_ROWS + and all(isinstance(slot, int) and 0 <= slot < num_slots for slot in rows) + ) + # Live: this step's slot. Not live: the graph (if one replays) writes nothing; the slot word keeps its value. + ctl = [rows[0], 1] if live else [self._bp_ctl_host[0], 0] + if ctl != self._bp_ctl_host: + # Ordered after the previous graph's tail (which read the last word) and before this step's replay. + self._bp_ctl.copy_( + torch.tensor(ctl, dtype=torch.long, pin_memory=prefer_pinned()), non_blocking=True + ) + self._bp_ctl_host = ctl + if not live: + return False + self._bp_rows = rows + if staged != rows: + return False + if self.bp_gathers == 0: + logger.info( + "[M3 dispatch] BOUNDARY-PACK BP-1: sampler stores + step gather in the step graph (taken)" + ) + self.bp_gathers += 1 + return True + + def sampler_rows_stored(self, outputs: dict, num_skip: int, seq_slots: list) -> bool: + self._bp_staged = None + mark = outputs.get(_BOUNDARY_PACK_MARK) if isinstance(outputs, dict) else None + if mark is None or mark[0] != self._bp_token or self._bp_rows is None: + # Unmarked, or the replay was not live (it wrote nothing): the sampler stores. + return False + rows, gathered = mark[1], mark[2] + if num_skip != 0 or len(seq_slots) != rows or list(seq_slots) != self._bp_rows: + # The graph stored the live rows, which are not the sampler's: it stores its own. + uncovered( + "BOUNDARY_PACK", + "a step graph stored rows that are not the sampler's (skipped context rows or other slot ids); " + "the sampler stores them", + ) + return False + self.bp_stores += 1 + if gathered: + self._bp_staged = list(seq_slots) + return True + + def graph_tail(self, outputs, attn_metadata): + """BOUNDARY_PACK, at the end of a decode step's forward: under capture of a batch-1 generation graph, the + sampler's slot stores and the next step's gather, gated by the live word (see the class docstring); returns + the outputs, marked when the graph has the tail. Anything else returns them unchanged (the eager path).""" + if not self.boundary_pack or not torch.cuda.is_current_stream_capturing(): + return outputs + path = "BOUNDARY_PACK" + if not isinstance(outputs, dict) or self._bp_ctl is None: + uncovered( + path, + "the sampler store or the engine buffers are not bound: eager stores and gathers", + ) + return outputs + if attn_metadata.num_contexts != 0: + return outputs + names = ("new_tokens", "next_new_tokens", "new_tokens_lens", "next_draft_tokens") + srcs = [outputs.get(name) for name in names] + if any( + not isinstance(t, torch.Tensor) or t.dtype != torch.int32 or not t.is_cuda for t in srcs + ) or [t.dim() for t in srcs] != [2, 2, 1, 2]: + uncovered( + path, + "the worker outputs are not int32 [rows, width] / [rows]: eager stores and gathers", + ) + return outputs + rows = srcs[0].shape[0] + if any(t.shape[0] != rows for t in srcs): + uncovered(path, "the worker outputs disagree on their rows: eager stores and gathers") + return outputs + if not 0 < rows <= BOUNDARY_PACK_MAX_ROWS: + uncovered( + path, + f"a {rows}-row decode graph (only batch-1 graphs, which are never padded): eager stores and gathers", + ) + return outputs + src = [t[:1] for t in srcs] + if not m3_bp_tail.aligned(*src): + uncovered(path, "a worker output is not 16-byte aligned: eager stores and gathers") + return outputs + store = self.store + # The next step's gather for the same row: the engine's host ints of a steady step (no rows before it, no + # padding), T fixed by the static draft length. + num_tokens = int(attn_metadata.num_tokens) + t = num_tokens // rows + gather = ( + self.tokens_per_step is not None + and t == self.tokens_per_step + and num_tokens == rows * t + and t <= store.next_new_tokens.shape[0] + and t - 1 <= store.next_draft_tokens.shape[1] + ) + if not gather: + uncovered( + path, + f"the next step's tokens per request may differ from this graph's {t} (a draft length that can " + f"change, or {num_tokens} tokens over {rows} rows): the graph stores, the engine gathers", + ) + # The previous step's side-stream copies of the store have read it (an external wait node). + torch.cuda.current_stream().wait_event(self.store_copies_done) + m3_bp_tail.bp_tail(self._bp_ctl, src, store, self.step_inputs, t, gather, self._bp_block) + marked = dict(outputs) + marked[_BOUNDARY_PACK_MARK] = (self._bp_token, rows, gather) + return marked + + +def _static_tokens_per_step(spec_config) -> Optional[int]: + """BOUNDARY_PACK: the tokens per request of every decode step when the draft length cannot change at runtime (no + draft-length schedule, no acceptance-rate gate that turns speculation off), else None.""" + if getattr(spec_config, "draft_len_schedule", None) is not None: + return None + if ( + getattr(spec_config, "acceptance_rate_window_size", None) + and getattr(spec_config, "acceptance_rate_threshold", None) is not None + ): + return None + from tensorrt_llm._torch.speculative.utils import get_static_draft_len + + return int(spec_config.get_runtime_tokens_per_gen_step(get_static_draft_len(spec_config))) + + +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 ValueError( + "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 ValueError( + 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: + return False + if attn_metadata.num_seqs * (draft_len + 1) > DECODE_MAX_TOKENS: + uncovered("spec_tail", f"more than {DECODE_MAX_TOKENS} verify rows") + 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 attn_metadata.num_seqs * (draft_len + 1) > m3_sampled_accept_kernel.MAX_ROWS: + uncovered( + "spec_tail", + f"a sampled verify of more than {m3_sampled_accept_kernel.MAX_ROWS} rows", + ) + 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), or of one-model standalone DSpark (``dspark.build_dspark``: the in-tree drafter + and worker).""" + spec = model_config.spec_config + if _dspark.is_standalone_dspark(spec): + shell.model.embed_tokens = _dspark.M3EmbedTokens(shell.model) + return _dspark.build_dspark(shell, model_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 ValueError( + "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; + under DSpark with ``BOUNDARY_PACK``, both inside a batch-1 decode graph's tail. + * 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. + * HOSTTAIL (kill switches above ``HOSTTAIL_*``): ``sampler_d2h_side_stream`` (the sampler's store copies on its + D2H side stream), ``spec_small_h2d`` (static slot ids, packed accepted counts), ``spec_batch_map_changed_only`` + (the drafter's batch_indices / batch-to-slot uploads only on change), ``skip_accepted_draft_upload`` (no reader + of the accepted counts here), ``pack_spec_rng_uploads`` (Philox uploads in the pack), ``kv_lens_update_once`` + (one MSA KV-length re-derivation per speculative overlap step). + * ``pack_parallel_scatter`` (``PACK_PSCATTER``): the pack's non-overlapping copies scatter in parallel. + * ``pack_in_graph`` (``PACK_IN_GRAPH``) / ``pack_early_copy`` (``PACK_EARLY_COPY``): the decode graph's prologue + issues the pack, from a window uploaded on a side copy stream while the previous graph runs (no H2D copy or + scatter between the graphs). ``h2d_changed_span`` (``H2D_CHANGED_SPAN``): page-crossing steps stage only the + changed page-table entries, so their pack still fits the prologue's window. + """ + return RuntimeHooks( + pack_step_inputs=True, + pack_parallel_scatter=PACK_PSCATTER, + pack_in_graph=PACK_IN_GRAPH, + pack_early_copy=PACK_EARLY_COPY, + skip_unchanged_h2d=True, + h2d_changed_span=H2D_CHANGED_SPAN, + loop_on_execution_stream=True, + merge_tuned_tactics_across_ranks=True, + skip_greedy_rng_upload=True, + spec_step_fold=spec_fold, + kv_lens_update_once=HOSTTAIL_KV_LENS_ONCE, + pack_spec_rng_uploads=HOSTTAIL_PHILOX_PACK, + spec_small_h2d=HOSTTAIL_SMALL_H2D, + skip_accepted_draft_upload=HOSTTAIL_SKIP_ACCEPTED_UPLOAD, + spec_batch_map_changed_only=HOSTTAIL_BATCH_CHANGED_ONLY, + sampler_d2h_side_stream=HOSTTAIL_D2H_SIDE, + 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 ModelingV2MiniMaxM31Nvfp4Sm103Tp4( + 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 + and _dspark.draft_kv_override_applies(spec_config) + ): + # 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[nn.Module] = None + self.draft_config: Optional[ModelConfig] = None + self.spec_worker: Optional[Union[M3Eagle3Worker, _dspark.M3DSparkWorker]] = 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) + # The DSpark worker keeps the in-tree sampler slot stores; only the engine's input gathers fold (DSFOLD). + if not _dspark.is_standalone_dspark(spec_config): + spec_fold = M3SpecStepFold() + elif DSFOLD: + spec_fold = M3DSparkStepFold( + boundary_pack=BOUNDARY_PACK, + tokens_per_step=_static_tokens_per_step(spec_config) if BOUNDARY_PACK else None, + ) + # BOUNDARY_PACK: the fold whose graph_tail ends each decode forward (None: the forward returns the worker's). + self.boundary_pack_fold = ( + spec_fold + if isinstance(spec_fold, M3DSparkStepFold) and spec_fold.boundary_pack + else None + ) + self.runtime_hooks = _runtime_hooks(spec_fold) + _install_topp_det() + _install_softmax_mc() + _install_rej_mc() + + 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.""" + if _dspark.is_standalone_dspark(self.spec_config): + _dspark.load_dspark_draft_weights(self, weights, weight_mapper) + return + _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() + _prebuild_topp_det() + _prebuild_softmax_mc() + _prebuild_rej_mc() + 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: + # A worker may produce the gathered target logits itself (bitwise to the processor); None = the processor. + target_logits = getattr(worker, "target_logits", None) + logits = ( + target_logits(hidden_states, attn_metadata, spec_metadata, self.lm_head) + if target_logits is not None + else None + ) + if logits is None: + 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) + outputs = 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, + ) + if self.boundary_pack_fold is not None and spec_metadata is not None: + # BOUNDARY_PACK: a batch-1 decode graph ends with the sampler's stores and the next step's gather. + outputs = self.boundary_pack_fold.graph_tail(outputs, attn_metadata) + return outputs diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/mxdraft.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/mxdraft.py new file mode 100644 index 000000000000..da8a096aaa15 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/mxdraft.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MXDRAFT: the DSpark drafter's decode projections on this target's MXFP8 GEMV. + +The drafter checkpoint ships its linears as MXFP8 (E4M3 codes, one UE8M0 scale per 32 columns) and the stock drafter +dequantizes them to BF16 at load, so every decode pass streams twice the bytes it needs. MXDRAFT recovers the MXFP8 +form of each BF16 weight (exactly: the recovery is checked to dequantize back to the same BF16 bits, else the +projection is left alone) and routes the projection's GEMM to ``m3_mx_gemv_bf16`` for decode inputs it serves; any +other call runs the Linear unchanged. The activation is quantized to MXFP8 inside the GEMV, so the output is not the +stock BF16 GEMM's: drafter-only precision, gated on natural AL. Row-parallel projections keep +their all-reduce (only the Linear's GEMM is replaced). +""" + +from __future__ import annotations + +from typing import Iterable, Optional, Tuple + +import torch + +from .ops.comm.m3_mc_allreduce_norm import swizzle_sf +from .ops.mx_gemv import op as _gemv + +E4M3_MAX = 448.0 +BLOCK = 32 + + +def mxfp8_from_bf16(weight: torch.Tensor) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """The MXFP8 form (E4M3 codes [N, K], UE8M0 scale bytes [N, K / 32]) of a BF16 weight that is exactly an MXFP8 + dequantization, or None. Each block takes the smallest power-of-two scale that fits its amax in E4M3, which is at + most the shipped scale, so the codes are exact whenever the weight came from MXFP8.""" + if weight.dtype != torch.bfloat16 or weight.dim() != 2 or weight.shape[1] % BLOCK != 0: + return None + n, k = weight.shape + blocks = weight.float().view(n, k // BLOCK, BLOCK) + amax = blocks.abs().amax(-1) + exp = torch.ceil(torch.log2(amax.clamp_min(2.0**-126) / E4M3_MAX)).clamp(-127, 127) + scale = torch.exp2(exp).unsqueeze(-1) + codes = (blocks / scale).to(torch.float8_e4m3fn) + back = (codes.float() * scale).view(n, k).to(torch.bfloat16) + if not torch.equal(back, weight): + return None + return codes.view(n, k).contiguous(), (exp + 127).to(torch.uint8).contiguous() + + +def _projections(draft_model) -> Iterable[torch.nn.Module]: + for layer in getattr(getattr(draft_model, "model", None), "layers", []) or []: + attn, mlp = getattr(layer, "self_attn", None), getattr(layer, "mlp", None) + for lin in ( + getattr(attn, "qkv_proj", None), + getattr(attn, "o_proj", None), + getattr(mlp, "gate_up_proj", None), + getattr(mlp, "down_proj", None), + ): + if lin is not None: + yield lin + + +def route_mx_linears(draft_model, rows: Iterable[int]) -> int: + """Route each drafter projection whose BF16 weight is an exact MXFP8 dequantization to the MXFP8 GEMV for 2-D + contiguous bf16 inputs of the given row counts (the decode passes'); compiles those kernels now (outside capture). + Returns how many projections were routed.""" + rows = sorted({int(m) for m in rows if 0 < int(m) <= _gemv.MAX_TOKENS}) + n_routed = 0 + for lin in _projections(draft_model): + weight = getattr(lin, "weight", None) + if ( + weight is None + or getattr(lin, "bias", None) is not None + or getattr(lin, "_m31_mxdraft", None) is not None + ): + continue + rec = mxfp8_from_bf16(weight.data) + if rec is None: + continue + n, k = weight.shape + served = tuple(m for m in rows if _gemv.serves(n, k, m, True)) + if not served: + continue + w8, sf = rec[0], swizzle_sf(rec[1]) + for m in served: + _gemv.precompile(w8, sf, m, True) + original = lin.apply_linear + + def apply_linear( + input, + bias, + lora_params=None, + layer_idx=None, + *, + _o=original, + _w8=w8, + _sf=sf, + _served=served, + ): + if ( + bias is None + and not lora_params + and isinstance(input, torch.Tensor) + and input.dtype == torch.bfloat16 + and input.dim() == 2 + and input.shape[0] in _served + and input.is_contiguous() + ): + out = _gemv.m3_mx_gemv_bf16(input, _w8, _sf) + if out is not None: + return out + return _o(input, bias, lora_params, layer_idx) + + lin._m31_mxdraft = (w8, sf, served) + lin.apply_linear = apply_linear + n_routed += 1 + return n_routed diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/activation/m3_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/activation/m3_swq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/activation/m3_swq.py new file mode 100644 index 000000000000..95efc345ea4a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/activation/m3_swq.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SWQ: the shared expert's SwiGLU-OAI written directly as the down projection's MXFP8 activation.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +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 + + +@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 _swq_kernel( + x_ptr, + x_stride, + q_ptr, + sf_ptr, + M, + INTER: tl.constexpr, + NB: tl.constexpr, + NB_PAD: tl.constexpr, + P2: tl.constexpr, + swiglu_limit: tl.constexpr, + swiglu_alpha: tl.constexpr, + swiglu_beta: tl.constexpr, + HAS_LIMIT: tl.constexpr, +): + # One program per row of the 128-row-padded scale layout; rows past M only zero their scale bytes. + r = tl.program_id(0) + blk = tl.arange(0, P2) + blk_ok = blk < NB_PAD + # The 128x4 swizzle of trtllm::mxfp8_quantize(x, True) (quantization.cuh get_sf_out_offset_128x4). + sf_off = ( + (blk % 4) + + (blk // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * NB_PAD) + ) + if r >= M: + tl.store(sf_ptr + sf_off, tl.zeros([P2], dtype=tl.uint8), mask=blk_ok) + return + cols = tl.arange(0, P2 * 32) + valid = cols < INTER + row = x_ptr + r.to(tl.int64) * x_stride + # silu_and_mul_kernel's arithmetic (modules/swiglu.py), then its bf16 store. + a = tl.load(row + cols, mask=valid, other=0.0).to(tl.float32) + b = tl.load(row + cols + INTER, mask=valid, other=0.0).to(tl.float32) + if HAS_LIMIT: + a = tl.minimum(a, swiglu_limit) + b = tl.clamp(b, -swiglu_limit, swiglu_limit) + v = (a * tl.sigmoid(swiglu_alpha * a) * (b + swiglu_beta)).to(tl.bfloat16).to(tl.float32) + # cvt_warp_fp16_to_mxfp8: SF = amax * rcp.approx(448), UE8M0 toward +inf, scale = rcp.approx(SF), E4M3 RN. + v2 = tl.reshape(v, (P2, 32)) + amax = tl.max(tl.abs(v2), axis=1) + sf = amax * _rcp_approx_ftz(tl.full([P2], 448.0, tl.float32)) + bits = sf.to(tl.int32, bitcast=True) + exp = (bits >> 23) & 0xFF + mant = bits & 0x7FFFFF + e = tl.where(exp > 0, exp + (mant != 0).to(tl.int32), tl.where(bits > 0x400000, 1, 0)) + e = tl.minimum(e, 254) + # at code 0 the scale 2^-127 is a float subnormal, so rcp.approx.ftz returns inf (native publishes + # inf/NaN there). Use its exact reciprocal 2^127 instead; every other code keeps + # native's rcp.approx (exact on a normal power of two), so each block native emits finite stays bitwise. + pinv = tl.where( + e == 0, 1.7014118346046923e38, _rcp_approx_ftz((e << 23).to(tl.float32, bitcast=True)) + ) + inv = tl.where(amax != 0.0, pinv, 0.0) + q = tl.reshape((v2 * inv[:, None]).to(tl.float8e4nv), (P2 * 32,)) + tl.store(q_ptr + r.to(tl.int64) * INTER + cols, q, mask=valid) + tl.store(sf_ptr + sf_off, tl.where(blk < NB, e, 0).to(tl.uint8), mask=blk_ok) + + +class M3Swq(OpWrapper): + """``silu_and_mul`` (SwiGLU-OAI) over ``x = [g | u]`` (bf16 [M, 2I]) rounded to bf16, then MXFP8-quantized as + ``trtllm::mxfp8_quantize(h, True)``: (E4M3 [M, I], UE8M0 scales in the 128x4-swizzled layout), bit for bit, in one + launch. The shared expert's down projection reads it at the rows its GEMV does not serve (``SWQ_ROWS``).""" + + ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert, the smallest row count it serves", spec=dict(m=17, inter=768)), + Cell(why="shared expert, a 32-row decode verify (c8)", spec=dict(m=32, inter=768)), + ) + + note = """ + Bitwise to the two-kernel path (silu_and_mul's fp32 arithmetic and bf16 rounding, then mxfp8_quantize's + rcp.approx scale rule): q bytes and every scale byte, padding included. + One deliberate exception: a block whose scale code is 0 (amax / 448 <= 2^-127), where native's + rcp.approx.ftz of the subnormal scale is inf and its e4m3 bytes encode inf/NaN, is quantized with the exact + reciprocal 2^127 instead, so it stays finite; its scale byte is native's. + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> tuple[torch.Tensor, torch.Tensor]: + m, i = x.shape[0], x.shape[1] // 2 + nb = i // 32 + nb_pad = (nb + 3) // 4 * 4 + q = torch.empty(m, i, dtype=torch.float8_e4m3fn, device=x.device) + sf = torch.empty((m + 127) // 128 * 128 * nb_pad, dtype=torch.uint8, device=x.device) + _swq_kernel[((m + 127) // 128 * 128,)]( + x, + x.stride(0), + q, + sf, + m, + INTER=i, + NB=nb, + NB_PAD=nb_pad, + P2=triton.next_power_of_2(nb_pad), + swiglu_limit=limit or 0.0, + swiglu_alpha=alpha if alpha is not None else 1.0, + swiglu_beta=beta if beta is not None else 0.0, + HAS_LIMIT=limit is not None and limit > 0.0, + ) + return q, sf + + def reference( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> tuple[torch.Tensor, torch.Tensor]: + """The two kernels it replaces (bitwise is the contract).""" + h = torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + q, sf = torch.ops.trtllm.mxfp8_quantize(h, True) + return q.view(h.shape), sf + + def is_valid( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> None: + assert x.dtype == torch.bfloat16 and x.stride(-1) == 1, "bf16 [M, 2I] with contiguous rows" + assert (x.shape[1] // 2) % 32 == 0, "I must be a multiple of the 32-element MXFP8 block" + + +m3_swq = M3Swq() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/attention/m3_attn_c1.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/attention/m3_attn_mega.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/attention/m3_dense_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/attention/m3_dense_decode.py new file mode 100644 index 000000000000..38d5f1342c78 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/attention/m3_dense_decode.py @@ -0,0 +1,187 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' decode attention (trtllm-gen over the full page table) with the KV split this target sizes per +request count (``dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH``).""" + +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 +#: The split this target hands trtllm-gen at one request (dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH), and a model +#: maximum the in-tree step passes otherwise. +CAPPED = 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): + """``minimax_m3_trtllm_gen_dense_decode`` as ``MsaDecodeFmha`` calls it for a dense layer, with the max_seq_len + ``msa_decode.dense_decode_max_seq_len`` resolves: this target's per-request-count cap, else the step's bound. + + Decode attention of every query row over its request's full causal context in the paged FP8 pool, written into + ``output`` (BF16 [T, heads, head_dim]) in place. ``max_seq_len`` only sizes trtllm-gen's multi-CTA KV split: + every CTA still covers its share of each request's actual length, which may exceed it. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over a 20000-token context, the split sized for 16384", + spec=dict(kv_lens=(20000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request, one row, over a 65536-token context (an AgentX-length cache) at the capped split", + spec=dict(kv_lens=(65536,), qlen=1, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request verifying 1 + 3 rows over a 3000-token context: the cap above the length", + spec=dict(kv_lens=(3000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="two requests keep the model maximum (no cap at two requests)", + spec=dict(kv_lens=(20000, 5000), qlen=4, max_seq_len=MODEL_MAX, capped=False), + ), + ) + + note = """ + The split only reorders the merge of the CTAs' partial softmax sums, so the capped split's output is the uncapped + one's to a few BF16 ulps (the test holds it there), and both are held to the fp32 reference within the FP8-KV + band 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" + # 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_m31_nvfp4__sm_103__tp4/ops/attention/m3_draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/attention/m3_fused_decode_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/attention/m3_index_score.py new file mode 100644 index 000000000000..89deacf68826 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/attention/m3_index_score.py @@ -0,0 +1,167 @@ +# 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="one request verifying 1 + 7 rows over 20000 tokens: the early-issue kernel", + spec=dict(seq_lens=(20000,), qlen=8, early=True), + ), + Cell( + why="two requests verifying 1 + 7 rows: the prefetching kernel", + spec=dict(seq_lens=(20000, 3000), qlen=8, early=False), + ), + 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_m31_nvfp4__sm_103__tp4/ops/attention/m3_msa_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/attention/m3_msa_attention.py new file mode 100644 index 000000000000..dac983d8982b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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 at their capped KV split). ``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_m31_nvfp4__sm_103__tp4/ops/attention/m3_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/attention/m3_producer.py new file mode 100644 index 000000000000..9b63def4c494 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/attention/m3_producer.py @@ -0,0 +1,202 @@ +# 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 K7 c1 verify: 8 tokens", spec=dict(tokens=8, 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 before R1: 48 tokens", spec=dict(tokens=48, scale=1.0)), + Cell( + why="the widest decode step (R1): a K7 c8 verify, 64 tokens, a skipped and an out-of-pool slot", + spec=dict(tokens=64, scale=40.0, bad_slots=True), + ), + 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_m31_nvfp4__sm_103__tp4/ops/attention/m3_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/attention/m3_qkv_insert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/attention/m3_select_blocks.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/attention/m3_select_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode_lc.py new file mode 100644 index 000000000000..b5a1f9b750fd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode_lc.py @@ -0,0 +1,487 @@ +# 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 + +import os +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 +from tensorrt_llm._torch._experimental.modeling_v2.catalog.fakes import ( + m31_fake_ops as _m31_fake_ops, +) + +from ...m31_switches import m31_disabled + +# LC4 (LongAttn): this target's copy of the LC kernel (``minimax_m31_tp4_sparse_decode_lc``, +# kernels/minimaxM31/sparseDecodeLcTp4) with two exact options, rows and hand-off bit-identical to the M3 op. +# LC4_BTSMEM stages the request's block-table row in shared memory before the grid wait, on steps of at most +# ``LC4_BTSMEM_MAX_ROWS`` rows (TP4 K3 chain microbench: -0.16 us/layer at 4 rows and -0.39 at 8, but +3.3 at 16 and +# +7.9 at 32: the extra 32 KB of shared memory per CTA costs the larger grids their occupancy). LC4_SYNC3 reads the +# softmax row max / sum in every thread (two of five __syncthreads go; -0.10 us/layer at 4 rows, -0.05 at 32). +# A build without the op in its library loads it from ``tensorrt_llm/libs/libm31_lc4.so`` when that file exists (an +# overlay); otherwise, or with no option on for the step, the M3 op runs. Kill switches: +# TRTLLM_M31_DISABLE_LC4_BTSMEM=1, TRTLLM_M31_DISABLE_LC4_SYNC3=1. +LC4_BTSMEM = not m31_disabled("LC4_BTSMEM") +LC4_SYNC3 = not m31_disabled("LC4_SYNC3") +LC4_BTSMEM_MAX_ROWS = 8 # capacity: fixed: measured crossover (wins at 8 rows, loses at 16) +_LC4_STATE = {} + + +def _lc4_op(): + """``minimax_m31_tp4_sparse_decode_lc`` when an LC4 option is on and the op is built (or overlaid), else None + (the M3 op runs). Decided once per process.""" + if "op" not in _LC4_STATE: + op = None + if LC4_BTSMEM or LC4_SYNC3: + op = getattr(torch.ops.trtllm, "minimax_m31_tp4_sparse_decode_lc", None) + if op is None: + import tensorrt_llm + + lib = os.path.join(os.path.dirname(tensorrt_llm.__file__), "libs", "libm31_lc4.so") + if os.path.exists(lib): + torch.ops.load_library(lib) + _m31_fake_ops.register() + op = getattr(torch.ops.trtllm, "minimax_m31_tp4_sparse_decode_lc", None) + from tensorrt_llm.logger import logger + + if op is not None: + reason = f"engaged (btsmem={LC4_BTSMEM}, sync3={LC4_SYNC3})" + else: + reason = ( + "switched off" + if not (LC4_BTSMEM or LC4_SYNC3) + else "not taken: the op is not built" + ) + logger.info(f"[M3 dispatch] LC4 sparse decode options: {reason}") + _LC4_STATE["op"] = op + return _LC4_STATE["op"] + + +def _lc4_btsmem(rows: int) -> bool: + """Whether a step of ``rows`` decode rows stages its block-table rows (LC4_BTSMEM, at most LC4_BTSMEM_MAX_ROWS).""" + return LC4_BTSMEM and rows <= LC4_BTSMEM_MAX_ROWS + + +#: 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="7-draft verify at one request: T=8, the log-domain recurrence the target keeps for local batch 1", + spec=dict(seq_lens=(20000,), decode_query_len=8, linear_softmax=False, mxfp8=True), + ), + Cell( + why="7-draft verify at two requests: T=16 on the linear denominator", + spec=dict(seq_lens=(5000, 12345), decode_query_len=8, 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: + lc4 = _lc4_op() + btsmem = _lc4_btsmem(int(q.shape[0])) + if lc4 is not None and (btsmem or LC4_SYNC3): + lc4( + 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, + btsmem, + LC4_SYNC3, + False, # Preserve the measured normal LSE merge; no opt-in one-division path. + ) + return out + 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_m31_nvfp4__sm_103__tp4/ops/attention/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/comm/m3_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/embedding/m3_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/gemm/m3_bf16_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/gemm/m3_fp32_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_gemm.py new file mode 100644 index 000000000000..8671a81c2227 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_gemm.py @@ -0,0 +1,76 @@ +# 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 ( + flashinfer_mxfp8_gemm_cute_dsl, + 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() + + +def m3_mxfp8_gemm_cute_dsl( + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, +) -> torch.Tensor: + """The same product on FlashInfer's CuTe DSL backend, fixed (its heuristic config), for the projections the + target pins to it: the AutoTuner, profiling warm weights per start, can pick the CUTLASS backend there, which + streams the cold decode weight about 2x slower.""" + return flashinfer_mxfp8_gemm_cute_dsl( + act, act_scale, weight, weight_scale, output_dtype, sweep_configs=False + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_pinned_quantize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_pinned_quantize.py new file mode 100644 index 000000000000..a141b7d639c4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_pinned_quantize.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-graph MXFP8 activation quantize at a fixed tactic (QPIN), ahead of the tuned ``m3_mxfp8_gemm``.""" + +from __future__ import annotations + +from typing import Dict, Tuple + +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 MXFP8QuantizeRunner + +_RUNNERS: Dict[torch.dtype, object] = {} + + +def _runner(dtype: torch.dtype): + runner = _RUNNERS.get(dtype) + if runner is None: + runner = _RUNNERS[dtype] = MXFP8QuantizeRunner(dtype) + return runner + + +class M3Mxfp8PinnedQuantize(OpWrapper): + """A bf16 activation [M, K] to MXFP8: (e4m3 [M, K], swizzled UE8M0 scales), on the runner's ``tactic``.""" + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="a 16-row decode verify on a K 6144 input (c4)", spec=dict(m=16, k=6144)), + Cell(why="dense-MLP down projection input, 32 rows", spec=dict(m=32, k=3072)), + ) + + note = """ + Bit-exact: `reference` is the stock MXFP8 quantize (``torch.ops.trtllm.mxfp8_quantize``, swizzled scales) of the + same input; every tactic computes the same block maxima and roundings. + """ + + def __call__(self, x: torch.Tensor, tactic: int) -> Tuple[torch.Tensor, torch.Tensor]: + return _runner(x.dtype).forward([x], tactic=tactic) + + def reference(self, x: torch.Tensor, tactic: int) -> Tuple[torch.Tensor, torch.Tensor]: + """The stock quantize; ``tactic`` is ignored.""" + return torch.ops.trtllm.mxfp8_quantize(x, True) + + +m3_mxfp8_pinned_quantize = M3Mxfp8PinnedQuantize() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/gemm/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/gemm/mxfp8_low_m_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/gqamix_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/gqamix_sparse_decode.py new file mode 100644 index 000000000000..e6b824de1b96 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/gqamix_sparse_decode.py @@ -0,0 +1,1311 @@ +# 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. +""" + +from __future__ import annotations + +import os +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.memory_buffer_utils import get_memory_buffers + +# One sparse block is exactly one KV page. +SPARSE_BLOCK_SIZE = 128 + +_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 + +# 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 + + +def _pdl_enabled() -> bool: + return os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +@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"]), + } +) +@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 + 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 + 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 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() + + 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 _ 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) + + # 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.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.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, +): + pid_b, pid_h = tl.program_id(0), tl.program_id(1) + + if USE_PDL: + 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 USE_PDL: + 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: tuple[int, int], +) -> Optional[int]: + """Select GB300-measured split-K factors for M3's SM100-family decode shape.""" + 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: tuple[int, int], +) -> bool: + """Use the GB300-measured lower-register recurrence on validated M3 shapes.""" + 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: tuple[int, int], +) -> int: + """Group Eagle queries using the GB300-measured shared-page policy.""" + 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: tuple[int, int], +) -> 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 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: tuple[int, int], +) -> dict[str, int]: + """Select GB300-measured one-warp merges for M3's SM100-family graphs.""" + 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}." + ) + + +@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, +) -> 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 + + 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, + capability=capability, + ) + 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, + capability=capability, + ) + 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) + ) + elif num_topk_chunks & (num_topk_chunks - 1): + raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") + + 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 = _pdl_enabled() + 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, + capability=capability, + ) + 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, + capability=capability, + ) + 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 + ) + pair_kernel_options = ( + {"QUERY_GROUP_SIZE": query_group_size, "MAX_TOPK": max_topk} + if query_group_size == 2 + 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, + 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), + 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, + **pair_kernel_options, + **pdl_launch, + **launch_options, + ) + if not direct_output: + 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, + capability=capability, + ) + 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, + **pdl_launch, + **merge_launch_options, + ) + + +__all__ = [ + "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_m31_nvfp4__sm_103__tp4/ops/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/m3_attn_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/m3_attn_mega.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_attn_mega.py new file mode 100644 index 000000000000..964a138afd67 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_attn_mega.py @@ -0,0 +1,206 @@ +# 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: + if not hasattr(torch.ops.trtllm, "minimax_m3_attn_mega"): + raise RuntimeError( + "torch.ops.trtllm.minimax_m3_attn_mega is not registered: this TensorRT-LLM build lacks " + "cpp/tensorrt_llm/kernels/minimaxM3/attnMega.cu" + ) + _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_m31_nvfp4__sm_103__tp4/ops/m3_attn_mega_c1.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_attn_mega_c1.py new file mode 100644 index 000000000000..82c3b1cc3593 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_attn_mega_c1.py @@ -0,0 +1,365 @@ +# 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: + if not hasattr(torch.ops.trtllm, "minimax_m3_attn_c1"): + raise RuntimeError( + "torch.ops.trtllm.minimax_m3_attn_c1 is not registered: this TensorRT-LLM build lacks " + "cpp/tensorrt_llm/kernels/minimaxM3/attnC1.cu" + ) + _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_m31_nvfp4__sm_103__tp4/ops/m3_bp_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_bp_tail.py new file mode 100644 index 000000000000..419a5baef898 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_bp_tail.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""BOUNDARY_PACK BP-1's graph tail for a batch-1 decode step, in one launch: the spec sampler's four slot stores and +the engine's next-step gathers, both gated by a device-side live word. + +``ctl`` (int64 [2]) holds (slot, live), written by the step fold's ``begin_step`` before the replay. With ``live`` 0 +the launch writes nothing (a step the fold did not take: a dummy-only replay, a step it could not identify). With +``live`` 1, for the one row: + +* the stores, as the sampler's ``_store_outputs`` (zero padding / truncation to the store widths, then index_copy_ + at ``slot``):: + + new_tokens[slot][c] = src_new[c] if c < src width else 0 (c < store width), likewise + next_new_tokens[slot][c] , next_draft_tokens[slot][c]; new_tokens_lens[slot] = src_lens + +* with ``GATHER``, the engine's gathers for the next step of the same row (P = 1, no rows before it), as + ``step_gather`` reads them back from the store just written:: + + ids[t] = next_new_tokens[slot][t] (t < T) draft[d] = next_draft_tokens[slot][d] (d < T - 1) + pos_off[i] = lens[slot] if i < T else 0 kv_off[i] = lens[slot] - T if i < 1 else 0 + (every i of pos_off / kv_off: the engine zeroes both buffers before its gathers) + +Integer copies only. Program 0 does the row; every program zero-fills its ``ZBLOCK`` chunk of the offsets. The int +arguments are not specialized and every pointer is 16-byte aligned (the caller checks), so the two ``GATHER`` +variants compiled by ``warm_up`` outside capture serve every graph. +""" + +import torch +import triton +import triton.language as tl + +ZBLOCK = 1024 +_INTS = ( + "s_new_s1", + "s_new_w", + "d_new_s0", + "d_new_s1", + "d_new_w", + "s_nnew_s1", + "s_nnew_w", + "d_nnew_s0", + "d_nnew_s1", + "d_nnew_w", + "d_lens_s0", + "s_draft_s1", + "s_draft_w", + "d_draft_s0", + "d_draft_s1", + "d_draft_w", + "n_pos", + "n_kv", + "T", +) + + +@triton.jit(do_not_specialize=list(_INTS)) +def _bp_tail_kernel( + ctl, + s_new, + s_new_s1, + s_new_w, + d_new, + d_new_s0, + d_new_s1, + d_new_w, + s_nnew, + s_nnew_s1, + s_nnew_w, + d_nnew, + d_nnew_s0, + d_nnew_s1, + d_nnew_w, + s_lens, + d_lens, + d_lens_s0, + s_draft, + s_draft_s1, + s_draft_w, + d_draft, + d_draft_s0, + d_draft_s1, + d_draft_w, + ids, + draft, + pos_off, + n_pos, + kv_off, + n_kv, + T, + GATHER: tl.constexpr, + BLOCK: tl.constexpr, + ZBLOCK_: tl.constexpr, +): + slot = tl.load(ctl) + live = tl.load(ctl + 1) + if live != 0: + pid = tl.program_id(0) + lens = tl.load(s_lens) + if pid == 0: + c = tl.arange(0, BLOCK) + new = tl.load(s_new + c * s_new_s1, mask=c < s_new_w, other=0) + tl.store(d_new + slot * d_new_s0 + c * d_new_s1, new, mask=c < d_new_w) + nnew = tl.load(s_nnew + c * s_nnew_s1, mask=c < s_nnew_w, other=0) + tl.store(d_nnew + slot * d_nnew_s0 + c * d_nnew_s1, nnew, mask=c < d_nnew_w) + dr = tl.load(s_draft + c * s_draft_s1, mask=c < s_draft_w, other=0) + tl.store(d_draft + slot * d_draft_s0 + c * d_draft_s1, dr, mask=c < d_draft_w) + tl.store(d_lens + slot * d_lens_s0, lens) + if GATHER: + tl.store(ids + c, nnew, mask=c < T) + tl.store(draft + c, dr, mask=c < T - 1) + if GATHER: + z = pid * ZBLOCK_ + tl.arange(0, ZBLOCK_) + tl.store(pos_off + z, tl.where(z < T, lens, 0), mask=z < n_pos) + tl.store(kv_off + z, tl.where(z < 1, lens - T, 0), mask=z < n_kv) + + +def block_for(store) -> int: + """The row block of a sampler store: every store width fits it (fixed per store, so per fold).""" + widths = ( + store.new_tokens.shape[0], + store.next_new_tokens.shape[0], + store.next_draft_tokens.shape[1], + ) + return triton.next_power_of_2(max(*widths, 1)) + + +def aligned(*tensors: torch.Tensor) -> bool: + """Every pointer the launch takes is 16-byte aligned (the compiled variants' assumption).""" + return all(t.data_ptr() % 16 == 0 for t in tensors) + + +def bp_tail(ctl, src, store, engine, t: int, gather: bool, block: int) -> None: + """One launch: ``src`` = (new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens) row 0 of the worker's + int32 outputs ([1, w] / [1]); ``store`` the sampler's store; ``engine`` = (input_ids, draft_tokens, + position_id_offsets, kv_lens_offsets) int32 contiguous; ``t`` tokens per request (gather only).""" + s_new, s_nnew, s_lens, s_draft = src + d_new = store.new_tokens.squeeze(-1).T + d_nnew = store.next_new_tokens.squeeze(-1).T + d_lens, d_draft = store.new_tokens_lens, store.next_draft_tokens + ids, draft, pos_off, kv_off = engine + n_pos, n_kv = pos_off.numel(), kv_off.numel() + grid = (max(1, triton.cdiv(max(n_pos, n_kv), ZBLOCK)) if gather else 1,) + _bp_tail_kernel[grid]( + ctl, + s_new, + s_new.stride(1), + s_new.shape[1], + d_new, + d_new.stride(0), + d_new.stride(1), + d_new.shape[1], + s_nnew, + s_nnew.stride(1), + s_nnew.shape[1], + d_nnew, + d_nnew.stride(0), + d_nnew.stride(1), + d_nnew.shape[1], + s_lens, + d_lens, + d_lens.stride(0), + s_draft, + s_draft.stride(1), + s_draft.shape[1], + d_draft, + d_draft.stride(0), + d_draft.stride(1), + d_draft.shape[1], + ids, + draft, + pos_off, + n_pos, + kv_off, + n_kv, + t, + GATHER=gather, + BLOCK=block, + ZBLOCK_=ZBLOCK, + num_warps=4, + ) + + +def warm_up(ctl, store, engine, block: int) -> None: + """Compile both variants outside capture, with ``ctl`` live 0 (the launch writes nothing).""" + src = ( + store.new_tokens.squeeze(-1).T[:1], + store.next_new_tokens.squeeze(-1).T[:1], + store.new_tokens_lens[:1], + store.next_draft_tokens[:1], + ) + for gather in (False, True): + bp_tail(ctl, src, store, engine, 2, gather, block) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/m3_fused_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_fused_attn.py new file mode 100644 index 000000000000..fefd30acf333 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_fused_attn.py @@ -0,0 +1,352 @@ +# 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: + if not hasattr(torch.ops.trtllm, "minimax_m3_fused_decode_attn"): + 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" + ) + _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_m31_nvfp4__sm_103__tp4/ops/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_index_score.py new file mode 100644 index 000000000000..812410b130ed --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_index_score.py @@ -0,0 +1,645 @@ +# 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.cute_dsl_kernels.blackwell.cute_ptx_utils import ( + EVICT_FIRST, + fp8x4_to_fp16x4, + mma_sync, + simple_tma_copy, +) +from tensorrt_llm._torch.cute_dsl_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.cute_dsl_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.cute_dsl_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_m31_nvfp4__sm_103__tp4/ops/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_mc_allreduce.py new file mode 100644 index 000000000000..fc992b83d152 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_mc_allreduce.py @@ -0,0 +1,376 @@ +# 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 + + +def max_rows() -> int: + """The most rows one boundary all-reduce call (plain or MoE-finalize) serves: the decode plans and the + MC_ALLREDUCE / MC_FINALIZE ranges read this, so a wider bucket or a row split here widens them with no change in + the modeling.""" + return MAX_TOKENS + + +# 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 ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + 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() + 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): + 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_m31_nvfp4__sm_103__tp4/ops/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_moe/op.py new file mode 100644 index 000000000000..87605e245fe3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_moe/op.py @@ -0,0 +1,622 @@ +# 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 + +from ...m31_switches import m31_disabled + +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) +# Early group-0 stages warp 4 issues right after its top-k (``early0_n``; 0 = the kernel's default, every ring stage). +# Fewer at 4 and 8 tokens measured faster on GB300 (MOE_EARLY0N_READ: 6 at 4 tokens -0.37 us, 3 at 8 tokens -0.15 us +# per call, bit-identical outputs); other token counts were flat or unmeasured. Kill switch +# ``TRTLLM_M31_DISABLE_MOE_EARLY0N`` restores the default everywhere. +EARLY0_N_BY_M: Tuple[int, ...] = (0, 0, 0, 6, 0, 0, 0, 3) +# Per-token-count overrides of the tables above, measured together: at 4 tokens the B-image FC1 -> FC2 hand-off plus 8 +# ring stages run 1.62 us/call faster on GB300 with bit-identical outputs (MOE_CLK64_MR6C_READ; each alone ~0.8). +# Kill switch ``TRTLLM_M31_DISABLE_MOE_M4CFG`` restores the tables' values. +M4_CFG_BY_M: Dict[int, Dict[str, int]] = {4: {"b2img": 1, "stages": 8}} +# 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] + if EARLY0_N_BY_M[num_tokens - 1] > 0 and not m31_disabled("MOE_EARLY0N"): + self.config["early0_n"] = EARLY0_N_BY_M[num_tokens - 1] + if num_tokens in M4_CFG_BY_M and not m31_disabled("MOE_M4CFG"): + self.config.update(M4_CFG_BY_M[num_tokens]) + 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_m31_nvfp4__sm_103__tp4/ops/m3_moe/op16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_sampled_accept.py new file mode 100644 index 000000000000..59c58336c5d7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_sampled_accept.py @@ -0,0 +1,159 @@ +# 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 ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + 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() + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + 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_m31_nvfp4__sm_103__tp4/ops/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_shard_accept.py new file mode 100644 index 000000000000..2b7d05ad920f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/m3_shard_accept.py @@ -0,0 +1,176 @@ +# 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 ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + 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() + if not _mnnvl_workspace_all_succeeded(comm, True): + 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_m31_nvfp4__sm_103__tp4/ops/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/moe/m3_moe16_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/moe/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/moe/m3_router_gemv_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/moe/m3_router_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/mxfp8_low_m_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/sampling/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/sampling/__init__.py new file mode 100644 index 000000000000..b3611a4b7866 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/sampling/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sampling entries this target alone calls: the spec sampler's top-p renorm override (ST-01 TOPP_DET), its softmax +override (ST-01-SM SOFTMAX_MC) and its chain speculative sampling override (ST-01-REJ REJ_MC).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/sampling/m31_chain_spec_mc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/sampling/m31_chain_spec_mc.py new file mode 100644 index 000000000000..2d51fd37b047 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/sampling/m31_chain_spec_mc.py @@ -0,0 +1,266 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ST-01-REJ REJ_MC: the spec sampler's chain speculative sampling (rejection) on many CTAs per request, exact vs +``torch.ops.trtllm.chain_speculative_sampling`` (``trtllm::minimax_m31_chain_speculative_sampling_mc``, +kernels/minimaxM31/chainSpecSamplingMcSm103.cu). + +The target installs ``chain_spec_mc`` as the rejection sampler's override +(``speculative/interface.set_chain_speculative_sampling_override``); it returns None, and the sampler keeps the stock +op, whenever the op is not built or the input is outside the kernel's contract. + +Workspace: one zero-initialized byte buffer per device, ``MAX_BATCH`` requests of +``minimax_m31_chain_speculative_sampling_mc_ws_bytes()``, kept for the life of the process (CUDA graphs captured with +it keep its address). The kernel leaves it ready for the next launch, so it is zeroed once. ``prebuild`` allocates it +outside any CUDA-graph capture (the target calls it from ``cache_derived_state``, before the executor's warmup and +captures); a call that finds no workspace allocates one when the stream is not capturing, and inside a capture returns +None (the stock op for that call) rather than allocating from the graph's private pool. Calls that share it must be +stream-ordered (the sampler's are). +""" + +from __future__ import annotations + +import os +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 tensorrt_llm._torch._experimental.modeling_v2.catalog.fakes import ( + m31_fake_ops as _m31_fake_ops, +) + +# The largest CTA count per request (the grid is batch x ctas_per_row 1024-thread CTAs, all co-resident). +MAX_CTAS_PER_ROW = 16 + +# The requests the workspace holds and the override takes: the TP8 / TP4 K7 calls up to c8. Microbenchmarks +# gated it bitwise at batch 1 / 2 / 4 / 8 and measured 31.2 vs 41.2 us per call at batch 1; larger calls keep the stock +# kernel. +MAX_BATCH = 8 + +# The widest vocab the kernel takes (64 tiles of 4096; kChainSpecMcMaxVocab). +MAX_VOCAB = 64 * 4096 + +# The op when it is built (or overlaid as tensorrt_llm/libs/libm31_chain_spec_mc.so), decided once per process. +_STATE: dict = {} +# Per device index: the persistent workspace and the SM count. +_WORKSPACES: dict = {} +_SMS: dict = {} + + +def chain_spec_mc_op(): + """``trtllm::minimax_m31_chain_speculative_sampling_mc`` if this build carries it (in-tree, or the overlay + library), else None.""" + if "op" not in _STATE: + op = getattr(torch.ops.trtllm, "minimax_m31_chain_speculative_sampling_mc", None) + if op is None: + import tensorrt_llm + + lib = os.path.join( + os.path.dirname(tensorrt_llm.__file__), "libs", "libm31_chain_spec_mc.so" + ) + if os.path.exists(lib): + torch.ops.load_library(lib) + _m31_fake_ops.register() + op = getattr(torch.ops.trtllm, "minimax_m31_chain_speculative_sampling_mc", None) + _STATE["op"] = op + _STATE["ws_bytes"] = ( + int(torch.ops.trtllm.minimax_m31_chain_speculative_sampling_mc_ws_bytes()) + if op is not None + else 0 + ) + return _STATE["op"] + + +def ws_bytes_per_row() -> int: + """Workspace bytes per request (0 when the op is not built).""" + chain_spec_mc_op() + return _STATE["ws_bytes"] + + +def sm_count(device: torch.device) -> int: + """The device's SM count, cached per device.""" + idx = torch.device(device).index + idx = torch.cuda.current_device() if idx is None else idx + if idx not in _SMS: + _SMS[idx] = torch.cuda.get_device_properties(idx).multi_processor_count + return _SMS[idx] + + +def ctas_per_row(batch: int, sms: int) -> int: + """The largest power of two <= MAX_CTAS_PER_ROW with ctas x batch <= sms (1 past that: one CTA per request needs + no co-residency).""" + g = MAX_CTAS_PER_ROW + while g > 1 and g * batch > sms: + g //= 2 + return g + + +def prebuild(device: Optional[torch.device] = None) -> Optional[torch.Tensor]: + """Allocate (once) the device's zeroed workspace for MAX_BATCH requests; must run outside a CUDA-graph capture. + Returns it, or None when the op is not built or the stream is capturing.""" + if chain_spec_mc_op() is None: + return None + dev = ( + torch.device("cuda", torch.cuda.current_device()) + if device is None + else torch.device(device) + ) + if dev.index is None: + dev = torch.device("cuda", torch.cuda.current_device()) + if dev.index not in _WORKSPACES: + if torch.cuda.is_current_stream_capturing(): + return None + _WORKSPACES[dev.index] = torch.zeros( + MAX_BATCH * ws_bytes_per_row(), dtype=torch.uint8, device=dev + ) + return _WORKSPACES[dev.index] + + +class M31ChainSpecMc(OpWrapper): + """``trtllm::chain_speculative_sampling(draft_probs, draft_token_ids, target_probs, seed, offset)`` bit for bit + (accepted tokens [batch, D + 1] and token counts [batch]) in ONE launch of batch x ctas_per_row CTAs (stock: one + CTA per request). ``workspace`` is the persistent zeroed buffer of ``prebuild``. ``__call__`` returns None (the + caller keeps the stock op) for a non-fp32, non-contiguous or misaligned probs, a vocab not divisible by 4 or above + MAX_VOCAB, mismatched shapes, non-int32 draft tokens, non-int64 or mis-sized seed / offset, inputs on different + devices, more requests than the workspace holds, a batch x ctas_per_row above the SM count, or when the op is not + built.""" + + ARCHS = frozenset({Arch.SM_103}) + + # The K7 calls were checked bitwise on GB300 over 5 cases x 64 seeds at ctas_per_row 4 / 8 / 16. + CELLS: tuple[Cell, ...] = ( + Cell(why="K7 at c1 (vocab 200064, 16 CTAs)", spec=dict(batch=1, k=7, v=200064)), + Cell(why="K7 at c2 (vocab 200064)", spec=dict(batch=2, k=7, v=200064)), + Cell(why="K7 at c4 (vocab 200064)", spec=dict(batch=4, k=7, v=200064)), + Cell(why="K7 at c8 (vocab 200064, 128 CTAs)", spec=dict(batch=8, k=7, v=200064)), + ) + + note = """ + Exact by construction: every CTA of a request runs the stock Philox acceptance loop, the (target, residual) mass is + the stock 1024 per-thread chains into the same cub BlockReduce, each tile sum is the stock per-tile + cub reduce, and one CTA runs the stock tile-order prefix and in-tile cub BlockScan search. Default flags, as the + stock kernel. + """ + + def __call__( + self, + draft_probs: torch.Tensor, + draft_token_ids: torch.Tensor, + target_probs: torch.Tensor, + seed: torch.Tensor, + offset: torch.Tensor, + workspace: Optional[torch.Tensor], + ctas_per_row: int, + ) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + if ( + workspace is None + or not all( + isinstance(t, torch.Tensor) + for t in (draft_probs, draft_token_ids, target_probs, seed, offset) + ) + or not draft_probs.is_cuda + or draft_probs.dtype != torch.float32 + or target_probs.dtype != torch.float32 + or draft_token_ids.dtype != torch.int32 + or seed.dtype != torch.int64 + or offset.dtype != torch.int64 + or draft_probs.dim() != 3 + or target_probs.dim() != 3 + or draft_token_ids.dim() != 2 + ): + return None + batch, k, v = draft_probs.shape + if ( + batch < 1 + or tuple(draft_token_ids.shape) != (batch, k) + or tuple(target_probs.shape) != (batch, k + 1, v) + or seed.numel() != batch + or offset.numel() != batch + or v % 4 != 0 + or v > MAX_VOCAB + or any( + not t.is_contiguous() or t.device != draft_probs.device + for t in (draft_token_ids, target_probs, seed, offset, workspace) + ) + or not draft_probs.is_contiguous() + or draft_probs.data_ptr() % 16 != 0 + or target_probs.data_ptr() % 16 != 0 + or ctas_per_row < 1 + or 1024 % ctas_per_row != 0 + or (ctas_per_row > 1 and batch * ctas_per_row > sm_count(draft_probs.device)) + ): + return None + op = chain_spec_mc_op() + if op is None or workspace.numel() * workspace.element_size() < batch * ws_bytes_per_row(): + return None + tokens, num_tokens = op( + draft_probs, draft_token_ids, target_probs, seed, offset, workspace, int(ctas_per_row) + ) + return tokens, num_tokens + + def reference( + self, + draft_probs, + draft_token_ids, + target_probs, + seed, + offset, + workspace=None, + ctas_per_row=None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """The stock op (workspace and ctas_per_row are ignored).""" + return torch.ops.trtllm.chain_speculative_sampling( + draft_probs, draft_token_ids, target_probs, seed, offset + ) + + def is_valid( + self, + draft_probs, + draft_token_ids, + target_probs, + seed, + offset, + workspace=None, + ctas_per_row=None, + **kwargs, + ) -> None: + assert draft_probs.dtype == torch.float32 and draft_probs.dim() == 3, ( + "draft_probs: fp32 [batch, D, vocab]" + ) + + +m31_chain_spec_mc = M31ChainSpecMc() + + +def chain_spec_mc( + draft_probs: torch.Tensor, + draft_token_ids: torch.Tensor, + target_probs: torch.Tensor, + seed: torch.Tensor, + offset: torch.Tensor, +) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + """The rejection sampler's override: the stock op's (accepted tokens, token counts), or None (the sampler runs the + stock op) when the op is not built, the batch exceeds MAX_BATCH, no workspace can be had (first use inside a + capture) or the input is outside the contract.""" + if ( + not isinstance(draft_probs, torch.Tensor) + or not draft_probs.is_cuda + or draft_probs.dim() != 3 + ): + return None + batch = draft_probs.shape[0] + if batch < 1 or batch > MAX_BATCH or chain_spec_mc_op() is None: + return None + workspace = prebuild(draft_probs.device) + if workspace is None: + return None + return m31_chain_spec_mc( + draft_probs, + draft_token_ids, + target_probs, + seed, + offset, + workspace, + ctas_per_row(batch, sm_count(draft_probs.device)), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/sampling/m31_softmax_mc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/sampling/m31_softmax_mc.py new file mode 100644 index 000000000000..ac4d6dc8add2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/sampling/m31_softmax_mc.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ST-01-SM SOFTMAX_MC: the spec sampler's softmax with its write pass on many CTAs per row, exact vs +``flashinfer.sampling.softmax(logits, temperature)`` (``trtllm::minimax_m31_softmax_mc``, +kernels/minimaxM31/softmaxMcSm103.cu). + +The target installs ``softmax_mc`` as the sampler's softmax override (``sampler/ops/flashinfer.set_softmax_override``); +it returns None, and the sampler keeps the stock FlashInfer call, whenever the op is not built or the input is outside +the kernel's contract (FlashInfer's Path A shapes: rows <= 128 and vocab >= 24576; here also rows <= ``MAX_ROWS``). + +Workspace: one partials buffer per device (``MAX_ROWS`` x ceil(``MAX_VOCAB`` / 8192) float2, 8 KB), kept for the life +of the process (CUDA graphs captured with it keep its address). Every launch overwrites what it reads, so it needs no +initialization. ``prebuild`` allocates it outside any CUDA-graph capture (the target calls it from +``cache_derived_state``, before the executor's warmup and captures); a call that finds no workspace allocates one when +the stream is not capturing, and inside a capture returns None (stock FlashInfer for that call) rather than allocating +from the graph's private pool. Calls that share it must be stream-ordered (the sampler's are). +""" + +from __future__ import annotations + +import os +from typing import Optional, Union + +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.fakes import ( + m31_fake_ops as _m31_fake_ops, +) + +# The write pass's CTAs per row (microbenchmark: 9.2 vs FlashInfer 21.8 us per call at 7-8 rows; bitwise at +# 4 / 8 / 16). +WRITE_CTAS = 16 + +# The rows the override takes: the TP8 / TP4 K7 c1 calls (the draft's 7 rows and the verify's 8), where the MB measured +# it. Larger calls keep the stock kernel. +MAX_ROWS = 8 + +# FlashInfer's Path A bounds (OnlineSoftmax: batch <= 128 and d >= 24576); outside them FlashInfer takes its +# single-block Path B, which this kernel does not reproduce. +PATH_A_MAX_ROWS = 128 +PATH_A_MIN_VOCAB = 24576 +# FlashInfer's Path A slice (DEFAULT_SLICE_SIZE) and the widest vocab the workspace holds. +SLICE = 8192 +MAX_VOCAB = 1 << 20 + +# The op when it is built (or overlaid as tensorrt_llm/libs/libm31_softmax_mc.so), decided once per process. +_STATE: dict = {} +# Per device index: the persistent partials workspace. +_WORKSPACES: dict = {} + + +def softmax_mc_op(): + """``trtllm::minimax_m31_softmax_mc`` if this build carries it (in-tree, or the overlay library), else None.""" + if "op" not in _STATE: + op = getattr(torch.ops.trtllm, "minimax_m31_softmax_mc", None) + if op is None: + import tensorrt_llm + + lib = os.path.join( + os.path.dirname(tensorrt_llm.__file__), "libs", "libm31_softmax_mc.so" + ) + if os.path.exists(lib): + torch.ops.load_library(lib) + _m31_fake_ops.register() + op = getattr(torch.ops.trtllm, "minimax_m31_softmax_mc", None) + _STATE["op"] = op + return _STATE["op"] + + +def partials_numel(rows: int, d: int) -> int: + """fp32 elements of the partials workspace for rows x d (one float2 per row and 8192-wide slice).""" + return 2 * rows * ((d + SLICE - 1) // SLICE) + + +def prebuild(device: Optional[torch.device] = None) -> Optional[torch.Tensor]: + """Allocate (once) the device's partials workspace for MAX_ROWS rows of up to MAX_VOCAB; must run outside a + CUDA-graph capture. Returns it, or None when the op is not built or the stream is capturing.""" + if softmax_mc_op() is None: + return None + dev = ( + torch.device("cuda", torch.cuda.current_device()) + if device is None + else torch.device(device) + ) + if dev.index is None: + dev = torch.device("cuda", torch.cuda.current_device()) + if dev.index not in _WORKSPACES: + if torch.cuda.is_current_stream_capturing(): + return None + _WORKSPACES[dev.index] = torch.empty( + partials_numel(MAX_ROWS, MAX_VOCAB), dtype=torch.float32, device=dev + ) + return _WORKSPACES[dev.index] + + +class M31SoftmaxMc(OpWrapper): + """FlashInfer's ``softmax(logits, temperature)`` of fp32 logits [rows, d], bit for bit, on FlashInfer's Path A + shapes (rows <= 128, d >= 24576): FlashInfer's Map kernel, then its Reduce kernel's prologue in each of + ``write_ctas`` CTAs per row, each writing its share of the row (FlashInfer: one CTA writes the whole row). + ``temperature`` is None (1.0), a Python number (one temperature) or an fp32 [rows] tensor (one per row), as + FlashInfer takes it. ``partials`` is the persistent workspace of ``prebuild``. ``__call__`` returns None (the caller + keeps the stock call) for a non-fp32, non-contiguous or misaligned input, a width not divisible by 4, a shape + outside Path A or the workspace, a temperature tensor that is not a contiguous fp32 [rows] tensor on logits' + device, or when the op is not built.""" + + ARCHS = frozenset({Arch.SM_103}) + + # The c1 K7 calls were checked bitwise on GB300. + CELLS: tuple[Cell, ...] = ( + Cell(why="the K7 draft's 7 rows at c1 (vocab 200064)", spec=dict(rows=7, d=200064)), + Cell(why="the K7 verify's 8 rows at c1 (vocab 200064)", spec=dict(rows=8, d=200064)), + Cell(why="one row (the smallest grid)", spec=dict(rows=1, d=200064)), + ) + + note = """ + Exact by construction: the Map kernel is FlashInfer's, every write CTA reduces the same partials with the same cub + BlockReduce and Float2SoftmaxReduceOp as FlashInfer's single Reduce CTA, and the per-element write + keeps FlashInfer's source form. Built with --use_fast_math, as FlashInfer's JIT builds it. + """ + + def __call__( + self, + logits: torch.Tensor, + temperature: Union[torch.Tensor, float, None], + partials: Optional[torch.Tensor], + write_ctas: int, + ) -> Optional[torch.Tensor]: + if ( + partials is None + or not isinstance(logits, torch.Tensor) + or not logits.is_cuda + or logits.dtype != torch.float32 + or logits.dim() != 2 + or not logits.is_contiguous() + or logits.shape[0] < 1 + or logits.shape[0] > PATH_A_MAX_ROWS + or logits.shape[1] < PATH_A_MIN_VOCAB + or logits.shape[1] % 4 != 0 + or logits.data_ptr() % 16 != 0 + or partials.device != logits.device + or partials.numel() * partials.element_size() + < 4 * partials_numel(logits.shape[0], logits.shape[1]) + or write_ctas < 1 + or write_ctas > 65535 + ): + return None + if temperature is None: + temp_arr, temp_val = None, 1.0 + elif isinstance(temperature, torch.Tensor): + if ( + temperature.dtype != torch.float32 + or temperature.dim() != 1 + or temperature.shape[0] != logits.shape[0] + or not temperature.is_contiguous() + or temperature.device != logits.device + ): + return None + temp_arr, temp_val = temperature, 0.0 + elif isinstance(temperature, (int, float)) and not isinstance(temperature, bool): + temp_arr, temp_val = None, float(temperature) + else: + return None + op = softmax_mc_op() + if op is None: + return None + out = torch.empty_like(logits) + op(logits, temp_arr, temp_val, partials, out, int(write_ctas)) + return out + + def reference(self, logits, temperature, partials=None, write_ctas=None) -> torch.Tensor: + """FlashInfer's softmax (partials and write_ctas are ignored).""" + import flashinfer.sampling + + return flashinfer.sampling.softmax(logits, temperature) + + def is_valid(self, logits, temperature, partials=None, write_ctas=None, **kwargs) -> None: + assert logits.dtype == torch.float32 and logits.dim() == 2, "logits: fp32 [rows, d]" + + +m31_softmax_mc = M31SoftmaxMc() + + +def softmax_mc( + logits: torch.Tensor, temperature: Union[torch.Tensor, float, None] +) -> Optional[torch.Tensor]: + """The sampler's softmax override: FlashInfer's softmax of logits [rows, d] at temperature, or None (the sampler + runs stock FlashInfer) when the op is not built, rows exceed MAX_ROWS, d exceeds MAX_VOCAB, no workspace can be had + (first use inside a capture) or the input is outside the contract.""" + if not isinstance(logits, torch.Tensor) or not logits.is_cuda or logits.dim() != 2: + return None + rows = logits.shape[0] + if rows < 1 or rows > MAX_ROWS or logits.shape[1] > MAX_VOCAB or softmax_mc_op() is None: + return None + partials = prebuild(logits.device) + if partials is None: + return None + return m31_softmax_mc(logits, temperature, partials, WRITE_CTAS) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/sampling/m31_topp_det.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/sampling/m31_topp_det.py new file mode 100644 index 000000000000..0d28b1e2699c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/sampling/m31_topp_det.py @@ -0,0 +1,248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ST-01 TOPP_DET: the spec sampler's top-p renorm as one persistent C++ launch, exact vs FlashInfer's deterministic +``top_p_renorm_probs`` (``trtllm::minimax_m31_topp_renorm_det``, kernels/minimaxM31/toppRenormDetSm103.cu). + +TOPP_DET v3 (``trtllm::minimax_m31_topp_renorm_det_v3``, kernels/minimaxM31/toppRenormDetV3Sm103.cu) is the same +renorm, bitwise v1.2's output, with the redundant per-CTA selection and a programmatic dependent launch. It is taken +when it is built (in-tree, or overlaid as tensorrt_llm/libs/libm31_topp_det_v3.so) unless +``TRTLLM_M31_DISABLE_TOPPDET_V3`` is set; v1.2 otherwise. ``TRTLLM_M31_DISABLE_TOPPDET_T3`` keeps v3 on its global +radix passes (no candidate list). + +The target installs ``top_p_renorm_det`` as the sampler's top-p renorm override +(``sampler/ops/flashinfer.set_top_p_renorm_override``); it returns None, and the sampler keeps the stock FlashInfer +call, whenever the op is not built or the input is outside the kernel's contract. + +Workspace: one zero-initialized byte buffer per device, ``MAX_ROWS`` rows of +``minimax_m31_topp_renorm_det_ws_bytes()``, kept for the life of the process (CUDA graphs captured with it keep its +address). The kernel leaves it ready for the next launch, so it is zeroed once. ``prebuild`` allocates it outside any +CUDA-graph capture (the target calls it from ``cache_derived_state``, before the executor's warmup and captures); a +call that finds no workspace allocates one when the stream is not capturing, and inside a capture returns None +(stock FlashInfer for that call) rather than allocating from the graph's private pool. +""" + +from __future__ import annotations + +import hashlib +import os +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 tensorrt_llm._torch._experimental.modeling_v2.catalog.fakes import ( + m31_fake_ops as _m31_fake_ops, +) +from tensorrt_llm.logger import logger + +from ...m31_switches import m31_disabled + +# The largest CTA count per row (the persistent grid is rows x ctas_per_row 1024-thread CTAs, all co-resident). +MAX_CTAS_PER_ROW = 16 + +# The rows the workspace holds and the override takes: the TP8 / TP4 K7 c1 calls (the draft's 7 rows and the +# verify's 8), where the microbenchmarks measured it faster than stock FlashInfer (v1.2: 54.3 vs 104.2 us per +# call at 7 peaked rows, 54.8 vs 105.1 at 8). At 56 / 64 rows v1 was slower than FlashInfer's default mode (340 vs +# 187 / 199 us) and larger shapes are not re-measured, so larger calls keep the stock kernel. +MAX_ROWS = 8 + +# The op when it is built (or overlaid as tensorrt_llm/libs/libm31_topp_det.so), decided once per process. +_STATE: dict = {} +# Per device index: the persistent workspace and the SM count. +_WORKSPACES: dict = {} +_SMS: dict = {} + + +def _load_op(name: str, lib_name: str): + """``torch.ops.trtllm.`` if registered (in-tree), else after loading tensorrt_llm/libs/ when it + exists, else None. Returns (op, the library path loaded or None).""" + op = getattr(torch.ops.trtllm, name, None) + if op is not None: + return op, None + import tensorrt_llm + + lib = os.path.join(os.path.dirname(tensorrt_llm.__file__), "libs", lib_name) + if not os.path.exists(lib): + return None, None + torch.ops.load_library(lib) + _m31_fake_ops.register() + return getattr(torch.ops.trtllm, name, None), lib + + +def topp_det_op(): + """The renorm op this process uses, decided once: ``trtllm::minimax_m31_topp_renorm_det_v3`` when it is built + and TOPPDET_V3 is not disabled, else ``trtllm::minimax_m31_topp_renorm_det`` when built, else None.""" + if "op" not in _STATE: + op, ws, v3 = None, 0, False + if not m31_disabled("TOPPDET_V3"): + op, lib = _load_op("minimax_m31_topp_renorm_det_v3", "libm31_topp_det_v3.so") + if op is not None: + v3 = True + ws = int(torch.ops.trtllm.minimax_m31_topp_renorm_det_v3_ws_bytes()) + tag = "in-tree" + if lib is not None: + with open(lib, "rb") as f: + digest = hashlib.md5(f.read(), usedforsecurity=False).hexdigest()[:8] + tag = f"libm31_topp_det_v3.so md5 {digest}" + logger.info(f"[M3.1] TOPP_DET v3 private op engaged ({tag})") + if op is None: + op, _ = _load_op("minimax_m31_topp_renorm_det", "libm31_topp_det.so") + ws = ( + int(torch.ops.trtllm.minimax_m31_topp_renorm_det_ws_bytes()) + if op is not None + else 0 + ) + _STATE["op"], _STATE["ws_bytes"], _STATE["v3"] = op, ws, v3 + # v3's candidate-list radix passes (bitwise the same output) unless TOPPDET_T3 is disabled. + _STATE["cand_list"] = v3 and not m31_disabled("TOPPDET_T3") + if v3: + logger.info( + "[M3.1] TOPP_DET v3 T3 candidate-list fused passes: " + f"{'on' if _STATE['cand_list'] else 'off'} (cap 4096)" + ) + return _STATE["op"] + + +def is_v3() -> bool: + """Whether the engaged op is TOPP_DET v3.""" + topp_det_op() + return _STATE["v3"] + + +def ws_bytes_per_row() -> int: + """Workspace bytes per row (0 when the op is not built).""" + topp_det_op() + return _STATE["ws_bytes"] + + +def sm_count(device: torch.device) -> int: + """The device's SM count, cached per device.""" + idx = torch.device(device).index + idx = torch.cuda.current_device() if idx is None else idx + if idx not in _SMS: + _SMS[idx] = torch.cuda.get_device_properties(idx).multi_processor_count + return _SMS[idx] + + +def ctas_per_row(rows: int, sms: int) -> int: + """The largest power of two <= MAX_CTAS_PER_ROW with ctas x rows <= sms (1 past that: one CTA per row needs no + co-residency).""" + g = MAX_CTAS_PER_ROW + while g > 1 and g * rows > sms: + g //= 2 + return g + + +def prebuild(device: Optional[torch.device] = None) -> Optional[torch.Tensor]: + """Allocate (once) the device's zeroed workspace for MAX_ROWS rows; must run outside a CUDA-graph capture. + Returns it, or None when the op is not built or the stream is capturing.""" + if topp_det_op() is None: + return None + dev = ( + torch.device("cuda", torch.cuda.current_device()) + if device is None + else torch.device(device) + ) + if dev.index is None: + dev = torch.device("cuda", torch.cuda.current_device()) + if dev.index not in _WORKSPACES: + if torch.cuda.is_current_stream_capturing(): + return None + _WORKSPACES[dev.index] = torch.zeros( + MAX_ROWS * ws_bytes_per_row(), dtype=torch.uint8, device=dev + ) + return _WORKSPACES[dev.index] + + +class M31ToppDet(OpWrapper): + """FlashInfer's ``top_p_renorm_probs(probs, top_p, is_deterministic=True)`` of fp32 probs [rows, d] with fp32 + per-row top_p [rows], bit for bit on every row with top_p < 1, in ONE launch of rows x ctas_per_row CTAs + (FlashInfer's: five launches, apply on one CTA per row, slow on high-entropy rows). Rows with top_p >= 1 keep every + element (FlashInfer drops tail buckets there in ~43% of rows, in both of its modes); their output is + ``ftz(ftz(p) * norm)``. ``workspace`` is the persistent zeroed buffer of ``prebuild``. ``__call__`` returns None + (the caller keeps the stock call) for a non-fp32, non-contiguous or misaligned input, a width not divisible by 4, + a top_p that is not a contiguous fp32 [rows] tensor on probs' device, more rows than the workspace holds, or when + the op is not built.""" + + ARCHS = frozenset({Arch.SM_103}) + + # The c1 K7 calls were checked bitwise on GB300 over 8 row classes x 4 top_p modes x 16 seeds. + CELLS: tuple[Cell, ...] = ( + Cell(why="the K7 draft's 7 rows at c1 (vocab 200064)", spec=dict(rows=7, d=200064)), + Cell(why="the K7 verify's 8 rows at c1 (vocab 200064)", spec=dict(rows=8, d=200064)), + Cell(why="one row (the smallest grid, 16 CTAs)", spec=dict(rows=1, d=200064)), + ) + + note = """ + Exact by construction: the det histograms are integer (any aggregation order is FlashInfer's), the per-pass + selection is FlashInfer's last-block code on them, and the kept sum is FlashInfer's 1024 per-thread chains into + the same cub BlockReduce. Built with --use_fast_math, as FlashInfer's JIT builds it. + """ + + def __call__( + self, + probs: torch.Tensor, + top_p: torch.Tensor, + workspace: Optional[torch.Tensor], + ctas_per_row: int, + ) -> Optional[torch.Tensor]: + if ( + not isinstance(top_p, torch.Tensor) + or workspace is None + or not probs.is_cuda + or probs.dtype != torch.float32 + or probs.dim() != 2 + or not probs.is_contiguous() + or probs.shape[0] < 1 + or probs.shape[1] % 4 != 0 + or probs.data_ptr() % 16 != 0 + or top_p.dtype != torch.float32 + or top_p.dim() != 1 + or top_p.shape[0] != probs.shape[0] + or not top_p.is_contiguous() + or top_p.device != probs.device + or workspace.device != probs.device + or ctas_per_row < 1 + or 1024 % ctas_per_row != 0 + ): + return None + op = topp_det_op() + if ( + op is None + or workspace.numel() * workspace.element_size() < probs.shape[0] * ws_bytes_per_row() + ): + return None + out = torch.empty_like(probs) + if is_v3(): + op(probs, top_p, workspace, out, int(ctas_per_row), True, _STATE["cand_list"]) + else: + op(probs, top_p, workspace, out, int(ctas_per_row)) + return out + + def reference(self, probs, top_p, workspace=None, ctas_per_row=None) -> torch.Tensor: + """FlashInfer's deterministic renorm (rows with top_p < 1; workspace and ctas_per_row are ignored).""" + import flashinfer.sampling + + return flashinfer.sampling.top_p_renorm_probs(probs, top_p, is_deterministic=True) + + def is_valid(self, probs, top_p, workspace=None, ctas_per_row=None, **kwargs) -> None: + assert probs.dtype == torch.float32 and probs.dim() == 2, "probs: fp32 [rows, d]" + + +m31_topp_det = M31ToppDet() + + +def top_p_renorm_det(probs: torch.Tensor, top_p: torch.Tensor) -> Optional[torch.Tensor]: + """The sampler's top-p renorm override: the det renorm of probs [rows, d] at per-row top_p, or None (the sampler + runs stock FlashInfer) when the op is not built, rows exceed MAX_ROWS, no workspace can be had (first use inside + a capture) or the input is outside the contract.""" + if not isinstance(top_p, torch.Tensor) or not probs.is_cuda or probs.dim() != 2: + return None + rows = probs.shape[0] + if rows < 1 or rows > MAX_ROWS or topp_det_op() is None: + return None + workspace = prebuild(probs.device) + if workspace is None: + return None + return m31_topp_det(probs, top_p, workspace, ctas_per_row(rows, sm_count(probs.device))) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/spec/m3_capture_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/spec/m3_chunk_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/spec/m3_draft_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/spec/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/spec/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/spec/m3_sampled_accept_ws_words.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/spec/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/spec/m3_slot_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/spec/m3_step_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/ops/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/weights.py new file mode 100644 index 000000000000..7eaac1fb8d31 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/weights.py @@ -0,0 +1,910 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight tables and loaders: MiniMax-M3.1 (MXFP8 + packed NVFP4 experts) / 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.models.checkpoints.hf.minimaxm3_weight_mapper import ( + reciprocal_global_scale, +) +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 fused_projection(sparse_config) -> bool: + """Whether the sparse layers run the five-way packed QKV + index projection. + + On by default: this target requires the MSA implementation and an FP8 indexer, which the packed projection needs. + An explicit ``fuse_qkv_index_projection`` in the deployment's sparse attention config wins, so setting it to False + is the kill switch (the QKV projection plus a separate MXFP8 index projection). + """ + if sparse_config is None: + return False + # A config that does not record which fields were set (a plain namespace or dataclass) is taken at its word. + fields_set = getattr(sparse_config, "model_fields_set", None) + if fields_set is None or "fuse_qkv_index_projection" in fields_set: + return bool(getattr(sparse_config, "fuse_qkv_index_projection", False)) + return True + + +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 _divisor(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """A stored global scale (multiplier) -> the modelopt divisor: its FP32 reciprocal.""" + return reciprocal_global_scale("expert global scale", t) + + +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]`` as the checkpoint's MXFP8 (E4M3) values: this rank's index-Q heads and the replicated + index-K, the same rows the five-way projection packs after Q / K / V.""" + iq, ik = t + return torch.cat([_col(iq, d), _dev(ik, d)], dim=0).contiguous() + + +def _index_qk_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """The UE8M0 block scales of ``_index_qk_w``, 128x4-swizzled as every MXFP8 projection's.""" + iq, ik = t + return _swizzle(torch.cat([_col(iq, d), _dev(ik, d)], dim=0)) + + +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_mlp_layers(d): + """Layers with the dense SwiGLU-OAI MLP (``moe_layer_freq`` 0).""" + return d.dense_mlp + + +def _moe_layers(d): + return d.moe + + +def _sparse_attn_layers(d): + """Layers with MSA sparse attention (``sparse_attention_freq`` 1): every layer of MiniMax-M3.1.""" + return d.sparse_attn + + +def _three_way_qkv_layers(d): + """Layers whose QKV is the three-way [Q | K | V]: the dense-attention layers, and every layer without the + five-way packed projection (``fuse_qkv_index_projection=False``).""" + return d.dense_attn if d.fused else range(d.num_layers) + + +def _five_way_qkv_layers(d): + return d.sparse_attn if d.fused else range(0) + + +def _unfused_sparse_layers(d): + return range(0) if d.fused else d.sparse_attn + + +_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 index projection [index-Q | index-K] of the unfused sparse layers: the checkpoint's MXFP8 + # values with their block scales (the rows of the five-way projection's index part), run as an MXFP8 GEMM. + W( + "idx_w", + shape=lambda d: (d.index_heads * d.index_dim + d.index_dim, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(_ATTN + ".index_q_proj.weight", _ATTN + ".index_k_proj.weight"), + transform=_index_qk_w, + layers=_unfused_sparse_layers, + ), + W( + "idx_s", + shape=lambda d: ((d.index_heads * d.index_dim + d.index_dim) * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + _ATTN + ".index_q_proj.weight_scale_inv", _ATTN + ".index_k_proj.weight_scale_inv" + ), + transform=_index_qk_s, + layers=_unfused_sparse_layers, + ), + W( + "idx_q_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_q_norm.weight"), + layers=_sparse_attn_layers, + ), + W( + "idx_k_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_k_norm.weight"), + layers=_sparse_attn_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_mlp_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_mlp_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_mlp_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_mlp_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=_moe_layers, + ), + W( + "router_bias", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_one(_MOE + ".e_score_correction_bias"), + transform=_f32, + layers=_moe_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=_moe_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=_moe_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=_moe_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=_moe_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_packed", "w1.weight_packed"), + transform=_expert_w31, + layers=_moe_layers, + ), + W( + "e_w2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2), + dtype=torch.uint8, + src=_per_expert("w2.weight_packed"), + transform=_expert_w2, + layers=_moe_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=_moe_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=_moe_layers, + ), + # The per-tensor scalars of every expert: the checkpoint stores the multipliers into the quantized range + # (``*_global_scale``); the modelopt divisors modeling derives the kernels' scalars from are their FP32 + # reciprocals, taken by the same helper as the built-in loader. + *( + W( + f"e_{name}", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_per_expert(suffix), + transform=_divisor, + layers=_moe_layers, + ) + for name, suffix in ( + ("isc1", "w1.input_global_scale"), + ("isc3", "w3.input_global_scale"), + ("isc2", "w2.input_global_scale"), + ("ws2_1", "w1.weight_global_scale"), + ("ws2_3", "w3.weight_global_scale"), + ("ws2_2", "w2.weight_global_scale"), + ) + ), + # --- 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"]] + moe_freq = [ + int(f) for f in getattr(cfg, "moe_layer_freq", None) or [1] * cfg.num_hidden_layers + ] + 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=fused_projection(core.model_config.sparse_attention_config), + num_layers=cfg.num_hidden_layers, + sparse_attn=[i for i, f in enumerate(freq) if f], + dense_attn=[i for i, f in enumerate(freq) if not f], + moe=[i for i, f in enumerate(moe_freq) if f], + dense_mlp=[i for i, f in enumerate(moe_freq) if not f], + 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_m31_nvfp4__sm_103__tp8/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/__init__.py new file mode 100644 index 000000000000..2e340ff7284d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/__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 NVFP4 / sm_103 / tp8.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/dspark.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/dspark.py new file mode 100644 index 000000000000..71af3ca118cc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/dspark.py @@ -0,0 +1,829 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3.1 TP8 implementation.""" + +from __future__ import annotations + +import os +import time +from typing import Tuple + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm import ( + allgather_cast as _allgather_cast, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode import ( + strict_accept as _strict_accept, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_block_tables import ( + enable_parallel_refresh, + refresh_block_tables, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_ctx_advance import ( + advance_ctx_and_read_view, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_ctx_store import ( + fold_k_norm_weight, + fused_ctx_store, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_drafter_inputs import ( + drafter_inputs as _drafter_inputs, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_forward_indices import ( + forward_indices, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_qk_norm_rope import ( + qk_norm_rope, + qk_norm_rope_append, +) +from tensorrt_llm._torch.distributed import allgather +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_speculative import ( + external_drafter_config_kwargs, + get_draft_model, +) +from tensorrt_llm._torch.speculative.dspark import DSparkWorker +from tensorrt_llm._torch.speculative.interface import ( + _RNG_SLOT_SPAN, + should_use_separate_draft_kv_cache, + tail_bookkeep_draft_len, + tail_bookkeep_log_once, +) +from tensorrt_llm.logger import logger + +from .dspark_s4 import MAX_FAST_BATCH, TAILBK_SPEC, S4WorkerMixin +from .m31_switches import m31_disabled +from .mxdraft import route_mx_linears +from .ops.embedding.m3_embedding import m3_embedding + +MXDRAFT = not m31_disabled("MXDRAFT") +DRAFT_APPEND_FUSE = not m31_disabled("DRAFT_APPEND_FUSE") +DRAFT_ATTN_PDL = not m31_disabled("DRAFT_ATTN_PDL") +_APPEND_FUSE_LOGGED = [False] + + +def _qk_norm_rope_append_logged(*args): + """catalog qk_norm_rope_append, logging the engagement once per process on its first taken call (host only; a + captured step logs while it is captured).""" + out = qk_norm_rope_append(*args) + if out is not None and (not _APPEND_FUSE_LOGGED[0]): + _APPEND_FUSE_LOGGED[0] = True + logger.info( + "[M3 dispatch] DRAFT_APPEND_FUSE taken: drafter K / V append in the EX-11 launch" + ) + return out + + +def is_standalone_dspark(spec_config) -> bool: + """Whether the run speculates with a standalone DSpark drafter (its own checkpoint, ``decoding_type: DSpark``).""" + return bool( + spec_config is not None + and spec_config.spec_dec_mode.use_one_engine() + and spec_config.spec_dec_mode.is_dspark() + and (not getattr(spec_config, "draft_is_embedded_in_target", False)) + ) + + +class M3EmbedTokens(nn.Module): + """The target's token embedding as the module the in-tree drafter borrows (``load_weights_from_target_model``). + + The drafter embeds the bonus and mask tokens through ``embed_tokens(ids)``; the target keeps its replicated + table as a raw tensor and gathers it with ``m3_embedding``, so this module calls the same gather on the same + table (an exact gather, so the drafter sees the rows the target embeds). + """ + + def __init__(self, core: nn.Module) -> None: + super().__init__() + self._core = [core] # not a submodule: the table belongs to the target + + @property + def weight(self) -> torch.Tensor: + return self._core[0].w["embed"] + + def forward(self, ids: torch.Tensor) -> torch.Tensor: + return m3_embedding(ids, self.weight) + + +def s10c_c_drafter_inputs( + embed_tokens, + accepted, + num_accepted, + slots, + ctx_len, + max_ctx, + mask_token_id, + trained_mask_embedding, + block_size, + tokens_per_request, +): + """S10c-C: the drafter's step inputs in one launch (catalog/spec_decode/dspark_drafter_inputs), for this target's + embedding only: ``M3EmbedTokens`` is an exact gather of the replicated table, so the noise block is table rows (the + trained mask row if the drafter keeps one). None for any other embedding: the stock path runs.""" + if not isinstance(embed_tokens, M3EmbedTokens): + return None + table = embed_tokens.weight + mask_row = ( + trained_mask_embedding if trained_mask_embedding is not None else table[int(mask_token_id)] + ) + return _drafter_inputs( + accepted, + num_accepted, + slots, + ctx_len, + max_ctx, + table, + mask_row, + block_size, + tokens_per_request, + ) + + +N3A_DISABLED = m31_disabled("N3A") +N3A_KV_DISABLED = N3A_DISABLED or m31_disabled("N3A_KV") +N3A_KV_CHECK = os.environ.get("TRTLLM_M31_CHECK_N3A_KV", "").strip() not in ("", "0") +N3P_DISABLED = m31_disabled("N3P") +A1_DISABLED = m31_disabled("A1") +A1_MAX_ROWS = 64 +MCAG = not m31_disabled("MCAG") +A1_K7_DISABLED = m31_disabled("A1_K7") +_n3a_logged = False + + +def _n3a_log_once() -> None: + global _n3a_logged + if not _n3a_logged: + _n3a_logged = True + logger.info( + "[M3 dispatch] n3_strict_accept taken: fused strict acceptance + next-step input tokens" + ) + + +_n3p_logged = False +_n3a_kv_logged = False + + +def _n3a_kv_log_once() -> None: + global _n3a_kv_logged + if not _n3a_kv_logged: + _n3a_kv_logged = True + logger.info( + "[M3 dispatch] n3a_kv_rewind taken: draft-forward KV rewind + kv_lens bump in one launch" + ) + + +_a1_logged = set() + + +def _a1_log_once(key, message: str) -> None: + if key not in _a1_logged: + _a1_logged.add(key) + logger.info(message) + + +def _n3p_log_once() -> None: + global _n3p_logged + if not _n3p_logged: + _n3p_logged = True + logger.info( + "[M3 dispatch] n3p_target_logits taken: verify-row logits without the gather_ids copy and the concat" + ) + + +class M3DSparkWorker(S4WorkerMixin, DSparkWorker): + """The in-tree ``DSparkWorker`` with the calls this target's shell makes on its worker. W1: the one-launch Markov + chain (``S4WorkerMixin``, switch ``TRTLLM_M31_DRAFT_DISABLE_S4``). N3 A2: the strict acceptance and the next-step + input tokens in one launch each (switch ``TRTLLM_M31_DISABLE_N3A``). TAILBK R1 / Z1: the one-request Philox rows + from the per-token window and no padding-row zeroing in captured batch-1 graphs, bitwise (spec_metadata's + tail_bookkeep; switch ``TRTLLM_M31_DISABLE_TAILBK``). Everything else is the in-tree path.""" + + def capture_hidden_states( + self, spec_metadata, layer_id: int, hidden_states, residual=None + ) -> None: + """The in-tree DFlash capture: layer output (hidden + residual) into the static capture buffer.""" + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + + def arm_capture_norm( + self, attn_metadata, spec_metadata, lm_head: nn.Module, draft_model + ) -> None: + """No capture norm: DSpark's fc reads the raw captured layers.""" + return None + + def tail_eligible(self, attn_metadata, spec_metadata, lm_head: nn.Module) -> bool: + """Never: the target logits always go through the logits processor and the in-tree strict accept.""" + return False + + _a1 = None + + def prebuild(self, lm_head: nn.Module, draft_model) -> None: + """The one-launch Markov chain's shards, buffers, exchange workspace and compiles, then the N3 A2 compiles + (rank-local) and the N3 A1 workspace and compiles (collective), outside capture, every rank, in this order; + no fused verify, no coded target head, no fused draft loop.""" + rank = getattr(getattr(self, "mapping", None), "tp_rank", None) + lean = not m31_disabled("S10") + self.enable_lean_glue(draft_model, lean) + k = getattr(getattr(self, "spec_config", None), "max_draft_len", 0) or 0 + logger.info( + f"M3.1 DSpark S10: lean drafter glue {('on' if lean else 'off')} on rank {rank}" + ) + s10c_a = self.enable_fused_ctx_store( + None if m31_disabled("S10C") or m31_disabled("S10C_A") else fused_ctx_store + ) + tailbk_f = fold_k_norm_weight(bool(s10c_a) and TAILBK_SPEC) + logger.info( + ( + "M3.1 DSpark S10c-A: fused drafter ctx store " + f"{('on' if s10c_a else 'off')}" + "; TAILBK_SPEC F k_norm weight in the store kernel " + f"{('on' if tailbk_f else 'off')}" + " on rank " + f"{rank}" + ) + ) + s10c_c = self.enable_fused_drafter_inputs( + None if m31_disabled("S10C") or m31_disabled("S10C_C") else s10c_c_drafter_inputs + ) + logger.info( + f"M3.1 DSpark S10c-C: fused drafter step inputs {('on' if s10c_c else 'off')} on rank {rank}" + ) + s10c_d = self.enable_fused_ctx_advance( + None if m31_disabled("S10C") or m31_disabled("S10C_D") else advance_ctx_and_read_view + ) + logger.info( + f"M3.1 DSpark S10c-D: fused drafter context advance {('on' if s10c_d else 'off')} on rank {rank}" + ) + s10c_b = self.enable_fused_block_refresh( + None if m31_disabled("S10C") or m31_disabled("S10C_B") else refresh_block_tables + ) + s10c_b2 = enable_parallel_refresh(s10c_b and (not m31_disabled("S10C_B2"))) + logger.info( + ( + "M3.1 DSpark S10c-B: fused drafter block-table refresh " + f"{('on' if s10c_b else 'off')}" + " on rank " + f"{rank}" + " (S10c-B2 2-D grid " + f"{('on' if s10c_b2 else 'off')}" + ")" + ) + ) + s10c_e = self.enable_fused_forward_indices( + draft_model, None if m31_disabled("S10C") or m31_disabled("S10C_E") else forward_indices + ) + logger.info( + f"M3.1 DSpark S10c-E: fused drafter forward indices {('on' if s10c_e else 'off')} on rank {rank}" + ) + ex11 = self.enable_fused_qk_norm_rope( + draft_model, None if m31_disabled("EX11") else qk_norm_rope + ) + logger.info( + f"M3.1 DSpark EX-11: fused drafter q/k norm + RoPE {('on' if ex11 else 'off')} on rank {rank}" + ) + append_fuse = self.enable_fused_qk_norm_rope_append( + draft_model, _qk_norm_rope_append_logged if ex11 and DRAFT_APPEND_FUSE else None + ) + logger.info( + ( + "M3.1 DSpark DRAFT_APPEND_FUSE: drafter K / V append in the EX-11" + " launch " + f"{('on' if append_fuse else 'off')}" + " on rank " + f"{rank}" + ) + ) + attn_pdl = self.enable_attention_pdl(draft_model, DRAFT_ATTN_PDL) + logger.info( + f"M3.1 DSpark DRAFT_ATTN_PDL: drafter attention PDL {('on' if attn_pdl else 'off')} on rank {rank}" + ) + mx = ( + 0 + if not MXDRAFT or k <= 0 + else route_mx_linears(draft_model, [b * k for b in range(1, MAX_FAST_BATCH + 1)]) + ) + logger.info(f"M3.1 DSpark MXDRAFT: MXFP8 drafter GEMV on {mx} projections on rank {rank}") + self.prebuild_s4(lm_head, draft_model) + if not N3A_DISABLED and lm_head is not None: + _strict_accept.precompile((self.max_draft_len,), lm_head.weight.device) + self._a1 = self._prebuild_a1(lm_head) if lm_head is not None else None + self._mcag_cols = self._prebuild_mcag(lm_head) if self._a1 is not None and MCAG else 0 + mcag = self.enable_draft_logits_gather_op(self._mcag_gather if self._mcag_cols else None) + logger.info( + f"M3.1 MCAG: draft-logits all-gather on the multicast kernel {('on' if mcag else 'off')} on rank {rank}" + ) + n3a_kv = False + if not N3A_KV_DISABLED and lm_head is not None: + _strict_accept.precompile_kv_rewind(lm_head.weight.device) + self._n3a_kv_buf = torch.empty( + _strict_accept.max_requests(), dtype=torch.int32, device=lm_head.weight.device + ) + n3a_kv = self.enable_fused_kv_rewind(self._n3a_kv_rewind) + logger.info( + ( + "M3.1 N3A-KV: one-launch draft-forward KV bookkeeping " + f"{('on' if n3a_kv else 'off')}" + " on rank " + f"{getattr(getattr(self, 'mapping', None), 'tp_rank', None)}" + ) + ) + + _n3a_kv_buf = None + _n3a_kv_counts = None + + def _n3a_kv_rewind(self, num_accepted, kv_lens, nc: int, bs: int): + """N3A-KV for the DFlash hook: ``1 - num_accepted[nc:bs]`` (a view of the persistent buffer) with + ``kv_lens[nc:bs] += 1`` in one launch, or None (the stock lines) for operands or a variant it does not take.""" + n = bs - nc + buf = self._n3a_kv_buf + if buf is None or not _strict_accept.supports_kv_rewind(num_accepted, kv_lens, buf, nc, n): + return None + if not _strict_accept.is_compiled_kv_rewind(num_accepted, kv_lens): + return None + _n3a_kv_log_once() + if not N3A_KV_CHECK: + return _strict_accept.spec_kv_rewind(num_accepted, kv_lens, buf, nc, n) + want_rewind = 1 - num_accepted[nc:bs] + want_kv = kv_lens[nc:bs] + 1 + out = _strict_accept.spec_kv_rewind(num_accepted, kv_lens, buf, nc, n) + if self._n3a_kv_counts is None: + self._n3a_kv_counts = torch.zeros(2, dtype=torch.int64, device=kv_lens.device) + bad = (out != want_rewind).any() | (kv_lens[nc:bs] != want_kv).any() + self._n3a_kv_counts[0] += 1 + self._n3a_kv_counts[1] += bad.to(torch.int64) + if not torch.cuda.is_current_stream_capturing(): + calls, mism = self._n3a_kv_counts.tolist() + logger.info(f"M3.1 N3A-KV check: {calls} calls, {mism} mismatches") + return out + + def _a1_all_ok(self, local: bool) -> bool: + """Whether ``local`` holds on every rank of the TP group: A1 is set up on every rank or on none (a rank that + is not eligible, or whose compile failed, turns it off everywhere).""" + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _mnnvl_workspace_all_succeeded, + ) + + return _mnnvl_workspace_all_succeeded(_get_mnnvl_workspace_comm(self.mapping), local) + + def _prebuild_a1(self, lm_head: nn.Module): + """The N3 A1 workspace and its K3 verify-row compiles, or None (every rank alike): off by its switch, at + another draft length, without a column-parallel bf16 head over a TP group, or when a compile fails on any + rank.""" + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + mapping = self.mapping + if mapping.tp_size <= 1 or mapping.enable_attention_dp: + return None + cols = int(getattr(lm_head, "out_features", 0)) + sms = torch.cuda.get_device_properties(lm_head.weight.device).multi_processor_count + off = [ + why + for why, bad in ( + ("TRTLLM_M31_DISABLE_A1 set", A1_DISABLED), + ( + f"draft length {int(self.max_draft_len)} with TRTLLM_M31_DISABLE_A1_K7 set", + int(self.max_draft_len) != 3 and A1_K7_DISABLED, + ), + ( + f"draft length {int(self.max_draft_len)}: more than {A1_MAX_ROWS} verify rows at batch 1", + int(self.max_draft_len) + 1 > A1_MAX_ROWS, + ), + ( + "not a column-parallel, gathered, unpadded, bias-free head", + getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN + or not getattr(lm_head, "gather_output", False) + or getattr(lm_head, "padding_size", 0) != 0 + or (getattr(lm_head, "bias", None) is not None), + ), + (f"head dtype {lm_head.weight.dtype}", lm_head.weight.dtype != torch.bfloat16), + ( + f"{cols} shard columns (whole 16-byte vectors needed)", + cols <= 0 or cols % 8 != 0, + ), + ( + f"{sms} SMs (fewer than {_allgather_cast.MAX_GRID})", + sms < _allgather_cast.MAX_GRID, + ), + ) + if bad + ] + if not self._a1_all_ok(not off): + logger.info( + ( + "M3.1 N3 A1: off on rank " + f"{mapping.tp_rank}" + " (" + f"{'; '.join(off) or 'off on another rank'}" + "); the verify logits use the NCCL all-gather" + ) + ) + return None + k1 = int(self.max_draft_len) + 1 + rows = tuple((b * k1 for b in range(1, A1_MAX_ROWS // k1 + 1))) + start = time.perf_counter() + if not _allgather_cast.precompile( + lm_head.mapping, cols, rows, lm_head.weight.device, agree=self._a1_all_ok + ): + logger.warning( + "M3.1 N3 A1: a compile failed on some rank; every rank keeps the NCCL all-gather" + ) + return None + ws = _allgather_cast.workspace_args(lm_head.mapping, cols) + logger.info( + ( + "M3.1 N3 A1: one-launch verify all-gather set up on rank " + f"{mapping.tp_rank}" + "/" + f"{mapping.tp_size}" + " (shard " + f"{cols}" + " columns, rows " + f"{rows}" + ", workspace " + f"{ws['ws_uc'].numel() * 4}" + " bytes/rank, set-up " + f"{time.perf_counter() - start:.2f}" + " s)" + ) + ) + return ws + + def _prebuild_mcag(self, lm_head: nn.Module) -> int: + """Compile the fp32 draft-logits gather for B * K rows up to its capacity, on every rank (A1's workspace and + agreement); its shard columns, or 0 (every rank alike) when a compile fails anywhere.""" + cols = int(getattr(lm_head, "out_features", 0)) + k = int(self.max_draft_len) + if k <= 0 or cols % 4 != 0: + return 0 + rows = tuple((b * k for b in range(1, _allgather_cast.MAX_ROWS_FP32 // k + 1))) + if not _allgather_cast.precompile_f32( + lm_head.mapping, cols, rows, lm_head.weight.device, agree=self._a1_all_ok + ): + logger.warning( + "M3.1 MCAG: a compile failed on some rank; every rank keeps the NCCL all-gather" + ) + return 0 + return cols + + def _mcag_gather(self, logits: torch.Tensor): + """The draft logits gathered in one launch on the A1 workspace, or None (the stock all-gather) for a shape + that is not compiled; the decision depends on the shape only, the same on every rank.""" + rows = logits.numel() // max(int(logits.shape[-1]), 1) + if ( + logits.dtype != torch.float32 + or int(logits.shape[-1]) != self._mcag_cols + or ( + not _allgather_cast.is_compiled( + rows, self._mcag_cols, self._a1["world"], device=logits.device, fp32=True + ) + ) + ): + _a1_log_once( + ("mcag", rows, logits.dtype), + f"[M3 dispatch] MCAG draft-logits all-gather not taken: {rows} rows {logits.dtype}; NCCL all-gather", + ) + return None + _a1_log_once( + "mcag", + "[M3 dispatch] MCAG draft-logits all-gather taken: one launch on the A1 workspace", + ) + return _allgather_cast.gather_last_dim_f32(logits.contiguous(), **self._a1) + + def _n3a_eligible( + self, logits, draft_tokens, num_contexts: int, batch_size: int, spec_metadata + ) -> bool: + """Whether this step's strict acceptance takes the fused op: a decode-only, non-ragged step whose logits hold + exactly batch_size x (D + 1) verify rows and whose drafts the op takes. Host state and shapes only, so it is + fixed per captured graph and the same on every rank. Decided before the target tokens are sampled, so the + sampler runs exactly once either way.""" + if N3A_DISABLED or spec_metadata is None or spec_metadata.is_ragged_verify: + return False + if num_contexts != 0 or batch_size <= 0 or draft_tokens.dim() != 2: + return False + n_req, draft_len = draft_tokens.shape + if n_req != batch_size or draft_len <= 0 or n_req > _strict_accept.max_requests(): + return False + rows = logits.shape[0] if logits.dim() > 1 else 1 + return rows == batch_size * (draft_len + 1) and draft_tokens.is_contiguous() + + def _sample_and_accept_draft_tokens_base( + self, logits, draft_tokens, num_contexts, batch_size, spec_metadata + ): + """The in-tree strict acceptance with its epilogue in one launch (N3 A2): the same target tokens (the in-tree + sampler, called once), the same accepted tokens and counts, then the in-tree forced-acceptance override (its + RNG pool and counters advance as on an in-tree step).""" + if not self._n3a_eligible(logits, draft_tokens, num_contexts, batch_size, spec_metadata): + return super()._sample_and_accept_draft_tokens_base( + logits, draft_tokens, num_contexts, batch_size, spec_metadata + ) + if logits.dim() == 1: + logits = logits.unsqueeze(0) + target_tokens = self._sample_tokens_for_batch( + logits, spec_metadata, num_contexts, batch_size + ) + if _strict_accept.supports_strict_accept( + target_tokens, draft_tokens + ) and _strict_accept.is_compiled_strict_accept(target_tokens, draft_tokens): + accepted, num_accepted = _strict_accept.spec_strict_accept(target_tokens, draft_tokens) + _n3a_log_once() + else: + accepted, num_accepted = _strict_accept.spec_strict_accept.reference( + target_tokens, draft_tokens + ) + num_accepted = self._apply_force_accepted_tokens( + num_accepted, num_contexts, draft_tokens.shape[-1], spec_metadata=spec_metadata + ) + return (accepted, num_accepted) + + def target_logits(self, hidden_states, attn_metadata, spec_metadata, lm_head: nn.Module): + """The verify rows' gathered fp32 target logits (N3 P-gid + P-cat), or None for the logits processor. + + On a decode-only, non-ragged step every row is a verify row and ``gather_ids`` is ``arange(num_tokens)`` (each + generation request extends it with the next ``1 + D`` row indices), so the rows are a slice, not a gather. + The processor's head is the column-parallel BF16 GEMM, an all-gather along the vocab (NCCL, then a concat) + and ``.float()``; here the same all-gather runs along dim 0 (no concat) and one strided copy writes the + rank-major shards into the fp32 output: the same bytes, an exact widening. Shapes and host state only, so + the decision is fixed per captured graph and the same on every rank.""" + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + if ( + N3P_DISABLED + or spec_metadata is None + or getattr(spec_metadata, "is_ragged_verify", False) + ): + return None + if attn_metadata.num_contexts != 0 or spec_metadata.gather_ids is None: + return None + rows = hidden_states.shape[0] + if spec_metadata.gather_ids.shape[0] != rows or rows == 0: + return None + if ( + getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN + or not getattr(lm_head, "gather_output", False) + or getattr(lm_head, "padding_size", 0) != 0 + or (getattr(lm_head, "bias", None) is not None) + ): + return None + shard = lm_head.apply_linear(hidden_states, None) + _n3p_log_once() + tp = lm_head.tp_size + if tp == 1: + return shard.float() + if self._a1 is not None: + if rows <= A1_MAX_ROWS and _allgather_cast.is_compiled( + rows, shard.shape[-1], tp, 0, shard.device + ): + _a1_log_once( + "taken", + "[M3 dispatch] n3_a1 allgather_cast taken: verify logits gathered in one launch", + ) + return _allgather_cast.allgather_cast(shard, **self._a1) + _a1_log_once( + rows, + ( + "[M3 dispatch] n3_a1 allgather_cast not taken: " + f"{rows}" + " verify rows (compiled: verify rows up to " + f"{A1_MAX_ROWS}" + "); NCCL all-gather" + ), + ) + gathered = allgather(shard, lm_head.mapping, dim=0) + vocab_shard = shard.shape[-1] + out = torch.empty((rows, tp * vocab_shard), dtype=torch.float32, device=shard.device) + out.view(rows, tp, vocab_shard).copy_(gathered.view(tp, rows, vocab_shard).transpose(0, 1)) + return out + + def _prepare_next_new_tokens( + self, + accepted_tokens, + next_draft_tokens, + batch_indices_cuda, + batch_size, + num_accepted_tokens, + ): + """The next step's input tokens in one launch (N3 A2), else the in-tree index + concat.""" + batch_idx = batch_indices_cuda[:batch_size] + if ( + not N3A_DISABLED + and next_draft_tokens.dim() == 2 + and (next_draft_tokens.shape[0] == batch_size) + and _strict_accept.supports_next_new_tokens( + accepted_tokens, num_accepted_tokens, batch_idx, next_draft_tokens + ) + and _strict_accept.is_compiled_next_new_tokens( + accepted_tokens, num_accepted_tokens, batch_idx, next_draft_tokens + ) + ): + return _strict_accept.spec_next_new_tokens( + accepted_tokens, num_accepted_tokens, batch_idx, next_draft_tokens + ) + return super()._prepare_next_new_tokens( + accepted_tokens, next_draft_tokens, batch_indices_cuda, batch_size, num_accepted_tokens + ) + + _tailbk_win = None + + def _forward_impl(self, *args, **kwargs): + """The S4 / DFlash step; R1's offsets window lives for one step only (set by its accept call, read by its + draft block), so a step whose accept did not stash one recomputes it.""" + self._tailbk_win = None + try: + return super()._forward_impl(*args, **kwargs) + finally: + self._tailbk_win = None + + def _rng_state_per_request(self, spec_metadata, start=0, end=None, repeat=1, *, slot): + """TAILBK R1: at one decode request, the accept and draft-block Philox (seed, offset) rows read from the + per-token buffers (one int64 add per step instead of eight launches), else the in-tree expansion. + + ``_populate_request_rng_state`` writes row r (0..K) of the request's per-token window as + ``seeds[r] = seed`` and ``offsets[r] = o + r * SPAN`` (o = ``request_offsets[0]``; the host sum fits int64, + or its upload would have raised), uploaded with ``request_seeds`` / ``request_offsets`` on every sampled step. + The in-tree rows are ``(o + slot * SPAN) + k * SPAN``: draft row k at slot K + 1 is ``offsets[k] + (K + 1) * + SPAN`` and the accept row at slot 2K + 1 is ``offsets[K] + (K + 1) * SPAN``. Integer addition is associative + and commutative mod 2^64, so the window ``offsets[0:K + 1] + (K + 1) * SPAN`` holds exactly the in-tree + offsets, even where they wrap, and the seeds are the request's seed by construction. The returned tensors are + contiguous int64 views with one entry per row, as the in-tree ones. Preconditions are host state and shapes + (``tail_bookkeep_draft_len``), so the choice is fixed per captured graph.""" + state = self._tailbk_rng_window(spec_metadata, start, end, repeat, slot) + if state is not None: + return state + return super()._rng_state_per_request(spec_metadata, start, end, repeat, slot=slot) + + def _tailbk_rng_window(self, spec_metadata, start, end, repeat: int, slot: int): + """R1's (seed, offset) rows, or None (the in-tree expansion) outside its preconditions.""" + k = tail_bookkeep_draft_len(spec_metadata, start, end) + if k == 0: + return None + first = spec_metadata.rng_draft_slot(0) + if first != k + 1 or spec_metadata.rng_accept_slot != first + k: + return None + accept = repeat == 1 and slot == spec_metadata.rng_accept_slot + block = repeat == k and slot == first + if not (accept or block): + return None + stash = self._tailbk_win + if block and stash is not None and (stash[0] is spec_metadata): + win = stash[1] + else: + win = spec_metadata.offsets[: k + 1] + first * _RNG_SLOT_SPAN + tail_bookkeep_log_once( + "R1", "Philox seeds / offsets from the per-token window (one add per step)" + ) + if accept: + self._tailbk_win = (spec_metadata, win) + return (spec_metadata.request_seeds[0:1], win[k : k + 1]) + self._tailbk_win = None + return (spec_metadata.seeds[:k], win[:k]) + + def _zero_padding_rows( + self, + logits, + spec_metadata, + num_contexts: int, + batch_size: int, + rows_per_request: int, + *, + packed_verify: bool = False, + ): + """TAILBK Z1: no padding-row zeroing while a batch-1 decode graph is captured, else the in-tree zeroing. + + Why the in-tree result is a bitwise copy of ``logits`` on every replay of such a graph: + (i) a batch-1 graph replays exactly one real request. ``CUDAGraphRunner._get_padded_batch`` only appends + dummies after the batch's own requests, up to the graph's size, and the executor forwards no empty batch + without attention DP (``_can_queue``: batch_size > 0), so a size-1 graph never holds a dummy. Under + attention DP an empty rank does run a dummy, so attention DP is excluded. + (ii) a real request's slot is its ``py_seq_slot`` in ``[0, slot_capacity)`` and the padding marker is + ``dummy_slot_row == slot_capacity`` (``prepare_rejection_sampling_buffers``), so the padding mask is + all-false and the in-tree ``masked_fill`` copies ``logits`` unchanged. + (iii) every consumer is out-of-place (the non-fused probs: the softmax / top-k mask ops and the m31_softmax_mc + override allocate their outputs, top-p reads probs), so returning ``logits`` itself instead of that copy + exposes no write to the caller's verify logits. The fused modes are not audited and keep the copy. + (iv) the consumers see the layout the copy would have: the skip needs contiguous logits at a 256-byte aligned + address (the copy is a fresh contiguous allocation, which the caching allocator aligns at least that much), + since the m31_softmax_mc override takes contiguous, 16-byte aligned inputs only. + Eager steps (the graph warmup's dummy included) keep the in-tree zeroing; so do padded graphs (2..8), which + can hold dummies, and ragged verification (its layout check). Host state only, fixed per captured graph.""" + if self._tailbk_skips_padding_rows(logits, spec_metadata, num_contexts, batch_size): + tail_bookkeep_log_once("Z1", "padding-row zeroing skipped in captured batch-1 graphs") + return logits + return super()._zero_padding_rows( + logits, + spec_metadata, + num_contexts, + batch_size, + rows_per_request, + packed_verify=packed_verify, + ) + + def _tailbk_skips_padding_rows( + self, logits, spec_metadata, num_contexts: int, batch_size: int + ) -> bool: + """Z1's exact condition: tail bookkeeping on, one generation request and no context rows (num_gens == 1), + no ragged verification, a non-fused sampling mode, no attention DP, contiguous 256-byte aligned logits and a + capturing stream.""" + if not getattr(spec_metadata, "tail_bookkeep", False): + return False + if not logits.is_contiguous() or logits.data_ptr() % 256 != 0: + return False + if num_contexts != 0 or batch_size != 1 or getattr(spec_metadata, "is_ragged_verify", True): + return False + mode = getattr(spec_metadata, "advanced_sampling_mode", None) + if mode is None or mode.is_fused: + return False + mapping = getattr(self, "mapping", None) + if mapping is None or getattr(mapping, "enable_attention_dp", True): + return False + return torch.cuda.is_current_stream_capturing() + + +def build_dspark( + shell, model_config: ModelConfig +) -> Tuple[ModelConfig, nn.Module, M3DSparkWorker, bool]: + """The drafter config, the drafter and the worker of one-model standalone DSpark (mirrors + ``SpecDecOneEngineForCausalLM`` for an external drafter).""" + spec = model_config.spec_config + if not is_standalone_dspark(spec): + raise ValueError("build_dspark serves the standalone DSpark drafter only") + draft_config = ModelConfig.from_pretrained( + spec.speculative_model, **external_drafter_config_kwargs(model_config, spec) + ) + keep_drafter_kv_dtype(draft_config) + draft_config.extra_attrs = model_config.extra_attrs + use_separate = should_use_separate_draft_kv_cache(spec) + draft = get_draft_model(model_config, draft_config, shell.lm_head, shell.model) + draft.logits_processor = shell.logits_processor + worker = M3DSparkWorker(spec, model_config.mapping, use_separate_draft_kv_cache=use_separate) + worker.set_draft_model(draft) + shell.epilogue.append(draft) + shell.epilogue.append(worker) + logger.info( + ( + "M3.1 target DSpark: " + f"{type(worker).__name__}" + " over " + f"{type(draft).__name__}" + ", K=" + f"{spec.max_draft_len}" + ", separate draft KV cache=" + f"{use_separate}" + ) + ) + return (draft_config, draft, worker, use_separate) + + +def load_dspark_draft_weights(shell, weights, weight_mapper=None) -> None: + """The drafter's own checkpoint, then the target's embedding and head (mirrors + ``SpecDecOneEngineForCausalLM.load_draft_weights`` for DSpark). + + The drafter borrows the target's embedding and head, and ``M3EmbedTokens.weight`` IS the target's table, so a + drafter checkpoint that ships either would load into a module the target owns or be shadowed by it: refused. + """ + borrowed = sorted( + (k for k in weights.keys() if "embed_tokens." in k or k.split(".")[-2:-1] == ["lm_head"]) + ) + if borrowed: + raise ValueError( + f"the DSpark drafter checkpoint ships tensors this target lends it (embedding / head): {borrowed[:4]}" + ) + shell.draft_model.load_weights(weights=weights, weight_mapper=weight_mapper) + shell.draft_model.load_weights_from_target_model(shell) + if ( + shell.draft_model.lm_head is not shell.lm_head + or shell.draft_model.draft_model_full.model.embed_tokens is not shell.model.embed_tokens + ): + raise RuntimeError("the DSpark drafter does not borrow this target's embedding and head") + own = sorted( + (n for n, _ in shell.named_parameters() if not n.startswith(("model.w.", "draft_model."))) + ) + logger.info(f"M3.1 target DSpark: shell parameters outside the core and the drafter: {own}") + + +def keep_drafter_kv_dtype(draft_config: ModelConfig) -> None: + """Drop any KV quantization from the drafter's config, so its context pool is allocated in its projection dtype. + + The one-engine construction copies the target's KV algo onto the drafter; on this target that is NVFP4, which + allocates a packed int8 pool the worker refuses to bind (it falls back to its private arena, and a reused + prefix's drafter K/V is lost). The cached quant modes are dropped with it, as the draft KV budget does. + """ + for quant in [draft_config.quant_config, *(draft_config.quant_config_dict or {}).values()]: + quant.kv_cache_quant_algo = None + quant.__dict__.pop("quant_mode", None) + quant.__dict__.pop("layer_quant_mode", None) + + +def draft_kv_override_applies(spec_config) -> bool: + """The target's FP8 draft-KV override is an Eagle3 need; a DSpark drafter keeps its BF16 context pool.""" + return not is_standalone_dspark(spec_config) + + +__all__ = [ + "M3DSparkWorker", + "M3EmbedTokens", + "build_dspark", + "draft_kv_override_applies", + "is_standalone_dspark", + "keep_drafter_kv_dtype", + "load_dspark_draft_weights", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/dspark_s4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/dspark_s4.py new file mode 100644 index 000000000000..79bb6aa118e3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/dspark_s4.py @@ -0,0 +1,477 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3.1 TP8 implementation.""" + +from __future__ import annotations + +import atexit +import os +from typing import Dict, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode import ( + dspark_markov_bias as markov_bias, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode import ( + dspark_markov_chain as chain, +) +from tensorrt_llm._torch.models.modeling_speculative import dspark_markov_step_bias +from tensorrt_llm.logger import logger + +from .m31_switches import m31_disabled + +MARKOV_RANK = 256 +MAX_FAST_BATCH = 8 +DISABLE_ENV = "TRTLLM_M31_DRAFT_DISABLE_S4" +S4_COUNT_LOG_EVERY = 1000 +MARKOV_LL = not m31_disabled("MARKOV_LL") +TAILBK_SPEC = not m31_disabled("TAILBK_SPEC") +_DISABLED = os.environ.get(DISABLE_ENV, "").strip() not in ("", "0") + + +def disabled() -> bool: + return _DISABLED + + +class _S4State: + def __init__( + self, + mapping, + rows: int, + shard: int, + offset: int, + w1: torch.Tensor, + w2s: torch.Tensor, + block: int, + ws: dict, + ) -> None: + self.mapping = mapping + self.ws = ws + self.rows = rows + self.shard = shard + self.offset = offset + self.w1 = w1 + self.w2s = w2s + self.block = block + self.bufs: Dict[Tuple[int, int], torch.Tensor] = {} + + def buffer(self, batch: int) -> torch.Tensor: + return self.bufs[batch, self.block] + + +S4_LOGITS_BATCHED = True + + +def _biased_shard_logits( + shard_logits: torch.Tensor, + first_prev: torch.Tensor, + tokens: torch.Tensor, + w1: torch.Tensor, + w2_shard: torch.Tensor, + fp32: bool = True, +) -> torch.Tensor: + """The in-tree chain's biased block logits on this rank's vocab shard, given the chain's tokens: step k adds + bf16(w2 @ w1[prev_k]) with prev_0 = the anchor and prev_k = token k-1 (dspark_markov_step_bias). With ``fp32`` + (the in-tree arithmetic: LogitsProcessor hands the chain fp32 logits) the bf16 shard logits and the bias are + widened and added in fp32; without it, one bf16 add.""" + if ( + fp32 + and S4_LOGITS_BATCHED + and TAILBK_SPEC + and markov_bias.supports(shard_logits, first_prev, tokens, w1) + ): + B, K, N = shard_logits.shape + bias = F.linear(markov_bias.prev_embeddings(first_prev, tokens, w1), w2_shard) + return markov_bias.widen_add(shard_logits, bias.view(B, K, N)) + if fp32: + shard_logits = shard_logits.float() + B, K, N = shard_logits.shape + prevs = torch.cat([first_prev.view(B, 1), tokens[:, : K - 1].long()], dim=1) + if S4_LOGITS_BATCHED: + bias = dspark_markov_step_bias(prevs.reshape(-1), w1, w2_shard).view(B, K, N) + return shard_logits + bias.to(shard_logits.dtype) + out = torch.empty_like(shard_logits) + for k in range(K): + out[:, k] = shard_logits[:, k] + dspark_markov_step_bias(prevs[:, k], w1, w2_shard).to( + shard_logits.dtype + ) + return out + + +class _ShardedDraftLogits(nn.Module): + """The step's draft logits processor: this rank's vocab shard of ``lm_head`` over the gathered block rows, + copied into the padded buffer; any other call goes to the drafter's own processor.""" + + def __init__(self, inner: nn.Module, state: _S4State, batch: int) -> None: + super().__init__() + self.inner = inner + self._state = [state] + self.batch = batch + + def forward(self, hidden_states, lm_head, attn_metadata, return_context_logits: bool = False): + state = self._state[0] + if hidden_states.dim() != 2 or hidden_states.shape[0] != self.batch * state.block: + return self.inner(hidden_states, lm_head, attn_metadata, return_context_logits) + logits = lm_head.apply_linear(hidden_states, None) + buf = state.buffer(self.batch) + buf[:, :, : state.shard].copy_(logits.view(self.batch, state.block, state.shard)) + return logits + + +def _head_is_sharded(lm_head: nn.Module, vocab: int, tp: int) -> bool: + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + return ( + tp > 1 + and getattr(lm_head, "tp_mode", None) == TensorParallelMode.COLUMN + and (getattr(lm_head, "padding_size", 0) == 0) + and (getattr(lm_head, "bias", None) is None) + and (lm_head.out_features * tp == vocab) + and (lm_head.weight.dtype == torch.bfloat16) + ) + + +class S4WorkerMixin: + """Mixed into the target's ``DSparkWorker`` subclass (before it in the bases).""" + + _s4: Optional[_S4State] = None + _s4_step: Optional[_S4State] = None + _s4_tokens: Optional[torch.Tensor] = None + _s4_step_meta: Optional[Tuple[int, int, object]] = None + _s4_draft: Optional[nn.Module] = None + _s4_served: int = 0 + _s4_gen_steps: int = 0 + _s4_ll_logged: bool = False + _s4_step_attn = None + _tailbk_rewound: bool = False + _tailbk_h_logged: bool = False + _tailbk_s4l_logged: bool = False + _tailbk_c_logged: bool = False + + def prebuild_s4(self, lm_head: nn.Module, draft_model) -> bool: + """Set up the one-launch chain for this worker; False (and nothing set up) when it does not apply.""" + self._s4 = None + if disabled(): + logger.info(f"M3.1 DSpark S4: off ({DISABLE_ENV} set); the in-tree Markov chain runs") + return False + if not getattr(draft_model, "has_markov_head", False): + return False + w1 = draft_model.markov_w1 + w2 = draft_model.markov_w2 + mapping = self.mapping + tp = mapping.tp_size + vocab = w1.shape[0] + if ( + w1.dtype != torch.bfloat16 + or w2.dtype != torch.bfloat16 + or w1.shape[1] != MARKOV_RANK + or (tuple(w2.shape) != tuple(w1.shape)) + or mapping.enable_attention_dp + or (not _head_is_sharded(lm_head, vocab, tp)) + ): + return False + block = int(self.max_draft_len) + shard = vocab // tp + rows = chain.padded_rows(shard, MAX_FAST_BATCH, w1.device) + if not self._s4_ranks_agree(rows != 0 and block <= chain.MAX_BLOCK): + return False + offset = mapping.tp_rank * shard + w2s = chain.pad_w2_shard(w2[offset : offset + shard].contiguous(), rows) + state = _S4State( + mapping, rows, shard, offset, w1, w2s, block, chain.workspace_args(mapping) + ) + for batch in range(1, MAX_FAST_BATCH + 1): + state.bufs[batch, block] = chain.pad_logits_buffer(batch, block, rows, w1.device) + s4b8 = chain.enable_batch1_variant(not m31_disabled("S4B8")) + markov_ll = chain.enable_ll_exchange(MARKOV_LL) + tailbk_c = chain.enable_cached_inputs(TAILBK_SPEC) + self._s4_fp32 = chain.enable_fp32_sum(not m31_disabled("S4_FP32SUM")) + self._s4_logits_on = not m31_disabled("S4_LOGITS") + chain.precompile(mapping, rows, w1, w2s, offset, blocks=(block,), max_batch=MAX_FAST_BATCH) + self._s4 = state + self._s4_draft = draft_model + atexit.register(self._s4_log_count) + logger.info( + ( + "M3.1 DSpark S4: one-launch Markov chain set up on rank " + f"{mapping.tp_rank}" + "/" + f"{tp}" + " (shard " + f"{shard}" + " rows padded to " + f"{rows}" + ", K=" + f"{block}" + ", batches 1.." + f"{MAX_FAST_BATCH}" + "; S4B8 batch-1 kernel " + f"{('on' if s4b8 else 'off')}" + "; MARKOV_LL " + f"{('on' if markov_ll and s4b8 else 'off')}" + "; TAILBK_SPEC cached inputs " + f"{('on' if tailbk_c else 'off')}" + "; S4-logits " + f"{('on' if self._s4_logits_on else 'off')}" + "; fp32 sum " + f"{('on' if self._s4_fp32 else 'off')}" + ")" + ) + ) + return True + + def _s4_ranks_agree(self, local: bool) -> bool: + """``local`` if every rank of the TP group decided the same; raises when they differ (a rank going on into + the collective workspace alone would hang its peers).""" + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _mnnvl_workspace_all_succeeded, + ) + + comm = _get_mnnvl_workspace_comm(self.mapping) + all_on = _mnnvl_workspace_all_succeeded(comm, local) + all_off = _mnnvl_workspace_all_succeeded(comm, not local) + if not (all_on or all_off): + raise RuntimeError( + "m31 DSpark S4: the TP ranks disagree on setting up the one-launch Markov chain" + ) + return all_on + + def _s4_serves(self, num_gens: int, spec_metadata, draft_model) -> bool: + return self._s4_mode(num_gens, spec_metadata, draft_model) is not None + + def _s4_mode(self, num_gens: int, spec_metadata, draft_model): + """ "tokens" (greedy drafting: the chain's tokens are the draft), "logits" (advanced draft sampling with + S4-logits on: the chain's corrected logits go to the stock sampler), or None (the in-tree chain runs).""" + if not self._s4_common(num_gens, draft_model): + return None + if not getattr(spec_metadata, "wants_advanced_draft_sampling", False): + return "tokens" + return "logits" if getattr(self, "_s4_logits_on", False) else None + + def _s4_common(self, num_gens: int, draft_model) -> bool: + return ( + self._s4 is not None + and 1 <= num_gens <= MAX_FAST_BATCH + and (getattr(self, "guided_decoder", None) is None) + and (getattr(self, "_d2t", None) is None) + and (getattr(draft_model, "lilicorr", None) is None) + and (not getattr(draft_model, "has_candidate_selector", False)) + ) + + def _host_step_hook(self, spec_metadata) -> None: + """The engagement counter (host only, no CUDA): called by the spec metadata's prepare() on every iteration, + outside any captured graph, it evaluates the same eligibility the captured step was routed by and counts + the generation steps the chain serves. One line per rank every S4_COUNT_LOG_EVERY generation steps and at + exit: "M3.1 DSpark S4: served / generation steps". The counts start at server start, so they include + warmup and padding steps: read engagement as the difference of two lines logged after the server is ready.""" + if self._s4 is None: + return + num_gens = int(getattr(spec_metadata, "num_generations", 0) or 0) + if num_gens <= 0: + return + self._s4_gen_steps += 1 + self._s4_served += int(self._s4_serves(num_gens, spec_metadata, self._s4_draft)) + if self._s4_gen_steps % S4_COUNT_LOG_EVERY == 0: + self._s4_log_count() + + def _s4_log_count(self) -> None: + if self._s4 is not None and self._s4_gen_steps: + logger.info( + ( + "M3.1 DSpark S4: served " + f"{self._s4_served}" + "/" + f"{self._s4_gen_steps}" + " generation steps on rank " + f"{self.mapping.tp_rank}" + ) + ) + + def _forward_impl( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + num_gens = attn_metadata.num_seqs - attn_metadata.num_contexts + mode = self._s4_mode(num_gens, spec_metadata, draft_model) + served = mode is not None + if not served: + return super()._forward_impl( + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager, + ) + inner = draft_model.logits_processor + draft_model.logits_processor = _ShardedDraftLogits(inner, self._s4, num_gens) + self._s4_step = self._s4 + self._s4_step_mode = mode + self._s4_step_meta = (attn_metadata.num_contexts, attn_metadata.num_seqs, spec_metadata) + self._s4_step_attn = attn_metadata + try: + return super()._forward_impl( + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager, + ) + finally: + draft_model.logits_processor = inner + self._s4_step = None + self._s4_step_mode = None + self._s4_tokens = None + self._s4_step_meta = None + self._s4_step_attn = None + self._tailbk_rewound = False + + def _refine_block_logits(self, draft_model, gen_logits, inputs, spec_metadata): + state = self._s4_step + if state is not None and gen_logits.dim() == 3 and (gen_logits.shape[-1] == state.shard): + base = state.buffer(gen_logits.shape[0]) + if ( + chain.supports(base, state.w1, state.w2s) + and getattr(self, "_s4_step_mode", None) == "logits" + ): + first_prev = inputs["first_prev_tokens"].long() + tokens = self._s4_chain(state, base, first_prev, spec_metadata) + self._s4_log_ll(base.shape[0]) + biased = _biased_shard_logits( + gen_logits, + first_prev, + tokens, + state.w1, + state.w2s[: state.shard], + self._s4_fp32, + ) + self._s4l_log(gen_logits, first_prev, tokens, state.w1) + return biased + if chain.supports(base, state.w1, state.w2s): + tokens = self._s4_chain( + state, base, inputs["first_prev_tokens"].long(), spec_metadata + ) + self._s4_log_ll(base.shape[0]) + self._s4_tokens = tokens + return gen_logits + return super()._refine_block_logits(draft_model, gen_logits, inputs, spec_metadata) + + def _s4_chain( + self, state: _S4State, base: torch.Tensor, first_prev: torch.Tensor, spec_metadata + ): + """The chain launch; under TAILBK_SPEC (site H) it also applies the step's deferred KV-length rewind when + the stock _apply_kv_rewind_after_draft would (a real step: not the eager warmup of a graph batch, whose + kv_lens is restored wholesale), and marks it done for that method.""" + rewind = self._tailbk_rewind_args(base.shape[0], spec_metadata) + if rewind is None: + tokens = chain.dspark_markov_chain( + base, first_prev, state.w1, state.w2s, state.offset, **state.ws + ) + self._tailbk_c_log() + return tokens + kv_lens, amount, first = rewind + tokens = chain.dspark_markov_chain_rewind( + base, first_prev, state.w1, state.w2s, state.offset, kv_lens, amount, first, **state.ws + ) + self._tailbk_rewound = True + self._tailbk_c_log() + if not self._tailbk_h_logged: + self._tailbk_h_logged = True + logger.info( + ( + "[M3 dispatch] TAILBK_SPEC H: KV-length rewind in the Markov chai" + "n launch on rank " + f"{self.mapping.tp_rank}" + ) + ) + return tokens + + def _tailbk_c_log(self) -> None: + """The TAILBK_SPEC C engagement line: once per rank, after the first chain call that took its accepted inputs + from the precompile cache (host only; a captured step logs it while it is captured).""" + if self._tailbk_c_logged or not getattr(chain, "cached_inputs_taken", lambda: False)(): + return + self._tailbk_c_logged = True + logger.info( + ( + "[M3 dispatch] TAILBK_SPEC C: chain accepted inputs from precompi" + "le cache (taken) on rank " + f"{self.mapping.tp_rank}" + ) + ) + + def _tailbk_rewind_args(self, batch: int, spec_metadata): + """(kv_lens, rewind amounts, first row) for the chain to apply, or None (the stock rewind runs).""" + if not TAILBK_SPEC or not getattr(self, "_kv_rewind_pending", False): + return None + amount = getattr(self, "_kv_rewind_amount", None) + attn = self._s4_step_attn + if amount is None or attn is None or (not hasattr(attn, "kv_lens_cuda")): + return None + if spec_metadata.is_cuda_graph and (not torch.cuda.is_current_stream_capturing()): + return None + nc, bs = (self._kv_rewind_nc, self._kv_rewind_bs) + if bs - nc != batch or amount.numel() != batch: + return None + kv_lens = attn.kv_lens_cuda + if not chain.supports_rewind(kv_lens, amount, nc): + return None + return (kv_lens, amount, nc) + + def _apply_kv_rewind_after_draft(self, attn_metadata, spec_metadata): + if self._tailbk_rewound: + self._tailbk_rewound = False + self._kv_rewind_pending = False + return + super()._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) + + def _s4l_log(self, shard_logits, first_prev, tokens, w1) -> None: + """The TAILBK_SPEC S4L engagement line: once per rank, on the first step whose biased logits took the + two-launch glue (host only; a captured step logs it while it is captured).""" + if self._tailbk_s4l_logged or not (self._s4_fp32 and S4_LOGITS_BATCHED and TAILBK_SPEC): + return + if markov_bias.supports(shard_logits, first_prev, tokens, w1): + self._tailbk_s4l_logged = True + logger.info( + ( + "[M3 dispatch] TAILBK_SPEC S4L: biased block logits glue in two l" + "aunches on rank " + f"{self.mapping.tp_rank}" + ) + ) + + def _s4_log_ll(self, batch: int) -> None: + """The MARKOV_LL engagement line: once per rank, on the first served step whose chain call took the LL + exchange kernel (host only; a captured step logs it while it is captured).""" + if not self._s4_ll_logged and chain.ll_exchange_engaged(batch): + self._s4_ll_logged = True + logger.info( + ( + "[M3 dispatch] MARKOV_LL: k3_markov LL flag-in-payload exchange (" + "batch 1) on rank " + f"{self.mapping.tp_rank}" + ) + ) + + def sample_draft_tokens(self, logits, spec_metadata, batch_size, **kwargs): + tokens = self._s4_tokens + if tokens is not None and logits.dim() == 3 and (tokens.shape == logits.shape[:2]): + self._s4_tokens = None + return tokens if tokens.dtype == torch.int32 else tokens.to(torch.int32) + return super().sample_draft_tokens(logits, spec_metadata, batch_size, **kwargs) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/m31_switches.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/m31_switches.py new file mode 100644 index 000000000000..51f966c595fd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/m31_switches.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kill switches of this MiniMax-M3.1 target's default-on items (each target keeps its own copy: targets share +only the catalog). + +Every item is on by default; ``TRTLLM_M31_DISABLE_=1`` turns item ```` off and restores the path it replaced. +The environment is read once, at import, so every rank and every step of a process sees the same setting (a step must +not change its graph mid-run). Each item lists its id and what the switch restores where it reads it. + +Values (case-insensitive, surrounding blanks ignored): ``1`` / ``true`` / ``on`` / ``yes`` set a switch; empty / ``0`` / +``false`` / ``off`` / ``no`` leave it unset; anything else raises at import, naming the variable, so a typo fails loudly +instead of silently turning an item on or off. +""" + +import os +from typing import FrozenSet + +_PREFIX = "TRTLLM_M31_DISABLE_" +_SET = frozenset(("1", "true", "on", "yes")) +_UNSET = frozenset(("", "0", "false", "off", "no")) + + +def _is_set(name: str, value: str) -> bool: + """Whether the switch variable ``name`` is set to ``value``; raises on a value of neither kind.""" + word = value.strip().lower() + if word in _SET: + return True + if word in _UNSET: + return False + raise ValueError( + f"{name}={value!r}: expected 1 / true / on / yes (set) or 0 / false / off / no / empty (unset)" + ) + + +# The item ids switched off in this process, read once at import. +_DISABLED: FrozenSet[str] = frozenset( + name[len(_PREFIX) :] + for name, value in os.environ.items() + if name.startswith(_PREFIX) and _is_set(name, value) +) + + +def m31_disabled(item_id: str) -> bool: + """Whether ``TRTLLM_M31_DISABLE_`` turned the item off for this process.""" + return item_id in _DISABLED + + +# Items that are OFF by default (rule 1b without a passing gate c, or a failed gate c): ``TRTLLM_M31_ENABLE_=1`` +# opts in; the item's ``TRTLLM_M31_DISABLE_`` still wins. +_ENABLE_PREFIX = "TRTLLM_M31_ENABLE_" + +_ENABLED: FrozenSet[str] = frozenset( + name[len(_ENABLE_PREFIX) :] + for name, value in os.environ.items() + if name.startswith(_ENABLE_PREFIX) and _is_set(name, value) +) + + +def m31_enabled(item_id: str) -> bool: + """Whether ``TRTLLM_M31_ENABLE_`` opted a default-off item in for this process.""" + return item_id in _ENABLED diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/modeling.py new file mode 100644 index 000000000000..2ec5fe79f89b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/modeling.py @@ -0,0 +1,5107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3.1 TP8 implementation.""" + +from __future__ import annotations + +import copy +import dataclasses +import os +import threading +import time +from dataclasses import dataclass +from types import MappingProxyType +from typing import ( + Any, + Callable, + Dict, + FrozenSet, + Hashable, + Iterable, + List, + Mapping, + Optional, + Set, + Tuple, + Union, +) + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + ModelingV2Mode, + 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.fakes import ( + m31_fake_ops as _m31_fake_ops, +) +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 ( + 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, + require_msa_module, +) +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.pyexecutor.sampler.ops.flashinfer import ( + set_softmax_override, + set_top_p_renorm_override, +) +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, + set_chain_speculative_sampling_override, + should_use_separate_draft_kv_cache, +) +from tensorrt_llm._torch.utils import ActType_TrtllmGen +from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm.functional import AllReduceStrategy, PositionEmbeddingType +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantAlgo + +from . import dspark as _dspark +from . import weights as _weights +from .m31_switches import m31_disabled, m31_enabled +from .ops import gqamix_sparse_decode as _gqamix_kernel +from .ops import m3_bp_tail, m3_mc_allreduce, m3_prequant, m3_spec_tail +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_sampled_accept as m3_sampled_accept_kernel +from .ops import m3_select_long_dsl as _select_long_dsl_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.activation.m3_swq import m3_swq +from .ops.attention.m3_dense_decode import m3_dense_decode +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_DECODE_QUERY_LEN as FUSED_SCORE_SELECT_MAX_QUERY_LEN, +) +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_select_long_dsl import m3_select_long_dsl +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.sampling import m31_chain_spec_mc as _chain_spec_mc +from .ops.sampling import m31_softmax_mc as _softmax_mc +from .ops.sampling import m31_topp_det as _topp_det +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) +MSA_PROXY_PREBUILD = not m31_disabled("PROXYPB") +CTXNORMQ = not m31_disabled("CTXNORMQ") +CTXNORMQ_MIN_ROWS = 1024 +FUSED_CAPTURE = not m31_disabled("FUSED_CAPTURE") +AR64AUTO = not m31_disabled("AR64AUTO") +TSEL = not m31_disabled("TSEL") +TSEL_MIN_ROWS = 256 +TSEL_MAX_BLOCKS = 5856 +TSEL_CHECK = os.environ.get("TRTLLM_M31_CHECK_TSEL") == "1" +GQAMIX = not m31_disabled("GQAMIX") +DSFOLD = not m31_disabled("DSFOLD") +BOUNDARY_PACK = DSFOLD and (not m31_disabled("BOUNDARY_PACK")) +BOUNDARY_PACK_MAX_ROWS = 1 +_BOUNDARY_PACK_MARK = "m31_boundary_pack" +QKV_PRODUCER_FOLD = not m31_disabled("QKVPROD") +QKV_PRODUCER_FOLD_PATH = "m3_qkv_producer_fold" +_QKV_FOLD_ENGAGED: Set[int] = set() +DECODE_MAX_TOKENS = 48 if m31_disabled("R1") else 64 +# Fused-layer dispatch scope; keep the larger prebuild/workspace capacity above. +FUSED_DECODE_MAX_TOKENS = 8 +SWQ_SHARED_DOWN = (6144, 384) +SWQ_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = ( + {} if m31_disabled("SWQ") else {SWQ_SHARED_DOWN: frozenset(range(9, 49))} +) +_SWQ_LOGGED: set = set() +TOPP_DET = not m31_disabled("TOPP_DET") +_TOPP_DET_STATE: dict = {} + + +def _install_topp_det() -> None: + """Once per process: install ``_topp_det.top_p_renorm_det`` as the sampler's top-p renorm override when TOPP_DET + is on and the op is built, and log the engagement (or why it is not taken) once per rank.""" + if _TOPP_DET_STATE: + return + if not TOPP_DET: + state = "off (TRTLLM_M31_DISABLE_TOPP_DET=1); stock FlashInfer top_p_renorm_probs" + elif _topp_det.topp_det_op() is None: + state = "not taken: the op is not built; stock FlashInfer top_p_renorm_probs" + else: + set_top_p_renorm_override(_topp_det.top_p_renorm_det) + ctas = _topp_det.ctas_per_row(_topp_det.MAX_ROWS, _topp_det.sm_count(torch.device("cuda"))) + state = f"engaged (ctas_per_row={ctas}, rows <= {_topp_det.MAX_ROWS})" + _TOPP_DET_STATE["installed"] = state.startswith("engaged") + logger.info(f"[M3 dispatch] ST-01 TOPP_DET det top-p renorm: {state}") + + +def _prebuild_topp_det() -> None: + """The override's persistent workspace on this rank's device, allocated before warmup and any graph capture.""" + if _TOPP_DET_STATE.get("installed"): + _topp_det.prebuild() + + +SOFTMAX_MC = not m31_disabled("SOFTMAX_MC") +_SOFTMAX_MC_STATE: dict = {} + + +def _install_softmax_mc() -> None: + """Once per process: install ``_softmax_mc.softmax_mc`` as the sampler's softmax override when SOFTMAX_MC is on + and the op is built, and log the engagement (or why it is not taken) once per rank.""" + if _SOFTMAX_MC_STATE: + return + if not SOFTMAX_MC: + state = "off (TRTLLM_M31_DISABLE_SOFTMAX_MC=1); stock FlashInfer softmax" + elif _softmax_mc.softmax_mc_op() is None: + state = "not taken: the op is not built; stock FlashInfer softmax" + else: + set_softmax_override(_softmax_mc.softmax_mc) + state = f"engaged (write_ctas={_softmax_mc.WRITE_CTAS}, rows <= {_softmax_mc.MAX_ROWS})" + _SOFTMAX_MC_STATE["installed"] = state.startswith("engaged") + logger.info(f"[M3 dispatch] ST-01-SM softmax multi-CTA write: {state}") + + +def _prebuild_softmax_mc() -> None: + """The override's persistent partials workspace on this rank's device, allocated before warmup and any graph + capture.""" + if _SOFTMAX_MC_STATE.get("installed"): + _softmax_mc.prebuild() + + +REJ_MC = not m31_disabled("REJ_MC") +_REJ_MC_STATE: dict = {} + + +def _install_rej_mc() -> None: + """Once per process: install ``_chain_spec_mc.chain_spec_mc`` as the rejection sampler's override when REJ_MC is + on and the op is built, and log the engagement (or why it is not taken) once per rank.""" + if _REJ_MC_STATE: + return + if not REJ_MC: + state = "off (TRTLLM_M31_DISABLE_REJ_MC=1); stock trtllm::chain_speculative_sampling" + elif _chain_spec_mc.chain_spec_mc_op() is None: + state = "not taken: the op is not built; stock trtllm::chain_speculative_sampling" + else: + set_chain_speculative_sampling_override(_chain_spec_mc.chain_spec_mc) + ctas = _chain_spec_mc.ctas_per_row( + _chain_spec_mc.MAX_BATCH, _chain_spec_mc.sm_count(torch.device("cuda")) + ) + state = f"engaged (ctas_per_row={ctas}, batch <= {_chain_spec_mc.MAX_BATCH})" + _REJ_MC_STATE["installed"] = state.startswith("engaged") + logger.info(f"[M3 dispatch] ST-01-REJ chain speculative sampling multi-CTA: {state}") + + +def _prebuild_rej_mc() -> None: + """The override's persistent workspace on this rank's device, allocated before warmup and any graph capture.""" + if _REJ_MC_STATE.get("installed"): + _chain_spec_mc.prebuild() + + +_HOSTTAIL_OFF = m31_disabled("HOSTTAIL") +HOSTTAIL_D2H_SIDE = not (_HOSTTAIL_OFF or m31_disabled("HOST_D2H_SIDE")) +HOSTTAIL_SMALL_H2D = not (_HOSTTAIL_OFF or m31_disabled("SMALL_H2D")) +HOSTTAIL_BATCH_CHANGED_ONLY = not (_HOSTTAIL_OFF or m31_disabled("BATCH_CHANGED_ONLY")) +HOSTTAIL_SKIP_ACCEPTED_UPLOAD = not (_HOSTTAIL_OFF or m31_disabled("SKIP_ACCEPTED_UPLOAD")) +HOSTTAIL_PHILOX_PACK = not (_HOSTTAIL_OFF or m31_disabled("PHILOX_PACK")) +HOSTTAIL_KV_LENS_ONCE = not (_HOSTTAIL_OFF or m31_disabled("KV_LENS_ONCE")) +SLOT_IDS_CHANGED = not m31_disabled("SLOT_IDS_CHANGED") +PACK_PSCATTER = not m31_disabled("PACK_PSCATTER") +PACK_IN_GRAPH = PACK_PSCATTER and (not m31_disabled("PACK_IN_GRAPH")) +PACK_EARLY_COPY = PACK_IN_GRAPH and (not m31_disabled("PACK_EARLY_COPY")) +H2D_CHANGED_SPAN = not m31_disabled("H2D_CHANGED_SPAN") +BLOCKOFF_GUARD = not m31_disabled("BLOCKOFF_GUARD") +TAILBK = not m31_disabled("TAILBK") +_TAILBK_LOGGED: list = [] + + +def _log_tailbk(worker) -> None: + """The TAILBK engagement line, once per process when the model arms it: the items its worker can take (the DSpark + worker all five; another worker's accept only A1). Nothing when TAILBK is off or without speculation.""" + if not TAILBK or worker is None or _TAILBK_LOGGED: + return + _TAILBK_LOGGED.append(True) + items = "R1 B1 Z1 S1 A1" if isinstance(worker, _dspark.M3DSparkWorker) else "A1" + logger.info(f"[M3 host] tail bookkeeping: {items}") + + +@dataclass(frozen=True) +class DecodePlan: + """The layer-boundary collectives of one decode token bucket (GB300 measurements in the comments).""" + + name: str + max_tokens: int + attn_mc_ar: bool + moe_mc_ar: bool + qkv_prequant: bool + + +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), +) +DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({1: 16384}) +DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({1: 16384}) +PRODUCER_EARLY_TRIGGER = True +PRODUCER_SHORTFIX = not m31_disabled("PRODUCER_EARLY_TRIGGER_SHORTFIX") +_PRODUCER_SHORTFIX_LOGGED: list = [] + + +def _producer_early_trigger(step) -> bool: + """Whether this step's sparse-layer producers release their dependents early (see PRODUCER_SHORTFIX).""" + if not PRODUCER_EARLY_TRIGGER: + return False + no_selector = bool(step.get("msa_short_ok")) and step.get("msa_short_table") is not None + no_selector = no_selector or False + if not (PRODUCER_SHORTFIX and no_selector): + return True + if not _PRODUCER_SHORTFIX_LOGGED: + _PRODUCER_SHORTFIX_LOGGED.append(True) + logger.info( + ( + "[M3 dispatch] PRODUCER_SHORTFIX: producer early trigger off on s" + "teps without a selector (sparse decode loads q before its grid w" + "ait)" + ) + ) + return False + + +EMBED_FIRSTNORM_MAX_TOKENS = 8 +ATTN_PUSH_MAX_TOKENS = 16 +AR_PUSH_MODE = 2 +MOE_ENGINE_MAX_TOKENS = 8 +MOE_PUSH_MAX_TOKENS = MOE_ENGINE_MAX_TOKENS +MC_FINALIZE_MIN_TOKENS = MOE_ENGINE_MAX_TOKENS + 1 +ROUTER_GEMV_MAX_TOKENS = 1 +ROUTER_TC_MAX_TOKENS = 48 +SHARD_ACCEPT_DRAFT_LEN = 3 +SHARD_ACCEPT_MAX_REQUESTS = min(DECODE_MAX_TOKENS, m3_shard_accept_kernel.SHARD_ROWS) // ( + SHARD_ACCEPT_DRAFT_LEN + 1 +) +FC_EARLY_MAX_ROWS = 0 +DRAFT_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN +DRAFT_STEP_MAX_REQUESTS = DRAFT_STEP_KERNEL_MAX_TOKENS +SPARSE_DECODE_MQ_MIN_TOKENS = 17 +FUSED_SCORE_SELECT_MAX_REQUESTS = FUSED_SCORE_SELECT_KERNEL_MAX_REQUESTS +LONG_SELECT_MIN_BLOCKS = 129 +LONG_SELECT_MAX_BLOCKS = SELECT_LONG_MAX_BLOCKS +LONG_SELECT_KEY_WIDTHS = (2, 4, 9) +SELECT_EARLY_MAX_ROWS = 8 +SELECT_LONG_DSL = not m31_disabled("SELECT_DSL") +LONG_SELECT_DSL_KEY_WIDTHS = (4, 11, 24) +_SELECT_DSL_LOGGED: set = set() +LC8_PREWAIT_MAX_ROWS = 8 +INDEX_SCORE_EARLY_MAX_REQUESTS = 0 if m31_disabled("INDEX_EARLY") else 1 +_INDEX_EARLY_LOGGED: set = set() + + +def _index_score_early(rows: int, rank: int) -> bool: + """Whether a decode step of ``rows`` requests runs the index scorer's early-issue schedule. The first time it does, + logs the EARLY-ISSUE engagement line once per rank (host side: the graph capture runs this Python too).""" + early = rows <= INDEX_SCORE_EARLY_MAX_REQUESTS + if early and rank not in _INDEX_EARLY_LOGGED: + _INDEX_EARLY_LOGGED.add(rank) + logger.info( + f"[M3 dispatch] INDEX_EARLY index-scorer early-issue: engaged (rows={rows}) on rank {rank}" + ) + return early + + +_ALL_ROWS = frozenset(range(1, min(DECODE_MAX_TOKENS, mx_gemv.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, +} +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, +} +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,) + + +def _long_select_dsl_widths(spec: Mapping[str, Any]) -> Iterable[Key]: + """The 384-thread select's key widths a cell's rows reach (every width up to its blocks' for mixed rows).""" + blocks = spec["num_blocks"] + widest = 4 if blocks <= 4 * 384 else 11 if blocks <= 11 * 384 else 24 + if spec["valid"] == "mixed": + return tuple((w for w in LONG_SELECT_DSL_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, +) +SPARSE_DECODE_LC8 = RowRange( + "m3_sparse_decode_lc8", + frozenset({"prewait", "plain"}), + m3_sparse_decode_lc8, + lambda spec: ("prewait" if spec["prewait"] else "plain",), +) +PRODUCER = RowRange("m3_producer", frozenset({"nvfp4_tp8"}), m3_producer, _const("nvfp4_tp8")) +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", "early"}) if INDEX_SCORE_EARLY_MAX_REQUESTS else frozenset({"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 +) +LONG_SELECT_DSL = RowRange( + "long_select_dsl", + frozenset(LONG_SELECT_DSL_KEY_WIDTHS) if SELECT_LONG_DSL else frozenset(), + m3_select_long_dsl, + _long_select_dsl_widths, +) +MC_ALLREDUCE = RowRange( + "mc_allreduce_norm", + _rows(1, min(DECODE_MAX_TOKENS, m3_mc_allreduce.max_rows())), + 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, min(DECODE_MAX_TOKENS, m3_mc_allreduce.max_rows())), + 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"), +) +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"]),), +) +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"]),), +) +DENSE_DECODE_SPLIT = RowRange( + "m3_dense_decode", + frozenset(DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH), + m3_dense_decode, + lambda spec: (len(spec["kv_lens"]),) if spec["capped"] else (), +) +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, + DENSE_DECODE_SPLIT, + SPARSE_DECODE_PRODUCER, + SPARSE_DECODE_LC8, + PRODUCER, + FUSED_SCORE_SELECT, + INDEX_SCORE, + LONG_SELECT, + LONG_SELECT_DSL, + 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, +) + + +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 <= min(FUSED_DECODE_MAX_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 +PREQUANT_GEMM_PATH = "flashinfer_mxfp8_gemm" +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 + + +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: {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 long_select_dsl_keys() -> Tuple[Tuple[bool, bool, int], ...]: + """The (PDL launch, early trigger, trigger phase) keys of the 384-thread select a decode step can launch, compiled + before warmup (trigger phase 0: released at its exit when not early, as the 1024-thread kernel); none when + ``LONG_SELECT_DSL`` covers nothing.""" + if not LONG_SELECT_DSL.keys: + return () + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + pdl = get_env_enable_pdl() + keys = [(pdl, False, 0)] + if SELECT_EARLY_MAX_ROWS: + keys.append((pdl, True, 0)) + return tuple(keys) + + +def _note_select_dsl(rank: int) -> None: + """The SELECT_DSL engagement line, once per process and rank (host side: also during graph capture).""" + if rank not in _SELECT_DSL_LOGGED: + _SELECT_DSL_LOGGED.add(rank) + logger.info(f"[M3 dispatch] SELECT_DSL 384-thread long select: engaged on rank {rank}") + + +def reset_logged() -> None: + """Forget which (path, reason) pairs were logged.""" + with _lock: + _logged.clear() + + +MC_AR_MAX_TOKENS = 8 +MOE_FOLD_MIN_TOKENS = 1 +MOE_ENGINE_RESERVED_SMS = 24 +FINAL_HIDDEN_LAYER_ID = -1 +ROUTING_MINIMAX2 = int(RoutingMethodType.MiniMax2) +ACT_SWIGLU = int(ActType_TrtllmGen.SwiGlu) + + +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 ValueError( + "this MiniMax-M3 target requires sparse_attention_config.algorithm='minimax_m3' with implementation='msa'" + ) + if getattr(sparse, "indexer_kv_dtype", None) != "fp8": + raise ValueError( + "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 ValueError("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 + autotune: bool + tuned_quantize_gemm: bool + flashinfer: bool + native_tuning: bool + + +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, + ) + + +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``.""" + 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)" + + +class _DeferredGemv: + """QKVPROD (``QKV_PRODUCER_FOLD``): a sparse layer's QKV + index MXFP8 GEMV, recorded instead of launched. + + ``_linear(..., defer=True)`` returns one on the decode GEMV path of the 12-head packed projection at <= 8 rows with + an MXFP8 hand-off. The producer call site either runs the GEMV with the producer in its epilogue + (``_qkv_producer_fold``) or calls ``materialize`` (the plain GEMV) before anything reads the packed row; + ``placeholder`` (unwritten, the packed shape) stands for that row until then and is never read.""" + + def __init__(self, act, act_scale, weight, weight_scale, use_pdl): + self.placeholder = torch.empty( + act.shape[0], weight.shape[0], dtype=torch.bfloat16, device=act.device + ) + self.act, self.act_scale = (act, act_scale) + self.weight, self.weight_scale = (weight, weight_scale) + self.use_pdl = use_pdl + self.done = False + + def materialize(self) -> torch.Tensor: + """Launch the plain GEMV this record stands for; returns the packed row the projection would have returned.""" + self.done = True + rows, (n, k) = (self.act.shape[0], self.weight.shape) + output = m3_mx_gemv( + self.act, self.act_scale, self.weight, self.weight_scale, True, None, self.use_pdl, None + ) + if output is None: + raise RuntimeError( + f"the deferred [{n}, {k}] GEMV declined at {rows} rows" + _uncompiled_note() + ) + return output + + +def _linear( + core, x: torch.Tensor, wk: str, sk: str, push: Optional[tuple] = None, defer: bool = False +) -> Tuple[Any, 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: + prequant = _CTXNORMQ_MX.take(x) if CTXNORMQ and (not mode.tuned_quantize_gemm) else None + if prequant is not None: + return (_ctxnormq_mxfp8_gemm(core, prequant, w, s, mode), False) + 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) + output = None + if prequant is not None: + act, act_scale = prequant + if ( + defer + and path == MX_GEMV_PATH + and (n == mx_gemv.PROD_N) + and (rows <= 8) + and (push is None) + ): + return (_DeferredGemv(act, act_scale, w, s, use_pdl), False) + 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) + + +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_rows() 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 + 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, + ) + + +_CTXNORMQ_MX = m3_prequant.PrequantSlot() +_CTXNORMQ_FP4 = m3_prequant.PrequantSlot() +_CTXNORMQ_STATE = {} + + +def _ctxnormq_op(): + """``torch.ops.trtllm.minimax_m31_ctx_residual_rmsnorm_mxfp8``, or None. A build without the op in its library + loads it from ``tensorrt_llm/libs/libm31_ctxnormq.so`` when that file exists (an overlay); otherwise the boundary + keeps the fused all-reduce. Decided once per process.""" + if "op" not in _CTXNORMQ_STATE: + op = getattr(torch.ops.trtllm, "minimax_m31_ctx_residual_rmsnorm_mxfp8", None) + if op is None: + import tensorrt_llm + + lib = os.path.join(os.path.dirname(tensorrt_llm.__file__), "libs", "libm31_ctxnormq.so") + if os.path.exists(lib): + torch.ops.load_library(lib) + _m31_fake_ops.register() + op = getattr(torch.ops.trtllm, "minimax_m31_ctx_residual_rmsnorm_mxfp8", None) + _CTXNORMQ_STATE["op"] = op + logger.info( + ( + "[M3 dispatch] CTXNORMQ prefill boundary norm + quantize: " + f"{('engaged' if op is not None else 'not taken: the op is not built')}" + " (rows >= " + f"{CTXNORMQ_MIN_ROWS}" + ")" + ) + ) + return _CTXNORMQ_STATE["op"] + + +def _ctxnormq_allreduce_norm(core, x, residual, norm_w, fp4_scale): + """A prefill boundary as the plain all-reduce + the fused residual add, RMSNorm and quantize, or None (switched + off, fewer than ``CTXNORMQ_MIN_ROWS`` rows, op not built). Returns (normed, updated) like the fused all-reduce; + the MXFP8 activation (and with ``fp4_scale`` the NVFP4 one) wait in the CTXNORMQ slots for their consumers.""" + if not CTXNORMQ or x.dim() != 2 or x.shape[0] < CTXNORMQ_MIN_ROWS: + return None + op = _ctxnormq_op() + if op is None: + return None + reduced = m3_allreduce(x, core.allreduce, None) + normed, xq, xsf, xfp4, xfp4_sf = op( + reduced, residual.contiguous(), norm_w, core.rms_eps, True, fp4_scale, get_env_enable_pdl() + ) + _CTXNORMQ_MX.publish(normed, xq, xsf) + if fp4_scale is not None: + _CTXNORMQ_FP4.publish(normed, xfp4, xfp4_sf) + else: + _CTXNORMQ_FP4.clear() + return (normed, reduced) + + +def _ctxnormq_mxfp8_gemm(core, prequant, w, s, mode): + """The GEMM ``m3_mxfp8_linear`` runs after its own quantize, on the CTXNORMQ activation (bit-identical operands).""" + act, act_scale = prequant + if mode.flashinfer: + return torch.ops.trtllm.flashinfer_mm_mxfp8(act, act_scale, w, s, torch.bfloat16) + gemm = ( + torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned + if mode.native_tuning + else torch.ops.trtllm.mxfp8_mxfp8_gemm + ) + return gemm(act, act_scale, w, s, core.mx_alpha, torch.bfloat16) + + +_TSEL_STATE = {} + + +def _tsel_op(): + """``torch.ops.trtllm.minimax_m31_select_prefill_tile`` on sm_103, or None. A build without the op in its library + loads it from ``tensorrt_llm/libs/libm31_tsel.so`` when that file exists (an overlay). Decided once per process.""" + if "op" not in _TSEL_STATE: + op = None + if torch.cuda.get_device_capability() == (10, 3): + op = getattr(torch.ops.trtllm, "minimax_m31_select_prefill_tile", None) + if op is None: + import tensorrt_llm + + lib = os.path.join(os.path.dirname(tensorrt_llm.__file__), "libs", "libm31_tsel.so") + if os.path.exists(lib): + torch.ops.load_library(lib) + _m31_fake_ops.register() + op = getattr(torch.ops.trtllm, "minimax_m31_select_prefill_tile", None) + _TSEL_STATE["op"] = op + logger.info( + ( + "[M3 dispatch] TSEL prefill tile select: " + f"{('engaged' if op is not None else 'not taken: not sm_103 or the op is not built')}" + " (rows >= " + f"{TSEL_MIN_ROWS}" + ", blocks <= " + f"{TSEL_MAX_BLOCKS}" + ")" + ) + ) + return _TSEL_STATE["op"] + + +def _tsel_select(max_score_kv, n_valid_blocks, init_blocks, local_blocks): + """The MSA indexer's selector hook: the head-major [total_q, num_kv_heads, 16] table of + ``select_blocks_from_maxscore``, or None (outside the tile select's rows / blocks range, or no op) to keep it.""" + if ( + torch.cuda.is_current_stream_capturing() + or max_score_kv.dtype != torch.float32 + or max_score_kv.dim() != 3 + or (max_score_kv.stride(2) != 1) + or (max_score_kv.shape[2] < TSEL_MIN_ROWS) + or (max_score_kv.shape[1] > TSEL_MAX_BLOCKS) + ): + return None + op = _tsel_op() + if op is None: + return None + num_kv_heads, num_blocks, total_q = max_score_kv.shape + out = torch.empty( + (num_kv_heads, total_q, MSA_REQUIRED_TOPK), dtype=torch.int32, device=max_score_kv.device + ).permute(1, 0, 2) + nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous() + op(max_score_kv, nvb, out, init_blocks, local_blocks) + if TSEL_CHECK: + ref = torch.ops.trtllm.minimax_m3_select_blocks( + max_score_kv, nvb, MSA_REQUIRED_TOPK, init_blocks, local_blocks, True + ) + if not torch.equal(out, ref): + raise AssertionError( + ( + "TSEL check: tile table mismatch against minimax_m3_select_blo" + "cks (" + f"{int((out != ref).any(dim=-1).sum())}" + " rows; " + f"{num_blocks}" + " blocks, " + f"{total_q}" + " rows)" + ) + ) + _TSEL_STATE["checked"] = _TSEL_STATE.get("checked", 0) + 1 + if _TSEL_STATE["checked"] in (1, 100, 1000): + logger.info( + f"[M3 dispatch] TSEL check: {_TSEL_STATE['checked']} tile-selected steps equal" + ) + return out + + +def _tsel_register(backend) -> None: + """Register ``_tsel_select`` on the layer backend's MSA indexer (an indexer without the hook is left alone).""" + indexer = getattr(backend, "indexer", None) + if indexer is not None and getattr(indexer, "selector_hook", False) is None: + indexer.selector_hook = _tsel_select + + +_GQAMIX_STATE: dict = {} + + +def _gqamix_decode(*args, **kwargs) -> None: + """``minimax_m3_sparse_attn_decode`` for MsaDecodeFmha's sparse generation rows, through the target-local copy.""" + if not _GQAMIX_STATE.get("logged"): + _GQAMIX_STATE["logged"] = True + logger.info( + "[M3 dispatch] GQAMIX: mixed-step sparse decode through ops/gqamix_sparse_decode" + ) + _gqamix_kernel.minimax_m3_sparse_attn_decode(*args, **kwargs) + + +def _gqamix_register(backend) -> None: + """Register ``_gqamix_decode`` on the layer backend's MsaDecodeFmha (a library without the hook is left alone).""" + manager = getattr(backend, "_fmha_manager", None) + for fmha in getattr(manager, "fmha_libs", ()): + if getattr(fmha, "sparse_decode_hook", False) is None: + fmha.sparse_decode_hook = _gqamix_decode + + +def _native_allreduce_norm(core, x, residual, norm_w): + """The built-in fused all-reduce + residual + RMSNorm: with AR64AUTO, the strategy-MNNVL instance for the rows past + the multicast all-reduce's, up to the decode bound; the configured instance otherwise.""" + params = AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=norm_w, + eps=core.rms_eps, + trigger_completion_at_end=False, + ) + wide = getattr(core, "allreduce_wide", None) + rows = x.shape[0] + if wide is not None and m3_mc_allreduce.max_rows() < rows <= DECODE_MAX_TOKENS: + out = m3_allreduce(x.contiguous(), wide, params) + if out is not None: + _log_once( + logger.info, + "ar64auto", + "taken", + f"[M3.1] AR64AUTO taken: MNNVL all-reduce at {rows} rows", + ) + return out + _log_once( + logger.warning, + "ar64auto", + f"declined {rows}", + f"[M3.1] AR64AUTO declined: MNNVL refused {rows} rows; the configured all-reduce runs", + ) + return m3_allreduce(x, core.allreduce, params) + + +@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() + 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) + + +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, + deferred: Optional[_DeferredGemv] = None, + early_trigger: bool = PRODUCER_EARLY_TRIGGER, +): + """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. ``deferred``: the QKV GEMV ``_linear`` recorded instead of launching + (``QKV_PRODUCER_FOLD``; ``packed`` is then its unwritten placeholder). The ``m3_producer`` branch runs it with the + producer in its epilogue (``_qkv_producer_fold``); any other branch launches it first (``materialize``). A None + return leaves it unlaunched.""" + 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) + ): + if deferred is not None: + folded = _qkv_producer_fold( + core, + i, + deferred, + buffers, + scale_buffers, + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + cos_sin, + pos, + early_trigger, + ) + if folded is not None: + return folded + packed = deferred.materialize() + 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(), + early_trigger, + ) + else: + if layer_uses_nvfp4 and covered: + missed( + PRODUCER.path, + rows, + ( + "the NVFP4 pool " + f"{tuple(buffers.shape)}" + " / scale pool " + f"{tuple(scale_buffers.shape)}" + " is not the producer's layout" + ), + ) + if deferred is not None: + packed = deferred.materialize() + 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 _qkv_producer_fold( + core, + i: int, + deferred: _DeferredGemv, + buffers, + scale_buffers, + index_k_cache, + out_cache_loc, + kv_quant_scale, + cos_sin, + pos, + early_trigger: bool = True, +) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """QKVPROD: the deferred QKV GEMV with layer ``i``'s producer in its epilogue (``mx_gemv.m3_mx_gemv_producer``), + the same outputs and cache bytes as the GEMV + ``m3_producer``; None (nothing launched) when the kernel declines.""" + rows = int(deferred.act.shape[0]) + device = deferred.act.device + q = torch.empty(rows, core.heads, core.head_dim, dtype=torch.float8_e4m3fn, device=device) + index_q = torch.empty(rows, 1, core.head_dim, dtype=torch.float8_e4m3fn, device=device) + w = core.w + ok = mx_gemv.m3_mx_gemv_producer( + deferred.act, + deferred.act_scale, + deferred.weight, + deferred.weight_scale, + q, + index_q, + buffers.view(torch.uint8), + scale_buffers.view(torch.uint8), + index_k_cache, + out_cache_loc, + 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, + trigger_early=early_trigger, + use_pdl=deferred.use_pdl, + ) + if not ok: + missed( + QKV_PRODUCER_FOLD_PATH, rows, "the producer GEMV declined the step" + _uncompiled_note() + ) + return None + deferred.done = True + if rows not in _QKV_FOLD_ENGAGED: + _QKV_FOLD_ENGAGED.add(rows) + logger.info( + ( + "[M3 dispatch] QKVPROD producer in the QKV GEMV epilogue: engaged" + " (rows=" + f"{rows}" + ") on rank " + f"{int(core.mapping.tp_rank)}" + ) + ) + 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, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + 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( + f"MiniMax-M3 fused FP8 indexer requires an E4M3 index-K cache, got {index_k_cache.dtype}." + ) + rows = int(idx_qk.shape[0]) + 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) + + +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"] + + +_ENGINE_STATES: Dict[tuple, object] = {} +AR_AUTOTUNE_ENABLED = m31_enabled("AR_AUTOTUNE") + + +def _pin_allreduce_heuristic(rank: int) -> bool: + """Default the AllReduce autotune off for this process (see AR-DET); returns whether it is off. Logs one line.""" + if not AR_AUTOTUNE_ENABLED: + os.environ.setdefault("TLLM_DISABLE_ALLREDUCE_AUTOTUNE", "1") + off = os.environ.get("TLLM_DISABLE_ALLREDUCE_AUTOTUNE", "0") == "1" + logger.info( + f"M3.1 AR-DET: AllReduce autotune {('off (C++ size / world heuristic)' if off else 'on')} on rank {rank}" + ) + return off + + +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 + _pin_allreduce_heuristic(mapping.tp_rank) + 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 + 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()) + self.fused_projection = _weights.fused_projection(model_config.sparse_attention_config) + self.w = _weights.MODEL_WEIGHTS.declare(d) + self.mxfp8 = _Mxfp8GemmState() + self._use_flashinfer_mxfp8_decode_graph_default = True + self.allreduce = AllReduce( + mapping=mapping, strategy=model_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + self.allreduce_wide = None + self._allreduce_wide_owner = None + if AR64AUTO and DECODE_MAX_TOKENS > m3_mc_allreduce.max_rows(): + self._allreduce_wide_owner = AllReduce( + mapping=mapping, strategy=AllReduceStrategy.MNNVL, dtype=cfg.torch_dtype + ) + self.allreduce_wide = getattr(self._allreduce_wide_owner, "mnnvl_allreduce", None) + if self.allreduce_wide is not None: + logger.info( + ( + "[M3.1] AR64AUTO: boundaries of " + f"{m3_mc_allreduce.max_rows() + 1}" + ".." + f"{DECODE_MAX_TOKENS}" + " rows take the strategy-MNNVL all-reduce (configured: " + f"{model_config.allreduce_strategy}" + ")" + ) + ) + else: + logger.warning( + "[M3.1] AR64AUTO: MNNVL unavailable on this mapping; every boundary keeps the configured all-reduce" + ) + 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 + } + 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} + ) + 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 + ) + 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__}" + ) + if i not in self.sparse_layers: + backend.dense_decode_max_seq_len_by_batch = DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH + self.attn.append(backend) + 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) + + 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): + 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: + per_step = getattr(spec, "tokens_per_gen_step", None) + if isinstance(per_step, int) and per_step > 1: + lens.add(per_step) + 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) + if ( + QKV_PRODUCER_FOLD + and n == mx_gemv.PROD_N + and (not x_bf16) + and (ar is None) + and (rows <= 8) + ): + mx_gemv.precompile(weight, scale, rows, False, use_pdl=use_pdl, prod=True) + if PRODUCER_EARLY_TRIGGER and PRODUCER_SHORTFIX: + mx_gemv.precompile( + weight, + scale, + rows, + False, + use_pdl=use_pdl, + prod=True, + trigger_early=False, + ) + swq_shapes = {tuple(self.w[wk].shape) for wk, _, _ in self.projections()} & set(SWQ_ROWS) + for n, k in sorted(swq_shapes): + gu = torch.zeros( + max(SWQ_ROWS[n, k]), 2 * k, dtype=torch.bfloat16, device=self.w["embed"].device + ) + m3_swq(gu, self.swiglu_limit, self.swiglu_alpha, self.swiglu_beta) + logger.info(f"[M3 prebuild] SWQ: compiled for intermediate {k}") + 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 + start = time.perf_counter() + times["select_keys"] = [] + if long_select_dsl_keys(): + times["select_keys"] = _select_long_dsl_kernel.prebuild( + long_select_dsl_keys(), self.w["embed"].device + ) + torch.cuda.synchronize() + times["select_s"] = time.perf_counter() - start + start = time.perf_counter() + times["proxy_variants"] = self._prebuild_msa_proxy() if MSA_PROXY_PREBUILD else [] + times["proxy_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), MSA index" + "er proxy %.1f s (%d variants)" + ) + % ( + times["gemv_s"], + times["moe_s"], + times["scorer_s"], + len(times["scorer_keys"]), + times["proxy_s"], + len(times["proxy_variants"]), + ) + ) + return times + + 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, + ) + m3_moe_engine( + torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device), + engine, + ) + return engine + + 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_msa_proxy(self) -> list: + """Compile and load every fmha_sm100 variant the MSA indexer's proxy pass can pick (OnlyScore over the paged + index-K: one KV head, page 128, no split-KV), so no context chunk compiles one after the server is ready (a + 65..128-token chunk's variant did, stalling the steps behind it by 0.8-1.6 ms). The variant follows only the + longest query row: fmha_sm100 packs a row of up to 32 tokens by its pack factor, then takes the 128- or 256-row + query tile and the single-warpgroup kernel from the packed length, so query lengths 1..129 reach every variant. + Compile and load only, nothing runs: no numerics change. fmha_sm100 keys its cache by its installed sources, so + a source overlay builds its own variants. Returns the variant names.""" + require_msa_module() + from fmha_sm100 import api as fmha_api + from fmha_sm100 import jit as fmha_jit + + only_score, page = (2, 128) + names = [] + for i in sorted(self.sparse_layers): + backend = self.attn[i] + dtype = torch.float8_e4m3fn if backend.indexer_kv_dtype == "fp8" else torch.bfloat16 + heads = int(backend.m3_config.num_index_heads) + for qlen in range(1, 130): + pack = fmha_api._compute_pack_factor(qlen, heads, 1) + packed = qlen * pack + key = ( + fmha_jit._dlpack_dtype_code(dtype), + 128 if packed <= 128 else 256, + packed <= 64, + only_score, + page, + False, + pack, + ) + name, _ = fmha_jit._variant_key_from_runtime(*key) + if name not in names: + fmha_jit.get_fmha_variant(*key) + names.append(name) + return names + + 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 + + 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 + ) + + +def _swiglu_act(core, gu: torch.Tensor) -> torch.Tensor: + return m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + + +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 + 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) + 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 = _swiglu_act(core, gu) + 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 + fp4 = _CTXNORMQ_FP4.take(x) if CTXNORMQ else None + if num_chunks <= 1: + if fp4 is not None: + return _routed_experts(core, i, fp4[0], fp4[1], router_logits, True) + 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)) + sfs = [None] * num_chunks + if fp4 is not None: + xs = fp4[0].split(sizes) + sfs = [t.flatten() for t in fp4[1].view(rows, -1).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], sfs[c], logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], sfs[c], 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: + 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], + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if return_parts: + return (shared_out, routed_out) + 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) + if norm_w is not core.w["final_norm"]: + fused = _ctxnormq_allreduce_norm(core, x, residual, norm_w, None) + if fused is not None: + return fused + 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: + fused = _ctxnormq_allreduce_norm( + core, + hidden, + residual, + core.w[f"l{i}_post_norm"], + core.moe[i].fc31_input_scale if i in core.moe_layers else None, + ) + hidden, residual = ( + fused + if fused is not None + else _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: + 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", defer=QKV_PRODUCER_FOLD) + deferred = None + if isinstance(packed, _DeferredGemv): + deferred, packed = (packed, packed.placeholder) + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + deferred, + early_trigger=_producer_early_trigger(step), + ) + if horizontal is None: + raise RuntimeError( + ( + "MiniMax-M3 sparse layer " + f"{i}" + ": no horizontal producer takes this step's caches (an FP8 or NVF" + "P4 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 MXFP8 index projection + (the checkpoint's block scales applied) and the FP8 indexer kernel on the attention aux stream; then the MSA + backend (it writes the main KV cache) and the o_proj. The NVFP4 producer does not apply: it packs the + five-way projection's output.""" + core = self.core + if _producer_covered(core): + uncovered( + PRODUCER.path, + "fuse_qkv_index_projection is disabled: QKV + separate MXFP8 index projection", + ) + + 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, _ = _linear(core, hidden, f"l{i}_idx_w", f"l{i}_idx_s", True) + 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 + if TSEL and sparse: + _tsel_register(core.attn[i]) + if GQAMIX and sparse: + _gqamix_register(core.attn[i]) + 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 8 tokens with 1-D input ids from one request, per layer: attention (o_proj all-reduce deferred, its rows + pushed into the boundary all-reduce) -> attention boundary (all-reduce + residual + post-attention RMSNorm) -> + MoE (the 1..8-row engine, pushing its sum into the boundary all-reduce) 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] = {} + 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) + kernels = ( + inputs_embeds is None + and input_ids.ndim == 1 + and attn_metadata.num_seqs == 1 + and is_decode_step(input_ids.numel()) + ) + if kernels: + num_tokens = int(input_ids.numel()) + plan = plan_for(num_tokens) + hidden, residual = self._embed(input_ids, num_tokens) + else: + uncovered( + "decode_kernels", + f"outside single-request {FUSED_DECODE_MAX_TOKENS}-row fused scope, " + "embeddings or 2-D ids: generic layers", + ) + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target does not run under piecewise (breakable) CUDA graphs" + ) + plan = None + 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): + captured = spec_metadata is not None and spec_metadata.is_layer_capture(i) + if not kernels or residual is None or (captured and (not FUSED_CAPTURE)): + 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) + if captured: + capture(spec_metadata, i, residual, None) + _log_once( + logger.info, + "B-4", + "fused capture", + ( + "[M3 dispatch] B-4 fused capture taken: capture layer " + f"{i}" + " runs the decode-kernel layer (logged once)" + ), + ) + 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) + 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 and (not MC_ALLREDUCE.covers(x.shape[0])): + uncovered( + MC_ALLREDUCE.path, + f"more than {m3_mc_allreduce.max_rows()} rows: the built-in all-reduce runs", + ) + use_mc = False + 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] + 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") + quant = False + return self._boundary(plan.attn_mc_ar, out, residual, norm_w, quant=quant) + quant = plan.qkv_prequant and i + 1 < core.num_layers + path = moe_path(num_tokens) + if path is None: + uncovered( + MC_FINALIZE.path, + f"more than {m3_mc_allreduce.max_rows()} rows: the finalized MoE and the boundary all-reduce run", + ) + moe_push = None + if path == MOE_ENGINE.path: + 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 + ) + 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: + 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") + n, k = core.w[f"{prefix}_dn_w"].shape + rows = gu.shape[0] + if rows in SWQ_ROWS.get((n, k), ()) and _mx_mode(core.mxfp8).decode_graph: + act, act_scale = m3_swq(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + m3_prequant.SLOT.publish(gu, act, act_scale) + if rows not in _SWQ_LOGGED: + _SWQ_LOGGED.add(rows) + logger.info(f"[M3 dispatch] SWQ: SwiGLU + MXFP8 producer rows={rows} [{n}, {k}]") + return _linear(core, gu, f"{prefix}_dn_w", f"{prefix}_dn_s", push) + act = _swiglu_act(core, gu) + 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: + 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], + 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) + 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: + 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", defer=QKV_PRODUCER_FOLD) + deferred = None + if isinstance(packed, _DeferredGemv): + deferred, packed = (packed, packed.placeholder) + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + deferred, + early_trigger=_producer_early_trigger(step), + ) + if horizontal is None: + raise RuntimeError( + ( + "MiniMax-M3 sparse layer " + f"{i}" + ": no horizontal producer takes this step's caches (an FP8 or NVF" + "P4 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 MXFP8 index projection + (the checkpoint's block scales applied) and the FP8 indexer kernel on the attention aux stream; then the MSA + backend (it writes the main KV cache) and the o_proj. The NVFP4 producer does not apply: it packs the + five-way projection's output. ``push`` is the attention + boundary exchange the o_proj GEMV writes its rows into, as in the five-way path.""" + core = self.core + if _producer_covered(core): + uncovered( + PRODUCER.path, + "fuse_qkv_index_projection is disabled: QKV + separate MXFP8 index projection", + ) + + 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, _ = _linear(core, hidden, f"l{i}_idx_w", f"l{i}_idx_s", True) + 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 + 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 or False + 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: + for _rep in range(1): + 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 = _index_score_early(rows, int(core.mapping.tp_rank)) + group = config.num_index_heads // config.num_kv_heads + 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: + missed( + INDEX_SCORE.path, + num_tokens, + f"the index scorer {key} was not compiled before warmup", + ) + score_view = max_score.transpose(1, 2) + m3_index_score( + idx_q_view, + idx_k_cache.squeeze(1), + block_table, + seq_lens, + score_view, + qlen, + INDEX_SCORE_SPLIT_K, + early, + ) + n_valid_blocks = step["msa_n_valid_blocks"][:num_tokens] + 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 + ) + early_trigger = num_tokens <= SELECT_EARLY_MAX_ROWS + dsl = False + if LONG_SELECT_DSL.keys: + dsl_key = _select_long_dsl_kernel.kernel_key( + get_env_enable_pdl(), early_trigger, 0 + ) + dsl = dsl_key in _select_long_dsl_kernel.kernel_cache + if dsl: + _note_select_dsl(int(core.mapping.tp_rank)) + else: + missed( + LONG_SELECT_DSL.path, + num_tokens, + f"the 384-thread select {dsl_key} was not compiled before warmup", + ) + 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=early_trigger, + dsl=dsl, + ) + if table is None: + uncovered( + LONG_SELECT.path, + ( + f"{max_score_kv.shape[1]}" + " KV blocks, outside " + f"{LONG_SELECT_MIN_BLOCKS}" + ".." + f"{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, + ) + 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 + if qlen > FUSED_SCORE_SELECT_MAX_QUERY_LEN: + uncovered( + path, + f"{qlen} query rows per request (more than {FUSED_SCORE_SELECT_MAX_QUERY_LEN}): 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 + 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], + ) + + +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 M3DSparkStepFold(M3SpecStepFold): + """The step fold under the standalone DSpark worker (DSFOLD): the engine's previous-batch gathers in one launch + (``step_gather``). ``taken`` / ``declined`` count the steps the gather ran / fell back to the engine's own copies. + + With ``boundary_pack`` (BOUNDARY_PACK BP-1) a batch-1 decode graph also ends with the sampler's slot stores and the + next step's gather (``graph_tail``: one ``ops/m3_bp_tail`` launch), so the boundary between two replays runs + neither: + + * ``begin_step`` (before every replay) writes the device control word (slot, live) when it changed: live 1 with + the step's slot when the step is one row the fold identifies (a slot inside the store), else live 0, so a replay + the fold did not take (a dummy-only replay, an unidentified step) writes nothing at all; it reports the inputs + staged when the previous graph gathered them for exactly these rows; + * the graph's tail waits for the sampler's previous side-stream store copies (``store_copies_done``), then with live + 1 stores the row as the sampler's ``_store_outputs`` (same output row, same store, zero padding / truncation to + the store widths, at the slot) and, with a static draft length, writes the engine's next-step gathers of that row + (P = 1, ids0 = d0 = b0 = 0, both offset buffers zeroed elsewhere), and marks its outputs; + * ``sampler_rows_stored`` (after the replay) accepts a marked step whose rows are the live ones, and arms the next + ``begin_step``. + + Every other step (context or mixed steps, batches over one row, eager forwards, a draft length that can change, + first steps, slots the fold could not identify) keeps the sampler's stores and the engine's gathers. + ``bp_stores`` / ``bp_gathers`` count the steps whose stores / gathers the graph took.""" + + store_copies_done: Optional[torch.cuda.Event] = None + + def __init__(self, boundary_pack: bool = False, tokens_per_step: Optional[int] = None) -> None: + super().__init__() + self.taken = 0 + self.declined = 0 + self.boundary_pack = boundary_pack + self.tokens_per_step = tokens_per_step + self.bp_stores = 0 + self.bp_gathers = 0 + self._bp_ctl: Optional[torch.Tensor] = None + self._bp_ctl_host: Optional[list] = None + self._bp_block = 0 + self._bp_rows: Optional[list] = None + self._bp_staged: Optional[list] = None + self._bp_token = id(self) + + def bind_engine_buffers( + self, input_ids, draft_tokens, position_id_offsets, kv_lens_offsets, max_batch_size + ): + super().bind_engine_buffers( + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets, max_batch_size + ) + self._bp_reset() + + def bind_sampler_store(self, store: object) -> None: + super().bind_sampler_store(store) + self._bp_reset() + + def _bp_reset(self) -> None: + """A new binding (an engine or a sampler built again): BOUNDARY_PACK's buffers are made again for it.""" + self._bp_ctl = self._bp_ctl_host = None + self._bp_rows = self._bp_staged = None + self.store_copies_done = None + + def gather(self, *args, **kwargs) -> bool: + if not super().gather(*args, **kwargs): + self.declined += 1 + uncovered( + "DSFOLD", + "step_gather declined the inputs (dtype / layout); the engine's copies run", + ) + return False + if self.taken == 0: + logger.info( + "[M3 dispatch] DSFOLD DSpark step-input gather: engaged (one step_gather launch)" + ) + self.taken += 1 + return True + + def scatter( + self, seq_slots, new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens + ) -> bool: + return False + + def _bp_ready(self) -> bool: + """BOUNDARY_PACK: the control word (live 0), the tail kernel's two variants and the store-copies event, made + once outside capture when the sampler store and the engine buffers are bound and the tail can take them.""" + if not self.boundary_pack: + return False + if self._bp_ctl is not None: + return True + store = self.store + names = ("new_tokens", "next_new_tokens", "new_tokens_lens", "next_draft_tokens") + if ( + self.step_inputs is None + or not all((isinstance(getattr(store, name, None), torch.Tensor) for name in names)) + or torch.cuda.is_current_stream_capturing() + ): + return False + tensors = [getattr(store, name) for name in names] + list(self.step_inputs) + if ( + any((t.dtype != torch.int32 or not t.is_cuda for t in tensors)) + or not all((t.is_contiguous() for t in self.step_inputs)) + or (not m3_bp_tail.aligned(*tensors)) + ): + uncovered( + "BOUNDARY_PACK", + "the store or the engine buffers are not aligned int32 CUDA tensors", + ) + return False + ctl = torch.zeros(2, dtype=torch.long, device=store.new_tokens_lens.device) + block = m3_bp_tail.block_for(store) + m3_bp_tail.warm_up(ctl, store, self.step_inputs, block) + done = torch.cuda.Event(external=True) + done.record() + self.store_copies_done = done + self._bp_ctl, self._bp_ctl_host, self._bp_block = (ctl, [0, 0], block) + return True + + def begin_step(self, seq_slots: Optional[list], num_rows: int) -> bool: + staged, self._bp_staged = (self._bp_staged, None) + self._bp_rows = None + if not self._bp_ready(): + return False + num_slots = self.store.new_tokens_lens.shape[0] + rows = list(seq_slots) if seq_slots is not None else None + live = ( + rows is not None + and len(rows) == num_rows + and (0 < num_rows <= BOUNDARY_PACK_MAX_ROWS) + and all((isinstance(slot, int) and 0 <= slot < num_slots for slot in rows)) + ) + ctl = [rows[0], 1] if live else [self._bp_ctl_host[0], 0] + if ctl != self._bp_ctl_host: + self._bp_ctl.copy_( + torch.tensor(ctl, dtype=torch.long, pin_memory=prefer_pinned()), non_blocking=True + ) + self._bp_ctl_host = ctl + if not live: + return False + self._bp_rows = rows + if staged != rows: + return False + if self.bp_gathers == 0: + logger.info( + "[M3 dispatch] BOUNDARY-PACK BP-1: sampler stores + step gather in the step graph (taken)" + ) + self.bp_gathers += 1 + return True + + def sampler_rows_stored(self, outputs: dict, num_skip: int, seq_slots: list) -> bool: + self._bp_staged = None + mark = outputs.get(_BOUNDARY_PACK_MARK) if isinstance(outputs, dict) else None + if mark is None or mark[0] != self._bp_token or self._bp_rows is None: + return False + rows, gathered = (mark[1], mark[2]) + if num_skip != 0 or len(seq_slots) != rows or list(seq_slots) != self._bp_rows: + uncovered( + "BOUNDARY_PACK", + ( + "a step graph stored rows that are not the sampler's (skipped con" + "text rows or other slot ids); the sampler stores them" + ), + ) + return False + self.bp_stores += 1 + if gathered: + self._bp_staged = list(seq_slots) + return True + + def graph_tail(self, outputs, attn_metadata): + """BOUNDARY_PACK, at the end of a decode step's forward: under capture of a batch-1 generation graph, the + sampler's slot stores and the next step's gather, gated by the live word (see the class docstring); returns + the outputs, marked when the graph has the tail. Anything else returns them unchanged (the eager path).""" + if not self.boundary_pack or not torch.cuda.is_current_stream_capturing(): + return outputs + path = "BOUNDARY_PACK" + if not isinstance(outputs, dict) or self._bp_ctl is None: + uncovered( + path, + "the sampler store or the engine buffers are not bound: eager stores and gathers", + ) + return outputs + if attn_metadata.num_contexts != 0: + return outputs + names = ("new_tokens", "next_new_tokens", "new_tokens_lens", "next_draft_tokens") + srcs = [outputs.get(name) for name in names] + if any( + ( + not isinstance(t, torch.Tensor) or t.dtype != torch.int32 or (not t.is_cuda) + for t in srcs + ) + ) or [t.dim() for t in srcs] != [2, 2, 1, 2]: + uncovered( + path, + "the worker outputs are not int32 [rows, width] / [rows]: eager stores and gathers", + ) + return outputs + rows = srcs[0].shape[0] + if any((t.shape[0] != rows for t in srcs)): + uncovered(path, "the worker outputs disagree on their rows: eager stores and gathers") + return outputs + if not 0 < rows <= BOUNDARY_PACK_MAX_ROWS: + uncovered( + path, + f"a {rows}-row decode graph (only batch-1 graphs, which are never padded): eager stores and gathers", + ) + return outputs + src = [t[:1] for t in srcs] + if not m3_bp_tail.aligned(*src): + uncovered(path, "a worker output is not 16-byte aligned: eager stores and gathers") + return outputs + store = self.store + num_tokens = int(attn_metadata.num_tokens) + t = num_tokens // rows + gather = ( + self.tokens_per_step is not None + and t == self.tokens_per_step + and (num_tokens == rows * t) + and (t <= store.next_new_tokens.shape[0]) + and (t - 1 <= store.next_draft_tokens.shape[1]) + ) + if not gather: + uncovered( + path, + ( + "the next step's tokens per request may differ from this graph's " + f"{t}" + " (a draft length that can change, or " + f"{num_tokens}" + " tokens over " + f"{rows}" + " rows): the graph stores, the engine gathers" + ), + ) + torch.cuda.current_stream().wait_event(self.store_copies_done) + m3_bp_tail.bp_tail(self._bp_ctl, src, store, self.step_inputs, t, gather, self._bp_block) + marked = dict(outputs) + marked[_BOUNDARY_PACK_MARK] = (self._bp_token, rows, gather) + return marked + + +def _static_tokens_per_step(spec_config) -> Optional[int]: + """BOUNDARY_PACK: the tokens per request of every decode step when the draft length cannot change at runtime (no + draft-length schedule, no acceptance-rate gate that turns speculation off), else None.""" + if getattr(spec_config, "draft_len_schedule", None) is not None: + return None + if ( + getattr(spec_config, "acceptance_rate_window_size", None) + and getattr(spec_config, "acceptance_rate_threshold", None) is not None + ): + return None + from tensorrt_llm._torch.speculative.utils import get_static_draft_len + + return int(spec_config.get_runtime_tokens_per_gen_step(get_static_draft_len(spec_config))) + + +_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(): + 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) + 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 + 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 ValueError( + ( + "this MiniMax-M3 target drafts with a one-layer Llama-style Eagle" + "3 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 ValueError( + f"the drafter's attention backend {type(self.attn).__name__} must fuse RoPE and QKV" + ) + self.attn.update_quant_config(quant_config) + 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.""" + + 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"]) + 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:" + " num_tokens=" + f"{num_tokens}" + ", capacity=" + f"{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: num_tokens=" + f"{num_tokens}" + ", hidden_states=" + f"{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 + ) + 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 + self._m3_fc_norm_w: Optional[tuple] = None + self._m3_capture_norm: Optional[tuple] = None + self._m3_aux_stream: Optional[torch.cuda.Stream] = None + self._m3_fc_early_drafter: Optional[M3Eagle3Draft] = None + self._m3_fc_early_stream: Optional[torch.cuda.Stream] = None + self._m3_force_vec_warm = False + self._m3_draft_head = None + self._m3_next_new: Optional[torch.Tensor] = None + self._m3_next_new_tokens: Optional[torch.Tensor] = None + 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 + 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: + 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) + 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) + 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: + self._prepare_attn_metadata_for_spec_dec(attn_metadata) + 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, + ) + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + 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, + ) + self._restore_attn_metadata_from_spec_dec(attn_metadata) + if original_all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + if self._m3_next_new_tokens is not None: + 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 + 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: + return False + if attn_metadata.num_seqs * (draft_len + 1) > DECODE_MAX_TOKENS: + uncovered("spec_tail", f"more than {DECODE_MAX_TOKENS} verify rows") + 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 attn_metadata.num_seqs * (draft_len + 1) > m3_sampled_accept_kernel.MAX_ROWS: + uncovered( + "spec_tail", + f"a sampled verify of more than {m3_sampled_accept_kernel.MAX_ROWS} rows", + ) + 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 + 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(): + 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) + 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] + 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) + 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()), + ) + + 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, and cap the drafter attention's trtllm-gen KV + split per request count (``DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH``). Called before the executor's warmup, on every + rank. Returns the number of low-M calls compiled.""" + if draft_model is None: + return 0 + draft_model.attn.dense_decode_max_seq_len_by_batch = DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH + 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: + 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: + 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), or of one-model standalone DSpark (``dspark.build_dspark``: the in-tree drafter + and worker).""" + spec = model_config.spec_config + if _dspark.is_standalone_dspark(spec): + shell.model.embed_tokens = _dspark.M3EmbedTokens(shell.model) + return _dspark.build_dspark(shell, model_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 ValueError( + "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) + + +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; + under DSpark with ``BOUNDARY_PACK``, both inside a batch-1 decode graph's tail. + * 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. + * HOSTTAIL (kill switches above ``HOSTTAIL_*``): ``sampler_d2h_side_stream`` (the sampler's store copies on its + D2H side stream), ``spec_small_h2d`` (static slot ids, packed accepted counts), ``spec_batch_map_changed_only`` + (the drafter's batch_indices / batch-to-slot uploads only on change), ``skip_accepted_draft_upload`` (no reader + of the accepted counts here), ``pack_spec_rng_uploads`` (Philox uploads in the pack), ``kv_lens_update_once`` + (one MSA KV-length re-derivation per speculative overlap step). + * ``pack_parallel_scatter`` (``PACK_PSCATTER``): the pack's non-overlapping copies scatter in parallel. + * ``pack_in_graph`` (``PACK_IN_GRAPH``) / ``pack_early_copy`` (``PACK_EARLY_COPY``): the decode graph's prologue + issues the pack, from a window uploaded on a side copy stream while the previous graph runs (no H2D copy or + scatter between the graphs). ``h2d_changed_span`` (``H2D_CHANGED_SPAN``): page-crossing steps stage only the + changed page-table entries, so their pack still fits the prologue's window. + * ``spec_tail_bookkeep`` (``TAILBK``): the sampler tail's no-op bookkeeping launches dropped, bitwise. + """ + return RuntimeHooks( + pack_step_inputs=True, + pack_parallel_scatter=PACK_PSCATTER, + pack_in_graph=PACK_IN_GRAPH, + pack_early_copy=PACK_EARLY_COPY, + skip_unchanged_h2d=True, + h2d_changed_span=H2D_CHANGED_SPAN, + loop_on_execution_stream=True, + merge_tuned_tactics_across_ranks=True, + skip_greedy_rng_upload=True, + spec_step_fold=spec_fold, + kv_lens_update_once=HOSTTAIL_KV_LENS_ONCE, + pack_spec_rng_uploads=HOSTTAIL_PHILOX_PACK, + spec_small_h2d=HOSTTAIL_SMALL_H2D, + skip_accepted_draft_upload=HOSTTAIL_SKIP_ACCEPTED_UPLOAD, + spec_batch_map_changed_only=HOSTTAIL_BATCH_CHANGED_ONLY, + spec_slot_ids_changed_only=SLOT_IDS_CHANGED, + sampler_d2h_side_stream=HOSTTAIL_D2H_SIDE, + skip_unchanged_draft_block_offsets=BLOCKOFF_GUARD, + spec_tail_bookkeep=TAILBK, + 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 ModelingV2MiniMaxM31Nvfp4Sm103Tp8( + 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, ( + "this target is built for sm_" + f"{_SM[0]}" + f"{_SM[1]}" + ", the device is sm_" + f"{''.join(map(str, torch.cuda.get_device_capability()))}" + ) + 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 + and _dspark.draft_kv_override_applies(spec_config) + ): + 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[nn.Module] = None + self.draft_config: Optional[ModelConfig] = None + self.spec_worker: Optional[Union[M3Eagle3Worker, _dspark.M3DSparkWorker]] = 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) + if not _dspark.is_standalone_dspark(spec_config): + spec_fold = M3SpecStepFold() + elif DSFOLD: + spec_fold = M3DSparkStepFold( + boundary_pack=BOUNDARY_PACK, + tokens_per_step=_static_tokens_per_step(spec_config) if BOUNDARY_PACK else None, + ) + self.boundary_pack_fold = ( + spec_fold + if isinstance(spec_fold, M3DSparkStepFold) and spec_fold.boundary_pack + else None + ) + self.runtime_hooks = _runtime_hooks(spec_fold) + _log_tailbk(self.spec_worker) + _install_topp_det() + _install_softmax_mc() + _install_rej_mc() + + 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.""" + if _dspark.is_standalone_dspark(self.spec_config): + _dspark.load_dspark_draft_weights(self, weights, weight_mapper) + return + _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() + _prebuild_topp_det() + _prebuild_softmax_mc() + _prebuild_rej_mc() + 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": + 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 + ) + if worker.tail_eligible(attn_metadata, spec_metadata, self.lm_head): + logits = None + else: + target_logits = getattr(worker, "target_logits", None) + logits = ( + target_logits(hidden_states, attn_metadata, spec_metadata, self.lm_head) + if target_logits is not None + else None + ) + if logits is None: + 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) + outputs = 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, + ) + if self.boundary_pack_fold is not None and spec_metadata is not None: + outputs = self.boundary_pack_fold.graph_tail(outputs, attn_metadata) + return outputs diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/mxdraft.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/mxdraft.py new file mode 100644 index 000000000000..da8a096aaa15 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/mxdraft.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MXDRAFT: the DSpark drafter's decode projections on this target's MXFP8 GEMV. + +The drafter checkpoint ships its linears as MXFP8 (E4M3 codes, one UE8M0 scale per 32 columns) and the stock drafter +dequantizes them to BF16 at load, so every decode pass streams twice the bytes it needs. MXDRAFT recovers the MXFP8 +form of each BF16 weight (exactly: the recovery is checked to dequantize back to the same BF16 bits, else the +projection is left alone) and routes the projection's GEMM to ``m3_mx_gemv_bf16`` for decode inputs it serves; any +other call runs the Linear unchanged. The activation is quantized to MXFP8 inside the GEMV, so the output is not the +stock BF16 GEMM's: drafter-only precision, gated on natural AL. Row-parallel projections keep +their all-reduce (only the Linear's GEMM is replaced). +""" + +from __future__ import annotations + +from typing import Iterable, Optional, Tuple + +import torch + +from .ops.comm.m3_mc_allreduce_norm import swizzle_sf +from .ops.mx_gemv import op as _gemv + +E4M3_MAX = 448.0 +BLOCK = 32 + + +def mxfp8_from_bf16(weight: torch.Tensor) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """The MXFP8 form (E4M3 codes [N, K], UE8M0 scale bytes [N, K / 32]) of a BF16 weight that is exactly an MXFP8 + dequantization, or None. Each block takes the smallest power-of-two scale that fits its amax in E4M3, which is at + most the shipped scale, so the codes are exact whenever the weight came from MXFP8.""" + if weight.dtype != torch.bfloat16 or weight.dim() != 2 or weight.shape[1] % BLOCK != 0: + return None + n, k = weight.shape + blocks = weight.float().view(n, k // BLOCK, BLOCK) + amax = blocks.abs().amax(-1) + exp = torch.ceil(torch.log2(amax.clamp_min(2.0**-126) / E4M3_MAX)).clamp(-127, 127) + scale = torch.exp2(exp).unsqueeze(-1) + codes = (blocks / scale).to(torch.float8_e4m3fn) + back = (codes.float() * scale).view(n, k).to(torch.bfloat16) + if not torch.equal(back, weight): + return None + return codes.view(n, k).contiguous(), (exp + 127).to(torch.uint8).contiguous() + + +def _projections(draft_model) -> Iterable[torch.nn.Module]: + for layer in getattr(getattr(draft_model, "model", None), "layers", []) or []: + attn, mlp = getattr(layer, "self_attn", None), getattr(layer, "mlp", None) + for lin in ( + getattr(attn, "qkv_proj", None), + getattr(attn, "o_proj", None), + getattr(mlp, "gate_up_proj", None), + getattr(mlp, "down_proj", None), + ): + if lin is not None: + yield lin + + +def route_mx_linears(draft_model, rows: Iterable[int]) -> int: + """Route each drafter projection whose BF16 weight is an exact MXFP8 dequantization to the MXFP8 GEMV for 2-D + contiguous bf16 inputs of the given row counts (the decode passes'); compiles those kernels now (outside capture). + Returns how many projections were routed.""" + rows = sorted({int(m) for m in rows if 0 < int(m) <= _gemv.MAX_TOKENS}) + n_routed = 0 + for lin in _projections(draft_model): + weight = getattr(lin, "weight", None) + if ( + weight is None + or getattr(lin, "bias", None) is not None + or getattr(lin, "_m31_mxdraft", None) is not None + ): + continue + rec = mxfp8_from_bf16(weight.data) + if rec is None: + continue + n, k = weight.shape + served = tuple(m for m in rows if _gemv.serves(n, k, m, True)) + if not served: + continue + w8, sf = rec[0], swizzle_sf(rec[1]) + for m in served: + _gemv.precompile(w8, sf, m, True) + original = lin.apply_linear + + def apply_linear( + input, + bias, + lora_params=None, + layer_idx=None, + *, + _o=original, + _w8=w8, + _sf=sf, + _served=served, + ): + if ( + bias is None + and not lora_params + and isinstance(input, torch.Tensor) + and input.dtype == torch.bfloat16 + and input.dim() == 2 + and input.shape[0] in _served + and input.is_contiguous() + ): + out = _gemv.m3_mx_gemv_bf16(input, _w8, _sf) + if out is not None: + return out + return _o(input, bias, lora_params, layer_idx) + + lin._m31_mxdraft = (w8, sf, served) + lin.apply_linear = apply_linear + n_routed += 1 + return n_routed diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/activation/m3_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/activation/m3_swq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/activation/m3_swq.py new file mode 100644 index 000000000000..5c8b970549f1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/activation/m3_swq.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SWQ: the shared expert's SwiGLU-OAI written directly as the down projection's MXFP8 activation.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +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 + + +@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 _swq_kernel( + x_ptr, + x_stride, + q_ptr, + sf_ptr, + M, + INTER: tl.constexpr, + NB: tl.constexpr, + NB_PAD: tl.constexpr, + P2: tl.constexpr, + swiglu_limit: tl.constexpr, + swiglu_alpha: tl.constexpr, + swiglu_beta: tl.constexpr, + HAS_LIMIT: tl.constexpr, +): + # One program per row of the 128-row-padded scale layout; rows past M only zero their scale bytes. + r = tl.program_id(0) + blk = tl.arange(0, P2) + blk_ok = blk < NB_PAD + # The 128x4 swizzle of trtllm::mxfp8_quantize(x, True) (quantization.cuh get_sf_out_offset_128x4). + sf_off = ( + (blk % 4) + + (blk // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * NB_PAD) + ) + if r >= M: + tl.store(sf_ptr + sf_off, tl.zeros([P2], dtype=tl.uint8), mask=blk_ok) + return + cols = tl.arange(0, P2 * 32) + valid = cols < INTER + row = x_ptr + r.to(tl.int64) * x_stride + # silu_and_mul_kernel's arithmetic (modules/swiglu.py), then its bf16 store. + a = tl.load(row + cols, mask=valid, other=0.0).to(tl.float32) + b = tl.load(row + cols + INTER, mask=valid, other=0.0).to(tl.float32) + if HAS_LIMIT: + a = tl.minimum(a, swiglu_limit) + b = tl.clamp(b, -swiglu_limit, swiglu_limit) + v = (a * tl.sigmoid(swiglu_alpha * a) * (b + swiglu_beta)).to(tl.bfloat16).to(tl.float32) + # cvt_warp_fp16_to_mxfp8: SF = amax * rcp.approx(448), UE8M0 toward +inf, scale = rcp.approx(SF), E4M3 RN. + v2 = tl.reshape(v, (P2, 32)) + amax = tl.max(tl.abs(v2), axis=1) + sf = amax * _rcp_approx_ftz(tl.full([P2], 448.0, tl.float32)) + bits = sf.to(tl.int32, bitcast=True) + exp = (bits >> 23) & 0xFF + mant = bits & 0x7FFFFF + e = tl.where(exp > 0, exp + (mant != 0).to(tl.int32), tl.where(bits > 0x400000, 1, 0)) + e = tl.minimum(e, 254) + # at code 0 the scale 2^-127 is a float subnormal, so rcp.approx.ftz returns inf (native publishes + # inf/NaN there). Use its exact reciprocal 2^127 instead; every other code keeps + # native's rcp.approx (exact on a normal power of two), so each block native emits finite stays bitwise. + pinv = tl.where( + e == 0, 1.7014118346046923e38, _rcp_approx_ftz((e << 23).to(tl.float32, bitcast=True)) + ) + inv = tl.where(amax != 0.0, pinv, 0.0) + q = tl.reshape((v2 * inv[:, None]).to(tl.float8e4nv), (P2 * 32,)) + tl.store(q_ptr + r.to(tl.int64) * INTER + cols, q, mask=valid) + tl.store(sf_ptr + sf_off, tl.where(blk < NB, e, 0).to(tl.uint8), mask=blk_ok) + + +class M3Swq(OpWrapper): + """``silu_and_mul`` (SwiGLU-OAI) over ``x = [g | u]`` (bf16 [M, 2I]) rounded to bf16, then MXFP8-quantized as + ``trtllm::mxfp8_quantize(h, True)``: (E4M3 [M, I], UE8M0 scales in the 128x4-swizzled layout), bit for bit, in one + launch. The TP8 shared expert's down projection (intermediate 384, no GEMV row table) reads it at ``SWQ_ROWS``.""" + + ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert, the smallest row count it serves", spec=dict(m=9, inter=384)), + Cell(why="shared expert, a 32-row decode verify (c8)", spec=dict(m=32, inter=384)), + Cell( + why="shared expert, the widest decode row count it serves", spec=dict(m=48, inter=384) + ), + ) + + note = """ + Bitwise to the two-kernel path (silu_and_mul's fp32 arithmetic and bf16 rounding, then mxfp8_quantize's + rcp.approx scale rule): q bytes and every scale byte, padding included. + One deliberate exception: a block whose scale code is 0 (amax / 448 <= 2^-127), where native's + rcp.approx.ftz of the subnormal scale is inf and its e4m3 bytes encode inf/NaN, is quantized with the exact + reciprocal 2^127 instead, so it stays finite; its scale byte is native's. + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> tuple[torch.Tensor, torch.Tensor]: + m, i = x.shape[0], x.shape[1] // 2 + nb = i // 32 + nb_pad = (nb + 3) // 4 * 4 + q = torch.empty(m, i, dtype=torch.float8_e4m3fn, device=x.device) + sf = torch.empty((m + 127) // 128 * 128 * nb_pad, dtype=torch.uint8, device=x.device) + _swq_kernel[((m + 127) // 128 * 128,)]( + x, + x.stride(0), + q, + sf, + m, + INTER=i, + NB=nb, + NB_PAD=nb_pad, + P2=triton.next_power_of_2(nb_pad), + swiglu_limit=limit or 0.0, + swiglu_alpha=alpha if alpha is not None else 1.0, + swiglu_beta=beta if beta is not None else 0.0, + HAS_LIMIT=limit is not None and limit > 0.0, + ) + return q, sf + + def reference( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> tuple[torch.Tensor, torch.Tensor]: + """The two kernels it replaces (bitwise is the contract).""" + h = torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + q, sf = torch.ops.trtllm.mxfp8_quantize(h, True) + return q.view(h.shape), sf + + def is_valid( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> None: + assert x.dtype == torch.bfloat16 and x.stride(-1) == 1, "bf16 [M, 2I] with contiguous rows" + assert (x.shape[1] // 2) % 32 == 0, "I must be a multiple of the 32-element MXFP8 block" + + +m3_swq = M3Swq() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/attention/m3_dense_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/attention/m3_dense_decode.py new file mode 100644 index 000000000000..525de997aaf9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/attention/m3_dense_decode.py @@ -0,0 +1,187 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' decode attention (trtllm-gen over the full page table) with the KV split this target sizes per +request count (``dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH``).""" + +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 +#: The split this target hands trtllm-gen at one request (dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH), and a model +#: maximum the in-tree step passes otherwise. +CAPPED = 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): + """``minimax_m3_trtllm_gen_dense_decode`` as ``MsaDecodeFmha`` calls it for a dense layer, with the max_seq_len + ``msa_decode.dense_decode_max_seq_len`` resolves: this target's per-request-count cap, else the step's bound. + + Decode attention of every query row over its request's full causal context in the paged FP8 pool, written into + ``output`` (BF16 [T, heads, head_dim]) in place. ``max_seq_len`` only sizes trtllm-gen's multi-CTA KV split: + every CTA still covers its share of each request's actual length, which may exceed it. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over a 20000-token context, the split sized for 16384", + spec=dict(kv_lens=(20000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request, one row, over a 65536-token context (an AgentX-length cache) at the capped split", + spec=dict(kv_lens=(65536,), qlen=1, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request verifying 1 + 3 rows over a 3000-token context: the cap above the length", + spec=dict(kv_lens=(3000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="two requests keep the model maximum (no cap at two requests)", + spec=dict(kv_lens=(20000, 5000), qlen=4, max_seq_len=MODEL_MAX, capped=False), + ), + ) + + note = """ + The split only reorders the merge of the CTAs' partial softmax sums, so the capped split's output is the uncapped + one's to a few BF16 ulps (the test holds it there), and both are held to the fp32 reference within the FP8-KV + band 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" + # 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_m31_nvfp4__sm_103__tp8/ops/attention/m3_draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/attention/m3_fused_score_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/attention/m3_index_score.py new file mode 100644 index 000000000000..139b216086eb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/attention/m3_index_score.py @@ -0,0 +1,165 @@ +# 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="one request verifying 1 + 7 rows (a 7-draft verify the fused score + select does not take)", + spec=dict(seq_lens=(20000,), qlen=8, early=False), + ), + Cell( + why="two requests verifying 1 + 7 rows, one longer than 512 blocks (the CTAs loop)", + spec=dict(seq_lens=(70000, 3000), qlen=8, early=False), + ), + Cell( + why="one request verifying 1 + 7 rows on the early-issue schedule (the K7 batch-1 decode step)", + spec=dict(seq_lens=(131072,), qlen=8, early=True), + ), + 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_m31_nvfp4__sm_103__tp8/ops/attention/m3_msa_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/attention/m3_msa_attention.py new file mode 100644 index 000000000000..60990ce4a415 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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 at their capped KV split). ``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_m31_nvfp4__sm_103__tp8/ops/attention/m3_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/attention/m3_producer.py new file mode 100644 index 000000000000..783304219ae5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/attention/m3_producer.py @@ -0,0 +1,202 @@ +# 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 K7 c1 verify: 8 tokens", spec=dict(tokens=8, 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 before R1: 48 tokens", spec=dict(tokens=48, scale=1.0)), + Cell( + why="the widest decode step (R1): a K7 c8 verify, 64 tokens, a skipped and an out-of-pool slot", + spec=dict(tokens=64, scale=40.0, bad_slots=True), + ), + 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_m31_nvfp4__sm_103__tp8/ops/attention/m3_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/attention/m3_qkv_insert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/attention/m3_select_blocks.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/attention/m3_select_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/attention/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/attention/m3_select_long_dsl.py new file mode 100644 index 000000000000..425fe62da388 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/attention/m3_select_long_dsl.py @@ -0,0 +1,126 @@ +# 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, one row per 384-thread CTA (CuTe DSL).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_select_long import _M3, MAX_BLOCKS, M3SelectLong + + +class M3SelectLongDsl(OpWrapper): + """``m3_select_long``'s contract, cells and reference, run by ``ops/m3_select_long_dsl``: the same selection + bit for bit as the 1024-thread ``minimax_m3_select_long`` (also on rows where more than 2048 blocks tie + at the 16th-largest score), from one 384-thread CTA per row. Returns ``out``.""" + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 4 rows over 165521 tokens (1294 blocks, the AgentX ISL): the 2-key path", + spec=dict(num_blocks=1294, total_q=4, valid="near_full", scores="random", **_M3), + ), + Cell( + why="5888-block rows at two requests x 4 rows: the 9-key path", + spec=dict(num_blocks=5888, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the longest covered rows (9216 blocks, 1.18M tokens), 2 requests x 4 rows: the 9-key path at its end", + spec=dict(num_blocks=9216, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="~1500 blocks tied at the 16th score: the in-tree op also takes the smaller blocks (bit for bit)", + spec=dict(num_blocks=2000, total_q=8, valid="near_full", scores="ties1500", **_M3), + ), + Cell( + why="two requests x 4 rows of mixed lengths, 4500-block engine: every key width, short and padding rows", + spec=dict(num_blocks=4500, total_q=8, valid="mixed", scores="random", **_M3), + ), + Cell( + why="duplicated scores: ties broken to the smaller block, and > 2048 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), + ), + ) + + 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, + trigger_phase: int = 0, + ) -> torch.Tensor: + """``m3_select_long``'s reference (which takes no trigger phase): the selection is the same bit for bit.""" + return M3SelectLong.reference( + self, scores, n_valid_blocks, out, init_blocks, local_blocks, use_pdl, early_trigger + ) + + 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, + trigger_phase: int = 0, + ) -> None: + """``m3_select_long``'s domain.""" + M3SelectLong.is_valid( + self, scores, n_valid_blocks, out, init_blocks, local_blocks, use_pdl, early_trigger + ) + + note = """ + Certified at one KV head per rank (TP8, GB300) with M3's forcing + (`init_blocks=0`, `local_blocks=1`), and bit for bit against the + 1024-thread `m3_select_long` on every cell. The kernel takes any head count + and forcing; other forcing is not driven here. + + The kernel is compiled per (PDL launch, early trigger, trigger phase) with + symbolic shapes and strides; the target compiles its keys before warmup + (`ops/m3_select_long_dsl.prebuild`), and a call outside them compiles + on first use. + """ + + 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, + trigger_phase: int = 0, + ) -> torch.Tensor: + from ..m3_select_long_dsl import select_long_dsl + + if scores.shape[1] > MAX_BLOCKS: + raise RuntimeError( + f"m3_select_long_dsl supports at most {MAX_BLOCKS} blocks, got {scores.shape[1]}" + ) + select_long_dsl( + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + return out + + +m3_select_long_dsl = M3SelectLongDsl() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode_lc8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode_lc8.py new file mode 100644 index 000000000000..9c5912667cde --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode_lc8.py @@ -0,0 +1,499 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3.1 TP8 implementation.""" + +from __future__ import annotations + +import os +from typing import Optional, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.fakes import ( + m31_fake_ops as _m31_fake_ops, +) + +from ...m31_switches import m31_disabled + +LC8_SYNC3 = not m31_disabled("LC8_SYNC3") +_LC8_STATE = {} + + +def _lc8_launcher(): + """The sparse decode op the entry launches, decided once per process: the SYNC3 copy when switched on and built + (or overlaid), else ``minimax_m3_tp8_sparse_decode_lc8``.""" + if "op" not in _LC8_STATE: + op = None + if LC8_SYNC3: + op = getattr(torch.ops.trtllm, "minimax_m31_tp8_sparse_decode_lc8", None) + if op is None: + import tensorrt_llm + + lib = os.path.join( + os.path.dirname(tensorrt_llm.__file__), "libs", "libm31_lc8sync3.so" + ) + if os.path.exists(lib): + torch.ops.load_library(lib) + _m31_fake_ops.register() + op = getattr(torch.ops.trtllm, "minimax_m31_tp8_sparse_decode_lc8", None) + from tensorrt_llm.logger import logger + + reason = ( + "engaged" + if op is not None + else "switched off" + if not LC8_SYNC3 + else "not taken: the op is not built" + ) + logger.info(f"[M3 dispatch] LC8_SYNC3 sparse decode softmax: {reason}") + if op is not None: + logger.info(f"[M3 dispatch] LC8 one-division merge: one_div={False}") + _LC8_STATE["m31"] = op is not None + _LC8_STATE["op"] = ( + op if op is not None else torch.ops.trtllm.minimax_m3_tp8_sparse_decode_lc8 + ) + return _LC8_STATE["op"] + + +def _lc8_onediv(rows: int) -> bool: + _lc8_launcher() + return False + + +NUM_HEADS = 8 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +SF_VEC = 16 +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) + nibbles = torch.stack([rows & 15, 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 & 255) + (bits & 8388607 != 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="7-draft verify at one request: T=8, log-domain recurrence, prewait at its row limit", + spec=dict( + seq_lens=(20000,), + decode_query_len=8, + linear_softmax=False, + mxfp8=True, + prewait=True, + ), + ), + Cell( + why="7-draft verify at two requests: T=16 on the linear denominator, past the prewait row limit", + spec=dict( + seq_lens=(5000, 12345), + decode_query_len=8, + linear_softmax=True, + mxfp8=True, + prewait=False, + ), + ), + 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 = ( + "\n Certified: `linear_softmax` x hand-off on / off x `prewait`, the " + "variants\n the cells drive. The hand-off cells are certified " + "bit-exactly against\n `mxfp8_handoff_reference` applied to the op's " + "own BF16 rows; the BF16 rows\n themselves within the band the test " + "derives, and the two `prewait` variants\n against each other bit for" + " bit.\n\n `topk` is read for the first `min(16, ceil(kv_len / 128))` " + "entries of each\n token; a -1 among them skips that CTA's block, " + "entries past them are never\n read. `use_pdl` only changes the " + "launch attribute. The hand-off assumes\n `num_sf_k_tiles = 8` (a " + "1024-wide o_proj input) and buffers with at least\n as many rows as " + "the step; `MERGE_MXFP8` sizes them at 128 rows.\n " + ) + + 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: + one_div = _lc8_onediv(int(q.shape[0])) + _lc8_launcher()( + 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, + **{"one_div": True} if one_div else {}, + ) + 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()), ( + "token " + f"{tok}" + ": selected block ids must be -1 or below the request's block cou" + "nt; the block table is indexed with them unchecked" + ) + if q8 is not None: + assert num_sf_k_tiles == WIDTH // 128, ( + "num_sf_k_tiles must be " + f"{WIDTH // 128}" + ": the hand-off rows are " + f"{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_m31_nvfp4__sm_103__tp8/ops/attention/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/comm/m3_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/embedding/m3_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/gemm/m3_bf16_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/gemm/m3_fp32_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/gemm/m3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/gemm/m3_mxfp8_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/gqamix_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/gqamix_sparse_decode.py new file mode 100644 index 000000000000..e6b824de1b96 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/gqamix_sparse_decode.py @@ -0,0 +1,1311 @@ +# 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. +""" + +from __future__ import annotations + +import os +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.memory_buffer_utils import get_memory_buffers + +# One sparse block is exactly one KV page. +SPARSE_BLOCK_SIZE = 128 + +_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 + +# 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 + + +def _pdl_enabled() -> bool: + return os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +@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"]), + } +) +@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 + 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 + 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 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() + + 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 _ 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) + + # 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.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.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, +): + pid_b, pid_h = tl.program_id(0), tl.program_id(1) + + if USE_PDL: + 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 USE_PDL: + 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: tuple[int, int], +) -> Optional[int]: + """Select GB300-measured split-K factors for M3's SM100-family decode shape.""" + 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: tuple[int, int], +) -> bool: + """Use the GB300-measured lower-register recurrence on validated M3 shapes.""" + 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: tuple[int, int], +) -> int: + """Group Eagle queries using the GB300-measured shared-page policy.""" + 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: tuple[int, int], +) -> 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 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: tuple[int, int], +) -> dict[str, int]: + """Select GB300-measured one-warp merges for M3's SM100-family graphs.""" + 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}." + ) + + +@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, +) -> 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 + + 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, + capability=capability, + ) + 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, + capability=capability, + ) + 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) + ) + elif num_topk_chunks & (num_topk_chunks - 1): + raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") + + 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 = _pdl_enabled() + 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, + capability=capability, + ) + 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, + capability=capability, + ) + 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 + ) + pair_kernel_options = ( + {"QUERY_GROUP_SIZE": query_group_size, "MAX_TOPK": max_topk} + if query_group_size == 2 + 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, + 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), + 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, + **pair_kernel_options, + **pdl_launch, + **launch_options, + ) + if not direct_output: + 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, + capability=capability, + ) + 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, + **pdl_launch, + **merge_launch_options, + ) + + +__all__ = [ + "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_m31_nvfp4__sm_103__tp8/ops/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/long_select.py new file mode 100644 index 000000000000..8a028e1cc288 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/long_select.py @@ -0,0 +1,77 @@ +# 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, or +with ``dsl`` the 384-thread CuTe DSL kernel (``ops/m3_select_long_dsl``, entry ``m3_select_long_dsl``): the same table. +""" + +from typing import Optional + +import torch + +from ..ops.attention.m3_select_long import MAX_BLOCKS, m3_select_long +from ..ops.attention.m3_select_long_dsl import m3_select_long_dsl + + +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, + dsl: 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). ``dsl`` runs the 384-thread kernel (``SELECT_LONG_DSL``; + the same table), released at its exit unless ``early_trigger``. + """ + 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) + if dsl: + return m3_select_long_dsl( + max_score_kv, + n_valid_blocks.contiguous(), + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + 0, + ) + 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_m31_nvfp4__sm_103__tp8/ops/m3_attn_lc8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/m3_bp_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_bp_tail.py new file mode 100644 index 000000000000..419a5baef898 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_bp_tail.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""BOUNDARY_PACK BP-1's graph tail for a batch-1 decode step, in one launch: the spec sampler's four slot stores and +the engine's next-step gathers, both gated by a device-side live word. + +``ctl`` (int64 [2]) holds (slot, live), written by the step fold's ``begin_step`` before the replay. With ``live`` 0 +the launch writes nothing (a step the fold did not take: a dummy-only replay, a step it could not identify). With +``live`` 1, for the one row: + +* the stores, as the sampler's ``_store_outputs`` (zero padding / truncation to the store widths, then index_copy_ + at ``slot``):: + + new_tokens[slot][c] = src_new[c] if c < src width else 0 (c < store width), likewise + next_new_tokens[slot][c] , next_draft_tokens[slot][c]; new_tokens_lens[slot] = src_lens + +* with ``GATHER``, the engine's gathers for the next step of the same row (P = 1, no rows before it), as + ``step_gather`` reads them back from the store just written:: + + ids[t] = next_new_tokens[slot][t] (t < T) draft[d] = next_draft_tokens[slot][d] (d < T - 1) + pos_off[i] = lens[slot] if i < T else 0 kv_off[i] = lens[slot] - T if i < 1 else 0 + (every i of pos_off / kv_off: the engine zeroes both buffers before its gathers) + +Integer copies only. Program 0 does the row; every program zero-fills its ``ZBLOCK`` chunk of the offsets. The int +arguments are not specialized and every pointer is 16-byte aligned (the caller checks), so the two ``GATHER`` +variants compiled by ``warm_up`` outside capture serve every graph. +""" + +import torch +import triton +import triton.language as tl + +ZBLOCK = 1024 +_INTS = ( + "s_new_s1", + "s_new_w", + "d_new_s0", + "d_new_s1", + "d_new_w", + "s_nnew_s1", + "s_nnew_w", + "d_nnew_s0", + "d_nnew_s1", + "d_nnew_w", + "d_lens_s0", + "s_draft_s1", + "s_draft_w", + "d_draft_s0", + "d_draft_s1", + "d_draft_w", + "n_pos", + "n_kv", + "T", +) + + +@triton.jit(do_not_specialize=list(_INTS)) +def _bp_tail_kernel( + ctl, + s_new, + s_new_s1, + s_new_w, + d_new, + d_new_s0, + d_new_s1, + d_new_w, + s_nnew, + s_nnew_s1, + s_nnew_w, + d_nnew, + d_nnew_s0, + d_nnew_s1, + d_nnew_w, + s_lens, + d_lens, + d_lens_s0, + s_draft, + s_draft_s1, + s_draft_w, + d_draft, + d_draft_s0, + d_draft_s1, + d_draft_w, + ids, + draft, + pos_off, + n_pos, + kv_off, + n_kv, + T, + GATHER: tl.constexpr, + BLOCK: tl.constexpr, + ZBLOCK_: tl.constexpr, +): + slot = tl.load(ctl) + live = tl.load(ctl + 1) + if live != 0: + pid = tl.program_id(0) + lens = tl.load(s_lens) + if pid == 0: + c = tl.arange(0, BLOCK) + new = tl.load(s_new + c * s_new_s1, mask=c < s_new_w, other=0) + tl.store(d_new + slot * d_new_s0 + c * d_new_s1, new, mask=c < d_new_w) + nnew = tl.load(s_nnew + c * s_nnew_s1, mask=c < s_nnew_w, other=0) + tl.store(d_nnew + slot * d_nnew_s0 + c * d_nnew_s1, nnew, mask=c < d_nnew_w) + dr = tl.load(s_draft + c * s_draft_s1, mask=c < s_draft_w, other=0) + tl.store(d_draft + slot * d_draft_s0 + c * d_draft_s1, dr, mask=c < d_draft_w) + tl.store(d_lens + slot * d_lens_s0, lens) + if GATHER: + tl.store(ids + c, nnew, mask=c < T) + tl.store(draft + c, dr, mask=c < T - 1) + if GATHER: + z = pid * ZBLOCK_ + tl.arange(0, ZBLOCK_) + tl.store(pos_off + z, tl.where(z < T, lens, 0), mask=z < n_pos) + tl.store(kv_off + z, tl.where(z < 1, lens - T, 0), mask=z < n_kv) + + +def block_for(store) -> int: + """The row block of a sampler store: every store width fits it (fixed per store, so per fold).""" + widths = ( + store.new_tokens.shape[0], + store.next_new_tokens.shape[0], + store.next_draft_tokens.shape[1], + ) + return triton.next_power_of_2(max(*widths, 1)) + + +def aligned(*tensors: torch.Tensor) -> bool: + """Every pointer the launch takes is 16-byte aligned (the compiled variants' assumption).""" + return all(t.data_ptr() % 16 == 0 for t in tensors) + + +def bp_tail(ctl, src, store, engine, t: int, gather: bool, block: int) -> None: + """One launch: ``src`` = (new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens) row 0 of the worker's + int32 outputs ([1, w] / [1]); ``store`` the sampler's store; ``engine`` = (input_ids, draft_tokens, + position_id_offsets, kv_lens_offsets) int32 contiguous; ``t`` tokens per request (gather only).""" + s_new, s_nnew, s_lens, s_draft = src + d_new = store.new_tokens.squeeze(-1).T + d_nnew = store.next_new_tokens.squeeze(-1).T + d_lens, d_draft = store.new_tokens_lens, store.next_draft_tokens + ids, draft, pos_off, kv_off = engine + n_pos, n_kv = pos_off.numel(), kv_off.numel() + grid = (max(1, triton.cdiv(max(n_pos, n_kv), ZBLOCK)) if gather else 1,) + _bp_tail_kernel[grid]( + ctl, + s_new, + s_new.stride(1), + s_new.shape[1], + d_new, + d_new.stride(0), + d_new.stride(1), + d_new.shape[1], + s_nnew, + s_nnew.stride(1), + s_nnew.shape[1], + d_nnew, + d_nnew.stride(0), + d_nnew.stride(1), + d_nnew.shape[1], + s_lens, + d_lens, + d_lens.stride(0), + s_draft, + s_draft.stride(1), + s_draft.shape[1], + d_draft, + d_draft.stride(0), + d_draft.stride(1), + d_draft.shape[1], + ids, + draft, + pos_off, + n_pos, + kv_off, + n_kv, + t, + GATHER=gather, + BLOCK=block, + ZBLOCK_=ZBLOCK, + num_warps=4, + ) + + +def warm_up(ctl, store, engine, block: int) -> None: + """Compile both variants outside capture, with ``ctl`` live 0 (the launch writes nothing).""" + src = ( + store.new_tokens.squeeze(-1).T[:1], + store.next_new_tokens.squeeze(-1).T[:1], + store.new_tokens_lens[:1], + store.next_draft_tokens[:1], + ) + for gather in (False, True): + bp_tail(ctl, src, store, engine, 2, gather, block) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/m3_fused_score_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_index_score.py new file mode 100644 index 000000000000..c6a7f6d80080 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_index_score.py @@ -0,0 +1,700 @@ +# 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 four 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. +* ``prestage_old`` (with ``early_issue``, ``INDEX_SCORE_PRESTAGE_OLD``): when the CTA's first block holds only + tokens older than this step (block < (seq_lens hint - decode query length) / 128), its K tile is TMA-loaded into its + stage before the wait; Q (and a first block that is not all old) is still loaded after it. +* ``early_trigger`` (``INDEX_SCORE_SET``): every CTA releases the dependent block selector right after its + setup, so the selector launches and makes its pre-wait reads while this grid scores; the selector still waits for + this whole grid before it reads a score, so the scores and the selection are unchanged. + +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.cute_dsl_kernels.blackwell.cute_ptx_utils import ( + EVICT_FIRST, + fp8x4_to_fp16x4, + mma_sync, + simple_tma_copy, +) +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import ( + TRTLLM_ENABLE_PDL, + griddepcontrol_launch_dependents, + griddepcontrol_wait, +) +from tensorrt_llm.logger import logger + +from ..m31_switches import m31_disabled + +# PRESTAGE_OLD (ported from the sm_107 TP8 target, the M3 Rubin Indexer's G1): with the early issue, a CTA whose first +# block holds only tokens older than this step also TMA-loads that block's index-K tile before the grid wait +# (bit-identical scores). Kill switch: TRTLLM_M31_DISABLE_PRESTAGE_OLD=1. +INDEX_SCORE_PRESTAGE_OLD = not m31_disabled("PRESTAGE_OLD") +# SET (ported from the sm_107 TP8 target): every CTA releases the dependent block selector right after its setup +# instead of at its exit (the selector's own grid wait still returns only once the whole scoring grid is complete, so +# the picks are unchanged), at every key. Kill switch: TRTLLM_M31_DISABLE_SCORE_SET=1. +INDEX_SCORE_SET = not m31_disabled("SCORE_SET") + +__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 # capacity: fixed: the MSA block size of the checkpoint + 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, + prestage_old: bool = False, + early_trigger: bool = False, + ): + # 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 + # prestage_old (with early_issue): when this CTA's first block holds only tokens older than this step, its K + # tile is TMA-loaded into its stage BEFORE the grid dependency wait (the predecessor writes only this step's + # tokens; the block table and the seq_lens hint are final before the grid starts). Q stays after the wait. + self.prestage_old = prestage_old and early_issue + # early_trigger: release the dependent block selector right after this CTA's setup (see the module docstring). + self.early_trigger = early_trigger + 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.early_trigger): + griddepcontrol_launch_dependents() + + 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)), + ) + + # Without early_trigger the dependent block selector is released as the CTAs exit. + early_q = False + early_k = False + early_page = Int32(0) + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + seqlen_hint = seq_lens[batch_id] + hint_blocks = cute.ceil_div(seqlen_hint, 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] + if cutlass.const_expr(self.prestage_old): + # Stage the first K tile now when every token of the block predates this step. + if split_id < (seqlen_hint - decode_query_len) // BLOCK_K: + early_k = True + 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, + ) + + 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) + if not early_k: + 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, + ) + + # The length that bounds this grid's work is read after the wait (the + # pre-wait reads above are prefetch / early-issue hints only, clamped to + # the table width). seq_lens 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; ``early_trigger`` is ``INDEX_SCORE_SET`` for every key of the process.""" + + 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.cute_dsl_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.cute_dsl_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, + prestage_old=early and INDEX_SCORE_PRESTAGE_OLD, + early_trigger=INDEX_SCORE_SET, + ) + if INDEX_SCORE_SET and not cls.kernel_cache: + logger.info( + "[M3 dispatch] SCORE_SET tp8 index scorer: the block selector is released after the scorer's setup" + ) + 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_m31_nvfp4__sm_103__tp8/ops/m3_layer/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_layer/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_layer/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_layer/phase_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_layer/phase_producer.py new file mode 100644 index 000000000000..02f3f91b7ed8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_layer/phase_producer.py @@ -0,0 +1,489 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode producer as a phase of the TP8 per-layer kernel (m3_layer): q / k / index-q / index-k RMSNorm + NeoX +partial RoPE, NVFP4 K / V and FP8 index-K cache insert, FP8 q / index-q out. + +Bit for bit ``torch.ops.trtllm.minimax_m3_producer_pw`` (cpp/tensorrt_llm/kernels/minimaxM3/producerPw.cu, +m3_producer_kernel<8, *>) at the TP8 geometry: 8 query heads, 1 KV head, 1 index head, head_dim 128, rotary_dim 64, +page 128. One work item = one decode token; the CTA's warp w (< 12) owns packed head w of it, lane l holds elements +[4l, 4l + 4) exactly as the C++ kernel, so every reduction (the RMS sum of squares, the RoPE pair exchange, the +NVFP4 block max) runs the same shuffle tree in the same order. Every floating-point operation is written as its PTX +instruction (mul.rn / add.rn / fma.rn / rsqrt.approx / rcp.approx.ftz / cvt) so the DSL compiler cannot contract +or reassociate differently from the library build (--use_fast_math). The places where that build contracted are +marked CONTRACT and follow the shipped SASS (msa/results/sassdump-*). + +Packed row layout per token: [8 q | k | v | index_q | index_k] x 128 BF16 (the QKV GEMV output rounded to BF16). +""" + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op + +HEADS_Q = 8 # capacity: fixed: the tp8 query heads per rank +PACKED_HEADS = HEADS_Q + 4 # q x 8, k, v, index_q, index_k +HEAD_DIM = 128 +ROTARY_DIM = 64 # capacity: fixed: the checkpoint's rotary dim +PAGE = 128 +ELEMS = HEAD_DIM // 32 # capacity: fixed: elements per lane (per lane) +K_HEAD, V_HEAD, IQ_HEAD, IK_HEAD = HEADS_Q, HEADS_Q + 1, HEADS_Q + 2, HEADS_Q + 3 + + +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, + ) + + +# ---- exact float arithmetic (no contraction) --------------------------------------------------------------------- +@dsl_user_op +def f_mul(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", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def f_add(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", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def f_fma(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", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def f_max(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "max.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def f_abs(a, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "abs.f32 $0, $1;", "=f,f", loc=loc, ip=ip) + ) + + +@dsl_user_op +def f_neg(a, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "neg.f32 $0, $1;", "=f,f", loc=loc, ip=ip) + ) + + +@dsl_user_op +def f_rsqrt(a, *, loc=None, ip=None): + # NOT ftz on this target: the sm_103 library build of producerPw.cu runs rsqrt.approx.f32 with denormal scaling + # (MSA's SASS read; edge row denormal_rms with eps 0, differs under .ftz). The sm_107 copy keeps .ftz. + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rsqrt.approx.f32 $0, $1;", "=f,f", loc=loc, ip=ip) + ) + + +@dsl_user_op +def f_rcp(a, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", loc=loc, ip=ip) + ) + + +@dsl_user_op +def bf16_bits_to_f32(bits16, *, loc=None, ip=None): + """The float value of a BF16 given in the low 16 bits of an Int32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(bits16)], + "{\n\t.reg .b32 t;\n\tshl.b32 t, $1, 16;\n\tmov.b32 $0, t;\n\t}", + "=f,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def round_bf16(a, *, loc=None, ip=None): + """__bfloat162float(__float2bfloat16_rn(a)).""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a)], + "{\n\t.reg .b16 h;\n\t.reg .b32 t;\n\tcvt.rn.bf16.f32 h, $1;\n\tmov.b32 t, {0, h};\n\tmov.b32 $0, t;\n\t}", + "=f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def e4m3x4_word(e0, e1, e2, e3, *, loc=None, ip=None): + """store_fp8x4: bytes [e0, e1, e2, e3] (cvt.rn.satfinite.e4m3x2.f32 puts its first source in the high byte).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(e0), cutlass.Float32(e1), cutlass.Float32(e2), cutlass.Float32(e3)], + "{\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", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def e4m3_byte_and_value(x, *, loc=None, ip=None): + """(byte, float value) of __nv_fp8_e4m3(x): the encoded scale and static_cast of it.""" + byte = cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n\t.reg .b16 p;\n\t.reg .b32 w;\n\tcvt.rn.satfinite.e4m3x2.f32 p, 0f00000000, $1;\n\t" + "cvt.u32.u16 w, p;\n\tand.b32 $0, w, 0xFF;\n\t}", + "=r,f", + loc=loc, + ip=ip, + ) + ) + value = cutlass.Float32( + _asm( + _T.f32(), + [byte], + "{\n\t.reg .b16 p;\n\t.reg .b32 h2;\n\t.reg .b16 h;\n\tcvt.u16.u32 p, $1;\n\t" + "cvt.rn.f16x2.e4m3x2 h2, p;\n\tmov.b32 {h, _}, h2;\n\tcvt.f32.f16 $0, h;\n\t}", + "=f,r", + loc=loc, + ip=ip, + ) + ) + return byte, value + + +@dsl_user_op +def e2m1x4_halfword(q0, q1, q2, q3, *, loc=None, ip=None): + """The low 16 bits of fp32x8_to_e2m1 for q[0..3] (q[4..7] = 0): byte0 = (q1, q0), byte1 = (q3, q2).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(q0), cutlass.Float32(q1), cutlass.Float32(q2), cutlass.Float32(q3)], + "{\n\t.reg .b8 b0, b1;\n\t.reg .b32 w;\n\tcvt.rn.satfinite.e2m1x2.f32 b0, $2, $1;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 b1, $4, $3;\n\tmov.b32 w, {b0, b1, 0, 0};\n\tmov.b32 $0, w;\n\t}", + "=r,f,f,f,f", + loc=loc, + ip=ip, + ) + ) + + +# ---- global memory ----------------------------------------------------------------------------------------------- +@dsl_user_op +def ld_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm(_T.i64(), [cutlass.Int64(addr)], "ld.global.u64 $0, [$1];", "=l,l", True, loc, ip) + ) + + +@dsl_user_op +def ld_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm(_T.i32(), [cutlass.Int64(addr)], "ld.global.s32 $0, [$1];", "=r,l", True, loc, ip) + ) + + +@dsl_user_op +def ld_f32(addr, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Int64(addr)], "ld.global.f32 $0, [$1];", "=f,l", True, loc, ip) + ) + + +@dsl_user_op +def st_u32(addr, v, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(v)], + "st.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def st_u16(addr, v, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(v)], + "{\n\t.reg .b16 h;\n\tcvt.u16.u32 h, $1;\n\tst.global.u16 [$0], h;\n\t}", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def st_u8(addr, v, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(v)], + "{\n\t.reg .b16 h;\n\tcvt.u16.u32 h, $1;\n\tst.global.u8 [$0], h;\n\t}", + "l,r", + True, + loc, + ip, + ) + + +# ---- CONTRACT points (follow the shipped SASS; see the module docstring) ------------------------------------------- +def sumsq_pair(acc, x, y): + """sumSquares += x * x + y * y (C++: one pair of the packed row). CONTRACT: fma(y, y, x * x) then add.""" + return f_add(acc, f_fma(y, y, f_mul(x, x))) + + +def rms_reciprocal(sum_squares, eps): + """rsqrtf(sumSquares / 128 + eps). CONTRACT: the / 128 is an exact * 2^-7; fused with + eps into one fma.""" + return f_rsqrt(f_fma(sum_squares, cutlass.Float32(1.0 / HEAD_DIM), eps)) + + +class ProducerArgs: + """Raw device addresses (int64) + strides, as the C++ Params (producerPw.h) carries them.""" + + def __init__( + self, + packed, + q_out, + index_q_out, + index_k, + kv_data, + kv_scale, + out_cache_loc, + kv_quant_scale, + position_ids, + rotary_cos_sin, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + eps, + num_pages, + index_page_stride, + index_token_stride, + data_page_stride, + data_plane_stride, + data_token_stride, + scale_page_stride, + scale_plane_stride, + ): + self.__dict__.update(locals()) + del self.__dict__["self"] + + +@cute.jit +def phase_producer(tok, a: ProducerArgs): + """Producer work item ``tok``: warp w (< 12) of the calling CTA does packed head w of decode token tok. + The caller guarantees tok < rows and that warps >= 12 skip. Stores only (no CTA barrier inside).""" + tid = cute.arch.thread_idx()[0] + lane = tid % 32 + head = tid // 32 + in_off = (cutlass.Int64(tok) * PACKED_HEADS + cutlass.Int64(head)) * HEAD_DIM + cutlass.Int64( + lane * ELEMS + ) + pk = ld_u64(a.packed + in_off * 2) + x0 = bf16_bits_to_f32(cutlass.Int32(pk & cutlass.Int64(0xFFFF))) + x1 = bf16_bits_to_f32(cutlass.Int32((pk >> cutlass.Int64(16)) & cutlass.Int64(0xFFFF))) + x2 = bf16_bits_to_f32(cutlass.Int32((pk >> cutlass.Int64(32)) & cutlass.Int64(0xFFFF))) + x3 = bf16_bits_to_f32(cutlass.Int32((pk >> cutlass.Int64(48)) & cutlass.Int64(0xFFFF))) + producer_head(tok, head, lane, x0, x1, x2, x3, a) + + +@cute.jit +def producer_head(tok, head, lane, x0, x1, x2, x3, a: ProducerArgs): + """Everything after the packed-row load for packed head ``head`` (warp-uniform; dynamic Int32 or Constexpr) of + decode token ``tok``, on the calling warp: lane ``lane`` holds the head's elements [4 lane, 4 lane + 4) as the + float values x0..x3 of the BF16 packed row (the QKV GEMV output rounded to BF16). Norm weights, KV quant scale, + position and RoPE loads, the RMS tree, RoPE, and the q / index-q / index-K / K-V NVFP4 stores. Shared by + phase_producer and the QKV GEMV owner epilogue (MK-A), so the producer math has one implementation. + + Data flow is branch-free (every load is in bounds for every head: V reads the q norm weight, a non-rotating lane + reads its token's RoPE row, and the results are then discarded with select); only the stores branch, on + warp-uniform conditions, so every shuffle runs on the full warp as in the C++ kernel.""" + l4 = cutlass.Int64(lane * ELEMS) + is_v = head == V_HEAD + zero = cutlass.Float32(0.0) + + # Norm weights (BF16): q / k / index-q / index-k weight by head (V: q's, unused). + w_base = cutlass.select_( + head == K_HEAD, + a.k_weight, + cutlass.select_( + head == IQ_HEAD, + a.index_q_weight, + cutlass.select_(head == IK_HEAD, a.index_k_weight, a.q_weight), + ), + ) + ww = ld_u64(w_base + l4 * 2) + w0 = bf16_bits_to_f32(cutlass.Int32(ww & cutlass.Int64(0xFFFF))) + w1 = bf16_bits_to_f32(cutlass.Int32((ww >> cutlass.Int64(16)) & cutlass.Int64(0xFFFF))) + w2 = bf16_bits_to_f32(cutlass.Int32((ww >> cutlass.Int64(32)) & cutlass.Int64(0xFFFF))) + w3 = bf16_bits_to_f32(cutlass.Int32((ww >> cutlass.Int64(48)) & cutlass.Int64(0xFFFF))) + quant = ld_f32(a.kv_quant_scale + cutlass.select_(is_v, cutlass.Int64(8), cutlass.Int64(4))) + + # Position id, RoPE row. + pos = ld_s32(a.position_ids + cutlass.Int64(tok) * 4) + row = cutlass.Int64(pos) * ROTARY_DIM + coef = cutlass.Int64((lane * ELEMS) % (ROTARY_DIM // 2)) + c0 = ld_f32(a.rotary_cos_sin + (row + coef) * 4) + c1 = ld_f32(a.rotary_cos_sin + (row + coef + 1) * 4) + c2 = ld_f32(a.rotary_cos_sin + (row + coef + 2) * 4) + c3 = ld_f32(a.rotary_cos_sin + (row + coef + 3) * 4) + s0 = ld_f32(a.rotary_cos_sin + (row + ROTARY_DIM // 2 + coef) * 4) + s1 = ld_f32(a.rotary_cos_sin + (row + ROTARY_DIM // 2 + coef + 1) * 4) + s2 = ld_f32(a.rotary_cos_sin + (row + ROTARY_DIM // 2 + coef + 2) * 4) + s3 = ld_f32(a.rotary_cos_sin + (row + ROTARY_DIM // 2 + coef + 3) * 4) + + # RMSNorm (all heads compute it; V keeps x). + ss = sumsq_pair(zero, x0, x1) + ss = sumsq_pair(ss, x2, x3) + ss = f_add(ss, cute.arch.shuffle_sync_bfly(ss, offset=16)) + ss = f_add(ss, cute.arch.shuffle_sync_bfly(ss, offset=8)) + ss = f_add(ss, cute.arch.shuffle_sync_bfly(ss, offset=4)) + ss = f_add(ss, cute.arch.shuffle_sync_bfly(ss, offset=2)) + ss = f_add(ss, cute.arch.shuffle_sync_bfly(ss, offset=1)) + r = rms_reciprocal(ss, a.eps) + one = cutlass.Float32(1.0) + n0 = f_mul(x0, f_mul(r, f_add(one, w0))) + n1 = f_mul(x1, f_mul(r, f_add(one, w1))) + n2 = f_mul(x2, f_mul(r, f_add(one, w2))) + n3 = f_mul(x3, f_mul(r, f_add(one, w3))) + # NeoX partial RoPE: lane l's dim 4l + i pairs with lane l ^ 8 (dim +-32) within the first 64 dims. + pair_off = (ROTARY_DIM // 2) // ELEMS + p0 = cute.arch.shuffle_sync_bfly(n0, offset=pair_off) + p1 = cute.arch.shuffle_sync_bfly(n1, offset=pair_off) + p2 = cute.arch.shuffle_sync_bfly(n2, offset=pair_off) + p3 = cute.arch.shuffle_sync_bfly(n3, offset=pair_off) + rot = lane * ELEMS < ROTARY_DIM + low = lane * ELEMS < ROTARY_DIM // 2 + p0 = cutlass.select_(low, f_neg(p0), p0) + p1 = cutlass.select_(low, f_neg(p1), p1) + p2 = cutlass.select_(low, f_neg(p2), p2) + p3 = cutlass.select_(low, f_neg(p3), p3) + r0 = cutlass.select_(rot, f_fma(p0, s0, f_mul(n0, c0)), n0) + r1 = cutlass.select_(rot, f_fma(p1, s1, f_mul(n1, c1)), n1) + r2 = cutlass.select_(rot, f_fma(p2, s2, f_mul(n2, c2)), n2) + r3 = cutlass.select_(rot, f_fma(p3, s3, f_mul(n3, c3)), n3) + e0 = cutlass.select_(is_v, x0, r0) + e1 = cutlass.select_(is_v, x1, r1) + e2 = cutlass.select_(is_v, x2, r2) + e3 = cutlass.select_(is_v, x3, r3) + + if head < HEADS_Q: + st_u32( + a.q_out + (cutlass.Int64(tok) * HEADS_Q + cutlass.Int64(head)) * HEAD_DIM + l4, + e4m3x4_word(e0, e1, e2, e3), + ) + if head == IQ_HEAD: + st_u32(a.index_q_out + cutlass.Int64(tok) * HEAD_DIM + l4, e4m3x4_word(e0, e1, e2, e3)) + if head >= K_HEAD: + if head != IQ_HEAD: + slot = ld_s32(a.out_cache_loc + cutlass.Int64(tok) * 4) + page = slot // PAGE + if slot >= 0: + if page < a.num_pages: + within = slot & (PAGE - 1) + if head == IK_HEAD: + st_u32( + a.index_k + + cutlass.Int64(page) * a.index_page_stride + + cutlass.Int64(within) * a.index_token_stride + + l4, + e4m3x4_word(e0, e1, e2, e3), + ) + else: + plane = cutlass.select_(is_v, cutlass.Int64(1), cutlass.Int64(0)) + col = lane // 4 + sw = cutlass.select_( + is_v, + cutlass.Int64((within // 4) * 32 + col * 4 + within % 4), + cutlass.Int64(within * 8 + col), + ) + s_off = ( + cutlass.Int64(page) * a.scale_page_stride + + plane * a.scale_plane_stride + + sw + ) + # quantize_fp4x4: BF16-rounded values, 16-element block max over the 4-lane group. + q0 = round_bf16(e0) + q1 = round_bf16(e1) + q2 = round_bf16(e2) + q3 = round_bf16(e3) + lmax = f_max( + f_max(f_max(f_max(zero, f_abs(q0)), f_abs(q1)), f_abs(q2)), f_abs(q3) + ) + lmax = f_max(lmax, cute.arch.shuffle_sync_bfly(lmax, offset=1)) + lmax = f_max(lmax, cute.arch.shuffle_sync_bfly(lmax, offset=2)) + sbyte, sval = e4m3_byte_and_value( + f_mul(quant, f_mul(lmax, f_rcp(cutlass.Float32(6.0)))) + ) + if lane % 4 == 0: + st_u8(a.kv_scale + s_off, sbyte) + outs = cutlass.select_(lmax != zero, f_rcp(f_mul(sval, f_rcp(quant))), zero) + hw = e2m1x4_halfword( + f_mul(q0, outs), f_mul(q1, outs), f_mul(q2, outs), f_mul(q3, outs) + ) + d_off = ( + cutlass.Int64(page) * a.data_page_stride + + plane * a.data_plane_stride + + cutlass.Int64(within) * a.data_token_stride + + cutlass.Int64(lane * 2) + ) + st_u16(a.kv_data + d_off, hw) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_mc_allreduce.py new file mode 100644 index 000000000000..d5266a50dd3c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_mc_allreduce.py @@ -0,0 +1,376 @@ +# 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 + + +def max_rows() -> int: + """The most rows one boundary all-reduce call (plain or MoE-finalize) serves: the decode plans and the + MC_ALLREDUCE / MC_FINALIZE ranges read this, so a wider bucket or a row split here widens them with no change in + the modeling.""" + return MAX_TOKENS + + +# 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 ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + 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() + 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): + 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_m31_nvfp4__sm_103__tp8/ops/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_moe/m3_moe_kernel.py new file mode 100644 index 000000000000..97cef6b7ea10 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_moe/m3_moe_kernel.py @@ -0,0 +1,3673 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3.1 TP8 implementation.""" + +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 +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 +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 +NGEMV = (NROUTE + NCTA - 1) // NCTA +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +XVDUMP = bool(_cfg("xvdump", 0)) +YDUMP = bool(_cfg("ydump", 0)) +LAG = _cfg("lag", 2) +G_CAP = TOP_K * M_MAX +NG_CAP = M_MAX +XW = _cfg("xw", 5 if M_MAX > 4 else 4 if M_MAX > 2 else 1) +THREADS = 256 + 32 * (XW - 1) +GEMV_THREADS = 256 +SHARED = bool(_cfg("shared", 0)) +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +AR_PRE = bool(_cfg("ar_pre", 1)) +TAIL_TRIGGER = bool(_cfg("tail_trigger", 0)) +assert True, "early_ar and tail_trigger are alternative release points" +EARLY_EPOCH = TAIL_TRIGGER +XSTAGE = bool(_cfg("xstage", 1)) and bool(_cfg("fused_router", 0)) +B1WIDE = bool(_cfg("b1wide", 0)) +B1W_BAR = 10 +NSH = 24 if SHARED else 0 +SH_I = _cfg("sh_i", 768) +SH_ROWS = SH_I // 24 +SH_SPLIT = 128 // (2 * SH_ROWS) +SH_NB = 8 * SH_SPLIT +SH_KB = H // 32 +SH_DKB = SH_I // 32 +SH_DKT = SH_I // 128 +SH_GT = SH_KB // 4 // SH_SPLIT +SH_STAGES_GU = SH_GT +SH_STAGES = SH_GT + 2 * SH_DKT +SHX_BAR = 6 +SH_PRETILE = bool(_cfg("sh_pretile", 1 if M_MAX <= 2 or SH_SPLIT != 2 else 0)) +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 +N1 = 16 if M_MAX <= 4 else 32 +QS = N1 // 4 +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 +K1_TILES = H // 256 +S1 = K1_TILES // C1 +NUM_BYTES_A = MMA_M * KT_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 +R2 = 32 +NB2 = H // R2 +KT2 = (I_TP + 255) // 256 +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) +NUM_BYTES_B2 = N2 * KT_BYTES +NB2_CAP = _cfg("nb2_cap", 2) +UNITS_ROW = TWO_I // R1 +H_ROW = UNITS_ROW * 16 +HB_REPL = _cfg("hb_repl", 1) +assert 1 <= HB_REPL <= 32, HB_REPL +HB_STRIDE = (G_CAP * M_MAX * UNITS_ROW * 4 + 255) // 256 * 256 + 32 +PARTIAL_EMPTY = -1 +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +B2_BAR = 5 +TD_STRIDE = 64 +TOPK_BAR = 4 +B1_AFTER_TOPK = bool(_cfg("b1_after_topk", 0)) +TOPK2_BAR = 7 +EARLY0 = bool(_cfg("early0", 1)) +EARLY0_N = _cfg("early0_n", 99) +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 +NRW = NRW0 + RW_EXTRA +LG_REPL = _cfg("lg_repl", 32 if M_MAX > 4 else 128) +GEMV_SMEM = bool(_cfg("gemv_smem", 1)) +LG_STRIDE = M_MAX * E_LOCAL + 32 +POLL_NS = _cfg("poll_ns", 150) +B2_POLL_NS = _cfg("b2_poll_ns", 200) +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) +B2_DLY = _cfg("b2_dly", 0) +W_WARP = 4 +X_WARP = 5 +S_WARP = 6 +MMA_WARP = 7 +E0_WARP = MMA_WARP +E0_BAR = 8 +E1_BAR = 9 +E0_OFFLOAD = EARLY0 and NRW0 <= 3 and bool(_cfg("e0_off", 1)) +EVICT_FIRST = 1364590687093260288 +_LOG2E = 1.4426950408889634 +LD_POW2 = 1 if M_MAX == 1 else 2 if M_MAX == 2 else 4 if M_MAX <= 4 else 8 +_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 = 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) +if B1WIDE: + assert ( + XSTAGE + and FUSED_ROUTER + and (W_WARP == 4) + and (NRW0 == 4) + and (RW_EXTRA == THREADS // 32 - 8) + ), ( + "b1wide: every warp of the CTA must be warps 0..3 or a routing warp", + NRW0, + RW_EXTRA, + THREADS, + ) + assert M_MAX * (H // 16) % THREADS == 0, (M_MAX, THREADS) +assert M_MAX * H * 2 <= STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA), ( + "x staging must fit the ring" +) +_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 +_SB_A = 0 +_SB_RAW = None +_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 +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 +SH_TMEM = 256 +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) + + +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, + ) + + +@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, + ) + ) + + +@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, + ) + + +@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\tand.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl " + "and.b32 w, w, 0xFFFF0000;\n\tshr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h," + " 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\tmov.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, + ) + ) + + +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(255) + + +@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.f16x" + "2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\ncvt.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(255) + + +@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(2147483647), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(4194304), 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 _b1_quant( + qtid, + NTHR: int, + xb32: cutlass.Array, + sfsc: cutlass.Array, + sB1: cutlass.Array, + s_xsf: cutlass.Array, +): + """(B1WIDE) the quantize loop of _build_b1_fused over NTHR threads (thread qtid): blocks qtid, qtid + NTHR, ..""" + 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) // NTHR + QB = 3 if NBLK % 3 == 0 else 2 if NBLK % 2 == 0 else 1 + for rg in cutlass.range(NBLK // QB, unroll=1): + wall = [] + for rr in cutlass.range_constexpr(QB): + bq = qtid + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(NTHR) + 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 = qtid + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(NTHR) + 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(255) + 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) + + +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).""" + _b1_quant(tidx, EPI_THREADS, xb32, sfsc, sB1, s_xsf) + 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.4028234663852886e38) + 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.4426950216293335))) + 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.1666666716337204))) + 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(65535), 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() + elif 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)) + ) + 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 + 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) + 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 & 65535 + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): + _sh_issue(lane, sid, cutlass.Int32(g), s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + if cutlass.const_expr(SH_PRETILE): + 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) + if cutlass.const_expr(EARLY_EPOCH): + prims.barrier_cta_sync(0) + if tidx == 0: + epochs.store(ep + cutlass.Int32(1), idx=bx) + tag16 = ep & cutlass.Int32(32767) | cutlass.Int32(32768) + NQ = THREADS - 64 + 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: + 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): + 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)) + 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: + 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) + 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) + _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): + 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: + 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): + 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) + 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.7020000219345093)), + cutlass.Float32(1.4426950216293335), + ) + ) + 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) + 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(65535) == 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) + _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 cutlass.const_expr(not EARLY_EPOCH): + 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], + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_gu_desc: cutlass.GridConstant[cuda.TensorMap], + tma_dn_desc: cutlass.GridConstant[cuda.TensorMap], + logits: cutlass.Array, + bias: cutlass.Array, + x32: cutlass.Array, + xsf32: cutlass.Array, + escal: cutlass.Array, + plan: cutlass.Array, + partials: cutlass.Array, + hbuf32: cutlass.Array, + epochs: cutlass.Array, + out: cutlass.Array, + route_dbg: cutlass.Array, + tdbg: cutlass.Array, + xb32: cutlass.Array, + rw: cutlass.Array, + sfsc: cutlass.Array, + lgbuf: cutlass.Array, + guw: cutlass.Array, + gusf: cutlass.Array, + dnw: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + s1w: cutlass.Array, + s2w: cutlass.Array, + img32: cutlass.Array, + b2cnt: 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 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) + 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) + img_bar = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + 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_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: + 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 & 65535 + 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(127) | cutlass.Int32(128)) << cutlass.Int32(24) + tag16 = ep & cutlass.Int32(32767) | cutlass.Int32(32768) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + if cutlass.const_expr(FUSED_ROUTER): + 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(4294967295) + _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) + if tidx == 0: + _stamp(tdbg, bx, 21) + 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) + ) + if cutlass.const_expr(B1WIDE): + _wait(x_bar.subview(0), 0) + xw32 = cutlass.Array( + s_big.data_ptr(0), shape=(M_MAX * H // 2,), dtype=cutlass.Int32, alignment=16 + ) + _b1_quant(tidx, THREADS, xw32, sfsc, sB1, s_xsf) + _bar_arrive(B1W_BAR, THREADS) + 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(4294967295, 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) + 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): + 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): + 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(65535) + 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) + 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) + 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(65535) + 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) + 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(4294967295, 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): + 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): + 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, + ) + if cutlass.const_expr(B1WIDE): + _b1_quant(tidx, THREADS, xs32, sfsc, sB1, s_xsf) + cute.arch.barrier(barrier_id=B1W_BAR, number_of_threads=THREADS) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + if cutlass.const_expr(not B1WIDE): + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xs32, sfsc, sB1, s_xsf, b1_ready + ) + prims.mbarrier_arrive(x_bar.subview(1)) + 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: + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ROUTE_BAR, number_of_threads=THREADS) + if cutlass.const_expr(EARLY_EPOCH): + if tidx == 0: + epochs.store(ep + cutlass.Int32(1), idx=bx) + 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(65535) + nb = pr2 >> cutlass.Int32(16) + if warp == W_WARP: + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + 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) + 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) + if cutlass.const_expr(TAIL_TRIGGER): + if ar_on: + cute.arch.griddepcontrol_launch_dependents() + if (warp == S_WARP) & cutlass.Boolean(not SFA_CP): + 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) + elif 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): + 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() + 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) + 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)): + 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 + hb_off = cutlass.Int32(0) + if cutlass.const_expr(HB_REPL > 1): + hb_off = bx % cutlass.Int32(HB_REPL) * cutlass.Int32(HB_STRIDE) + NZ = (nrec_cap + 32 * XW - 1) // (32 * XW) + 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, + ) + 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) + 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(4294967295, 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) + c = qi - rr * cutlass.Int32(UNITS_ROW) + t2 = c >> cutlass.Int32(4) + cc = c >> cutlass.Int32(1) & cutlass.Int32(7) + lo = w0 & cutlass.Int32(16777215) | w1 << cutlass.Int32(24) + hi = w1 >> cutlass.Int32(8) & cutlass.Int32(65535) | ( + w2 & cutlass.Int32(65535) + ) << 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, + ) + sSFB2_8.store( + cutlass.Int8(w2 >> cutlass.Int32(16) & cutlass.Int32(255)), + 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), + ) + 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)) + 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: + 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): + 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): + 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(255) + 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(255)) << cutlass.Int32(16) + | tagv + ) + if cutlass.const_expr(B2IMG): + if lane == 0: + 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(255)), + 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)), + ) + elif 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, + ) + elif lane < HB_REPL: + 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))) + 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) + 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): + 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): + wy = wys[bt] + while wy >> cutlass.Int32(16) & cutlass.Int32(65535) != 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, + ) + if tidx == 0: + _stamp(tdbg, bx, 7) + if cutlass.const_expr(not EARLY_EPOCH): + 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, + 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_m31_nvfp4__sm_103__tp8/ops/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_moe/op.py new file mode 100644 index 000000000000..3f69c9392eb1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_moe/op.py @@ -0,0 +1,670 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3.1 TP8 implementation.""" + +from __future__ import annotations + +import importlib.util +import os +import sys +from typing import Dict, Optional, Tuple + +import torch + +from ...m31_switches import m31_disabled + +HIDDEN = 6144 +INTER = 384 +NUM_EXPERTS = 128 +SHARED_INTER = 384 +SHARED_CTAS = 24 +TOP_K = 4 +U1 = 2 * INTER // 32 +MIN_CTAS = 64 +_UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) +_ROW2_BYTES = INTER // 2 + INTER // 16 +_S1 = 6 +_ITEM_BYTES = _UNIT_BYTES // _S1 +_modules: Dict[tuple, object] = {} +_compiled: Dict[tuple, object] = {} +SFA_CP_BY_M: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1, 1) +B2IMG_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 1, 1) +AR_PRE_BY_M: Tuple[int, ...] = (1, 0, 1, 1, 1, 1, 1, 1) +STAGES_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 0, 6) +HB_REPL_BY_M: Tuple[int, ...] = (32, 32, 8, 8, 1, 1, 1, 1) +M8_CFG_BY_M: Dict[int, Dict[str, int]] = {8: {"b2img": 0}} +B1WIDE_TOKENS = frozenset({8}) +LG_REPL = 32 +LG_REPL_MAX_TOKENS = 4 +LG_REPL_BY_M: Dict[int, int] = {8: 16} +FINE_PLAN = True +PLAN_FC2_W = 2.0 +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 - dependency availability check + 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): + 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 + out = torch.empty(n_cta, gt + 2 * dkt, 16384 + 512, dtype=torch.uint8, device=dev) + for s in range(n_cta): + nrow = (octant & 1) * I_ + P * s + r % P + sl = octant >> 1 + for g in range(gt): + kt = sl * gt + g + tile = gu[nrow[:, None], kt[:, None] * 128 + torch.arange(128, device=dev)[None, :]] + img = tile.view(128, 8, 16)[r[:, None], swz] + 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) + 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) + + +_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: + 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, + fc2_w: float = 1.0, +): + """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 = int(round(32 * _ROW2_BYTES * fc2_w)) + 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))): + 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: + 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( + ( + "m3_moe (shared): num_ctas + " + f"{SHARED_CTAS}" + " shared-expert CTAs must fit the " + f"{sms_all}" + " SMs (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 + plan, nb_cap = make_phase_plan( + self.num_ctas, + num_tokens, + M3_MOE_LAG, + 0.0 if num_tokens == 1 else 1.0, + True, + fc2_w=1.0 if m31_disabled("PLANW") else PLAN_FC2_W, + ) + if not m31_disabled("PLANW") and PLAN_FC2_W != 1.0: + from tensorrt_llm.logger import logger as _planw_log + + _planw_log.info( + ( + "MiniMax-M3.1 MoE engine PLAN-W: FC2 bytes weighted " + f"{PLAN_FC2_W}" + " in the phase plan (tokens=" + f"{num_tokens}" + ")" + ) + ) + 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 + elif num_tokens in LG_REPL_BY_M and (not m31_disabled("MOE_LGREPL16")): + self.config["lg_repl"] = LG_REPL_BY_M[num_tokens] + from tensorrt_llm.logger import logger as _lg_log + + _lg_log.info( + ( + "MiniMax-M3.1 MoE engine MOE_LGREPL16: " + f"{LG_REPL_BY_M[num_tokens]}" + " router-logit exchange replicas (tokens=" + f"{num_tokens}" + ")" + ) + ) + if shared: + self.config["sh_i"] = SHARED_INTER + self.config["shared"] = 1 + self.shared = shared + self.config["sfa_cp"] = SFA_CP_BY_M[num_tokens - 1] + self.config["b2img"] = B2IMG_BY_M[num_tokens - 1] + 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] + if num_tokens in M8_CFG_BY_M and (not m31_disabled("MOE_M8CFG")): + self.config.update(M8_CFG_BY_M[num_tokens]) + if ( + num_tokens in B1WIDE_TOKENS + and self.config.get("fused_router") + and (not m31_disabled("MOE_B1WIDE")) + ): + self.config["b1wide"] = 1 + from tensorrt_llm.logger import logger as _b1w_log + + _b1w_log.info( + f"MiniMax-M3.1 MoE engine B1WIDE: B1 quantize on every warp (tokens={num_tokens})" + ) + if not m31_disabled("MOE_TAIL_TRIGGER"): + self.config["tail_trigger"] = 1 + for key in ("tail_trigger",): + if self.config.get(key): + from tensorrt_llm.logger import logger as _rel_log + + _rel_log.info( + f"MiniMax-M3.1 MoE engine dependent release: {key} (tokens={num_tokens}, epoch stored early)" + ) + try: + self.mod = mod = _kernel_module(self.config, _ENGINE_PATH) + except AssertionError as err: + raise ValueError( + ( + "m3_moe: num_tokens=" + f"{num_tokens}" + ", num_ctas=" + f"{self.num_ctas}" + " does not fit the shared memory / TMEM budget (fewer CTAs means " + "more FC2 blocks per CTA): " + f"{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 + ) + 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) + 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) + 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: + 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 + shared_mxfp8 = None + + 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_m31_nvfp4__sm_103__tp8/ops/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_sampled_accept.py new file mode 100644 index 000000000000..59c58336c5d7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_sampled_accept.py @@ -0,0 +1,159 @@ +# 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 ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + 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() + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + 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_m31_nvfp4__sm_103__tp8/ops/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_select_long_dsl.py new file mode 100644 index 000000000000..b56c93b93a2b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_select_long_dsl.py @@ -0,0 +1,367 @@ +# 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 up to 9216 blocks, one row per 384-thread CTA (CuTe DSL). + +Bit for bit the selection of this target's ``m3_select_long`` (``torch.ops.trtllm.minimax_m3_select_long``, +one 1024-thread CTA per row), from a CTA of 384 threads: + +* Every block's score is forced (the first ``init_blocks`` and the last ``local_blocks`` of the row's valid blocks) and + mapped to an order key held in registers: thread t owns blocks t, t + 384, ... (4, 11 or 24 of them by the row's + valid count). Larger scores get smaller keys; NaN sorts above +inf. The key is the 1024-thread kernel's unsigned + key with its top bit flipped, so signed Int32 order here is that kernel's order. Every thread loads at a clamped + block and masks the result, so a thread's loads issue back to back. +* 384 is a multiple of 16, so the 16 interleaved groups ``block % 16`` are the lanes' ``lane % 16``. The weakest of + the 16 group maxima bounds the 16th-largest score, so the top 16 is among the keys at or above it. +* The candidates are compacted by a CTA prefix count, each as one int64 (key << 32 | block), and ranked exactly by + that value (one signed compare per pair); more than 1024 of them (keys tied at the bound) rank every candidate + against the whole row instead. +* Warp 0 emits the 16 picked block ids sorted ascending, -1 for an empty or -inf pick. + +The kernel waits on its grid dependency first; it releases the dependent sparse decode right after that wait +(``early_trigger``), once every score of its row is loaded (trigger phase 1) or at its exit (phase 0). Kernels are +compiled per (PDL launch, early trigger, trigger phase) with symbolic shapes and strides; ``prebuild`` compiles and +loads them before the executor's warmup, so a decode step never compiles one. +""" + +from typing import Dict, Tuple + +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_fake_tensor +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import ( + griddepcontrol_launch_dependents, + griddepcontrol_wait, +) + +# The 384-thread select reads its row's valid-block count before its grid wait instead of after it, which takes one +# dependent global load off the post-wait path; the selection is the same either way. Exact: the count +# (MsaMetadata.msa_n_valid_blocks) is written only before the step's first layer, behind a plain launch. +SELECT_N_VALID_PREWAIT = True + +SELECT_THREADS = 384 +SELECT_WARPS = SELECT_THREADS // 32 +# 1M tokens + draft tokens at 128 tokens per block. +MAX_BLOCKS = 9216 # capacity: fixed: the kernel's built blocks per row +MAX_CAND = 1024 # capacity: fixed: the kernel's candidate staging +TOPK = 16 # capacity: fixed: MSA's top-k +NO_KEY = 0x7FFFFFFF # capacity: fixed: a sentinel, not a size (an unused key slot (the reference's 0xffffffff)) +# The key of -inf (the reference's 0xff800000): never emitted. +NEG_INF_KEY = 0x7F800000 # capacity: fixed: a key encoding, not a size +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 +# (keys per thread, rows up to this many valid blocks) +ITEM_CLASSES = ((4, 4 * SELECT_THREADS), (11, 11 * SELECT_THREADS), (24, MAX_BLOCKS)) + +# Shared memory of the phase, int32 words: candidates [MAX_CAND] int64 (key << 32 | block, so one signed compare +# orders them by (key, block)), picked keys / blocks [TOPK] each, per-warp group minima [SELECT_WARPS][16], group +# minima [16], per-warp candidate counts [SELECT_WARPS]. +_CAND = 0 +_PICK_KEY = _CAND + 2 * MAX_CAND +_PICK_BLOCK = _PICK_KEY + TOPK +_WARP_MIN = _PICK_BLOCK + TOPK +_GROUP_MIN = _WARP_MIN + SELECT_WARPS * 16 +_WARP_CNT = _GROUP_MIN + 16 +SMEM_WORDS = _WARP_CNT + SELECT_WARPS +SMEM_BYTES = 4 * SMEM_WORDS + + +def _key(score): + """Order key of a float score: smaller key = larger score, NaN first; signed order = the reference's unsigned.""" + bits = cutlass.Float32(score).bitcast(cutlass.Int32) + return cutlass.select_(bits < cutlass.Int32(0), bits & cutlass.Int32(0x7FFFFFFF), ~bits) + + +def _imin(a, b): + return cutlass.select_(b < a, b, a) + + +def _imax(a, b): + return cutlass.select_(b > a, b, a) + + +def _forced_key(scores, head, block, query, init_blocks, local_start): + raw = scores[head, block, query] + s = cutlass.select_( + block >= local_start, + cutlass.Float32(LOCAL_SCORE), + cutlass.select_(block < init_blocks, cutlass.Float32(INIT_SCORE), cutlass.Float32(raw)), + ) + return _key(s) + + +@cute.jit +def _select_row( + row, + scores: cute.Tensor, # float32 [num_kv_heads, num_blocks, total_q], any strides + n_valid: cute.Tensor, # int32 [>= total_q] + out: cute.Tensor, # int32 [total_q, num_kv_heads, 16] + init_blocks, + local_blocks, + num_kv_heads, + smem: cute.Tensor, # int32 [SMEM_WORDS], shared memory + bar_id: cutlass.Constexpr, + loaded_trigger: cutlass.Constexpr, + raw_prewait, + n_valid_prewait: cutlass.Constexpr, +): + """Select row ``row`` (= query * num_kv_heads + head) with this CTA's SELECT_THREADS threads. Uses named barrier + ``bar_id`` over SELECT_THREADS threads. ``loaded_trigger`` releases the dependent grid once every score of the + row is loaded (trigger phase 1). With ``n_valid_prewait`` the row's valid-block count is ``raw_prewait``, read by + the caller before its grid wait; otherwise it is read here.""" + tid, _, _ = cute.arch.thread_idx() + lane = tid % cutlass.Int32(32) + warp = tid // cutlass.Int32(32) + query = row // num_kv_heads + head = row % num_kv_heads + if cutlass.const_expr(n_valid_prewait): + raw = raw_prewait + else: + raw = n_valid[query] + valid = _imax(_imin(raw, cutlass.Int32(scores.shape[1])), cutlass.Int32(0)) + local_start = _imax(raw - local_blocks, cutlass.Int32(0)) + cand = cute.make_tensor( + cute.recast_ptr(smem.iterator + _CAND, dtype=cutlass.Int64), cute.make_layout(MAX_CAND) + ) + + if tid < TOPK: + smem[_PICK_KEY + tid] = cutlass.Int32(NO_KEY) + smem[_PICK_BLOCK + tid] = cutlass.Int32(0) + + if valid <= TOPK: + # Every valid block is picked (thread t holds block t). + if tid < valid: + smem[_PICK_KEY + tid] = _forced_key(scores, head, tid, query, init_blocks, local_start) + smem[_PICK_BLOCK + tid] = tid + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + for ci in cutlass.range_constexpr(len(ITEM_CLASSES)): + items = ITEM_CLASSES[ci][0] + lo = TOPK if ci == 0 else ITEM_CLASSES[ci - 1][1] + hi = ITEM_CLASSES[ci][1] + if (valid > lo) & (valid <= hi): + key = cute.make_rmem_tensor(items, cutlass.Int32) + local_min = cutlass.Int32(NO_KEY) + for s in cutlass.range_constexpr(items): + # Every thread loads (a clamped block past the row) so the loads issue back to back, unbranched. + block = tid + cutlass.Int32(s * SELECT_THREADS) + k = _forced_key( + scores, head, _imin(block, valid - 1), query, init_blocks, local_start + ) + key[s] = cutlass.select_(block < valid, k, cutlass.Int32(NO_KEY)) + local_min = _imin(local_min, key[s]) + local_min = _imin(local_min, cute.arch.shuffle_sync_bfly(local_min, 16)) + if lane < 16: + smem[_WARP_MIN + warp * cutlass.Int32(16) + lane] = local_min + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + if tid < 16: + g = cutlass.Int32(NO_KEY) + for w in cutlass.range_constexpr(SELECT_WARPS): + g = _imin(g, smem[_WARP_MIN + w * 16 + tid]) + smem[_GROUP_MIN + tid] = g + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + bound = cutlass.Int32( + -2147483648 + ) # the smallest key: signed keys, unlike the reference's unsigned ones + for g in cutlass.range_constexpr(16): + bound = _imax(bound, smem[_GROUP_MIN + g]) + + # Candidates (key <= bound; unused slots hold NO_KEY > any bound): a CTA prefix count gives the slots. + cnt = cutlass.Int32(0) + for s in cutlass.range_constexpr(items): + cnt += cutlass.select_(key[s] <= bound, cutlass.Int32(1), cutlass.Int32(0)) + incl = cnt + for d in cutlass.range_constexpr(5): + up = cute.arch.shuffle_sync( + incl, _imax(lane - cutlass.Int32(1 << d), cutlass.Int32(0)) + ) + incl += cutlass.select_(lane >= cutlass.Int32(1 << d), up, cutlass.Int32(0)) + if lane == 31: + smem[_WARP_CNT + warp] = incl + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + before = cutlass.Int32(0) + total = cutlass.Int32(0) + for w in cutlass.range_constexpr(SELECT_WARPS): + c = smem[_WARP_CNT + w] + before += cutlass.select_(cutlass.Int32(w) < warp, c, cutlass.Int32(0)) + total += c + slot = before + incl - cnt + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + if slot < MAX_CAND: + cand[slot] = (cutlass.Int64(key[s]) << cutlass.Int64(32)) | cutlass.Int64( + tid + cutlass.Int32(s * SELECT_THREADS) + ) + slot += 1 + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if total <= MAX_CAND: + for i in cutlass.range(tid, total, SELECT_THREADS, unroll=1): + mine = cand[i] + rank = cutlass.Int32(0) + for j in cutlass.range(total, unroll=8): + rank += cutlass.select_(cand[j] < mine, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = cutlass.Int32(mine >> cutlass.Int64(32)) + smem[_PICK_BLOCK + rank] = cutlass.Int32(mine & cutlass.Int64(0xFFFFFFFF)) + else: + # More than MAX_CAND keys tied at the bound: rank each candidate against the whole row. + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + block = tid + cutlass.Int32(s * SELECT_THREADS) + rank = cutlass.Int32(0) + for j in cutlass.range(valid, unroll=1): + kj = _forced_key(scores, head, j, query, init_blocks, local_start) + less = (kj < key[s]) | ((kj == key[s]) & (j < block)) + rank += cutlass.select_(less, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = key[s] + smem[_PICK_BLOCK + rank] = block + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if warp == 0: + # The 16 picked block ids ascending (NO_KEY = empty or -inf, emitted as -1), bitonic over the warp. + v = cutlass.Int32(NO_KEY) + if lane < TOPK: + pk = smem[_PICK_KEY + lane] + if (pk != cutlass.Int32(NO_KEY)) & (pk != cutlass.Int32(NEG_INF_KEY)): + v = smem[_PICK_BLOCK + lane] + for size_log in cutlass.range_constexpr(1, 6): + for stride_log in cutlass.range_constexpr(size_log - 1, -1, -1): + other = cute.arch.shuffle_sync_bfly(v, 1 << stride_log) + ascending = (lane & cutlass.Int32(1 << size_log)) == cutlass.Int32(0) + lower = (lane & cutlass.Int32(1 << stride_log)) == cutlass.Int32(0) + v = cutlass.select_(lower == ascending, _imin(v, other), _imax(v, other)) + if lane < TOPK: + out[query, head, lane] = cutlass.select_( + v == cutlass.Int32(NO_KEY), cutlass.Int32(-1), v + ) + + +class _SelectLongDslKernel: + """One row per CTA: the grid wait, then ``_select_row``. ``n_valid_prewait`` reads the row's valid-block count + before the grid wait, so that dependent load leaves the post-wait critical path. That is exact because + ``n_valid`` is written only by ``on_update_kv_lens`` and the host copy in ``prepare()``, both before the step's + first layer and behind a plain (non-PDL) launch, so no grid this select can overlap writes it + (``attention/test_modeling_v2_m3_n_valid_contract.py``).""" + + def __init__( + self, use_pdl: bool, early_trigger: bool, trigger_phase: int, n_valid_prewait: bool + ): + self.use_pdl = use_pdl + self.early_trigger = early_trigger + self.loaded_trigger = (not early_trigger) and trigger_phase == 1 + self.n_valid_prewait = n_valid_prewait + + @cute.kernel + def kernel(self, scores, n_valid, out, init_blocks, local_blocks, num_kv_heads): + row, _, _ = cute.arch.block_idx() + smem = SmemAllocator() + words = smem.allocate_tensor(cutlass.Int32, cute.make_layout(SMEM_WORDS), byte_alignment=16) + raw_prewait = cutlass.Int32(0) + if cutlass.const_expr(self.n_valid_prewait): + raw_prewait = n_valid[row // num_kv_heads] + griddepcontrol_wait() + if cutlass.const_expr(self.early_trigger): + griddepcontrol_launch_dependents() + _select_row( + row, + scores, + n_valid, + out, + init_blocks, + local_blocks, + num_kv_heads, + words, + bar_id=1, + loaded_trigger=self.loaded_trigger, + raw_prewait=raw_prewait, + n_valid_prewait=self.n_valid_prewait, + ) + + @cute.jit + def __call__(self, scores, n_valid, out, init_blocks, local_blocks, stream): + num_kv_heads = scores.shape[0] + self.kernel(scores, n_valid, out, init_blocks, local_blocks, num_kv_heads).launch( + grid=(scores.shape[2] * num_kv_heads, 1, 1), + block=(SELECT_THREADS, 1, 1), + stream=stream, + use_pdl=self.use_pdl, + ) + + +#: Compiled kernels by (PDL launch, early trigger, trigger phase, valid count read before the wait). +kernel_cache: Dict[Tuple[bool, bool, int, bool], object] = {} + + +def kernel_key( + use_pdl: bool, early_trigger: bool, trigger_phase: int +) -> Tuple[bool, bool, int, bool]: + """The compile key of a launch (the trigger phase does not matter with the early trigger), with this target's + ``SELECT_N_VALID_PREWAIT``.""" + return ( + bool(use_pdl), + bool(early_trigger), + 0 if early_trigger else int(trigger_phase), + bool(SELECT_N_VALID_PREWAIT), + ) + + +def _compile(key: Tuple[bool, bool, int, bool]): + if key not in kernel_cache: + heads, blocks, rows = cute.sym_int(), cute.sym_int(), cute.sym_int() + scores = make_fake_tensor( + cutlass.Float32, + (heads, blocks, rows), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + n_valid = make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(),), stride_order=(0,)) + out = make_fake_tensor( + cutlass.Int32, (rows, heads, TOPK), stride=(cute.sym_int64(), cute.sym_int64(), 1) + ) + kernel_cache[key] = cute.compile( + _SelectLongDslKernel(*key), + scores, + n_valid, + out, + cutlass.Int32(0), + cutlass.Int32(0), + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + return kernel_cache[key] + + +def select_long_dsl( + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool, + early_trigger: bool, + trigger_phase: int, +) -> None: + """Launch the select on the current stream (compiling its key first if ``prebuild`` did not).""" + _compile(kernel_key(use_pdl, early_trigger, trigger_phase))( + scores, n_valid_blocks, out, init_blocks, local_blocks + ) + + +def prebuild(keys, device) -> list: + """Compile every key and launch each once on a small dummy row (loading its module). Returns the keys.""" + scores = torch.zeros( + 1, 2048, 4, dtype=torch.float32, device=device + ) # capacity: fixed: a dummy prebuild row, not a size + n_valid = torch.full((4,), 2048, dtype=torch.int32, device=device) + out = torch.empty(1, 4, TOPK, dtype=torch.int32, device=device).permute(1, 0, 2) + built = [] + for key in keys: + key = kernel_key(*key) + _compile(key)(scores, n_valid, out, 0, 1) + built.append(key) + return built diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_shard_accept.py new file mode 100644 index 000000000000..2b7d05ad920f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/m3_shard_accept.py @@ -0,0 +1,176 @@ +# 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 ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + 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() + if not _mnnvl_workspace_all_succeeded(comm, True): + 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_m31_nvfp4__sm_103__tp8/ops/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/moe/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/moe/m3_router_gemv_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/moe/m3_router_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py new file mode 100644 index 000000000000..3ba6fb0d7c8c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py @@ -0,0 +1,1472 @@ +# 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. + +``prod`` (#117, the sparse QKV + index projection at split 8, at most 8 tokens): instead of storing the packed row, the +owners stage their bf16 rows, the CTAs exchange them by st.async so that rank t holds token t's 128-dim row of this +tile's packed head, and rank t's warp 4 runs that head's producer (``m3_layer.phase_producer.producer_head``: RMSNorm, +RoPE, FP8 q / index-q / index-K and NVFP4 K / V stores) on it, bit for bit what the separate producer kernel does with +the stored row. + +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 + +from ..m3_layer.phase_producer import ProducerArgs, bf16_bits_to_f32, producer_head + +CTA_M = 128 # capacity: fixed: a kernel tile constant (weight rows per CTA = the MMA's M) +MMA_N = 8 # token columns +CTA_K = 128 # capacity: fixed: a kernel tile constant (one k-tile: 128 E4M3 bytes per row = one 128-byte swizzle row) +MMA_INST_K = 32 +# 4 MMAs per k-tile, one per scale byte of the atom. +K_BLOCKS = CTA_K // MMA_INST_K # 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 _ld_shared_v4_b32(smem_ptr, *, loc=None, ip=None): + """ld.shared.v4.b32 of four int32 words 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.i32()] * 4), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", "=r,=r,=r,=r,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + return [cutlass.Int32(_llvm.extractvalue(_T.i32(), 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) + + +# The swigluoai activation of MiniMax-M3's dense MLP, gate * sigmoid(1.702 * gate) * (up + 1) with gate clamped to <= 7 +# and up to [-7, 7], in exactly the instruction sequence the Triton silu_and_mul kernel compiles to (modules/swiglu.py +# on SM107): the gate clamp on the packed bf16 pair, the up clamp as min.xorsign.abs, two rounded multiplies into +# ex2.approx, div.full for the reciprocal, two rounded multiplies, round to bf16. No multiply feeds an add, so the +# explicit .rn forms compile to what the Triton kernel's contractable forms do. +SWIGLU_OAI_ALPHA = 1.702 +SWIGLU_OAI_LIMIT = 7.0 + + +@dsl_user_op +def _swiglu_oai_bf16x2(gate_word, up_word, *, loc=None, ip=None): + """The activation of one packed bf16 pair of gate and up values: its two bf16 results as one int32 word (the low + half from the low halves).""" + 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.Int32(w).ir_value(loc=loc, ip=ip) for w in (gate_word, up_word)], + "{\n\t.reg .b32 g, glim, gl, gh, ul, uh, one, lim, t0, t1, e0, e1, s0, s1, r0, r1;\n\t.reg .b16 hl, hh;\n\t" + "mov.b32 glim, 0x40E040E0;\n\tmin.bf16x2 g, $1, glim;\n\t" + "mov.b32 {hl, hh}, g;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, $2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mov.b32 lim, 0f40E00000;\n\tmin.xorsign.abs.f32 ul, ul, lim;\n\tmin.xorsign.abs.f32 uh, uh, lim;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 ul, ul, one;\n\tadd.rn.f32 uh, uh, one;\n\t" + "mul.rn.f32 t0, gl, 0fBFD9DB23;\n\tmul.rn.f32 t1, gh, 0fBFD9DB23;\n\t" + "mul.rn.f32 t0, t0, 0f3FB8AA3B;\n\tmul.rn.f32 t1, t1, 0f3FB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "add.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 r0, s0, gl;\n\tmul.rn.f32 r1, s1, gh;\n\t" + "mul.rn.f32 r0, ul, r0;\n\tmul.rn.f32 r1, uh, r1;\n\t" + "cvt.rn.bf16x2.f32 $0, r1, r0;\n\t}", "=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 _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], + x_swiglu: cutlass.Constexpr[ + bool + ], # x_bf16: x_words is [M, 2K] gate | up; the activation is swigluoai(gate, up) + # prod: the producer's addresses and strides (ProducerArgs fields; `packed` is unused), then num_pages and eps. + pr_q: cutlass.Int64, + pr_iq: cutlass.Int64, + pr_ik: cutlass.Int64, + pr_kvd: cutlass.Int64, + pr_kvs: cutlass.Int64, + pr_loc: cutlass.Int64, + pr_quant: cutlass.Int64, + pr_pos: cutlass.Int64, + pr_rope: cutlass.Int64, + pr_qw: cutlass.Int64, + pr_kw: cutlass.Int64, + pr_iqw: cutlass.Int64, + pr_ikw: cutlass.Int64, + pr_ips: cutlass.Int64, + pr_its: cutlass.Int64, + pr_dps: cutlass.Int64, + pr_dls: cutlass.Int64, + pr_dts: cutlass.Int64, + pr_sps: cutlass.Int64, + pr_sls: cutlass.Int64, + pr_pages: cutlass.Int32, + pr_eps: cutlass.Float32, + prod: 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, + ) + # prod: this CTA's bf16 owner rows by token ([8][rows_per_rank]), and the packed-head row of token `rank` gathered + # from the cluster (128 bf16 as 64 words, completed on pr_full by 8 x rows_per_rank * 2 bytes of st.async). + pr_stage = None + pr_row = None + pr_full = None + if cutlass.const_expr(prod): + pr_stage = cutlass.Array( + cutlass.Uint16, 8 * rows_per_rank, space=cutlass.AddressSpace.smem, alignment=16 + ) + pr_row = cutlass.Array(cutlass.Int32, 64, space=cutlass.AddressSpace.smem, alignment=16) + pr_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + # 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 cutlass.const_expr(prod): + # Completed by the 8 ranks' st.async (this one's included) only on a rank that receives a token row. + prims.mbarrier_init(pr_full, 1) + if rank < num_tokens: + prims.mbarrier_arrive_expect_tx(pr_full, split * rows_per_rank * 2) + 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 if x_swiglu else 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) + ] + if cutlass.const_expr(x_swiglu): + # The up half of the same columns, K values after the gate in the row. + qz_wvecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(k_in // 2 + 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]) + if cutlass.const_expr(x_swiglu): + qz_word = _swiglu_oai_bf16x2( + qz_word, cutlass.Int32(qz_wvecs[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 flags are rotated by the previous all-reduce after its dependent release; this load is ordered by + # the wait above, which waits for ALL in-flight upstream grids, not only the immediate predecessor (PTX + # griddepcontrol; MoE test_pdl_transitive_wait.py). Keep the wait before this load. + # 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.const_expr(prod): + # The bf16 the packed row would hold (RN, as the y store below), staged by token. + if cutlass.Int32(4 * g + j) < num_tokens: + pr_stage.store( + cutlass.Uint16( + (_bf16_bits(sums[j]) >> cutlass.Int32(16)) + & cutlass.Int32(0xFFFF) + ), + idx=cutlass.Int32((4 * g + j) * rows_per_rank) + local, + ) + else: + if 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) + # The flags are rotated by the previous all-reduce after its dependent release; this load is ordered + # by the wait above, which waits for ALL in-flight upstream grids, not only the immediate + # predecessor (PTX griddepcontrol; MoE test_pdl_transitive_wait.py). Keep the wait before this load. + 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(prod): + # Transpose: thread e < 2 T sends half e % 2 (16 bytes = 8 dims) of token e / 2's staged rows to cluster + # rank e / 2, at dims [rank * rows_per_rank + 8 (e % 2), ...) of that rank's packed-head row. + pr_e = tx - cutlass.Int32(128) + if pr_e < num_tokens * cutlass.Int32(2): + pr_t = pr_e // cutlass.Int32(2) + pr_h = pr_e % cutlass.Int32(2) + pr_w = _ld_shared_v4_b32( + pr_stage.subview( + pr_t * cutlass.Int32(rows_per_rank) + pr_h * cutlass.Int32(8) + ).data_ptr() + ) + _st_async_v4( + _mapa_u32( + pr_row.subview( + rank * cutlass.Int32(rows_per_rank // 2) + pr_h * cutlass.Int32(4) + ).data_ptr(), + pr_t, + ), + pr_w[0], + pr_w[1], + pr_w[2], + pr_w[3], + _mapa_u32(pr_full.data_ptr(), pr_t), + ) + if warp_id == 4: + if rank < num_tokens: + # The producer reads step inputs (positions, slots) and writes the caches: ordered after the + # upstream grids by this thread's own grid wait (already satisfied: the activation landed). + prims.griddepcontrol(prims.GridDepAction.WAIT) + while not _test_wait_cluster(pr_full.data_ptr(), 0): + pass + pr_lo = pr_row.load(idx=lane * cutlass.Int32(2)) + pr_hi = pr_row.load(idx=lane * cutlass.Int32(2) + cutlass.Int32(1)) + mask = cutlass.Int32(0xFFFF) + producer_head( + rank, + n_tile, + lane, + bf16_bits_to_f32(pr_lo & mask), + bf16_bits_to_f32((pr_lo >> cutlass.Int32(16)) & mask), + bf16_bits_to_f32(pr_hi & mask), + bf16_bits_to_f32((pr_hi >> cutlass.Int32(16)) & mask), + ProducerArgs( + packed=cutlass.Int64(0), + q_out=pr_q, + index_q_out=pr_iq, + index_k=pr_ik, + kv_data=pr_kvd, + kv_scale=pr_kvs, + out_cache_loc=pr_loc, + kv_quant_scale=pr_quant, + position_ids=pr_pos, + rotary_cos_sin=pr_rope, + q_weight=pr_qw, + k_weight=pr_kw, + index_q_weight=pr_iqw, + index_k_weight=pr_ikw, + eps=pr_eps, + num_pages=pr_pages, + index_page_stride=pr_ips, + index_token_stride=pr_its, + data_page_stride=pr_dps, + data_plane_stride=pr_dls, + data_token_stride=pr_dts, + scale_page_stride=pr_sps, + scale_plane_stride=pr_sls, + ), + ) + 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], + x_swiglu: cutlass.Constexpr[bool], + # prod: see m3_mx_gemv_kernel (zeros when off). + pr_q: cutlass.Int64, + pr_iq: cutlass.Int64, + pr_ik: cutlass.Int64, + pr_kvd: cutlass.Int64, + pr_kvs: cutlass.Int64, + pr_loc: cutlass.Int64, + pr_quant: cutlass.Int64, + pr_pos: cutlass.Int64, + pr_rope: cutlass.Int64, + pr_qw: cutlass.Int64, + pr_kw: cutlass.Int64, + pr_iqw: cutlass.Int64, + pr_ikw: cutlass.Int64, + pr_ips: cutlass.Int64, + pr_its: cutlass.Int64, + pr_dps: cutlass.Int64, + pr_dls: cutlass.Int64, + pr_dts: cutlass.Int64, + pr_sps: cutlass.Int64, + pr_sls: cutlass.Int64, + pr_pages: cutlass.Int32, + pr_eps: cutlass.Float32, + prod: cutlass.Constexpr[bool], +) -> None: + # x_swiglu's up half rides the slot din uses for its norm weights, and it activates the bf16 x. + assert not (din and x_swiglu) and (x_bf16 or not x_swiglu), ( + "x_swiglu needs x_bf16 and excludes din" + ) + 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, x_swiglu, + pr_q, pr_iq, pr_ik, pr_kvd, pr_kvs, pr_loc, pr_quant, pr_pos, pr_rope, pr_qw, pr_kw, pr_iqw, pr_ikw, pr_ips, + pr_its, pr_dps, pr_dls, pr_dts, pr_sps, pr_sls, pr_pages, pr_eps, prod, + ).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_m31_nvfp4__sm_103__tp8/ops/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/mx_gemv/op.py new file mode 100644 index 000000000000..882fc18b3431 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/mx_gemv/op.py @@ -0,0 +1,649 @@ +# 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). The split count sets the fp32 association of the K partials, so splits +# 8, 4, 2 and 1 are NOT bitwise equal in general (wide-range block scales show it): an output +# row's bits depend on the split its row count selects (batch invariance holds within a split band only, see +# test_modeling_v2_m3_mx_gemv_split_bands.py). +_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) + + +# #117: the producer epilogue serves the sparse-layer QKV + index projection, 12 packed heads of 128 (8 q, k, v, +# index-q, index-k), one head per split-8 cluster. +PROD_N = 12 * 128 # capacity: fixed: the producer's packed heads +# The kernel's producer arguments on a call without the producer epilogue (20 addresses / strides, pages, eps). +PROD_OFF = (0,) * 21 + (0.0,) + + +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], + x_swiglu: bool = False, + prod: bool = False, +) -> 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 + if prod and not ( + n == PROD_N + and split == 8 + and m <= split + and st_async + and not x_bf16 + and not x_swiglu + and ar is None + ): + # The producer epilogue: the sparse QKV + index projection (one 128-row packed head per cluster of 8), at most + # one token per cluster rank, the st.async exchange, MXFP8 input, no all-reduce push. + return None + 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, + x_swiglu, + prod, + ) + return key, (ar_ws, ar_flags, ar_rank) + + +def _compile( + key: tuple, args: tuple, m: int, runtime_ar: tuple, stream, extra: tuple, prod_args: 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, + x_swiglu, + prod, + ) = 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, + x_swiglu, + *prod_args, + prod, + ) + return fn + + +def _views( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight, + weight_scale, + y, + x_bf16: bool, + x_swiglu: bool = False, +) -> tuple: + # x_swiglu: x is [M, 2K] (gate | up); the GEMV's K is half its width. + k = x.shape[1] // 2 if x_swiglu else 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, + x_swiglu: bool = False, + prod_args: 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 + if x_swiglu: + k //= 2 + n = weight.shape[0] + if weight.shape[1] != k: + return None + prod = prod_args is not None + resolved = _resolve(m, n, k, x_bf16, trigger_early, cfg, use_pdl, ar, x_swiglu, prod) + if resolved is None: + return None + key, runtime_ar = resolved + if not prod: + prod_args = PROD_OFF + 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" + ) + # prod: no packed row is stored (the producer epilogue consumes it); y is the placeholder. + y = torch.empty(0 if prod else m, n, dtype=torch.bfloat16, device=x.device) + args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) + 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, prod_args) + ar_ws, ar_flags, ar_rank = runtime_ar + fn(*args, m, ar_ws, ar_flags, ar_rank, stream, *extra, *prod_args) + 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, + x_swiglu: bool = False, + prod: bool = False, + trigger_early: bool = True, +) -> 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. ``trigger_early`` is part of the key (the early PDL release is a compile-time choice). + ``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) + # prod: the producer epilogue needs the split-8 cluster; the TP8 table's default (N, K) config is split 8 at every + # row count. The plain GEMV takes split 4 at 2 rows, which is not bitwise split 8 (see _CONFIGS). + resolved = _resolve( + m, + n, + k, + x_bf16, + trigger_early, + get_config(n, k) if prod else None, + use_pdl, + push, + x_swiglu, + prod, + ) + 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, + 2 * k if x_swiglu else 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, x_swiglu) + stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) + _compile(key, args, m, runtime_ar, stream, _extra_args(dev), PROD_OFF) + 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) + + +def m3_mx_gemv_swiglu( + gate_up: 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]: + """``m3_mx_gemv_bf16(swiglu(gate_up), ...)`` in one launch: ``gate_up`` bf16 [M <= 8, 2K] (contiguous, gate then + up) is activated in the kernel's quantize prologue with the swigluoai arithmetic of the Triton silu_and_mul kernel + (alpha 1.702, limit 7, the (up + 1) offset; bit for bit), then quantized as the bf16 path does.""" + if ( + gate_up.dtype != torch.bfloat16 + or gate_up.dim() != 2 + or not gate_up.is_contiguous() + or gate_up.data_ptr() % 16 != 0 + or gate_up.shape[1] % 2 != 0 + ): + return None + return _run(gate_up, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar, True) + + +def m3_mx_gemv_producer( + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + q: torch.Tensor, + index_q: 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, + trigger_early: bool = True, + use_pdl: Optional[bool] = None, +) -> bool: + """#117: ``m3_producer(m3_mx_gemv(xq, xsf, weight, weight_scale), ...)`` in one launch, bit for bit: the QKV + + index GEMV of the sparse layer (``weight`` [12 x 128, K]) whose owner epilogue hands each token's packed head to + the producer phase (``m3_layer.phase_producer.producer_head``). Writes ``q`` FP8 [T, 8, 128] and ``index_q`` FP8 + [T, 1, 128] and the token's cache slot (NVFP4 K / V + scales, FP8 index-K) exactly as the producer does; the + caches and tensors are those ``minimax_m3_producer_pw`` takes (kv_data [pages, 2, 1, 128, 64] u8, kv_scale [pages, + 2, 1, 128, 8] u8, index_k [pages, 1, 128, 128] e4m3). Returns False (nothing launched) when the call is declined. + """ + m = xq.shape[0] + if ( + m > 8 + or weight.shape[0] != PROD_N + or q.shape != (m, 8, 128) + or index_q.shape != (m, 1, 128) + or not (q.is_contiguous() and index_q.is_contiguous()) + or kv_data.dim() != 5 + or kv_scale.dim() != 5 + or index_k.dim() != 4 + ): + return False + tensors = ( + q, + index_q, + index_k, + kv_data, + kv_scale, + out_cache_loc, + kv_quant_scale, + position_ids, + rotary_cos_sin, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + ) + prod_args = ( + *(t.data_ptr() for t in tensors), + index_k.stride(0), + index_k.stride(2), + kv_data.stride(0), + kv_data.stride(1), + kv_data.stride(3), + kv_scale.stride(0), + kv_scale.stride(1), + kv_data.shape[0], + float(eps), + ) + out = _run( + xq, + xsf, + weight, + weight_scale, + False, + trigger_early, + get_config(weight.shape[0], weight.shape[1]), # split 8 at every row count (see precompile) + use_pdl, + None, + False, + prod_args, + ) + return out is not None + + +_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_m31_nvfp4__sm_103__tp8/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/sampling/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/sampling/__init__.py new file mode 100644 index 000000000000..b3611a4b7866 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/sampling/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sampling entries this target alone calls: the spec sampler's top-p renorm override (ST-01 TOPP_DET), its softmax +override (ST-01-SM SOFTMAX_MC) and its chain speculative sampling override (ST-01-REJ REJ_MC).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/sampling/m31_chain_spec_mc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/sampling/m31_chain_spec_mc.py new file mode 100644 index 000000000000..a0b26091aced --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/sampling/m31_chain_spec_mc.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ST-01-REJ REJ_MC: the spec sampler's chain speculative sampling (rejection) on many CTAs per request, exact vs +``torch.ops.trtllm.chain_speculative_sampling`` (``trtllm::minimax_m31_chain_speculative_sampling_mc``, +kernels/minimaxM31/chainSpecSamplingMcSm103.cu). + +The target installs ``chain_spec_mc`` as the rejection sampler's override +(``speculative/interface.set_chain_speculative_sampling_override``); it returns None, and the sampler keeps the stock +op, whenever the op is not built or the input is outside the kernel's contract. + +Workspace: one zero-initialized byte buffer per device, ``MAX_BATCH`` requests of +``minimax_m31_chain_speculative_sampling_mc_ws_bytes()``, kept for the life of the process (CUDA graphs captured with +it keep its address). The kernel leaves it ready for the next launch, so it is zeroed once. ``prebuild`` allocates it +outside any CUDA-graph capture (the target calls it from ``cache_derived_state``, before the executor's warmup and +captures); a call that finds no workspace allocates one when the stream is not capturing, and inside a capture returns +None (the stock op for that call) rather than allocating from the graph's private pool. Calls that share it must be +stream-ordered (the sampler's are). +""" + +from __future__ import annotations + +import os +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 tensorrt_llm._torch._experimental.modeling_v2.catalog.fakes import ( + m31_fake_ops as _m31_fake_ops, +) + +# The largest CTA count per request (the grid is batch x ctas_per_row 1024-thread CTAs, all co-resident). REJ v2 MB +#: one request 32.3 us at 16 (v1) -> 24.8 us at 64 (v2), 25.8 at 32, 25.5 at 128; bitwise at every count. +MAX_CTAS_PER_ROW = 64 + +# The requests the workspace holds and the override takes: the TP8 / TP4 K7 calls up to c8. Microbenchmarks +# gated it bitwise at batch 1 / 2 / 4 / 8 and measured 31.2 vs 41.2 us per call at batch 1; larger calls keep the stock +# kernel. +MAX_BATCH = 8 + +# The widest vocab the kernel takes (64 tiles of 4096; kChainSpecMcMaxVocab). +MAX_VOCAB = 64 * 4096 +MAX_DRAFT = ( + 32 # capacity: fixed: the kernel's shared acceptance probabilities (kChainSpecMcMaxDraft) +) + +# The op when it is built (or overlaid as tensorrt_llm/libs/libm31_chain_spec_mc.so), decided once per process. +_STATE: dict = {} +# Per device index: the persistent workspace and the SM count. +_WORKSPACES: dict = {} +_SMS: dict = {} + + +def chain_spec_mc_op(): + """``trtllm::minimax_m31_chain_speculative_sampling_mc`` if this build carries it (in-tree, or the overlay + library), else None.""" + if "op" not in _STATE: + op = getattr(torch.ops.trtllm, "minimax_m31_chain_speculative_sampling_mc", None) + if op is None: + import tensorrt_llm + + lib = os.path.join( + os.path.dirname(tensorrt_llm.__file__), "libs", "libm31_chain_spec_mc.so" + ) + if os.path.exists(lib): + torch.ops.load_library(lib) + _m31_fake_ops.register() + op = getattr(torch.ops.trtllm, "minimax_m31_chain_speculative_sampling_mc", None) + _STATE["op"] = op + _STATE["ws_bytes"] = ( + int(torch.ops.trtllm.minimax_m31_chain_speculative_sampling_mc_ws_bytes()) + if op is not None + else 0 + ) + return _STATE["op"] + + +def ws_bytes_per_row() -> int: + """Workspace bytes per request (0 when the op is not built).""" + chain_spec_mc_op() + return _STATE["ws_bytes"] + + +def sm_count(device: torch.device) -> int: + """The device's SM count, cached per device.""" + idx = torch.device(device).index + idx = torch.cuda.current_device() if idx is None else idx + if idx not in _SMS: + _SMS[idx] = torch.cuda.get_device_properties(idx).multi_processor_count + return _SMS[idx] + + +def ctas_per_row(batch: int, sms: int) -> int: + """The largest power of two <= MAX_CTAS_PER_ROW with ctas x batch <= sms (1 past that: one CTA per request needs + no co-residency).""" + g = MAX_CTAS_PER_ROW + while g > 1 and g * batch > sms: + g //= 2 + return g + + +def prebuild(device: Optional[torch.device] = None) -> Optional[torch.Tensor]: + """Allocate (once) the device's zeroed workspace for MAX_BATCH requests; must run outside a CUDA-graph capture. + Returns it, or None when the op is not built or the stream is capturing.""" + if chain_spec_mc_op() is None: + return None + dev = ( + torch.device("cuda", torch.cuda.current_device()) + if device is None + else torch.device(device) + ) + if dev.index is None: + dev = torch.device("cuda", torch.cuda.current_device()) + if dev.index not in _WORKSPACES: + if torch.cuda.is_current_stream_capturing(): + return None + _WORKSPACES[dev.index] = torch.zeros( + MAX_BATCH * ws_bytes_per_row(), dtype=torch.uint8, device=dev + ) + return _WORKSPACES[dev.index] + + +class M31ChainSpecMc(OpWrapper): + """``trtllm::chain_speculative_sampling(draft_probs, draft_token_ids, target_probs, seed, offset)`` bit for bit + (accepted tokens [batch, D + 1] and token counts [batch]) in ONE launch of batch x ctas_per_row CTAs (stock: one + CTA per request). ``workspace`` is the persistent zeroed buffer of ``prebuild``. ``__call__`` returns None (the + caller keeps the stock op) for a non-fp32, non-contiguous or misaligned probs, a vocab not divisible by 4 or above + MAX_VOCAB, mismatched shapes, non-int32 draft tokens, non-int64 or mis-sized seed / offset, inputs on different + devices, more requests than the workspace holds, a batch x ctas_per_row above the SM count, or when the op is not + built.""" + + ARCHS = frozenset({Arch.SM_103}) + + # The K7 calls were checked bitwise on GB300 over 5 cases x 64 seeds at ctas_per_row 4 / 8 / 16. + CELLS: tuple[Cell, ...] = ( + Cell(why="K7 at c1 (vocab 200064, 16 CTAs)", spec=dict(batch=1, k=7, v=200064)), + Cell(why="K7 at c2 (vocab 200064)", spec=dict(batch=2, k=7, v=200064)), + Cell(why="K7 at c4 (vocab 200064)", spec=dict(batch=4, k=7, v=200064)), + Cell(why="K7 at c8 (vocab 200064, 128 CTAs)", spec=dict(batch=8, k=7, v=200064)), + ) + + note = """ + Exact by construction: every CTA of a request runs the stock Philox acceptance loop, the (target, residual) mass is + the stock 1024 per-thread chains into the same cub BlockReduce, each tile sum is the stock per-tile + cub reduce, and one CTA runs the stock tile-order prefix and in-tile cub BlockScan search. Default flags, as the + stock kernel. + """ + + def __call__( + self, + draft_probs: torch.Tensor, + draft_token_ids: torch.Tensor, + target_probs: torch.Tensor, + seed: torch.Tensor, + offset: torch.Tensor, + workspace: Optional[torch.Tensor], + ctas_per_row: int, + ) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + if ( + workspace is None + or not all( + isinstance(t, torch.Tensor) + for t in (draft_probs, draft_token_ids, target_probs, seed, offset) + ) + or not draft_probs.is_cuda + or draft_probs.dtype != torch.float32 + or target_probs.dtype != torch.float32 + or draft_token_ids.dtype != torch.int32 + or seed.dtype != torch.int64 + or offset.dtype != torch.int64 + or draft_probs.dim() != 3 + or target_probs.dim() != 3 + or draft_token_ids.dim() != 2 + ): + return None + batch, k, v = draft_probs.shape + if ( + batch < 1 + or tuple(draft_token_ids.shape) != (batch, k) + or tuple(target_probs.shape) != (batch, k + 1, v) + or seed.numel() != batch + or offset.numel() != batch + or v % 4 != 0 + or v > MAX_VOCAB + or k > MAX_DRAFT + or any( + not t.is_contiguous() or t.device != draft_probs.device + for t in (draft_token_ids, target_probs, seed, offset, workspace) + ) + or not draft_probs.is_contiguous() + or draft_probs.data_ptr() % 16 != 0 + or target_probs.data_ptr() % 16 != 0 + or ctas_per_row < 1 + or 1024 % ctas_per_row != 0 + or (ctas_per_row > 1 and batch * ctas_per_row > sm_count(draft_probs.device)) + ): + return None + op = chain_spec_mc_op() + if op is None or workspace.numel() * workspace.element_size() < batch * ws_bytes_per_row(): + return None + tokens, num_tokens = op( + draft_probs, draft_token_ids, target_probs, seed, offset, workspace, int(ctas_per_row) + ) + return tokens, num_tokens + + def reference( + self, + draft_probs, + draft_token_ids, + target_probs, + seed, + offset, + workspace=None, + ctas_per_row=None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """The stock op (workspace and ctas_per_row are ignored).""" + return torch.ops.trtllm.chain_speculative_sampling( + draft_probs, draft_token_ids, target_probs, seed, offset + ) + + def is_valid( + self, + draft_probs, + draft_token_ids, + target_probs, + seed, + offset, + workspace=None, + ctas_per_row=None, + **kwargs, + ) -> None: + assert draft_probs.dtype == torch.float32 and draft_probs.dim() == 3, ( + "draft_probs: fp32 [batch, D, vocab]" + ) + + +m31_chain_spec_mc = M31ChainSpecMc() + + +def chain_spec_mc( + draft_probs: torch.Tensor, + draft_token_ids: torch.Tensor, + target_probs: torch.Tensor, + seed: torch.Tensor, + offset: torch.Tensor, +) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + """The rejection sampler's override: the stock op's (accepted tokens, token counts), or None (the sampler runs the + stock op) when the op is not built, the batch exceeds MAX_BATCH, no workspace can be had (first use inside a + capture) or the input is outside the contract.""" + if ( + not isinstance(draft_probs, torch.Tensor) + or not draft_probs.is_cuda + or draft_probs.dim() != 3 + ): + return None + batch = draft_probs.shape[0] + if batch < 1 or batch > MAX_BATCH or chain_spec_mc_op() is None: + return None + workspace = prebuild(draft_probs.device) + if workspace is None: + return None + return m31_chain_spec_mc( + draft_probs, + draft_token_ids, + target_probs, + seed, + offset, + workspace, + ctas_per_row(batch, sm_count(draft_probs.device)), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/sampling/m31_softmax_mc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/sampling/m31_softmax_mc.py new file mode 100644 index 000000000000..ac4d6dc8add2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/sampling/m31_softmax_mc.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ST-01-SM SOFTMAX_MC: the spec sampler's softmax with its write pass on many CTAs per row, exact vs +``flashinfer.sampling.softmax(logits, temperature)`` (``trtllm::minimax_m31_softmax_mc``, +kernels/minimaxM31/softmaxMcSm103.cu). + +The target installs ``softmax_mc`` as the sampler's softmax override (``sampler/ops/flashinfer.set_softmax_override``); +it returns None, and the sampler keeps the stock FlashInfer call, whenever the op is not built or the input is outside +the kernel's contract (FlashInfer's Path A shapes: rows <= 128 and vocab >= 24576; here also rows <= ``MAX_ROWS``). + +Workspace: one partials buffer per device (``MAX_ROWS`` x ceil(``MAX_VOCAB`` / 8192) float2, 8 KB), kept for the life +of the process (CUDA graphs captured with it keep its address). Every launch overwrites what it reads, so it needs no +initialization. ``prebuild`` allocates it outside any CUDA-graph capture (the target calls it from +``cache_derived_state``, before the executor's warmup and captures); a call that finds no workspace allocates one when +the stream is not capturing, and inside a capture returns None (stock FlashInfer for that call) rather than allocating +from the graph's private pool. Calls that share it must be stream-ordered (the sampler's are). +""" + +from __future__ import annotations + +import os +from typing import Optional, Union + +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.fakes import ( + m31_fake_ops as _m31_fake_ops, +) + +# The write pass's CTAs per row (microbenchmark: 9.2 vs FlashInfer 21.8 us per call at 7-8 rows; bitwise at +# 4 / 8 / 16). +WRITE_CTAS = 16 + +# The rows the override takes: the TP8 / TP4 K7 c1 calls (the draft's 7 rows and the verify's 8), where the MB measured +# it. Larger calls keep the stock kernel. +MAX_ROWS = 8 + +# FlashInfer's Path A bounds (OnlineSoftmax: batch <= 128 and d >= 24576); outside them FlashInfer takes its +# single-block Path B, which this kernel does not reproduce. +PATH_A_MAX_ROWS = 128 +PATH_A_MIN_VOCAB = 24576 +# FlashInfer's Path A slice (DEFAULT_SLICE_SIZE) and the widest vocab the workspace holds. +SLICE = 8192 +MAX_VOCAB = 1 << 20 + +# The op when it is built (or overlaid as tensorrt_llm/libs/libm31_softmax_mc.so), decided once per process. +_STATE: dict = {} +# Per device index: the persistent partials workspace. +_WORKSPACES: dict = {} + + +def softmax_mc_op(): + """``trtllm::minimax_m31_softmax_mc`` if this build carries it (in-tree, or the overlay library), else None.""" + if "op" not in _STATE: + op = getattr(torch.ops.trtllm, "minimax_m31_softmax_mc", None) + if op is None: + import tensorrt_llm + + lib = os.path.join( + os.path.dirname(tensorrt_llm.__file__), "libs", "libm31_softmax_mc.so" + ) + if os.path.exists(lib): + torch.ops.load_library(lib) + _m31_fake_ops.register() + op = getattr(torch.ops.trtllm, "minimax_m31_softmax_mc", None) + _STATE["op"] = op + return _STATE["op"] + + +def partials_numel(rows: int, d: int) -> int: + """fp32 elements of the partials workspace for rows x d (one float2 per row and 8192-wide slice).""" + return 2 * rows * ((d + SLICE - 1) // SLICE) + + +def prebuild(device: Optional[torch.device] = None) -> Optional[torch.Tensor]: + """Allocate (once) the device's partials workspace for MAX_ROWS rows of up to MAX_VOCAB; must run outside a + CUDA-graph capture. Returns it, or None when the op is not built or the stream is capturing.""" + if softmax_mc_op() is None: + return None + dev = ( + torch.device("cuda", torch.cuda.current_device()) + if device is None + else torch.device(device) + ) + if dev.index is None: + dev = torch.device("cuda", torch.cuda.current_device()) + if dev.index not in _WORKSPACES: + if torch.cuda.is_current_stream_capturing(): + return None + _WORKSPACES[dev.index] = torch.empty( + partials_numel(MAX_ROWS, MAX_VOCAB), dtype=torch.float32, device=dev + ) + return _WORKSPACES[dev.index] + + +class M31SoftmaxMc(OpWrapper): + """FlashInfer's ``softmax(logits, temperature)`` of fp32 logits [rows, d], bit for bit, on FlashInfer's Path A + shapes (rows <= 128, d >= 24576): FlashInfer's Map kernel, then its Reduce kernel's prologue in each of + ``write_ctas`` CTAs per row, each writing its share of the row (FlashInfer: one CTA writes the whole row). + ``temperature`` is None (1.0), a Python number (one temperature) or an fp32 [rows] tensor (one per row), as + FlashInfer takes it. ``partials`` is the persistent workspace of ``prebuild``. ``__call__`` returns None (the caller + keeps the stock call) for a non-fp32, non-contiguous or misaligned input, a width not divisible by 4, a shape + outside Path A or the workspace, a temperature tensor that is not a contiguous fp32 [rows] tensor on logits' + device, or when the op is not built.""" + + ARCHS = frozenset({Arch.SM_103}) + + # The c1 K7 calls were checked bitwise on GB300. + CELLS: tuple[Cell, ...] = ( + Cell(why="the K7 draft's 7 rows at c1 (vocab 200064)", spec=dict(rows=7, d=200064)), + Cell(why="the K7 verify's 8 rows at c1 (vocab 200064)", spec=dict(rows=8, d=200064)), + Cell(why="one row (the smallest grid)", spec=dict(rows=1, d=200064)), + ) + + note = """ + Exact by construction: the Map kernel is FlashInfer's, every write CTA reduces the same partials with the same cub + BlockReduce and Float2SoftmaxReduceOp as FlashInfer's single Reduce CTA, and the per-element write + keeps FlashInfer's source form. Built with --use_fast_math, as FlashInfer's JIT builds it. + """ + + def __call__( + self, + logits: torch.Tensor, + temperature: Union[torch.Tensor, float, None], + partials: Optional[torch.Tensor], + write_ctas: int, + ) -> Optional[torch.Tensor]: + if ( + partials is None + or not isinstance(logits, torch.Tensor) + or not logits.is_cuda + or logits.dtype != torch.float32 + or logits.dim() != 2 + or not logits.is_contiguous() + or logits.shape[0] < 1 + or logits.shape[0] > PATH_A_MAX_ROWS + or logits.shape[1] < PATH_A_MIN_VOCAB + or logits.shape[1] % 4 != 0 + or logits.data_ptr() % 16 != 0 + or partials.device != logits.device + or partials.numel() * partials.element_size() + < 4 * partials_numel(logits.shape[0], logits.shape[1]) + or write_ctas < 1 + or write_ctas > 65535 + ): + return None + if temperature is None: + temp_arr, temp_val = None, 1.0 + elif isinstance(temperature, torch.Tensor): + if ( + temperature.dtype != torch.float32 + or temperature.dim() != 1 + or temperature.shape[0] != logits.shape[0] + or not temperature.is_contiguous() + or temperature.device != logits.device + ): + return None + temp_arr, temp_val = temperature, 0.0 + elif isinstance(temperature, (int, float)) and not isinstance(temperature, bool): + temp_arr, temp_val = None, float(temperature) + else: + return None + op = softmax_mc_op() + if op is None: + return None + out = torch.empty_like(logits) + op(logits, temp_arr, temp_val, partials, out, int(write_ctas)) + return out + + def reference(self, logits, temperature, partials=None, write_ctas=None) -> torch.Tensor: + """FlashInfer's softmax (partials and write_ctas are ignored).""" + import flashinfer.sampling + + return flashinfer.sampling.softmax(logits, temperature) + + def is_valid(self, logits, temperature, partials=None, write_ctas=None, **kwargs) -> None: + assert logits.dtype == torch.float32 and logits.dim() == 2, "logits: fp32 [rows, d]" + + +m31_softmax_mc = M31SoftmaxMc() + + +def softmax_mc( + logits: torch.Tensor, temperature: Union[torch.Tensor, float, None] +) -> Optional[torch.Tensor]: + """The sampler's softmax override: FlashInfer's softmax of logits [rows, d] at temperature, or None (the sampler + runs stock FlashInfer) when the op is not built, rows exceed MAX_ROWS, d exceeds MAX_VOCAB, no workspace can be had + (first use inside a capture) or the input is outside the contract.""" + if not isinstance(logits, torch.Tensor) or not logits.is_cuda or logits.dim() != 2: + return None + rows = logits.shape[0] + if rows < 1 or rows > MAX_ROWS or logits.shape[1] > MAX_VOCAB or softmax_mc_op() is None: + return None + partials = prebuild(logits.device) + if partials is None: + return None + return m31_softmax_mc(logits, temperature, partials, WRITE_CTAS) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/sampling/m31_topp_det.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/sampling/m31_topp_det.py new file mode 100644 index 000000000000..0d28b1e2699c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/sampling/m31_topp_det.py @@ -0,0 +1,248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ST-01 TOPP_DET: the spec sampler's top-p renorm as one persistent C++ launch, exact vs FlashInfer's deterministic +``top_p_renorm_probs`` (``trtllm::minimax_m31_topp_renorm_det``, kernels/minimaxM31/toppRenormDetSm103.cu). + +TOPP_DET v3 (``trtllm::minimax_m31_topp_renorm_det_v3``, kernels/minimaxM31/toppRenormDetV3Sm103.cu) is the same +renorm, bitwise v1.2's output, with the redundant per-CTA selection and a programmatic dependent launch. It is taken +when it is built (in-tree, or overlaid as tensorrt_llm/libs/libm31_topp_det_v3.so) unless +``TRTLLM_M31_DISABLE_TOPPDET_V3`` is set; v1.2 otherwise. ``TRTLLM_M31_DISABLE_TOPPDET_T3`` keeps v3 on its global +radix passes (no candidate list). + +The target installs ``top_p_renorm_det`` as the sampler's top-p renorm override +(``sampler/ops/flashinfer.set_top_p_renorm_override``); it returns None, and the sampler keeps the stock FlashInfer +call, whenever the op is not built or the input is outside the kernel's contract. + +Workspace: one zero-initialized byte buffer per device, ``MAX_ROWS`` rows of +``minimax_m31_topp_renorm_det_ws_bytes()``, kept for the life of the process (CUDA graphs captured with it keep its +address). The kernel leaves it ready for the next launch, so it is zeroed once. ``prebuild`` allocates it outside any +CUDA-graph capture (the target calls it from ``cache_derived_state``, before the executor's warmup and captures); a +call that finds no workspace allocates one when the stream is not capturing, and inside a capture returns None +(stock FlashInfer for that call) rather than allocating from the graph's private pool. +""" + +from __future__ import annotations + +import hashlib +import os +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 tensorrt_llm._torch._experimental.modeling_v2.catalog.fakes import ( + m31_fake_ops as _m31_fake_ops, +) +from tensorrt_llm.logger import logger + +from ...m31_switches import m31_disabled + +# The largest CTA count per row (the persistent grid is rows x ctas_per_row 1024-thread CTAs, all co-resident). +MAX_CTAS_PER_ROW = 16 + +# The rows the workspace holds and the override takes: the TP8 / TP4 K7 c1 calls (the draft's 7 rows and the +# verify's 8), where the microbenchmarks measured it faster than stock FlashInfer (v1.2: 54.3 vs 104.2 us per +# call at 7 peaked rows, 54.8 vs 105.1 at 8). At 56 / 64 rows v1 was slower than FlashInfer's default mode (340 vs +# 187 / 199 us) and larger shapes are not re-measured, so larger calls keep the stock kernel. +MAX_ROWS = 8 + +# The op when it is built (or overlaid as tensorrt_llm/libs/libm31_topp_det.so), decided once per process. +_STATE: dict = {} +# Per device index: the persistent workspace and the SM count. +_WORKSPACES: dict = {} +_SMS: dict = {} + + +def _load_op(name: str, lib_name: str): + """``torch.ops.trtllm.`` if registered (in-tree), else after loading tensorrt_llm/libs/ when it + exists, else None. Returns (op, the library path loaded or None).""" + op = getattr(torch.ops.trtllm, name, None) + if op is not None: + return op, None + import tensorrt_llm + + lib = os.path.join(os.path.dirname(tensorrt_llm.__file__), "libs", lib_name) + if not os.path.exists(lib): + return None, None + torch.ops.load_library(lib) + _m31_fake_ops.register() + return getattr(torch.ops.trtllm, name, None), lib + + +def topp_det_op(): + """The renorm op this process uses, decided once: ``trtllm::minimax_m31_topp_renorm_det_v3`` when it is built + and TOPPDET_V3 is not disabled, else ``trtllm::minimax_m31_topp_renorm_det`` when built, else None.""" + if "op" not in _STATE: + op, ws, v3 = None, 0, False + if not m31_disabled("TOPPDET_V3"): + op, lib = _load_op("minimax_m31_topp_renorm_det_v3", "libm31_topp_det_v3.so") + if op is not None: + v3 = True + ws = int(torch.ops.trtllm.minimax_m31_topp_renorm_det_v3_ws_bytes()) + tag = "in-tree" + if lib is not None: + with open(lib, "rb") as f: + digest = hashlib.md5(f.read(), usedforsecurity=False).hexdigest()[:8] + tag = f"libm31_topp_det_v3.so md5 {digest}" + logger.info(f"[M3.1] TOPP_DET v3 private op engaged ({tag})") + if op is None: + op, _ = _load_op("minimax_m31_topp_renorm_det", "libm31_topp_det.so") + ws = ( + int(torch.ops.trtllm.minimax_m31_topp_renorm_det_ws_bytes()) + if op is not None + else 0 + ) + _STATE["op"], _STATE["ws_bytes"], _STATE["v3"] = op, ws, v3 + # v3's candidate-list radix passes (bitwise the same output) unless TOPPDET_T3 is disabled. + _STATE["cand_list"] = v3 and not m31_disabled("TOPPDET_T3") + if v3: + logger.info( + "[M3.1] TOPP_DET v3 T3 candidate-list fused passes: " + f"{'on' if _STATE['cand_list'] else 'off'} (cap 4096)" + ) + return _STATE["op"] + + +def is_v3() -> bool: + """Whether the engaged op is TOPP_DET v3.""" + topp_det_op() + return _STATE["v3"] + + +def ws_bytes_per_row() -> int: + """Workspace bytes per row (0 when the op is not built).""" + topp_det_op() + return _STATE["ws_bytes"] + + +def sm_count(device: torch.device) -> int: + """The device's SM count, cached per device.""" + idx = torch.device(device).index + idx = torch.cuda.current_device() if idx is None else idx + if idx not in _SMS: + _SMS[idx] = torch.cuda.get_device_properties(idx).multi_processor_count + return _SMS[idx] + + +def ctas_per_row(rows: int, sms: int) -> int: + """The largest power of two <= MAX_CTAS_PER_ROW with ctas x rows <= sms (1 past that: one CTA per row needs no + co-residency).""" + g = MAX_CTAS_PER_ROW + while g > 1 and g * rows > sms: + g //= 2 + return g + + +def prebuild(device: Optional[torch.device] = None) -> Optional[torch.Tensor]: + """Allocate (once) the device's zeroed workspace for MAX_ROWS rows; must run outside a CUDA-graph capture. + Returns it, or None when the op is not built or the stream is capturing.""" + if topp_det_op() is None: + return None + dev = ( + torch.device("cuda", torch.cuda.current_device()) + if device is None + else torch.device(device) + ) + if dev.index is None: + dev = torch.device("cuda", torch.cuda.current_device()) + if dev.index not in _WORKSPACES: + if torch.cuda.is_current_stream_capturing(): + return None + _WORKSPACES[dev.index] = torch.zeros( + MAX_ROWS * ws_bytes_per_row(), dtype=torch.uint8, device=dev + ) + return _WORKSPACES[dev.index] + + +class M31ToppDet(OpWrapper): + """FlashInfer's ``top_p_renorm_probs(probs, top_p, is_deterministic=True)`` of fp32 probs [rows, d] with fp32 + per-row top_p [rows], bit for bit on every row with top_p < 1, in ONE launch of rows x ctas_per_row CTAs + (FlashInfer's: five launches, apply on one CTA per row, slow on high-entropy rows). Rows with top_p >= 1 keep every + element (FlashInfer drops tail buckets there in ~43% of rows, in both of its modes); their output is + ``ftz(ftz(p) * norm)``. ``workspace`` is the persistent zeroed buffer of ``prebuild``. ``__call__`` returns None + (the caller keeps the stock call) for a non-fp32, non-contiguous or misaligned input, a width not divisible by 4, + a top_p that is not a contiguous fp32 [rows] tensor on probs' device, more rows than the workspace holds, or when + the op is not built.""" + + ARCHS = frozenset({Arch.SM_103}) + + # The c1 K7 calls were checked bitwise on GB300 over 8 row classes x 4 top_p modes x 16 seeds. + CELLS: tuple[Cell, ...] = ( + Cell(why="the K7 draft's 7 rows at c1 (vocab 200064)", spec=dict(rows=7, d=200064)), + Cell(why="the K7 verify's 8 rows at c1 (vocab 200064)", spec=dict(rows=8, d=200064)), + Cell(why="one row (the smallest grid, 16 CTAs)", spec=dict(rows=1, d=200064)), + ) + + note = """ + Exact by construction: the det histograms are integer (any aggregation order is FlashInfer's), the per-pass + selection is FlashInfer's last-block code on them, and the kept sum is FlashInfer's 1024 per-thread chains into + the same cub BlockReduce. Built with --use_fast_math, as FlashInfer's JIT builds it. + """ + + def __call__( + self, + probs: torch.Tensor, + top_p: torch.Tensor, + workspace: Optional[torch.Tensor], + ctas_per_row: int, + ) -> Optional[torch.Tensor]: + if ( + not isinstance(top_p, torch.Tensor) + or workspace is None + or not probs.is_cuda + or probs.dtype != torch.float32 + or probs.dim() != 2 + or not probs.is_contiguous() + or probs.shape[0] < 1 + or probs.shape[1] % 4 != 0 + or probs.data_ptr() % 16 != 0 + or top_p.dtype != torch.float32 + or top_p.dim() != 1 + or top_p.shape[0] != probs.shape[0] + or not top_p.is_contiguous() + or top_p.device != probs.device + or workspace.device != probs.device + or ctas_per_row < 1 + or 1024 % ctas_per_row != 0 + ): + return None + op = topp_det_op() + if ( + op is None + or workspace.numel() * workspace.element_size() < probs.shape[0] * ws_bytes_per_row() + ): + return None + out = torch.empty_like(probs) + if is_v3(): + op(probs, top_p, workspace, out, int(ctas_per_row), True, _STATE["cand_list"]) + else: + op(probs, top_p, workspace, out, int(ctas_per_row)) + return out + + def reference(self, probs, top_p, workspace=None, ctas_per_row=None) -> torch.Tensor: + """FlashInfer's deterministic renorm (rows with top_p < 1; workspace and ctas_per_row are ignored).""" + import flashinfer.sampling + + return flashinfer.sampling.top_p_renorm_probs(probs, top_p, is_deterministic=True) + + def is_valid(self, probs, top_p, workspace=None, ctas_per_row=None, **kwargs) -> None: + assert probs.dtype == torch.float32 and probs.dim() == 2, "probs: fp32 [rows, d]" + + +m31_topp_det = M31ToppDet() + + +def top_p_renorm_det(probs: torch.Tensor, top_p: torch.Tensor) -> Optional[torch.Tensor]: + """The sampler's top-p renorm override: the det renorm of probs [rows, d] at per-row top_p, or None (the sampler + runs stock FlashInfer) when the op is not built, rows exceed MAX_ROWS, no workspace can be had (first use inside + a capture) or the input is outside the contract.""" + if not isinstance(top_p, torch.Tensor) or not probs.is_cuda or probs.dim() != 2: + return None + rows = probs.shape[0] + if rows < 1 or rows > MAX_ROWS or topp_det_op() is None: + return None + workspace = prebuild(probs.device) + if workspace is None: + return None + return m31_topp_det(probs, top_p, workspace, ctas_per_row(rows, sm_count(probs.device))) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/spec/m3_capture_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/spec/m3_chunk_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/spec/m3_draft_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/spec/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/spec/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/spec/m3_sampled_accept_ws_words.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/spec/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/spec/m3_slot_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/spec/m3_step_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/ops/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp8/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/weights.py new file mode 100644 index 000000000000..870f60035f74 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp8/weights.py @@ -0,0 +1,928 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight tables and loaders: MiniMax-M3.1 (MXFP8 + packed NVFP4 experts) / 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.models.checkpoints.hf.minimaxm3_weight_mapper import ( + reciprocal_global_scale, +) +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 fused_projection(sparse_config) -> bool: + """Whether the sparse layers run the five-way packed QKV + index projection. + + On by default: this target requires the MSA implementation and an FP8 indexer, which the packed projection needs. + An explicit ``fuse_qkv_index_projection`` in the deployment's sparse attention config wins, so setting it to False + is the kill switch (the QKV projection plus a separate MXFP8 index projection). + """ + if sparse_config is None: + return False + # A config that does not record which fields were set (a plain namespace or dataclass) is taken at its word. + fields_set = getattr(sparse_config, "model_fields_set", None) + if fields_set is None or "fuse_qkv_index_projection" in fields_set: + return bool(getattr(sparse_config, "fuse_qkv_index_projection", False)) + return True + + +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 _divisor(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """A stored global scale (multiplier) -> the modelopt divisor: its FP32 reciprocal.""" + return reciprocal_global_scale("expert global scale", t) + + +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]`` as the checkpoint's MXFP8 (E4M3) values: this rank's index-Q head (rank r holds head + r // 2, as the five-way projection's ``_heads`` split) and the replicated index-K, the same rows it packs.""" + iq, ik = t + return torch.cat([_heads(iq, d, d.index_dim), _dev(ik, d)], dim=0).contiguous() + + +def _index_qk_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """The UE8M0 block scales of ``_index_qk_w``, 128x4-swizzled as every MXFP8 projection's.""" + iq, ik = t + return _swizzle(torch.cat([_heads(iq, d, d.index_dim), _dev(ik, d)], dim=0)) + + +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_mlp_layers(d): + """Layers with the dense SwiGLU-OAI MLP (``moe_layer_freq`` 0).""" + return d.dense_mlp + + +def _moe_layers(d): + return d.moe + + +def _sparse_attn_layers(d): + """Layers with MSA sparse attention (``sparse_attention_freq`` 1): every layer of MiniMax-M3.1.""" + return d.sparse_attn + + +def _three_way_qkv_layers(d): + """Layers whose QKV is the three-way [Q | K | V]: the dense-attention layers, and every layer without the + five-way packed projection (``fuse_qkv_index_projection=False``).""" + return d.dense_attn if d.fused else range(d.num_layers) + + +def _five_way_qkv_layers(d): + return d.sparse_attn if d.fused else range(0) + + +def _unfused_sparse_layers(d): + return range(0) if d.fused else d.sparse_attn + + +_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 index projection [index-Q | index-K] of the unfused sparse layers: the checkpoint's MXFP8 + # values with their block scales (the rows of the five-way projection's index part), run as an MXFP8 GEMM. + W( + "idx_w", + shape=lambda d: (d.index_heads * d.index_dim + d.index_dim, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(_ATTN + ".index_q_proj.weight", _ATTN + ".index_k_proj.weight"), + transform=_index_qk_w, + layers=_unfused_sparse_layers, + ), + W( + "idx_s", + shape=lambda d: ((d.index_heads * d.index_dim + d.index_dim) * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + _ATTN + ".index_q_proj.weight_scale_inv", _ATTN + ".index_k_proj.weight_scale_inv" + ), + transform=_index_qk_s, + layers=_unfused_sparse_layers, + ), + W( + "idx_q_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_q_norm.weight"), + layers=_sparse_attn_layers, + ), + W( + "idx_k_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_k_norm.weight"), + layers=_sparse_attn_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_mlp_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_mlp_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_mlp_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_mlp_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=_moe_layers, + ), + W( + "router_bias", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_one(_MOE + ".e_score_correction_bias"), + transform=_f32, + layers=_moe_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=_moe_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=_moe_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=_moe_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=_moe_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_packed", "w1.weight_packed"), + transform=_expert_w31, + layers=_moe_layers, + ), + W( + "e_w2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2), + dtype=torch.uint8, + src=_per_expert("w2.weight_packed"), + transform=_expert_w2, + layers=_moe_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=_moe_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=_moe_layers, + ), + # The per-tensor scalars of every expert: the checkpoint stores the multipliers into the quantized range + # (``*_global_scale``); the modelopt divisors modeling derives the kernels' scalars from are their FP32 + # reciprocals, taken by the same helper as the built-in loader. + *( + W( + f"e_{name}", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_per_expert(suffix), + transform=_divisor, + layers=_moe_layers, + ) + for name, suffix in ( + ("isc1", "w1.input_global_scale"), + ("isc3", "w3.input_global_scale"), + ("isc2", "w2.input_global_scale"), + ("ws2_1", "w1.weight_global_scale"), + ("ws2_3", "w3.weight_global_scale"), + ("ws2_2", "w2.weight_global_scale"), + ) + ), + # --- 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"]] + moe_freq = [ + int(f) for f in getattr(cfg, "moe_layer_freq", None) or [1] * cfg.num_hidden_layers + ] + 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=fused_projection(core.model_config.sparse_attention_config), + num_layers=cfg.num_hidden_layers, + sparse_attn=[i for i, f in enumerate(freq) if f], + dense_attn=[i for i, f in enumerate(freq) if not f], + moe=[i for i, f in enumerate(moe_freq) if f], + dense_mlp=[i for i, f in enumerate(moe_freq) if not f], + 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..03977da8e8af --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/routing.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Route the MiniMax-M3.1 sm_103 TP4/TP8 targets; other shapes decline.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Optional, Tuple + +from tensorrt_llm.quantization.mode import QuantAlgo + +from ..._router_index import NoTarget + +TEXT_ONLY = True +_SM = (10, 3) +_SHAPE = (60, 6144, 128, 3072, 64, 4) +_M31_ROPE_THETA = 10_000_000 +_M31_SPARSE_ATTENTION_FREQ = (1,) * 60 +_M31_EXPERTS = ("NVFP4", "nvfp4-pack-quantized") +_M31_DSPARK_ARCH = "DSparkMiniMaxDraftModel" +_M31_DSPARK_GEOMETRY = dict( + hidden_size=6144, vocab_size=200064, num_key_value_heads=4, head_dim=128 +) +_M31_DSPARK_TARGET_LAYERS = [3, 17, 31, 45, 59] + + +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 _rope_theta(text_config: Any) -> Any: + 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" + ) + return theta + + +def _m31_mismatch( + pretrained_config: Any, text_config: Any, quant_algo: Optional[QuantAlgo] +) -> Optional[str]: + """Why the checkpoint is not the MiniMax-M3.1 export the m31 targets build, or None.""" + if text_config is None: + return "no text_config" + reasons = [] + theta = _rope_theta(text_config) + if theta != _M31_ROPE_THETA: + reasons.append(f"rope_theta {theta}, not {_M31_ROPE_THETA}") + sparse = _field(text_config, "sparse_attention_config") + freq = None if sparse is None else _field(sparse, "sparse_attention_freq") + enabled = sparse is not None and _field(sparse, "use_sparse_attention") is not False + if not enabled or freq is None or tuple(int(f) for f in freq) != _M31_SPARSE_ATTENTION_FREQ: + reasons.append("sparse attention not on every layer") + declared = getattr(pretrained_config, "quantization_config", None) or {} + experts = (_field(declared, "moe_quant_algo"), _field(declared, "moe_quant_format")) + if quant_algo != QuantAlgo.MXFP8 or experts != _M31_EXPERTS: + reasons.append( + f"quantization {quant_algo} with experts {experts}, not MXFP8 with {_M31_EXPERTS}" + ) + 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 widths 768 at TP4 and 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 _m31_speculation_mismatch(spec_config: Optional[Any]) -> Optional[str]: + """Why the speculation is not one the M3.1 targets serve, or None when it is (none, or a standalone MiniMax DSpark + drafter). Read from the drafter checkpoint's ``config.json``.""" + if spec_config is None: + return None + mode = spec_config.spec_dec_mode + if not (mode.use_one_engine() and mode.is_dspark()): + return f"speculative decoding {mode.name}, not standalone DSpark" + if getattr(spec_config, "draft_is_embedded_in_target", False): + return "a DSpark drafter embedded in the target checkpoint" + path = getattr(spec_config, "speculative_model", None) + if path is None: + return "no DSpark drafter checkpoint (speculative_model)" + try: + with open(Path(path) / "config.json") as f: + cfg = json.load(f) + except (OSError, ValueError) as err: + return f"the DSpark drafter's config.json is not readable ({err})" + text = cfg.get("text_config") or {} + + def field(name): + return text.get(name, cfg.get(name)) + + arches = cfg.get("architectures") or text.get("architectures") or [] + if not arches or arches[0] != _M31_DSPARK_ARCH: + return f"drafter architecture {arches[:1]}, not {_M31_DSPARK_ARCH}" + if list(field("dspark_target_layer_ids") or []) != _M31_DSPARK_TARGET_LAYERS: + return f"drafter target layers {field('dspark_target_layer_ids')}, not {_M31_DSPARK_TARGET_LAYERS}" + for name, want in _M31_DSPARK_GEOMETRY.items(): + if field(name) != want: + return f"drafter {name} {field(name)}, not {want}" + return None + + +def route( + pretrained_config: Any, + mapping: Any, + compute_capability: Tuple[int, int], + spec_config: Optional[Any], + kv_cache_quant_algo: Optional[Any], +) -> type: + """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. + + ``kv_cache_quant_algo`` is the KV-cache algorithm the engine will run with + (see ``_router_index._kv_cache_quant_algo``), not the checkpoint's alone. + + Each target is imported where it is returned, so declining costs no import + of modeling code. + """ + 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) + kv_cache = kv_cache_quant_algo in (QuantAlgo.FP8, QuantAlgo.NVFP4) + m31_miss = _m31_mismatch(pretrained_config, text_config, quant_algo) + m31_spec_miss = _m31_speculation_mismatch(spec_config) if m31_miss is None else None + + if ( + compute_capability == _SM + and shape == _SHAPE + and m31_miss is None + and parallel in ("tp4", "tp8") + and m31_spec_miss is None + and kv_cache + ): + if parallel == "tp4": + from .minimax_m31_nvfp4__sm_103__tp4.modeling import ModelingV2MiniMaxM31Nvfp4Sm103Tp4 + + return ModelingV2MiniMaxM31Nvfp4Sm103Tp4 + from .fake_ops import register + + register() + from .minimax_m31_nvfp4__sm_103__tp8.modeling import ModelingV2MiniMaxM31Nvfp4Sm103Tp8 + + return ModelingV2MiniMaxM31Nvfp4Sm103Tp8 + + raise NoTarget( + f"MiniMax-M3.1 requires sm_103 TP4 or TP8/EP1, shape {_SHAPE}, MXFP8/NVFP4 experts, " + f"FP8 or NVFP4 KV, and standalone MiniMax DSpark or no speculation; " + f"got SM={compute_capability}, shape={shape}, topology={parallel}, " + f"KV={kv_cache_quant_algo}, checkpoint={m31_miss}, speculation={m31_spec_miss}" + ) 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 4563331e9754..93af5a432f8c 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 @@ -60,6 +60,37 @@ from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig +class _UnchangedCopy: + """Remembers the inputs of the last decode-only block-table copy into one destination. + + In steady decode the block tables change only when a request grows into a + new page, so a copy whose destination, batch order, decode-only shape and + host block-offset table equal the previous copy's can be skipped. + """ + + def __init__(self) -> None: + self._key: Optional[tuple] = None + self._host: Optional[torch.Tensor] = None + + def skip(self, dst, request_ids, beam_width, num_contexts, num_seqs, max_blocks, host) -> bool: + if num_contexts != 0: + self._key = None + return False + key = (dst.data_ptr(), tuple(request_ids), beam_width, num_seqs, max_blocks) + return key == self._key and self._host is not None and torch.equal(host, self._host) + + def record( + self, dst, request_ids, beam_width, num_contexts, num_seqs, max_blocks, host + ) -> None: + if num_contexts != 0: + return + self._key = (dst.data_ptr(), tuple(request_ids), beam_width, num_seqs, max_blocks) + if self._host is None or self._host.shape != host.shape: + self._host = host.clone() + else: + self._host.copy_(host) + + class MiniMaxM3SparseIndexCache: """Plain-tensor side cache for the M3 sparse index branch. @@ -304,6 +335,10 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): own pool. """ + # 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.NHD # Virtual attention-op pools for shared draft layers: # (attention_op_pool_id, source_storage_pool_id) plus their copy parameters. @@ -370,6 +405,7 @@ def __init__( f"{self.indexer_kv_dtype!r}." ) super().__init__(*args, **kwargs) + self._unchanged_copies: dict[int, _UnchangedCopy] = {} if self.dtype == DataType.NVFP4: dense_target_layers = [ @@ -618,6 +654,20 @@ def copy_batch_block_offsets( num_seqs: int, max_blocks: Optional[int] = None, ) -> None: + if self.skip_unchanged_block_offsets: + guard = self._unchanged_copies.setdefault(dst_tensor.data_ptr(), _UnchangedCopy()) + args = ( + dst_tensor, + request_ids, + beam_width, + num_contexts, + num_seqs, + max_blocks, + self.host_kv_cache_block_offsets, + ) + if guard.skip(*args): + return + guard.record(*args) super().copy_batch_block_offsets( dst_tensor, request_ids, beam_width, num_contexts, num_seqs, max_blocks=max_blocks ) @@ -827,6 +877,20 @@ def _build_kv_buffers( full_view = convert_to_torch_tensor(TensorWrapper(addr_key, torch_dtype, full_slot_shape)) return full_view[:, :2] + 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 + # The feature base already caches every pool view in _cached_pool_view. + self.reuse_pool_views = reuse_pool_views + def _kv_scale_slot_geometry( self, layer_idx: int, kv_layout: Optional[str] ) -> Tuple[int, torch.dtype, int, int, List[int]]: @@ -1063,6 +1127,7 @@ def __init__(self, manager: MiniMaxM3KVCacheManagerV2, draft_layer_ids: Sequence f"got {len(draft_layer_ids)}. Multiple layers require separate pool roots." ) self._manager = manager + self._unchanged_copy = _UnchangedCopy() layer_id = draft_layer_ids[0] local = int(manager.layer_offsets[layer_id]) self._local_layer_idx = local @@ -1127,6 +1192,19 @@ def copy_batch_block_offsets( """ assert beam_width == 1, "beam_width must be 1 for KVCacheManagerV2" manager = self._manager + if manager.skip_unchanged_block_offsets: + args = ( + dst_tensor, + request_ids, + beam_width, + num_contexts, + num_seqs, + max_blocks, + self.host_kv_cache_block_offsets, + ) + if self._unchanged_copy.skip(*args): + return + self._unchanged_copy.record(*args) copy_idx = manager.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) assert copy_idx.shape[0] == num_seqs copy_batch_block_offsets_to_device( 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 b5022eb77256..fc9ec623cc0e 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 @@ -39,6 +39,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 @@ -103,6 +106,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. @@ -159,6 +173,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 @@ -202,6 +223,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 @@ -644,6 +671,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() @@ -841,6 +902,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]) @@ -901,6 +985,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. @@ -931,7 +1017,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() @@ -962,13 +1050,28 @@ 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. - self._msa_prefill_proxy_plan = plan_for( - self._msa_proxy_plan_rows(), - num_qo_heads=num_index_heads, - num_kv_heads=1, - output_maxscore=True, + # A short pure-prefill step selects every valid block, so it needs none. + self._msa_prefill_proxy_plan = ( + None + if short_prefill + else plan_for( + self._msa_proxy_plan_rows(), + num_qo_heads=num_index_heads, + num_kv_heads=1, + output_maxscore=True, + ) ) attn_rows = self._msa_attn_plan_rows() # Sparse layers: kv_block_num=topk limits attention to top-k blocks. @@ -988,14 +1091,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. @@ -1013,11 +1117,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 @@ -1034,6 +1138,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 @@ -1075,41 +1181,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()) @@ -1144,7 +1271,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, @@ -1168,16 +1295,20 @@ def _build_msa_fields(self) -> None: intra_np = ( np.arange(total_new_tokens, dtype=np.int64) - np.repeat(starts_np, qo_np) ).astype(np.int32) - self.msa_q_batch_row[:total_new_tokens].copy_( - maybe_pin_memory(torch.from_numpy(batch_row_np)), non_blocking=True + copy_h2d_if_changed( + self.msa_q_batch_row[:total_new_tokens], + maybe_pin_memory(torch.from_numpy(batch_row_np)), + pack=True, ) - self.msa_q_intra[:total_new_tokens].copy_( - maybe_pin_memory(torch.from_numpy(intra_np)), non_blocking=True + copy_h2d_if_changed( + self.msa_q_intra[:total_new_tokens], + maybe_pin_memory(torch.from_numpy(intra_np)), + pack=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: @@ -1470,6 +1601,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 d29def811bc5..8a262843975e 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -39,7 +39,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 .fmha.manager import FmhaManager @@ -997,8 +999,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( @@ -1051,9 +1053,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. @@ -1160,7 +1161,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/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 01b9f9cace7c..7ac21c784c47 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -849,6 +849,65 @@ def mxfp8_quantize_gemm_autotuned( gemm_inputs, 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, + tactic=_choose_mxfp8_tactic( + _FLASHINFER_MXFP8_GEMM_AUTOTUNED_OP, gemm_runner, + 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/checkpoints/hf/minimaxm3_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py index 261559f186d4..7ce7a546f0d7 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py @@ -60,6 +60,11 @@ def _read_global_scale(name: str, value) -> Tensor: return scale +def reciprocal_global_scale(name: str, value: Tensor) -> Tensor: + """Validate a packed expert scale and return its FP32 modelopt divisor.""" + return nvfp4_weight_scale_2(_read_global_scale(name, value)) + + def translate_packed_nvfp4_experts(weights: Mapping[str, Tensor]) -> Mapping[str, Tensor]: """Rename packed NVFP4 routed-expert tensors to the NVFP4 MoE contract. diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index f670cdb7937f..d1d80a29efe2 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -509,6 +509,13 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "AUTO"): f"{self.dflash_attention_backend!r} but loads no op set for it." ) self._dflash_trtllm_gen_workspace = None + # Optional target-installed hooks. Unarmed draft models retain the stock path. + self._dflash_lean_glue = False + self._dflash_cache_rows = None + self._dflash_fused_indices = None + self._dflash_fused_qk_norm_rope = None + self._dflash_fused_qk_norm_rope_append = None + self._dflash_attn_pdl = False self._dflash_trtllm_gen_counters = None self.register_buffer("_dflash_batch_indices", None, persistent=False) self.register_buffer("_dflash_block_offsets", None, persistent=False) @@ -1047,7 +1054,8 @@ def precompute_context_kv( # after the select() splits — no extra copy required. kv = kv_flat.view(N, L, 2, nkv, hd) k = kv[:, :, 0].contiguous() - v = kv[:, :, 1].contiguous() + # V only feeds the paged append and an in-place mask multiply, both stride-aware. + v = kv[:, :, 1] if self._dflash_lean_glue else kv[:, :, 1].contiguous() if self._k_norm_stacked is not None: # Fuse L per-layer RMSNorms into one. k is [N, L, nkv, hd]; @@ -1535,10 +1543,26 @@ def dflash_forward( q_rope_cos, q_rope_sin = self._get_rope_cos_sin(query_positions, dtype=rope_dtype) _rope = RotaryEmbedding.apply_rotary_pos_emb - # cache_seqlens BEFORE the block's K/V are stored. - cache_seqlens_i32 = num_ctx_per_req[:B].to(torch.int32) - cache_batch_idx_i32 = ctx_cache_batch_idx.to(torch.int32) - seq_lens_after = cache_seqlens_i32 + block_size + # A target's one-launch index prelude for the paged decode: the same int32 tensors as the stock lines, or None. + fused_indices = ( + self._dflash_fused_indices( + num_ctx_per_req[:B], + ctx_cache_batch_idx, + query_positions, + block_size, + ctx_page_table.size(1), + ctx_kv_cache[0].size(-2), + ) + if is_paged and self._dflash_fused_indices is not None + else None + ) + if fused_indices is not None: + cache_seqlens_i32, cache_batch_idx_i32, seq_lens_after = fused_indices[:3] + else: + # cache_seqlens BEFORE the block's K/V are stored. + cache_seqlens_i32 = num_ctx_per_req[:B].to(torch.int32) + cache_batch_idx_i32 = ctx_cache_batch_idx.to(torch.int32) + seq_lens_after = cache_seqlens_i32 + block_size if is_paged: max_batch_size = ctx_page_table.size(0) @@ -1554,28 +1578,38 @@ def dflash_forward( ) else: # FA4 needs no workspace or counter buffers. self._prepare_dflash_index_buffers(hidden_states.device, max_batch_size, block_size) - block_tables = ctx_page_table.index_select(0, cache_batch_idx_i32.long()) + rows = self._dflash_cache_rows if self._dflash_lean_glue else None + if rows is not None: + block_tables = ctx_page_table[rows[0] : rows[1]] + else: + block_tables = ctx_page_table.index_select(0, cache_batch_idx_i32.long()) pages_per_slot = block_tables.size(1) # Index the layer first: ctx_kv_cache is an [L, ...] tensor for the # private arena but a per-layer list when bound to the draft KV # cache manager's pool. page_size sits at -2 either way. page_size = ctx_kv_cache[0].size(-2) kv_indices = block_tables.flatten() - kv_indptr = torch.arange( - 0, - (B + 1) * pages_per_slot, - pages_per_slot, - dtype=torch.int32, - device=hidden_states.device, - ) - kv_last_page_len = ((seq_lens_after - 1) % page_size) + 1 + if fused_indices is not None: + kv_indptr, kv_last_page_len = fused_indices[3:5] + else: + kv_indptr = torch.arange( + 0, + (B + 1) * pages_per_slot, + pages_per_slot, + dtype=torch.int32, + device=hidden_states.device, + ) + kv_last_page_len = ((seq_lens_after - 1) % page_size) + 1 batch_indices = self._dflash_batch_indices append_batch_indices = batch_indices[:B].reshape(-1) - append_positions = ( - (cache_seqlens_i32.view(-1, 1) + self._dflash_block_offsets) - .reshape(-1) - .contiguous() - ) + if fused_indices is not None: + append_positions = fused_indices[5] + else: + append_positions = ( + (cache_seqlens_i32.view(-1, 1) + self._dflash_block_offsets) + .reshape(-1) + .contiguous() + ) else: # VANILLA # Slot and column that each block K/V row occupies self._prepare_dflash_index_buffers( @@ -1591,8 +1625,11 @@ def dflash_forward( (cache_seqlens_i32.view(-1, 1) + self._dflash_block_offsets).reshape(-1).long() ) - # Flatten query positions once for the fused QK-norm-RoPE kernel. - query_positions_flat_i32 = query_positions.reshape(-1).to(torch.int32) + if fused_indices is not None: + query_positions_flat_i32 = fused_indices[6] + else: + # Flatten query positions once for the fused QK-norm-RoPE kernel. + query_positions_flat_i32 = query_positions.reshape(-1).to(torch.int32) residual = None has_block_conv = self.has_block_conv @@ -1622,6 +1659,9 @@ def dflash_forward( # QKV projection on normed query tokens (2D) qkv_query = attn_mod.qkv_proj(hs_normed_flat) # [B*blk, qkv_size] + kv_appended = ( + False # the fused q / k norm + RoPE also wrote this layer's K / V pages (see below) + ) if use_fused_qk_norm_rope: # One kernel does q_norm + k_norm + RoPE in-place on qkv. # Only safe when the drafter's rope params don't use YaRN / @@ -1635,21 +1675,66 @@ def dflash_forward( k_noise_bshd = k_noise_2d.reshape(B, block_size, num_kv_heads_per_rank, head_dim) v_noise_bshd = v_noise_2d.reshape(B, block_size, num_kv_heads_per_rank, head_dim) elif use_fused_rope: - # Per-head RMSNorm on q/k (returns new contiguous tensors), - # then flashinfer in-place RoPE sharing the same cos/sin cache - # as precompute_context_kv. - q = attn_mod.q_norm(qkv_query[:, :q_size].reshape(-1, head_dim)).view(-1, q_size) - k = attn_mod.k_norm( - qkv_query[:, q_size : q_size + kv_size].reshape(-1, head_dim) - ).view(-1, kv_size) - _flashinfer_rope( - query_positions_flat_i32, - q, - k, - head_dim, - self._get_cos_sin_cache(), - self._is_neox, + # A target's one-launch q / k norm + RoPE: the same new q / k as the lines below, or None. With the + # append form, the same launch also writes the block's K / V into this layer's paged cache. + fused_qk = ( + self._dflash_fused_qk_norm_rope_append( + qkv_query, + q_size, + kv_size, + head_dim, + attn_mod.q_norm, + attn_mod.k_norm, + query_positions_flat_i32, + self._get_cos_sin_cache(), + self._is_neox, + ctx_kv_cache[layer_idx], + append_batch_indices, + append_positions, + kv_indices, + kv_indptr, + ) + if self._dflash_fused_qk_norm_rope_append is not None + and self.dflash_attention_backend == "TRTLLM" + else None ) + kv_appended = fused_qk is not None + if fused_qk is None: + fused_qk = ( + self._dflash_fused_qk_norm_rope( + qkv_query, + q_size, + kv_size, + head_dim, + attn_mod.q_norm, + attn_mod.k_norm, + query_positions_flat_i32, + self._get_cos_sin_cache(), + self._is_neox, + ) + if self._dflash_fused_qk_norm_rope is not None + else None + ) + if fused_qk is not None: + q, k = fused_qk + else: + # Per-head RMSNorm on q/k (returns new contiguous tensors), + # then flashinfer in-place RoPE sharing the same cos/sin cache + # as precompute_context_kv. + q = attn_mod.q_norm(qkv_query[:, :q_size].reshape(-1, head_dim)).view( + -1, q_size + ) + k = attn_mod.k_norm( + qkv_query[:, q_size : q_size + kv_size].reshape(-1, head_dim) + ).view(-1, kv_size) + _flashinfer_rope( + query_positions_flat_i32, + q, + k, + head_dim, + self._get_cos_sin_cache(), + self._is_neox, + ) Q_bshd = q.view(B, block_size, num_heads_per_rank, head_dim) k_noise_bshd = k.view(B, block_size, num_kv_heads_per_rank, head_dim) v_noise_bshd = qkv_query[:, q_size + kv_size :].reshape( @@ -1694,21 +1779,26 @@ def dflash_forward( causal, window_size = self._resolve_block_attention(layer_idx) if self.dflash_attention_backend == "TRTLLM": layer_cache = ctx_kv_cache[layer_idx] - trtllm_gen_ops.append_paged_kv_cache( - append_key=k_noise_bshd.reshape( - -1, num_kv_heads_per_rank, head_dim - ).contiguous(), - append_value=v_noise_bshd.reshape( - -1, num_kv_heads_per_rank, head_dim - ).contiguous(), - batch_indices=append_batch_indices, - positions=append_positions, - paged_kv_cache=layer_cache, - kv_indices=kv_indices, - kv_indptr=kv_indptr, - kv_last_page_len=kv_last_page_len, - kv_layout="HND", - ) + if not kv_appended: + trtllm_gen_ops.append_paged_kv_cache( + append_key=k_noise_bshd.reshape( + -1, num_kv_heads_per_rank, head_dim + ).contiguous(), + append_value=( + v_noise_bshd.reshape(-1, num_kv_heads_per_rank, head_dim) + if self._dflash_lean_glue + else v_noise_bshd.reshape( + -1, num_kv_heads_per_rank, head_dim + ).contiguous() + ), + batch_indices=append_batch_indices, + positions=append_positions, + paged_kv_cache=layer_cache, + kv_indices=kv_indices, + kv_indptr=kv_indptr, + kv_last_page_len=kv_last_page_len, + kv_layout="HND", + ) out = torch.empty_like(Q_bshd) q_flat = Q_bshd.reshape(-1, num_heads_per_rank, head_dim) out_flat = out.reshape(-1, num_heads_per_rank, head_dim) @@ -1726,7 +1816,7 @@ def dflash_forward( window_left=window_left, out=out_flat, sinks=None, - enable_pdl=False, + enable_pdl=self._dflash_attn_pdl, kv_layout="HND", backend="trtllm-gen", q_len_per_req=block_size, @@ -1767,7 +1857,7 @@ def dflash_forward( window_left=window_left, out=out_flat, sinks=None, - enable_pdl=False, + enable_pdl=self._dflash_attn_pdl, kv_layout="HND", kv_cache_sf=None, uses_shared_paged_kv_idx=True, @@ -1780,9 +1870,11 @@ def dflash_forward( append_key=k_noise_bshd.reshape( -1, num_kv_heads_per_rank, head_dim ).contiguous(), - append_value=v_noise_bshd.reshape( - -1, num_kv_heads_per_rank, head_dim - ).contiguous(), + append_value=( + v_noise_bshd.reshape(-1, num_kv_heads_per_rank, head_dim) + if self._dflash_lean_glue + else v_noise_bshd.reshape(-1, num_kv_heads_per_rank, head_dim).contiguous() + ), batch_indices=append_batch_indices, positions=append_positions, paged_kv_cache=layer_cache, diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 0b229654a59c..ad794d6b494d 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -3366,6 +3366,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 b177b8c50894..cf91917ec042 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -53,6 +53,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 ..utils import is_gdn_replay_enabled from . import config_utils from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer, @@ -77,6 +78,7 @@ from .resource_manager import (KVCacheCompressionManager, KVCacheManager, PeftCacheManager, ResourceManager, ResourceManagerType) +from .runtime_hooks import runtime_hooks_of from .sampler import EarlyStopSampler, EarlyStopWithMMResult, TorchSampler from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler, KVCacheV2Scheduler, MultimodalEagerEncoderScheduler, @@ -4626,7 +4628,18 @@ 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) + runtime_hooks = runtime_hooks_of(engine.model) + if isinstance(sampler, SpecSampler): + if runtime_hooks.spec_step_fold is not None: + sampler.bind_spec_step_fold(runtime_hooks.spec_step_fold) + sampler.d2h_side_stream_copies = runtime_hooks.sampler_d2h_side_stream + sampler.static_slot_ids = runtime_hooks.spec_small_h2d + if runtime_hooks.sampler_d2h_side_stream: + logger.info( + "[M3 host] spec sampler store copies on the D2H side stream" + ) + 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 19ed990f33a1..3aea5f022a99 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -27,6 +27,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, ResourceManager, ResourceManagerType) @@ -149,6 +150,15 @@ class CUDAGraphRunnerConfig: kv_cache_manager_key: Any 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 @@ -162,6 +172,12 @@ 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 + # The graphs captured with a step-input pack's in-graph prologue, and that + # pack (step_inputs.capture_prologue). None until the first such capture: a + # class-level default, so a runner without one replays exactly as before. + _step_input_prologue_keys: Optional[Dict[KeyType, Any]] = None def __init__(self, config: CUDAGraphRunnerConfig): self.config = config @@ -174,6 +190,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) @@ -203,7 +220,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 @@ -301,6 +320,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, @@ -313,8 +346,7 @@ def _get_seq_len_mode( that the model engine 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 @@ -356,7 +388,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 @@ -771,6 +803,10 @@ def capture(self, # setup/capture; release its reference before entering. output = None with torch.cuda.graph(graph, pool=self.memory_pool): + # A pre-replay hook reads the static inputs before the replay, + # so their copies must land before it: no prologue then. + if not self._pre_replay_hooks: + self._capture_step_input_prologue(key) output = forward_fn(capture_inputs) # Snapshot before postprocessing so the check compares against # the addresses the graph actually baked in. @@ -786,6 +822,35 @@ def capture(self, self.memory_pool = graph.pool() return graph_output + def _capture_step_input_prologue(self, key: KeyType) -> None: + """Inside the capture of ``key``: the active step-input pack's in-graph + prologue, if it has one (step_inputs.capture_prologue).""" + prologue_pack = step_inputs.capture_prologue() + if prologue_pack is not None: + if self._step_input_prologue_keys is None: + self._step_input_prologue_keys = {} + self._step_input_prologue_keys[key] = prologue_pack + elif self._step_input_prologue_keys: + # A recapture of the key without the prologue. + self._step_input_prologue_keys.pop(key, None) + + def _issue_step_inputs(self, key: KeyType) -> None: + """Issue the step's packed input copies before the replay of ``key`` + (a no-op without an open pack). A graph captured with a pack's + prologue issues that pack itself at the start of the replay.""" + prologue_keys = self._step_input_prologue_keys + prologue_pack = prologue_keys.get(key) if prologue_keys else None + if prologue_pack is None: + step_inputs.flush() + elif step_inputs.is_active(prologue_pack): + step_inputs.flush(in_graph=True) + else: + # Replayed outside the pack's context: another (or no) pack is + # active, which the prologue does not read, so it is issued + # eagerly, and nothing filled the prologue pack's window. + step_inputs.flush() + prologue_pack.quiesce_graph_window() + def replay( self, key: KeyType, @@ -816,7 +881,7 @@ def replay( f"replay() got {seqlen} tokens for key {key}, but the graph " f"was captured for {expected_num_tokens} tokens. A shorter " "input_ids leaves the tail of the static input buffer stale.") - static_tensors["input_ids"][:seqlen].copy_(input_ids) + step_inputs.copy_(static_tensors["input_ids"][:seqlen], input_ids) position_ids = current_inputs["position_ids"] if self.config.use_mrope: @@ -863,7 +928,8 @@ def replay( f"for key {key}, but expected {expected_position_ids_shape}. " "torch.Tensor.copy_() silently broadcasts mismatched shapes, " "which would corrupt the static input buffer.") - static_tensors["position_ids"][:, :seqlen].copy_(position_ids) + step_inputs.copy_(static_tensors["position_ids"][:, :seqlen], + position_ids) num_encoder_tokens = key.num_encoder_tokens if num_encoder_tokens: @@ -887,6 +953,10 @@ 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; a graph captured + # with the pack's prologue issues them itself at the start of the replay. + self._issue_step_inputs(key) if self._pre_replay_hooks: hook_tensors = static_tensors if position_id_offsets is not None: @@ -1338,6 +1408,7 @@ def clear(self): self.graph_outputs.clear() self.graph_metadata.clear() self._pre_replay_position_ids = None + self._step_input_prologue_keys = None self.padding_dummy_requests = {} del self.memory_pool self.memory_pool = None diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 3b8d7fa2da91..b281c15d0f27 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -34,6 +34,7 @@ ) 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, @@ -42,6 +43,7 @@ get_mrope_dummy_seq_slot, ) from tensorrt_llm._torch.pyexecutor.guided_decoder import CapturableGuidedDecoder +from tensorrt_llm._torch.pyexecutor.h2d_cache import H2DSkipCache, copy_h2d_if_changed from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( BaseMambaCacheManager, @@ -54,11 +56,23 @@ ResourceManager, ResourceManagerType, ) +from tensorrt_llm._torch.pyexecutor.runtime_hooks import ( + DEFAULT_RUNTIME_HOOKS, + MsaRuntimePolicy, + RuntimeHooks, + runtime_hooks_of, +) from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import warmup_sampling_module +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.scheduler.scheduler import RemoteTailPhase +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack from tensorrt_llm._torch.pyexecutor.trace_log_utils import log_mem_snapshot from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer from tensorrt_llm._torch.pyexecutor.workspace import EagerWorkspaceReclaimer @@ -91,6 +105,7 @@ from tensorrt_llm.llmapi.llm_args import ( BaseSparseAttentionConfig, DecodingBaseConfig, + MiniMaxM3SparseAttentionConfig, PrefillCudaGraphBackend, SeqLenAwareSparseAttentionConfig, ) @@ -296,6 +311,15 @@ def __getattr__(self, name: str) -> Any: class DecoderRunner(ScheduledModelRunner): """Run scheduled decoder generation with its own buffers, graphs and metadata.""" + # 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 _context_graph_promotion_supported = True @@ -347,6 +371,32 @@ def __init__( self._mm_item_scheduler = mm_item_scheduler self._config = config + self.runtime_hooks = runtime_hooks_of(model) + self._h2d_cache = ( + H2DSkipCache(changed_span=self.runtime_hooks.h2d_changed_span) + if self.runtime_hooks.skip_unchanged_h2d + else None + ) + self._step_pack = ( + StepInputPack( + torch.device("cuda", torch.cuda.current_device()), + parallel_scatter=self.runtime_hooks.pack_parallel_scatter, + in_graph=self.runtime_hooks.pack_in_graph, + early_copy=self.runtime_hooks.pack_early_copy, + ) + if self.runtime_hooks.pack_step_inputs + else None + ) + if self._step_pack is not None and self._step_pack.parallel_scatter: + logger.info("[M3 host] step-input pack: parallel scatter for non-overlapping copies") + if self._step_pack is not None and self._step_pack.in_graph: + logger.info("[M3 host] step-input pack issued by the decode graph's prologue") + if self._step_pack is not None and self._step_pack.early_copy: + logger.info( + "[M3 host] step-input pack: window uploaded ahead of the replay on a side copy stream" + ) + if self._h2d_cache is not None and self._h2d_cache.changed_span: + logger.info("[M3 host] changed metadata uploads stage only each row's changed span") self.guided_decoder: CapturableGuidedDecoder | None = None self.lora_model_config: LoraModelConfig | None = None @@ -391,7 +441,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 warmup(self, resource_manager: ResourceManager) -> None: kv_cache_manager = resource_manager.get_resource_manager(self.kv_cache_manager_key) @@ -472,6 +530,14 @@ def _allocate_decoder_buffers(self) -> None: self.position_ids_cuda = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, device="cuda" ) + 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.input_ids_cuda, + self.draft_tokens_cuda, + self.previous_pos_id_offsets_cuda, + self.previous_kv_lens_offsets_cuda, + self._config.max_batch_size, + ) # Host counter of the steady-state generation prepare cache. self._steady_gen_positions_pinned = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, pin_memory=prefer_pinned() @@ -550,9 +616,33 @@ def _cuda_graph_runner_config(self) -> CUDAGraphRunnerConfig: dist=self.dist, kv_cache_manager_key=self.kv_cache_manager_key, 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, ) + 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.cuda_graph_runner is None @@ -1643,6 +1733,11 @@ def run_autotuner_pass(autotune_context: Any, synchronize_trtllm_cache: bool) -> ran_flashinfer_forward = run_autotuner_pass( flashinfer_mxfp8_autotune(), synchronize_trtllm_cache=False ) + if self.runtime_hooks.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: @@ -2037,6 +2132,10 @@ def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): ) with flashinfer_autotune_context, flashinfer_mxfp8_decode_graph_capture(): self._capture_generation_cuda_graphs(resource_manager) + if is_warmup_only and self.runtime_hooks.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. @@ -2081,11 +2180,8 @@ def _capture_generation_cuda_graphs(self, resource_manager: ResourceManager): if self.mapping is not None and 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._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. @@ -2093,7 +2189,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: @@ -2776,15 +2872,84 @@ def _set_up_attn_metadata( 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) + self._arm_draft_block_offset_guard(draft_kv_cache_manager) return self.attn_metadata - def _set_up_spec_metadata(self, spec_resource_manager: BaseResourceManager | None): + def _arm_draft_block_offset_guard( + self, draft_kv_cache_manager: KVCacheManager | KVCacheManagerV2 | None + ) -> None: + """Turn on the separate draft KV cache's unchanged block-table skip the model opted into.""" + if not self.runtime_hooks.skip_unchanged_draft_block_offsets: + return + if not isinstance(draft_kv_cache_manager, KVCacheManagerV2): + return + if not draft_kv_cache_manager.skip_unchanged_block_offsets: + draft_kv_cache_manager.skip_unchanged_block_offsets = True + logger.info("[M3 host] draft block-offset copy skipped when unchanged") + + def _arm_msa_policy(self, kv_cache_manager: Any) -> 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 _set_up_spec_metadata( + self, spec_resource_manager: BaseResourceManager | None + ) -> SpecMetadata | None: if self.spec_metadata is not None: return self.spec_metadata self.spec_metadata = create_spec_metadata( self._config, self.model.config, spec_resource_manager ) + if self.spec_metadata is None: + return None + self.spec_metadata.skip_greedy_rng_upload = self.runtime_hooks.skip_greedy_rng_upload + self.spec_metadata.pack_rng_uploads = self.runtime_hooks.pack_spec_rng_uploads + self.spec_metadata.batch_map_changed_only = self.runtime_hooks.spec_batch_map_changed_only + self.spec_metadata.tail_bookkeep = self.runtime_hooks.spec_tail_bookkeep + self.spec_metadata.slot_ids_changed_only = self.runtime_hooks.spec_slot_ids_changed_only + if self.runtime_hooks.pack_spec_rng_uploads: + logger.info("[M3 host] sampled-step Philox uploads staged into the step-input pack") + if self.runtime_hooks.spec_small_h2d: + logger.info( + "[M3 host] spec step small uploads off the boundary (accepted counts packed, static slot ids)" + ) + if self.runtime_hooks.skip_accepted_draft_upload: + logger.info("[M3 host] spec step accepted-draft-count upload skipped (no reader)") + if self.runtime_hooks.spec_batch_map_changed_only: + logger.info( + "[M3 host] drafter batch_indices / batch_to_slot uploaded only when changed" + ) + if self.runtime_hooks.spec_slot_ids_changed_only: + logger.info( + "[M3 host] spec metadata slot ids (batch_slot_ids) uploaded only when changed" + ) + if self.runtime_hooks.kv_lens_update_once: + logger.info( + "[M3 host] spec step KV-length metadata re-derived once, after the KV-length correction" + ) return self.spec_metadata def _release_decoder_graphs(self) -> None: @@ -2804,6 +2969,20 @@ def _should_use_full_generation_page_table( self._config.disable_overlap_scheduler, spec_config, attn_metadata ) + def _kv_lens_corrected_in_preprocess( + self, inputs: dict[str, Any], enable_spec_decode: bool + ) -> bool: + """Whether preprocess_inputs applies the overlap scheduler's KV-length offsets to kv_lens_cuda and then calls + on_update_kv_lens (the branch below).""" + attn_meta = inputs.get("attn_metadata") + return ( + enable_spec_decode + and not self._config.disable_overlap_scheduler + and attn_meta is not None + and attn_meta.kv_cache_manager is not None + and hasattr(attn_meta, "kv_lens_cuda") + ) + def _preprocess_inputs( self, inputs: dict[str, Any], *, enable_spec_decode: bool, runtime_draft_len: int ): @@ -2811,8 +2990,14 @@ def _preprocess_inputs( Make some changes to the device inputs and avoid blocking the async data transfer """ attn_meta = inputs.get("attn_metadata") - # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step. - if attn_meta is not None: + # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step. With + # RuntimeHooks.kv_lens_update_once, a step that corrects the KV lengths below (and then calls it again) calls + # it only there: the second call rewrites everything the first one writes, and nothing reads the metadata + # in between. + if attn_meta is not None and not ( + self.runtime_hooks.kv_lens_update_once + and self._kv_lens_corrected_in_preprocess(inputs, enable_spec_decode) + ): attn_meta.on_update_kv_lens() if enable_spec_decode and not self._config.disable_overlap_scheduler: @@ -4002,7 +4187,7 @@ def previous_seq_slots_device(): previous_batch_indices, dtype=torch.int, pin_memory=prefer_pinned() ) previous_slots = self.previous_batch_indices_cuda[:previous_batch_len] - previous_slots.copy_(previous_batch_indices_host, non_blocking=True) + copy_h2d_if_changed(previous_slots, previous_batch_indices_host) return previous_slots if scheduled_requests.is_attention_dp_phase_idle: @@ -4023,21 +4208,77 @@ def previous_seq_slots_device(): if num_draft_tokens > 0: draft_tokens = torch.tensor(draft_tokens, dtype=torch.int, pin_memory=prefer_pinned()) self.draft_tokens_cuda[: len(draft_tokens)].copy_(draft_tokens, non_blocking=True) - if self._config.is_spec_decode and len(num_accepted_draft_tokens) > 0: + if ( + self._config.is_spec_decode + and len(num_accepted_draft_tokens) > 0 + and not self.runtime_hooks.skip_accepted_draft_upload + ): num_accepted_draft_tokens = torch.tensor( num_accepted_draft_tokens, dtype=torch.int, pin_memory=prefer_pinned() ) - self.num_accepted_draft_tokens_cuda[: len(num_accepted_draft_tokens)].copy_( - num_accepted_draft_tokens, non_blocking=True + copy_h2d_if_changed( + self.num_accepted_draft_tokens_cuda[: len(num_accepted_draft_tokens)], + num_accepted_draft_tokens, + pack=self.runtime_hooks.spec_small_h2d, + ) + step_fold = self.runtime_hooks.spec_step_fold if self._config.is_spec_decode 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 num_ctx_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: - # Initialize these two values to zeros - self.previous_pos_id_offsets_cuda *= 0 - self.previous_kv_lens_offsets_cuda *= 0 runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step(runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 + # 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_host = torch.tensor( + previous_pos_indices, dtype=torch.int, pin_memory=prefer_pinned() + ) + copy_h2d_if_changed( + self.previous_pos_indices_cuda[ + : previous_batch_len * runtime_tokens_per_gen_step + ], + previous_pos_indices_host, + ) + num_extend_reqeust_wo_dummy = len(extend_requests) - len(extend_dummy_requests) + inputs_staged = step_fold.gather( + new_tokens_device, + next_draft_tokens_device, + new_tokens_lens_device, + previous_slots, + self.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, + ) + if not inputs_staged: + # Initialize these two values to zeros + self.previous_pos_id_offsets_cuda *= 0 + self.previous_kv_lens_offsets_cuda *= 0 - if previous_batch_len > 0: + if previous_batch_len > 0 and not inputs_staged: previous_slots = previous_seq_slots_device() # previous input ids previous_batch_tokens = previous_batch_len * runtime_tokens_per_gen_step @@ -4064,8 +4305,9 @@ def previous_seq_slots_device(): previous_pos_indices_host = torch.tensor( previous_pos_indices, dtype=torch.int, pin_memory=prefer_pinned() ) - self.previous_pos_indices_cuda[0:previous_batch_tokens].copy_( - previous_pos_indices_host, non_blocking=True + copy_h2d_if_changed( + self.previous_pos_indices_cuda[0:previous_batch_tokens], + previous_pos_indices_host, ) # The order of requests in a batch: [context requests, generation requests] @@ -4164,13 +4406,13 @@ def previous_seq_slots_device(): ) final_position_ids = self.mrope_position_ids_cuda[:, :, :total_num_tokens] else: - self.position_ids_cuda[:total_num_tokens].copy_(host_position_ids, non_blocking=True) + step_inputs.copy_(self.position_ids_cuda[:total_num_tokens], host_position_ids) final_position_ids = self.position_ids_cuda[:total_num_tokens].unsqueeze(0) if enable_spec_decode: - self.gather_ids_cuda[: len(gather_ids)].copy_( + copy_h2d_if_changed( + self.gather_ids_cuda[: len(gather_ids)], torch.tensor(gather_ids, dtype=torch.int, pin_memory=prefer_pinned()), - non_blocking=True, ) if self.mapping.has_cp_helix(): @@ -4413,9 +4655,11 @@ def previous_seq_slots_device(): # num_generations / num_tokens / seq_lens are set above, before the # attention-DP allgather that must agree with prepare(). spec_metadata.host_position_ids = host_position_ids - spec_metadata.num_accepted_draft_tokens = self.num_accepted_draft_tokens_cuda[ - : len(num_accepted_draft_tokens) - ] + spec_metadata.num_accepted_draft_tokens = ( + None + if self.runtime_hooks.skip_accepted_draft_upload + else self.num_accepted_draft_tokens_cuda[: len(num_accepted_draft_tokens)] + ) if context_prompt_lookahead is not None: spec_metadata.populate_context_prompt_lookahead(context_prompt_lookahead) # No-op for non 1-model @@ -4790,6 +5034,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, @@ -4921,9 +5175,14 @@ def capture_postprocess_fn(inputs: dict[str, Any]): finally: restore_attn_metadata_after_draft_replay(attn_metadata, saved_draft) + # No-op unless a pack was opened above and the replay did not issue it. + step_inputs.flush() + 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) if not isinstance(outputs, dict): @@ -5044,6 +5303,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.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 _execute_logit_post_processors(self, scheduled_requests: ScheduledRequests, outputs: dict): """Apply logit post processors (in-place modify outputs Tensors) if any.""" @@ -5055,6 +5347,10 @@ def _execute_logit_post_processors(self, scheduled_requests: ScheduledRequests, 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 logits_row_offset = 0 request_groups = ( diff --git a/tensorrt_llm/_torch/pyexecutor/h2d_cache.py b/tensorrt_llm/_torch/pyexecutor/h2d_cache.py new file mode 100644 index 000000000000..cd7edb316bef --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/h2d_cache.py @@ -0,0 +1,204 @@ +# 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.logger import logger + +from . import step_inputs + + +class H2DSkipCache: + """One engine's host shadows of its skip-routed destinations (see the module docstring). + + ``changed_span`` (``RuntimeHooks.h2d_changed_span``): a packed copy into a 2-D view with contiguous rows whose + shadow covers the same rows (same first element and row stride) uploads, per row, only the span from its first to + its last changed column, plus any columns the shadow does not cover. A page table that gained one page then stages + a few bytes instead of every padded row, so the pack still fits the decode graph's prologue window. The span is + sound under the same invariant as the skip itself: the shadow describes the device bytes outside the span. + """ + + # A changed-span upload is used only while it moves at most this fraction of the view's bytes. + CHANGED_SPAN_MAX_FRACTION = 0.5 + # Class-level defaults: a cache built without __init__ never stages a span. + changed_span = False + # Engagement counters: changed-span uploads, the bytes they staged and the bytes the whole views would have staged. + span_uploads = 0 + span_bytes = 0 + span_full_bytes = 0 + + def __init__(self, changed_span: bool = False) -> None: + self.changed_span = changed_span + # 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, torch.Tensor]] = {} + # 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].dtype == src.dtype + and entry[2].shape == src.shape + and torch.equal(entry[2], src) + ): + return False + if not ( + pack + and self.changed_span + and entry is not None + and self._copy_changed_span(dst, src, entry[1:]) + ): + if pack: + step_inputs.copy_(dst, src) + else: + dst.copy_(src, non_blocking=True) + self._shadow[addr] = (storage, key, src.clone()) + return True + + def _copy_changed_span(self, dst: torch.Tensor, src: torch.Tensor, shadow: tuple) -> bool: + """Stage only the changed span of each row of ``dst`` (see the class docstring); False if not applicable.""" + (offset, shape, stride, dtype), old = shadow + if ( + dst.dim() != 2 + or dst.stride(1) != 1 + or offset != dst.storage_offset() + or len(shape) != 2 + or shape[0] != dst.shape[0] + or stride != tuple(dst.stride()) + or dtype != dst.dtype + or old.dtype != src.dtype + or tuple(src.shape) != tuple(dst.shape) + or not step_inputs.is_open() + ): + return False + cols = dst.shape[1] + known = min(shape[1], cols) + changed = src[:, :known] != old[:, :known] + if cols > known: + # Columns the shadow does not cover hold unknown device values. + changed = torch.cat([changed, changed.new_ones(dst.shape[0], cols - known)], dim=1) + rows = changed.any(dim=1).nonzero().flatten().tolist() + spans = [] + for row in rows: + columns = changed[row].nonzero().flatten() + spans.append((row, int(columns[0]), int(columns[-1]) + 1)) + staged = sum(hi - lo for _, lo, hi in spans) + if staged > self.CHANGED_SPAN_MAX_FRACTION * dst.numel(): + return False + for row, lo, hi in spans: + step_inputs.copy_(dst[row, lo:hi], src[row, lo:hi]) + self.span_uploads += 1 + self.span_bytes += staged * dst.element_size() + self.span_full_bytes += dst.numel() * dst.element_size() + if self.span_uploads & (self.span_uploads - 1) == 0: + logger.info( + f"[M3 host] changed-span uploads: {self.span_uploads}, {self.span_bytes} bytes staged " + f"of {self.span_full_bytes} for the whole views" + ) + 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/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 8ee3b7dff9e8..080c05303dec 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -1284,6 +1284,14 @@ class KVCacheManagerV2(BaseResourceManager): # The lender sharing.attach_lender installs, one per manager for its life; None without one. # On the class so it exists without running __init__. _sharing = None + # Opt-in (RuntimeHooks.skip_unchanged_draft_block_offsets): copy_batch_block_offsets skips a decode-only copy + # whose destination, copied rows and the rows' host page indices equal the previous copy's into that + # destination (see _block_offsets_unchanged). On the class so it exists without running __init__. + skip_unchanged_block_offsets: bool = False + _block_offset_guard: Optional[dict] = None + _block_offset_host: Optional[np.ndarray] = None + block_offset_copies: int = 0 + block_offset_skips: int = 0 def __init__( self, @@ -6589,6 +6597,13 @@ def copy_batch_block_offsets( ) assert copy_idx.shape[0] == num_seqs + if ( + self.skip_unchanged_block_offsets + and not self._use_per_layer_page_tables + and self._block_offsets_unchanged(dst_tensor, copy_idx, num_contexts) + ): + return + if self._use_per_layer_page_tables: self._copy_batch_block_offsets_per_layer( dst_tensor, request_ids, copy_idx, num_contexts, num_seqs @@ -6604,6 +6619,59 @@ def copy_batch_block_offsets( self._stream.cuda_stream, ) + def _block_offsets_unchanged( + self, dst_tensor: torch.Tensor, copy_idx: torch.Tensor, num_contexts: int + ) -> bool: + """Whether the copy into ``dst_tensor`` would write exactly what the previous copy into it wrote. + + copy_batch_block_offsets_to_device writes, for each copied row i, dst[:, i, 0/1, :] from field 0 of host + row copy_idx[i] alone (the V entry is derived from the same page index) and the per-pool scales and offsets, + which are fixed at construction. So the copy can be skipped when the destination view (storage, offset, + shape, strides), the copied rows and those rows' field-0 page indices equal the last copy's into that + storage: the device bytes are then already the ones the copy would write. Every write into these + destinations goes through this method, and a copy into another view of the same storage replaces the + record. A step with context requests always copies (its rows are being filled). The compare is a + single-threaded numpy compare of the copied rows only (a torch compare of the whole table costs tens of us + on the host per step). + """ + guard = self._block_offset_guard + if guard is None: + guard = self._block_offset_guard = {} + storage = dst_tensor.untyped_storage().data_ptr() + if num_contexts != 0: + guard.pop(storage, None) + return self._count_block_offset_copy(skipped=False) + rows = copy_idx.tolist() + key = ( + dst_tensor.data_ptr(), + tuple(dst_tensor.shape), + tuple(dst_tensor.stride()), + tuple(rows), + ) + host = self._block_offset_host + if host is None: + host = self._block_offset_host = self.host_kv_cache_block_offsets.numpy() + current = host[:, rows, 0, :] + last = guard.get(storage) + if last is not None and last[0] == key and np.array_equal(last[1], current): + return self._count_block_offset_copy(skipped=True) + guard[storage] = (key, current) + return self._count_block_offset_copy(skipped=False) + + def _count_block_offset_copy(self, skipped: bool) -> bool: + """Count one guarded copy (``skipped``: not issued) and log the counts whenever the calls reach a power of + two; returns ``skipped``.""" + if skipped: + self.block_offset_skips += 1 + else: + self.block_offset_copies += 1 + calls = self.block_offset_skips + self.block_offset_copies + if calls & (calls - 1) == 0: + logger.info( + f"[M3 host] block-offset copies skipped when unchanged: {self.block_offset_skips} of {calls}" + ) + return skipped + @staticmethod def _derive_reuse_salt(cache_salt: str | None) -> int | None: """Derive ``ReuseScope.salt`` (int|None) from the ``cache_salt`` string. diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 221f9145b17c..d87da7f1e6be 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -854,6 +854,11 @@ def _init_cuda_graph_lora_manager(self, lora_config: LoraConfig): if isinstance(self._runner, DecoderRunner): self._runner.init_cuda_graph_lora_manager(lora_config) + 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"): diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 9d179a02f620..87d960f4d2e5 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -114,6 +114,7 @@ get_from_waiting_queue, merge_requests) from .resource_manager import (NoFreeSlotsError, ResourceManager, ResourceManagerType, request_context) +from .runtime_hooks import runtime_hooks_of from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState, SampleStateTensors) from .scheduler import (RequestScheduler, ScheduledRequests, @@ -379,6 +380,13 @@ class PendingEncoderStep: class PyExecutor: + # True while the overlap loop runs with the execution stream current + # (_overlap_loop_stream_scope). + _loop_on_execution_stream = False + # Whether the model opted into running the overlap loop on the execution + # stream (RuntimeHooks.loop_on_execution_stream); read once at construction. + _execution_stream_loop_armed = False + # Minimum number of async micro batches for async PP execution. # This is a trade-off between memory usage and performance. # If the number of micro batches is too small, the executor will spend too much time in synchronization. @@ -490,6 +498,8 @@ def __init__( # related modules self.resource_manager = resource_manager self.model_engine = model_engine + self._execution_stream_loop_armed = runtime_hooks_of( + model_engine.model).loop_on_execution_stream self._enable_adp_dummy_fixes = getattr(model_engine, "_enable_adp_dummy_fixes", False) self._enable_scheduler_aware_adp_dummy = getattr( @@ -5602,11 +5612,42 @@ def _wait_for_model_engine_input_copy(self): if wait_for_input_copy is not None: wait_for_input_copy() + @contextmanager + def _overlap_loop_stream_scope(self): + """Run the overlap loop with the execution stream current, when the model opted in. + + _forward_step makes the execution stream wait for the current stream + before every forward and the current stream wait for the execution + stream after it. On the thread's default stream, which in steady decode + carries only these waits and timing events, step N+1's inputs wait for + decode graph N through two cross-stream event hops. With the execution + stream current the loop's stream work is queued in the same order on one + stream and _forward_step skips both waits. + + Not taken with a two-model drafter, a guided decoder, a KV connector, + an encoder-decoder model or multimodal encoder item scheduling, whose + side streams synchronize with the current stream. + """ + if not (self._execution_stream_loop_armed and self.drafter is None and + self.draft_model_engine is None and self.guided_decoder is None + and self.kv_connector_manager is None + and not self.is_encoder_decoder + and not self._mm_encoder_item_scheduling_enabled): + yield + return + with torch.cuda.stream(self.execution_stream): + self._loop_on_execution_stream = True + try: + yield + finally: + self._loop_on_execution_stream = False + def _executor_loop_overlap(self): torch.cuda.set_device(self.device_id) # ensure the context is created, otherwise, some MPI calls will fail. CUASSERT(cudart.cudaSetDevice(self.device_id)) - with self._profiler() as profile_step, self.hang_detector: + with self._profiler() as profile_step, self.hang_detector, \ + self._overlap_loop_stream_scope(): iter_start_time = time.time() iter_stats = None target_inputs = None @@ -8311,8 +8352,10 @@ def forward(scheduled_requests, resource_manager, new_tensors_device, cache_indirection_buffer = self.sampler.get_cache_indirection() # Run model forward on the execution stream for proper synchronization - # with KVCacheTransferManager's onboard/offload operations. - self.execution_stream.wait_stream(torch.cuda.current_stream()) + # with KVCacheTransferManager's onboard/offload operations. With the + # loop on the execution stream both waits are program order. + if not self._loop_on_execution_stream: + self.execution_stream.wait_stream(torch.cuda.current_stream()) self._attach_encoder_output_to_execution_stream(scheduled_requests) with torch.cuda.stream(self.execution_stream): outputs = forward(scheduled_requests, self.resource_manager, @@ -8326,7 +8369,8 @@ def forward(scheduled_requests, resource_manager, new_tensors_device, # Ensure the default stream waits for execution_stream to complete # before downstream operations use the outputs. - torch.cuda.current_stream().wait_stream(self.execution_stream) + if not self._loop_on_execution_stream: + torch.cuda.current_stream().wait_stream(self.execution_stream) self._kv_connector_wait_for_save() if route_capture is not None: @@ -8455,9 +8499,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 348328d75d96..84ef85352cf6 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -971,6 +971,7 @@ def allocation_scope(current_stage: ExecutorMemoryType): ): # GuidedDecoder is applicable to non-speculative decoding and two-model 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 new file mode 100644 index 000000000000..651777934a81 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py @@ -0,0 +1,255 @@ +# 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 once by the engine. + +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. 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 TYPE_CHECKING, 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 + + +@dataclass(frozen=True) +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 + # Issue the step-input pack's host copies with one parallel scatter (a program per copy) instead of the + # serial one-CTA scatter, when their destinations do not overlap (pack_step_inputs). + pack_parallel_scatter: bool = False + # Issue the step-input pack inside the replayed graph (a prologue copies a pinned window and scatters it), so no + # operation is left between the decode graphs (needs pack_parallel_scatter). + pack_in_graph: bool = False + # With pack_in_graph: a side copy stream uploads the pinned window into the prologue's own device buffer while the + # previous graph still runs, so the prologue only waits on that copy and scatters device to device (no host copy + # at the head of the replay). + pack_early_copy: bool = False + # Skip host-to-device metadata copies whose values the destination already holds. + skip_unchanged_h2d: bool = False + # With skip_unchanged_h2d: a changed packed 2-D upload stages only each row's changed column span (a page table that + # gained a page re-sends a few bytes, not every padded row), so the pack keeps fitting the in-graph window. + h2d_changed_span: bool = False + # A speculative step under the overlap scheduler re-derives its KV-length metadata (on_update_kv_lens) once, + # after the forward's KV-length correction, instead of also before it: the first pass's device writes are all + # rewritten by the second. + kv_lens_update_once: 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 + # 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 + # Stage the one-model speculative sampler's per-step Philox seed and offset + # uploads into the step-input pack (pack_step_inputs), so they ride its one + # upload instead of each being a copy the next graph waits for. + pack_spec_rng_uploads: bool = False + # The one-model speculative step's small per-step uploads: the accepted-draft + # counts join the step-input pack, and the spec sampler keeps its slot ids in + # a device buffer it re-uploads only when they change. + spec_small_h2d: bool = False + # Skip the per-step accepted-draft-count upload (spec_metadata.num_accepted_draft_tokens is then None). Only for a + # model whose forward and drafter never read that tensor. + skip_accepted_draft_upload: bool = False + # The one-model speculative metadata re-uploads its per-request slot-id table (batch_slot_ids, read by rejection + # sampling and the penalties) only when the batch's slot ids change (steady decode keeps them), instead of + # building a pinned tensor and copying it every step. + spec_slot_ids_changed_only: bool = False + # The DFlash / standalone DSpark drafter re-uploads its batch_indices and batch-to-slot map only when they change + # (steady decode keeps both), instead of copying both every step. + spec_batch_map_changed_only: bool = False + # Speculative step-input staging folded into the forward (see SpecStepFold). + spec_step_fold: Optional[SpecStepFold] = None + # Copy the one-model speculative sampler's per-slot store to the host on the sampler's D2H side + # stream, so those copies overlap the next step's input staging instead of preceding it. + sampler_d2h_side_stream: bool = False + # MiniMax-M3 MSA metadata preparation. + msa: MsaRuntimePolicy = MsaRuntimePolicy() + # A separate draft KV cache manager (KVCacheManagerV2) skips a decode-only block-table copy whose destination, + # copied rows and those rows' page indices equal the previous copy's (the standalone DSpark drafter's table is + # unchanged on most decode steps; see KVCacheManagerV2._block_offsets_unchanged). + skip_unchanged_draft_block_offsets: bool = False + # The one-model speculative sampler's tail bookkeeping without its no-op launches, bitwise the same outputs + # (copied to spec_metadata.tail_bookkeep): at one decode request the draft block reads its sampling parameters as + # views of the per-token buffers, a draft block with no d2t map skips the int32 -> int64 -> int32 round trip, and + # a decode-only rejection accept returns the acceptance kernel's outputs instead of copying them into fresh ones. A + # model's own worker may drop more of them under the same flag. + spec_tail_bookkeep: bool = False + # Under disaggregated serving, ship a standalone drafter's context window (its separate draft + # KV pool) alongside the target's KV, so the generation worker drafts from the prompt instead + # of an empty window. Takes effect only with a separate draft KV cache manager on both sides + # (see disaggregation/draft_kv.py). + disagg_draft_kv: bool = False + + @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 + or self.skip_unchanged_draft_block_offsets + ) + + +DEFAULT_RUNTIME_HOOKS = RuntimeHooks() + + +def runtime_hooks_of(model: object) -> RuntimeHooks: + """The hooks ``model`` declares, or the all-off default when it declares none. + + Declaring hooks is optional for every model class, so absence is the common + case and means the default. + """ + hooks = getattr(model, "runtime_hooks", None) + if hooks is None: + return DEFAULT_RUNTIME_HOOKS + if not isinstance(hooks, RuntimeHooks): + raise TypeError( + f"{type(model).__name__}.runtime_hooks must be a RuntimeHooks, got {type(hooks).__name__}" + ) + return hooks diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py index ae43820dd1df..7b46e2a09c6f 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py @@ -63,6 +63,40 @@ def __getattr__(self, name: str) -> Any: SeedOrTensor = Union[int, torch.Tensor] +# An optional replacement of ``flashinfer.sampling.top_p_renorm_probs`` with +# per-row ``top_p`` tensors, installed by a model at load time: it returns the +# renormed probs, or None to keep the stock call for that input. None (the +# default) leaves every op below unchanged. +_TOP_P_RENORM_OVERRIDE: Optional[Callable[[torch.Tensor, torch.Tensor], Optional[torch.Tensor]]] = ( + None +) + + +def set_top_p_renorm_override( + fn: Optional[Callable[[torch.Tensor, torch.Tensor], Optional[torch.Tensor]]], +) -> None: + """Install (or, with None, remove) the top-p renorm override used by + ``top_p_renorm_probs_op`` and ``compute_probs_from_logits_op``.""" + global _TOP_P_RENORM_OVERRIDE + _TOP_P_RENORM_OVERRIDE = fn + + +# An optional replacement of ``flashinfer.sampling.softmax(logits, temperature)``, +# installed by a model at load time: it returns the probs, or None to keep the +# stock call for that input. None (the default) leaves every op below unchanged. +_SOFTMAX_OVERRIDE: Optional[ + Callable[[torch.Tensor, Optional[torch.Tensor]], Optional[torch.Tensor]] +] = None + + +def set_softmax_override( + fn: Optional[Callable[[torch.Tensor, Optional[torch.Tensor]], Optional[torch.Tensor]]], +) -> None: + """Install (or, with None, remove) the softmax override used by + ``softmax_op`` and ``compute_probs_from_logits_op``.""" + global _SOFTMAX_OVERRIDE + _SOFTMAX_OVERRIDE = fn + @_compiler_disable def radix_topk_op( @@ -241,6 +275,10 @@ def softmax_op( logits: torch.Tensor, temperature: Optional[torch.Tensor], ) -> torch.Tensor: + if _SOFTMAX_OVERRIDE is not None: + overridden = _SOFTMAX_OVERRIDE(logits, temperature) + if overridden is not None: + return overridden probs: torch.Tensor = flashinfer.sampling.softmax( logits, temperature, enable_pdl=get_env_enable_pdl() ) @@ -270,6 +308,10 @@ def top_p_renorm_probs_op( probs: torch.Tensor, top_p: torch.Tensor, ) -> torch.Tensor: + if _TOP_P_RENORM_OVERRIDE is not None: + overridden = _TOP_P_RENORM_OVERRIDE(probs, top_p) + if overridden is not None: + return overridden renormed: torch.Tensor = flashinfer.sampling.top_p_renorm_probs(probs, top_p) return renormed @@ -289,11 +331,21 @@ def compute_probs_from_logits_op( """ if top_k is not None: logits = flashinfer.sampling.top_k_mask_logits(logits, top_k) - probs: torch.Tensor = flashinfer.sampling.softmax( - logits, temperatures, enable_pdl=get_env_enable_pdl() + softmaxed = _SOFTMAX_OVERRIDE(logits, temperatures) if _SOFTMAX_OVERRIDE is not None else None + probs: torch.Tensor = ( + softmaxed + if softmaxed is not None + else flashinfer.sampling.softmax(logits, temperatures, enable_pdl=get_env_enable_pdl()) ) if top_p is not None: - probs = flashinfer.sampling.top_p_renorm_probs(probs, top_p) + overridden = ( + _TOP_P_RENORM_OVERRIDE(probs, top_p) if _TOP_P_RENORM_OVERRIDE is not None else None + ) + probs = ( + overridden + if overridden is not None + else flashinfer.sampling.top_p_renorm_probs(probs, top_p) + ) return probs diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 2464827a1a9a..5bc63d9b0e5a 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -145,6 +145,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_strategy(request, vocab_size=2**31) == GREEDY + ) + def _fast_strategy_params(strategy: Strategy) -> tuple[float, int, float, float]: """Flatten a fast-tier strategy into per-row ``(temperature, top_k, top_p, min_p)``.""" @@ -2513,7 +2541,11 @@ 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 = model_outputs.get("logits") + greedy_tokens = model_outputs.get(GREEDY_TOKENS_KEY) + rows_cuda = raw_logits_cuda if greedy_tokens is None else greedy_tokens + assert rows_cuda is not None, "Expected logits or in-graph greedy tokens" generation_requests = scheduled_requests.generation_requests request_ids = [request.py_request_id for request in generation_requests] @@ -2526,7 +2558,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 model_outputs.get("d2t") is None and all( not request.is_dummy and get_draft_token_length(request) == 0 @@ -2563,6 +2595,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 # A draft request reuses its target's py_request_id, so the id # list alone cannot tell a drafter batch from the target batch @@ -2660,13 +2695,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, - None, - ) + 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: + 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, + None, + ) new_tokens_host = self._copy_to_host(next_tokens) return ( generation_requests, @@ -2681,6 +2723,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..4207c341cce7 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/step_inputs.py @@ -0,0 +1,471 @@ +# 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. + +With ``parallel_scatter`` (``RuntimeHooks.pack_parallel_scatter``) a pack of host copies only, whose destinations do +not overlap, is issued with one program per copy (``_parallel_scatter_kernel``) instead of the serial scatter: the same +bytes land in the same places, and with no overlap their order cannot matter. Any other pack takes the serial scatter. + +With ``in_graph`` (``RuntimeHooks.pack_in_graph``, on top of ``parallel_scatter``) the issue itself moves into the +replayed graph: a graph captured with ``capture_prologue`` starts by copying a fixed pinned window (``ghost``) to the +device staging buffer, records ``ghost_free``, and runs the parallel scatter with the copy counts read from the staged +header (one program per descriptor slot; the slots past the counts write nothing). ``flush(in_graph=True)`` before +such a replay then only fills the window on the host, after ``ghost_free`` shows the previous replay has copied it: no +operation is left between the graphs. A pack that does not fit the window or is not order-free is issued eagerly as +before, with a zero-count window so the graph's prologue writes nothing. + +With ``early_copy`` (``RuntimeHooks.pack_early_copy``, on top of ``in_graph``) the window reaches the device before the +replay: ``flush(in_graph=True)`` fills the pinned window and enqueues its upload on a side copy stream, into a device +buffer of the prologue's own (``gdev``), while the previous graph may still run. The prologue then waits for that upload +(``staged_ready``) and scatters device to device; it zeroes the window's header after the scatter and records +``prologue_done``, which the next upload waits for before it overwrites ``gdev``. The next fill of the pinned window +waits for ``ghost_free``, which the upload records once it has read the window. + +Every in-graph attribute has a class-level default (off), so a pack built without its ``__init__`` and every pack built +with both flags off takes exactly the path above (``flush()`` and ``flush(in_graph=True)`` are then the same call). +""" + +from __future__ import annotations + +import contextlib +from contextvars import ContextVar +from typing import Iterator, Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm.logger import logger + + +# first (the copy index offset) varies with the pack's host-copy count: not specialized on its value, so the two +# DEVICE_SRC variants compiled by the warm-up serve every pack (no compile on a serving step). +@triton.jit(do_not_specialize=["first"]) +def _parallel_scatter_kernel( + stage_u8, stage_i64, first, DEVICE_SRC: tl.constexpr, BLOCK: tl.constexpr +): + """Program i writes copy ``first + i`` ([dst, src, bytes] descriptors after the 2-word header). ``src`` is an offset + into the staged buffer for a host copy, a device address for a device copy (``DEVICE_SRC``).""" + d = 2 + 3 * (first + tl.program_id(0)) + dst = tl.load(stage_i64 + d).to(tl.pointer_type(tl.uint8)) + src_word = tl.load(stage_i64 + d + 1) + if DEVICE_SRC: + src = src_word.to(tl.pointer_type(tl.uint8)) + else: + src = stage_u8 + src_word + nbytes = tl.load(stage_i64 + d + 2) + for start in range(0, nbytes, BLOCK): + o = start + tl.arange(0, BLOCK) + tl.store(dst + o, tl.load(src + o, mask=o < nbytes), mask=o < nbytes) + + +# The in-graph prologue's scatter: no integer argument (the counts are read from the staged header and the grid is the +# fixed MAX_DESC), so the two DEVICE_SRC variants its warm-up compiles serve every replay and every capture. +@triton.jit +def _graph_scatter_kernel(stage_u8, stage_i64, DEVICE_SRC: tl.constexpr, BLOCK: tl.constexpr): + """Descriptor slot i of the host copies (``DEVICE_SRC`` False) or of the device copies, with the counts read from + the staged header; a slot past its count writes nothing.""" + i = tl.program_id(0) + num_h2d = tl.load(stage_i64) + if DEVICE_SRC: + count = tl.load(stage_i64 + 1) + d = 2 + 3 * (num_h2d + i) + else: + count = num_h2d + d = 2 + 3 * i + if i < count: + dst = tl.load(stage_i64 + d).to(tl.pointer_type(tl.uint8)) + src_word = tl.load(stage_i64 + d + 1) + if DEVICE_SRC: + src = src_word.to(tl.pointer_type(tl.uint8)) + else: + src = stage_u8 + src_word + nbytes = tl.load(stage_i64 + d + 2) + for start in range(0, nbytes, BLOCK): + o = start + tl.arange(0, BLOCK) + tl.store(dst + o, tl.load(src + o, mask=o < nbytes), mask=o < nbytes) + + +def _disjoint(ranges: list[tuple[int, int]]) -> bool: + """Whether the [start, start + bytes) ranges are pairwise disjoint.""" + end = -1 + for start, nbytes in sorted(ranges): + if start < end: + return False + end = start + nbytes + return True + + +def _parallel_ok(h2d: list[tuple[int, int, int]], d2d: list[tuple[int, int, int]]) -> bool: + """Whether issuing the host copies in parallel and then the device copies in parallel writes what the serial + scatter (every copy in order) writes: the host copies' destinations are disjoint, and the device copies' + destinations are disjoint and none of them is a device copy's source (a host copy may still feed a device copy, + as the launches run in order).""" + if not _disjoint([(dst, n) for dst, _, n in h2d]): + return False + if not d2d: + return True + dsts = [(dst, n) for dst, _, n in d2d] + if not _disjoint(dsts): + return False + return all(_disjoint(dsts + [(src, n)]) for _, src, n in d2d) + + +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 + # The pinned window a graph's prologue copies (in_graph): the header, the descriptors and up to ~10 KB of data. + GRAPH_WINDOW = 1 << 14 + + # Class-level defaults: a pack built without __init__ (or with the flags off) never takes an in-graph branch. + parallel_scatter = False + in_graph = False + early_copy = False + # Engagement counters of this pack (one per engine and rank): packs the replay's prologue issued, packs issued + # eagerly before a prologue graph (over the window or not order-free), and side-stream window uploads + # (early_copy). Logged whenever the in-graph flushes (windows + fallbacks) reach a power of two. + graph_windows = 0 + eager_fallbacks = 0 + early_uploads = 0 + + def __init__( + self, + device: torch.device, + parallel_scatter: bool = False, + in_graph: bool = False, + early_copy: bool = False, + ) -> None: + # Only an opted-in pack needs the native scatter fake registration. + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import ( # noqa: F401 + fake_ops as _m3_fake_ops, # registers trtllm::step_inputs_scatter's fake impl + ) + + self.parallel_scatter = parallel_scatter + self.in_graph = in_graph and parallel_scatter + self.early_copy = early_copy and self.in_graph + if self.early_copy: + # The prologue's own device window: an eager flush (self.dev) can never leave a header in it. + self.gdev = torch.zeros(self.GRAPH_WINDOW, dtype=torch.uint8, device=device) + self.copy_stream = torch.cuda.Stream(device=device) + # Recorded by the side upload once it landed (the prologue waits on it) ... + self.staged_ready = torch.cuda.Event(external=True) + # ... and by every replayed prologue after its scatter (the next upload waits on it). + self.prologue_done = torch.cuda.Event(external=True) + # A CUDA event is created on its first record, and a wait on one never created is skipped: a capture + # would then leave the prologue's wait out of the graph. Create both now. + self.staged_ready.record(self.copy_stream) + self.prologue_done.record(self.copy_stream) + if self.in_graph: + self.ghost = torch.zeros( + self.GRAPH_WINDOW, dtype=torch.uint8, pin_memory=prefer_pinned() + ) + # Recorded by every replayed prologue right after it copied the window (an event node), or by the side + # upload once it read the window (early_copy). + self.ghost_free = torch.cuda.Event(external=True) + 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 + if parallel_scatter: + self._warm_up(device) + if self.in_graph: + self._warm_up_graph_kernel() + + def _warm_up(self, device: torch.device) -> None: + """Compile the parallel scatter now, not on a serving step's flush.""" + dst = torch.zeros(8, dtype=torch.int32, device=device) + self.begin() + self.stage(dst[:4], torch.arange(4, dtype=torch.int32)) + self.stage(dst[4:], dst[:4]) + self.flush() + torch.cuda.synchronize(device) + + def _warm_up_graph_kernel(self) -> None: + """Compile both variants of the prologue's scatter now, not inside a capture (a zero-count header: it writes + nothing).""" + self.dev[: self.HEADER_BYTES].zero_() + self._launch_graph_scatter(self.dev) + torch.cuda.synchronize(self.dev.device) + + def _launch_graph_scatter(self, window: torch.Tensor) -> None: + stage_i64 = window.view(torch.int64) + for device_src in (False, True): + _graph_scatter_kernel[(self.MAX_DESC,)]( + window, stage_i64, DEVICE_SRC=device_src, BLOCK=256, num_warps=1 + ) + + def capture_prologue(self) -> bool: + """Inside a graph capture: the prologue that issues the window ``flush(in_graph=True)`` fills (see the module + docstring). Returns whether the graph has it.""" + if not self.in_graph: + return False + if self.early_copy: + torch.cuda.current_stream().wait_event(self.staged_ready) + self._launch_graph_scatter(self.gdev) + # Consumed: a replay with no new upload scatters nothing. + self.gdev[: self.HEADER_BYTES].zero_() + self.prologue_done.record() + return True + self.dev[: self.GRAPH_WINDOW].copy_(self.ghost, non_blocking=True) + self.ghost_free.record() + self._launch_graph_scatter(self.dev) + return True + + def _upload_window(self, nbytes: int) -> None: + """early_copy: upload the filled pinned window into gdev on the side stream, once the last replayed prologue + has consumed gdev; the window is fully written on the host before this is enqueued.""" + stream = self.copy_stream + stream.wait_event(self.prologue_done) + with torch.cuda.stream(stream): + self.gdev[:nbytes].copy_(self.ghost[:nbytes], non_blocking=True) + self.ghost_free.record(stream) + self.staged_ready.record(stream) + + def _ghost_writable(self) -> None: + # The last replayed prologue (or side upload) has read the window (a no-op before the first one). + self.ghost_free.synchronize() + + def quiesce_graph_window(self) -> None: + """Before replaying a graph with the prologue without a flush: a zero-count window, so it writes nothing.""" + if self.in_graph and self.ghost[: self.HEADER_BYTES].any(): + self._ghost_writable() + self.ghost[: self.HEADER_BYTES].zero_() + if self.early_copy: + # A window uploaded but never replayed would still sit in gdev: overwrite its header. + self._upload_window(self.HEADER_BYTES) + + def _count(self, in_window: bool) -> None: + """Count one in-graph flush (and its early upload); log the counters whenever the number of in-graph flushes + reaches a power of two (per rank: one pack per engine).""" + if in_window: + self.graph_windows += 1 + self.early_uploads += int(self.early_copy) + else: + self.eager_fallbacks += 1 + flushes = self.graph_windows + self.eager_fallbacks + if flushes & (flushes - 1) == 0: + logger.info( + f"[M3 host] step-input pack engagement: {self.graph_windows} in-graph windows, " + f"{self.eager_fallbacks} eager fallbacks, {self.early_uploads} early uploads" + ) + + 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, in_graph: bool = False) -> None: + """Issue the pack. ``in_graph``: the next replay is of a graph captured with ``capture_prologue``.""" + if not self.open: + if in_graph: + self.quiesce_graph_window() + return + self.open = False + if in_graph and self.in_graph: + if self._fill_graph_window(): + if self.early_copy: + self._upload_window(self.used) + self._count(in_window=True) + self._clear() + return + self.quiesce_graph_window() + self._count(in_window=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) + if self.parallel_scatter and _parallel_ok(self.h2d, self.d2d): + stage_i64 = self.dev.view(torch.int64) + if self.h2d: + _parallel_scatter_kernel[(len(self.h2d),)]( + self.dev, stage_i64, 0, DEVICE_SRC=False, BLOCK=256, num_warps=2 + ) + if self.d2d: + _parallel_scatter_kernel[(len(self.d2d),)]( + self.dev, stage_i64, len(self.h2d), DEVICE_SRC=True, BLOCK=256, num_warps=2 + ) + else: + 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 _fill_graph_window(self) -> bool: + """Write the pack into the pinned window the replay's prologue issues; False when it does not fit or is not + order-free (the caller issues it eagerly).""" + if self.used > self.GRAPH_WINDOW or not _parallel_ok(self.h2d, self.d2d): + return False + desc = [len(self.h2d), len(self.d2d)] + [ + value for copy in self.h2d + self.d2d for value in copy + ] + self._ghost_writable() + host = self.host[self.parity] + self.ghost[: self.used].copy_(host[: self.used]) + self.ghost[: 8 * len(desc)].view(torch.int64).copy_(torch.tensor(desc, dtype=torch.int64)) + return True + + 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_active(pack: StepInputPack) -> bool: + """Whether ``pack`` is the calling context's pack.""" + return _active.get() is pack + + +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(in_graph: bool = False) -> None: + """Issue the active open pack (one upload and one scatter launch on the current stream) and close it. + + ``in_graph``: the next replay is of a graph whose prologue issues the active pack (``capture_prologue``).""" + pack = _active.get() + if pack is not None: + pack.flush(in_graph=in_graph) + + +def capture_prologue() -> Optional[StepInputPack]: + """Inside a graph capture: add the active pack's in-graph prologue; the pack whose window the graph issues, or + None when the graph has no prologue (no active pack, or one without ``in_graph``: nothing is captured).""" + pack = _active.get() + return pack if pack is not None and pack.capture_prologue() else None diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 0db5892872bd..6c9dc7ab938e 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -576,12 +576,16 @@ class DFlashSpecMetadata(SpecMetadata): dtype: torch.dtype = torch.bfloat16 captured_hidden_states: Optional[torch.Tensor] = None + # Highest arange prefix uploaded to this metadata instance. + _batch_indices_uploaded = 0 + def __post_init__(self): self.batch_indices_cuda = torch.empty( [self.max_num_requests], dtype=torch.int, device="cuda", ) + self._batch_indices_uploaded = 0 self.is_spec_dec_tree = False self.is_spec_dec_dynamic_tree = False @@ -622,10 +626,17 @@ def prepare(self): assert self.request_ids is not None num_seqs = len(self.request_ids) - batch_indices = torch.arange( - num_seqs, dtype=torch.int, device="cpu", pin_memory=prefer_pinned() - ) - self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) + # getattr: unit tests drive prepare() with stand-in metadata and workers. + changed_only = getattr(self, "batch_map_changed_only", False) + # The batch indices are an arange, so with batch_map_changed_only a buffer already holding a long enough + # prefix is not re-uploaded (nothing else writes batch_indices_cuda). + if not (changed_only and num_seqs <= self._batch_indices_uploaded): + batch_indices = torch.arange( + num_seqs, dtype=torch.int, device="cpu", pin_memory=prefer_pinned() + ) + self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) + if changed_only: + self._batch_indices_uploaded = num_seqs # Update slot mapping for DFlash context buffers worker = getattr(self, "_dflash_worker", None) @@ -660,13 +671,25 @@ def prepare(self): # Route unknown request IDs (cuda graph padding or warmup dummies) # to dummy slot to avoid corrupting real request's context - mapping = torch.tensor( - [worker._req_to_slot.get(rid, worker._dummy_slot) for rid in self.request_ids], - dtype=torch.long, - device="cpu", - pin_memory=prefer_pinned(), - ) - worker._batch_to_slot[:num_seqs].copy_(mapping, non_blocking=True) + slots = [worker._req_to_slot.get(rid, worker._dummy_slot) for rid in self.request_ids] + # With batch_map_changed_only the map is re-uploaded only when it differs from the last upload (see + # DFlashWorker._batch_to_slot_uploaded). + if not (changed_only and getattr(worker, "_batch_to_slot_uploaded", None) == slots): + mapping = torch.tensor( + slots, + dtype=torch.long, + device="cpu", + pin_memory=prefer_pinned(), + ) + worker._batch_to_slot[:num_seqs].copy_(mapping, non_blocking=True) + if changed_only: + worker._batch_to_slot_uploaded = slots + + # A target's per-iteration host hook (None in stock), e.g. an engagement counter. prepare() runs on the host + # every iteration, outside any captured graph, after the batch's sampling mode is set. + hook = getattr(worker, "_host_step_hook", None) + if hook is not None: + hook(self) def is_layer_capture(self, layer_id: int) -> bool: return layer_id in self._capture_layer_set @@ -735,6 +758,10 @@ class DFlashWorker(SpecWorkerBase): Reference: https://arxiv.org/pdf/2602.06036 """ + # The map DFlashSpecMetadata.prepare last uploaded into _batch_to_slot[:len(map)] under batch_map_changed_only, or + # None when unknown. Every other writer of _batch_to_slot (its allocation, the context-step rebuild) resets it. + _batch_to_slot_uploaded: Optional[List[int]] = None + def __init__( self, spec_config: "DFlashDecodingConfig", @@ -764,6 +791,14 @@ def __init__( self._ctx_len_restore_pending = False # Deferred kv_lens_cuda rewind state (see _prepare_kv_for_draft_forward, # _apply_kv_rewind_after_draft, _ensure_spec_dec_state_restored). + self._lean_glue = False + self._lean_aranges = {} + self._fused_ctx_store = None + self._fused_block_refresh = None + self._fused_drafter_inputs = None + self._fused_ctx_advance = None + self._fused_kv_rewind = None + self._host_step_hook = None self._kv_rewind_pending = False self._kv_rewind_amount = None self._kv_rewind_nc = None @@ -1070,6 +1105,15 @@ def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool: # pool, and the drafter addresses blocks, not planes. # clone(): .to() is a no-op for an int64 source, and the in-place ops # below would then edit attn_metadata's own tensor. + if self._fused_block_refresh is not None and self._fused_block_refresh( + src, + self._ctx_pool_idx, + num_seqs, + self._ctx_block_divisor, + self._ctx_block_tables, + self._ctx_block_counts, + ): + return True encoded = src[self._ctx_pool_idx, :num_seqs, 0].to(torch.int64).clone() # Before the clamp erases the placeholders. Informative only on the # python V2 backend; the cpp one already mapped BAD_PAGE_INDEX to 0. @@ -1165,6 +1209,7 @@ def _lazy_init_ctx_buffers( self._ctx_base_pages = torch.zeros(num_slots, dtype=torch.long, device="cuda") self._ctx_base_host = [0] * num_slots self._batch_to_slot = torch.zeros(max_batch, dtype=torch.long, device="cuda") + self._batch_to_slot_uploaded = None self._free_slots = deque(range(max_batch)) self._req_to_slot = {} @@ -1399,6 +1444,124 @@ def _get_ctx_paged_append(self) -> Callable[..., None]: self._ctx_paged_append = get_dflash_paged_append() return self._ctx_paged_append + def _static_arange(self, start: int, end: int) -> torch.Tensor: + """torch.arange(start, end) (long, cuda). With lean glue on, cached per (start, end) outside CUDA-graph capture + (read-only uses only); never cached during capture, so no graph-pool tensor outlives its graph.""" + if not self._lean_glue: + return torch.arange(start, end, dtype=torch.long, device="cuda") + key = (start, end) + t = self._lean_aranges.get(key) + if t is None: + t = torch.arange(start, end, dtype=torch.long, device="cuda") + if not torch.cuda.is_current_stream_capturing(): + self._lean_aranges[key] = t + return t + + def enable_lean_glue(self, draft_model, enabled: bool = True) -> None: + """Skip the drafter's copies whose consumer reads strided views (worker and draft model); off by default. + The values are bitwise the same; only the copies go.""" + self._lean_glue = bool(enabled) + if hasattr(draft_model, "_dflash_lean_glue"): + draft_model._dflash_lean_glue = bool(enabled) + + def enable_fused_ctx_store(self, store) -> bool: + """Install a target's one-launch context K / V store (None, the default, keeps the stock path). + + ``store(draft_model, projected, positions, num_accepted, request_rows, cols, block_table, kv_buf, + tokens_per_request)`` must write exactly what precompute_context_kv + the accept mask + _store_context_kv_paged + write, and return False (writing nothing) for a drafter or cache it does not reproduce. It is first called with + ``cols`` None and the keywords ``capacity_counts`` (bound pool: pages per block-table row) / ``max_col`` (the + private arena's last column), to derive each row's column from its position as the stock clamp does.""" + self._fused_ctx_store = store + return store is not None + + def enable_fused_ctx_advance(self, advance) -> bool: + """Install a target's one-launch context advance + read view on a bound draft pool (None, the default, keeps + the stock path). + + ``advance(slots, rows, num_accepted, ctx_len, max_ctx, dummy_slot, block_counts, base_pages, tables, + tables_read, page_size, block_size)`` must leave ``ctx_len`` and ``tables_read[rows]`` exactly as + _advance_ctx_len and the read view below do and return the advertised context lengths, or None (writing + nothing) for a layout it does not reproduce.""" + self._fused_ctx_advance = advance + return advance is not None + + def enable_fused_forward_indices(self, draft_model, indices) -> bool: + """Install a target's one-launch index prelude for the draft model's paged block decode (None, the default, + keeps the stock lines). + + ``indices(num_ctx, batch_idx, query_positions, block_size, pages_per_slot, page_size)`` must return the int32 + (cache_seqlens, cache_batch_idx, seq_lens_after, kv_indptr, kv_last_page_len, append_positions, + query_positions_flat) exactly as the stock lines of dflash_forward compute them, or None for a layout it does + not reproduce.""" + if not hasattr(draft_model, "_dflash_fused_indices"): + return False + draft_model._dflash_fused_indices = indices + return indices is not None + + def enable_fused_qk_norm_rope(self, draft_model, fused) -> bool: + """Install a target's one-launch q / k RMSNorm + RoPE for the draft model's flashinfer-RoPE branch (None, the + default, keeps the stock lines). + + ``fused(qkv, q_size, kv_size, head_dim, q_norm, k_norm, positions, cos_sin_cache, is_neox)`` must return new + (q, k) exactly as the stock q_norm / k_norm + flashinfer RoPE lines leave them, or None for a layout it does + not reproduce.""" + if not hasattr(draft_model, "_dflash_fused_qk_norm_rope"): + return False + draft_model._dflash_fused_qk_norm_rope = fused + return fused is not None + + def enable_fused_qk_norm_rope_append(self, draft_model, fused) -> bool: + """Install a target's one-launch q / k RMSNorm + RoPE that also appends the block's K / V into the layer's + paged cache, for the draft model's flashinfer-RoPE branch on the TRTLLM backend (None, the default, keeps the + stock lines). + + ``fused(qkv, q_size, kv_size, head_dim, q_norm, k_norm, positions, cos_sin_cache, is_neox, paged_kv_cache, + append_batch_indices, append_positions, kv_indices, kv_indptr)`` must return new (q, k) exactly as + ``enable_fused_qk_norm_rope``'s op does AND leave the cache exactly as flashinfer's append_paged_kv_cache of + that K / V would; or return None (nothing computed or written), and the stock lines run.""" + if not hasattr(draft_model, "_dflash_fused_qk_norm_rope_append"): + return False + draft_model._dflash_fused_qk_norm_rope_append = fused + return fused is not None + + def enable_attention_pdl(self, draft_model, enabled: bool) -> bool: + """Launch the draft model's TRTLLM-backend block attention (trtllm-gen decode / context) with programmatic + dependent launch: the same kernel and inputs, launched to overlap the previous kernel's tail (it waits for that + grid before reading). Off by default.""" + if not hasattr(draft_model, "_dflash_attn_pdl"): + return False + draft_model._dflash_attn_pdl = bool(enabled) + return draft_model._dflash_attn_pdl + + def enable_fused_drafter_inputs(self, inputs) -> bool: + """Install a target's one-launch drafter step inputs (None, the default, keeps the stock path). + + ``inputs(embed_tokens, accepted_tokens, num_accepted, slots, ctx_len, max_ctx, mask_token_id, + trained_mask_embedding, block_size, tokens_per_request)`` must return exactly the stock (bonus, + query_position_ids, ctx_position_ids, noise block) of prepare_1st_drafter_inputs' generation branch, or None + (computing nothing) for an embedding or layout it does not reproduce.""" + self._fused_drafter_inputs = inputs + return inputs is not None + + def enable_fused_kv_rewind(self, rewind) -> bool: + """Install a target's one-launch draft-forward KV bookkeeping (None, the default, keeps the stock lines). + + ``rewind(num_accepted, kv_lens, nc, bs)`` must leave ``kv_lens[nc:bs]`` incremented by one and return the + tensor ``1 - num_accepted[nc:bs]`` (the same values and dtype), exactly as _prepare_kv_for_draft_forward's + two lines, or None (writing nothing) for operands it does not take.""" + self._fused_kv_rewind = rewind + return rewind is not None + + def enable_fused_block_refresh(self, refresh) -> bool: + """Install a target's one-launch block-table refresh (None, the default, keeps the stock path). + + ``refresh(offsets, pool_idx, num_seqs, divisor, tables, counts)`` must write exactly what + _refresh_ctx_block_tables writes into ``tables[:num_seqs]`` and ``counts[:num_seqs]``, and return False + (writing nothing) for a layout it does not reproduce.""" + self._fused_block_refresh = refresh + return refresh is not None + def _store_context_kv_paged( self, k: torch.Tensor, @@ -1430,9 +1593,11 @@ def _store_context_kv_paged( positions_i32 = positions.to(torch.int32) kv_indices, kv_indptr = self._ctx_paged_index_args() for layer_idx in range(k.size(1)): + # Lean glue: flashinfer's append is stride-aware (last dim contiguous), so the per-layer slices go + # in as views; the values written are bitwise the same. append_paged_kv_cache( - append_key=k[:, layer_idx].contiguous(), - append_value=v[:, layer_idx].contiguous(), + append_key=k[:, layer_idx] if self._lean_glue else k[:, layer_idx].contiguous(), + append_value=v[:, layer_idx] if self._lean_glue else v[:, layer_idx].contiguous(), batch_indices=rows_i32, positions=positions_i32, paged_kv_cache=self._ctx_kv_buf[layer_idx], @@ -1460,11 +1625,22 @@ def _prepare_kv_for_draft_forward( ): """Adjust kv_lens_cuda so the draft model sees correct RoPE positions.""" if hasattr(attn_metadata, "kv_lens_cuda"): - self._kv_rewind_amount = 1 - num_accepted_tokens[num_contexts:batch_size] + # A target's one-launch bookkeeping: the same rewind amount and kv_lens increment as the lines below. + fused = ( + self._fused_kv_rewind( + num_accepted_tokens, attn_metadata.kv_lens_cuda, num_contexts, batch_size + ) + if self._fused_kv_rewind is not None and batch_size > num_contexts + else None + ) + if fused is not None: + self._kv_rewind_amount = fused + else: + self._kv_rewind_amount = 1 - num_accepted_tokens[num_contexts:batch_size] self._kv_rewind_nc = num_contexts self._kv_rewind_bs = batch_size - if batch_size > num_contexts: + if fused is None and batch_size > num_contexts: attn_metadata.kv_lens_cuda[num_contexts:batch_size] += 1 self._kv_rewind_pending = True @@ -1828,6 +2004,7 @@ def _forward_impl( ), non_blocking=True, ) + self._batch_to_slot_uploaded = None inputs = self.prepare_1st_drafter_inputs( input_ids=input_ids, @@ -1843,6 +2020,12 @@ def _forward_impl( if num_gens > 0: with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + # Lean glue: in bound mode the cache rows are the contiguous gen rows, so the block table is a slice. + draft_model._dflash_cache_rows = ( + (num_contexts, num_contexts + num_gens) + if self._lean_glue and self._ctx_block_tables is not None + else None + ) hidden_states_out = draft_model.dflash_forward( noise_embedding=inputs["noise_embedding"], query_positions=inputs["query_positions"], @@ -1859,13 +2042,22 @@ def _forward_impl( # Which block slots carry them is a drafter-family convention, # resolved through _draft_slot_ids. block_size = self._compute_block_size - gen_gather_ids = self._draft_slot_ids(draft_model, num_gens, block_size, K) - # Shields only the last request: at block_size == K with - # shift_label off, slots run 1..K, so every request reads the - # next one's slot 0 and the last overruns. Degrades, never raises. - gen_gather_ids = gen_gather_ids.clamp(max=hidden_states_out.shape[0] - 1) - - gen_hidden_states = hidden_states_out[gen_gather_ids] + if ( + self._lean_glue + and getattr(draft_model, "_dspark_shift_label", False) + and block_size == K + and hidden_states_out.shape[0] == num_gens * block_size + ): + # Lean glue: under shift_label with block_size == K the slot ids are arange(num_gens * K), so the + # gather is the identity. + gen_hidden_states = hidden_states_out + else: + gen_gather_ids = self._draft_slot_ids(draft_model, num_gens, block_size, K) + # Shields only the last request: at block_size == K with + # shift_label off, slots run 1..K, so every request reads the + # next one's slot 0 and the last overruns. Degrades, never raises. + gen_gather_ids = gen_gather_ids.clamp(max=hidden_states_out.shape[0] - 1) + gen_hidden_states = hidden_states_out[gen_gather_ids] gen_logits = draft_model.logits_processor( gen_hidden_states, draft_model.lm_head, attn_metadata, True ) @@ -2196,39 +2388,61 @@ def prepare_1st_drafter_inputs( # Get slots for gen requests from pre-computed mapping slots = self._batch_to_slot[num_contexts : num_contexts + num_gens] - gen_rows_out = torch.arange( - num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" - ) + gen_rows_out = self._static_arange(num_contexts, num_contexts + num_gens) K_plus_1 = K + 1 - bonus_idx = (gen_num_accepted - 1).clamp_min(0).long().unsqueeze(1) - bonus = gen_accepted_tokens.gather(1, bonus_idx).squeeze(1).long() - - ctx_len_gen = self._ctx_len[slots] - j_block = torch.arange(query_tokens_per_req, dtype=torch.long, device="cuda") - offsets_kp1 = torch.arange(K_plus_1, dtype=torch.long, device="cuda") - - # _ctx_len is clamped to _max_ctx only AFTER this step's accepted - # tokens are folded in (see the update below), so the running length - # used here has to be clamped on its own -- otherwise a request that - # already sits at the ceiling indexes num_accepted positions past - # any position the sequence can legitimately reach. - ctx_len_now = (ctx_len_gen + gen_num_accepted.long()).clamp_(max=self._max_ctx) - query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) - ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) - - # The drafter's own trained mask row, if it kept one for the mask id in use. - trained_mask_embedding = ( - getattr(draft_model, "mask_token_embedding", None) - if mask_token_id == getattr(draft_model, "mask_token_id", None) + # A target's one-launch step inputs: the same bonus, positions and noise block as the stock lines below. + fused_inputs = ( + self._fused_drafter_inputs( + embed_tokens, + gen_accepted_tokens, + gen_num_accepted, + slots, + self._ctx_len, + self._max_ctx, + mask_token_id, + getattr(draft_model, "mask_token_embedding", None) + if mask_token_id == getattr(draft_model, "mask_token_id", None) + else None, + query_tokens_per_req, + K_plus_1, + ) + if self._fused_drafter_inputs is not None else None ) - noise_embed_2d = dflash_noise_block_embedding( - embed_tokens, bonus, mask_token_id, query_tokens_per_req, trained_mask_embedding - ) + if fused_inputs is not None: + bonus, query_position_ids, ctx_position_ids, noise_embed_2d = fused_inputs + offsets_kp1 = self._static_arange(0, K_plus_1) + else: + bonus_idx = (gen_num_accepted - 1).clamp_min(0).long().unsqueeze(1) + bonus = gen_accepted_tokens.gather(1, bonus_idx).squeeze(1).long() + + ctx_len_gen = self._ctx_len[slots] + j_block = self._static_arange(0, query_tokens_per_req) + offsets_kp1 = self._static_arange(0, K_plus_1) + + # _ctx_len is clamped to _max_ctx only AFTER this step's accepted + # tokens are folded in (see the update below), so the running length + # used here has to be clamped on its own -- otherwise a request that + # already sits at the ceiling indexes num_accepted positions past + # any position the sequence can legitimately reach. + ctx_len_now = (ctx_len_gen + gen_num_accepted.long()).clamp_(max=self._max_ctx) + query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) + ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) + + # The drafter's own trained mask row, if it kept one for the mask id in use. + trained_mask_embedding = ( + getattr(draft_model, "mask_token_embedding", None) + if mask_token_id == getattr(draft_model, "mask_token_id", None) + else None + ) + noise_embed_2d = dflash_noise_block_embedding( + embed_tokens, bonus, mask_token_id, query_tokens_per_req, trained_mask_embedding + ) # Accumulate new accepted features into context buffers + fused_num_ctx = None if has_target_features: gen_start = attn_metadata.num_ctx_tokens # Target now processes exactly K+1 tokens per gen req, so the @@ -2242,81 +2456,142 @@ def prepare_1st_drafter_inputs( projected_rows = projected_to_store.reshape(num_gens, K_plus_1, -1) lilicorr_anchor = last_accepted_hidden(projected_rows, gen_num_accepted) gen_num_accepted_long = gen_num_accepted.long() - col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) - write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) - if self._ctx_block_tables is not None: - # Bound by what the manager handed THIS request: the masked - # entries below are still written (fixed-size, for graph - # safety) and a placeholder page resolves to another's. - gen_capacity = self._ctx_block_counts[gen_rows_out] * self._ctx_page_size - col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_(min=0) - else: - col_idx = col_idx.clamp(max=self._max_ctx - 1) - - # Fixed-size writes for CUDA graph compatibility: - # Write ALL entries but zero out invalid ones. Invalid - # writes land at clamped column indices beyond valid - # _ctx_len range, so they're harmless. - slot_flat = slots.unsqueeze(1).expand(-1, K + 1).reshape(-1) - col_flat = col_idx.reshape(-1) - proj_flat = projected_to_store # already [num_gens*(K+1), proj_dim] - pos_flat = ctx_position_ids.long().reshape(-1) - mask_1d = write_mask.reshape(-1) - - # Fast path: store the pre-projected/pre-RoPE'd K/V. - # dflash_forward reads these directly via cache_batch_idx. - cache_dtype = ( - self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype - ) - k_new, v_new = draft_model.precompute_context_kv( - proj_flat.to(cache_dtype), pos_flat - ) - mask_bc = mask_1d.view(-1, 1, 1, 1).to(k_new.dtype) - k_new.mul_(mask_bc) - if v_new is not None: - v_new.mul_(mask_bc) - slot_long = slot_flat.long() - col_long = col_flat.long() - if self._ctx_paged: + # S10c-A2: a target's fused store derives each row's column from its position itself (the stock clamp + # below), so the column glue runs only on the stock path. + fused_stored = False + if self._fused_ctx_store is not None and self._ctx_paged: + bound = self._ctx_block_tables is not None + fused_stored = self._fused_ctx_store( + draft_model, + projected_to_store.to(self._ctx_kv_buf[0].dtype), + ctx_position_ids.long().reshape(-1), + gen_num_accepted_long, + gen_rows_out if bound else slots, + None, + self._ctx_block_tables if bound else self._ctx_page_table, + self._ctx_kv_buf, + K + 1, + capacity_counts=self._ctx_block_counts if bound else None, + max_col=None if bound else self._max_ctx - 1, + ) + if not fused_stored: + col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) + write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) if self._ctx_block_tables is not None: - # Batch positions of the gen requests, matching the - # per-request block table's row order. - rows_long = gen_rows_out.unsqueeze(1).expand(-1, K + 1).reshape(-1) + # Bound by what the manager handed THIS request: the masked + # entries below are still written (fixed-size, for graph + # safety) and a placeholder page resolves to another's. + gen_capacity = self._ctx_block_counts[gen_rows_out] * self._ctx_page_size + col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_( + min=0 + ) else: - rows_long = slot_long - self._store_context_kv_paged(k_new, v_new, rows_long, col_long) - else: # VANILLA DFlash backend (FlashAttention) - self._ctx_k_buf[slot_long, :, col_long] = k_new - if v_new is not None: - self._ctx_v_buf[slot_long, :, col_long] = v_new - - self._advance_ctx_len(slots, gen_num_accepted_long) - - num_ctx_per_req_t = self._ctx_len[slots] - ctx_page_table = self._ctx_page_table - if self._ctx_block_tables is not None: - # The write above clamps columns to the request's allocation, so - # a context that outruns it has its tail written over the last - # valid slot. Truncate what is advertised to match, or the read - # attends that clobbered value -- or another request's page. - allocated = dflash_allocated_ctx_limit( - self._ctx_block_counts, self._ctx_page_size, block_size - ) - num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) - # Pages below a slot's base hold no drafter K/V (a prefix the - # target reused): the read numbers from the base instead, see - # _ctx_base_for_store and dflash_rebased_ctx_view. Writes above - # keep the absolute table. - base_pages = self._ctx_base_pages[slots] - self._ctx_block_tables_read[gen_rows_out], num_ctx_per_req_t = ( - dflash_rebased_ctx_view( - self._ctx_block_tables[gen_rows_out], - base_pages, - num_ctx_per_req_t, + col_idx = col_idx.clamp(max=self._max_ctx - 1) + + # Fixed-size writes for CUDA graph compatibility: + # Write ALL entries but zero out invalid ones. Invalid + # writes land at clamped column indices beyond valid + # _ctx_len range, so they're harmless. + slot_flat = slots.unsqueeze(1).expand(-1, K + 1).reshape(-1) + col_flat = col_idx.reshape(-1) + proj_flat = projected_to_store # already [num_gens*(K+1), proj_dim] + pos_flat = ctx_position_ids.long().reshape(-1) + mask_1d = write_mask.reshape(-1) + + # Fast path: store the pre-projected/pre-RoPE'd K/V. + # dflash_forward reads these directly via cache_batch_idx. + cache_dtype = ( + self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype + ) + # A target's fused store does the K / V epilogue, the accept mask and the paged write in one launch. + fused_stored = ( + self._fused_ctx_store is not None + and self._ctx_paged + and self._fused_ctx_store( + draft_model, + proj_flat.to(cache_dtype), + pos_flat, + gen_num_accepted_long, + gen_rows_out if self._ctx_block_tables is not None else slots, + col_flat, + self._ctx_block_tables + if self._ctx_block_tables is not None + else self._ctx_page_table, + self._ctx_kv_buf, + K + 1, + ) + ) + if not fused_stored: + k_new, v_new = draft_model.precompute_context_kv( + proj_flat.to(cache_dtype), pos_flat + ) + mask_bc = mask_1d.view(-1, 1, 1, 1).to(k_new.dtype) + k_new.mul_(mask_bc) + if v_new is not None: + v_new.mul_(mask_bc) + slot_long = slot_flat.long() + col_long = col_flat.long() + if self._ctx_paged: + if self._ctx_block_tables is not None: + # Batch positions of the gen requests, matching the + # per-request block table's row order. + rows_long = gen_rows_out.unsqueeze(1).expand(-1, K + 1).reshape(-1) + else: + rows_long = slot_long + self._store_context_kv_paged(k_new, v_new, rows_long, col_long) + else: # VANILLA DFlash backend (FlashAttention) + self._ctx_k_buf[slot_long, :, col_long] = k_new + if v_new is not None: + self._ctx_v_buf[slot_long, :, col_long] = v_new + + # A target's one-launch advance + read view: the stock lines' context lengths and read rows, or None. + if self._fused_ctx_advance is not None and self._ctx_block_tables is not None: + fused_num_ctx = self._fused_ctx_advance( + slots, + gen_rows_out, + gen_num_accepted_long, + self._ctx_len, + self._max_ctx, + self._dummy_slot, + self._ctx_block_counts, + self._ctx_base_pages, + self._ctx_block_tables, + self._ctx_block_tables_read, self._ctx_page_size, + block_size, ) - ) + if fused_num_ctx is None: + self._advance_ctx_len(slots, gen_num_accepted_long) + + if fused_num_ctx is not None: + num_ctx_per_req_t = fused_num_ctx ctx_page_table = self._ctx_block_tables_read + else: + num_ctx_per_req_t = self._ctx_len[slots] + ctx_page_table = self._ctx_page_table + if self._ctx_block_tables is not None: + # The write above clamps columns to the request's allocation, so + # a context that outruns it has its tail written over the last + # valid slot. Truncate what is advertised to match, or the read + # attends that clobbered value -- or another request's page. + allocated = dflash_allocated_ctx_limit( + self._ctx_block_counts, self._ctx_page_size, block_size + ) + num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) + # Pages below a slot's base hold no drafter K/V (a prefix the + # target reused): the read numbers from the base instead, see + # _ctx_base_for_store and dflash_rebased_ctx_view. Writes above + # keep the absolute table. + base_pages = self._ctx_base_pages[slots] + self._ctx_block_tables_read[gen_rows_out], num_ctx_per_req_t = ( + dflash_rebased_ctx_view( + self._ctx_block_tables[gen_rows_out], + base_pages, + num_ctx_per_req_t, + self._ctx_page_size, + ) + ) + ctx_page_table = self._ctx_block_tables_read noise_embedding = noise_embed_2d query_positions = query_position_ids.long() diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index fc29a0633f41..c766c5027624 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 @@ -352,8 +353,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 @@ -374,8 +374,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 724b2aa15e36..d0bfd76c28cc 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -19,7 +19,7 @@ from contextlib import contextmanager from dataclasses import dataclass, field from enum import IntEnum -from typing import TYPE_CHECKING, List, Optional, Protocol, Type +from typing import TYPE_CHECKING, Callable, List, Optional, Protocol, Type import torch from packaging.version import Version @@ -32,6 +32,7 @@ from ..attention.backends.trtllm import (AttentionBackend, TrtllmAttention, TrtllmAttentionMetadata) from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE +from ..pyexecutor import step_inputs from ..pyexecutor.resource_manager import ResourceManagerType if TYPE_CHECKING: @@ -48,14 +49,31 @@ from ..pyexecutor.sampler.ops import flashinfer as flashinfer_sampling from ..pyexecutor.sampler.ops.vanilla import greedy_search_sampling_batch +# An optional replacement of the tensor-seeded deterministic chain sampler, +# installed by a model at load time: it takes the op's arguments and returns +# its (accepted_tokens, num_tokens), or None to keep the stock op for that +# input. None (the default) leaves rejection_sampling_one_model unchanged. +_ChainSpecOverride = Callable[ + [torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + Optional[tuple[torch.Tensor, torch.Tensor]]] +_CHAIN_SPEC_OVERRIDE: Optional[_ChainSpecOverride] = None + + +def set_chain_speculative_sampling_override( + fn: Optional[_ChainSpecOverride]) -> None: + """Install (or, with None, remove) the chain speculative sampling override + used by ``rejection_sampling_one_model``.""" + global _CHAIN_SPEC_OVERRIDE + _CHAIN_SPEC_OVERRIDE = fn + def rejection_sampling_one_model( draft_probs: torch.Tensor, draft_token_ids: torch.Tensor, target_probs: torch.Tensor, deterministic: bool = True, - seed: Optional[int] = None, - offset: Optional[int] = None, + seed: Optional[int | torch.Tensor] = None, + offset: Optional[int | torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor]: # chain_speculative_sampling requires flashinfer>=0.6.4. This entry point can # be reached independently of SpecWorkerBase.__init__ (e.g. via @@ -66,6 +84,15 @@ def rejection_sampling_one_model( raise RuntimeError( "Rejection sampling for one-model speculative decoding requires flashinfer>=0.6.4" ) + # The inherited override accepts per-request tensor Philox state and has no + # nondeterministic mode. Keep feature-base scalar/default RNG calls stock. + if (_CHAIN_SPEC_OVERRIDE is not None and deterministic + and isinstance(seed, torch.Tensor) + and isinstance(offset, torch.Tensor)): + overridden = _CHAIN_SPEC_OVERRIDE(draft_probs, draft_token_ids, + target_probs, seed, offset) + if overridden is not None: + return overridden batch_size = draft_token_ids.shape[0] device = draft_token_ids.device output_accepted_token_num = torch.zeros(batch_size, @@ -466,6 +493,79 @@ def from_string(name: Optional[str]) -> "SpeculativeDecodingMode": # slot leaves them room before the next one starts. _RNG_SLOT_SPAN = 32 +# RuntimeHooks.spec_tail_bookkeep (``SpecMetadata.tail_bookkeep``): the +# per-token buffers a single decode request's rows are read from, with the +# dtype ``populate_sampling_params_for_one_model`` / +# ``_populate_request_rng_state`` allocate them with. +_TAIL_BOOKKEEP_ROW_BUFFERS = ( + ("temperatures", torch.float32), + ("top_ks", torch.int32), + ("top_ps", torch.float32), + ("min_ps", torch.float32), + ("seeds", torch.int64), + ("offsets", torch.int64), +) +# The tail bookkeeping items that already logged their first use. +_TAIL_BOOKKEEP_LOGGED: set = set() + + +def tail_bookkeep_log_once(item: str, message: str) -> None: + """Log a tail bookkeeping item's engagement line once per process. + + Host only, so a captured step logs it while it is captured. + """ + if item not in _TAIL_BOOKKEEP_LOGGED: + _TAIL_BOOKKEEP_LOGGED.add(item) + logger.info(f"[M3 host] tail bookkeeping {item}: {message}") + + +def tail_bookkeep_draft_len(spec_metadata, num_contexts: int, + batch_size: Optional[int]) -> int: + """The draft length when one request owns the per-token rows, else 0. + + ``K`` when this step's one request owns rows ``[0, K]`` of the per-token + sampling and Philox buffers. + + ``SpecMetadata._scan_one_model_sampling`` lays the per-token buffers out as + ``[ctx: 1 row, gen: 1 + runtime_draft_len rows]`` per request, and + ``populate_sampling_params_for_one_model`` / + ``_populate_request_rng_state`` fill a generation request's rows with its + per-request values: the same normalized temperature / top-k / top-p / min-p + as ``request_*``, its seed, and offsets ``request_offsets + r * + _RNG_SLOT_SPAN`` for row ``r``. Both layouts are kept current on every + sampled step (the Philox state is uploaded together; the parameters' + signature check only skips a refill that would write the same values). So + with tail bookkeeping on, a decode-only step of one request (no tree, no + ragged verification, the static draft length) can read its draft block's + parameters and Philox state from those rows instead of expanding the + per-request buffers. + + Host state and shapes only (no device read), so it is fixed per captured + graph. + """ + if not getattr(spec_metadata, "tail_bookkeep", False): + return 0 + if num_contexts != 0 or batch_size != 1: + return 0 + if (spec_metadata.is_spec_dec_tree or spec_metadata.is_spec_dec_dynamic_tree + or spec_metadata.is_ragged_verify): + return 0 + draft_len = spec_metadata.max_draft_len + if draft_len < 1 or spec_metadata.runtime_draft_len != draft_len: + return 0 + for name, dtype in _TAIL_BOOKKEEP_ROW_BUFFERS: + buf = getattr(spec_metadata, name, None) + if (not isinstance(buf, torch.Tensor) or buf.dim() != 1 + or buf.dtype != dtype or buf.numel() < draft_len + 1 + or not buf.is_contiguous()): + return 0 + for name in ("request_seeds", "request_offsets"): + buf = getattr(spec_metadata, name, None) + if (not isinstance(buf, torch.Tensor) or buf.dim() != 1 + or buf.dtype != torch.int64 or buf.numel() < 1): + return 0 + return draft_len + @dataclass class SpecMetadata: @@ -584,6 +684,19 @@ 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 + # Stage the per-step Philox seed/offset uploads into the engine's step-input + # pack (RuntimeHooks.pack_spec_rng_uploads) instead of copying each one. + pack_rng_uploads: bool = False + # Re-upload a drafter's batch_indices / batch-to-slot map only when they + # change (RuntimeHooks.spec_batch_map_changed_only; read by DFlash). + batch_map_changed_only: bool = False + # Drop the sampler tail's no-op bookkeeping launches, bitwise the same + # outputs (RuntimeHooks.spec_tail_bookkeep; see tail_bookkeep_draft_len). + tail_bookkeep: 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 @@ -730,6 +843,13 @@ def batch_uses_penalty(self, value: bool) -> None: # Used to scatter draft probs by slot at write time and gather them by slot # at the next iter's verify. Shape: [max_num_requests], dtype=long. batch_slot_ids: Optional[torch.Tensor] = None + # Re-upload batch_slot_ids only when the batch's slot ids change (RuntimeHooks.spec_slot_ids_changed_only). + slot_ids_changed_only: bool = False + # [buffer, ids]: the slot ids last copied and the buffer they went into. Nothing else writes batch_slot_ids, so + # equal ids into the same buffer mean equal device values. A list shared with the shallow copies of + # create_cuda_graph_metadata, which share batch_slot_ids too (same reasoning as _sampling_params_signature). + _batch_slot_ids_last: list = field(default_factory=lambda: [None, None], + repr=False) # Draft probs expanded to the target vocab size. Zero-filled once at # prepare(); each rejection iter overwrites only the d2t-selected positions # (or [:draft_vocab] when there is no d2t). @@ -820,6 +940,23 @@ def _take_rng_window_offsets(self, requests: list["LlmRequest"], offsets.append(base + step * window) return offsets + def _upload_batch_slot_ids(self, per_request_slot_ids) -> None: + """Copy the batch's slot ids into batch_slot_ids; with slot_ids_changed_only, skip the copy when the same ids + were last copied into this same buffer (nothing else writes it, so it already holds them). With the hook off + this is the plain copy, and no record is kept.""" + if self.slot_ids_changed_only: + key = tuple(per_request_slot_ids) + last = self._batch_slot_ids_last + if last[0] is self.batch_slot_ids and last[1] == key: + return + last[0], last[1] = self.batch_slot_ids, key + self.batch_slot_ids[:len(per_request_slot_ids)].copy_( + torch.tensor(per_request_slot_ids, + dtype=torch.long, + pin_memory=prefer_pinned()), + non_blocking=True, + ) + def _populate_request_rng_state( self, requests: list["LlmRequest"], per_request_normalized: list[tuple[float, int, float, float, @@ -894,11 +1031,19 @@ def _populate_request_rng_state( device='cuda') def _upload(dst: torch.Tensor, values: list[int]) -> None: - dst[:len(values)].copy_(torch.tensor(values, - dtype=torch.int64, - pin_memory=prefer_pinned()), - non_blocking=True) + src = torch.tensor(values, + dtype=torch.int64, + pin_memory=prefer_pinned()) + if self.pack_rng_uploads: + # The forward is the first reader of these buffers, so the copy + # may move to the pack's flush (a plain copy when no pack is open). + step_inputs.copy_(dst[:len(values)], src) + else: + dst[:len(values)].copy_(src, non_blocking=True) + if self.skip_greedy_rng_upload and self.is_all_greedy_sample: + # The offset window counters above still advanced. + return _upload(self.request_seeds, request_seeds) _upload(self.request_offsets, request_offsets) _upload(self.seeds, flat_seeds) @@ -1435,12 +1580,7 @@ def populate_sampling_params_for_one_model( # same table to map logits rows back to their slot, so they enable it too. if (self.use_rejection_sampling or self.enable_penalty) and self.batch_slot_ids is not None: - self.batch_slot_ids[:len(per_request_slot_ids)].copy_( - torch.tensor(per_request_slot_ids, - dtype=torch.long, - pin_memory=prefer_pinned()), - non_blocking=True, - ) + self._upload_batch_slot_ids(per_request_slot_ids) # Penalties are independent of the greedy/advanced split -- they rewrite the # logits before sampling, so an all-greedy batch is penalized too (its argmax @@ -1644,6 +1784,10 @@ class SpecWorkerBase(nn.Module, ABC): Provides common functionality for sampling and token handling. """ + # Per-worker target-installed collective; None retains the stock gather. + _draft_logits_gather_op: Optional[Callable[[torch.Tensor], + Optional[torch.Tensor]]] = None + def __init__(self, use_separate_draft_kv_cache: bool = False): super().__init__() self.guided_decoder: Optional["CapturableGuidedDecoder"] = None @@ -2262,9 +2406,25 @@ def maybe_gather_sharded_draft_logits(self, # Only plain TP (no attention DP) reaches here -- ADP variants return # early via _draft_logits_are_sharded, since their ranks are # data-parallel and must not gather draft logits across ranks. + op = self._draft_logits_gather_op + if op is not None: + gathered = op(logits) + if gathered is not None: + return gathered from ..distributed.ops import allgather return allgather(logits, self.mapping, dim=-1) + def enable_draft_logits_gather_op(self, op) -> bool: + """Install a target's all-gather of the vocab-sharded draft logits + (None, the default, keeps ``allgather(logits, mapping, dim=-1)``). + + ``op(logits)`` must return ``torch.cat(all_gather(logits), -1)`` bit + for bit, identical on every TP rank, or None for operands it does not + take (decided by shapes / dtypes only, so every rank agrees), and then + the stock all-gather runs. Collective: every rank installs the same.""" + self._draft_logits_gather_op = op + return op is not None + @staticmethod def _sample_from_logits( advanced_sampling_mode: AdvancedSamplingMode, @@ -2588,12 +2748,18 @@ def _sample_and_accept_draft_tokens_rejection( if logits.dim() == 1: logits = logits.unsqueeze(0) - accepted_tokens = torch.empty((batch_size, runtime_draft_len + 1), - dtype=torch.int, - device=device) - num_accepted_tokens = torch.ones(batch_size, - dtype=torch.int, - device=device) + # Tail bookkeeping A1: on a decode-only step the acceptance kernel's + # outputs are the whole result, so they are returned as they are (see + # below) and the result buffers are allocated only if they cannot be. + defer_outputs = (getattr(spec_metadata, "tail_bookkeep", False) + and num_contexts == 0 and num_gens > 0) + if not defer_outputs: + accepted_tokens = torch.empty((batch_size, runtime_draft_len + 1), + dtype=torch.int, + device=device) + num_accepted_tokens = torch.ones(batch_size, + dtype=torch.int, + device=device) # === Context subset: sample target's first token directly === if num_contexts > 0: @@ -2729,8 +2895,30 @@ def _sample_and_accept_draft_tokens_rejection( gen_accepted[:, 1:runtime_draft_len + 1].copy_(full_draft_tokens) - accepted_tokens[num_contexts:] = gen_accepted - num_accepted_tokens[num_contexts:] = gen_num_accepted + if defer_outputs and self._tail_bookkeep_takes_outputs( + gen_accepted, gen_num_accepted, batch_size, + runtime_draft_len, device): + # Exact: with no context rows the two copies below overwrite + # every element of same-shape, same-dtype buffers (the ones + # fill included), and the kernel's outputs are fresh tensors + # nothing else holds, so the in-place updates after this point + # (forced acceptance, the ragged clamp) see the same values. + tail_bookkeep_log_once( + "A1", "rejection outputs returned as the result (no " + "ones fill + 2 copies)") + accepted_tokens = gen_accepted + num_accepted_tokens = gen_num_accepted + else: + if defer_outputs: + accepted_tokens = torch.empty( + (batch_size, runtime_draft_len + 1), + dtype=torch.int, + device=device) + num_accepted_tokens = torch.ones(batch_size, + dtype=torch.int, + device=device) + accepted_tokens[num_contexts:] = gen_accepted + num_accepted_tokens[num_contexts:] = gen_num_accepted num_accepted_tokens = self._apply_force_accepted_tokens( num_accepted_tokens, @@ -2748,6 +2936,26 @@ def _sample_and_accept_draft_tokens_rejection( spec_metadata.verify_lens.to(num_accepted_tokens.dtype)) return accepted_tokens, num_accepted_tokens + @staticmethod + def _tail_bookkeep_takes_outputs(gen_accepted, gen_num_accepted, + batch_size: int, runtime_draft_len: int, + device: torch.device) -> bool: + """Whether the acceptance kernel's outputs can be the accept's result. + + True when they are exactly the buffers the decode-only rejection accept + would copy them into: contiguous int32 ``[batch_size, + runtime_draft_len + 1]`` tokens and ``[batch_size]`` counts on + ``device``. Shapes only. + """ + outputs = ((gen_accepted, (batch_size, runtime_draft_len + 1)), + (gen_num_accepted, (batch_size, ))) + for t, shape in outputs: + if (not isinstance(t, torch.Tensor) or t.dtype != torch.int + or tuple(t.shape) != shape or not t.is_contiguous() + or t.device != device): + return False + return True + def _rng_state_per_request( self, spec_metadata: SpecMetadata, @@ -2910,16 +3118,31 @@ def advanced_sample_draft_block(self, gen_logits: torch.Tensor, dtype=torch.int32, device=gen_logits.device) - # Take the gen slice and repeat each request's value K times to line up - # with the flattened [num_gens*K, vocab] logits (K rows per request). - temps = spec_metadata.request_temperatures[ - num_contexts:batch_size].repeat_interleave(K) - top_ks = spec_metadata.request_top_ks[ - num_contexts:batch_size].repeat_interleave(K) - top_ps = spec_metadata.request_top_ps[ - num_contexts:batch_size].repeat_interleave(K) - min_ps = spec_metadata.request_min_ps[ - num_contexts:batch_size].repeat_interleave(K) + if (num_gens == 1 and tail_bookkeep_draft_len( + spec_metadata, num_contexts, batch_size) == K): + # Tail bookkeeping B1: the one request's per-token rows already + # hold its values K times (tail_bookkeep_draft_len), so read-only + # views replace the four repeats. The rejection accept reads the + # same buffers' views. + tail_bookkeep_log_once( + "B1", "draft-block sampling params as per-token views (no " + "repeat copies)") + temps = spec_metadata.temperatures[:K] + top_ks = spec_metadata.top_ks[:K] + top_ps = spec_metadata.top_ps[:K] + min_ps = spec_metadata.min_ps[:K] + else: + # Take the gen slice and repeat each request's value K times to + # line up with the flattened [num_gens*K, vocab] logits (K rows per + # request). + temps = spec_metadata.request_temperatures[ + num_contexts:batch_size].repeat_interleave(K) + top_ks = spec_metadata.request_top_ks[ + num_contexts:batch_size].repeat_interleave(K) + top_ps = spec_metadata.request_top_ps[ + num_contexts:batch_size].repeat_interleave(K) + min_ps = spec_metadata.request_min_ps[ + num_contexts:batch_size].repeat_interleave(K) flat_logits = gen_logits.reshape(num_gens * K, vocab) # A block sampler emits all K draft positions in one launch; they take @@ -3062,7 +3285,18 @@ def sample_draft_tokens(self, logits, spec_metadata, mapping_lm_head_tp) if is_block: tokens = self.advanced_sample_draft_block( - logits, spec_metadata, num_contexts, batch_size).long() + logits, spec_metadata, num_contexts, batch_size) + if (getattr(spec_metadata, "tail_bookkeep", False) + and self._d2t is None and tokens.dtype == torch.int32 + and tokens.is_contiguous()): + # Tail bookkeeping S1: with no d2t map the int64 widening + # only feeds the int32 cast below (an identity on int32), + # which then returns the sampler's fresh tokens as is. + tail_bookkeep_log_once( + "S1", "int32 draft tokens without the int64 round " + "trip") + else: + tokens = tokens.long() else: tokens = self.advanced_sample_draft(logits, spec_metadata, diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 5fb5b185dc0b..1c81e24be184 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -27,9 +27,11 @@ import torch +from ..._utils import prefer_pinned from ...sampling_params import SamplingParams 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, @@ -165,6 +167,16 @@ 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 + # Copy the per-slot store to the host on the D2H side stream (RuntimeHooks.sampler_d2h_side_stream). + d2h_side_stream_copies: bool = False + # Keep the stored rows' slot ids in a device buffer re-uploaded only when they change + # (RuntimeHooks.spec_small_h2d). + static_slot_ids: bool = False + _slot_ids_host: Optional[list[int]] = None + _slot_ids_cuda: Optional[torch.Tensor] = None + SampleState = SampleStateSpec def is_generation_model(self) -> bool: @@ -313,6 +325,76 @@ def __init__( new_tokens_lens=int_tensor((seq_slots,)), ) + def _device_slot_ids(self, seq_slots: list[int]) -> torch.Tensor: + """The slot ids on the device, uploaded only when they differ from the last upload. + + The buffer is this sampler's own: only this method writes it, so the last uploaded + list describes its contents, and the stream orders each upload after the scatter that + read the previous ids. + """ + n = len(seq_slots) + if self._slot_ids_cuda is None or self._slot_ids_cuda.numel() < n: + self._slot_ids_cuda = torch.empty(max(n, 8), dtype=torch.long, device="cuda") + self._slot_ids_host = None + if self._slot_ids_host != list(seq_slots): + self._slot_ids_cuda[:n].copy_( + torch.tensor(seq_slots, dtype=torch.long, pin_memory=prefer_pinned()), + non_blocking=True, + ) + self._slot_ids_host = list(seq_slots) + return self._slot_ids_cuda[:n] + + def _copy_outputs_to_host( + self, + o_verify_lens: Optional[torch.Tensor], + store_copies_on_side_stream: bool, + store_copies_done: Optional[torch.cuda.Event] = None, + ) -> tuple[SampleStateTensorsSpec, Optional[torch.cuda.Event]]: + """Host copies of the step's stored outputs, and the side-stream event to await them by (or None). + + With ``d2h_side_stream_copies`` the per-slot store is copied on the D2H side stream, so the copies + overlap the next step's input staging instead of preceding it on the main stream. That is sound + only while nothing rewrites the store before the sampler event is synchronized host-side, which + holds when this sampler stored the rows itself: the store is next written by the following + ``sample_async``, after the loop has synchronized this step. ``verify_lens`` is a forward output + buffer the next replay rewrites, so it is always copied on the main stream. + + ``store_copies_done`` (a step fold's ``store_copies_done``) is recorded on the side stream after the copies: + a forward that stores into the store waits on it first, which keeps the side-stream copies sound when the + next replay rather than the next ``sample_async`` writes the store. + """ + if not ( + self.d2h_side_stream_copies + and store_copies_on_side_stream + and not self._async_worker_active() + ): + # The stock copies, in the stock order. + host = SampleStateTensorsSpec( + new_tokens=self._copy_to_host(self.store.new_tokens), + new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens), + next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens), + verify_lens=( + self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None + ), + ) + return host, None + verify_lens = self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None + with self._make_side_stream_copier() as copier: + host = SampleStateTensorsSpec( + new_tokens=copier.stage_copy_to_host(self.store.new_tokens), + new_tokens_lens=copier.stage_copy_to_host(self.store.new_tokens_lens), + next_draft_tokens=copier.stage_copy_to_host(self.store.next_draft_tokens), + verify_lens=verify_lens, + ) + if store_copies_done is not None and copier.event is not None: + store_copies_done.record(self._d2h_side_stream) + return host, copier.event + + 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, @@ -396,6 +478,57 @@ def update_requests( req.py_rewind_len = verified_len - req.py_num_accepted_draft_tokens self._request_common_handling(req, next_draft_tokens_list, runtime_draft_len) + def _store_outputs( + self, + seq_slots: list[int], + o_new_tokens: torch.Tensor, + o_next_new_tokens: torch.Tensor, + o_new_tokens_lens: torch.Tensor, + o_next_draft_tokens: torch.Tensor, + ) -> None: + """Store one step's worker outputs into the per-slot store.""" + if self.static_slot_ids: + slots = self._device_slot_ids(seq_slots) + else: + slots = torch.as_tensor(seq_slots, dtype=torch.long) + slots = slots.to(device="cuda", non_blocking=True) + step_fold = self.spec_step_fold + if step_fold is not None and step_fold.scatter( + slots, 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. + new_tokens_width = self.store.new_tokens.shape[0] + next_new_tokens_width = self.store.next_new_tokens.shape[0] + draft_tokens_width = self.store.next_draft_tokens.shape[1] + if o_new_tokens.shape[1] < new_tokens_width: + o_new_tokens = torch.nn.functional.pad( + o_new_tokens, (0, new_tokens_width - o_new_tokens.shape[1]) + ) + elif o_new_tokens.shape[1] > new_tokens_width: + o_new_tokens = o_new_tokens[:, :new_tokens_width] + if o_next_draft_tokens.shape[1] < draft_tokens_width: + o_next_draft_tokens = torch.nn.functional.pad( + o_next_draft_tokens, (0, draft_tokens_width - o_next_draft_tokens.shape[1]) + ) + elif o_next_draft_tokens.shape[1] > draft_tokens_width: + o_next_draft_tokens = o_next_draft_tokens[:, :draft_tokens_width] + if o_next_new_tokens.shape[1] < next_new_tokens_width: + o_next_new_tokens = torch.nn.functional.pad( + o_next_new_tokens, (0, next_new_tokens_width - o_next_new_tokens.shape[1]) + ) + elif o_next_new_tokens.shape[1] > next_new_tokens_width: + o_next_new_tokens = o_next_new_tokens[:, :next_new_tokens_width] + + # Use index_copy_ for efficient copying (slots are unique) + self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens) + self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens) + self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens) + self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens) + def sample_async( self, scheduled_requests: ScheduledRequests, @@ -442,8 +575,7 @@ def sample_async( for r in sampling_requests ] - slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long) - slots = slots.to(device="cuda", non_blocking=True) + seq_slots = [r.py_seq_slot for r in sampling_requests] o_new_tokens = outputs["new_tokens"][num_skip : num_skip + num_sampling_requests] o_new_tokens_lens = outputs["new_tokens_lens"][num_skip : num_skip + num_sampling_requests] @@ -489,36 +621,14 @@ def sample_async( raise ValueError("verify_lens must contain one int32 token window per output row") o_verify_lens = o_verify_lens[num_skip : num_skip + num_sampling_requests] - # 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. - new_tokens_width = self.store.new_tokens.shape[0] - next_new_tokens_width = self.store.next_new_tokens.shape[0] - draft_tokens_width = self.store.next_draft_tokens.shape[1] - if o_new_tokens.shape[1] < new_tokens_width: - o_new_tokens = torch.nn.functional.pad( - o_new_tokens, (0, new_tokens_width - o_new_tokens.shape[1]) - ) - elif o_new_tokens.shape[1] > new_tokens_width: - o_new_tokens = o_new_tokens[:, :new_tokens_width] - if o_next_draft_tokens.shape[1] < draft_tokens_width: - o_next_draft_tokens = torch.nn.functional.pad( - o_next_draft_tokens, (0, draft_tokens_width - o_next_draft_tokens.shape[1]) - ) - elif o_next_draft_tokens.shape[1] > draft_tokens_width: - o_next_draft_tokens = o_next_draft_tokens[:, :draft_tokens_width] - if o_next_new_tokens.shape[1] < next_new_tokens_width: - o_next_new_tokens = torch.nn.functional.pad( - o_next_new_tokens, (0, next_new_tokens_width - o_next_new_tokens.shape[1]) + step_fold = self.spec_step_fold + stored_here = step_fold is None or not step_fold.sampler_rows_stored( + outputs, num_skip, seq_slots + ) + if stored_here: + self._store_outputs( + seq_slots, o_new_tokens, o_next_new_tokens, o_new_tokens_lens, o_next_draft_tokens ) - elif o_next_new_tokens.shape[1] > next_new_tokens_width: - o_next_new_tokens = o_next_new_tokens[:, :next_new_tokens_width] - - # Use index_copy_ for efficient copying (slots are unique) - self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens) - self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens) - self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens) - self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens) # Create sample state with async D2H copy device_tensors = SampleStateTensorsSpec( @@ -527,13 +637,20 @@ def sample_async( next_draft_tokens=self.store.next_draft_tokens, ) - host_tensors = SampleStateTensorsSpec( - new_tokens=self._copy_to_host(self.store.new_tokens), - new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens), - next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens), - verify_lens=(self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None), + # A fold whose forward waits for the store copies (store_copies_done) lets them go on the side stream + # on the steps its forward stored the rows too. + store_copies_done = getattr(step_fold, "store_copies_done", None) if step_fold else None + host_tensors, side_stream_event = self._copy_outputs_to_host( + o_verify_lens, + store_copies_on_side_stream=stored_here or store_copies_done is not None, + store_copies_done=store_copies_done, + ) + # Without side-stream copies the call is the stock one (no argument). + sampler_event = ( + self._record_sampler_event(side_stream_event) + if side_stream_event is not None + else self._record_sampler_event() ) - sampler_event = self._record_sampler_event() # Add dummy draft tokens to context requests for KV cache preparation for request in finished_context_requests: 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..fb2360a86083 --- /dev/null +++ b/tests/unittest/_torch/executor/test_h2d_cache.py @@ -0,0 +1,580 @@ +# 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() + + +# --------------------------------------------------------------------------- +# changed_span (RuntimeHooks.h2d_changed_span): a changed packed 2-D upload stages only each row's changed span +# --------------------------------------------------------------------------- + +ROWS, WIDTH = 4, 1024 + + +def _span_pack(): + from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack + + return StepInputPack(torch.device("cuda"), parallel_scatter=True) + + +def _table(seed, rows=ROWS, width=WIDTH): + generator = torch.Generator().manual_seed(seed) + return torch.randint( + 0, 1 << 20, (rows, width), generator=generator, dtype=torch.int32 + ).pin_memory() + + +def _upload(pack, span_cache, dst, src): + """One step: open the pack, route the copy through the cache, issue the pack; returns the staged copies.""" + from tensorrt_llm._torch.pyexecutor import step_inputs + + with h2d_cache.active(span_cache), step_inputs.active(pack): + step_inputs.begin() + copied = copy_h2d_if_changed(dst, src, pack=True) + staged = list(pack.h2d) + step_inputs.flush() + torch.cuda.synchronize() + return copied, staged + + +def test_changed_span_is_off_by_default_and_without_init(): + assert not H2DSkipCache().changed_span + assert not object.__new__(H2DSkipCache).changed_span + + +@pytest.mark.parametrize("changed_span", [False, True]) +def test_changed_span_device_rows_equal_the_source_after_every_step(changed_span): + """Bitwise: whatever is staged, the device view holds the source after each step.""" + pack = _span_pack() + span_cache = H2DSkipCache(changed_span=changed_span) + buf = torch.full((8, 2 * WIDTH), -7, dtype=torch.int32, device="cuda") + dst = buf[:ROWS, :WIDTH] + src = _table(0) + generator = torch.Generator().manual_seed(1) + for step in range(40): + src = src.clone().pin_memory() + # A request crosses into a new page: one entry of one row changes (sometimes several, in several rows). + for _ in range(1 + step % 3): + row = int(torch.randint(0, ROWS, (1,), generator=generator)) + col = int(torch.randint(0, WIDTH, (1,), generator=generator)) + src[row, col] = step + _upload(pack, span_cache, dst, src) + assert torch.equal(dst.cpu(), src) + # Rows and columns outside the view were never written. + assert bool((buf[ROWS:] == -7).all()) and bool((buf[:, WIDTH:] == -7).all()) + assert (span_cache.span_uploads > 0) is changed_span + + +def test_changed_span_one_changed_entry_stages_one_element(): + pack = _span_pack() + span_cache = H2DSkipCache(changed_span=True) + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(2) + _upload(pack, span_cache, dst, src) + + src = src.clone().pin_memory() + src[2, 517] += 1 + copied, staged = _upload(pack, span_cache, dst, src) + + assert copied is True + assert staged == [(dst[2, 517].data_ptr(), staged[0][1], 4)] + assert torch.equal(dst.cpu(), src) + assert (span_cache.span_uploads, span_cache.span_bytes) == (1, 4) + assert span_cache.span_full_bytes == ROWS * WIDTH * 4 + + +def test_changed_span_each_changed_row_stages_its_first_to_last_changed_column(): + pack = _span_pack() + span_cache = H2DSkipCache(changed_span=True) + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(3) + _upload(pack, span_cache, dst, src) + + src = src.clone().pin_memory() + src[0, 10] += 1 + src[0, 20] += 1 + src[3, 1000] += 1 + _, staged = _upload(pack, span_cache, dst, src) + + assert [(ptr, nbytes) for ptr, _, nbytes in staged] == [ + (dst[0, 10].data_ptr(), 11 * 4), + (dst[3, 1000].data_ptr(), 4), + ] + assert torch.equal(dst.cpu(), src) + + +def test_changed_span_empty_span_stages_nothing_and_keeps_the_device_values(): + """PR-CTX (d): an empty span. A narrower view of the same rows whose values all equal the shadow's: no column + changed, so nothing is staged (the copy still counts as issued, and the shadow follows the new view).""" + pack = _span_pack() + span_cache = H2DSkipCache(changed_span=True) + buf = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(10) + _upload(pack, span_cache, buf, src) + + narrow = src[:, : WIDTH // 2].contiguous().pin_memory() + copied, staged = _upload(pack, span_cache, buf[:, : WIDTH // 2], narrow) + + assert copied is True + assert staged == [] + assert torch.equal(buf.cpu(), src) + assert (span_cache.span_uploads, span_cache.span_bytes) == (1, 0) + # The shadow is the narrow view now: the same values again are skipped. + copied, staged = _upload(pack, span_cache, buf[:, : WIDTH // 2], narrow) + assert copied is False and staged == [] + + +def test_changed_span_whole_buffer_span_is_staged_whole(): + """PR-CTX (d): a whole-buffer span. The first and last column of every row changed, so every row's span is the + whole row: over the fraction, the view is staged whole (one descriptor), not as per-row spans.""" + pack = _span_pack() + span_cache = H2DSkipCache(changed_span=True) + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(11) + _upload(pack, span_cache, dst, src) + + src = src.clone().pin_memory() + src[:, 0] += 1 + src[:, -1] += 1 + _, staged = _upload(pack, span_cache, dst, src) + + assert [nbytes for _, _, nbytes in staged] == [ROWS * WIDTH * 4] + assert torch.equal(dst.cpu(), src) + assert span_cache.span_uploads == 0 + + +def test_changed_span_at_the_fraction_is_still_a_span(): + pack = _span_pack() + span_cache = H2DSkipCache(changed_span=True) + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(12) + _upload(pack, span_cache, dst, src) + + src = src.clone().pin_memory() + half = int(H2DSkipCache.CHANGED_SPAN_MAX_FRACTION * WIDTH) + src[:, :half] += 1 # exactly the fraction of the view's elements + _, staged = _upload(pack, span_cache, dst, src) + + assert [nbytes for _, _, nbytes in staged] == [half * 4] * ROWS + assert torch.equal(dst.cpu(), src) + + +def test_changed_span_wider_view_uploads_the_columns_the_shadow_does_not_cover(): + pack = _span_pack() + span_cache = H2DSkipCache(changed_span=True) + buf = torch.full((ROWS, WIDTH), -1, dtype=torch.int32, device="cuda") + src = _table(4) + _upload(pack, span_cache, buf[:, : WIDTH // 2], src[:, : WIDTH // 2].contiguous().pin_memory()) + + _, staged = _upload( + pack, + span_cache, + buf[:, : WIDTH // 2 + 3], + src[:, : WIDTH // 2 + 3].contiguous().pin_memory(), + ) + + assert [nbytes for _, _, nbytes in staged] == [3 * 4] * ROWS + assert torch.equal(buf[:, : WIDTH // 2 + 3].cpu(), src[:, : WIDTH // 2 + 3]) + + +def test_changed_span_a_mostly_changed_table_is_staged_whole(): + pack = _span_pack() + span_cache = H2DSkipCache(changed_span=True) + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + _upload(pack, span_cache, dst, _table(5)) + + src = _table(6) + _, staged = _upload(pack, span_cache, dst, src) + + assert [nbytes for _, _, nbytes in staged] == [ROWS * WIDTH * 4] + assert torch.equal(dst.cpu(), src) + + +def test_changed_span_off_by_default_the_whole_view_is_staged(): + pack = _span_pack() + span_cache = H2DSkipCache() + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(7) + _upload(pack, span_cache, dst, src) + + src = src.clone().pin_memory() + src[1, 1] += 1 + _, staged = _upload(pack, span_cache, dst, src) + + assert [nbytes for _, _, nbytes in staged] == [ROWS * WIDTH * 4] + + +def test_changed_span_an_unpacked_copy_copies_the_whole_view(): + span_cache = H2DSkipCache(changed_span=True) + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(8) + with h2d_cache.active(span_cache): + copy_h2d_if_changed(dst, src) + src = src.clone().pin_memory() + src[1, 1] += 1 + assert copy_h2d_if_changed(dst, src) is True + torch.cuda.synchronize() + assert torch.equal(dst.cpu(), src) + assert span_cache.span_uploads == 0 + + +def test_changed_span_a_view_with_another_row_stride_is_staged_whole(): + """The shadow describes other device bytes: nothing of it may be trusted.""" + pack = _span_pack() + span_cache = H2DSkipCache(changed_span=True) + buf = torch.zeros((2 * ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(9) + _upload(pack, span_cache, buf[:ROWS], src) + + _, staged = _upload(pack, span_cache, buf[::2], src) + + assert [nbytes for _, _, nbytes in staged] == [WIDTH * 4] * ROWS + assert torch.equal(buf[::2].cpu(), src) + + +@pytest.mark.parametrize("changed_span", [False, True]) +def test_changed_span_keeps_a_page_crossing_step_in_the_graph_window(changed_span): + """A 16 KB page table (over the in-graph window) gains one page: with changed_span the step's pack still fits the + window and the replay's prologue issues it; without, the step falls back to an eager upload and scatter.""" + from tensorrt_llm._torch.pyexecutor import step_inputs + from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack + + pack = StepInputPack(torch.device("cuda"), parallel_scatter=True, in_graph=True) + span_cache = H2DSkipCache(changed_span=changed_span) + table = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + out = torch.zeros(ROWS, dtype=torch.int32, device="cuda") + graph = torch.cuda.CUDAGraph() + with step_inputs.active(pack): + torch.cuda.synchronize() + with torch.cuda.graph(graph): + assert step_inputs.capture_prologue() is pack + out.copy_(table[:, -1]) + src = _table(13) + table.copy_(src) + with h2d_cache.active(span_cache): + copy_h2d_if_changed(table, src) # the shadow of the table the device holds + torch.cuda.synchronize() + src = src.clone().pin_memory() + src[1, -1] = 4242 + with h2d_cache.active(span_cache), step_inputs.active(pack): + step_inputs.begin() + copy_h2d_if_changed(table, src, pack=True) + step_inputs.flush(in_graph=True) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(table.cpu(), src) + assert out.tolist() == src[:, -1].tolist() + expect = (1, 0) if changed_span else (0, 1) + assert (pack.graph_windows, pack.eager_fallbacks) == expect diff --git a/tests/unittest/_torch/executor/test_instantiate_sampler_hooks.py b/tests/unittest/_torch/executor/test_instantiate_sampler_hooks.py new file mode 100644 index 000000000000..8c87369c9eca --- /dev/null +++ b/tests/unittest/_torch/executor/test_instantiate_sampler_hooks.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""instantiate_sampler hands a one-model speculative sampler the model's host-tail hooks +(RuntimeHooks.sampler_d2h_side_stream / spec_small_h2d) and logs the D2H side-stream engagement line only when it +is on.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import _util +from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks +from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler +from tensorrt_llm.mapping import Mapping + +pytestmark = pytest.mark.cpu_only + +_D2H_LINE = "[M3 host] spec sampler store copies on the D2H side stream" + + +class _Fold: + """A SpecStepFold stand-in; binding hands it the sampler's store.""" + + store = None + + def bind_sampler_store(self, store): + self.store = store + + +def _instantiate(monkeypatch, hooks, sampler): + lines = [] + monkeypatch.setattr(_util, "confidential_compute_enabled", lambda: False) + monkeypatch.setattr(_util, "resolve_max_num_sequences", lambda *args, **kwargs: 4) + monkeypatch.setattr(_util, "get_spec_decoder", lambda sampler_args, spec_config: sampler) + monkeypatch.setattr(_util.logger, "info", lines.append) + + class _Model(torch.nn.Module): + runtime_hooks = hooks + + engine = SimpleNamespace( + max_seq_len=64, + model=_Model(), + spec_config=SimpleNamespace( + spec_dec_mode=SimpleNamespace(has_spec_decoder=lambda: True), max_draft_len=3 + ), + ) + llm_args = SimpleNamespace( + sampler_force_async_worker=False, + disable_overlap_scheduler=False, + enable_speculative_beam_history_d2h=False, + ) + result = _util.instantiate_sampler( + engine, + llm_args, + Mapping(), + max_batch_size=4, + max_beam_width=1, + mm_encoder_only=False, + speculative_config=None, + max_num_sequences=4, + ) + assert result is sampler + return lines + + +def _spec_sampler(): + sampler = object.__new__(SpecSampler) + sampler.store = object() + return sampler + + +def test_default_hooks_leave_the_sampler_stock_and_log_nothing(monkeypatch): + sampler = _spec_sampler() + + lines = _instantiate(monkeypatch, RuntimeHooks(), sampler) + + assert sampler.d2h_side_stream_copies is False + assert sampler.static_slot_ids is False + assert sampler.spec_step_fold is None + assert lines == [] + + +@pytest.mark.parametrize( + "side, static", + [(True, False), (False, True), (True, True)], + ids=["d2h_side_stream", "spec_small_h2d", "both"], +) +def test_the_hooks_reach_the_sampler(monkeypatch, side, static): + sampler = _spec_sampler() + fold = _Fold() + hooks = RuntimeHooks(sampler_d2h_side_stream=side, spec_small_h2d=static, spec_step_fold=fold) + + lines = _instantiate(monkeypatch, hooks, sampler) + + assert sampler.d2h_side_stream_copies is side + assert sampler.static_slot_ids is static + assert sampler.spec_step_fold is fold and fold.store is sampler.store + # spec_small_h2d's line is DecoderMetadata.setup_speculative's, not this one. + assert lines == ([_D2H_LINE] if side else []) + + +def test_a_sampler_that_is_not_a_spec_sampler_is_left_alone(monkeypatch): + other = SimpleNamespace() + + lines = _instantiate( + monkeypatch, RuntimeHooks(sampler_d2h_side_stream=True, spec_small_h2d=True), other + ) + + assert not hasattr(other, "d2h_side_stream_copies") + assert not hasattr(other, "static_slot_ids") + assert lines == [] diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_block_offset_guard.py b/tests/unittest/_torch/executor/test_kv_cache_v2_block_offset_guard.py new file mode 100644 index 000000000000..581cc255684a --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_block_offset_guard.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KVCacheManagerV2's opt-in unchanged block-table skip (``skip_unchanged_block_offsets``) and the RuntimeHooks field +that arms it on a separate draft KV cache manager (``skip_unchanged_draft_block_offsets``). + +The guard is driven on a manager built without ``__init__`` (every attribute it reads is declared on the class or set +here). The CUDA test checks the bytes: a guarded sequence of copies (the in-tree copy op issued only when the guard +does not skip) leaves the destination equal to always copying, after every step of a random host-table history. +""" + +import random +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks + +POOLS, CAPACITY, WIDTH = 2, 6, 16 + + +def _manager(pin: bool = False) -> KVCacheManagerV2: + manager = object.__new__(KVCacheManagerV2) + manager.host_kv_cache_block_offsets = torch.zeros( + POOLS, CAPACITY, 2, WIDTH, dtype=torch.int32, pin_memory=pin + ) + manager.host_kv_cache_block_offsets[:, :, 0, :] = torch.arange( + CAPACITY * WIDTH, dtype=torch.int32 + ).view(CAPACITY, WIDTH) + manager.skip_unchanged_block_offsets = True + return manager + + +def _rows(*rows: int) -> torch.Tensor: + return torch.tensor(rows, dtype=torch.int32) + + +@pytest.mark.cpu_only +def test_the_guard_is_off_and_unarmed_by_default(): + assert KVCacheManagerV2.skip_unchanged_block_offsets is False + assert RuntimeHooks().skip_unchanged_draft_block_offsets is False + assert not RuntimeHooks().skips_unchanged_metadata_uploads + assert RuntimeHooks(skip_unchanged_draft_block_offsets=True).skips_unchanged_metadata_uploads + + +@pytest.mark.cpu_only +def test_an_unchanged_decode_copy_is_skipped(): + manager, dst = _manager(), torch.zeros(POOLS, CAPACITY, 2, WIDTH, dtype=torch.int32) + + assert not manager._block_offsets_unchanged(dst, _rows(3), 0) + assert manager._block_offsets_unchanged(dst, _rows(3), 0) + assert manager._block_offsets_unchanged(dst, _rows(3), 0) + assert (manager.block_offset_copies, manager.block_offset_skips) == (1, 2) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("pool", range(POOLS)) +def test_a_changed_page_index_in_a_copied_row_copies(pool): + manager, dst = _manager(), torch.zeros(POOLS, CAPACITY, 2, WIDTH, dtype=torch.int32) + assert not manager._block_offsets_unchanged(dst, _rows(1, 3), 0) + + manager.host_kv_cache_block_offsets[pool, 3, 0, WIDTH - 1] += 1 + assert not manager._block_offsets_unchanged(dst, _rows(1, 3), 0) + assert manager._block_offsets_unchanged(dst, _rows(1, 3), 0) + + +@pytest.mark.cpu_only +def test_field_one_and_rows_not_copied_do_not_force_a_copy(): + """The copy reads field 0 of the copied rows only.""" + manager, dst = _manager(), torch.zeros(POOLS, CAPACITY, 2, WIDTH, dtype=torch.int32) + assert not manager._block_offsets_unchanged(dst, _rows(2), 0) + + manager.host_kv_cache_block_offsets[:, 2, 1, :] = 7 + manager.host_kv_cache_block_offsets[:, 4, 0, :] = 9 + assert manager._block_offsets_unchanged(dst, _rows(2), 0) + + +@pytest.mark.cpu_only +def test_another_row_set_or_order_copies(): + manager, dst = _manager(), torch.zeros(POOLS, CAPACITY, 2, WIDTH, dtype=torch.int32) + assert not manager._block_offsets_unchanged(dst, _rows(1, 3), 0) + + assert not manager._block_offsets_unchanged(dst, _rows(3, 1), 0) + assert not manager._block_offsets_unchanged(dst, _rows(3), 0) + assert manager._block_offsets_unchanged(dst, _rows(3), 0) + + +@pytest.mark.cpu_only +def test_a_step_with_context_requests_copies_and_drops_the_record(): + manager, dst = _manager(), torch.zeros(POOLS, CAPACITY, 2, WIDTH, dtype=torch.int32) + assert not manager._block_offsets_unchanged(dst, _rows(3), 0) + + assert not manager._block_offsets_unchanged(dst, _rows(3), 1) + assert not manager._block_offsets_unchanged(dst, _rows(3), 0) + assert manager._block_offsets_unchanged(dst, _rows(3), 0) + + +@pytest.mark.cpu_only +def test_another_view_of_the_same_storage_copies_and_replaces_the_record(): + manager = _manager() + storage = torch.zeros(POOLS, CAPACITY + 2, 2, WIDTH, dtype=torch.int32) + whole, prefix = storage, storage[:, :CAPACITY] + assert not manager._block_offsets_unchanged(whole, _rows(3), 0) + + assert not manager._block_offsets_unchanged(prefix, _rows(3), 0) + assert not manager._block_offsets_unchanged(whole, _rows(3), 0) + assert manager._block_offsets_unchanged(whole, _rows(3), 0) + + +@pytest.mark.cpu_only +def test_destinations_in_different_storages_keep_their_own_records(): + manager = _manager() + first, second = (torch.zeros(POOLS, CAPACITY, 2, WIDTH, dtype=torch.int32) for _ in range(2)) + assert not manager._block_offsets_unchanged(first, _rows(3), 0) + assert not manager._block_offsets_unchanged(second, _rows(3), 0) + + assert manager._block_offsets_unchanged(first, _rows(3), 0) + assert manager._block_offsets_unchanged(second, _rows(3), 0) + + +@pytest.mark.cpu_only +def test_the_counts_are_logged_at_powers_of_two(monkeypatch): + lines = [] + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2.logger.info", lines.append + ) + manager, dst = _manager(), torch.zeros(POOLS, CAPACITY, 2, WIDTH, dtype=torch.int32) + for _ in range(9): + manager._block_offsets_unchanged(dst, _rows(3), 0) + + assert lines == [ + f"[M3 host] block-offset copies skipped when unchanged: {skips} of {calls}" + for skips, calls in ((0, 1), (1, 2), (3, 4), (7, 8)) + ] + + +@pytest.mark.cpu_only +def test_the_hook_arms_only_a_separate_v2_draft_manager(monkeypatch): + lines = [] + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner.logger.info", lines.append + ) + arm = DecoderRunner._arm_draft_block_offset_guard + draft = object.__new__(KVCacheManagerV2) + + arm(SimpleNamespace(runtime_hooks=RuntimeHooks()), draft) + assert draft.skip_unchanged_block_offsets is False + + armed = SimpleNamespace(runtime_hooks=RuntimeHooks(skip_unchanged_draft_block_offsets=True)) + arm(armed, None) + arm(armed, SimpleNamespace(skip_unchanged_block_offsets=False)) + assert lines == [] + + arm(armed, draft) + arm(armed, draft) + assert draft.skip_unchanged_block_offsets is True + assert KVCacheManagerV2.skip_unchanged_block_offsets is False + assert lines == ["[M3 host] draft block-offset copy skipped when unchanged"] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="runs the in-tree block-offset copy op") +@pytest.mark.parametrize("seed", range(4)) +def test_guarded_copies_leave_the_bytes_of_always_copying(seed): + from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( + copy_batch_block_offsets_to_device, + ) + + rng = random.Random(seed) + manager = _manager(pin=True) + host = manager.host_kv_cache_block_offsets + scales = torch.tensor([2, 3], dtype=torch.int32, pin_memory=True) + offsets = torch.tensor([1, 5], dtype=torch.int32, pin_memory=True) + guarded = torch.full((POOLS, CAPACITY, 2, WIDTH), -1, dtype=torch.int32, device="cuda") + reference = guarded.clone() + stream = torch.cuda.current_stream().cuda_stream + rows = [1, 3] + skips = 0 + for step in range(200): + event = rng.random() + if event < 0.15: + host[rng.randrange(POOLS), rng.choice(rows), 0, rng.randrange(WIDTH)] = rng.randrange( + 1000 + ) + elif event < 0.25: + host[rng.randrange(POOLS), rng.randrange(CAPACITY), 1, rng.randrange(WIDTH)] = ( + rng.randrange(1000) + ) + elif event < 0.30: + rows = rng.sample(range(CAPACITY), rng.randint(1, 3)) + num_contexts = 1 if rng.random() < 0.05 else 0 + copy_idx = torch.tensor(rows, dtype=torch.int32, pin_memory=True) + copy_batch_block_offsets_to_device(host, reference, copy_idx, scales, offsets, stream) + if manager._block_offsets_unchanged(guarded, copy_idx, num_contexts): + skips += 1 + else: + copy_batch_block_offsets_to_device(host, guarded, copy_idx, scales, offsets, stream) + torch.cuda.synchronize() + assert torch.equal(guarded, reference), f"seed {seed} step {step}" + assert skips > 50 diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py new file mode 100644 index 000000000000..2e392661d69f --- /dev/null +++ b/tests/unittest/_torch/executor/test_runtime_hooks.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""How the engine reads the step-loop hooks a model declares (``runtime_hooks_of``).""" + +import dataclasses + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.runtime_hooks import ( + DEFAULT_RUNTIME_HOOKS, + MsaRuntimePolicy, + RuntimeHooks, + runtime_hooks_of, +) + +pytestmark = pytest.mark.cpu_only + + +class _PlainModel(torch.nn.Module): + """A model that declares no hooks, like every model outside the opt-in.""" + + +class _ClassHooksModel(torch.nn.Module): + runtime_hooks = RuntimeHooks(pack_step_inputs=True, skip_unchanged_h2d=True) + + +def test_model_without_hooks_gets_the_all_off_default(): + hooks = runtime_hooks_of(_PlainModel()) + + 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 not hooks.pack_spec_rng_uploads + assert hooks.spec_step_fold is None + assert not hooks.sampler_d2h_side_stream + assert not hooks.spec_small_h2d + assert not hooks.skip_accepted_draft_upload + assert not hooks.spec_batch_map_changed_only + assert not hooks.kv_lens_update_once + assert not hooks.pack_parallel_scatter + assert not hooks.pack_in_graph + assert not hooks.pack_early_copy + assert not hooks.h2d_changed_span + 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_hooks_declared_as_none_mean_the_default(): + model = _PlainModel() + model.runtime_hooks = None + + assert runtime_hooks_of(model) is DEFAULT_RUNTIME_HOOKS + + +def test_class_attribute_hooks_are_returned_as_declared(): + model = _ClassHooksModel() + + assert runtime_hooks_of(model) is _ClassHooksModel.runtime_hooks + + +def test_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 = _PlainModel() + model.runtime_hooks = hooks + + assert runtime_hooks_of(model) is hooks + + +def test_instance_attribute_overrides_class_attribute(): + hooks = RuntimeHooks(skip_greedy_rng_upload=True) + model = _ClassHooksModel() + 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_hooks_of_the_wrong_type_raise(declared): + model = _PlainModel() + model.runtime_hooks = declared + + with pytest.raises(TypeError, match=r"_PlainModel\.runtime_hooks must be a RuntimeHooks"): + runtime_hooks_of(model) + + +def test_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_every_metadata_upload_skip_asks_for_private_metadata_buffers(hooks, skips): + assert hooks.skips_unchanged_metadata_uploads is skips + + +HOSTTAIL_FIELDS = ( + "sampler_d2h_side_stream", + "spec_small_h2d", + "pack_spec_rng_uploads", + "skip_accepted_draft_upload", + "spec_batch_map_changed_only", + "kv_lens_update_once", +) + + +@pytest.mark.parametrize("field", HOSTTAIL_FIELDS) +def test_each_host_tail_hook_defaults_off_and_is_a_plain_bool(field): + default = next(f for f in dataclasses.fields(RuntimeHooks) if f.name == field) + + assert default.default is False + assert getattr(RuntimeHooks(), field) is False + assert getattr(RuntimeHooks(**{field: True}), field) is True + + +def test_the_default_hooks_still_equal_a_fresh_instance(): + assert RuntimeHooks() == DEFAULT_RUNTIME_HOOKS + assert RuntimeHooks(**{field: False for field in HOSTTAIL_FIELDS}) == DEFAULT_RUNTIME_HOOKS + + +@pytest.mark.parametrize("field", HOSTTAIL_FIELDS) +def test_no_host_tail_hook_asks_for_private_metadata_buffers(field): + """skips_unchanged_metadata_uploads keys only on the metadata-upload skips; the host-tail hooks leave it alone.""" + assert RuntimeHooks(**{field: True}).skips_unchanged_metadata_uploads is False + assert ( + RuntimeHooks(skip_unchanged_h2d=True, **{field: True}).skips_unchanged_metadata_uploads + is True + ) + assert ( + RuntimeHooks(**{name: True for name in HOSTTAIL_FIELDS}).skips_unchanged_metadata_uploads + is False + ) 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..ceb8e36847e2 --- /dev/null +++ b/tests/unittest/_torch/executor/test_step_inputs.py @@ -0,0 +1,1202 @@ +# 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 + + +def _random_pack_run(parallel: bool, seed: int, overlap: bool): + """Stage a fixed random mix of host copies (2-D page-table rows included) and device copies that read a host + copy's destination; return every destination after the flush.""" + g = torch.Generator().manual_seed(seed) + pack = StepInputPack(_device(), parallel_scatter=parallel) + big = _zeros(4096) + table = torch.zeros((8, 64), dtype=torch.int32, device=_device()) + out = _zeros(256) + with step_inputs.active(pack): + step_inputs.begin() + for k in range(24): + n = int(torch.randint(1, 64, (1,), generator=g)) + step_inputs.copy_( + big[k * 128 : k * 128 + n], + torch.randint(0, 1 << 30, (n,), generator=g, dtype=torch.int32), + ) + step_inputs.copy_( + table[:3, :40], torch.randint(0, 1 << 30, (3, 40), generator=g, dtype=torch.int32) + ) + if overlap: + # A later host copy overwrites part of an earlier one: order matters, so the serial scatter must run. + step_inputs.copy_(big[0:4], torch.full((4,), -7, dtype=torch.int32)) + step_inputs.copy_(out[:64], big[:64]) + step_inputs.copy_(out[64:96], big[128:160]) + step_inputs.flush() + torch.cuda.synchronize() + return big.cpu(), table.cpu(), out.cpu() + + +@pytest.mark.parametrize("overlap", [False, True], ids=["disjoint", "overlapping"]) +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_parallel_scatter_writes_what_the_serial_scatter_writes(seed, overlap): + serial = _random_pack_run(False, seed, overlap) + parallel = _random_pack_run(True, seed, overlap) + for a, b in zip(serial, parallel): + assert torch.equal(a, b) + + +def test_parallel_scatter_only_for_order_free_packs(): + assert step_inputs._parallel_ok([(0, 0, 8), (8, 0, 8)], []) + assert not step_inputs._parallel_ok([(0, 0, 8), (4, 0, 8)], []) + # A device copy may read a host copy's destination (the launches run in order) ... + assert step_inputs._parallel_ok([(0, 0, 8)], [(100, 0, 8)]) + # ... but not another device copy's destination, and device destinations must be disjoint. + assert not step_inputs._parallel_ok([], [(100, 0, 8), (200, 100, 8)]) + assert not step_inputs._parallel_ok([], [(100, 0, 8), (104, 300, 8)]) + + +def test_parallel_scatter_is_off_by_default(): + assert not StepInputPack(_device()).parallel_scatter + + +def _parallel_scatter_compiles() -> int: + """How many variants of the parallel scatter kernel Triton holds compiled (any device).""" + kernel = step_inputs._parallel_scatter_kernel + caches = getattr(kernel, "device_caches", None) + if caches is not None: + return sum(len(entry[0]) for entry in caches.values()) + return sum(len(entry) for entry in kernel.cache.values()) + + +def test_parallel_scatter_compiles_nothing_after_the_warm_up(): + # The warm-up compiles both DEVICE_SRC variants; packs with any number of host copies (which moves the device + # copies' offset `first`) must reuse them, so no serving-step flush compiles. + pack = StepInputPack(_device(), parallel_scatter=True) + compiled = _parallel_scatter_compiles() + dst = _zeros(4096) + out = _zeros(64) + with step_inputs.active(pack): + for n_host in (1, 2, 3, 5, 15, 16, 17, 24): + step_inputs.begin() + for k in range(n_host): + step_inputs.copy_(dst[k * 64 : k * 64 + 4], torch.full((4,), k, dtype=torch.int32)) + step_inputs.copy_(out[:4], dst[:4]) + step_inputs.flush() + torch.cuda.synchronize() + assert _parallel_scatter_compiles() == compiled + assert out[:4].tolist() == [0] * 4 + assert dst[23 * 64 : 23 * 64 + 4].tolist() == [23] * 4 + + +# --------------------------------------------------------------------------- +# Flags off: the GB path, byte for byte +# --------------------------------------------------------------------------- + + +def _gb_flush(pack): + """StepInputPack.flush as of the GB base (1af1e3272c), verbatim: the reference the flags-off path must match.""" + if not pack.open: + return + pack.open = False + if not pack.h2d and not pack.d2d: + return + host = pack.host[pack.parity] + desc = [len(pack.h2d), len(pack.d2d)] + [ + value for copy in pack.h2d + pack.d2d for value in copy + ] + host[: 8 * len(desc)].view(torch.int64).copy_(torch.tensor(desc, dtype=torch.int64)) + pack.dev[: pack.used].copy_(host[: pack.used], non_blocking=True) + if pack.parallel_scatter and step_inputs._parallel_ok(pack.h2d, pack.d2d): + stage_i64 = pack.dev.view(torch.int64) + if pack.h2d: + step_inputs._parallel_scatter_kernel[(len(pack.h2d),)]( + pack.dev, stage_i64, 0, DEVICE_SRC=False, BLOCK=256, num_warps=2 + ) + if pack.d2d: + step_inputs._parallel_scatter_kernel[(len(pack.d2d),)]( + pack.dev, stage_i64, len(pack.h2d), DEVICE_SRC=True, BLOCK=256, num_warps=2 + ) + else: + torch.ops.trtllm.step_inputs_scatter(pack.dev) + event = pack.events[pack.parity] + if event is None: + event = pack.events[pack.parity] = torch.cuda.Event() + event.record() + pack.parity = (pack.parity + 1) % pack.RING + pack._clear() + + +def _flags_off_run(parallel: bool, seed: int, overlap: bool, how: str): + """``_random_pack_run`` issued by ``flush()``, ``flush(in_graph=True)`` or the GB reference; returns every + destination and the uploaded staging bytes.""" + g = torch.Generator().manual_seed(seed) + pack = StepInputPack(_device(), parallel_scatter=parallel) + big = _zeros(4096) + table = torch.zeros((8, 64), dtype=torch.int32, device=_device()) + out = _zeros(256) + with step_inputs.active(pack): + step_inputs.begin() + for k in range(24): + n = int(torch.randint(1, 64, (1,), generator=g)) + step_inputs.copy_( + big[k * 128 : k * 128 + n], + torch.randint(0, 1 << 30, (n,), generator=g, dtype=torch.int32), + ) + step_inputs.copy_( + table[:3, :40], torch.randint(0, 1 << 30, (3, 40), generator=g, dtype=torch.int32) + ) + if overlap: + step_inputs.copy_(big[0:4], torch.full((4,), -7, dtype=torch.int32)) + step_inputs.copy_(out[:64], big[:64]) + step_inputs.copy_(out[64:96], big[128:160]) + used = pack.used + if how == "flush": + step_inputs.flush() + elif how == "flush_in_graph": + step_inputs.flush(in_graph=True) + else: + _gb_flush(pack) + torch.cuda.synchronize() + assert not pack.open and pack.staged == [] + assert (pack.graph_windows, pack.eager_fallbacks, pack.early_uploads) == (0, 0, 0) + return big.cpu(), table.cpu(), out.cpu(), _uploaded(pack, used, (big, table, out)), pack.parity + + +def _uploaded(pack, used, tensors): + """The staging bytes the flush uploaded, with each device address written as (tensor index, byte offset) so two + runs compare: the header, every descriptor, and every host copy's staged data bytes.""" + staged = pack.dev[:used].cpu() + words = staged[: pack.data_start].view(torch.int64).tolist() + num_h2d, num_d2d = words[0], words[1] + spans = [(t.data_ptr(), t.numel() * t.element_size()) for t in tensors] + + def where(addr): + for i, (base, nbytes) in enumerate(spans): + if base <= addr < base + nbytes: + return (i, addr - base) + raise AssertionError(f"address {addr:#x} outside the test tensors") + + copies = [] + for k in range(num_h2d + num_d2d): + dst, src, nbytes = words[2 + 3 * k : 5 + 3 * k] + if k < num_h2d: + copies.append((where(dst), src, nbytes, staged[src : src + nbytes].tolist())) + else: + copies.append((where(dst), where(src), nbytes)) + return num_h2d, num_d2d, copies + + +@pytest.mark.parametrize("parallel", [False, True], ids=["serial", "parallel"]) +@pytest.mark.parametrize("overlap", [False, True], ids=["disjoint", "overlapping"]) +@pytest.mark.parametrize("seed", [0, 1]) +def test_flags_off_flush_is_the_gb_flush_byte_for_byte(parallel, overlap, seed): + """With pack_in_graph off, flush() and flush(in_graph=True) upload the same staging bytes, write the same + destinations and advance the ring like the GB base's flush.""" + reference = _flags_off_run(parallel, seed, overlap, "gb") + for how in ("flush", "flush_in_graph"): + got = _flags_off_run(parallel, seed, overlap, how) + for a, b in zip(reference[:3], got[:3]): + assert torch.equal(a, b), how + assert got[3] == reference[3], how + assert got[4] == reference[4], how + + +def test_flags_off_pack_allocates_no_window_and_has_no_prologue(): + for kwargs in ({}, {"parallel_scatter": True}): + pack = StepInputPack(_device(), **kwargs) + assert not pack.in_graph and not pack.early_copy + for name in ("ghost", "ghost_free", "gdev", "copy_stream", "staged_ready", "prologue_done"): + assert not hasattr(pack, name), name + with step_inputs.active(pack): + assert step_inputs.capture_prologue() is None + # A closed flags-off pack: flush(in_graph=True) is a no-op, like flush(). + pack.flush(in_graph=True) + pack.quiesce_graph_window() + assert step_inputs.capture_prologue() is None + + +def test_a_pack_built_without_init_keeps_every_in_graph_item_off(): + """The class-level defaults (object.__new__ executors in tests never run __init__).""" + pack = object.__new__(StepInputPack) + assert not pack.parallel_scatter and not pack.in_graph and not pack.early_copy + assert (pack.graph_windows, pack.eager_fallbacks, pack.early_uploads) == (0, 0, 0) + pack.quiesce_graph_window() + assert not pack.capture_prologue() + + +# --------------------------------------------------------------------------- +# In-graph prologue (pack_in_graph) and early copy (pack_early_copy) +# --------------------------------------------------------------------------- + + +def _graph_case(early=False): + """A pack with the in-graph prologue and a graph that starts with it and then reads the destinations.""" + pack = StepInputPack(_device(), parallel_scatter=True, in_graph=True, early_copy=early) + a, b, c = _zeros(16), _zeros(16), _zeros(16) + out = _zeros(48) + graph = torch.cuda.CUDAGraph() + with step_inputs.active(pack): + torch.cuda.synchronize() + with torch.cuda.graph(graph): + assert step_inputs.capture_prologue() is pack + out.copy_(torch.cat([a, b, c])) + return pack, graph, (a, b, c), out + + +def _step(pack, graph, dsts, values, overlap=False, flush=True): + a, b, c = dsts + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(a, _host(values[0])) + step_inputs.copy_(b, _host(values[1])) + if overlap: + step_inputs.copy_(a[:4], _host([-1, -2, -3, -4])) + step_inputs.copy_(c, a) # a device copy reading a host copy's destination + if flush: + step_inputs.flush(in_graph=True) + else: + step_inputs.discard() + step_inputs.flush(in_graph=True) + graph.replay() + + +def test_in_graph_pack_is_issued_by_the_replay(): + pack, graph, dsts, out = _graph_case() + for step in range(5): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + _step(pack, graph, dsts, (va, vb)) + torch.cuda.synchronize() + assert out.tolist() == va + vb + va, step + # Nothing was issued outside the graph: the window carried the copies. + assert int(pack.ghost[:8].view(torch.int64)[0]) == 2 + assert (pack.graph_windows, pack.eager_fallbacks, pack.early_uploads) == (5, 0, 0) + + +def test_in_graph_pack_not_order_free_is_issued_eagerly(): + pack, graph, dsts, out = _graph_case() + va = list(range(16)) + vb = list(range(100, 116)) + _step(pack, graph, dsts, (va, vb), overlap=True) + torch.cuda.synchronize() + expect_a = [-1, -2, -3, -4] + va[4:] + assert out.tolist() == expect_a + vb + expect_a + # The window was zeroed, so the prologue wrote nothing over the eager result. + assert not pack.ghost[:16].any() + assert (pack.graph_windows, pack.eager_fallbacks) == (0, 1) + + +@pytest.mark.parametrize("early", [False, True], ids=["prologue_copy", "early_copy"]) +def test_in_graph_pack_over_the_window_is_issued_eagerly(early): + pack, graph, dsts, out = _graph_case(early) + big = _zeros(StepInputPack.GRAPH_WINDOW // 4) + va = list(range(16)) + vb = list(range(200, 216)) + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(dsts[0], _host(va)) + step_inputs.copy_(dsts[1], _host(vb)) + step_inputs.copy_(big, torch.arange(big.numel(), dtype=torch.int32)) + step_inputs.copy_(dsts[2], dsts[0]) + assert pack.used > StepInputPack.GRAPH_WINDOW + step_inputs.flush(in_graph=True) + graph.replay() + torch.cuda.synchronize() + assert out.tolist() == va + vb + va + assert big.cpu().tolist() == list(range(big.numel())) + assert (pack.graph_windows, pack.eager_fallbacks, pack.early_uploads) == (0, 1, 0) + + +def test_in_graph_replay_without_a_pack_reapplies_nothing(): + pack, graph, dsts, out = _graph_case() + _step(pack, graph, dsts, (list(range(16)), list(range(16)))) + torch.cuda.synchronize() + dsts[0].fill_(7) + _step(pack, graph, dsts, ([0] * 16, [0] * 16), flush=False) + torch.cuda.synchronize() + assert dsts[0].tolist() == [7] * 16 + + +@pytest.mark.parametrize("early", [False, True], ids=["prologue_copy", "early_copy"]) +def test_in_graph_pack_lands_every_step(early): + pack, graph, dsts, out = _graph_case(early) + assert pack.early_copy is early + for step in range(6): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + _step(pack, graph, dsts, (va, vb), overlap=step % 2 == 1) + torch.cuda.synchronize() + head = [-1, -2, -3, -4] if step % 2 == 1 else [] + expect_a = head + va[len(head) :] + assert out.tolist() == expect_a + vb + expect_a, step + assert (pack.graph_windows, pack.eager_fallbacks) == (3, 3) + assert pack.early_uploads == (3 if early else 0) + + +def _history_run(pack, graph, dsts, out, steps, before_replay=None, before_upload=None): + """Steps back to back without host syncs; each replay's output is copied to a history row behind it (same + stream), so a later upload that overwrote the window too early shows up as the wrong step's values.""" + hist = _zeros(steps, out.numel()) + expect = [] + for step in range(steps): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + expect.append(va + vb + va) + a, b, c = dsts + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(a, _host(va)) + step_inputs.copy_(b, _host(vb)) + step_inputs.copy_(c, a) + if before_upload is not None: + before_upload(step) + step_inputs.flush(in_graph=True) + if before_replay is not None: + before_replay(step) + graph.replay() + hist[step].copy_(out) + torch.cuda.synchronize() + return hist.tolist(), expect + + +def test_early_copy_waits_for_the_previous_prologue_before_overwriting_the_window(): + """A slow GPU: the replay of step N (and its prologue) is held back behind a long kernel while the host already + uploads step N+1; that upload must wait for step N's prologue to consume the window.""" + pack, graph, dsts, out = _graph_case(early=True) + hist, expect = _history_run( + pack, graph, dsts, out, 4, before_replay=lambda step: torch.cuda._sleep(_SLEEP_CYCLES // 4) + ) + assert hist == expect + + +def test_early_copy_prologue_waits_for_a_slow_upload(): + """A slow copy stream: each step's upload is held back behind a long kernel on the copy stream; the prologue must + wait for it (staged_ready), and the next host fill must wait for the upload to have read the window.""" + pack, graph, dsts, out = _graph_case(early=True) + + def slow_copy(step): + with torch.cuda.stream(pack.copy_stream): + torch.cuda._sleep(_SLEEP_CYCLES // 4) + + hist, expect = _history_run(pack, graph, dsts, out, 4, before_upload=slow_copy) + assert hist == expect + + +def test_early_copy_with_a_slow_host(): + """A slow host: the GPU idles between steps (the window lands long before the replay).""" + import time + + pack, graph, dsts, out = _graph_case(early=True) + hist, expect = _history_run( + pack, graph, dsts, out, 3, before_replay=lambda step: time.sleep(0.05) + ) + assert hist == expect + + +def test_early_copy_snapshots_the_host_bytes_before_a_stalled_upload(): + """PR-CTX (d)(1): stall the copy stream, issue the early copy, overwrite the host source, then sync: the device + holds the pre-overwrite values. The next step's host fill also runs while the first upload is still stalled (it + must wait for that upload to have read the pinned window before it overwrites it).""" + pack, graph, dsts, out = _graph_case(early=True) + a, b, c = dsts + hist = _zeros(2, out.numel()) + va, vb = list(range(16)), list(range(500, 516)) + src_a, src_b = _host(va), _host(vb) + with torch.cuda.stream(pack.copy_stream): + torch.cuda._sleep(1 << 30) + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(a, src_a) + step_inputs.copy_(b, src_b) + step_inputs.copy_(c, a) + step_inputs.flush(in_graph=True) + # The upload is queued behind the stall: it has not read the window yet. + assert not pack.staged_ready.query() + src_a.fill_(-9) + src_b.fill_(-9) + graph.replay() + hist[0].copy_(out) + # Step 2's fill blocks on ghost_free (the stalled upload) before it writes the pinned window. + va2, vb2 = [7] * 16, [8] * 16 + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(a, _host(va2)) + step_inputs.copy_(b, _host(vb2)) + step_inputs.copy_(c, a) + step_inputs.flush(in_graph=True) + graph.replay() + hist[1].copy_(out) + torch.cuda.synchronize() + assert hist[0].tolist() == va + vb + va + assert hist[1].tolist() == va2 + vb2 + va2 + assert (pack.graph_windows, pack.early_uploads) == (2, 2) + + +@pytest.mark.parametrize("early", [False, True], ids=["prologue_copy", "early_copy"]) +def test_in_graph_replay_after_quiesce_reapplies_nothing(early): + pack, graph, dsts, out = _graph_case(early) + _step(pack, graph, dsts, (list(range(16)), list(range(16)))) + torch.cuda.synchronize() + dsts[0].fill_(7) + pack.quiesce_graph_window() + graph.replay() + torch.cuda.synchronize() + assert dsts[0].tolist() == [7] * 16 + + +def test_early_copy_needs_in_graph(): + assert not StepInputPack(_device(), parallel_scatter=True, early_copy=True).early_copy + + +def test_in_graph_needs_parallel_scatter_and_is_off_by_default(): + assert not StepInputPack(_device()).in_graph + assert not StepInputPack(_device(), in_graph=True).in_graph + + +def test_in_graph_replay_outside_the_pack_context_reapplies_nothing(): + pack, graph, dsts, out = _graph_case() + _step(pack, graph, dsts, (list(range(16)), list(range(16)))) + torch.cuda.synchronize() + dsts[0].fill_(7) + # What the graph runner does for a replay outside the pack's context. + assert not step_inputs.is_active(pack) + pack.quiesce_graph_window() + graph.replay() + torch.cuda.synchronize() + assert dsts[0].tolist() == [7] * 16 + + +def _kernel_compiles(kernel) -> int: + """How many variants of a Triton kernel are compiled (any device).""" + caches = getattr(kernel, "device_caches", None) + if caches is not None: + return sum(len(entry[0]) for entry in caches.values()) + return sum(len(entry) for entry in kernel.cache.values()) + + +@pytest.mark.parametrize("early", [False, True], ids=["prologue_copy", "early_copy"]) +def test_in_graph_kernels_compile_nothing_after_the_warm_up(early): + """The pack's warm-up compiles both scatter kernels' variants; captures and steps with any number of copies (in + the window, over it, not order-free, quiesced) compile nothing more.""" + pack = StepInputPack(_device(), parallel_scatter=True, in_graph=True, early_copy=early) + kernels = (step_inputs._parallel_scatter_kernel, step_inputs._graph_scatter_kernel) + compiled = [_kernel_compiles(kernel) for kernel in kernels] + dst = _zeros(8192) + out = _zeros(64) + graph = torch.cuda.CUDAGraph() + with step_inputs.active(pack): + torch.cuda.synchronize() + with torch.cuda.graph(graph): + assert step_inputs.capture_prologue() is pack + out.add_(1) + for n_host in (1, 2, 3, 5, 15, 16, 17, 24, 120): + with step_inputs.active(pack): + step_inputs.begin() + for k in range(n_host): + step_inputs.copy_(dst[k * 64 : k * 64 + 4], torch.full((4,), k, dtype=torch.int32)) + for k in range(n_host % 7): + step_inputs.copy_(out[k * 4 : k * 4 + 4], dst[k * 64 : k * 64 + 4]) + if n_host == 17: + step_inputs.copy_( + dst[4096:], torch.ones(4096, dtype=torch.int32) + ) # over the window + step_inputs.flush(in_graph=True) + graph.replay() + if n_host == 5: + pack.quiesce_graph_window() + graph.replay() + torch.cuda.synchronize() + assert [_kernel_compiles(kernel) for kernel in kernels] == compiled + assert dst[119 * 64 : 119 * 64 + 4].tolist() == [119] * 4 + assert pack.eager_fallbacks == 1 and pack.graph_windows == 8 + + +# --------------------------------------------------------------------------- +# pack_in_graph through the real CUDAGraphRunner capture -> replay +# --------------------------------------------------------------------------- + + +def _runner_case(early, in_graph=True): + from _torch.helpers import create_mock_cuda_graph_runner + + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import KeyType + + batch_size = 4 + runner = create_mock_cuda_graph_runner(batch_size, max_num_tokens=128) + key = KeyType(batch_size=batch_size, draft_len=0) + num_tokens = runner._get_num_tokens_for_key(key) + pack = StepInputPack(_device(), parallel_scatter=True, in_graph=in_graph, early_copy=early) + meta = _zeros(8) + out = torch.zeros(2 * num_tokens + meta.numel(), dtype=torch.int64, device="cuda") + attn_metadata = object() + + def forward_fn(inputs): + out.copy_( + torch.cat( + [ + inputs["input_ids"].to(torch.int64), + inputs["position_ids"].flatten().to(torch.int64), + meta.to(torch.int64), + ] + ) + ) + return out + + def inputs(value): + return { + "attn_metadata": attn_metadata, + "input_ids": torch.full((num_tokens,), value, dtype=torch.int32, device="cuda"), + "position_ids": torch.arange(num_tokens, dtype=torch.int32, device="cuda").view(1, -1) + + value, + } + + with step_inputs.active(pack): + runner.capture(key, forward_fn, inputs(0)) + return runner, key, pack, meta, out, inputs, num_tokens + + +def _runner_step(runner, key, pack, meta, inputs, value): + with step_inputs.active(pack): + step_inputs.begin() + # A metadata upload the engine stages before the replay (as the MSA metadata does). + step_inputs.copy_(meta, _host([value * 10 + i for i in range(meta.numel())])) + runner.replay(key, inputs(value)) + + +@pytest.mark.parametrize("early", [False, True], ids=["prologue_copy", "early_copy"]) +def test_in_graph_pack_through_the_graph_runner_capture_and_replay(early): + """capture() adds the active pack's prologue to the graph; replay() with fresh inputs only fills the window (the + static input copies and the metadata upload are issued by the replay's prologue), every step lands.""" + runner, key, pack, meta, out, inputs, n = _runner_case(early) + assert runner._step_input_prologue_keys == {key: pack} + for value in range(1, 6): + _runner_step(runner, key, pack, meta, inputs, value) + torch.cuda.synchronize() + expect = [value] * n + [value + i for i in range(n)] + [value * 10 + i for i in range(8)] + assert out.tolist() == expect, value + assert pack.graph_windows == value and pack.eager_fallbacks == 0 + assert pack.early_uploads == (5 if early else 0) + # Replayed outside the pack's context: the inputs are plain copies and the prologue re-issues nothing. + meta.fill_(-3) + runner.replay(key, inputs(9)) + torch.cuda.synchronize() + assert out.tolist() == [9] * n + [9 + i for i in range(n)] + [-3] * 8 + runner.clear() + assert runner._step_input_prologue_keys is None + + +def test_graph_runner_flags_off_replays_through_the_stock_flush(monkeypatch): + """A pack without in_graph: the capture adds no prologue and replay() issues the pack with flush() exactly as on + the GB base (no in_graph argument), before the graph replay.""" + runner, key, pack, meta, out, inputs, n = _runner_case(False, in_graph=False) + assert runner._step_input_prologue_keys is None + calls = [] + real_flush = step_inputs.flush + + def recording_flush(*args, **kwargs): + calls.append((args, kwargs)) + return real_flush(*args, **kwargs) + + monkeypatch.setattr(step_inputs, "flush", recording_flush) + _runner_step(runner, key, pack, meta, inputs, 3) + torch.cuda.synchronize() + assert calls == [((), {})] + assert out.tolist() == [3] * n + [3 + i for i in range(n)] + [30 + i for i in range(8)] + assert (pack.graph_windows, pack.eager_fallbacks) == (0, 0) diff --git a/tests/unittest/_torch/modeling_v2/comm/_allgather_cast_op_matrix.py b/tests/unittest/_torch/modeling_v2/comm/_allgather_cast_op_matrix.py new file mode 100644 index 000000000000..e39734090c28 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/comm/_allgather_cast_op_matrix.py @@ -0,0 +1,266 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Certify `comm/allgather_cast`: every rank's fp32 rows equal `reference` (the all-gather, the vocab concat and +`.float()`) bit for bit, on every cell certified at this run's world size. + +Every rank builds every rank's shard from one seed, so each rank checks its own output against the reference over +all shards. Per cell: random logits with adversarial bits planted (-0.0, +-inf, NaN payloads, denormals, the +0x80000000 pair word), several seeds, inside `validating`. Also: a captured call replayed with fresh contents copied +into its static input before every replay (the epoch and the buffer rotation advance on the device); call-unique +contents while the grid alternates between 16 CTAs and the picked grid; `supports` refuses more rows than the +workspace holds, fp32 shards and partial 16-byte vectors. + +World 4 (one GB300 tray) is spawned from one process: + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _allgather_cast_op_matrix.py + +World 8 spans two trays and runs as one srun task per rank, like `_mnnvl_allreduce_op_matrix.py`: + + MASTER_ADDR= MASTER_PORT= python _allgather_cast_op_matrix.py --srun + +under ``srun --ntasks=8 --ntasks-per-node=4 --kill-on-bad-exit=1``. +""" + +import importlib.util +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Optional + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "allgather_cast requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 +SEEDS = 4 +GPUS_PER_NODE = 4 +"""A GB300 tray. The mapping needs it to see a world-8 group span two nodes.""" +# bf16 bit patterns planted into every shard: -0.0, +inf, -inf, NaN payloads, denormals. +SPECIAL_BITS = (-32768, 0x7F80, -128, 0x7FC1, -63, 0x0001, -32767, 0x7FFF) + +RANK = 0 +WORLD = 1 +# Bound inside the rank body: `entry` is the module, `op` its `allgather_cast` instance (the call surface, +# `reference`, `is_valid` and `CELLS`), `WS` the workspace arguments of every call, per shard width. +entry: Any = None +op: Any = None +validating: Any = None +WS: dict = {} + + +def _shards(seed: int, rows: int, cols: int): + """Every rank's bf16 [rows, cols] shard (the same on every rank), adversarial bits planted.""" + g = torch.Generator(device="cuda").manual_seed(seed) + out = [] + for _ in range(WORLD): + x = (torch.randn(rows, cols, device="cuda", generator=g) * 3).bfloat16() + w = x.view(torch.int16) + # Distinct columns: scatter_ with a repeated index keeps an arbitrary writer, which can differ between + # devices, and every rank must build the same shards. + cols_hit = torch.randperm(cols, generator=g, device="cuda")[:64] + special = torch.tensor(SPECIAL_BITS, dtype=torch.int16, device="cuda").repeat(8) + w[:, cols_hit] = special + w[:, 0] = 0 + w[:, 1] = -32768 # the int32 word 0x80000000 + out.append(x) + return out + + +def _call(shard: torch.Tensor, grid: int = 0) -> torch.Tensor: + return op(shard, **WS[shard.shape[1]], grid=grid) + + +def _bits_equal(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and torch.equal(a.view(torch.int32), b.view(torch.int32)) + + +def check_certified_cells() -> None: + cells = [c for c in op.CELLS if c.spec["world"] == WORLD] + assert cells, f"no cell is certified at world size {WORLD}" + for cell in cells: + rows, cols = cell.spec["rows"], cell.spec["cols"] + for seed in range(SEEDS): + shards = _shards(seed, rows, cols) + with validating(op): + got = _call(shards[RANK]) + assert _bits_equal(got, op.reference(shards)), ( + f"{cell.why}: seed {seed} differs on rank {RANK}" + ) + dist.barrier() + if RANK == 0: + print(f"[cell] world={WORLD} {cell.why}: bitwise", flush=True) + + +def check_graph_replays_fresh_contents() -> None: + cell = [c for c in op.CELLS if c.spec["world"] == WORLD][0] + rows, cols = cell.spec["rows"], cell.spec["cols"] + static = _shards(100, rows, cols)[RANK].clone() + _call(static) + torch.cuda.synchronize() + graph_out = {} + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + g = torch.cuda.CUDAGraph() + with torch.cuda.stream(s): + with torch.cuda.graph(g): + graph_out["t"] = _call(static) + torch.cuda.current_stream().wait_stream(s) + for i in range(6): + shards = _shards(200 + i, rows, cols) + static.copy_(shards[RANK]) + g.replay() + torch.cuda.synchronize() + assert _bits_equal(graph_out["t"], op.reference(shards)), ( + f"replay {i} returned stale or wrong rows" + ) + dist.barrier() + + +def check_alternating_grids() -> None: + cell = [c for c in op.CELLS if c.spec["world"] == WORLD][0] + rows, cols = cell.spec["rows"], cell.spec["cols"] + for i in range(24): + shards = _shards(300 + i, rows, cols) + got = _call(shards[RANK], grid=16 if i % 2 else 0) + assert _bits_equal(got, op.reference(shards)), ( + f"call {i} (grid {16 if i % 2 else 'picked'}) differs" + ) + dist.barrier() + + +def check_supports_refuses() -> None: + cols = [c for c in op.CELLS if c.spec["world"] == WORLD][0].spec["cols"] + shard = torch.zeros(entry.MAX_ROWS + 1, cols, dtype=torch.bfloat16, device="cuda") + assert not entry.supports(shard, WORLD), "more rows than the workspace holds" + assert not entry.supports(shard[:4].float(), WORLD), "an fp32 shard" + assert not entry.supports( + torch.zeros(4, cols - 4, dtype=torch.bfloat16, device="cuda"), WORLD + ), "a partial 16-byte vector" + assert not entry.supports(shard[:4, :-8], WORLD), "a strided view" + dist.barrier() + + +CHECKS = ( + check_certified_cells, + check_graph_replays_fresh_contents, + check_alternating_grids, + check_supports_refuses, +) + + +def _run_one_rank(device: int) -> int: + global RANK, WORLD, WS, entry, op, validating + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(device) + dist.init_process_group( + "cuda:nccl,cpu:gloo", init_method="env://", device_id=torch.device("cuda", device) + ) + RANK, WORLD = rank, world + + from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm import ( + allgather_cast as entry_module, + ) + from tensorrt_llm.mapping import Mapping + + entry = entry_module + op = entry.allgather_cast + mapping = Mapping(world_size=world, rank=rank, gpus_per_node=GPUS_PER_NODE, tp_size=world) + for cols in sorted({c.spec["cols"] for c in op.CELLS if c.spec["world"] == world}): + WS[cols] = entry.workspace_args(mapping, cols, dist.group.WORLD) + + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + from _validating import validating as _validating_cm + + validating = _validating_cm + + 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() + os._exit(13) + print(f"[rank {RANK}] {check.__name__} passed", flush=True) + dist.barrier() + torch.cuda.synchronize() + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict(os.environ, MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port), WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1") # fmt: skip + 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, + ) # fmt: skip + 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}") + + +def _rank_job(): + """`_rank_job.py` beside this file, by path (a sibling directory ships its own module of that name).""" + path = Path(__file__).resolve().with_name("_rank_job.py") + spec = importlib.util.spec_from_file_location("_comm_rank_job", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +if __name__ == "__main__": + if "--srun" in sys.argv: + sys.exit(_run_one_rank(device=_rank_job().adopt_srun_env())) + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank(device=int(os.environ["RANK"]))) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/comm/_rank_job.py index 3784a1a98361..f5aabba7f1e0 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_rank_job.py +++ b/tests/unittest/_torch/modeling_v2/comm/_rank_job.py @@ -65,9 +65,10 @@ def _load_launcher_constants(launcher: Path): return module -def run(entry: str) -> None: +def run(entry: str, *, launcher_dir: Path | None = None) -> None: """Run ``__op_matrix``'s launcher over ``WORLD_SIZE`` devices.""" - launcher = Path(__file__).resolve().with_name(f"_{entry}_op_matrix.py") + directory = Path(__file__).resolve().parent if launcher_dir is None else launcher_dir.resolve() + launcher = directory / 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 diff --git a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_allgather_cast_op_matrix.py b/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_allgather_cast_op_matrix.py new file mode 100644 index 000000000000..0d5b6c41ab6e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_allgather_cast_op_matrix.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for `comm/allgather_cast`'s certification matrix (``_allgather_cast_op_matrix.py`` beside +this file, which drives the entry's ``.CELLS``). + +World 4 is collected here. The world-8 cells span two GB300 trays and run as one srun task per rank (see the +matrix's docstring), like ``_mnnvl_allreduce_op_matrix.py``. +""" + +import importlib.util +from pathlib import Path + +import pytest +import torch + +assert torch.cuda.is_available(), "allgather_cast requires CUDA devices" + + +def _rank_job(): + # By path: other test directories ship a module of the same name, and `import _rank_job` would get whichever + # one a collected sibling imported first. + path = Path(__file__).resolve().with_name("_rank_job.py") + spec = importlib.util.spec_from_file_location("_comm_rank_job", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +# Spawns 4 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_allgather_cast_op_matrix() -> None: + _rank_job().run("allgather_cast") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/activation/test_modeling_v2_m31_tp4_swq.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/activation/test_modeling_v2_m31_tp4_swq.py new file mode 100644 index 000000000000..bc688eeb3307 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/activation/test_modeling_v2_m31_tp4_swq.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 m3_swq entry of the MiniMax-M3.1 TP4 target: SwiGLU-OAI written as MXFP8, bitwise to +silu_and_mul + mxfp8_quantize.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.ops.activation.m3_swq import ( # noqa: E501 + m3_swq as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_swq requires a CUDA device" + +OAI = (7.0, 1.702, 1.0) + + +def _assert_bitwise(x): + q, sf = op(x, *OAI) + rq, rsf = op.reference(x, *OAI) + assert q.shape == rq.shape and q.dtype == torch.float8_e4m3fn + assert torch.equal(q.view(torch.uint8), rq.view(torch.uint8)) + assert sf.numel() <= rsf.numel() and torch.equal(sf, rsf.view(-1)[: sf.numel()]) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = (torch.randn(cell.spec["m"], 2 * cell.spec["inter"], generator=g, device="cuda") * 4.0).to( + torch.bfloat16 + ) + with validating(op): + _assert_bitwise(x) + + +@pytest.mark.parametrize("m", [17, 24, 32]) +@pytest.mark.parametrize("scale", [1e-3, 1.0, 1e3]) +def test_bitwise_across_magnitudes_and_edges(m, scale) -> None: + """Scaled inputs, gates past the clamp, an all-zero row and a -0.0 row: q bytes and every scale byte.""" + g = torch.Generator(device="cuda").manual_seed(m * 1000 + int(scale * 10) % 997) + x = torch.randn(m, 2 * 768, generator=g, device="cuda") * scale + x[1, :768] = x[1, :768].abs() + OAI[0] + x[2] = 0.0 + x[3] = -0.0 + _assert_bitwise(x.to(torch.bfloat16)) + + +def _scale_bytes(sf, m, nb): + """Un-swizzle the 128x4 scale layout to [m, nb] bytes.""" + nbp = (nb + 3) // 4 * 4 + r = torch.arange(m, device=sf.device).view(-1, 1) + c = torch.arange(nb, device=sf.device).view(1, -1) + idx = ( + (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (r // 128) * (128 * nbp) + ) + return sf.view(-1)[idx] + + +@pytest.mark.parametrize("m", [17, 32]) +def test_tiny_amax_blocks_stay_finite(m) -> None: + """a block whose scale code is 0 (amax / 448 <= 2^-127) is native's inf/NaN; SWQ keeps native's scale + byte, writes data quantized with the exact reciprocal 2^127, and stays bitwise to native on every other block.""" + g = torch.Generator(device="cuda").manual_seed(151 + m) + gate = torch.randn(m, 768, generator=g, device="cuda") + up = torch.randn(m, 768, generator=g, device="cuda") + pattern = torch.tensor([1.0, -0.75, 0.5, -0.25, 0.375, 0.0, -1.0, 0.625] * 4, device="cuda") + for r in range(m): + for j, c in enumerate([2.0**-130, 2.0**-127, 448 * 2.0**-128, 448 * 2.0**-127, 2.0**-118]): + blk = (r * 5 + j * 3) % 24 + gate[r, blk * 32 : (blk + 1) * 32] = pattern * (2 * c) + up[r, blk * 32 : (blk + 1) * 32] = pattern * (2 * c) + x = torch.cat([gate, up], 1).to(torch.bfloat16) + q, sf = op(x, *OAI) + rq, rsf = op.reference(x, *OAI) + h = torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=OAI[0], swiglu_alpha=OAI[1], swiglu_beta=OAI[2] + ).float() + byte, rbyte = _scale_bytes(sf, m, 24), _scale_bytes(rsf, m, 24) + assert torch.equal(byte, rbyte), "scale bytes are native's on every block" + deq = q.float().view(m, 24, 32) * torch.exp2(byte.float() - 127.0).unsqueeze(-1) + rdeq = rq.float().view(m, 24, 32) * torch.exp2(rbyte.float() - 127.0).unsqueeze(-1) + native_finite = torch.isfinite(rdeq).all(-1) + code0 = (byte == 0) & (h.view(m, 24, 32).abs().amax(-1) != 0) + assert bool(code0.any()), "the planted blocks reach scale code 0" + assert not bool(native_finite[code0].any()), ( + "native is inf/NaN on code-0 blocks (the bug this guards)" + ) + assert bool(torch.isfinite(deq).all()), "SWQ is finite everywhere" + qb, rqb = q.view(torch.uint8).view(m, 24, 32), rq.view(torch.uint8).view(m, 24, 32) + assert bool((qb == rqb).all(-1)[native_finite].all()), ( + "bitwise to native wherever native is finite" + ) + exact = (h.view(m, 24, 32) * 2.0**127).to(torch.float8_e4m3fn).view(torch.uint8) + assert bool((qb == exact).all(-1)[code0].all()), "code-0 data = RN-e4m3(h * 2^127)" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/attention/_m3_nvfp4_cache.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/attention/test_modeling_v2_m31_gqamix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/attention/test_modeling_v2_m31_gqamix.py new file mode 100644 index 000000000000..c3aa7c93c210 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/attention/test_modeling_v2_m31_gqamix.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GQAMIX: a mixed step's sparse generation rows (MsaDecodeFmha) through ops/gqamix_sparse_decode. + +The copy must be the in-tree kernel module except for the two decode kernels' specialization (no row count, no slice +alignment) and the one import it cannot make relatively, and must write the in-tree kernel's output bit for bit on a +mixed step's [num_ctx:] slices. The shared MsaDecodeFmha hook must change nothing while it is unset. +""" + +import inspect +import os +import subprocess +import sys +import types +from pathlib import Path + +import pytest +import torch + +import tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode as in_tree +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4 import ( + modeling as m, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.ops import ( + gqamix_sparse_decode as copy, +) +from tensorrt_llm._torch.attention.backends.fmha.msa_decode import MsaDecodeFmha + +HEADS = 16 # query heads per rank at tp4, one KV head +D, PAGE, TOPK, QLEN, SF_VEC = 128, 128, 16, 4, 16 + +_JIT_OLD = '@triton.jit(do_not_specialize=["decode_query_len"])' +_JIT_NEW = ( + "@triton.jit(\n" + ' do_not_specialize=["decode_query_len", "total_q"],\n' + ' do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"],\n' + ")" +) +_IMPORT_OLD = "from .msa_utils import check_decode_span_shape" +_IMPORT_NEW = ( + "from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import (\n" + " check_decode_span_shape,\n" + ")" +) + + +def _code_lines(text: str) -> list: + return [line for line in text.splitlines() if line.strip()] + + +def test_the_copy_is_the_in_tree_module_but_for_its_specialization() -> None: + """Line for line (blank lines aside: the import sorts into a different block) the in-tree module.""" + ours = Path(copy.__file__).read_text() + theirs = Path(in_tree.__file__).read_text() + assert ( + ours.count(_JIT_NEW) == 2 + and ours.count(_IMPORT_NEW) == 1 + and theirs.count(_IMPORT_OLD) == 1 + ) + ours = ours.replace(_JIT_NEW, _JIT_OLD).replace(_IMPORT_NEW, "") + assert _code_lines(ours) == _code_lines(theirs.replace(_IMPORT_OLD, "")), ( + "ops/gqamix_sparse_decode.py must stay a copy of the in-tree triton_sparse_decode.py: re-copy it" + ) + + +def test_register_sets_the_hook_once_and_only_where_it_exists() -> None: + lib = types.SimpleNamespace(sparse_decode_hook=None) + other_lib = types.SimpleNamespace() # an FMHA library without the hook + manager = types.SimpleNamespace(fmha_libs=[other_lib, lib]) + m._gqamix_register(types.SimpleNamespace(_fmha_manager=manager)) + assert lib.sparse_decode_hook is m._gqamix_decode + assert not hasattr(other_lib, "sparse_decode_hook") + mine = object() + lib.sparse_decode_hook = mine + m._gqamix_register(types.SimpleNamespace(_fmha_manager=manager)) + assert lib.sparse_decode_hook is mine + m._gqamix_register(types.SimpleNamespace()) # a backend without an FMHA manager + + +def test_the_kill_switch_leaves_the_pre_hook_decode() -> None: + """TRTLLM_M31_DISABLE_GQAMIX=1: GQAMIX is off, so no layer registers the hook and MsaDecodeFmha runs the in-tree + decode (whose output the copy matches bit for bit below).""" + target = "tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4" + probe = f"import {target}.modeling as m; print(m.GQAMIX)" + env = dict(os.environ, TRTLLM_M31_DISABLE_GQAMIX="1") + out = subprocess.run( + [sys.executable, "-c", probe], env=env, capture_output=True, text=True, check=True + ) + assert out.stdout.strip().splitlines()[-1] == "False" + source = inspect.getsource(m.PrefillTarget) + assert "if GQAMIX and sparse:\n _gqamix_register(core.attn[i])" in source + assert source.count("_gqamix_register(") == 1 + + +class _Attn: + is_mla_enable = False + kv_lora_rank = None + v_head_dim = None + head_dim = D + num_heads = HEADS + layer_idx = 3 + q_scaling = 1.0 + + +def _params(attn, num_tokens): + return types.SimpleNamespace( + attn=attn, + num_tokens=num_tokens, + meta=types.SimpleNamespace(kv_cache_manager=object()), + fwd=types.SimpleNamespace(kv_scale_quant_orig=None), + query_input=torch.zeros(num_tokens, HEADS * D), + output=torch.zeros(num_tokens, HEADS * D), + token_offset=2, + input_seq_length=QLEN, + ) + + +@pytest.mark.parametrize("hooked", [False, True], ids=["unset", "registered"]) +def test_the_shared_decode_hook_is_inert_unset(monkeypatch, hooked) -> None: + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels import msa_utils + + calls = [] + monkeypatch.setattr(msa_utils, "msa_paged_kv", lambda manager, layer: ("k", "v")) + monkeypatch.setattr( + in_tree, "minimax_m3_sparse_attn_decode", lambda *a, **k: calls.append(("in_tree", a, k)) + ) + attn = _Attn() + fmha = MsaDecodeFmha(attn) + assert fmha.sparse_decode_hook is None + if hooked: + fmha.sparse_decode_hook = lambda *a, **k: calls.append(("hook", a, k)) + indexes = torch.zeros(10, 1, TOPK, dtype=torch.int32) + fmha._run_sparse(_params(attn, QLEN), indexes, "bt", "sl") + assert [c[0] for c in calls] == ["hook" if hooked else "in_tree"] + _, args, kwargs = calls[0] + assert args[1:3] == ("k", "v") and args[4:] == ("bt", "sl") + assert args[3].shape == (1, QLEN, TOPK) + assert kwargs["decode_query_len"] == QLEN and kwargs["sm_scale"] == D**-0.5 + + +def _mixed_step(nreq, seq_len, row_first, seed): + """A mixed step's generation rows: NVFP4 caches, FP8 q, and block-table / length slices past row_first.""" + g = torch.Generator(device="cuda").manual_seed(seed) + nb = seq_len // PAGE + pages = (nreq + row_first) * nb + block_table = torch.randperm(pages, device="cuda", generator=g).to(torch.int32).view(-1, nb) + seq_lens = torch.full((nreq + row_first,), seq_len, device="cuda", dtype=torch.int32) + total_q = nreq * QLEN + q = torch.randn(total_q, HEADS, D, device="cuda", generator=g).to(torch.float8_e4m3fn) + kv = [ + torch.randint( + 0, 256, (pages, 1, PAGE, D // 2), device="cuda", generator=g, dtype=torch.uint8 + ) + for _ in "kv" + ] + sc = [ + (torch.rand(pages, 1, PAGE, D // SF_VEC, device="cuda", generator=g) + 0.5).to( + torch.float8_e4m3fn + ) + for _ in "kv" + ] + topk = torch.stack( + [torch.randperm(nb, device="cuda", generator=g)[:TOPK] for _ in range(total_q)] + ) + args = (q, kv[0], kv[1], topk.to(torch.int32).view(total_q, 1, TOPK).permute(1, 0, 2)) + args += (block_table[row_first:], seq_lens[row_first:]) + kwargs = dict( + sm_scale=D**-0.5, + decode_query_len=QLEN, + k_block_scale=sc[0], + v_block_scale=sc[1], + k_global_scale=torch.tensor([0.02], device="cuda"), + v_global_scale=torch.tensor([0.03], device="cuda"), + ) + return args, kwargs, total_q + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +@pytest.mark.parametrize("nreq,row_first", [(1, 0), (1, 1), (3, 2), (5, 3), (8, 1)]) +@pytest.mark.parametrize("seq_len", [16384, 131072]) +def test_the_copy_writes_the_in_tree_output_on_a_mixed_step(nreq, row_first, seq_len) -> None: + args, kwargs, total_q = _mixed_step(nreq, seq_len, row_first, 31 * nreq + row_first) + ref = torch.empty(total_q, HEADS, D, device="cuda", dtype=torch.bfloat16) + out = torch.empty_like(ref) + in_tree.minimax_m3_sparse_attn_decode(*args, output=ref, **kwargs) + copy.minimax_m3_sparse_attn_decode(*args, output=out, **kwargs) + torch.cuda.synchronize() + assert torch.isfinite(ref.float()).all() + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/attention/test_modeling_v2_m31_lc4.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/attention/test_modeling_v2_m31_lc4.py new file mode 100644 index 000000000000..06eaf7caa11f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/attention/test_modeling_v2_m31_lc4.py @@ -0,0 +1,170 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for LC4: this target's copy of the long-context sparse decode +(``torch.ops.trtllm.minimax_m31_tp4_sparse_decode_lc``) with its two exact options, block-table staging (btsmem) and +the softmax sync trim (sync3), must write the same bits as the M3 op (``minimax_m3_sparse_decode_lc``): the BF16 rows, +the MXFP8 hand-off and its UE8M0 scales, for every option combination over the certified cells' shapes, a row longer +than the staged block table, repeated runs and a CUDA-graph replay. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.ops.attention import ( # noqa: E501 + m3_sparse_decode_lc as entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.ops.attention.m3_sparse_decode_lc import ( # noqa: E501 + WIDTH, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_nvfp4_cache import build_step # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "LC4 requires a CUDA device" + +_HANDOFF_ROWS = 128 +_OPTIONS = ((True, False), (False, True), (True, True)) + + +def _lc4(): + op = getattr(torch.ops.trtllm, "minimax_m31_tp4_sparse_decode_lc", None) + if op is None and entry._lc4_op() is not None: + op = torch.ops.trtllm.minimax_m31_tp4_sparse_decode_lc + if op is None: + pytest.skip("minimax_m31_tp4_sparse_decode_lc is not built") + return op + + +def _run(fn, case, linear_softmax, mxfp8, *opts): + total_q = case.q.shape[0] + out = torch.full((total_q, WIDTH), float("nan"), dtype=torch.bfloat16, device="cuda") + q8 = torch.zeros(_HANDOFF_ROWS, WIDTH, dtype=torch.uint8, device="cuda") if mxfp8 else None + sf = ( + torch.zeros(_HANDOFF_ROWS * (WIDTH // 32), dtype=torch.uint8, device="cuda") + if mxfp8 + else None + ) + fn( + 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, + linear_softmax, + False, + True, + *opts, + ) + return out, q8, sf + + +def _assert_same(a, b, what): + assert torch.equal(a[0].view(torch.int16), b[0].view(torch.int16)), f"{what}: BF16 rows differ" + if a[1] is not None: + assert torch.equal(a[1], b[1]), f"{what}: MXFP8 rows differ" + assert torch.equal(a[2], b[2]), f"{what}: UE8M0 scales differ" + + +@pytest.mark.parametrize( + "cell", + entry.m3_sparse_decode_lc.CELLS, + ids=[c.why[:44] for c in entry.m3_sparse_decode_lc.CELLS], +) +@pytest.mark.parametrize("linear_softmax", (False, True)) +@pytest.mark.parametrize("mxfp8", (False, True)) +@pytest.mark.parametrize("btsmem,sync3", _OPTIONS) +def test_lc4_matches_the_m3_op_bit_for_bit(cell, linear_softmax, mxfp8, btsmem, sync3) -> None: + lc4 = _lc4() + spec = cell.spec + case = build_step( + spec["seq_lens"], + spec["decode_query_len"], + identity=False, + seed=zlib.crc32(cell.why.encode()) % 2**31, + ) + ref = _run(torch.ops.trtllm.minimax_m3_sparse_decode_lc, case, linear_softmax, mxfp8) + got = _run(lc4, case, linear_softmax, mxfp8, btsmem, sync3) + _assert_same(got, ref, f"btsmem={btsmem} sync3={sync3}") + + +def test_lc4_row_longer_than_the_staged_block_table() -> None: + """A request past 8192 blocks reads its block table from global memory under btsmem: same bits.""" + lc4 = _lc4() + case = build_step((8193 * 128 + 5,), 1, identity=False, seed=31) + ref = _run(torch.ops.trtllm.minimax_m3_sparse_decode_lc, case, False, True) + for btsmem, sync3 in _OPTIONS: + _assert_same( + _run(lc4, case, False, True, btsmem, sync3), + ref, + f"long row btsmem={btsmem} sync3={sync3}", + ) + + +def test_lc4_graph_replay_and_repeats_are_bit_identical() -> None: + """Three eager runs and a CUDA-graph replay of the shipped placement (both options, hand-off on) agree.""" + lc4 = _lc4() + case = build_step((33000,), 4, identity=False, seed=21) + first = _run(lc4, case, False, True, True, True) + for _ in range(2): + _assert_same(_run(lc4, case, False, True, True, True), first, "repeat") + total_q = case.q.shape[0] + out = torch.empty(total_q, WIDTH, dtype=torch.bfloat16, device="cuda") + q8 = torch.zeros(_HANDOFF_ROWS, WIDTH, dtype=torch.uint8, device="cuda") + sf = torch.zeros(_HANDOFF_ROWS * (WIDTH // 32), dtype=torch.uint8, device="cuda") + args = ( + 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, + False, + True, + True, + True, + True, + ) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + lc4(*args) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + lc4(*args) + out.fill_(float("nan")) + q8.zero_() + sf.zero_() + graph.replay() + torch.cuda.synchronize() + _assert_same((out, q8, sf), first, "graph replay") + + +def test_btsmem_is_scoped_to_small_steps() -> None: + """btsmem costs larger grids their occupancy: the entry enables it only up to LC4_BTSMEM_MAX_ROWS rows.""" + assert entry.LC4_BTSMEM_MAX_ROWS == 8 + if entry.LC4_BTSMEM: + assert entry._lc4_btsmem(4) and entry._lc4_btsmem(8) + assert not entry._lc4_btsmem(9) and not entry._lc4_btsmem(32) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/attention/test_modeling_v2_m31_tsel.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/attention/test_modeling_v2_m31_tsel.py new file mode 100644 index 000000000000..27ceeb61fb8d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/attention/test_modeling_v2_m31_tsel.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for TSEL: the PrefillTarget's top-16 block select through the #20629 prefill tile select +(torch.ops.trtllm.minimax_m31_select_prefill_tile), registered on the MSA indexer's selector hook. + +The table must equal minimax_m3_select_blocks' (head-major backing) bit for bit, ties included; outside its rows / +blocks range the hook returns None so the built-in select runs. Skipped when the op is not built (no overlay library). +""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4 import ( + modeling as m, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + select_blocks_from_maxscore, +) + +assert torch.cuda.is_available(), "TSEL requires a CUDA device" + +BLOCK = 128 +op = m._tsel_op() +needs_op = pytest.mark.skipif( + op is None, reason="minimax_m31_select_prefill_tile is not built (or not sm_103)" +) + + +def _inputs(ctx_end, rows, ties, seed, heads=1): + g = torch.Generator(device="cuda").manual_seed(seed) + blocks = (ctx_end + BLOCK - 1) // BLOCK + pos = torch.arange(ctx_end - rows, ctx_end, device="cuda") + nvb = ((pos + BLOCK) // BLOCK).to(torch.int32) + if ties: + scores = torch.randint(0, 8, (heads, blocks, rows), generator=g, device="cuda").float() + else: + scores = torch.randn(heads, blocks, rows, generator=g, device="cuda") + return scores.contiguous(), nvb + + +def _ref(scores, nvb, init_blocks, local_blocks): + return select_blocks_from_maxscore( + scores, + topk=16, + n_valid_blocks=nvb, + init_blocks=init_blocks, + local_blocks=local_blocks, + head_major_output=True, + ) + + +@needs_op +@pytest.mark.parametrize( + "ctx_end,rows", [(16384, 16384), (65536, 4096), (163840, 2048), (165000, 1160), (1024, 256)] +) +@pytest.mark.parametrize("ties", [False, True], ids=["random", "ties"]) +@pytest.mark.parametrize("init_blocks,local_blocks", [(0, 1), (2, 3)]) +def test_table_equals_select_blocks(ctx_end, rows, ties, init_blocks, local_blocks) -> None: + scores, nvb = _inputs(ctx_end, rows, ties, ctx_end + rows) + got = m._tsel_select(scores, nvb, init_blocks, local_blocks) + assert got is not None + ref = _ref(scores, nvb, init_blocks, local_blocks) + assert ( + got.shape == ref.shape + and got.stride() == ref.stride() + and got.dtype == ref.dtype == torch.int32 + ) + assert torch.equal(got, ref) + + +@needs_op +def test_two_kv_heads_and_minus_one_tail() -> None: + """Rows with fewer than 16 valid blocks pad with -1; two KV heads keep the head-major backing.""" + scores, nvb = _inputs(2048, 2048, False, 7, heads=2) + ref = _ref(scores, nvb, 0, 1) + got = m._tsel_select(scores, nvb, 0, 1) + assert got.shape == ref.shape and got.stride() == ref.stride() and got.dtype == ref.dtype + assert torch.equal(got, ref) + assert (got[:100] == -1).any() + + +@needs_op +def test_host_n_valid_blocks_are_converted() -> None: + """Callers hand either the device int32 step view or the host per_token_valid_blocks() (int64) tensor.""" + scores, nvb = _inputs(8192, 512, False, 5) + got = m._tsel_select(scores, nvb.cpu().long(), 0, 1) + assert torch.equal(got, _ref(scores, nvb, 0, 1)) + + +def test_outside_the_gate_keeps_the_built_in_select() -> None: + scores, nvb = _inputs(4096, m.TSEL_MIN_ROWS - 1, False, 3) + assert m._tsel_select(scores, nvb, 0, 1) is None + wide = torch.zeros(1, m.TSEL_MAX_BLOCKS + 1, 512, device="cuda") + assert m._tsel_select(wide, torch.ones(512, dtype=torch.int32, device="cuda"), 0, 1) is None + assert m._tsel_select(scores.transpose(1, 2).contiguous().transpose(1, 2), nvb, 0, 1) is None + + +def test_register_sets_the_hook_once_and_only_where_it_exists() -> None: + indexer = types.SimpleNamespace(selector_hook=None) + m._tsel_register(types.SimpleNamespace(indexer=indexer)) + assert indexer.selector_hook is m._tsel_select + other = object() + indexer.selector_hook = other + m._tsel_register(types.SimpleNamespace(indexer=indexer)) + assert indexer.selector_hook is other + m._tsel_register( + types.SimpleNamespace(indexer=types.SimpleNamespace()) + ) # an indexer without the hook + m._tsel_register(types.SimpleNamespace()) # a backend without an indexer + + +def test_the_shared_indexer_hook_is_inert_unset_or_declining() -> None: + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_indexer import MsaIndexer + + cfg = types.SimpleNamespace(init_blocks=0, local_blocks=1, num_index_heads=1, num_kv_heads=1) + indexer = MsaIndexer(cfg) + assert indexer.selector_hook is None + scores, nvb = _inputs(4096, 300, False, 11) + ref = _ref(scores, nvb, 0, 1) + assert torch.equal(indexer._select(scores, nvb), ref) + indexer.selector_hook = lambda *a: None + assert torch.equal(indexer._select(scores, nvb), ref) + + +@needs_op +def test_fake_impl_matches_the_op() -> None: + """The registered fake (catalog m31_fake_ops) returns what the op returns (nothing) and leaves out as is.""" + from torch._subclasses.fake_tensor import FakeTensorMode + + from tensorrt_llm._torch._experimental.modeling_v2.catalog.fakes import m31_fake_ops + + assert "minimax_m31_select_prefill_tile" in m31_fake_ops._REGISTERED + scores, nvb = _inputs(16384, 1160, False, 7) + out = torch.empty(1160, 1, 16, dtype=torch.int32, device="cuda") # (rows, heads, 16) + assert op(scores, nvb, out, 0, 1) is None + with FakeTensorMode(allow_non_fake_inputs=True) as mode: + fout = mode.from_tensor(out) + assert op(mode.from_tensor(scores), mode.from_tensor(nvb), fout, 0, 1) is None + assert (tuple(fout.shape), fout.dtype) == (tuple(out.shape), out.dtype) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/gemm/_m3_mxfp8_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m31_tp4_mxfp8_pinned_quantize.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m31_tp4_mxfp8_pinned_quantize.py new file mode 100644 index 000000000000..4b9ee50fc211 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m31_tp4_mxfp8_pinned_quantize.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_pinned_quantize entry of the MiniMax-M3.1 TP4 target (QPIN): every tactic is bitwise the +stock MXFP8 quantize.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.ops.gemm.m3_mxfp8_pinned_quantize import ( # noqa: E501 + m3_mxfp8_pinned_quantize as op, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import MXFP8QuantizeRunner as QR + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_pinned_quantize requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("tactic", [QR.CUTE_DSL, QR.TRTLLM], ids=["cute_dsl", "trtllm"]) +def test_certified_cells_bitwise(cell, tactic) -> None: + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = (torch.randn(cell.spec["m"], cell.spec["k"], generator=g, device="cuda") * 3.0).to( + torch.bfloat16 + ) + with validating(op): + q, sf = op(x, tactic) + rq, rsf = op.reference(x, tactic) + assert torch.equal(q.view(torch.uint8).view(-1), rq.view(torch.uint8).view(-1)) + assert torch.equal(sf.view(-1), rsf.view(-1)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m31_tp4_shdown_pin.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m31_tp4_shdown_pin.py new file mode 100644 index 000000000000..fa7fd4ef66d5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m31_tp4_shdown_pin.py @@ -0,0 +1,238 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The SHDOWN pin of the MiniMax-M3.1 TP4 target: the shared down projection at the rows the decode GEMV does not +serve runs FlashInfer's CuTe DSL MXFP8 GEMM, fixed, instead of the per-start tuned backend.""" + +import os +import subprocess +import sys +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4 import ( + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.ops.gemm.m3_mxfp8_gemm import ( # noqa: E501 + m3_mxfp8_gemm, + m3_mxfp8_gemm_cute_dsl, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import ULP, assert_within_ulp # noqa: E402 — needs the path declared above + +_N, _K = modeling.SHARED_DOWN + + +def test_the_pin_covers_exactly_the_rows_the_gemv_leaves_to_the_gemm(): + pinned = modeling.SHDOWN_PIN_ROWS[modeling.SHARED_DOWN] + gemv = modeling.MX_GEMV_ROWS[modeling.SHARED_DOWN] + assert pinned == frozenset(range(17, 33)) + assert not pinned & gemv + assert set(modeling.SHDOWN_PIN_ROWS) == {modeling.SHARED_DOWN} + + +def test_the_kill_switch_restores_the_tuned_backend(): + code = ( + "from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3." + "minimax_m31_nvfp4__sm_103__tp4 import modeling; print(len(modeling.SHDOWN_PIN_ROWS))" + ) + env = dict(os.environ, TRTLLM_M31_DISABLE_SHDOWN_PIN="1") + out = subprocess.run( + [sys.executable, "-c", code], env=env, capture_output=True, text=True, check=True + ) + assert out.stdout.strip().splitlines()[-1] == "0" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@pytest.mark.parametrize("m", [17, 24, 32]) +def test_the_pinned_backend_is_the_mxfp8_product(m): + """Same E4M3 operands and power-of-two scales as the reference; accumulation order then one bf16 rounding: the + band of the target's MXFP8 GEMM test (2 ulp of the row scale, RMS 0.25 ulp). Also reports the largest difference + from the CUTLASS backend the AutoTuner can pick here, in ulp of the row scale (report only).""" + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(f"shdown{m}".encode()) % 2**31) + act, act_sf = mxfp8(m, _K, g) + w, w_sf = mxfp8(_N, _K, g) + out = m3_mxfp8_gemm_cute_dsl(act, act_sf, w, w_sf, torch.bfloat16) + ref = m3_mxfp8_gemm.reference(act, act_sf, w, w_sf, torch.bfloat16, False) + assert out.shape == (m, _N) and out.dtype == torch.bfloat16 + assert_within_ulp(out, ref, element_ulp=2.0, rms_ulp=0.25) + from tensorrt_llm._torch.custom_ops.torch_custom_ops import FlashInferMXFP8GemmRunner as R + + cutlass = R(torch.bfloat16)([act, act_sf, w, w_sf], tactic=R.CUTLASS) + scale = ref.float().abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + ulps = ((out.float() - cutlass.float()).abs() / scale / ULP[torch.bfloat16]).max().item() + print(f"[shdown-pin] m={m}: max |cute_dsl - cutlass| = {ulps:.2f} ulp of the row scale") + + +class _GenericTaken(Exception): + pass + + +def _decode_graph_dispatch(monkeypatch): + """``_linear`` as inside a decode CUDA graph, with no producer-published MXFP8 activation (the shared expert's + down projection reads the SwiGLU's bf16 output) and the tuned quantize + GEMM replaced by a sentinel.""" + mode = modeling.MxMode( + decode_graph=True, + autotune=False, + tuned_quantize_gemm=True, + flashinfer=True, + native_tuning=False, + ) + monkeypatch.setattr(modeling, "_mx_mode", lambda state: mode) + monkeypatch.setattr( + modeling.m3_prequant, "SLOT", type("Slot", (), {"take": staticmethod(lambda x: None)})() + ) + + def generic(*args, **kwargs): + raise _GenericTaken + + monkeypatch.setattr(modeling, "_mxfp8_linear", generic) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@pytest.mark.parametrize("m", [17, 24, 32]) +def test_the_bf16_shared_down_activation_takes_the_pin(monkeypatch, m): + """The production shared down projection: a bf16 activation at a pinned row count runs the native MXFP8 quantize + and the CuTe DSL GEMM (bitwise), never the tuned quantize + GEMM a persisted tactic cache steers.""" + _decode_graph_dispatch(monkeypatch) + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(f"shdown-bf16-{m}".encode()) % 2**31) + w, w_sf = mxfp8(_N, _K, g) + core = type("Core", (), {"w": {"dn_w": w, "dn_s": w_sf}, "mxfp8": None})() + x = torch.randn(m, _K, generator=g, device="cuda", dtype=torch.bfloat16) + out, pushed = modeling._linear(core, x, "dn_w", "dn_s", False) + act, act_sf = torch.ops.trtllm.mxfp8_quantize(x, True) + expect = m3_mxfp8_gemm_cute_dsl(act.view(m, _K), act_sf, w, w_sf, torch.bfloat16) + assert not pushed + assert torch.equal(out, expect) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_unpinned_bf16_rows_keep_the_tuned_backend(monkeypatch): + _decode_graph_dispatch(monkeypatch) + g = torch.Generator(device="cuda").manual_seed(33) + w, w_sf = mxfp8(_N, _K, g) + core = type("Core", (), {"w": {"dn_w": w, "dn_s": w_sf}, "mxfp8": None})() + x = torch.randn(33, _K, generator=g, device="cuda", dtype=torch.bfloat16) + monkeypatch.setattr(modeling, "QPIN", False) # QPIN's own tests below + with pytest.raises(_GenericTaken): + modeling._linear(core, x, "dn_w", "dn_s", False) + + +class _SwqCore: + """The decode target's core as ``_swiglu_mlp`` reads it: weights, the MXFP8 state (unused: ``_mx_mode`` is + patched) and the SwiGLU-OAI scalars.""" + + def __init__(self, gu_out, w, w_sf): + self.w = {"sh_gu_w": gu_out, "sh_gu_s": None, "sh_dn_w": w, "sh_dn_s": w_sf} + self.mxfp8 = None + self.swiglu_limit, self.swiglu_alpha, self.swiglu_beta = 7.0, 1.702, 1.0 + + +def _swq_dispatch(monkeypatch, m, seed): + """Run DecodeTarget._swiglu_mlp at m rows with the gate_up projection replaced by a fixed bf16 output; returns + (output, the pinned GEMM's activation sources, the expected product).""" + _decode_graph_dispatch(monkeypatch) + monkeypatch.setattr( + modeling, "QPIN", False + ) # the gate_up stub below stands in for the generic path + # A real slot: SWQ publishes the MXFP8 activation that the down projection takes. + monkeypatch.setattr(modeling.m3_prequant, "SLOT", modeling.m3_prequant.PrequantSlot()) + g = torch.Generator(device="cuda").manual_seed(seed) + gu = (torch.randn(m, 2 * _K, generator=g, device="cuda") * 4.0).to(torch.bfloat16) + w, w_sf = mxfp8(_N, _K, g) + core = _SwqCore(gu, w, w_sf) + + def gate_up_or_raise(core_, x, w_, s_, mode): + if w_ is core_.w["sh_gu_w"]: + return gu + raise _GenericTaken + + monkeypatch.setattr(modeling, "_mxfp8_linear", gate_up_or_raise) + sources = [] + pinned = modeling._shdown_pinned + + def recording(act, act_scale, w_, s_, source): + sources.append(source) + return pinned(act, act_scale, w_, s_, source) + + monkeypatch.setattr(modeling, "_shdown_pinned", recording) + h = torch.zeros(m, 6144, device="cuda", dtype=torch.bfloat16) + target = type("Target", (), {"core": core})() + out, pushed = modeling.DecodeTarget._swiglu_mlp(target, h, "sh", False) + act = torch.ops.trtllm.silu_and_mul(gu, swiglu_limit=7.0, swiglu_alpha=1.702, swiglu_beta=1.0) + aq, asf = torch.ops.trtllm.mxfp8_quantize(act, True) + expect = m3_mxfp8_gemm_cute_dsl(aq.view(m, _K), asf, w, w_sf, torch.bfloat16) + return out, pushed, sources, expect + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@pytest.mark.parametrize("m", [17, 24, 32]) +def test_swq_reaches_the_pinned_gemm_on_its_mxfp8_activation(monkeypatch, m): + """SWQ at 17..32 rows: the down projection takes the prequant branch (the pin's '(mxfp8 activation)' source) and + its product is bitwise the bf16 SwiGLU + quantize + pinned GEMM path it replaces.""" + out, pushed, sources, expect = _swq_dispatch( + monkeypatch, m, zlib.crc32(f"swq{m}".encode()) % 2**31 + ) + assert sources == ["mxfp8"] and not pushed + assert torch.equal(out, expect) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_swq_kill_switch_keeps_the_bf16_pin_path(monkeypatch): + """SWQ off: the same rows take the bf16 SwiGLU and reach the pin with its '(bf16 activation)' source, bitwise + the same product.""" + monkeypatch.setattr(modeling, "SWQ_ROWS", {}) + out, pushed, sources, expect = _swq_dispatch(monkeypatch, 24, 24) + assert sources == ["bf16"] and not pushed + assert torch.equal(out, expect) + + +def test_swq_rows_are_the_pinned_rows(): + assert modeling.SWQ_ROWS[modeling.SHARED_DOWN] == frozenset(range(17, 33)) + assert modeling.SWQ_ROWS[modeling.SHARED_DOWN] <= modeling.SHDOWN_PIN_ROWS[modeling.SHARED_DOWN] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@pytest.mark.parametrize( + "m, k, tactic", [(16, 6144, "cute"), (32, 3072, "cute"), (33, 768, "cute"), (32, 768, "trt")] +) +def test_qpin_quantize_is_fixed_and_the_product_bitwise(monkeypatch, m, k, tactic): + """QPIN: a bf16 projection off the decode kernels' rows quantizes on the fixed tactic (CuTe DSL; TRT-LLM at + (32, 768)) and runs the tuned GEMM; the product is bitwise the tuned quantize + GEMM's (both quantize tactics + give the same bytes).""" + from tensorrt_llm._torch.custom_ops.torch_custom_ops import MXFP8QuantizeRunner as QR + + _decode_graph_dispatch(monkeypatch) + monkeypatch.setattr(modeling, "SHDOWN_PIN_ROWS", {}) + seen = [] + pinned_q = modeling.m3_mxfp8_pinned_quantize + + def recording(x, t): + seen.append(t) + return pinned_q(x, t) + + monkeypatch.setattr(modeling, "m3_mxfp8_pinned_quantize", recording) + g = torch.Generator(device="cuda").manual_seed(m * 7 + k) + w, w_sf = mxfp8(6144, k, g) + core = type("Core", (), {"w": {"dn_w": w, "dn_s": w_sf}, "mxfp8": None})() + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + out, pushed = modeling._linear(core, x, "dn_w", "dn_s", False) + assert seen == [QR.TRTLLM if tactic == "trt" else QR.CUTE_DSL] and not pushed + aq, asf = torch.ops.trtllm.mxfp8_quantize(x, True) + expect = m3_mxfp8_gemm(aq.view(m, k), asf, w, w_sf, torch.bfloat16, False) + assert torch.equal(out, expect) + + +def test_qpin_kill_switch(): + code = ( + "from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3." + "minimax_m31_nvfp4__sm_103__tp4 import modeling; print(modeling.QPIN)" + ) + env = dict(os.environ, TRTLLM_M31_DISABLE_QPIN="1") + out = subprocess.run( + [sys.executable, "-c", code], env=env, capture_output=True, text=True, check=True + ) + assert out.stdout.strip().splitlines()[-1] == "False" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m31_unfused_index.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m31_unfused_index.py new file mode 100644 index 000000000000..a3a5098cb9c3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m31_unfused_index.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test: the unfused sparse layers' separate MXFP8 index projection (``fuse_qkv_index_projection=False``) of the +tp4 target equals the index columns of the fused five-way projection on the same checkpoint tensors. + +Its weight rows are the five-way weight's index rows byte for byte, its block scales are applied (not dropped), and +on every engine MXFP8 backend its output is bitwise the five-way GEMM's index columns. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4 import ( + weights as tw, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.ops.gemm.m3_mxfp8_linear import ( # noqa: E501 + m3_mxfp8_linear as op, +) + +assert torch.cuda.is_available(), "the MXFP8 index projection requires a CUDA device" + +TP = 4 +HIDDEN = 6144 +BLOCK = 32 +# Text-model rows of the attention projections (64 query / 4 KV heads, 4 index heads, head and index dim 128). +ROWS = {"q": 64 * 128, "k": 4 * 128, "v": 4 * 128, "index_q": 4 * 128, "index_k": 128} +_FLAGS = { + "native": (False, False, False, False), + "flashinfer": (False, False, True, False), + "tuned": (True, False, False, False), +} + + +def _checkpoint(seed: int): + """E4M3 weights and UE8M0 block scales per projection; the scales spread over 2^-6 .. 2^5 (bytes 121..132).""" + g = torch.Generator(device="cuda").manual_seed(seed) + out = {} + for name, rows in ROWS.items(): + w = (torch.randn(rows, HIDDEN, generator=g, device="cuda") * 2).to(torch.float8_e4m3fn) + s = torch.randint( + 121, 133, (rows, HIDDEN // BLOCK), generator=g, device="cuda", dtype=torch.uint8 + ) + out[name] = (w.cpu(), s.cpu()) + return out + + +def _ns(rank: int) -> SimpleNamespace: + return SimpleNamespace(tp=TP, rank=rank, dev=torch.device("cuda"), head_dim=128, index_dim=128) + + +def _packs(ck, rank): + d = _ns(rank) + names = ("q", "k", "v", "index_q", "index_k") + w5 = tw._qkv_sparse_w(tuple(ck[n][0] for n in names), d) + s5 = tw._qkv_sparse_s( + tuple(ck[n][1] for n in names), d + ) # UE8M0 bytes, as the checkpoint stores them + wi = tw._index_qk_w((ck["index_q"][0], ck["index_k"][0]), d) + si = tw._index_qk_s((ck["index_q"][1], ck["index_k"][1]), d) + return w5, s5.view(torch.uint8), wi, si.view(torch.uint8) + + +@pytest.mark.parametrize("rank", [0, TP - 1]) +def test_weight_rows_are_the_five_way_index_rows(rank) -> None: + ck = _checkpoint(1234 + rank) + w5, _, wi, _ = _packs(ck, rank) + n = wi.shape[0] + assert wi.dtype == torch.float8_e4m3fn and wi.shape == (n, HIDDEN) + assert torch.equal(wi.view(torch.uint8), w5[-n:].view(torch.uint8)) + + +@pytest.mark.parametrize("path", list(_FLAGS)) +@pytest.mark.parametrize("rows", [1, 4, 8, 32, 64, 256]) +def test_unfused_output_is_bitwise_the_five_way_index_columns(path, rows) -> None: + if path == "flashinfer" and getattr(torch.ops.trtllm, "flashinfer_mm_mxfp8", None) is None: + pytest.skip("this build has no trtllm::flashinfer_mm_mxfp8") + ck = _checkpoint(77) + w5, s5, wi, si = _packs(ck, 0) + n = wi.shape[0] + g = torch.Generator(device="cuda").manual_seed(rows) + x = torch.randn(rows, HIDDEN, generator=g, device="cuda").to(torch.bfloat16) + alpha = torch.ones(1, dtype=torch.float32, device="cuda") + out5 = op(x, w5, s5, alpha, torch.bfloat16, *_FLAGS[path]) + outi = op(x, wi, si, alpha, torch.bfloat16, *_FLAGS[path]) + assert outi.shape == (rows, n) + assert torch.equal(outi, out5[:, -n:]), ( + f"{path}: max |diff| {(outi.float() - out5[:, -n:].float()).abs().max().item()}" + ) + + +def test_block_scales_are_applied() -> None: + """The output follows the op's reference with the checkpoint scales and differs from a unit-scale (dropped) + projection by far more than rounding.""" + ck = _checkpoint(5) + _, _, wi, si = _packs(ck, 0) + unit = tw._index_qk_s( + (torch.full_like(ck["index_q"][1], 127), torch.full_like(ck["index_k"][1], 127)), _ns(0) + ).view(torch.uint8) + x = torch.randn(8, HIDDEN, device="cuda").to(torch.bfloat16) + alpha = torch.ones(1, dtype=torch.float32, device="cuda") + out = op(x, wi, si, alpha, torch.bfloat16, *_FLAGS["native"]) + ref = op.reference(x, wi, si, alpha, torch.bfloat16, *_FLAGS["native"]) + dropped = op(x, wi, unit, alpha, torch.bfloat16, *_FLAGS["native"]) + scale = ref.float().abs().mean() + # The engine GEMM rounds once to BF16 after an fp32 accumulation in its own order: within 2 BF16 ulps (2^-7 + # relative) of the reference per element. + assert bool( + ((out.float() - ref.float()).abs() <= 2.0**-7 * ref.float().abs() + 1e-3 * scale).all() + ) + assert (out.float() - dropped.float()).abs().mean() > 0.5 * scale diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/moe/test_modeling_v2_m31_tp4_moe_early0n.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/moe/test_modeling_v2_m31_tp4_moe_early0n.py new file mode 100644 index 000000000000..56adf9d0869d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/moe/test_modeling_v2_m31_tp4_moe_early0n.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The TP4 MoE engine's early group-0 issue depth (MOE_EARLY0N): 6 stages at 4 tokens and 3 at 8, measured faster with +bit-identical outputs; every other token count keeps the kernel default; TRTLLM_M31_DISABLE_MOE_EARLY0N restores the +default everywhere.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.ops.m3_moe import ( + op as moe_op, +) + + +def test_early0n_table_scope() -> None: + assert moe_op.EARLY0_N_BY_M == (0, 0, 0, 6, 0, 0, 0, 3) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the engine state compiles the kernel") +@pytest.mark.parametrize(("num_tokens", "expected"), [(1, None), (4, 6), (8, 3)]) +@pytest.mark.parametrize("killed", [False, True]) +def test_early0n_config_and_kill_switch( + monkeypatch: pytest.MonkeyPatch, num_tokens: int, expected, killed: bool +) -> None: + monkeypatch.setattr(moe_op, "m31_disabled", lambda item_id: killed and item_id == "MOE_EARLY0N") + dev = torch.device("cuda") + sms = torch.cuda.get_device_properties(dev).multi_processor_count + state = moe_op.M3MoeState(dev, num_tokens=num_tokens, num_ctas=sms - 24, shared=True) + assert state.config.get("early0_n") == (None if killed else expected) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/moe/test_modeling_v2_m31_tp4_moe_m4cfg.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/moe/test_modeling_v2_m31_tp4_moe_m4cfg.py new file mode 100644 index 000000000000..f28f315a1365 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/moe/test_modeling_v2_m31_tp4_moe_m4cfg.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The TP4 MoE engine's 4-token overrides (MOE_M4CFG): the B-image hand-off plus 8 ring stages, measured faster with +bit-identical outputs; other token counts keep the tables; TRTLLM_M31_DISABLE_MOE_M4CFG restores the tables.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.ops.m3_moe import ( + op as moe_op, +) + + +def test_m4cfg_scope() -> None: + assert moe_op.M4_CFG_BY_M == {4: {"b2img": 1, "stages": 8}} + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the engine state compiles the kernel") +@pytest.mark.parametrize("killed", [False, True]) +def test_m4cfg_config_and_kill_switch(monkeypatch: pytest.MonkeyPatch, killed: bool) -> None: + monkeypatch.setattr(moe_op, "m31_disabled", lambda item_id: killed and item_id == "MOE_M4CFG") + dev = torch.device("cuda") + sms = torch.cuda.get_device_properties(dev).multi_processor_count + state = moe_op.M3MoeState(dev, num_tokens=4, num_ctas=sms - 24, shared=True) + if killed: + assert state.config["b2img"] == moe_op.B2IMG_BY_M[3] and "stages" not in state.config + else: + assert state.config["b2img"] == 1 and state.config["stages"] == 8 + assert state.mod.B2IMG and state.mod.STAGES == 8 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/norm/test_modeling_v2_m31_ctxnormq.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/norm/test_modeling_v2_m31_ctxnormq.py new file mode 100644 index 000000000000..48a1b49b0e96 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/norm/test_modeling_v2_m31_ctxnormq.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for CTXNORMQ: the prefill boundary's residual add + RMSNorm that also writes the next projection's MXFP8 +activation and the routed experts' NVFP4 input (torch.ops.trtllm.minimax_m31_ctx_residual_rmsnorm_mxfp8). + +The quantized outputs must equal mxfp8_quantize / fp4_quantize of the op's own norm bit for bit; the norm and the +updated residual are checked against an fp32 reference. Skipped when the op is not built (no overlay library). +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4 import ( + modeling as m, +) + +assert torch.cuda.is_available(), "CTXNORMQ requires a CUDA device" + +H = 6144 +EPS = 1e-6 +FP4_SCALE = 448.0 * 6.0 / 8.0 + +op = m._ctxnormq_op() +pytestmark = pytest.mark.skipif( + op is None, reason="minimax_m31_ctx_residual_rmsnorm_mxfp8 is not built" +) + + +def _operands(rows, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + inter = torch.randn(rows, H, generator=g, device="cuda").to(torch.bfloat16) + residual = (2 * torch.randn(rows, H, generator=g, device="cuda")).to(torch.bfloat16) + weight = (1.0 + 0.1 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + return inter, residual, weight + + +@pytest.mark.parametrize("rows", [1024, 1160, 4096]) +@pytest.mark.parametrize("fp4", [False, True], ids=["attn", "moe"]) +def test_quantized_outputs_equal_quantizing_the_norm(rows, fp4) -> None: + inter, residual, weight = _operands(rows, 20261021 + rows) + scale = torch.tensor([FP4_SCALE], device="cuda", dtype=torch.float32) if fp4 else None + ref_updated = (inter.float() + residual.float()).to(torch.bfloat16) + normed, q, sf, f4, f4_sf = op(inter, residual, weight, EPS, True, scale, False) + # The updated residual is written in place into inter: one bf16 rounding of the fp32 sum, exact. + assert torch.equal(inter.view(torch.int16), ref_updated.view(torch.int16)) + u = ref_updated.float() + ref_normed = u * torch.rsqrt(u.pow(2).mean(-1, keepdim=True) + EPS) * weight.float() + # fp32 sum of squares + fast rsqrt vs torch's: at most one bf16 rounding apart. + torch.testing.assert_close(normed.float(), ref_normed, rtol=2**-7, atol=2**-7) + q_ref, sf_ref = torch.ops.trtllm.mxfp8_quantize(normed, True) + assert torch.equal(q.view(torch.uint8), q_ref.view(torch.uint8)) + assert torch.equal(sf.view(torch.uint8)[: sf_ref.numel()], sf_ref.view(torch.uint8).flatten()) + if fp4: + f4_ref, f4_sf_ref = torch.ops.trtllm.fp4_quantize(normed, scale, 16, False, False) + assert torch.equal(f4.view(torch.uint8).flatten(), f4_ref.view(torch.uint8).flatten()) + assert torch.equal( + f4_sf.view(torch.uint8).flatten(), + f4_sf_ref.view(torch.uint8).flatten()[: f4_sf.numel()], + ) + # _routed_chunks splits the NVFP4 scales per row chunk. + assert f4_sf.numel() == rows * H // 16 + + +def test_deterministic() -> None: + outs = [] + for _ in range(3): + inter, residual, weight = _operands(2048, 7) + scale = torch.tensor([FP4_SCALE], device="cuda", dtype=torch.float32) + outs.append((inter,) + tuple(op(inter, residual, weight, EPS, True, scale, False))) + for o in outs[1:]: + assert all( + torch.equal(a.view(torch.uint8), b.view(torch.uint8)) for a, b in zip(outs[0], o) + ) + + +def test_small_boundaries_keep_the_fused_allreduce() -> None: + """Below CTXNORMQ_MIN_ROWS the boundary is not taken (the microbench crossover); decided before any collective.""" + x = torch.zeros(m.CTXNORMQ_MIN_ROWS - 1, H, dtype=torch.bfloat16, device="cuda") + assert m._ctxnormq_allreduce_norm(None, x, x, x[0], None) is None + + +@pytest.mark.parametrize("fp4", [False, True], ids=["attn", "moe"]) +@pytest.mark.parametrize("write_norm", [False, True]) +def test_fake_impl_matches_the_op(fp4, write_norm) -> None: + """The registered fake (catalog m31_fake_ops) gives every output the op's shape and dtype.""" + from torch._subclasses.fake_tensor import FakeTensorMode + + from tensorrt_llm._torch._experimental.modeling_v2.catalog.fakes import m31_fake_ops + + assert "minimax_m31_ctx_residual_rmsnorm_mxfp8" in m31_fake_ops._REGISTERED + inter, residual, weight = _operands(1160, 7) + scale = torch.tensor([FP4_SCALE], device="cuda") if fp4 else None + real = op(inter.clone(), residual, weight, EPS, write_norm, scale, False) + with FakeTensorMode(allow_non_fake_inputs=True) as mode: + fake = op( + mode.from_tensor(inter), + mode.from_tensor(residual), + mode.from_tensor(weight), + EPS, + write_norm, + None if scale is None else mode.from_tensor(scale), + False, + ) + assert [(tuple(t.shape), t.dtype) for t in fake] == [(tuple(t.shape), t.dtype) for t in real] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/sampling/test_modeling_v2_m31_chain_spec_mc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/sampling/test_modeling_v2_m31_chain_spec_mc.py new file mode 100644 index 000000000000..ed53e4304c7b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/sampling/test_modeling_v2_m31_chain_spec_mc.py @@ -0,0 +1,286 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for REJ_MC: the spec sampler's chain speculative sampling as +``torch.ops.trtllm.minimax_m31_chain_speculative_sampling_mc`` through ``ops/sampling/m31_chain_spec_mc``, installed as +the rejection sampler's override. + +It must equal the stock ``torch.ops.trtllm.chain_speculative_sampling`` bit for bit, tokens and counts +(the microbenchmark generators: realistic / same / far / ties / onehot drafts and targets, K7). The shared hook in +``speculative/interface.rejection_sampling_one_model`` must be stock while no override is set, and fall back to the +stock op whenever the override declines. TRTLLM_M31_DISABLE_REJ_MC=1 installs nothing. +""" + +import os +import subprocess +import sys +import types +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.ops.sampling import ( # noqa: E501 + m31_chain_spec_mc as entry, +) +from tensorrt_llm._torch.speculative import interface as spec_interface + +TARGET = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4" +) +CASES = ("realistic", "same", "far", "ties", "onehot") +op = entry.m31_chain_spec_mc + + +def _gpu(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3): + pytest.skip("REJ_MC is an sm_103 kernel") + if entry.chain_spec_mc_op() is None: + pytest.skip("minimax_m31_chain_speculative_sampling_mc is not built") + if not hasattr(torch.ops.trtllm, "chain_speculative_sampling"): + pytest.skip("the stock chain_speculative_sampling op is not built") + pytest.importorskip("flashinfer.sampling") + + +def _probs(logits, top_p=0.95): + """softmax then FlashInfer's top-p renorm, as the sampler feeds the rejection step.""" + import flashinfer.sampling + + v = logits.shape[-1] + p = torch.softmax(logits, -1) + renormed = flashinfer.sampling.top_p_renorm_probs(p.view(-1, v).contiguous(), top_p) + return renormed.view(logits.shape).contiguous() + + +def _make(case, batch, k, v, seed): + """The microbenchmark input generators.""" + g = torch.Generator(device="cuda").manual_seed(seed) + + def rl(shape, scale=3.0): + return torch.randn(*shape, v, generator=g, device="cuda") * scale + + dl = rl((batch, k)) + if case == "realistic": # independent draft / target + tl = torch.cat([dl, rl((batch, 1))], 1) + rl((batch, k + 1), 1.0) + elif case == "same": # target == draft: all-accept and zero-residual paths + tl = torch.cat([dl, rl((batch, 1))], 1) + elif case == "far": # target peaked elsewhere: early reject + tl = rl((batch, k + 1)) + tl[:, :, :8] += 20 + elif case == "ties": + dl = torch.round(dl * 2) / 2 + tl = torch.round(torch.cat([dl, rl((batch, 1))], 1) * 2) / 2 + elif case == "onehot": + tl = torch.full((batch, k + 1, v), -1e4, device="cuda") + tl[:, :, 1234] = 0 + else: + raise ValueError(case) + dp, tp = _probs(dl), _probs(tl) + ids = torch.multinomial(dp.view(-1, v), 1, generator=g).view(batch, k).to(torch.int32) + seeds = torch.randint(0, 2**31, (batch,), generator=g, device="cuda", dtype=torch.int64) + offsets = torch.randint(0, 2**20, (batch,), generator=g, device="cuda", dtype=torch.int64) + return dp, ids.contiguous(), tp, seeds, offsets + + +def _stock(*args): + return torch.ops.trtllm.chain_speculative_sampling(*args) + + +def _same(got, ref, what): + assert torch.equal(got[0], ref[0]), f"{what}: accepted tokens differ" + assert torch.equal(got[1], ref[1]), f"{what}: token counts differ" + + +@pytest.mark.parametrize("case", CASES) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:40] for c in op.CELLS]) +def test_bitwise_stock_op(cell, case) -> None: + """At each certified (batch, K, vocab): the override equals the stock op (tokens and counts), and the explicit + entry call at ctas_per_row 1 / 4 / 8 gives the same result again.""" + _gpu() + batch, k, v = cell.spec["batch"], cell.spec["k"], cell.spec["v"] + for seed in range(4): + args = _make(case, batch, k, v, 1000 * seed + batch) + got = entry.chain_spec_mc(*args) + assert got is not None, "an in-contract input must be served" + _same(got, _stock(*args), f"{case} batch={batch} seed={seed}") + for ctas in (1, 4, 8): + _same(op(*args, entry.prebuild(args[0].device), ctas), got, f"ctas_per_row={ctas}") + + +def test_graph_replay_equals_eager() -> None: + """Captured in a CUDA graph (workspace prebuilt outside the capture), the replay gives the eager result.""" + _gpu() + args = _make("realistic", 8, 7, 200064, 7) + entry.prebuild(args[0].device) + ref = entry.chain_spec_mc(*args) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + entry.chain_spec_mc(*args) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out = entry.chain_spec_mc(*args) + assert out is not None, "a prebuilt workspace must serve a captured call" + for _ in range(3): + out[0].fill_(-7) + out[1].fill_(-7) + graph.replay() + torch.cuda.synchronize() + _same(out, ref, "graph replay") + + +def test_outside_the_contract_the_override_declines() -> None: + """Non-fp32 probs, non-int32 draft tokens, non-int64 or mis-sized seed / offset, a vocab not divisible by 4, a + misaligned or non-contiguous probs, mismatched shapes, CPU inputs or more requests than MAX_BATCH return None (the + sampler runs the stock op).""" + _gpu() + dp, ids, tp, seeds, offsets = _make("realistic", 2, 3, 32768, 1) + assert entry.chain_spec_mc(dp, ids, tp, seeds, offsets) is not None + assert entry.chain_spec_mc(dp.double(), ids, tp, seeds, offsets) is None + assert entry.chain_spec_mc(dp, ids, tp.half(), seeds, offsets) is None + assert entry.chain_spec_mc(dp, ids.long(), tp, seeds, offsets) is None + assert entry.chain_spec_mc(dp, ids, tp, seeds.int(), offsets) is None + assert entry.chain_spec_mc(dp, ids, tp, seeds[:1], offsets) is None + assert entry.chain_spec_mc(dp, ids, tp, seeds, offsets[:1]) is None + assert entry.chain_spec_mc(dp, ids[:, :2].contiguous(), tp, seeds, offsets) is None + assert entry.chain_spec_mc(dp, ids, tp[:, :3].contiguous(), seeds, offsets) is None + assert entry.chain_spec_mc(dp[..., :-4], ids, tp[..., :-4], seeds, offsets) is None + assert ( + entry.chain_spec_mc( + dp[..., :-2].contiguous(), ids, tp[..., :-2].contiguous(), seeds, offsets + ) + is None + ) + flat = torch.empty(dp.numel() + 1, device="cuda") + assert entry.chain_spec_mc(flat[1:].view(dp.shape), ids, tp, seeds, offsets) is None + assert entry.chain_spec_mc(dp.cpu(), ids.cpu(), tp.cpu(), seeds.cpu(), offsets.cpu()) is None + b = entry.MAX_BATCH + 1 + wide = _make("realistic", b, 1, 32768, 2) + assert entry.chain_spec_mc(*wide) is None + assert entry.ctas_per_row(1, 148) == 16 and entry.ctas_per_row(8, 148) == 16 + assert entry.ctas_per_row(16, 148) == 8 and entry.ctas_per_row(64, 148) == 2 + assert entry.ctas_per_row(200, 148) == 1 + + +@pytest.fixture +def _fake_stock_op(monkeypatch): + """Record the feature FlashInfer fallback; retain real torch tensor semantics.""" + stock = (torch.zeros(2, 4, dtype=torch.int32), torch.ones(2, dtype=torch.int32)) + fake_op = mock.Mock(return_value=(stock[0], None, stock[1] - 1)) + fake_fi = types.SimpleNamespace( + __version__="0.6.4", sampling=types.SimpleNamespace(chain_speculative_sampling=fake_op) + ) + monkeypatch.setattr(spec_interface, "flashinfer", fake_fi, raising=False) + monkeypatch.setattr(spec_interface, "IS_FLASHINFER_AVAILABLE", True) + spec_interface.set_chain_speculative_sampling_override(None) + yield fake_op, stock + spec_interface.set_chain_speculative_sampling_override(None) + + +def _cpu_args(): + return ( + torch.rand(2, 3, 8), + torch.zeros(2, 3, dtype=torch.int32), + torch.rand(2, 4, 8), + torch.zeros(2, dtype=torch.int64), + torch.zeros(2, dtype=torch.int64), + ) + + +def test_hook_unset_is_stock(_fake_stock_op) -> None: + """No override: rejection_sampling_one_model makes exactly the stock op call and returns its result.""" + fake_op, stock = _fake_stock_op + args = _cpu_args() + assert spec_interface._CHAIN_SPEC_OVERRIDE is None + result = spec_interface.rejection_sampling_one_model(*args[:3], seed=args[3], offset=args[4]) + assert result[0] is stock[0] + assert torch.equal(result[1], stock[1]) + fake_op.assert_called_once() + assert all(a is b for a, b in zip(fake_op.call_args.args, args[:3])) + assert fake_op.call_args.kwargs["seed"] is args[3] + assert fake_op.call_args.kwargs["offset"] is args[4] + assert fake_op.call_args.kwargs["deterministic"] is True + assert fake_op.call_args.kwargs["generator"] is None + + +def test_hook_override_result_is_used(_fake_stock_op) -> None: + """An override's result replaces the stock op; the stock op does not run.""" + fake_op, _ = _fake_stock_op + mine = (torch.ones(2, 4, dtype=torch.int32), torch.full((2,), 4, dtype=torch.int32)) + override = mock.Mock(return_value=mine) + spec_interface.set_chain_speculative_sampling_override(override) + args = _cpu_args() + assert ( + spec_interface.rejection_sampling_one_model(*args[:3], seed=args[3], offset=args[4]) is mine + ) + override.assert_called_once_with(*args) + fake_op.assert_not_called() + + +def test_hook_override_declining_falls_back(_fake_stock_op) -> None: + """An override that returns None leaves the stock op call in place.""" + fake_op, stock = _fake_stock_op + override = mock.Mock(return_value=None) + spec_interface.set_chain_speculative_sampling_override(override) + args = _cpu_args() + result = spec_interface.rejection_sampling_one_model(*args[:3], seed=args[3], offset=args[4]) + assert result[0] is stock[0] + assert torch.equal(result[1], stock[1]) + assert override.call_count == 1 and fake_op.call_count == 1 + + +@pytest.mark.parametrize("state", ["none", "scalar", "mixed", "nondeterministic"]) +def test_ineligible_rng_state_stays_stock(_fake_stock_op, state) -> None: + fake_op, stock = _fake_stock_op + args = _cpu_args() + override = mock.Mock() + spec_interface.set_chain_speculative_sampling_override(override) + seed, offset = args[3:] + if state == "none": + seed = offset = None + elif state == "scalar": + seed, offset = 17, 0 + elif state == "mixed": + offset = 0 + result = spec_interface.rejection_sampling_one_model( + *args[:3], deterministic=state != "nondeterministic", seed=seed, offset=offset + ) + override.assert_not_called() + fake_op.assert_called_once() + assert result[0] is stock[0] + assert torch.equal(result[1], stock[1]) + + +def test_the_target_installs_the_override() -> None: + """With REJ_MC on (the default) and the op built, the target installs ``chain_spec_mc`` once per process.""" + _gpu() + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4 import ( # noqa: E501 + modeling, + ) + + if not modeling.REJ_MC: + pytest.skip("REJ_MC is switched off in this environment") + saved = dict(modeling._REJ_MC_STATE) + modeling._REJ_MC_STATE.clear() + try: + modeling._install_rej_mc() + assert spec_interface._CHAIN_SPEC_OVERRIDE is entry.chain_spec_mc + assert modeling._REJ_MC_STATE["installed"] is True + finally: + spec_interface.set_chain_speculative_sampling_override(None) + modeling._REJ_MC_STATE.clear() + modeling._REJ_MC_STATE.update(saved) + + +def test_the_kill_switch_installs_nothing() -> None: + """TRTLLM_M31_DISABLE_REJ_MC=1 (read at import): REJ_MC is off and the target leaves the rejection hook unset.""" + probe = ( + f"import {TARGET}.modeling as m\n" + "from tensorrt_llm._torch.speculative import interface as i\n" + "m._install_rej_mc()\n" + "print(m.REJ_MC, i._CHAIN_SPEC_OVERRIDE is None, m._REJ_MC_STATE['installed'])" + ) + env = dict(os.environ, TRTLLM_M31_DISABLE_REJ_MC="1") + out = subprocess.run( + [sys.executable, "-c", probe], env=env, capture_output=True, text=True, check=True + ) + assert out.stdout.strip().splitlines()[-1] == "False True False" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/sampling/test_modeling_v2_m31_softmax_mc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/sampling/test_modeling_v2_m31_softmax_mc.py new file mode 100644 index 000000000000..8b2de70ecddc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/sampling/test_modeling_v2_m31_softmax_mc.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for SOFTMAX_MC: the spec sampler's softmax as +``torch.ops.trtllm.minimax_m31_softmax_mc`` through ``ops/sampling/m31_softmax_mc``, installed as the sampler's softmax +override. + +It must equal ``flashinfer.sampling.softmax(logits, temperature)`` bit for bit (the microbenchmark generators: randn / +peaked / ties / near_uniform / all_equal / masked / wide rows, temperature 1 / per-row mixed / scalar). The shared +sampler hook must be byte-for-byte stock while no override is set, and fall back to the stock call whenever the +override declines. TRTLLM_M31_DISABLE_SOFTMAX_MC=1 installs nothing. +""" + +import os +import subprocess +import sys +import types +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.ops.sampling import ( # noqa: E501 + m31_softmax_mc as entry, +) +from tensorrt_llm._torch.pyexecutor.sampler.ops import flashinfer as fi_ops + +TARGET = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4" +) +CLASSES = ("randn", "peaked", "ties", "near_uniform", "all_equal", "masked", "wide") +TEMPS = ("t1", "mixed", "scalar") +op = entry.m31_softmax_mc + + +def _gpu(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3): + pytest.skip("SOFTMAX_MC is an sm_103 kernel") + if entry.softmax_mc_op() is None: + pytest.skip("minimax_m31_softmax_mc is not built") + pytest.importorskip("flashinfer.sampling") + + +def _make(cls, rows, d, seed): + """The microbenchmark logit generators.""" + g = torch.Generator(device="cuda").manual_seed(seed) + if cls == "randn": + return (torch.randn(rows, d, generator=g, device="cuda") * 3).contiguous() + if cls == "peaked": # few dominant tokens, long tail (real-like rows) + x = torch.randn(rows, d, generator=g, device="cuda") * 2 + idx = torch.randint(0, d, (rows, 8), generator=g, device="cuda") + x.scatter_(1, idx, 14 + 3 * torch.rand(rows, 8, generator=g, device="cuda")) + return x.contiguous() + if cls == "ties": + return (torch.round(torch.randn(rows, d, generator=g, device="cuda") * 2) / 2).contiguous() + if cls == "near_uniform": + return (torch.randn(rows, d, generator=g, device="cuda") * 1e-3).contiguous() + if cls == "all_equal": + return torch.zeros(rows, d, device="cuda") + if cls == "masked": + x = torch.randn(rows, d, generator=g, device="cuda") * 3 + x[:, d // 2 :] = -1e4 + return x.contiguous() + if cls == "wide": + return (torch.randn(rows, d, generator=g, device="cuda") * 30).contiguous() + raise ValueError(cls) + + +def _temperature(mode, rows, seed): + """FlashInfer's temperature argument: an fp32 [rows] tensor, or a Python float.""" + g = torch.Generator(device="cuda").manual_seed(seed + 3) + if mode == "t1": + return torch.ones(rows, device="cuda") + if mode == "mixed": + return (0.3 + 1.2 * torch.rand(rows, generator=g, device="cuda")).contiguous() + if mode == "scalar": + return 0.7 + raise ValueError(mode) + + +def _fi(logits, temperature): + import flashinfer.sampling + + return flashinfer.sampling.softmax(logits, temperature) + + +def _same(a, b, what): + """Bit for bit.""" + assert torch.equal(a.view(torch.int32), b.view(torch.int32)), f"{what}: fp32 bits differ" + + +@pytest.mark.parametrize("mode", TEMPS) +@pytest.mark.parametrize("cls", CLASSES) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:40] for c in op.CELLS]) +def test_bitwise_flashinfer_softmax(cell, cls, mode) -> None: + """At each certified (rows, d): the override equals FlashInfer's softmax bit for bit, and the explicit entry call + at write_ctas 4 / 8 gives the same bits again.""" + _gpu() + rows, d = cell.spec["rows"], cell.spec["d"] + for seed in range(2): + logits = _make(cls, rows, d, 100 * seed + rows) + temperature = _temperature(mode, rows, seed) + got = entry.softmax_mc(logits, temperature) + assert got is not None, "an in-contract input must be served" + _same(got, _fi(logits, temperature), f"{cls} {mode} rows={rows} seed={seed}") + for write_ctas in (4, 8): + again = op(logits, temperature, entry.prebuild(logits.device), write_ctas) + _same(again, got, f"write_ctas={write_ctas}") + + +def test_no_temperature_is_one() -> None: + """A None temperature (softmax_op's Optional) is FlashInfer's default, 1.0.""" + _gpu() + logits = _make("peaked", 8, 200064, 11) + got = entry.softmax_mc(logits, None) + assert got is not None + _same(got, _fi(logits, None), "temperature None") + + +def test_graph_replay_equals_eager() -> None: + """Captured in a CUDA graph (workspace prebuilt outside the capture), the replay gives the eager bits.""" + _gpu() + logits = _make("peaked", 8, 200064, 5) + temperature = _temperature("mixed", 8, 5) + entry.prebuild(logits.device) + ref = entry.softmax_mc(logits, temperature) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + entry.softmax_mc(logits, temperature) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out = entry.softmax_mc(logits, temperature) + assert out is not None, "a prebuilt workspace must serve a captured call" + for _ in range(3): + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + _same(out, ref, "graph replay") + + +def test_outside_the_contract_the_override_declines() -> None: + """A non-fp32 / misaligned / non-contiguous logits, a d % 4 width, a vocab below FlashInfer's Path A, a mis-sized, + non-fp32 or 2-D temperature, a non-number temperature or more rows than MAX_ROWS return None (the sampler runs + stock FlashInfer).""" + _gpu() + d = 200064 + logits = _make("randn", 4, d, 1) + temperature = torch.full((4,), 0.9, device="cuda") + assert entry.softmax_mc(logits.half(), temperature) is None + assert entry.softmax_mc(logits, temperature.double()) is None + assert entry.softmax_mc(logits, temperature[:3]) is None + assert entry.softmax_mc(logits, temperature.view(4, 1)) is None + assert entry.softmax_mc(logits, "0.9") is None + assert entry.softmax_mc(logits[:, : d - 4], temperature) is None + assert entry.softmax_mc(logits[:, : d - 2].contiguous(), temperature) is None + assert ( + entry.softmax_mc(logits[:, : entry.PATH_A_MIN_VOCAB - 4].contiguous(), temperature) is None + ) + flat = torch.empty(4 * d + 1, device="cuda") + assert entry.softmax_mc(flat[1:].view(4, d), temperature) is None + wide = _make("randn", entry.MAX_ROWS + 1, d, 2) + assert entry.softmax_mc(wide, torch.full((entry.MAX_ROWS + 1,), 0.9, device="cuda")) is None + assert entry.softmax_mc(logits.cpu(), temperature.cpu()) is None + + +@pytest.fixture +def _fake_flashinfer(monkeypatch): + """The sampler ops module with a stand-in flashinfer whose calls are recorded, and no override afterwards.""" + stock = torch.full((2, 8), 0.125) + sampling = types.SimpleNamespace( + softmax=mock.Mock(return_value=stock), + top_p_renorm_probs=mock.Mock(side_effect=lambda probs, top_p: probs), + top_k_mask_logits=mock.Mock(side_effect=lambda logits, top_k: logits), + ) + monkeypatch.setattr(fi_ops, "flashinfer", types.SimpleNamespace(sampling=sampling)) + fi_ops.set_softmax_override(None) + fi_ops.set_top_p_renorm_override(None) + yield sampling, stock + fi_ops.set_softmax_override(None) + + +def test_hook_unset_is_stock(_fake_flashinfer) -> None: + """No override: both ops make exactly the stock flashinfer call and return its result.""" + sampling, stock = _fake_flashinfer + logits, temperature = torch.rand(2, 8), torch.ones(2) + assert fi_ops._SOFTMAX_OVERRIDE is None + assert fi_ops.softmax_op(logits, temperature) is stock + sampling.softmax.assert_called_once_with(logits, temperature, enable_pdl=mock.ANY) + assert fi_ops.compute_probs_from_logits_op(logits, temperature, None, None) is stock + assert sampling.softmax.call_count == 2 + + +def test_hook_override_result_is_used(_fake_flashinfer) -> None: + """An override's tensor replaces the stock call in both ops; the stock kernel does not run, and a top_p renorm + after it gets the override's probs.""" + sampling, _ = _fake_flashinfer + mine = torch.zeros(2, 8) + override = mock.Mock(return_value=mine) + fi_ops.set_softmax_override(override) + logits, temperature = torch.rand(2, 8), torch.ones(2) + assert fi_ops.softmax_op(logits, temperature) is mine + override.assert_called_once_with(logits, temperature) + assert fi_ops.compute_probs_from_logits_op(logits, temperature, None, None) is mine + top_p = torch.full((2,), 0.9) + assert fi_ops.compute_probs_from_logits_op(logits, temperature, None, top_p) is mine + sampling.top_p_renorm_probs.assert_called_once_with(mine, top_p) + assert override.call_count == 3 + sampling.softmax.assert_not_called() + + +def test_hook_override_declining_falls_back(_fake_flashinfer) -> None: + """An override that returns None leaves the stock call in place in both ops.""" + sampling, stock = _fake_flashinfer + override = mock.Mock(return_value=None) + fi_ops.set_softmax_override(override) + logits, temperature = torch.rand(2, 8), torch.ones(2) + assert fi_ops.softmax_op(logits, temperature) is stock + assert fi_ops.compute_probs_from_logits_op(logits, temperature, None, None) is stock + assert override.call_count == 2 and sampling.softmax.call_count == 2 + + +def test_the_target_installs_the_override() -> None: + """With SOFTMAX_MC on (the default) and the op built, the target installs ``softmax_mc`` once per process.""" + _gpu() + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4 import ( # noqa: E501 + modeling, + ) + + if not modeling.SOFTMAX_MC: + pytest.skip("SOFTMAX_MC is switched off in this environment") + saved = dict(modeling._SOFTMAX_MC_STATE) + modeling._SOFTMAX_MC_STATE.clear() + try: + modeling._install_softmax_mc() + assert fi_ops._SOFTMAX_OVERRIDE is entry.softmax_mc + assert modeling._SOFTMAX_MC_STATE["installed"] is True + finally: + fi_ops.set_softmax_override(None) + modeling._SOFTMAX_MC_STATE.clear() + modeling._SOFTMAX_MC_STATE.update(saved) + + +def test_the_kill_switch_installs_nothing() -> None: + """TRTLLM_M31_DISABLE_SOFTMAX_MC=1 (read at import): SOFTMAX_MC is off and the target leaves the sampler hook + unset.""" + probe = ( + f"import {TARGET}.modeling as m\n" + "from tensorrt_llm._torch.pyexecutor.sampler.ops import flashinfer as f\n" + "m._install_softmax_mc()\n" + "print(m.SOFTMAX_MC, f._SOFTMAX_OVERRIDE is None, m._SOFTMAX_MC_STATE['installed'])" + ) + env = dict(os.environ, TRTLLM_M31_DISABLE_SOFTMAX_MC="1") + out = subprocess.run( + [sys.executable, "-c", probe], env=env, capture_output=True, text=True, check=True + ) + assert out.stdout.strip().splitlines()[-1] == "False True False" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/sampling/test_modeling_v2_m31_topp_det.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/sampling/test_modeling_v2_m31_topp_det.py new file mode 100644 index 000000000000..600defd2bff5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/sampling/test_modeling_v2_m31_topp_det.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for TOPP_DET: the spec sampler's top-p renorm as +``torch.ops.trtllm.minimax_m31_topp_renorm_det`` through ``ops/sampling/m31_topp_det``, installed as the sampler's +top-p renorm override. + +On rows with top_p < 1 it must equal FlashInfer's ``top_p_renorm_probs(probs, top_p, is_deterministic=True)`` bit for +bit (the microbenchmark generators: randn / peaked / ties / near_uniform / all_equal rows); +rows with top_p >= 1 keep every token. The shared sampler hook must be byte-for-byte stock while no override is set, and +fall back to the stock call whenever the override declines. TRTLLM_M31_DISABLE_TOPP_DET=1 installs nothing. +""" + +import os +import subprocess +import sys +import types +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.ops.sampling import ( # noqa: E501 + m31_topp_det as entry, +) +from tensorrt_llm._torch.pyexecutor.sampler.ops import flashinfer as fi_ops + +TARGET = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4" +) +CLASSES = ("randn", "peaked", "ties", "near_uniform", "all_equal") +FLT_MIN = torch.finfo(torch.float32).tiny +op = entry.m31_topp_det + + +def _gpu(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3): + pytest.skip("TOPP_DET is an sm_103 kernel") + if entry.topp_det_op() is None: + pytest.skip("minimax_m31_topp_renorm_det is not built") + pytest.importorskip("flashinfer.sampling") + + +def _make(cls, rows, d, seed): + """The microbenchmark row generators: softmax of the class's logits.""" + g = torch.Generator(device="cuda").manual_seed(seed) + if cls == "randn": + logits = torch.randn(rows, d, generator=g, device="cuda") * 3 + elif cls == "peaked": # few dominant tokens, long tail (real-like draft rows) + logits = torch.randn(rows, d, generator=g, device="cuda") * 2 + idx = torch.randint(0, d, (rows, 8), generator=g, device="cuda") + logits.scatter_(1, idx, 14 + 3 * torch.rand(rows, 8, generator=g, device="cuda")) + elif cls == "ties": + logits = torch.round(torch.randn(rows, d, generator=g, device="cuda") * 2) / 2 + elif cls == "near_uniform": + logits = torch.randn(rows, d, generator=g, device="cuda") * 1e-3 + elif cls == "all_equal": + logits = torch.zeros(rows, d, device="cuda") + else: + raise ValueError(cls) + return torch.softmax(logits, -1).contiguous() + + +def _top_p(mode, rows, seed): + g = torch.Generator(device="cuda").manual_seed(seed + 7) + if mode == "p95": + return torch.full((rows,), 0.95, device="cuda") + if mode == "mixed": + return (0.05 + 0.95 * torch.rand(rows, generator=g, device="cuda")).contiguous() + raise ValueError(mode) + + +def _fi_det(probs, top_p): + import flashinfer.sampling + + return flashinfer.sampling.top_p_renorm_probs(probs, top_p, is_deterministic=True) + + +def _same(a, b, what): + """Bit for bit.""" + assert torch.equal(a.view(torch.int32), b.view(torch.int32)), f"{what}: fp32 bits differ" + + +@pytest.mark.parametrize("mode", ["p95", "mixed"]) +@pytest.mark.parametrize("cls", CLASSES) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:40] for c in op.CELLS]) +def test_bitwise_flashinfer_det_below_one(cell, cls, mode) -> None: + """Rows with top_p < 1 at each certified (rows, d): the override equals FlashInfer's det renorm bit for bit, and + the explicit entry call gives the same bits again (deterministic).""" + _gpu() + rows, d = cell.spec["rows"], cell.spec["d"] + for seed in range(2): + probs = _make(cls, rows, d, 1000 * rows + seed) + top_p = _top_p(mode, rows, seed) + got = entry.top_p_renorm_det(probs, top_p) + assert got is not None, "an in-contract input must be served" + ref = _fi_det(probs, top_p) + _same(got, ref, f"{cls} {mode} rows={rows} seed={seed}") + again = op( + probs, + top_p, + entry.prebuild(probs.device), + entry.ctas_per_row(rows, entry.sm_count(probs.device)), + ) + _same(again, got, "repeat") + + +@pytest.mark.parametrize("cls", CLASSES) +def test_keep_all_rows_in_a_mixed_batch(cls) -> None: + """A batch of 8 rows, a third at top_p 1: the top_p 1 rows keep every element whose renormed value is a normal float + (output nonzero <=> ftz(p * norm) != 0, the --use_fast_math semantics FlashInfer's apply shares) and sum to 1 within + 1e-5; the top_p < 1 rows are FlashInfer's det renorm bit for bit.""" + _gpu() + rows, d = 8, 200064 + probs = _make(cls, rows, d, 77) + top_p = _top_p("mixed", rows, 3) + keep = torch.zeros(rows, dtype=torch.bool, device="cuda") + keep[::3] = True + top_p[keep] = 1.0 + got = entry.top_p_renorm_det(probs, top_p) + assert got is not None + ref = _fi_det(probs, top_p) + _same(got[~keep], ref[~keep], f"{cls}: top_p < 1 rows") + for r in keep.nonzero().flatten().tolist(): + p64 = probs[r].double() + norm = 1.0 / float(p64[probs[r] >= FLT_MIN].sum()) + # kept <=> ftz(p * norm) != 0; elements within 1e-6 of the FLT_MIN boundary are left to the rounding + live = p64 * norm >= FLT_MIN * (1 + 1e-6) + dead = probs[r] < FLT_MIN + assert bool((got[r][live] != 0).all()), f"{cls} row {r}: a kept token was dropped" + assert bool((got[r][dead] == 0).all()), f"{cls} row {r}: a sub-FLT_MIN token is nonzero" + total = float(got[r].double().sum()) + assert abs(total - 1.0) <= 1e-5, f"{cls} row {r}: sum {total}" + + +def test_graph_replay_equals_eager() -> None: + """Captured in a CUDA graph (workspace prebuilt outside the capture), the replay gives the eager bits.""" + _gpu() + probs = _make("peaked", 8, 200064, 5) + top_p = _top_p("mixed", 8, 5) + entry.prebuild(probs.device) + ref = entry.top_p_renorm_det(probs, top_p) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + entry.top_p_renorm_det(probs, top_p) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out = entry.top_p_renorm_det(probs, top_p) + assert out is not None, "a prebuilt workspace must serve a captured call" + for _ in range(3): + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + _same(out, ref, "graph replay") + + +def test_outside_the_contract_the_override_declines() -> None: + """A float top_p, a non-fp32 / misaligned / non-contiguous probs, a d % 4 width, a mis-sized top_p or more rows + than the workspace holds return None (the sampler runs stock FlashInfer).""" + _gpu() + d = 200064 + probs = _make("randn", 4, d, 1) + top_p = torch.full((4,), 0.9, device="cuda") + assert entry.top_p_renorm_det(probs, 0.9) is None + assert entry.top_p_renorm_det(probs.half(), top_p) is None + assert entry.top_p_renorm_det(probs, top_p.double()) is None + assert entry.top_p_renorm_det(probs, top_p[:3]) is None + assert entry.top_p_renorm_det(probs[:, : d - 2], top_p) is None + assert entry.top_p_renorm_det(probs[:, : d - 2].contiguous(), top_p) is None + flat = torch.empty(4 * d + 1, device="cuda") + assert entry.top_p_renorm_det(flat[1:].view(4, d), top_p) is None + wide = _make("randn", entry.MAX_ROWS + 1, d, 2) + assert ( + entry.top_p_renorm_det(wide, torch.full((entry.MAX_ROWS + 1,), 0.9, device="cuda")) is None + ) + assert entry.ctas_per_row(1, 148) == 16 and entry.ctas_per_row(8, 148) == 16 + assert entry.ctas_per_row(16, 148) == 8 and entry.ctas_per_row(64, 148) == 2 + assert entry.ctas_per_row(200, 148) == 1 + + +@pytest.fixture +def _fake_flashinfer(monkeypatch): + """The sampler ops module with a stand-in flashinfer whose calls are recorded, and no override afterwards.""" + stock = torch.full((2, 8), 0.125) + sampling = types.SimpleNamespace( + top_p_renorm_probs=mock.Mock(return_value=stock), + softmax=mock.Mock(side_effect=lambda logits, temps, enable_pdl=False: logits.softmax(-1)), + top_k_mask_logits=mock.Mock(side_effect=lambda logits, top_k: logits), + ) + monkeypatch.setattr(fi_ops, "flashinfer", types.SimpleNamespace(sampling=sampling)) + fi_ops.set_top_p_renorm_override(None) + yield sampling, stock + fi_ops.set_top_p_renorm_override(None) + + +def test_hook_unset_is_stock(_fake_flashinfer) -> None: + """No override: both ops make exactly the stock flashinfer call and return its result.""" + sampling, stock = _fake_flashinfer + probs, top_p = torch.rand(2, 8), torch.full((2,), 0.9) + assert fi_ops._TOP_P_RENORM_OVERRIDE is None + assert fi_ops.top_p_renorm_probs_op(probs, top_p) is stock + sampling.top_p_renorm_probs.assert_called_once_with(probs, top_p) + assert ( + fi_ops.compute_probs_from_logits_op(torch.rand(2, 8), torch.ones(2), None, top_p) is stock + ) + assert sampling.top_p_renorm_probs.call_count == 2 + + +def test_hook_override_result_is_used(_fake_flashinfer) -> None: + """An override's tensor replaces the stock call in both ops; the stock kernel does not run.""" + sampling, _ = _fake_flashinfer + mine = torch.zeros(2, 8) + override = mock.Mock(return_value=mine) + fi_ops.set_top_p_renorm_override(override) + probs, top_p = torch.rand(2, 8), torch.full((2,), 0.9) + assert fi_ops.top_p_renorm_probs_op(probs, top_p) is mine + override.assert_called_once_with(probs, top_p) + assert fi_ops.compute_probs_from_logits_op(torch.rand(2, 8), torch.ones(2), None, top_p) is mine + assert override.call_count == 2 + sampling.top_p_renorm_probs.assert_not_called() + + +def test_hook_override_declining_falls_back(_fake_flashinfer) -> None: + """An override that returns None leaves the stock call in place; without top_p the override is never asked.""" + sampling, stock = _fake_flashinfer + override = mock.Mock(return_value=None) + fi_ops.set_top_p_renorm_override(override) + probs, top_p = torch.rand(2, 8), torch.full((2,), 0.9) + assert fi_ops.top_p_renorm_probs_op(probs, top_p) is stock + assert ( + fi_ops.compute_probs_from_logits_op(torch.rand(2, 8), torch.ones(2), None, top_p) is stock + ) + assert override.call_count == 2 and sampling.top_p_renorm_probs.call_count == 2 + fi_ops.compute_probs_from_logits_op(torch.rand(2, 8), torch.ones(2), None, None) + assert override.call_count == 2 + + +def test_the_target_installs_the_override() -> None: + """With TOPP_DET on (the default) and the op built, the target installs ``top_p_renorm_det`` once per process.""" + _gpu() + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4 import ( # noqa: E501 + modeling, + ) + + if not modeling.TOPP_DET: + pytest.skip("TOPP_DET is switched off in this environment") + saved = dict(modeling._TOPP_DET_STATE) + modeling._TOPP_DET_STATE.clear() + try: + modeling._install_topp_det() + assert fi_ops._TOP_P_RENORM_OVERRIDE is entry.top_p_renorm_det + assert modeling._TOPP_DET_STATE["installed"] is True + finally: + fi_ops.set_top_p_renorm_override(None) + modeling._TOPP_DET_STATE.clear() + modeling._TOPP_DET_STATE.update(saved) + + +def test_the_kill_switch_installs_nothing() -> None: + """TRTLLM_M31_DISABLE_TOPP_DET=1 (read at import): TOPP_DET is off and the target leaves the sampler hook unset.""" + probe = ( + f"import {TARGET}.modeling as m\n" + "from tensorrt_llm._torch.pyexecutor.sampler.ops import flashinfer as f\n" + "m._install_topp_det()\n" + "print(m.TOPP_DET, f._TOP_P_RENORM_OVERRIDE is None, m._TOPP_DET_STATE['installed'])" + ) + env = dict(os.environ, TRTLLM_M31_DISABLE_TOPP_DET="1") + out = subprocess.run( + [sys.executable, "-c", probe], env=env, capture_output=True, text=True, check=True + ) + assert out.stdout.strip().splitlines()[-1] == "False True False" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/spec/_m3_spec_inputs.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/spec/_m3_spec_inputs.py new file mode 100644 index 000000000000..1ab56b80f960 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/spec/test_modeling_v2_m31_dspark_n3a.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/spec/test_modeling_v2_m31_dspark_n3a.py new file mode 100644 index 000000000000..ad765523c7f6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/spec/test_modeling_v2_m31_dspark_n3a.py @@ -0,0 +1,466 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""N3 A2 in this target's DSpark worker: the fused strict acceptance and next-step tokens equal the in-tree path +(N3 P and A1, the verify logits, at the end). + +The worker's two overrides are driven on a bare worker (no model) against the in-tree methods of its base class on +the same inputs: +- eligible steps (K3 / K7, padded drafts, graph-padding dummy rows) are bitwise; +- forced acceptance (fractional, so the RNG pool path runs) over a run that alternates eligible and ineligible + steps keeps the same accepted counts and the same RNG counter as an all-in-tree run; +- the sampler runs exactly once per step; +- the switch and an ineligible step take the in-tree path. +""" + +import importlib +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode import strict_accept +from tensorrt_llm._torch.speculative.dspark import DSparkWorker + +TARGET = "minimax_m31_nvfp4__sm_103__tp4" +dspark = importlib.import_module( + f"tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.{TARGET}.dspark" +) + +VOCAB = 200064 + +assert torch.cuda.is_available(), "the N3 A2 worker test requires a CUDA device" + + +def _worker(force: float = 0.0): + w = dspark.M3DSparkWorker.__new__(dspark.M3DSparkWorker) + nn.Module.__init__(w) + w.force_num_accepted_tokens = force + w._force_accept_rng_counter = None + w._force_accept_rng_pool = None + w.calls = 0 + return w + + +def _metadata(ragged: bool = False, graph: bool = False): + return SimpleNamespace(is_ragged_verify=ragged, is_cuda_graph=graph) + + +def _fixed_sampler(tokens: torch.Tensor): + """The in-tree sampler stands in as a fixed function of the logits: the overrides must call it once.""" + + def sample(self, logits, spec_metadata, num_contexts, batch_size): + self.calls += 1 + return tokens + + return sample + + +def _step(n_req: int, draft_len: int, gen: torch.Generator, pad: bool = False, int64: bool = False): + k1 = draft_len + 1 + tokens = torch.randint(0, VOCAB, (n_req * k1,), generator=gen, device="cuda") + tokens = tokens if int64 else tokens.to(torch.int32) + draft = tokens.view(n_req, k1)[:, :draft_len].clone().to(torch.int32) + cut = torch.randint(0, draft_len + 1, (n_req,), generator=gen, device="cuda") + pos = torch.arange(draft_len, device="cuda").unsqueeze(0) + draft = torch.where(pos == cut.unsqueeze(1), (draft + 1) % VOCAB, draft) + if pad: + draft[::2, draft_len // 2 :] = -1 + logits = torch.zeros( + n_req * k1, 8, device="cuda" + ) # shape only: the sampler stand-in ignores the values + return logits, draft, tokens + + +def _in_tree(w, logits, draft, num_contexts, batch_size, md): + return DSparkWorker._sample_and_accept_draft_tokens_base( + w, logits, draft, num_contexts, batch_size, md + ) + + +@pytest.fixture(autouse=True) +def _compiled(): + strict_accept.precompile((3, 7), torch.device("cuda")) + + +@pytest.mark.parametrize( + "n_req,draft_len,pad,int64", + [ + (1, 3, False, False), + (8, 3, False, False), + (12, 3, True, False), + (8, 7, False, False), + (8, 3, False, True), + ], +) +def test_eligible_steps_are_bitwise(monkeypatch, n_req, draft_len, pad, int64): + gen = torch.Generator(device="cuda").manual_seed(3) + logits, draft, tokens = _step(n_req, draft_len, gen, pad=pad, int64=int64) + monkeypatch.setattr(dspark.M3DSparkWorker, "_sample_tokens_for_batch", _fixed_sampler(tokens)) + w, ref = _worker(), _worker() + md = _metadata() + assert w._n3a_eligible(logits, draft, 0, n_req, md) + acc, num = w._sample_and_accept_draft_tokens_base(logits, draft, 0, n_req, md) + racc, rnum = _in_tree(ref, logits, draft, 0, n_req, md) + assert w.calls == 1 and ref.calls == 1 + assert acc.dtype == racc.dtype and num.dtype == rnum.dtype + assert torch.equal(acc, racc) and torch.equal(num, rnum) + + +def test_forced_acceptance_alternating_steps_matches_in_tree(monkeypatch): + """Fractional forced AL (the RNG pool path) over eligible, ragged (ineligible) and graph-warmup steps: the + counts and the RNG counter follow the all-in-tree run step for step.""" + gen = torch.Generator(device="cuda").manual_seed(5) + w, ref = _worker(force=1.6), _worker(force=1.6) + served = 0 + for step in range(12): + logits, draft, tokens = _step(8, 3, gen) + num_contexts, batch_size = 0, 8 + if step % 3 == 2: + # An ineligible step: one context row ahead of the 8 generation requests (the in-tree path runs). + ctx_token = torch.randint(0, VOCAB, (1,), generator=gen, device="cuda").to(tokens.dtype) + tokens = torch.cat([ctx_token, tokens]) + logits = torch.zeros(tokens.numel(), 8, device="cuda") + num_contexts, batch_size = 1, 9 + monkeypatch.setattr( + dspark.M3DSparkWorker, "_sample_tokens_for_batch", _fixed_sampler(tokens) + ) + md = _metadata(graph=(step == 4)) + eligible = w._n3a_eligible(logits, draft, num_contexts, batch_size, md) + assert eligible == (num_contexts == 0), step + served += int(eligible) + acc, num = w._sample_and_accept_draft_tokens_base( + logits, draft, num_contexts, batch_size, md + ) + racc, rnum = _in_tree(ref, logits, draft, num_contexts, batch_size, md) + # A context row defines only column 0 (the in-tree path leaves the rest uninitialized). + assert torch.equal(acc[:num_contexts, 0], racc[:num_contexts, 0]), step + assert torch.equal(acc[num_contexts:], racc[num_contexts:]) and torch.equal(num, rnum), step + if ref._force_accept_rng_counter is not None: + assert torch.equal(w._force_accept_rng_counter, ref._force_accept_rng_counter), step + assert served == 8 + # A ragged verify step is declined (the in-tree ragged path keeps it). + logits, draft, _ = _step(8, 3, gen) + assert not w._n3a_eligible(logits, draft, 0, 8, _metadata(ragged=True)) + + +def test_ineligible_and_switch_take_the_in_tree_path(monkeypatch): + gen = torch.Generator(device="cuda").manual_seed(9) + logits, draft, tokens = _step(4, 3, gen) + w = _worker() + # A context row, a row-count mismatch, too many requests: declined. + assert not w._n3a_eligible(logits, draft, 1, 4, _metadata()) + assert not w._n3a_eligible(logits[:-1], draft, 0, 4, _metadata()) + big_logits, big_draft, _ = _step(strict_accept.max_requests() + 1, 3, gen) + assert not w._n3a_eligible(big_logits, big_draft, 0, big_draft.shape[0], _metadata()) + monkeypatch.setattr(dspark, "N3A_DISABLED", True) + assert not w._n3a_eligible(logits, draft, 0, 4, _metadata()) + + +@pytest.mark.parametrize( + "draft_len,wide,shuffle", [(3, False, False), (3, True, True), (7, False, True)] +) +def test_next_new_tokens_matches_in_tree(draft_len, wide, shuffle): + gen = torch.Generator(device="cuda").manual_seed(13) + n_req = 8 + width = draft_len + 1 + (4 if wide else 0) + accepted = torch.randint( + 0, VOCAB, (n_req, width), generator=gen, device="cuda", dtype=torch.int32 + ) + num = torch.randint(1, draft_len + 2, (n_req,), generator=gen, device="cuda", dtype=torch.int32) + batch_idx = torch.arange(32, device="cuda", dtype=torch.int32) + if shuffle: + batch_idx[:n_req] = torch.randperm(n_req, generator=gen, device="cuda").to(torch.int32) + next_draft = torch.randint( + 0, VOCAB, (n_req, draft_len), generator=gen, device="cuda", dtype=torch.int32 + ) + w = _worker() + out = w._prepare_next_new_tokens(accepted, next_draft, batch_idx, n_req, num) + ref = DSparkWorker._prepare_next_new_tokens(w, accepted, next_draft, batch_idx, n_req, num) + assert out.dtype == ref.dtype and torch.equal(out, ref) + + +# N3 P-gid + P-cat: the worker's target logits equal the logits processor's. + + +@pytest.mark.parametrize("tp", [2, 4, 8]) +def test_rank_major_strided_copy_equals_vocab_concat_and_cast(tp): + """The all-gather along dim 0 + one strided copy writes what the vocab concat + .float() writes.""" + gen = torch.Generator(device="cuda").manual_seed(tp) + rows, shard = 32, 25008 + shards = [ + torch.randn(rows, shard, generator=gen, device="cuda").to(torch.bfloat16) for _ in range(tp) + ] + gathered = torch.cat(shards, dim=0) # what the dim-0 all-gather returns, rank-major + out = torch.empty((rows, tp * shard), dtype=torch.float32, device="cuda") + out.view(rows, tp, shard).copy_(gathered.view(tp, rows, shard).transpose(0, 1)) + assert torch.equal(out, torch.cat(shards, dim=-1).float()) + + +def _lm_head(vocab=50016, hidden=6144): + from tensorrt_llm._torch.modules.embedding import LMHead + from tensorrt_llm._torch.modules.linear import TensorParallelMode + from tensorrt_llm.mapping import Mapping + + head = LMHead( + vocab, + hidden, + dtype=torch.bfloat16, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + ).cuda() + with torch.no_grad(): + head.weight.copy_( + torch.randn_like(head.weight, dtype=torch.float32).to(torch.bfloat16) * 0.02 + ) + return head + + +def test_target_logits_equal_the_processor_at_tp1(): + from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor + + head = _lm_head() + rows = 32 + hidden = torch.randn(rows, 6144, device="cuda").to(torch.bfloat16) + md = SimpleNamespace( + is_ragged_verify=False, gather_ids=torch.arange(rows, device="cuda", dtype=torch.int32) + ) + attn = SimpleNamespace(num_contexts=0) + out = _worker().target_logits(hidden, attn, md, head) + ref = LogitsProcessor().forward(hidden[md.gather_ids], head, attn, True) + assert out is not None and out.dtype == ref.dtype and torch.equal(out, ref) + + +def test_target_logits_decline_to_the_processor(monkeypatch): + head = _lm_head(vocab=1024, hidden=256) + hidden = torch.randn(8, 256, device="cuda").to(torch.bfloat16) + ids = torch.arange(8, device="cuda", dtype=torch.int32) + w = _worker() + # A context row, a gather that is not every row, a ragged verify, a padded head: the processor runs. + assert ( + w.target_logits( + hidden, + SimpleNamespace(num_contexts=1), + SimpleNamespace(is_ragged_verify=False, gather_ids=ids), + head, + ) + is None + ) + assert ( + w.target_logits( + hidden, + SimpleNamespace(num_contexts=0), + SimpleNamespace(is_ragged_verify=False, gather_ids=ids[:4]), + head, + ) + is None + ) + assert ( + w.target_logits( + hidden, + SimpleNamespace(num_contexts=0), + SimpleNamespace(is_ragged_verify=True, gather_ids=ids), + head, + ) + is None + ) + monkeypatch.setattr(head, "padding_size", 3) + assert ( + w.target_logits( + hidden, + SimpleNamespace(num_contexts=0), + SimpleNamespace(is_ragged_verify=False, gather_ids=ids), + head, + ) + is None + ) + monkeypatch.setattr(head, "padding_size", 0) + monkeypatch.setattr(dspark, "N3P_DISABLED", True) + assert ( + w.target_logits( + hidden, + SimpleNamespace(num_contexts=0), + SimpleNamespace(is_ragged_verify=False, gather_ids=ids), + head, + ) + is None + ) + + +# N3 A1: the verify logits' all-gather + concat + cast in one launch (catalog/comm/allgather_cast) when set up. + + +def _a1_step(monkeypatch, rows: int, compiled: bool = True): + """A TP4-shaped call: the head reports tp 4, A1 is set up, the op and the NCCL all-gather are recorded.""" + head = _lm_head(vocab=1024, hidden=256) + monkeypatch.setattr(head, "tp_size", 4) + calls = {"a1": 0, "nccl": 0} + sentinel = torch.full((rows, 4 * 1024), 7.0, device="cuda") + + def fake_a1(shard, **ws): + assert ws == {"ws": "a1"} and shard.shape == (rows, 1024) + calls["a1"] += 1 + return sentinel + + def fake_allgather(shard, mapping, dim=0): + calls["nccl"] += 1 + return torch.cat([shard] * 4, dim=0) + + monkeypatch.setattr(dspark._allgather_cast, "allgather_cast", fake_a1) + monkeypatch.setattr(dspark._allgather_cast, "is_compiled", lambda *a, **k: compiled) + monkeypatch.setattr(dspark, "allgather", fake_allgather) + w = _worker() + w._a1 = {"ws": "a1"} + hidden = torch.randn(rows, 256, device="cuda").to(torch.bfloat16) + md = SimpleNamespace( + is_ragged_verify=False, gather_ids=torch.arange(rows, device="cuda", dtype=torch.int32) + ) + out = w.target_logits(hidden, SimpleNamespace(num_contexts=0), md, head) + return out, sentinel, calls + + +@pytest.mark.parametrize("rows", [4, 32, 64]) +def test_target_logits_take_a1_when_set_up(monkeypatch, rows): + out, sentinel, calls = _a1_step(monkeypatch, rows) + assert out is sentinel and calls == {"a1": 1, "nccl": 0} + + +@pytest.mark.parametrize("rows,compiled", [(68, True), (4, False)]) +def test_target_logits_fall_back_to_nccl_past_a1(monkeypatch, rows, compiled): + """More rows than the workspace holds, or a shape never compiled: the NCCL all-gather + strided copy.""" + out, _, calls = _a1_step(monkeypatch, rows, compiled=compiled) + assert ( + calls == {"a1": 0, "nccl": 1} + and out.shape == (rows, 4 * 1024) + and out.dtype == torch.float32 + ) + + +def _a1_prebuild_worker(monkeypatch, tp: int = 4, draft_len: int = 3): + w = _worker() + w.mapping = SimpleNamespace(tp_size=tp, tp_rank=0, enable_attention_dp=False) + w.spec_config = SimpleNamespace(max_draft_len=draft_len) # max_draft_len reads it + # Eligibility needs every SM up to the workspace's grid; any CUDA device passes with a 1-CTA workspace grid. + monkeypatch.setattr(dspark._allgather_cast, "MAX_GRID", 1) + monkeypatch.setattr(dspark.M3DSparkWorker, "_a1_all_ok", lambda self, local: local) + return w + + +def test_prebuild_a1_sets_up_k3_rows(monkeypatch): + head = _lm_head(vocab=1024, hidden=256) + seen = {} + + def fake_precompile(mapping, cols, rows, device, group=None, agree=None): + seen.update(cols=cols, rows=rows, agreed=agree(True)) + return True + + ws = {"ws_uc": torch.zeros(4, dtype=torch.int32), "rank": 0} + monkeypatch.setattr(dspark._allgather_cast, "precompile", fake_precompile) + monkeypatch.setattr( + dspark._allgather_cast, "workspace_args", lambda mapping, cols, group=None: ws + ) + assert _a1_prebuild_worker(monkeypatch)._prebuild_a1(head) is ws + # K3 verify rows of batches 1..16, all within the workspace's 64 rows. + assert seen == {"cols": 1024, "rows": tuple(range(4, 65, 4)), "agreed": True} + + +@pytest.mark.parametrize("draft_len", [3, 7]) +def test_prebuild_a1_sets_up_the_verify_rows_of_k3_and_k7(monkeypatch, draft_len): + """K3: rows 4..64 (batches 1..16); K7, the headline: rows 8..64 (batches 1..8).""" + head = _lm_head(vocab=1024, hidden=256) + seen = {} + + def fake_precompile(mapping, cols, rows, device, group=None, agree=None): + seen.update(rows=rows, agreed=agree(True)) + return True + + ws = {"ws_uc": torch.zeros(4, dtype=torch.int32), "rank": 0} + monkeypatch.setattr(dspark._allgather_cast, "precompile", fake_precompile) + monkeypatch.setattr( + dspark._allgather_cast, "workspace_args", lambda mapping, cols, group=None: ws + ) + assert _a1_prebuild_worker(monkeypatch, draft_len=draft_len)._prebuild_a1(head) is ws + k1 = draft_len + 1 + assert seen == {"rows": tuple(range(k1, 65, k1)), "agreed": True} + + +def test_prebuild_a1_off(monkeypatch): + """The switch, another draft length, one rank, a rank that is not eligible, or a compile that failed on any rank: + no A1, no workspace.""" + head = _lm_head(vocab=1024, hidden=256) + built = {"n": 0} + + def fake_precompile(*a, **k): + built["n"] += 1 + return False # a failed compile somewhere + + monkeypatch.setattr(dspark._allgather_cast, "precompile", fake_precompile) + monkeypatch.setattr(dspark, "A1_K7_DISABLED", True) + assert _a1_prebuild_worker(monkeypatch, draft_len=7)._prebuild_a1(head) is None + monkeypatch.setattr(dspark, "A1_K7_DISABLED", False) + assert ( + _a1_prebuild_worker(monkeypatch, draft_len=64)._prebuild_a1(head) is None + ) # 65 rows at batch 1 + assert _a1_prebuild_worker(monkeypatch, tp=1)._prebuild_a1(head) is None + assert built["n"] == 0 + assert _a1_prebuild_worker(monkeypatch)._prebuild_a1(head) is None + assert built["n"] == 1 + # Another rank not eligible: off here too, nothing compiled. + w = _a1_prebuild_worker(monkeypatch) + monkeypatch.setattr(dspark.M3DSparkWorker, "_a1_all_ok", lambda self, local: False) + assert w._prebuild_a1(head) is None and built["n"] == 1 + monkeypatch.setattr(dspark, "A1_DISABLED", True) + assert _a1_prebuild_worker(monkeypatch)._prebuild_a1(head) is None + assert built["n"] == 1 + + +# N3A-KV: the draft-forward KV bookkeeping in one launch, through the DFlash hook (bitwise to the stock lines). + + +def _kv_step(nc: int, bs: int, gen: torch.Generator): + num = torch.randint(1, 5, (bs + 2,), generator=gen, device="cuda", dtype=torch.int32) + kv = torch.randint(0, 1 << 20, (bs + 2,), generator=gen, device="cuda", dtype=torch.int32) + return num, kv + + +def _kv_worker(fused: bool): + w = _worker() + w._fused_kv_rewind = None + if fused: + strict_accept.precompile_kv_rewind(torch.device("cuda")) + w._n3a_kv_buf = torch.empty(strict_accept.max_requests(), dtype=torch.int32, device="cuda") + w._fused_kv_rewind = w._n3a_kv_rewind + return w + + +@pytest.mark.parametrize("nc,bs", [(0, 1), (0, 8), (2, 6), (3, 3)]) +def test_kv_rewind_hook_equals_the_stock_lines(nc, bs): + gen = torch.Generator(device="cuda").manual_seed(5) + num, kv = _kv_step(nc, bs, gen) + outs = [] + for fused in (False, True): + w = _kv_worker(fused) + attn = SimpleNamespace(kv_lens_cuda=kv.clone(), update_for_spec_dec=lambda: None) + w._prepare_kv_for_draft_forward(attn, num, nc, bs) + outs.append( + ( + w._kv_rewind_amount, + attn.kv_lens_cuda, + w._kv_rewind_nc, + w._kv_rewind_bs, + w._kv_rewind_pending, + ) + ) + (r0, k0, *rest0), (r1, k1, *rest1) = outs + assert torch.equal(k0, k1) and rest0 == rest1 + assert r0.dtype == r1.dtype and r0.shape == r1.shape and torch.equal(r0, r1) + + +def test_kv_rewind_hook_declines_unsupported_operands(): + w = _kv_worker(True) + num64 = torch.ones(4, dtype=torch.int64, device="cuda") # not the buffer's dtype + kv = torch.zeros(4, dtype=torch.int32, device="cuda") + assert w._n3a_kv_rewind(num64, kv, 0, 4) is None and torch.equal(kv, torch.zeros_like(kv)) + w._n3a_kv_buf = None + assert w._n3a_kv_rewind(torch.ones(4, dtype=torch.int32, device="cuda"), kv, 0, 4) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/spec/test_modeling_v2_m31_dspark_s4.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/spec/test_modeling_v2_m31_dspark_s4.py new file mode 100644 index 000000000000..fe633ff9a2c1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/spec/test_modeling_v2_m31_dspark_s4.py @@ -0,0 +1,347 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp4 target's worker-side routing of the one-launch Markov chain (``dspark_s4.S4WorkerMixin``), without GPUs +or a model: which steps it serves, that the kill switch sets nothing up, that a declined step is the in-tree path, that +the step's logits processor only takes the block rows, and that the chain's tokens replace the in-tree sampling for +that one call. The chain itself is the catalog's ``spec_decode/dspark_markov_chain``, certified by +``spec_decode/_dspark_markov_chain_op_matrix.py``; the in-model equality by the call-site gate.""" + +import importlib +import os +import subprocess +import sys +from types import SimpleNamespace + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4 import ( # noqa: E501 + dspark_s4 as w1, +) + + +class _Base: + def __init__(self): + self.calls = [] + + def _forward_impl(self, *args): + self.calls.append("forward") + return "in-tree" + + def _refine_block_logits(self, draft_model, gen_logits, inputs, spec_metadata): + self.calls.append("refine") + return gen_logits + + def sample_draft_tokens(self, logits, spec_metadata, batch_size, **kwargs): + self.calls.append("sample") + return torch.full(logits.shape[:2], -1, dtype=torch.int32) + + def _apply_kv_rewind_after_draft(self, attn_metadata, spec_metadata): + self.calls.append("stock_rewind") + self._kv_rewind_pending = False + + +class _Worker(w1.S4WorkerMixin, _Base): + pass + + +def _state(block=3, shard=8, rows=8): + st = w1._S4State(None, rows, shard, 0, None, None, block, {}) + for b in range(1, w1.MAX_FAST_BATCH + 1): + st.bufs[(b, block)] = torch.full((b, block, rows), float("-inf")) + return st + + +def _meta(advanced=False): + return SimpleNamespace(wants_advanced_draft_sampling=advanced) + + +def test_serves_only_eligible_steps(monkeypatch): + wk = _Worker() + draft = SimpleNamespace(lilicorr=None, has_candidate_selector=False) + assert not wk._s4_serves(1, _meta(), draft) # not set up + wk._s4 = _state() + assert wk._s4_serves(1, _meta(), draft) + assert wk._s4_serves(w1.MAX_FAST_BATCH, _meta(), draft) + assert not wk._s4_serves(0, _meta(), draft) + assert not wk._s4_serves(w1.MAX_FAST_BATCH + 1, _meta(), draft) + assert not wk._s4_serves(1, _meta(advanced=True), draft) + wk._d2t = torch.zeros(1) + assert not wk._s4_serves(1, _meta(), draft) + + +def test_kill_switch_sets_nothing_up(monkeypatch): + wk = _Worker() + wk._s4 = _state() + monkeypatch.setattr(w1, "_DISABLED", True) + assert not wk.prebuild_s4(None, SimpleNamespace(has_markov_head=True)) + assert wk._s4 is None + assert not wk._s4_serves( + 1, _meta(), SimpleNamespace(lilicorr=None, has_candidate_selector=False) + ) + + +def test_declined_step_is_the_in_tree_path(): + wk = _Worker() + attn = SimpleNamespace(num_seqs=12, num_contexts=0) # above MAX_FAST_BATCH + draft = SimpleNamespace(lilicorr=None, has_candidate_selector=False, logits_processor="proc") + wk._s4 = _state() + assert wk._forward_impl(None, None, None, None, attn, _meta(), draft) == "in-tree" + assert draft.logits_processor == "proc" + logits = torch.zeros(2, 3, 16) # full-vocab width: not this step's shard + assert wk._refine_block_logits(draft, logits, {}, None) is logits + assert wk.sample_draft_tokens(logits, None, 2).eq(-1).all() + assert wk.calls == ["forward", "refine", "sample"] + + +def test_sharded_logits_processor_takes_only_block_rows(): + st = _state(block=3, shard=8, rows=12) + inner_calls = [] + + def inner(h, head, meta, ctx): + inner_calls.append(h.shape) + return "inner" + + head = SimpleNamespace(apply_linear=lambda h, b: h[:, :8] * 2) + proc = w1._ShardedDraftLogits(inner, st, batch=2) + h = torch.arange(6 * 10, dtype=torch.float32).view(6, 10) + out = proc(h, head, None, True) + assert torch.equal(out, h[:, :8] * 2) + buf = st.buffer(2) + assert torch.equal(buf[:, :, :8], (h[:, :8] * 2).view(2, 3, 8)) + assert torch.isinf(buf[:, :, 8:]).all() + assert proc(torch.zeros(5, 10), head, None, True) == "inner" and inner_calls == [ + torch.Size([5, 10]) + ] + + +def test_chain_tokens_replace_sampling_once(monkeypatch): + wk = _Worker() + st = _state(block=3, shard=8, rows=8) + wk._s4_step = st + toks = torch.tensor([[1, 2, 3], [4, 5, 6]]) + monkeypatch.setattr(w1, "chain", SimpleNamespace(supports=lambda *a: True, + dspark_markov_chain=lambda *a, **k: toks.int(), + ll_exchange_engaged=lambda batch: False)) # fmt: skip + logits = torch.zeros(2, 3, 8) + assert wk._refine_block_logits(None, logits, {"first_prev_tokens": torch.zeros(2, dtype=torch.long)}, None) \ + is logits # fmt: skip + out = wk.sample_draft_tokens(logits, None, 2) + assert out.dtype == torch.int32 and torch.equal(out, toks.int()) + assert wk.sample_draft_tokens(logits, None, 2).eq(-1).all() # used once + assert "refine" not in wk.calls + + +def test_unsupported_shapes_fall_back_on_the_shards(monkeypatch): + wk = _Worker() + wk._s4_step = _state(block=3, shard=8, rows=8) + + def _never(*a, **k): + raise AssertionError("an unsupported step must not reach the op") + + monkeypatch.setattr( + w1, "chain", SimpleNamespace(supports=lambda *a: False, dspark_markov_chain=_never) + ) + logits = torch.zeros(2, 3, 8) + wk._refine_block_logits( + None, logits, {"first_prev_tokens": torch.zeros(2, dtype=torch.long)}, None + ) + assert wk.calls == ["refine"] + assert wk.sample_draft_tokens(logits, None, 2).eq(-1).all() + + +def test_switch_reads_any_value_but_empty_or_zero(monkeypatch): + """Read at import, like the targets' m31_switches: "true" turns S4 off, not only "1".""" + try: + for value, off in ( + ("1", True), + ("true", True), + ("yes", True), + ("0", False), + ("", False), + (" 0 ", False), + ): + monkeypatch.setenv(w1.DISABLE_ENV, value) + importlib.reload(w1) + assert w1.disabled() is off, value + finally: + monkeypatch.delenv(w1.DISABLE_ENV, raising=False) + importlib.reload(w1) + assert not w1.disabled() + + +def test_markov_ll_engagement_logged_once_per_rank(monkeypatch): + """MARKOV_LL's "[M3 dispatch]" line: once per rank, on the first served step that took the LL kernel.""" + wk = _Worker() + wk.mapping = SimpleNamespace(tp_rank=3) + wk._s4_step = _state(block=3, shard=8, rows=8) + lines = [] + monkeypatch.setattr(w1.logger, "info", lambda msg, *a, **k: lines.append(msg)) + monkeypatch.setattr(w1, "chain", SimpleNamespace(supports=lambda *a: True, + dspark_markov_chain=lambda *a, **k: torch.zeros( + a[0].shape[0], 3, dtype=torch.int32), + ll_exchange_engaged=lambda batch: batch == 1)) # fmt: skip + inputs2 = {"first_prev_tokens": torch.zeros(2, dtype=torch.long)} + wk._refine_block_logits(None, torch.zeros(2, 3, 8), inputs2, None) # batch 2: not the LL kernel + assert not [m for m in lines if "MARKOV_LL" in m] + inputs1 = {"first_prev_tokens": torch.zeros(1, dtype=torch.long)} + for _ in range(3): + wk._refine_block_logits(None, torch.zeros(1, 3, 8), inputs1, None) + ll = [m for m in lines if "MARKOV_LL" in m] + assert ll == [ + "[M3 dispatch] MARKOV_LL: k3_markov LL flag-in-payload exchange (batch 1) on rank 3" + ], ll + + +def _markov_ll_in_subprocess(value): + env = {k: v for k, v in os.environ.items() if k != "TRTLLM_M31_DISABLE_MARKOV_LL"} + if value is not None: + env["TRTLLM_M31_DISABLE_MARKOV_LL"] = value + code = ( + "from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3." + "minimax_m31_nvfp4__sm_103__tp4 import dspark_s4 as m; print('MARKOV_LL=' + str(m.MARKOV_LL))" + ) + out = subprocess.run( + [sys.executable, "-c", code], env=env, capture_output=True, text=True, timeout=600 + ) + assert out.returncode == 0, out.stderr + return [ln for ln in out.stdout.splitlines() if ln.startswith("MARKOV_LL=")][-1] + + +def test_markov_ll_kill_switch_subprocess(): + """MARKOV_LL is on by default; TRTLLM_M31_DISABLE_MARKOV_LL=1 (read at import) turns it off.""" + assert _markov_ll_in_subprocess(None) == "MARKOV_LL=True" + assert _markov_ll_in_subprocess("0") == "MARKOV_LL=True" + assert _markov_ll_in_subprocess("1") == "MARKOV_LL=False" + + +def _rewind_worker(monkeypatch, batch=2, nc=1, on=True): + """A worker in a served step with a pending KV-length rewind of ``batch`` rows from ``nc``, and a chain fake that + records which launch ran.""" + wk = _Worker() + wk.mapping = SimpleNamespace(tp_rank=5) + wk._s4_step = _state(block=3, shard=8, rows=8) + wk._s4_step_attn = SimpleNamespace(kv_lens_cuda=torch.zeros(nc + batch + 1, dtype=torch.int32)) + wk._kv_rewind_pending = True + wk._kv_rewind_amount = torch.zeros(batch, dtype=torch.int32) + wk._kv_rewind_nc, wk._kv_rewind_bs = nc, nc + batch + launches, lines = [], [] + toks = torch.zeros(batch, 3, dtype=torch.int32) + + def plain(*a, **k): + launches.append("chain") + return toks + + def with_rewind(base, first_prev, w1_, w2s, offset, kv_lens, amount, first, **k): + launches.append( + ( + "rewind", + kv_lens is wk._s4_step_attn.kv_lens_cuda, + amount is wk._kv_rewind_amount, + first, + ) + ) + return toks + + monkeypatch.setattr(w1, "TAILBK_SPEC", on) + monkeypatch.setattr(w1.logger, "info", lambda msg, *a, **k: lines.append(msg)) + monkeypatch.setattr(w1, "chain", SimpleNamespace(supports=lambda *a: True, supports_rewind=lambda *a: True, + dspark_markov_chain=plain, + dspark_markov_chain_rewind=with_rewind, + ll_exchange_engaged=lambda b: False)) # fmt: skip + return wk, launches, lines, {"first_prev_tokens": torch.zeros(batch, dtype=torch.long)} + + +def test_tailbk_h_rewind_in_chain_launch(monkeypatch): + """TAILBK_SPEC H: a real step's chain launch applies the pending rewind (this step's kv_lens, amounts, first + row) and the stock rewind lines are skipped for that step; the engagement line is logged once per rank.""" + wk, launches, lines, inputs = _rewind_worker(monkeypatch) + meta = SimpleNamespace(is_cuda_graph=False) + for _ in range(2): + wk._kv_rewind_pending = True + wk._refine_block_logits(None, torch.zeros(2, 3, 8), inputs, meta) + wk._apply_kv_rewind_after_draft(None, meta) + assert not wk._kv_rewind_pending and not wk._tailbk_rewound + assert launches == [("rewind", True, True, 1)] * 2, launches + assert "stock_rewind" not in wk.calls + assert [m for m in lines if "TAILBK_SPEC H" in m] == [ + "[M3 dispatch] TAILBK_SPEC H: KV-length rewind in the Markov chain launch on rank 5" + ] + + +def test_tailbk_h_warmup_keeps_stock_path(monkeypatch): + """The eager warmup of a graph batch (is_cuda_graph, not capturing): the stock method skips the rewind there, so + the chain must not apply it either.""" + wk, launches, _, inputs = _rewind_worker(monkeypatch) + meta = SimpleNamespace(is_cuda_graph=True) + wk._refine_block_logits(None, torch.zeros(2, 3, 8), inputs, meta) + wk._apply_kv_rewind_after_draft(None, meta) + assert launches == ["chain"] and wk.calls == ["stock_rewind"] + + +def test_tailbk_h_off_or_mismatch_keeps_stock_path(monkeypatch): + """Kill switch off, no pending rewind, or amounts that are not this chain call's rows: the stock rewind runs.""" + meta = SimpleNamespace(is_cuda_graph=False) + wk, launches, _, inputs = _rewind_worker(monkeypatch, on=False) + wk._refine_block_logits(None, torch.zeros(2, 3, 8), inputs, meta) + wk._apply_kv_rewind_after_draft(None, meta) + assert launches == ["chain"] and wk.calls == ["stock_rewind"] + wk, launches, _, inputs = _rewind_worker(monkeypatch) + wk._kv_rewind_pending = False + wk._refine_block_logits(None, torch.zeros(2, 3, 8), inputs, meta) + assert launches == ["chain"] + wk, launches, _, inputs = _rewind_worker(monkeypatch) + wk._kv_rewind_amount = torch.zeros(3, dtype=torch.int32) # not the 2 rows of this call + wk._refine_block_logits(None, torch.zeros(2, 3, 8), inputs, meta) + wk._apply_kv_rewind_after_draft(None, meta) + assert launches == ["chain"] and wk.calls == ["stock_rewind"] + + +def test_tailbk_h_unsupported_operands_keep_stock_path(monkeypatch): + """supports_rewind False (e.g. int64 amounts): the plain chain launch and the stock rewind.""" + wk, launches, _, inputs = _rewind_worker(monkeypatch) + monkeypatch.setattr(w1.chain, "supports_rewind", lambda *a: False) + meta = SimpleNamespace(is_cuda_graph=False) + wk._refine_block_logits(None, torch.zeros(2, 3, 8), inputs, meta) + wk._apply_kv_rewind_after_draft(None, meta) + assert launches == ["chain"] and wk.calls == ["stock_rewind"] + + +def _tailbk_in_subprocess(value): + env = {k: v for k, v in os.environ.items() if k != "TRTLLM_M31_DISABLE_TAILBK_SPEC"} + if value is not None: + env["TRTLLM_M31_DISABLE_TAILBK_SPEC"] = value + code = ( + "from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3." + "minimax_m31_nvfp4__sm_103__tp4 import dspark_s4 as m; print('TAILBK_SPEC=' + str(m.TAILBK_SPEC))" + ) + out = subprocess.run( + [sys.executable, "-c", code], env=env, capture_output=True, text=True, timeout=600 + ) + assert out.returncode == 0, out.stderr + return [ln for ln in out.stdout.splitlines() if ln.startswith("TAILBK_SPEC=")][-1] + + +def test_tailbk_spec_kill_switch_subprocess(): + """TAILBK_SPEC is on by default; TRTLLM_M31_DISABLE_TAILBK_SPEC=1 (read at import) turns it off.""" + assert _tailbk_in_subprocess(None) == "TAILBK_SPEC=True" + assert _tailbk_in_subprocess("0") == "TAILBK_SPEC=True" + assert _tailbk_in_subprocess("1") == "TAILBK_SPEC=False" + + +def test_tailbk_c_engagement_logged_once_per_rank(monkeypatch): + """TAILBK_SPEC C's "[M3 dispatch]" line: once per rank, only after a chain call took the cached inputs.""" + wk, launches, lines, inputs = _rewind_worker(monkeypatch) + wk._kv_rewind_pending = False # the plain chain launch + taken = [False] + monkeypatch.setattr(w1.chain, "cached_inputs_taken", lambda: taken[0], raising=False) + meta = SimpleNamespace(is_cuda_graph=False) + wk._refine_block_logits(None, torch.zeros(2, 3, 8), inputs, meta) + assert not [m for m in lines if "TAILBK_SPEC C" in m] + taken[0] = True + for _ in range(3): + wk._refine_block_logits(None, torch.zeros(2, 3, 8), inputs, meta) + wk._kv_rewind_pending = True # and through the rewind launch + wk._refine_block_logits(None, torch.zeros(2, 3, 8), inputs, meta) + assert [m for m in lines if "TAILBK_SPEC C" in m] == [ + "[M3 dispatch] TAILBK_SPEC C: chain accepted inputs from precompile cache (taken) on rank 5" + ] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/spec/test_modeling_v2_m31_dspark_w0.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/spec/test_modeling_v2_m31_dspark_w0.py new file mode 100644 index 000000000000..e1c9f72e7fcb --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/spec/test_modeling_v2_m31_dspark_w0.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-model DSpark on the tp4 target, phase W0 (``dspark.py``): the shell's calls on the DSpark worker take no +fused path, the capture is the in-tree DFlash capture, the drafter's borrowed embedding is the target's exact gather, +and the FP8 draft-KV override (an Eagle3 need) is never applied to a DSpark run, so the drafter's context pool keeps +the projection dtype (BF16) and binds to the separate draft KV cache. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4 import ( # noqa: E501 + dspark, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.ops import ( # noqa: E501 + m3_target_head, +) +from tensorrt_llm.quantization.mode import QuantAlgo + + +def _spec( + *, dspark_mode: bool, eagle3: bool = False, one_engine: bool = True, embedded: bool = False +): + mode = SimpleNamespace( + use_one_engine=lambda: one_engine, + is_dspark=lambda: dspark_mode, + is_eagle3_one_model=lambda: eagle3, + ) + return SimpleNamespace(spec_dec_mode=mode, draft_is_embedded_in_target=embedded) + + +def _worker() -> dspark.M3DSparkWorker: + # The shell-facing calls read no worker state, so the worker is not constructed (that needs a mapping and a + # DSpark spec config). + return dspark.M3DSparkWorker.__new__(dspark.M3DSparkWorker) + + +@pytest.mark.parametrize( + "spec, expected", + [ + (_spec(dspark_mode=True), True), + (_spec(dspark_mode=True, embedded=True), False), + (_spec(dspark_mode=True, one_engine=False), False), + (_spec(dspark_mode=False, eagle3=True), False), + (None, False), + ], +) +def test_standalone_dspark_detection(spec, expected): + assert dspark.is_standalone_dspark(spec) is expected + + +def test_fp8_draft_kv_override_is_eagle3_only(): + assert dspark.draft_kv_override_applies(_spec(dspark_mode=False, eagle3=True)) + assert dspark.draft_kv_override_applies(None) + assert not dspark.draft_kv_override_applies(_spec(dspark_mode=True)) + + +@pytest.mark.parametrize("target_algo", [QuantAlgo.NVFP4, QuantAlgo.FP8, None]) +def test_drafter_kv_pool_keeps_the_drafter_dtype(target_algo): + from tensorrt_llm.models.modeling_utils import QuantConfig + + whole, layer = ( + QuantConfig(kv_cache_quant_algo=target_algo), + QuantConfig(kv_cache_quant_algo=target_algo), + ) + _ = whole.quant_mode # cached before the drop, as a loaded config may have it + cfg = SimpleNamespace(quant_config=whole, quant_config_dict={"model.layers.0": layer}) + dspark.keep_drafter_kv_dtype(cfg) + assert whole.kv_cache_quant_algo is None and layer.kv_cache_quant_algo is None + assert not whole.quant_mode.has_fp8_kv_cache() and not whole.quant_mode.has_fp4_kv_cache() + + +def test_shell_calls_take_no_fused_path(monkeypatch): + def _forbidden(*args, **kwargs): + raise AssertionError("W0 must never reach the coded target head") + + monkeypatch.setattr(m3_target_head, "prepare_head", _forbidden) + monkeypatch.setattr(m3_target_head, "logits", _forbidden) + worker = _worker() + assert worker.tail_eligible(None, None, None) is False + assert worker.arm_capture_norm(None, None, None, None) is None + assert worker.prebuild(None, None) is None + + +def test_capture_is_the_in_tree_capture(): + calls = [] + meta = SimpleNamespace(maybe_capture_hidden_states=lambda *a: calls.append(a)) + hidden, residual = object(), object() + _worker().capture_hidden_states(meta, 17, hidden, residual) + _worker().capture_hidden_states(meta, 59, hidden) + assert calls == [(17, hidden, residual), (59, hidden, None)] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_embed_tokens_is_the_target_gather(): + gen = torch.Generator(device="cuda").manual_seed(31) + table = torch.randn(1000, 64, generator=gen, device="cuda").to(torch.bfloat16) + core = SimpleNamespace(w={"embed": table}) + embed = dspark.M3EmbedTokens(core) + ids = torch.tensor([0, 5, 999, 200, -1], dtype=torch.long, device="cuda") + out = embed(ids) + ref = table[ids.clamp(min=0)] * (ids >= 0).unsqueeze(-1).to(table.dtype) + assert embed.weight is table + assert torch.equal(out, ref) + assert not list(embed.parameters()), ( + "the table belongs to the target, not to the borrowed module" + ) + + +def test_borrowed_embedding_owns_no_parameters(): + # The weights walker asserts every shell parameter is fed: the borrowed module must add none. + model = torch.nn.Module() + model.w = torch.nn.ParameterDict( + {"embed": torch.nn.Parameter(torch.zeros(8, 4), requires_grad=False)} + ) + before = {n for n, _ in model.named_parameters()} + model.embed_tokens = dspark.M3EmbedTokens(model) + assert {n for n, _ in model.named_parameters()} == before == {"w.embed"} + assert model.embed_tokens.weight is model.w["embed"] + + +class _FakeDrafter: + def __init__(self, borrows: bool = True) -> None: + self.loaded = None + self.borrows = borrows + self.lm_head = None + self.draft_model_full = SimpleNamespace(model=SimpleNamespace(embed_tokens=None)) + + def load_weights(self, weights, weight_mapper=None): + self.loaded = dict(weights) + + def load_weights_from_target_model(self, target): + if self.borrows: + self.lm_head = target.lm_head + self.draft_model_full.model.embed_tokens = target.model.embed_tokens + + +def _shell(drafter): + return SimpleNamespace( + draft_model=drafter, + lm_head=object(), + model=SimpleNamespace(embed_tokens=object()), + named_parameters=lambda: iter([("lm_head.weight", None), ("model.w.embed", None)]), + ) + + +@pytest.mark.parametrize( + "name", + ["model.embed_tokens.weight", "language_model.model.embed_tokens.weight", "lm_head.weight"], +) +def test_draft_load_refuses_borrowed_tensors(name): + shell = _shell(_FakeDrafter()) + with pytest.raises(ValueError, match="lends it"): + dspark.load_dspark_draft_weights(shell, {name: torch.zeros(1), "fc.weight": torch.zeros(1)}) + assert shell.draft_model.loaded is None + + +def test_draft_load_borrows_target_modules(): + shell = _shell(_FakeDrafter()) + weights = { + "fc.weight": torch.zeros(1), + "markov_head.markov_w1.weight": torch.zeros(1), + "confidence_head.proj.weight": torch.zeros(1), + } + dspark.load_dspark_draft_weights(shell, weights) + assert shell.draft_model.loaded == weights + assert shell.draft_model.lm_head is shell.lm_head + with pytest.raises(RuntimeError, match="does not borrow"): + dspark.load_dspark_draft_weights(_shell(_FakeDrafter(borrows=False)), weights) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py new file mode 100644 index 000000000000..7279bbc33f9e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_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_m31_nvfp4__sm_103__tp4 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.ops import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_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_m31_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_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_ar_det.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_ar_det.py new file mode 100644 index 000000000000..bac700e045d8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_ar_det.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""AR-DET: the target defaults the AllReduce autotune off for its process (the C++ heuristic), keeps an explicit +TLLM_DISABLE_ALLREDUCE_AUTOTUNE as given, and leaves the autotune on with TRTLLM_M31_ENABLE_AR_AUTOTUNE=1.""" + +import os + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4 import ( # noqa: E501 + modeling, +) + +KEY = "TLLM_DISABLE_ALLREDUCE_AUTOTUNE" + + +def test_autotune_defaults_off(monkeypatch) -> None: + monkeypatch.setattr(os, "environ", {}) + monkeypatch.setattr(modeling, "AR_AUTOTUNE_ENABLED", False) + assert modeling._pin_allreduce_heuristic(0) is True + assert os.environ[KEY] == "1" + + +def test_escape_hatch_keeps_autotune_on(monkeypatch) -> None: + monkeypatch.setattr(os, "environ", {}) + monkeypatch.setattr(modeling, "AR_AUTOTUNE_ENABLED", True) + assert modeling._pin_allreduce_heuristic(0) is False + assert KEY not in os.environ + + +def test_explicit_setting_is_kept(monkeypatch) -> None: + monkeypatch.setattr(os, "environ", {KEY: "0"}) + monkeypatch.setattr(modeling, "AR_AUTOTUNE_ENABLED", False) + assert modeling._pin_allreduce_heuristic(0) is False + assert os.environ[KEY] == "0" + + +def test_escape_hatch_follows_the_switch_contract() -> None: + """TRTLLM_M31_ENABLE_AR_AUTOTUNE is parsed by m31_switches (1/true/on/yes set; ""/0/false/off/no unset), not by + a separate loose read.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4 import ( # noqa: E501 + m31_switches, + ) + + assert modeling.AR_AUTOTUNE_ENABLED == m31_switches.m31_enabled("AR_AUTOTUNE") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_fmha_sm100_fixes.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_fmha_sm100_fixes.py new file mode 100644 index 000000000000..b551103f0079 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_fmha_sm100_fixes.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The two arch-independent fmha_sm100 correctness fixes the M3.1 targets carry in +3rdparty/patches/msa_strided_paged_kv.patch (from #20629 cf8a77e816), checked on the installed package: + +1. Sm100FmhaFwdKernelTmaWarpspecialized is PDL-launched, so it must run griddepcontrol.wait before reading any + input (the upstream wait is commented out: a race with its predecessor; the M3.1 prefill index proxy pass runs it). +2. The NVFP4-KV forward zeroes the V block scales / values of rows at or past kv_valid_cols, so a never-written pool + slot cannot inject NaN through a zero probability. +""" + +import importlib.util +import re +from pathlib import Path + +import pytest + +_SPEC = importlib.util.find_spec("fmha_sm100") +pytestmark = pytest.mark.skipif(_SPEC is None, reason="fmha_sm100 is not installed") + + +def _pkg() -> Path: + return Path(_SPEC.origin).parent + + +def test_the_fwd_kernel_waits_on_its_grid_dependency_before_reading_inputs() -> None: + src = (_pkg() / "csrc/include/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp").read_text() + body = src[src.index("CUTLASS_DEVICE void operator()(const Params& params, char* smem) {") :] + first_lines = body.split("\n")[1:6] + active = [ + line + for line in first_lines + if "griddepcontrol.wait" in line and not line.lstrip().startswith("//") + ] + assert active, "griddepcontrol.wait is missing or commented out at the top of the fwd kernel" + + +def test_the_nvfp4_kv_forward_zeroes_v_rows_past_the_valid_columns() -> None: + src = (_pkg() / "cute/src/sm100/fwd/atten_fwd_nvfp4_kv.py").read_text() + assert re.search(r"scale_e4m3_lo = scale_e4m3_lo if row < kv_valid_cols else Int32\(0\)", src) + assert re.search(r"scale_e4m3_hi = scale_e4m3_hi if row < kv_valid_cols else Int32\(0\)", src) + assert re.search( + r"combined_bf16x2 = combined_bf16x2 if row < kv_valid_cols else Int32\(0\)", src + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_hosttail.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_hosttail.py new file mode 100644 index 000000000000..221e0d9a7e43 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_hosttail.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""HOSTTAIL's kill switches on this target: ``_runtime_hooks`` turns on all six host-tail RuntimeHooks fields by +default, ``TRTLLM_M31_DISABLE_`` turns off only its own field and ``TRTLLM_M31_DISABLE_HOSTTAIL`` all six. + +The ``HOSTTAIL_*`` expressions are evaluated from modeling.py's source against a private copy of m31_switches loaded +under each environment (as test_modeling_v2_m31_switches does), so neither the in-tree switches nor the heavy +modeling module is reloaded. CPU only, except the last test, which imports the in-tree module. +""" + +import ast +import importlib +import importlib.util +import os +from pathlib import Path + +import pytest +import torch + +import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 as _family +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import ( + MiniMaxM3KVCacheManagerV2, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + BaseMambaCacheManager, + MambaHybridCacheManager, +) +from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks + +TARGET = "minimax_m31_nvfp4__sm_103__tp4" +_DIR = Path(_family.__file__).resolve().parent / TARGET +_SWITCHES = _DIR / "m31_switches.py" +_MODELING = _DIR / "modeling.py" + +# Kill-switch id -> (the module constant, the RuntimeHooks field it sets). +ITEMS = { + "HOST_D2H_SIDE": ("HOSTTAIL_D2H_SIDE", "sampler_d2h_side_stream"), + "SMALL_H2D": ("HOSTTAIL_SMALL_H2D", "spec_small_h2d"), + "BATCH_CHANGED_ONLY": ("HOSTTAIL_BATCH_CHANGED_ONLY", "spec_batch_map_changed_only"), + "SKIP_ACCEPTED_UPLOAD": ("HOSTTAIL_SKIP_ACCEPTED_UPLOAD", "skip_accepted_draft_upload"), + "PHILOX_PACK": ("HOSTTAIL_PHILOX_PACK", "pack_spec_rng_uploads"), + "KV_LENS_ONCE": ("HOSTTAIL_KV_LENS_ONCE", "kv_lens_update_once"), +} +_SWITCH_IDS = ("HOSTTAIL", *ITEMS) + + +def _module_tree() -> ast.Module: + return ast.parse(_MODELING.read_text(), filename=str(_MODELING)) + + +def _hosttail_statements(tree: ast.Module) -> list[ast.stmt]: + """modeling.py's module-level ``_HOSTTAIL_OFF = ...`` and ``HOSTTAIL_* = ...`` assignments, in source order.""" + statements = [] + for node in tree.body: + if isinstance(node, ast.Assign) and len(node.targets) == 1: + target = node.targets[0] + if isinstance(target, ast.Name) and ( + target.id == "_HOSTTAIL_OFF" or target.id.startswith("HOSTTAIL_") + ): + statements.append(node) + return statements + + +def _switches(monkeypatch, env): + """A private copy of m31_switches imported under ``env`` (the HOSTTAIL variables cleared first).""" + for item in _SWITCH_IDS: + for prefix in ("TRTLLM_M31_DISABLE_", "TRTLLM_M31_ENABLE_"): + monkeypatch.delenv(prefix + item, raising=False) + for name, value in env.items(): + monkeypatch.setenv(name, value) + spec = importlib.util.spec_from_file_location(f"_m31_switches_hosttail_{id(env)}", _SWITCHES) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _hosttail(monkeypatch, env) -> dict[str, bool]: + """The six HOSTTAIL_* values modeling.py computes under ``env``, by item id.""" + switches = _switches(monkeypatch, env) + namespace = {"m31_disabled": switches.m31_disabled, "m31_enabled": switches.m31_enabled} + exec( + compile(ast.Module(_hosttail_statements(_module_tree()), []), str(_MODELING), "exec"), + namespace, + ) + return {item: namespace[constant] for item, (constant, _) in ITEMS.items()} + + +def test_the_switches_are_the_six_items(): + names = [node.targets[0].id for node in _hosttail_statements(_module_tree())] + + assert names == ["_HOSTTAIL_OFF"] + [constant for constant, _ in ITEMS.values()] + for _, field in ITEMS.values(): + assert getattr(RuntimeHooks(), field) is False, field + + +def test_every_item_is_on_by_default(monkeypatch): + assert _hosttail(monkeypatch, {}) == {item: True for item in ITEMS} + + +@pytest.mark.parametrize("item", sorted(ITEMS)) +def test_each_kill_switch_turns_off_only_its_own_item(monkeypatch, item): + values = _hosttail(monkeypatch, {"TRTLLM_M31_DISABLE_" + item: "1"}) + + assert values == {other: other != item for other in ITEMS} + + +@pytest.mark.parametrize("value", ["1", "true", "on", "yes"]) +def test_the_bundle_switch_turns_off_all_six(monkeypatch, value): + assert _hosttail(monkeypatch, {"TRTLLM_M31_DISABLE_HOSTTAIL": value}) == { + item: False for item in ITEMS + } + + +@pytest.mark.parametrize("value", ["", "0", "false", "off", "no"]) +def test_an_unset_bundle_switch_keeps_all_six(monkeypatch, value): + assert _hosttail(monkeypatch, {"TRTLLM_M31_DISABLE_HOSTTAIL": value}) == { + item: True for item in ITEMS + } + + +def test_an_enable_variable_does_not_turn_an_item_off_or_on(monkeypatch): + env = {"TRTLLM_M31_ENABLE_" + item: "1" for item in _SWITCH_IDS} + assert _hosttail(monkeypatch, env) == {item: True for item in ITEMS} + + +def test_runtime_hooks_wires_each_field_to_its_switch(): + """``_runtime_hooks`` passes each HOSTTAIL_* constant to its RuntimeHooks field, and nothing else sets those.""" + fn = next( + node + for node in _module_tree().body + if isinstance(node, ast.FunctionDef) and node.name == "_runtime_hooks" + ) + calls = [ + node + for node in ast.walk(fn) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "RuntimeHooks" + ] + assert len(calls) == 1 + keywords = {kw.arg: kw.value for kw in calls[0].keywords} + for constant, field in ITEMS.values(): + assert isinstance(keywords.get(field), ast.Name), field + assert keywords[field].id == constant, field + + +def test_the_target_kv_cache_is_not_a_mamba_cache(): + """skip_accepted_draft_upload is on only together with a non-mamba cache: the mamba state update is the one + consumer of accepted counts on the drafter path (it takes them as an argument, not from the spec metadata).""" + assert not issubclass(MiniMaxM3KVCacheManagerV2, BaseMambaCacheManager) + assert not issubclass(MiniMaxM3KVCacheManagerV2, MambaHybridCacheManager) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="imports the target's modeling module") +def test_the_in_tree_hooks_follow_the_switches(): + modeling = importlib.import_module( + f"tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.{TARGET}.modeling" + ) + hooks = modeling._runtime_hooks(None) + + for constant, field in ITEMS.values(): + assert getattr(hooks, field) is getattr(modeling, constant), field + + def disabled(item): + return os.environ.get("TRTLLM_M31_DISABLE_" + item, "").strip().lower() in ( + "1", + "true", + "on", + "yes", + ) + + for item, (_, field) in ITEMS.items(): + assert getattr(hooks, field) is not (disabled("HOSTTAIL") or disabled(item)), item diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_switches.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_switches.py new file mode 100644 index 000000000000..81a862f2a624 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_switches.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The value contract of this target's kill switches (m31_switches): 1 / true / on / yes set TRTLLM_M31_DISABLE_ / +TRTLLM_M31_ENABLE_, empty / 0 / false / off / no leave it unset (case-insensitive), anything else raises at import +naming the variable; the environment is read once, at import. + +Each case loads a private copy of the module from its file, so the in-tree module (read by every other module of the +target at import) is never reloaded. CPU only. +""" + +import ast +import importlib.util +import os +from pathlib import Path + +import pytest + +import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 as _family + +TARGET = "minimax_m31_nvfp4__sm_103__tp4" +_SWITCHES = Path(_family.__file__).resolve().parent / TARGET / "m31_switches.py" + +pytestmark = pytest.mark.cpu_only + +KINDS = [("TRTLLM_M31_DISABLE_", "m31_disabled"), ("TRTLLM_M31_ENABLE_", "m31_enabled")] +ITEM = "UTCONTRACT" # an id no item uses, so the real switches of the test process cannot interfere + + +def _load(monkeypatch, env): + """A private copy of m31_switches, imported under ``env`` (every other TRTLLM_M31_* variable cleared).""" + for name in list(os.environ): + if name.startswith(("TRTLLM_M31_DISABLE_", "TRTLLM_M31_ENABLE_")): + monkeypatch.delenv(name) + for name, value in env.items(): + monkeypatch.setenv(name, value) + spec = importlib.util.spec_from_file_location(f"_m31_switches_contract_{id(env)}", _SWITCHES) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize("prefix,query", KINDS) +@pytest.mark.parametrize( + "value", ["", " ", "0", " 0 ", "false", "FALSE", "off", " Off ", "no", "No"] +) +def test_unset_spellings_leave_the_switch_unset(monkeypatch, prefix, query, value): + assert not getattr(_load(monkeypatch, {}), query)(ITEM) + assert not getattr(_load(monkeypatch, {prefix + ITEM: value}), query)(ITEM) + + +@pytest.mark.parametrize("prefix,query", KINDS) +@pytest.mark.parametrize("value", ["1", " 1 ", "true", "TRUE", "on", "On", "yes", "YES"]) +def test_set_spellings_set_the_switch(monkeypatch, prefix, query, value): + assert getattr(_load(monkeypatch, {prefix + ITEM: value}), query)(ITEM) + + +@pytest.mark.parametrize("prefix", [k for k, _ in KINDS]) +@pytest.mark.parametrize("value", ["2", "-1", "1.0", "y", "n", "enable", "disabled", "truee", "of"]) +def test_any_other_value_fails_at_import_naming_the_variable(monkeypatch, prefix, value): + """Typos must fail loudly: invalid switch values must not silently disable a feature.""" + with pytest.raises(ValueError, match=prefix + ITEM): + _load(monkeypatch, {prefix + ITEM: value}) + + +def test_the_id_is_matched_exactly(monkeypatch): + switches = _load(monkeypatch, {"TRTLLM_M31_DISABLE_" + ITEM: "1"}) + assert switches.m31_disabled(ITEM) + for other in (ITEM.lower(), ITEM + "_X", ITEM[:-1], ""): + assert not switches.m31_disabled(other), other + + +def test_disable_and_enable_are_independent(monkeypatch): + """Both read: an item's own code applies "DISABLE wins" (e.g. B1, L-DS1); the module reports each separately.""" + switches = _load( + monkeypatch, {"TRTLLM_M31_DISABLE_" + ITEM: "1", "TRTLLM_M31_ENABLE_" + ITEM: "1"} + ) + assert switches.m31_disabled(ITEM) and switches.m31_enabled(ITEM) + + +def test_the_environment_is_read_once_at_import(monkeypatch): + switches = _load(monkeypatch, {}) + monkeypatch.setenv("TRTLLM_M31_DISABLE_" + ITEM, "1") + monkeypatch.setenv("TRTLLM_M31_ENABLE_" + ITEM, "1") + assert not switches.m31_disabled(ITEM) and not switches.m31_enabled(ITEM) + + +# The step-input pack items (VR2GB-01 / VR2GB-03) and the order each one needs: a killed prerequisite turns off +# every item that needs it. +_PACK_FLAGS = ("PACK_PSCATTER", "PACK_IN_GRAPH", "PACK_EARLY_COPY", "H2D_CHANGED_SPAN") +_PACK_HOOKS = { + "pack_parallel_scatter": "PACK_PSCATTER", + "pack_in_graph": "PACK_IN_GRAPH", + "pack_early_copy": "PACK_EARLY_COPY", + "h2d_changed_span": "H2D_CHANGED_SPAN", +} + + +def _modeling_tree(): + return ast.parse((_SWITCHES.parent / "modeling.py").read_text()) + + +def _pack_flags(disabled): + """The target's pack flags as its module computes them at import, with ``disabled`` the switched-off ids (the + flags' own statements, evaluated with a stand-in m31_disabled: no import of the target).""" + statements = [ + node + for node in _modeling_tree().body + if isinstance(node, ast.Assign) + and len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id in _PACK_FLAGS + ] + assert sorted(node.targets[0].id for node in statements) == sorted(_PACK_FLAGS) + namespace = {"m31_disabled": lambda item_id: item_id in disabled} + exec(compile(ast.Module(body=statements, type_ignores=[]), "", "exec"), namespace) + return {name: namespace[name] for name in _PACK_FLAGS} + + +@pytest.mark.parametrize( + "disabled,expect_on", + [ + ((), _PACK_FLAGS), + (("H2D_CHANGED_SPAN",), ("PACK_PSCATTER", "PACK_IN_GRAPH", "PACK_EARLY_COPY")), + (("PACK_EARLY_COPY",), ("PACK_PSCATTER", "PACK_IN_GRAPH", "H2D_CHANGED_SPAN")), + (("PACK_IN_GRAPH",), ("PACK_PSCATTER", "H2D_CHANGED_SPAN")), + (("PACK_PSCATTER",), ("H2D_CHANGED_SPAN",)), + (("PACK_PSCATTER", "H2D_CHANGED_SPAN"), ()), + ], +) +def test_pack_items_default_on_and_a_killed_prerequisite_turns_its_dependents_off( + disabled, expect_on +): + flags = _pack_flags(set(disabled)) + assert {name for name, on in flags.items() if on} == set(expect_on) + + +def test_runtime_hooks_take_the_pack_flags(): + """_runtime_hooks passes each pack flag to its RuntimeHooks field (not a literal True).""" + (function,) = [ + node + for node in _modeling_tree().body + if isinstance(node, ast.FunctionDef) and node.name == "_runtime_hooks" + ] + (call,) = [ + node + for node in ast.walk(function) + if isinstance(node, ast.Call) and getattr(node.func, "id", None) == "RuntimeHooks" + ] + passed = {kw.arg: kw.value for kw in call.keywords} + for field, flag in _PACK_HOOKS.items(): + assert isinstance(passed.get(field), ast.Name) and passed[field].id == flag, field + assert ( + isinstance(passed["skip_unchanged_h2d"], ast.Constant) + and passed["skip_unchanged_h2d"].value + ) + assert isinstance(passed["pack_step_inputs"], ast.Constant) and passed["pack_step_inputs"].value diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_tp4_ar64auto.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_tp4_ar64auto.py new file mode 100644 index 000000000000..1ad6d30cf627 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_tp4_ar64auto.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""AR64AUTO dispatch (CPU): the generation boundaries past the multicast all-reduce's rows, up to the decode bound, +take the strategy-MNNVL instance; every other row count, and a core without it (kill switch), the configured one.""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4 import ( + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.ops import ( + m3_mc_allreduce, +) + + +@pytest.fixture +def taken(monkeypatch): + seen = [] + monkeypatch.setattr( + dispatch, "m3_allreduce", lambda x, comm, params: seen.append(comm) or (x, x) + ) + return seen + + +# AllReduceParams(fusion_op=RESIDUAL_RMS_NORM) asserts a residual; the patched m3_allreduce never reads these. +_RESIDUAL = object() +_NORM_W = object() + + +class _Rows: + def __init__(self, m: int): + self.shape = (m, 6144) + + def contiguous(self): + return self + + +def _core(wide=True): + return SimpleNamespace( + allreduce="configured", allreduce_wide="auto" if wide else None, rms_eps=1e-6 + ) + + +def test_the_wide_rows_take_the_mnnvl_instance(taken) -> None: + lo, hi = m3_mc_allreduce.max_rows(), dispatch.DECODE_MAX_TOKENS + assert hi > lo, ( + "the decode bound must exceed the multicast all-reduce's rows for AR64AUTO to apply" + ) + for m, want in ((1, "configured"), (lo, "configured"), (lo + 1, "auto"), (hi, "auto"), (hi + 1, "configured"), + (2048, "configured")): # fmt: skip + dispatch._native_allreduce_norm(_core(), _Rows(m), _RESIDUAL, _NORM_W) + assert taken[-1] == want, (m, taken[-1]) + + +def test_without_the_auto_instance_every_row_count_takes_the_configured_one(taken) -> None: + for m in (1, m3_mc_allreduce.max_rows() + 1, dispatch.DECODE_MAX_TOKENS): + dispatch._native_allreduce_norm(_core(wide=False), _Rows(m), _RESIDUAL, _NORM_W) + assert taken[-1] == "configured", m + + +def test_a_declined_wide_call_runs_the_configured_instance(monkeypatch) -> None: + """MNNVL returning None (a shape it does not take) must not fall to AUTO: the configured instance runs.""" + seen = [] + + def fake(x, comm, params): + seen.append(comm) + return None if comm == "auto" else (x, x) + + monkeypatch.setattr(dispatch, "m3_allreduce", fake) + dispatch._native_allreduce_norm( + _core(), _Rows(m3_mc_allreduce.max_rows() + 1), _RESIDUAL, _NORM_W + ) + assert seen == ["auto", "configured"] + + +def test_the_kill_switch_is_read_at_import() -> None: + assert isinstance(dispatch.AR64AUTO, bool) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_tp4_fused_projection_default.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_tp4_fused_projection_default.py new file mode 100644 index 000000000000..90360a858a56 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_tp4_fused_projection_default.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The M3.1 TP4 target packs the sparse layers' QKV + index projection by default (USER rule: measured wins ship on +by default); an explicit ``fuse_qkv_index_projection`` in the deployment's config wins, so False is the kill switch.""" + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4 import ( + weights as tp_weights, +) +from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig + +_MSA = dict(implementation="msa", indexer_kv_dtype="fp8") + + +@pytest.mark.parametrize( + ("explicit", "expected"), + [ + ({}, True), + ({"fuse_qkv_index_projection": True}, True), + ({"fuse_qkv_index_projection": False}, False), + ], +) +def test_fused_projection_default_and_kill(explicit: dict, expected: bool) -> None: + config = MiniMaxM3SparseAttentionConfig(**_MSA, **explicit) + assert tp_weights.fused_projection(config) is expected + + +def test_no_sparse_config_is_unfused() -> None: + assert tp_weights.fused_projection(None) is False + + +@pytest.mark.parametrize("value", [True, False]) +def test_a_config_without_set_fields_is_taken_at_its_word(value: bool) -> None: + from types import SimpleNamespace + + assert tp_weights.fused_projection(SimpleNamespace(fuse_qkv_index_projection=value)) is value diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_tp4_weights_layout.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_tp4_weights_layout.py new file mode 100644 index 000000000000..8eb0dd6a700d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/test_modeling_v2_m31_tp4_weights_layout.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3.1 TP4 weight table: layer kinds from the config, packed-NVFP4 expert scalars as FP32 divisors.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4 import ( + weights as W, +) + +pytestmark = pytest.mark.cpu_only + + +def _core(sparse_freq, moe_freq, fused=True): + cfg = SimpleNamespace( + num_hidden_layers=len(sparse_freq), + moe_layer_freq=moe_freq, + sparse_attention_config=dict( + sparse_attention_freq=sparse_freq, sparse_num_index_heads=4, sparse_index_dim=128 + ), + head_dim=128, + num_attention_heads=64, + num_key_value_heads=4, + hidden_size=6144, + dense_intermediate_size=12288, + shared_intermediate_size=3072, + n_shared_experts=1, + intermediate_size=3072, + num_local_experts=128, + vocab_size=200064, + torch_dtype=torch.bfloat16, + ) + return SimpleNamespace( + model_config=SimpleNamespace( + pretrained_config=cfg, + mapping=SimpleNamespace(tp_size=4, tp_rank=0), + sparse_attention_config=SimpleNamespace(fuse_qkv_index_projection=fused), + ) + ) + + +def test_layer_kinds_come_from_the_config_independently(): + """Sparse attention on every layer while layers 0-1 keep the dense MLP.""" + d = W.MODEL_WEIGHTS.dims(_core([1] * 6, [0, 0, 1, 1, 1, 1])) + assert list(W._sparse_attn_layers(d)) == [0, 1, 2, 3, 4, 5] + assert list(W._dense_mlp_layers(d)) == [0, 1] + assert list(W._moe_layers(d)) == [2, 3, 4, 5] + assert list(W._five_way_qkv_layers(d)) == [0, 1, 2, 3, 4, 5] + assert list(W._three_way_qkv_layers(d)) == [] + + +def test_unfused_projection_uses_the_three_way_qkv_everywhere(): + d = W.MODEL_WEIGHTS.dims(_core([0, 0, 1, 1], [0, 0, 1, 1], fused=False)) + assert list(W._three_way_qkv_layers(d)) == [0, 1, 2, 3] + assert list(W._unfused_sparse_layers(d)) == [2, 3] + assert list(W._five_way_qkv_layers(d)) == [] + + +def test_expert_scalars_are_the_fp32_reciprocal_of_the_stored_multiplier(): + stored = torch.tensor(12345.678, dtype=torch.float32) + out = W._divisor(stored, None) + assert out.dtype == torch.float32 + assert out.view(torch.int32).item() == torch.reciprocal(stored).view(torch.int32).item() + with pytest.raises(ValueError): + W._divisor(torch.tensor(0.0), None) diff --git a/tests/unittest/_torch/modeling_v2/spec_decode/_dspark_markov_chain_ll_op_matrix.py b/tests/unittest/_torch/modeling_v2/spec_decode/_dspark_markov_chain_ll_op_matrix.py new file mode 100644 index 000000000000..1f314443f93a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/spec_decode/_dspark_markov_chain_ll_op_matrix.py @@ -0,0 +1,671 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Certify `spec_decode/dspark_markov_chain`'s LL exchange variant (``enable_ll_exchange``, kernel +``dspark_markov_chain_kernel_b1_ll``): its batch-1 draft tokens equal the S4B8 batch-1 kernel's +(``dspark_markov_chain_kernel_b1``) bit for bit, on real ranks over the multicast workspace. + +Every rank builds the same full-vocabulary tensors from one seed and passes its padded shard to the op; every rank +checks LL == b1 on its own tokens, rank 0 also checks both against ``reference``. Checks (batch 1, K = 3 and 7, the +M3.1 vocabulary): + +- inputs: random logits, logits on a coarse grid (many exact ties), all-equal values (zero ``markov_w2``: every + position's token is 0), NaN-free extremes (+-bf16 max, +-the smallest bf16 denormal, +-0.0, -inf); +- forced ties at fixed vocabulary columns (their ``markov_w2`` rows zeroed, so the value is the logit exactly): across + ranks at one local column, the last real column of rank r against the first of rank r + 1, the last row of CTA 0 + against the first of CTA 1, two warps of one CTA, the last real column of every rank; the tied value a large + finite value, the bf16 maximum, the smallest bf16 denormal (every other column -inf) or +-0.0 (every other column + -inf). The token must be the lowest tied column at every position; +- stale-tag safety: 64 back-to-back launches with changing inputs (no host sync), K = 7 and K = 3 interleaved with + batch-2 calls of the Lamport kernel, then 60 back-to-back K = 7 launches; every launch's tokens equal the b1 + kernel's on its inputs; +- a CUDA graph of four LL launches replayed 16 times with new inputs copied in before each replay, against eager b1; +- TAILBK_SPEC cached inputs (``enable_cached_inputs``): tokens with the cache equal tokens without it (LL batch 1 and + the Lamport kernel at batch 2, eager and in a CUDA graph); the cache holds ones / arange / zeros, is reused, never + written, and a shape first met under capture allocates without caching; +- TAILBK_SPEC KV-length rewind (``dspark_markov_chain_rewind``): its tokens equal ``dspark_markov_chain``'s and its + ``kv_lens`` equal the stock ``kv_lens[nc:bs] -= rewind; kv_lens[nc:bs].clamp_(min=0)`` (rows outside untouched), + for rewinds of both signs and clamped rows, eager and in a CUDA graph; ``supports_rewind`` rejects bad operands. + +World 4 (one GB300 tray) is spawned from one process: + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _dspark_markov_chain_ll_op_matrix.py + +World 8 spans two trays and runs as one srun task per rank: + + MASTER_ADDR= MASTER_PORT= python _dspark_markov_chain_ll_op_matrix.py --srun + +under ``srun --ntasks=8 --ntasks-per-node=4 --kill-on-bad-exit=1``. +""" + +import importlib.util +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, List, Optional + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "dspark_markov_chain requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 1500 +RANK_DIM = 256 +VOCAB = 200064 +SEEDS = 4 +BACK_TO_BACK = 64 +GPUS_PER_NODE = 4 +"""A GB300 tray. The mapping needs it to see a world-8 group span two nodes.""" + +BF16_MAX = torch.finfo(torch.bfloat16).max +BF16_DENORM_MIN = 2.0**-133 # the smallest positive bf16 subnormal + +RANK = 0 +WORLD = 1 +entry: Any = None +chain: Any = None +WS: dict = {} + + +def _set_ll(on: bool) -> None: + entry.enable_batch1_variant(True) + entry.enable_ll_exchange(on) + + +def _weights(vocab: int, zero_rows: Optional[List[int]] = None, zero_w2: bool = False): + g = torch.Generator(device="cuda").manual_seed(1234) + w1 = (torch.randn(vocab, RANK_DIM, device="cuda", generator=g) * 0.05).bfloat16() + w2 = (torch.randn(vocab, RANK_DIM, device="cuda", generator=g) * 0.05).bfloat16() + if zero_w2: + w2.zero_() + if zero_rows: + w2[torch.tensor(zero_rows, device="cuda")] = 0 + return w1, w2 + + +class _Shard: + """One vocabulary's shards: this rank's padded `w2` rows and logits, and every rank's for the reference.""" + + def __init__(self, vocab: int, **weights): + self.vocab = vocab + self.shard = vocab // WORLD + self.rows = entry.padded_rows(self.shard, entry.CERTIFIED_BATCH) + assert self.rows >= self.shard, f"no padded shard for {self.shard} rows" + self.w1, w2 = _weights(vocab, **weights) + self.w2s_all = [ + entry.pad_w2_shard(w2[r * self.shard : (r + 1) * self.shard].contiguous(), self.rows) + for r in range(WORLD) + ] + self.w2s = self.w2s_all[RANK] + + def base(self, full: torch.Tensor, rank: int) -> torch.Tensor: + b = entry.pad_logits_buffer(full.shape[0], full.shape[1], self.rows, full.device) + b[:, :, : self.shard] = full[:, :, rank * self.shard : (rank + 1) * self.shard] + return b + + def call(self, base: torch.Tensor, prev0: torch.Tensor) -> torch.Tensor: + return chain(base, prev0, self.w1, self.w2s, RANK * self.shard, **WS) + + def both(self, base: torch.Tensor, prev0: torch.Tensor): + _set_ll(False) + b1 = self.call(base, prev0) + _set_ll(True) + ll = self.call(base, prev0) + return b1, ll + + def reference(self, full: torch.Tensor, prev0: torch.Tensor) -> torch.Tensor: + bases = [self.base(full, r) for r in range(WORLD)] + return chain.reference(bases, prev0, self.w1, self.w2s_all, self.shard) + + def rows_per_cta(self, block: int) -> int: + grid = entry.pick_grid(self.rows, block, 1) + assert grid > 0 + return self.rows // grid + + +def _inputs(seed: int, batch: int, block: int, vocab: int, kind: str): + g = torch.Generator(device="cuda").manual_seed(seed) + base = torch.randn(batch, block, vocab, device="cuda", generator=g) * 3 + if kind == "ties": + base = torch.round(base * 2) / 2 + elif kind == "equal": + base = torch.ones_like(base) + elif kind == "extremes": + pool = torch.tensor( + [ + BF16_MAX, + -BF16_MAX, + BF16_DENORM_MIN, + -BF16_DENORM_MIN, + 0.0, + -0.0, + float("-inf"), + 1.0, + -1.0, + ], + device="cuda", + ) + base = pool[torch.randint(0, pool.numel(), base.shape, device="cuda", generator=g)] + prev0 = torch.randint(0, vocab, (batch,), device="cuda", generator=g) + return base.bfloat16(), prev0 + + +def _assert_equal(what: str, ll: torch.Tensor, b1: torch.Tensor) -> None: + assert torch.equal(ll, b1), f"{what}: LL {ll.tolist()} != b1 {b1.tolist()}" + + +def check_ll_matches_b1() -> None: + sh = _Shard(VOCAB) + sh_eq = _Shard(VOCAB, zero_w2=True) + for block in (3, 7): + for kind in ("random", "ties", "equal", "extremes"): + s = sh_eq if kind == "equal" else sh + for seed in range(SEEDS): + full, prev0 = _inputs(seed, 1, block, s.vocab, kind) + b1, ll = s.both(s.base(full, RANK), prev0) + _assert_equal(f"K {block} {kind} seed {seed}", ll, b1) + if kind == "equal": + assert ll.eq(0).all(), f"all-equal values: tokens {ll.tolist()} are not all 0" + if RANK == 0: + ref = s.reference(full, prev0) + _assert_equal(f"K {block} {kind} seed {seed} vs reference", ll, ref) + dist.barrier() + if RANK == 0: + print( + f"[ll] world={WORLD} LL == b1 (random, ties, all-equal, extremes; K 3 / 7)", flush=True + ) + + +def _tie_patterns(shard: int, rc: int) -> List[List[int]]: + """Global vocabulary columns that tie, one list per pattern (each pattern's token: its lowest column).""" + last = ( + shard - 1 + ) # the last real local column (in CTA last // rc, the last CTA holding real rows) + rows_per_warp = rc // 8 + pats = [ + [r * shard + 5 for r in range(1, WORLD)], # across ranks at one local column (not rank 0) + [0 * shard + last, 1 * shard + 0], # last real column of rank 0 vs first column of rank 1 + [ + (WORLD - 1) * shard + 0, + (WORLD - 2) * shard + last, + ], # ... of rank W - 2 vs rank W - 1, listed high first + [rc, rc - 1], # rank 0: CTA 0's last row vs CTA 1's first row + [ + 2 * shard + rc, + 2 * shard + rc - 1, + 3 % WORLD * shard + rc, + ], # the same across CTAs, on other ranks + [1 * shard + rows_per_warp + 3, 1 * shard + 3], # rank 1, CTA 0: warp 1 vs warp 0 + [r * shard + last for r in range(WORLD)][::-1], # every rank's last real column + [(WORLD - 1) * shard + 7, 7], # rank W - 1 vs rank 0 + ] + return [sorted(set(p)) for p in pats] + + +def check_forced_ties() -> None: + probe = _Shard.__new__(_Shard) + probe.shard = VOCAB // WORLD + probe.rows = entry.padded_rows(probe.shard, entry.CERTIFIED_BATCH) + rcs = {block: probe.rows_per_cta(block) for block in (3, 7)} + pats_by_block = {block: _tie_patterns(probe.shard, rcs[block]) for block in (3, 7)} + zero_rows = sorted({c for pats in pats_by_block.values() for p in pats for c in p}) + sh = _Shard(VOCAB, zero_rows=zero_rows) + values = ( + ("large", 100.0, False), + ("bf16max", BF16_MAX, False), + ("denormal", BF16_DENORM_MIN, True), + ("signed-zero", 0.0, True), + ) + checked = 0 + for block in (3, 7): + pats = pats_by_block[block] + for name, tie, others_neg_inf in values: + for offset in range(len(pats)): + full, prev0 = _inputs(1000 + offset, 1, block, sh.vocab, "random") + full = full.float() + expect = [] + for k in range(block): + cols = pats[(offset + k) % len(pats)] + if others_neg_inf: + full[0, k, :] = float("-inf") + idx = torch.tensor(cols, device="cuda") + if name == "signed-zero": + # -0.0 and +0.0 alternate over the tied columns: equal under the order. + signs = torch.tensor( + [(-0.0 if i % 2 == 0 else 0.0) for i in range(len(cols))], device="cuda" + ) + full[0, k, idx] = signs + else: + full[0, k, idx] = tie + expect.append(min(cols)) + full = full.bfloat16() + b1, ll = sh.both(sh.base(full, RANK), prev0) + what = f"forced ties {name} K {block} offset {offset}" + _assert_equal(what, ll, b1) + want = torch.tensor([expect], dtype=torch.int32, device=ll.device) + assert torch.equal(ll, want), ( + f"{what}: tokens {ll.tolist()} != lowest tied columns {expect}" + ) + if RANK == 0: + _assert_equal(f"{what} vs reference", ll, sh.reference(full, prev0)) + checked += 1 + dist.barrier() + if RANK == 0: + print( + f"[ll] world={WORLD} forced ties: {checked} cases, LL == b1 == lowest tied column", + flush=True, + ) + + +def check_back_to_back_stale_tags() -> None: + sh = _Shard(VOCAB) + # Warm every shape once (compiles run outside any sequence below). + for block, batch, ll in ( + (7, 1, True), + (3, 1, True), + (3, 2, True), + (7, 1, False), + (3, 1, False), + ): + full, prev0 = _inputs(1, batch, block, sh.vocab, "random") + _set_ll(ll) + sh.call(sh.base(full, RANK), prev0) + torch.cuda.synchronize() + dist.barrier() + + def run_sequence(shapes, seed0: int, tag: str) -> None: + launches = [] + for i, (block, batch) in enumerate(shapes): + full, prev0 = _inputs( + seed0 + i, batch, block, sh.vocab, "ties" if i % 3 == 0 else "random" + ) + launches.append((sh.base(full, RANK), prev0)) + torch.cuda.synchronize() + _set_ll(True) + outs = [sh.call(base, prev0) for base, prev0 in launches] # back to back, no host sync + torch.cuda.synchronize() + _set_ll(False) + refs = [sh.call(base, prev0) for base, prev0 in launches] + torch.cuda.synchronize() + for i, (got, ref) in enumerate(zip(outs, refs)): + _assert_equal(f"{tag} launch {i} (K {shapes[i][0]}, batch {shapes[i][1]})", got, ref) + + mixed = [((7, 1), (3, 1), (7, 1), (3, 2))[i % 4] for i in range(BACK_TO_BACK)] + run_sequence(mixed, 5000, "mixed back-to-back") + run_sequence([(7, 1)] * 60, 7000, "K7 back-to-back") + dist.barrier() + if RANK == 0: + print( + f"[ll] world={WORLD} back-to-back: {BACK_TO_BACK} mixed + 60 K7 launches equal b1", + flush=True, + ) + + +def check_cuda_graph() -> None: + sh = _Shard(VOCAB) + blocks = (7, 3, 7, 7) + static = [] + for j, block in enumerate(blocks): + full, prev0 = _inputs(300 + j, 1, block, sh.vocab, "random") + static.append((sh.base(full, RANK), prev0.clone())) + _set_ll(True) + for base, prev0 in static: # compiled outside capture + sh.call(base, prev0) + torch.cuda.synchronize() + dist.barrier() + graph_out = [] + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + g = torch.cuda.CUDAGraph() + with torch.cuda.stream(s): + with torch.cuda.graph(g): + for base, prev0 in static: + graph_out.append(sh.call(base, prev0)) + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + dist.barrier() + for rep in range(16): + for j, (block, (base, prev0)) in enumerate(zip(blocks, static)): + full, p = _inputs( + 400 + rep * 8 + j, 1, block, sh.vocab, "ties" if rep % 2 else "random" + ) + base.copy_(sh.base(full, RANK)) + prev0.copy_(p) + g.replay() + torch.cuda.synchronize() + got = [t.clone() for t in graph_out] + _set_ll(False) + refs = [sh.call(base, prev0) for base, prev0 in static] + torch.cuda.synchronize() + _set_ll(True) + for j, (a, b) in enumerate(zip(got, refs)): + _assert_equal(f"graph replay {rep} launch {j} (K {blocks[j]})", a, b) + dist.barrier() + if RANK == 0: + print( + f"[ll] world={WORLD} CUDA graph: 16 replays x 4 LL launches equal eager b1", flush=True + ) + + +def check_cached_inputs() -> None: + sh = _Shard(VOCAB) + shapes = ((7, 1), (3, 1), (7, 2)) + _set_ll(True) + entry._cached_taken[0] = False + entry.enable_cached_inputs(True) + for block, batch in shapes: # outside capture: compiles and fills the cache + full, prev0 = _inputs(11, batch, block, sh.vocab, "random") + sh.call(sh.base(full, RANK), prev0) + torch.cuda.synchronize() + dev = torch.device("cuda", torch.cuda.current_device()) + cached = {} + for block, batch in shapes: + key = (dev, batch, block) + assert key in entry._const_inputs, f"no cached inputs for batch {batch} K {block}" + nacc, rows, acc = entry._const_inputs[key] + assert nacc.dtype == rows.dtype == acc.dtype == torch.int32 + assert torch.equal(nacc, torch.ones(batch, dtype=torch.int32, device=dev)) + assert torch.equal(rows, torch.arange(batch, dtype=torch.int32, device=dev)) + assert torch.equal(acc, torch.zeros(batch, block + 1, dtype=torch.int32, device=dev)) + cached[key] = (nacc, rows, acc) + + assert not entry.cached_inputs_taken(), "a cache fill alone is not the engagement" + + def both(base, prev0): + entry.enable_cached_inputs(False) + off = sh.call(base, prev0) + entry.enable_cached_inputs(True) + on = sh.call(base, prev0) + return off, on + + for block, batch in shapes: + for kind in ("random", "ties", "extremes"): + for seed in range(SEEDS): + full, prev0 = _inputs(600 + seed, batch, block, sh.vocab, kind) + off, on = both(sh.base(full, RANK), prev0) + _assert_equal(f"cached inputs K {block} batch {batch} {kind} seed {seed}", on, off) + assert entry.cached_inputs_taken(), "a call with the cache filled must report the inputs taken" + # A shape compiled with the cache off, then first met with it on under capture: allocated, not cached. + miss = (dev, 3, 7) + entry._const_inputs.pop(miss, None) + entry.enable_cached_inputs(False) + full, prev0 = _inputs(13, 3, 7, sh.vocab, "random") + sh.call(sh.base(full, RANK), prev0) + torch.cuda.synchronize() + dist.barrier() + + graph_shapes = ((7, 1), (3, 1), (7, 2), (7, 3)) + static = [] + for j, (block, batch) in enumerate(graph_shapes): + full, prev0 = _inputs(700 + j, batch, block, sh.vocab, "random") + static.append((sh.base(full, RANK), prev0.clone())) + entry.enable_cached_inputs(True) + graph_out = [] + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + g = torch.cuda.CUDAGraph() + with torch.cuda.stream(s): + with torch.cuda.graph(g): + for base, prev0 in static: + graph_out.append(sh.call(base, prev0)) + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + dist.barrier() + assert miss not in entry._const_inputs, "a shape first met under capture was cached" + for key, tensors in cached.items(): + assert all(a is b for a, b in zip(entry._const_inputs[key], tensors)), ( + f"cache {key} replaced" + ) + for rep in range(8): + for j, ((block, batch), (base, prev0)) in enumerate(zip(graph_shapes, static)): + full, p = _inputs( + 800 + rep * 8 + j, batch, block, sh.vocab, "ties" if rep % 2 else "random" + ) + base.copy_(sh.base(full, RANK)) + prev0.copy_(p) + g.replay() + torch.cuda.synchronize() + got = [t.clone() for t in graph_out] + entry.enable_cached_inputs(False) + refs = [sh.call(base, prev0) for base, prev0 in static] + torch.cuda.synchronize() + entry.enable_cached_inputs(True) + for j, (a, b) in enumerate(zip(got, refs)): + _assert_equal(f"cached-inputs graph replay {rep} launch {j} {graph_shapes[j]}", a, b) + for key, (nacc, rows, acc) in cached.items(): # never written by the kernels + batch, block = key[1], key[2] + assert torch.equal(nacc, torch.ones(batch, dtype=torch.int32, device=dev)) + assert torch.equal(rows, torch.arange(batch, dtype=torch.int32, device=dev)) + assert torch.equal(acc, torch.zeros(batch, block + 1, dtype=torch.int32, device=dev)) + entry.enable_cached_inputs(False) + dist.barrier() + if RANK == 0: + print( + f"[ll] world={WORLD} cached inputs: eager + CUDA graph tokens equal uncached; cache intact", + flush=True, + ) + + +def _stock_rewind(kv_lens: torch.Tensor, amount: torch.Tensor, nc: int) -> torch.Tensor: + want = kv_lens.clone() + bs = nc + amount.numel() + want[nc:bs] -= amount + want[nc:bs].clamp_(min=0) + return want + + +def _rewind_case(seed: int, batch: int, nc: int): + g = torch.Generator(device="cuda").manual_seed(seed) + kv_lens = torch.randint(0, 40, (nc + batch + 3,), dtype=torch.int32, device="cuda", generator=g) + amount = torch.randint(-8, 9, (batch,), dtype=torch.int32, device="cuda", generator=g) + amount[0] = int(kv_lens[nc]) + 5 # a clamped row + return kv_lens, amount + + +def check_rewind() -> None: + sh = _Shard(VOCAB) + entry.enable_cached_inputs(False) + for ll in (True, False): + _set_ll(ll) + for block, batch, nc in ((7, 1, 0), (3, 1, 2), (7, 2, 1), (7, 3, 0)): + if batch > 1 and not ll: + continue + for seed in range(SEEDS): + full, prev0 = _inputs( + 900 + seed, batch, block, sh.vocab, "ties" if seed % 2 else "random" + ) + base = sh.base(full, RANK) + kv_lens, amount = _rewind_case(950 + seed * 7 + batch, batch, nc) + want = _stock_rewind(kv_lens, amount, nc) + ref = sh.call(base, prev0) + assert entry.supports_rewind(kv_lens, amount, nc) + got = entry.dspark_markov_chain_rewind(base, prev0, sh.w1, sh.w2s, RANK * sh.shard, kv_lens, + amount, nc, **WS) # fmt: skip + torch.cuda.synchronize() + tag = f"rewind ll={ll} K {block} batch {batch} nc {nc} seed {seed}" + _assert_equal(f"{tag} tokens", got, ref) + assert torch.equal(kv_lens, want), ( + f"{tag}: kv_lens {kv_lens.tolist()} != {want.tolist()}" + ) + i32 = torch.zeros(8, dtype=torch.int32, device="cuda") + assert not entry.supports_rewind(None, i32[:2], 0) + assert not entry.supports_rewind(i32, i32[:2].long(), 0) + assert not entry.supports_rewind(i32.long(), i32[:2], 0) + assert not entry.supports_rewind(i32, i32[:2], 7) # past the end + assert not entry.supports_rewind(i32, i32[:0], 0) + assert not entry.supports_rewind(i32[::2], i32[:2], 0) + + # A CUDA graph of the rewind launch (batch 1, K 7, nc 1), replayed with new inputs and kv_lens copied in. + _set_ll(True) + full, prev0 = _inputs(990, 1, 7, sh.vocab, "random") + base = sh.base(full, RANK) + kv_lens, amount = _rewind_case(991, 1, 1) + entry.dspark_markov_chain_rewind( + base, prev0, sh.w1, sh.w2s, RANK * sh.shard, kv_lens, amount, 1, **WS + ) + torch.cuda.synchronize() + dist.barrier() + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + g = torch.cuda.CUDAGraph() + with torch.cuda.stream(s): + with torch.cuda.graph(g): + out = entry.dspark_markov_chain_rewind(base, prev0, sh.w1, sh.w2s, RANK * sh.shard, kv_lens, + amount, 1, **WS) # fmt: skip + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + dist.barrier() + for rep in range(8): + f, p = _inputs(1000 + rep, 1, 7, sh.vocab, "ties" if rep % 2 else "random") + base.copy_(sh.base(f, RANK)) + prev0.copy_(p) + k, a = _rewind_case(1100 + rep, 1, 1) + kv_lens.copy_(k) + amount.copy_(a) + want = _stock_rewind(k, a, 1) + g.replay() + torch.cuda.synchronize() + got = out.clone() + ref = sh.call(base, prev0) + torch.cuda.synchronize() + _assert_equal(f"rewind graph replay {rep} tokens", got, ref) + assert torch.equal(kv_lens, want), ( + f"rewind graph replay {rep}: kv_lens {kv_lens.tolist()} != {want.tolist()}" + ) + dist.barrier() + if RANK == 0: + print( + f"[ll] world={WORLD} rewind: tokens equal the chain's, kv_lens equal the stock rewind", + flush=True, + ) + + +def check_switches() -> None: + _set_ll(True) + assert entry.ll_exchange_engaged(1) and not entry.ll_exchange_engaged(2) + entry.enable_batch1_variant(False) + assert not entry.ll_exchange_engaged(1), "LL needs the batch-1 variant" + _set_ll(False) + assert not entry.ll_exchange_engaged(1) + dist.barrier() + + +CHECKS = ( + check_switches, + check_ll_matches_b1, + check_forced_ties, + check_back_to_back_stale_tags, + check_cuda_graph, + check_cached_inputs, + check_rewind, +) + + +def _run_one_rank(device: int) -> int: + global RANK, WORLD, WS, entry, chain + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(device) + dist.init_process_group( + "cuda:nccl,cpu:gloo", init_method="env://", device_id=torch.device("cuda", device) + ) + RANK, WORLD = rank, world + + from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode import ( + dspark_markov_chain as entry_module, + ) + from tensorrt_llm.mapping import Mapping + + entry = entry_module + chain = entry.dspark_markov_chain + mapping = Mapping(world_size=world, rank=rank, gpus_per_node=GPUS_PER_NODE, tp_size=world) + WS = entry.workspace_args(mapping, dist.group.WORLD) + + 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() + os._exit(13) + print(f"[rank {RANK}] {check.__name__} passed", flush=True) + dist.barrier() + torch.cuda.synchronize() + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict(os.environ, MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port), WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1") # fmt: skip + 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, + ) # fmt: skip + 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}") + + +def _rank_job(): + """`_rank_job.py` beside this file, by path (a sibling directory ships its own module of that name).""" + path = Path(__file__).resolve().with_name("_rank_job.py") + spec = importlib.util.spec_from_file_location("_spec_decode_rank_job", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +if __name__ == "__main__": + if "--srun" in sys.argv: + sys.exit(_run_one_rank(device=_rank_job().adopt_srun_env())) + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank(device=int(os.environ["RANK"]))) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/spec_decode/_dspark_markov_chain_op_matrix.py b/tests/unittest/_torch/modeling_v2/spec_decode/_dspark_markov_chain_op_matrix.py new file mode 100644 index 000000000000..5f5692a621fa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/spec_decode/_dspark_markov_chain_op_matrix.py @@ -0,0 +1,287 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Certify `spec_decode/dspark_markov_chain`: the one-launch chain's draft tokens against the reference chain's, bit +for bit, on every cell certified at this run's world size. + +Every rank builds the same full-vocabulary tensors from one seed and passes its padded shard to the op; rank 0 checks +the tokens against `reference` over every rank's shard. Per cell: random logits and logits on a coarse grid (many +exact bf16 ties, so the tie order is exercised), each of the cell's batches, several seeds, inside `validating`. Also: +three repeated calls give identical tokens and a CUDA-graph replay gives the eager tokens; `supports` refuses a batch +above the kernel's capacity and fp32 logits; the guard refuses a padding column that is not -inf. + +World 4 (one GB300 tray) is spawned from one process: + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _dspark_markov_chain_op_matrix.py + +World 8 spans two trays and runs as one srun task per rank, like `comm/_mnnvl_allreduce_op_matrix.py`: + + MASTER_ADDR= MASTER_PORT= python _dspark_markov_chain_op_matrix.py --srun + +under ``srun --ntasks=8 --ntasks-per-node=4 --kill-on-bad-exit=1``. +""" + +import importlib.util +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Optional + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "dspark_markov_chain requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 +RANK_DIM = 256 +SEEDS = 8 +GPUS_PER_NODE = 4 +"""A GB300 tray. The mapping needs it to see a world-8 group span two nodes.""" + +RANK = 0 +WORLD = 1 +# Bound inside the rank body: `entry` is the module, `chain` its `dspark_markov_chain` instance (the call surface, +# `reference`, `is_valid` and `CELLS`), `WS` the workspace arguments of every call. +entry: Any = None +chain: Any = None +validating: Any = None +WS: dict = {} + + +def _weights(vocab: int): + g = torch.Generator(device="cuda").manual_seed(1234) + w1 = (torch.randn(vocab, RANK_DIM, device="cuda", generator=g) * 0.05).bfloat16() + w2 = (torch.randn(vocab, RANK_DIM, device="cuda", generator=g) * 0.05).bfloat16() + return w1, w2 + + +def _inputs(seed: int, batch: int, block: int, vocab: int, ties: bool): + g = torch.Generator(device="cuda").manual_seed(seed) + base = torch.randn(batch, block, vocab, device="cuda", generator=g) * 3 + if ties: + base = torch.round(base * 2) / 2 + prev0 = torch.randint(0, vocab, (batch,), device="cuda", generator=g) + return base.bfloat16(), prev0 + + +class _Shard: + """One vocabulary's shards: this rank's padded `w2` rows and logits, and every rank's for the reference.""" + + def __init__(self, vocab: int): + self.vocab = vocab + self.shard = vocab // WORLD + self.rows = entry.padded_rows(self.shard, entry.CERTIFIED_BATCH) + assert self.rows >= self.shard, f"no padded shard for {self.shard} rows" + self.w1, w2 = _weights(vocab) + self.w2s_all = [ + entry.pad_w2_shard(w2[r * self.shard : (r + 1) * self.shard].contiguous(), self.rows) + for r in range(WORLD) + ] + self.w2s = self.w2s_all[RANK] + + def base(self, full: torch.Tensor, rank: int) -> torch.Tensor: + b = entry.pad_logits_buffer(full.shape[0], full.shape[1], self.rows, full.device) + b[:, :, : self.shard] = full[:, :, rank * self.shard : (rank + 1) * self.shard] + return b + + def call(self, base: torch.Tensor, prev0: torch.Tensor) -> torch.Tensor: + return chain(base, prev0, self.w1, self.w2s, RANK * self.shard, **WS) + + def reference(self, full: torch.Tensor, prev0: torch.Tensor) -> torch.Tensor: + bases = [self.base(full, r) for r in range(WORLD)] + return chain.reference(bases, prev0, self.w1, self.w2s_all, self.shard) + + +def check_certified_cells() -> None: + cells = [c for c in chain.CELLS if c.spec["world"] == WORLD] + assert cells, f"no cell is certified at world size {WORLD}" + for cell in cells: + spec = cell.spec + sh = _Shard(spec["vocab"]) + for ties in (False, True): + for batch in spec["batches"]: + for seed in range(SEEDS): + full, prev0 = _inputs(seed, batch, spec["block"], sh.vocab, ties) + with validating(chain): + got = sh.call(sh.base(full, RANK), prev0) + if RANK == 0: + ref = sh.reference(full, prev0) + assert torch.equal(got, ref), ( + f"{cell.why}: ties {ties} batch {batch} seed {seed}: " + f"{(got != ref).any(1).sum().item()} request(s) differ" + ) + dist.barrier() + if RANK == 0: + print(f"[cell] world={WORLD} {cell.why}: bitwise", flush=True) + + +def check_deterministic_and_graph() -> None: + sh = _Shard(200064) + full, prev0 = _inputs(77, entry.CERTIFIED_BATCH, 7, sh.vocab, False) + base = sh.base(full, RANK) + outs = [sh.call(base, prev0) for _ in range(3)] + assert all(torch.equal(outs[0], o) for o in outs[1:]), "repeated calls differ" + graph_out = {} + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + g = torch.cuda.CUDAGraph() + with torch.cuda.stream(s): + with torch.cuda.graph(g): + graph_out["t"] = sh.call(base, prev0) + torch.cuda.current_stream().wait_stream(s) + for _ in range(3): + g.replay() + torch.cuda.synchronize() + assert torch.equal(graph_out["t"], outs[0]), "graph replay differs from eager" + dist.barrier() + + +def check_supports_refuses() -> None: + sh = _Shard(200064) + kern = entry._kernel_module() + full, prev0 = _inputs(5, kern.MAX_BATCH + 1, 3, sh.vocab, False) + assert not entry.supports(sh.base(full, RANK), sh.w1, sh.w2s), "a batch above MAX_BATCH" + full, prev0 = _inputs(5, 2, 3, sh.vocab, False) + assert not entry.supports(sh.base(full, RANK).float(), sh.w1, sh.w2s), "fp32 logits" + assert not entry.supports(sh.base(full, RANK)[:, :, :-64], sh.w1, sh.w2s), ( + "logits narrower than w2_shard" + ) + dist.barrier() + + +def check_the_guard_refuses_a_live_padding_column() -> None: + sh = _Shard(200064) + assert sh.rows > sh.shard, "the M3.1 shard is padded" + full, prev0 = _inputs(9, 2, 3, sh.vocab, False) + base = sh.base(full, RANK) + base[:, :, -1] = 0.0 + try: + chain.is_valid(base, prev0, sh.w1, sh.w2s, RANK * sh.shard, **WS) + except AssertionError: + pass + else: + raise AssertionError("the guard admitted a padding column that is not -inf") + dist.barrier() + + +CHECKS = ( + check_certified_cells, + check_deterministic_and_graph, + check_supports_refuses, + check_the_guard_refuses_a_live_padding_column, +) + + +def _run_one_rank(device: int) -> int: + global RANK, WORLD, WS, entry, chain, validating + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(device) + dist.init_process_group( + "cuda:nccl,cpu:gloo", init_method="env://", device_id=torch.device("cuda", device) + ) + RANK, WORLD = rank, world + + from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode import ( + dspark_markov_chain as entry_module, + ) + from tensorrt_llm.mapping import Mapping + + entry = entry_module + chain = entry.dspark_markov_chain + mapping = Mapping(world_size=world, rank=rank, gpus_per_node=GPUS_PER_NODE, tp_size=world) + WS = entry.workspace_args(mapping, dist.group.WORLD) + + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + from _validating import validating as _validating_cm + + validating = _validating_cm + + 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() + os._exit(13) + print(f"[rank {RANK}] {check.__name__} passed", flush=True) + dist.barrier() + torch.cuda.synchronize() + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict(os.environ, MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port), WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1") # fmt: skip + 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, + ) # fmt: skip + 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}") + + +def _rank_job(): + """`_rank_job.py` beside this file, by path (a sibling directory ships its own module of that name).""" + path = Path(__file__).resolve().with_name("_rank_job.py") + spec = importlib.util.spec_from_file_location("_spec_decode_rank_job", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +if __name__ == "__main__": + if "--srun" in sys.argv: + sys.exit(_run_one_rank(device=_rank_job().adopt_srun_env())) + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank(device=int(os.environ["RANK"]))) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/spec_decode/_rank_job.py b/tests/unittest/_torch/modeling_v2/spec_decode/_rank_job.py new file mode 100644 index 000000000000..47dc85fb582e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/spec_decode/_rank_job.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Use the existing four-rank launcher wrapper for spec-decode matrices.""" + +import importlib.util +from pathlib import Path + + +def run(entry: str, world_size: int) -> None: + if world_size != 4: + raise ValueError("This collected launcher supports WORLD4 only") + directory = Path(__file__).resolve().parent + path = directory.parent / "comm" / "_rank_job.py" + spec = importlib.util.spec_from_file_location("_spec_decode_comm_rank_job", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + module.run(entry, launcher_dir=directory) diff --git a/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_block_tables.py b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_block_tables.py new file mode 100644 index 000000000000..c521d8f78497 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_block_tables.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""catalog/spec_decode/dspark_block_tables: the one-launch draft block-table refresh writes, bit for bit, what +DFlashWorker._refresh_ctx_block_tables writes (decoded tables and placeholder counts), eager and in a CUDA graph, with +the per-row kernel and with S10c-B2's (row, chunk) grid.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_block_tables import ( + enable_parallel_refresh, + refresh_block_tables, +) +from tensorrt_llm._torch.speculative.dflash import DFlashWorker + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU (Triton)") + + +@pytest.fixture(params=[False, True], ids=["per_row", "grid2d"]) +def parallel(request): + enable_parallel_refresh(request.param) + yield request.param + enable_parallel_refresh(False) + + +def _worker(rows, width, pool_idx, divisor, fused): + worker = object.__new__(DFlashWorker) + worker._ctx_block_tables = torch.full((rows, width), 77, dtype=torch.int32, device="cuda") + worker._ctx_block_counts = torch.full((rows,), 77, dtype=torch.int64, device="cuda") + worker._ctx_pool_idx = pool_idx + worker._ctx_block_divisor = divisor + worker._fused_block_refresh = refresh_block_tables if fused else None + return worker + + +def _offsets(pools, rows, width, dtype, seed, high=1 << 20): + g = torch.Generator(device="cuda").manual_seed(seed) + off = torch.randint(0, high, (pools, rows, 2, width), generator=g, device="cuda").to(dtype) + # placeholders past each row's allocation (negative), at a random length per row + lens = torch.randint(0, width + 1, (pools, rows, 1, 1), generator=g, device="cuda") + cols = torch.arange(width, device="cuda").view(1, 1, 1, -1) + return torch.where(cols < lens, off, torch.full_like(off, -1)) + + +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("width", [7, 129, 1100, 4500]) +@pytest.mark.parametrize("divisor", [1, 2, 10]) +@pytest.mark.parametrize("pools,pool_idx", [(1, 0), (3, 2)]) +@pytest.mark.parametrize("high", [1 << 20, (1 << 31) - 1]) +def test_refresh_is_the_stock_refresh_bit_for_bit( + parallel, dtype, width, divisor, pools, pool_idx, high +): + rows = 9 + for seed in range(3): + src = _offsets(pools, rows, width, dtype, seed, high) + md = SimpleNamespace(draft_kv_cache_block_offsets=src) + for num_seqs in (1, 5, rows): + ref = _worker(rows, width, pool_idx, divisor, fused=False) + got = _worker(rows, width, pool_idx, divisor, fused=True) + assert DFlashWorker._refresh_ctx_block_tables(ref, md, num_seqs) + assert DFlashWorker._refresh_ctx_block_tables(got, md, num_seqs) + torch.cuda.synchronize() + assert torch.equal(ref._ctx_block_tables, got._ctx_block_tables), (seed, num_seqs) + assert torch.equal(ref._ctx_block_counts, got._ctx_block_counts), (seed, num_seqs) + assert torch.equal( + src, _offsets(pools, rows, width, dtype, seed, high) + ) # the metadata is not edited + + +@pytest.mark.parametrize("width", [64, 3000]) +def test_refresh_in_a_cuda_graph(parallel, width): + rows, num_seqs = 8, 6 + src = _offsets(2, rows, width, torch.int32, 0) + md = SimpleNamespace(draft_kv_cache_block_offsets=src) + got = _worker(rows, width, 1, 2, fused=True) + DFlashWorker._refresh_ctx_block_tables(got, md, num_seqs) # compile outside capture + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + g = torch.cuda.CUDAGraph() + with torch.cuda.stream(s): + with torch.cuda.graph(g, stream=s): + DFlashWorker._refresh_ctx_block_tables(got, md, num_seqs) + torch.cuda.current_stream().wait_stream(s) + src.copy_(_offsets(2, rows, width, torch.int32, 1)) # a later step's offsets, same buffer + got._ctx_block_tables.fill_(77) + g.replay() + ref = _worker(rows, width, 1, 2, fused=False) + DFlashWorker._refresh_ctx_block_tables(ref, md, num_seqs) + torch.cuda.synchronize() + assert torch.equal(ref._ctx_block_tables, got._ctx_block_tables) + assert torch.equal(ref._ctx_block_counts, got._ctx_block_counts) + + +def test_unsupported_layouts_keep_the_stock_path(): + src = _offsets(1, 4, 16, torch.int32, 0) + tables = torch.zeros(4, 16, dtype=torch.int32, device="cuda") + counts = torch.zeros(4, dtype=torch.int64, device="cuda") + assert refresh_block_tables(src, 0, 4, 2, tables, counts) + assert not refresh_block_tables(src, 0, 4, 2, tables[:, :8], counts) # width mismatch + assert not refresh_block_tables(src, 0, 4, 2, tables, counts.int()) # counts not int64 + assert not refresh_block_tables(src.float(), 0, 4, 2, tables, counts) # not integer offsets + assert not refresh_block_tables(src, 0, 4, 0, tables, counts) # no divisor diff --git a/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_ctx_advance.py b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_ctx_advance.py new file mode 100644 index 000000000000..15763bd310dc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_ctx_advance.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""catalog/spec_decode/dspark_ctx_advance: the one-launch context advance + read view leaves the context lengths and +the read block table exactly as DFlashWorker._advance_ctx_len and the generation branch's read-view lines +(dflash_allocated_ctx_limit, dflash_rebased_ctx_view) do, and returns the same advertised lengths; eager and in a CUDA +graph.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_ctx_advance import ( + advance_ctx_and_read_view, +) +from tensorrt_llm._torch.speculative.dflash import ( + DFlashWorker, + dflash_allocated_ctx_limit, + dflash_rebased_ctx_view, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU (Triton)") + +PAGE, BLOCK = 64, 3 + + +def _state(n, width, max_ctx, seed, dummy_rows): + g = torch.Generator(device="cuda").manual_seed(seed) + num_slots = n + 4 + dummy = num_slots - 1 + slots = torch.randperm(num_slots - 1, generator=g, device="cuda")[:n] + for r in range(dummy_rows): + slots[-1 - r] = dummy + ctx_len = torch.randint(0, max_ctx + 1, (num_slots,), generator=g, device="cuda").long() + ctx_len[slots[0]] = max_ctx - 1 # runs into the max_ctx clamp + ctx_len[dummy] = 0 + nacc = torch.randint(1, 5, (n,), generator=g, device="cuda").long() + rows_total = n + 2 + counts = torch.randint(0, width + 1, (rows_total,), generator=g, device="cuda").long() + base = torch.randint(0, 6, (num_slots,), generator=g, device="cuda").long() + base[slots[1 % n]] = (1 << 20) // PAGE # a context shorter than its base + tables = torch.randint(0, 1 << 16, (rows_total, width), generator=g, device="cuda").int() + read = torch.full_like(tables, -9) + rows = torch.arange(1, n + 1, device="cuda") + return dict( + slots=slots, + rows=rows, + nacc=nacc, + ctx_len=ctx_len, + max_ctx=max_ctx, + dummy=dummy, + counts=counts, + base=base, + tables=tables, + read=read, + ) + + +def _stock(st): + """_advance_ctx_len + the generation branch's read-view lines, verbatim.""" + w = object.__new__(DFlashWorker) + w._ctx_len, w._max_ctx, w._dummy_slot = st["ctx_len"], st["max_ctx"], st["dummy"] + DFlashWorker._advance_ctx_len(w, st["slots"], st["nacc"]) + num_ctx = w._ctx_len[st["slots"]] + allocated = dflash_allocated_ctx_limit(st["counts"], PAGE, BLOCK) + num_ctx = torch.minimum(num_ctx, allocated[st["rows"]]) + base_pages = st["base"][st["slots"]] + st["read"][st["rows"]], num_ctx = dflash_rebased_ctx_view( + st["tables"][st["rows"]], base_pages, num_ctx, PAGE + ) + return num_ctx + + +def _fused(st): + return advance_ctx_and_read_view( + st["slots"], + st["rows"], + st["nacc"], + st["ctx_len"], + st["max_ctx"], + st["dummy"], + st["counts"], + st["base"], + st["tables"], + st["read"], + PAGE, + BLOCK, + ) + + +def _copy(st): + return {k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in st.items()} + + +@pytest.mark.parametrize("width", [5, 33, 2112]) +@pytest.mark.parametrize("n,dummy_rows", [(1, 0), (4, 1), (8, 3), (16, 0)]) +def test_advance_is_the_stock_advance_bit_for_bit(width, n, dummy_rows): + for seed in range(4): + st = _state(n, width, 135168, seed, dummy_rows) + ref, got = _copy(st), _copy(st) + ref_ctx = _stock(ref) + got_ctx = _fused(got) + torch.cuda.synchronize() + assert got_ctx is not None and got_ctx.dtype == ref_ctx.dtype + assert torch.equal(ref_ctx, got_ctx), seed + assert torch.equal(ref["ctx_len"], got["ctx_len"]), seed + assert torch.equal(ref["read"], got["read"]), seed + + +def test_advance_in_a_cuda_graph(): + st = _state(8, 64, 4096, 0, 2) + got = _copy(st) + _fused(got) # compile outside capture + got = _copy(st) + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + g = torch.cuda.CUDAGraph() + with torch.cuda.stream(s): + with torch.cuda.graph(g, stream=s): + out = _fused(got) + torch.cuda.current_stream().wait_stream(s) + for k in ( + "ctx_len", + "read", + ): # the capture ran nothing: reset to the step's inputs, then replay + got[k].copy_(st[k]) + g.replay() + ref = _copy(st) + ref_ctx = _stock(ref) + torch.cuda.synchronize() + assert torch.equal(ref_ctx, out) + assert torch.equal(ref["ctx_len"], got["ctx_len"]) and torch.equal(ref["read"], got["read"]) + + +def test_unsupported_layouts_keep_the_stock_path(): + st = _state(2, 8, 4096, 0, 0) + assert _fused(_copy(st)) is not None + bad = _copy(st) + bad["read"] = bad["read"].long() + assert _fused(bad) is None + bad = _copy(st) + bad["ctx_len"] = bad["ctx_len"].int() + assert _fused(bad) is None diff --git a/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_ctx_store.py b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_ctx_store.py new file mode 100644 index 000000000000..d7a8ebddadc8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_ctx_store.py @@ -0,0 +1,266 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""catalog/spec_decode/dspark_ctx_store: the one-launch context store writes, bit for bit, what the stock path writes. + +Stock = DFlashForCausalLM.precompute_context_kv, the accept mask and DFlashWorker._store_context_kv_paged (flashinfer's +paged append per layer), run on a tiny real drafter (2 layers, 2 kv heads of 128, k_norm), with full RoPE and with +the MiniMax-M3.1 drafter's partial RoPE (rotary_dim 64 of head_dim 128: dims 64..127 are not rotated). Both write into +caches filled with a sentinel, compared in an int16 view, for the private arena (one stacked tensor) and the bound pool +(a list of per-layer tensors), with masked rows, clamped columns and a placeholder page. Cells written more than once +(see _single_writer) are excluded: stock's own value there is a race. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_ctx_store import ( + ctx_store_supported, + fold_k_norm_weight, + fused_ctx_store, +) +from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM +from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + +def _has_trtllm_gen(): + if not torch.cuda.is_available(): + return False + from tensorrt_llm._utils import is_sm_100f + + return is_sm_100f() + + +pytestmark = pytest.mark.skipif( + not _has_trtllm_gen(), reason="the paged draft cache path needs an SM100-family GPU" +) + +HIDDEN, LAYERS, HEADS, KV_HEADS, HEAD_DIM, INTER, VOCAB, PAGE = 64, 2, 4, 2, 128, 128, 512, 32 + + +def _drafter(partial_rotary_factor=1.0): + from transformers import Qwen3Config + + from tensorrt_llm._torch.model_config import ModelConfig + + cfg = Qwen3Config.from_dict( + dict( + architectures=["DFlashDraftModel"], + model_type="qwen3", + block_size=4, + hidden_size=HIDDEN, + num_hidden_layers=LAYERS, + num_attention_heads=HEADS, + num_key_value_heads=KV_HEADS, + head_dim=HEAD_DIM, + intermediate_size=INTER, + hidden_act="silu", + rms_norm_eps=1e-6, + vocab_size=VOCAB, + max_position_embeddings=8192, + rope_theta=1e7, + rope_scaling=None, + partial_rotary_factor=partial_rotary_factor, + attention_bias=False, + torch_dtype="bfloat16", + num_target_layers=4, + tie_word_embeddings=False, + dflash_config={"target_layer_ids": [0, 1], "mask_token_id": VOCAB - 2}, + ) + ) + model = DFlashForCausalLM( + ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM"), dflash_attention_backend="TRTLLM" + ).to("cuda") + g = torch.Generator().manual_seed(3) + + def rnd(*shape): + return (torch.randn(*shape, generator=g) * 0.05).to(torch.bfloat16) + + w = { + "fc.weight": rnd(HIDDEN, HIDDEN * 2), + "hidden_norm.weight": rnd(HIDDEN) + 1.0, + "norm.weight": rnd(HIDDEN) + 1.0, + } + for i in range(LAYERS): + p = f"layers.{i}." + w[p + "self_attn.q_proj.weight"] = rnd(HEADS * HEAD_DIM, HIDDEN) + w[p + "self_attn.k_proj.weight"] = rnd(KV_HEADS * HEAD_DIM, HIDDEN) + w[p + "self_attn.v_proj.weight"] = rnd(KV_HEADS * HEAD_DIM, HIDDEN) + w[p + "self_attn.o_proj.weight"] = rnd(HIDDEN, HEADS * HEAD_DIM) + w[p + "self_attn.q_norm.weight"] = rnd(HEAD_DIM) + 1.0 + w[p + "self_attn.k_norm.weight"] = rnd(HEAD_DIM) * 4 + 1.0 + w[p + "input_layernorm.weight"] = rnd(HIDDEN) + 1.0 + w[p + "post_attention_layernorm.weight"] = rnd(HIDDEN) + 1.0 + w[p + "mlp.gate_proj.weight"] = rnd(INTER, HIDDEN) + w[p + "mlp.up_proj.weight"] = rnd(INTER, HIDDEN) + w[p + "mlp.down_proj.weight"] = rnd(HIDDEN, INTER) + model.load_weights(w) + model._build_fused_kv_buffers() + assert model._get_cos_sin_cache().shape[-1] == int(HEAD_DIM * partial_rotary_factor) + return model + + +def _case(n_req, kp1, seed, placeholder): + g = torch.Generator(device="cuda").manual_seed(seed) + rows_total, pages_per_row = n_req + 2, 6 + cap = pages_per_row * PAGE + table = ( + torch.randperm(rows_total * pages_per_row, generator=g, device="cuda") + .int() + .view(rows_total, -1) + ) + if placeholder: + table[:, pages_per_row // 2 :] = ( + 0 # past each request's allocation: page 0, which another request owns + ) + rows = torch.arange(1, n_req + 1, device="cuda") + ctx_len = torch.randint(0, cap, (n_req,), generator=g, device="cuda") + ctx_len[0] = cap - 2 # the last request runs into the clamp + off = torch.arange(kp1, device="cuda") + cols = torch.minimum(ctx_len[:, None] + off, torch.tensor(cap - 1, device="cuda")) + nacc = torch.randint(1, kp1 + 1, (n_req,), generator=g, device="cuda") + proj = (torch.randn(n_req * kp1, HIDDEN, generator=g, device="cuda") * 2).bfloat16() + return dict( + table=table, + rows=rows, + pos=(ctx_len[:, None] + off).reshape(-1), + cols=cols.reshape(-1), + nacc=nacc, + proj=proj, + kp1=kp1, + pages=rows_total * pages_per_row, + rows_total=rows_total, + ppr=pages_per_row, + ) + + +def _stock(model, c, kv_buf): + """The worker's own code after the projection: precompute_context_kv, the mask, _store_context_kv_paged.""" + k, v = model.precompute_context_kv(c["proj"], c["pos"]) + mask = (torch.arange(c["kp1"], device="cuda")[None] < c["nacc"][:, None]).reshape(-1) + mb = mask.view(-1, 1, 1, 1).to(k.dtype) + k.mul_(mb) + v.mul_(mb) + worker = object.__new__(DFlashWorker) + worker._lean_glue = False + worker._ctx_paged_append = None + worker._ctx_block_tables = c["table"] + worker._ctx_block_indptr = torch.arange( + 0, (c["rows_total"] + 1) * c["ppr"], c["ppr"], device="cuda", dtype=torch.int32 + ) + worker._ctx_page_size = PAGE + worker._ctx_kv_buf = kv_buf + worker._ctx_kv_last_page_len = torch.full( + (c["rows_total"],), PAGE, device="cuda", dtype=torch.int32 + ) + rows = c["rows"][:, None].expand(-1, c["kp1"]).reshape(-1) + DFlashWorker._store_context_kv_paged(worker, k, v, rows, c["cols"]) + + +def _caches(c, layout): + full = torch.full( + (LAYERS, c["pages"], 2, KV_HEADS, PAGE, HEAD_DIM), 7.0, device="cuda" + ).bfloat16() + return full if layout == "stacked" else [full[i].clone() for i in range(LAYERS)] + + +def _single_writer(c): + """Cells exactly one row writes. A clamped column (several rows of one request) or a placeholder page (several + requests) is written more than once; stock's winner there is a race between the append's thread blocks and no + reader reads it (dflash_allocated_ctx_limit), so those cells are left out of the comparison.""" + rows = c["rows"].repeat_interleave(c["kp1"]) + key = c["table"][rows, c["cols"] // PAGE].long() * PAGE + c["cols"] % PAGE + uniq, counts = torch.unique(key, return_counts=True) + keep = torch.ones( + LAYERS, c["pages"], 2, KV_HEADS, PAGE, HEAD_DIM, dtype=torch.bool, device="cuda" + ) + for k in uniq[counts > 1].tolist(): + keep[:, k // PAGE, :, :, k % PAGE] = False + return keep + + +def _as_int16(kv_buf): + t = kv_buf if isinstance(kv_buf, torch.Tensor) else torch.stack(kv_buf) + return t.view(torch.int16) + + +@pytest.mark.parametrize("partial_rotary_factor", [1.0, 0.5]) +@pytest.mark.parametrize("layout", ["stacked", "list"]) +@pytest.mark.parametrize("n_req,kp1", [(1, 4), (3, 4), (2, 8)]) +@pytest.mark.parametrize("placeholder", [False, True]) +@pytest.mark.parametrize("columns", ["given", "from_counts", "from_max_col"]) +@pytest.mark.parametrize("fold", [False, True]) +def test_fused_store_is_the_stock_store_bit_for_bit( + layout, n_req, kp1, placeholder, partial_rotary_factor, columns, fold +): + """fold: TAILBK_SPEC F, the k_norm weight multiply inside the kernel (the drafter's weight is bf16, so it is + taken); the stores must equal stock's either way.""" + model = _drafter(partial_rotary_factor) + if fold: + assert ( + model._k_norm_stacked.dtype == torch.bfloat16 and model._k_norm_stacked.is_contiguous() + ) + fold_k_norm_weight(fold) + try: + _check_fused_store(model, layout, n_req, kp1, placeholder, columns) + finally: + fold_k_norm_weight(False) + + +def _check_fused_store(model, layout, n_req, kp1, placeholder, columns): + for seed in range(4): + c = _case(n_req, kp1, seed, placeholder) + ref, got = _caches(c, layout), _caches(c, layout) + assert ctx_store_supported(model, got) + _stock(model, c, ref) + # The case's columns are the positions clamped to the rows' pages (every row holds ppr pages), so the + # kernel's own derivation (S10c-A2: from the page counts, or from the arena's last column) must match them. + kwargs = {} + if columns == "from_counts": + kwargs["capacity_counts"] = torch.full( + (c["rows_total"],), c["ppr"], dtype=torch.int64, device="cuda" + ) + elif columns == "from_max_col": + kwargs["max_col"] = c["ppr"] * PAGE - 1 + cols = c["cols"] if columns == "given" else None + assert fused_ctx_store( + model, + c["proj"], + c["pos"], + c["nacc"], + c["rows"], + cols, + c["table"], + got, + c["kp1"], + **kwargs, + ) + torch.cuda.synchronize() + keep = _single_writer(c) + assert torch.equal(_as_int16(ref)[keep], _as_int16(got)[keep]), f"seed {seed}" + + +def test_unsupported_drafters_keep_the_stock_path(): + model = _drafter() + kv = torch.zeros(LAYERS, 4, 2, KV_HEADS, PAGE, HEAD_DIM, device="cuda", dtype=torch.bfloat16) + assert ctx_store_supported(model, kv) + assert not ctx_store_supported(model, kv.float()) # not a bf16 cache + assert not ctx_store_supported(model, kv[:, :, :1]) # not a [K | V] HND page + no_norm = SimpleNamespace( + **{k: getattr(model, k) for k in ("_k_norm_eps", "_is_neox", "_num_kv_heads", "_head_dim")} + ) + no_norm._k_norm_stacked = None + assert not ctx_store_supported(no_norm, kv) + assert not fused_ctx_store(no_norm, None, None, None, None, None, None, kv, 4) + # Derived columns need exactly one bound: the page counts (int64) or the arena's last column. + assert not fused_ctx_store(model, None, None, None, None, None, None, kv, 4) + # A rotary dim the kernel cannot split into power-of-two halves (48 of 128) keeps the stock path. + odd_rotary = SimpleNamespace( + **{ + k: getattr(model, k) + for k in ("_k_norm_eps", "_is_neox", "_num_kv_heads", "_head_dim", "_k_norm_stacked") + } + ) + odd_rotary._get_cos_sin_cache = lambda: torch.zeros(16, 48, device="cuda") + assert not ctx_store_supported(odd_rotary, kv) diff --git a/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_drafter_inputs.py b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_drafter_inputs.py new file mode 100644 index 000000000000..14e64c41d22c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_drafter_inputs.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""catalog/spec_decode/dspark_drafter_inputs through the M3.1 targets' S10c-C entry: the drafter's step inputs (bonus, +query / ctx positions, noise block) are bit for bit what the stock lines of DFlashWorker.prepare_1st_drafter_inputs +compute, with the target's own embedding (M3EmbedTokens), eager and in a CUDA graph.""" + +import importlib +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dflash import dflash_noise_block_embedding + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU (Triton)") + +_ROOT = "tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3" +VOCAB, HIDDEN, MASK_ID, SLOTS = 1000, 6144, 997, 20 + + +def _target(tp): + return importlib.import_module(f"{_ROOT}.minimax_m31_nvfp4__sm_103__{tp}.dspark") + + +def _stock(embed, accepted, nacc, slots, ctx_len, max_ctx, trained, block, kp1): + """The stock lines of prepare_1st_drafter_inputs' generation branch, verbatim.""" + bonus_idx = (nacc - 1).clamp_min(0).long().unsqueeze(1) + bonus = accepted.gather(1, bonus_idx).squeeze(1).long() + ctx_len_gen = ctx_len[slots] + j_block = torch.arange(0, block, dtype=torch.long, device="cuda") + offsets_kp1 = torch.arange(0, kp1, dtype=torch.long, device="cuda") + ctx_len_now = (ctx_len_gen + nacc.long()).clamp_(max=max_ctx) + query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) + ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) + noise = dflash_noise_block_embedding(embed, bonus, MASK_ID, block, trained) + return bonus, query_position_ids, ctx_position_ids, noise + + +def _case(n, k, acc_dtype, seed, max_ctx): + g = torch.Generator(device="cuda").manual_seed(seed) + kp1 = k + 1 + accepted = torch.randint(0, VOCAB, (n, kp1), generator=g, device="cuda").to(acc_dtype) + accepted[0, :] = -1 # a padding request: its bonus id is outside the table + nacc = torch.randint(1, kp1 + 1, (n,), generator=g, device="cuda").int() + slots = torch.randperm(SLOTS, generator=g, device="cuda")[:n] + ctx_len = torch.randint(0, max_ctx + 1, (SLOTS,), generator=g, device="cuda").long() + ctx_len[slots[-1]] = max_ctx - 1 # runs into the clamp + return accepted, nacc, slots, ctx_len + + +def _embed(tp): + g = torch.Generator(device="cuda").manual_seed(7) + table = torch.randn(VOCAB, HIDDEN, generator=g, device="cuda").bfloat16() + return _target(tp).M3EmbedTokens(SimpleNamespace(w={"embed": table})) + + +def _equal(a, b): + return all( + x.dtype == y.dtype and x.shape == y.shape and torch.equal(x, y) for x, y in zip(a, b) + ) + + +@pytest.mark.parametrize("tp", ["tp4", "tp8"]) +@pytest.mark.parametrize("k", [3, 7]) +@pytest.mark.parametrize("acc_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("trained", [False, True]) +def test_inputs_are_the_stock_inputs_bit_for_bit(tp, k, acc_dtype, trained): + embed = _embed(tp) + mask_row = (torch.randn(HIDDEN, device="cuda") * 3).bfloat16() if trained else None + for n in (1, 3, 8, 16): + for seed in range(3): + accepted, nacc, slots, ctx_len = _case(n, k, acc_dtype, seed, max_ctx=4096) + args = (embed, accepted, nacc, slots, ctx_len, 4096) + ref = _stock(*args, mask_row, k, k + 1) + got = _target(tp).s10c_c_drafter_inputs(*args, MASK_ID, mask_row, k, k + 1) + assert got is not None + torch.cuda.synchronize() + assert _equal(ref, got), (n, seed) + + +def test_inputs_in_a_cuda_graph(): + embed = _embed("tp8") + accepted, nacc, slots, ctx_len = _case(8, 3, torch.int32, 0, max_ctx=4096) + fn = _target("tp8").s10c_c_drafter_inputs + fn(embed, accepted, nacc, slots, ctx_len, 4096, MASK_ID, None, 3, 4) # compile outside capture + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + g = torch.cuda.CUDAGraph() + with torch.cuda.stream(s): + with torch.cuda.graph(g, stream=s): + got = fn(embed, accepted, nacc, slots, ctx_len, 4096, MASK_ID, None, 3, 4) + torch.cuda.current_stream().wait_stream(s) + a2, n2, s2, c2 = _case(8, 3, torch.int32, 1, max_ctx=4096) + accepted.copy_(a2), nacc.copy_(n2), slots.copy_(s2), ctx_len.copy_(c2) + g.replay() + ref = _stock(embed, accepted, nacc, slots, ctx_len, 4096, None, 3, 4) + torch.cuda.synchronize() + assert _equal(ref, got) + + +def test_other_embeddings_keep_the_stock_path(): + accepted, nacc, slots, ctx_len = _case(2, 3, torch.int32, 0, max_ctx=4096) + other = torch.nn.Embedding(VOCAB, HIDDEN).cuda().bfloat16() + fn = _target("tp8").s10c_c_drafter_inputs + assert fn(other, accepted, nacc, slots, ctx_len, 4096, MASK_ID, None, 3, 4) is None diff --git a/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_forward_indices.py b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_forward_indices.py new file mode 100644 index 000000000000..5b95d126b36e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_forward_indices.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""catalog/spec_decode/dspark_forward_indices: the paged dflash_forward's index prelude in one launch is bit for bit +the stock lines of DFlashForCausalLM.dflash_forward (cache_seqlens, cache_batch_idx, seq_lens_after, kv_indptr, +kv_last_page_len, append_positions, query_positions_flat), eager and in a CUDA graph.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_forward_indices import ( + forward_indices, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU (Triton)") + + +def _stock(num_ctx, batch_idx, query_positions, block_size, pages_per_slot, page_size): + """The stock lines of dflash_forward's paged branch, verbatim.""" + B = query_positions.shape[0] + cache_seqlens_i32 = num_ctx[:B].to(torch.int32) + cache_batch_idx_i32 = batch_idx.to(torch.int32) + seq_lens_after = cache_seqlens_i32 + block_size + kv_indptr = torch.arange( + 0, (B + 1) * pages_per_slot, pages_per_slot, dtype=torch.int32, device="cuda" + ) + kv_last_page_len = ((seq_lens_after - 1) % page_size) + 1 + block_offsets = torch.arange(block_size, dtype=torch.int32, device="cuda") + append_positions = (cache_seqlens_i32.view(-1, 1) + block_offsets).reshape(-1).contiguous() + query_positions_flat_i32 = query_positions.reshape(-1).to(torch.int32) + return ( + cache_seqlens_i32, + cache_batch_idx_i32, + seq_lens_after, + kv_indptr, + kv_last_page_len, + append_positions, + query_positions_flat_i32, + ) + + +def _case(B, block, dtype, seed, strided): + g = torch.Generator(device="cuda").manual_seed(seed) + num_ctx = torch.randint(0, 1 << 20, (B,), generator=g, device="cuda").to(dtype) + num_ctx[0] = 0 # an empty context: seq_lens_after = block + batch_idx = torch.randperm(B + 5, generator=g, device="cuda")[:B].to(dtype) + if strided: + q = torch.randint(0, 1 << 21, (block, B), generator=g, device="cuda").to(dtype).t() + else: + q = torch.randint(0, 1 << 21, (B, block), generator=g, device="cuda").to(dtype) + return num_ctx, batch_idx, q + + +def _equal(a, b): + return all( + x.dtype == y.dtype and x.shape == y.shape and torch.equal(x, y) for x, y in zip(a, b) + ) + + +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("block", [3, 4, 7, 8]) +@pytest.mark.parametrize("page_size,pages_per_slot", [(64, 66), (32, 2049), (16, 5)]) +@pytest.mark.parametrize("strided", [False, True]) +def test_indices_are_the_stock_lines_bit_for_bit(dtype, block, page_size, pages_per_slot, strided): + for B in (1, 2, 3, 8, 16, 33): + for seed in range(3): + num_ctx, batch_idx, q = _case(B, block, dtype, seed, strided) + ref = _stock(num_ctx, batch_idx, q, block, pages_per_slot, page_size) + got = forward_indices(num_ctx, batch_idx, q, block, pages_per_slot, page_size) + assert got is not None + torch.cuda.synchronize() + assert _equal(ref, got), (B, seed) + + +def test_indices_in_a_cuda_graph(): + num_ctx, batch_idx, q = _case(8, 4, torch.int64, 0, False) + forward_indices(num_ctx, batch_idx, q, 4, 66, 64) # compile outside capture + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + g = torch.cuda.CUDAGraph() + with torch.cuda.stream(s): + with torch.cuda.graph(g, stream=s): + got = forward_indices(num_ctx, batch_idx, q, 4, 66, 64) + torch.cuda.current_stream().wait_stream(s) + n2, b2, q2 = _case(8, 4, torch.int64, 1, False) + num_ctx.copy_(n2), batch_idx.copy_(b2), q.copy_(q2) + g.replay() + ref = _stock(num_ctx, batch_idx, q, 4, 66, 64) + torch.cuda.synchronize() + assert _equal(ref, got) + + +def test_unsupported_layouts_keep_the_stock_path(): + num_ctx, batch_idx, q = _case(2, 4, torch.int64, 0, False) + assert forward_indices(num_ctx, batch_idx, q, 4, 66, 64) is not None + assert forward_indices(num_ctx.float(), batch_idx, q, 4, 66, 64) is None + assert forward_indices(num_ctx, batch_idx, q[:, :3], 4, 66, 64) is None + assert forward_indices(num_ctx, batch_idx[:1], q, 4, 66, 64) is None diff --git a/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_markov_bias.py b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_markov_bias.py new file mode 100644 index 000000000000..0263a0fe6590 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_markov_bias.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""`spec_decode/dspark_markov_bias` (TAILBK_SPEC S4L): its two launches around the bias GEMM equal, bit for bit, the +in-tree batched biased-logits lines (prevs cat / casts, markov_prev_embeddings, F.linear, fp32 widen + add), on the +M3.1 TP8 shard (25008 columns, rank 256), K = 3 / 7, batches 1..8, with out-of-vocab anchors (-1, V, a large id), +extreme logits, eagerly and replayed from a CUDA graph; ``supports`` declines what the kernels do not take.""" + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode import ( + dspark_markov_bias as mb, +) +from tensorrt_llm._torch.models.modeling_speculative import dspark_markov_step_bias + +assert torch.cuda.is_available(), "dspark_markov_bias requires a CUDA device" + +V = 200064 +R = 256 +N = 25008 + + +def _weights(seed=7): + g = torch.Generator(device="cuda").manual_seed(seed) + w1 = (torch.randn(V, R, device="cuda", generator=g) * 0.05).bfloat16() + w2 = (torch.randn(N, R, device="cuda", generator=g) * 0.05).bfloat16() + return w1, w2 + + +def _case(seed, B, K, kind="random"): + g = torch.Generator(device="cuda").manual_seed(seed) + logits = (torch.randn(B, K, N, device="cuda", generator=g) * 4).bfloat16() + if kind == "extremes": + pool = torch.tensor([3.0e38, -3.0e38, 2.0**-133, -(2.0**-133), 0.0, -0.0, float("-inf"), 1.0], + device="cuda") # fmt: skip + logits = pool[ + torch.randint(0, pool.numel(), logits.shape, device="cuda", generator=g) + ].bfloat16() + first_prev = torch.randint(0, V, (B,), device="cuda", generator=g) + tokens = torch.randint(0, V, (B, K), device="cuda", generator=g, dtype=torch.int32) + if kind == "oov": + first_prev[0] = -1 + if B > 1: + first_prev[1] = V + tokens[0, 0] = -1 + tokens[-1, K - 2] = V + 12345 + return logits, first_prev, tokens + + +def _reference(logits, first_prev, tokens, w1, w2): + """The in-tree batched form (tp targets' _biased_shard_logits, fp32 + batched).""" + shard = logits.float() + B, K, n = shard.shape + prevs = torch.cat([first_prev.view(B, 1), tokens[:, : K - 1].long()], dim=1) + bias = dspark_markov_step_bias(prevs.reshape(-1), w1, w2).view(B, K, n) + return shard + bias.to(shard.dtype) + + +def _fused(logits, first_prev, tokens, w1, w2): + B, K, n = logits.shape + bias = F.linear(mb.prev_embeddings(first_prev, tokens, w1), w2) + return mb.widen_add(logits, bias.view(B, K, n)) + + +def _bits_equal(a, b): + return ( + a.dtype == b.dtype + and a.shape == b.shape + and torch.equal(a.view(torch.int32), b.view(torch.int32)) + ) + + +@pytest.mark.parametrize("K", [3, 7]) +@pytest.mark.parametrize("B", [1, 2, 8]) +@pytest.mark.parametrize("kind", ["random", "oov", "extremes"]) +def test_bitwise_in_tree(B, K, kind): + w1, w2 = _weights() + for seed in range(3): + logits, first_prev, tokens = _case(100 * seed + B * 10 + K, B, K, kind) + assert mb.supports(logits, first_prev, tokens, w1) + got = _fused(logits, first_prev, tokens, w1, w2) + want = _reference(logits, first_prev, tokens, w1, w2) + assert _bits_equal(got, want), f"B {B} K {K} {kind} seed {seed}" + + +def test_prev_embeddings_rows(): + """Each row is w1[prev] bit for bit, or +0.0 for an anchor outside [0, V); int32 and int64 ids.""" + w1, _ = _weights() + B, K = 3, 7 + first_prev = torch.tensor([5, -1, V], device="cuda") + tokens = torch.tensor([[1, 2, 3, -7, V - 1, 0, 9]] * B, device="cuda", dtype=torch.int32) + for fp, tk in ((first_prev, tokens), (first_prev.int(), tokens.long())): + e = mb.prev_embeddings(fp, tk, w1) + for b in range(B): + for k in range(K): + prev = int(fp[b]) if k == 0 else int(tk[b, k - 1]) + row = e[b * K + k] + if 0 <= prev < V: + assert torch.equal(row.view(torch.int16), w1[prev].view(torch.int16)), (b, k) + else: + assert torch.equal( + row.view(torch.int16), torch.zeros_like(row).view(torch.int16) + ), (b, k) + + +def test_cuda_graph_replay(): + w1, w2 = _weights() + B, K = 1, 7 + logits, first_prev, tokens = _case(1, B, K) + _fused(logits, first_prev, tokens, w1, w2) # compile outside capture + torch.cuda.synchronize() + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + g = torch.cuda.CUDAGraph() + with torch.cuda.stream(s): + with torch.cuda.graph(g): + out = _fused(logits, first_prev, tokens, w1, w2) + torch.cuda.current_stream().wait_stream(s) + for rep in range(6): + l2, f2, t2 = _case(500 + rep, B, K, "oov" if rep % 2 else "random") + logits.copy_(l2) + first_prev.copy_(f2) + tokens.copy_(t2) + g.replay() + torch.cuda.synchronize() + assert _bits_equal(out, _reference(logits, first_prev, tokens, w1, w2)), f"replay {rep}" + + +def test_supports_declines(): + w1, _ = _weights() + logits, first_prev, tokens = _case(3, 2, 7) + assert mb.supports(logits, first_prev, tokens, w1) + assert not mb.supports(logits.float(), first_prev, tokens, w1) + assert not mb.supports(logits.transpose(0, 1), first_prev, tokens, w1) + assert not mb.supports(logits, first_prev.float(), tokens, w1) + assert not mb.supports(logits, first_prev, tokens[:, :3], w1) + assert not mb.supports(logits, first_prev[:1], tokens, w1) + assert not mb.supports(logits, first_prev, tokens, w1[:, :200]) + assert not mb.supports(logits, first_prev, tokens, w1.float()) diff --git a/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_markov_chain_ll_op_matrix.py b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_markov_chain_ll_op_matrix.py new file mode 100644 index 000000000000..d08d25d97357 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_markov_chain_ll_op_matrix.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for `spec_decode/dspark_markov_chain`'s LL exchange variant +(``_dspark_markov_chain_ll_op_matrix.py`` beside this file: LL tokens == the S4B8 batch-1 kernel's, bitwise, on real +ranks, including forced cross-rank ties, back-to-back launches and CUDA-graph replays). + +World 4 is collected here. World 8 spans two GB300 trays and runs as one srun task per rank (see the matrix's +docstring). +""" + +import importlib.util +from pathlib import Path + +import pytest +import torch + +assert torch.cuda.is_available(), "dspark_markov_chain requires CUDA devices" + + +def _rank_job(): + # By path: other test directories ship a module of the same name. + path = Path(__file__).resolve().with_name("_rank_job.py") + spec = importlib.util.spec_from_file_location("_spec_decode_rank_job", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +# Spawns 4 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_dspark_markov_chain_ll_op_matrix() -> None: + _rank_job().run("dspark_markov_chain_ll", 4) diff --git a/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_markov_chain_op_matrix.py b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_markov_chain_op_matrix.py new file mode 100644 index 000000000000..505570c17ef5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_markov_chain_op_matrix.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for `spec_decode/dspark_markov_chain`'s certification matrix +(``_dspark_markov_chain_op_matrix.py`` beside this file, which drives the entry's ``.CELLS``). + +World 4 is collected here. The world-8 cells span two GB300 trays and run as one srun task per rank (see the +matrix's docstring), like ``comm/_mnnvl_allreduce_op_matrix.py``. +""" + +import importlib.util +from pathlib import Path + +import pytest +import torch + +assert torch.cuda.is_available(), "dspark_markov_chain requires CUDA devices" + + +def _rank_job(): + # By path: other test directories ship a module of the same name, and `import _rank_job` would get whichever + # one a collected sibling imported first. + path = Path(__file__).resolve().with_name("_rank_job.py") + spec = importlib.util.spec_from_file_location("_spec_decode_rank_job", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +# Spawns 4 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_dspark_markov_chain_op_matrix() -> None: + _rank_job().run("dspark_markov_chain", 4) diff --git a/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_qk_norm_rope.py new file mode 100644 index 000000000000..28d9486fce85 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_dspark_qk_norm_rope.py @@ -0,0 +1,244 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""catalog/spec_decode/dspark_qk_norm_rope: the drafter block's q / k RMSNorm + partial NeoX RoPE in one launch is bit +for bit the stock lines of DFlashForCausalLM.dflash_forward's flashinfer-RoPE branch (q_norm, k_norm, then +flashinfer_apply_rope_with_cos_sin_cache_inplace), at the M3.1 TP8 and TP4 rank head counts, 16 seeds, rows 3..56, +real-scale activations with tiny / huge / all-zero heads, eager and in a CUDA graph. DRAFT_APPEND_FUSE +(``qk_norm_rope_append``): the same launch's K / V pages are bit for bit flashinfer.page.append_paged_kv_cache's +(HND), with blocks crossing page boundaries, slots re-written, int32 / int64 indices, eager and in a CUDA graph.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.dspark_qk_norm_rope import ( + append_supported, + qk_norm_rope, + qk_norm_rope_append, +) +from tensorrt_llm._torch.custom_ops import ( + flashinfer_apply_rope_with_cos_sin_cache_inplace as fi_rope, +) +from tensorrt_llm._torch.modules.rms_norm import RMSNorm + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU (Triton)") + +HD, ROT, EPS, MAX_POS = 128, 64, 1e-6, 1 << 20 + + +def _cache(): + inv = 1.0 / (1e7 ** (torch.arange(0, ROT, 2, device="cuda", dtype=torch.float32) / ROT)) + ang = torch.arange(MAX_POS, device="cuda", dtype=torch.float32)[:, None] * inv[None, :] + return torch.cat([ang.cos(), ang.sin()], dim=1).contiguous() + + +_CS = None + + +def _cs(): + global _CS + if _CS is None: + _CS = _cache() + return _CS + + +def _norms(seed): + g = torch.Generator(device="cuda").manual_seed(1000 + seed) + qn = RMSNorm(hidden_size=HD, eps=EPS, dtype=torch.bfloat16, device="cuda") + kn = RMSNorm(hidden_size=HD, eps=EPS, dtype=torch.bfloat16, device="cuda") + with torch.no_grad(): + qn.weight.copy_((torch.randn(HD, generator=g, device="cuda") * 0.3 + 1).bfloat16()) + kn.weight.copy_((torch.randn(HD, generator=g, device="cuda") * 0.3 + 1).bfloat16()) + return qn, kn + + +def _qkv(T, hq, hkv, seed): + """Real-scale qkv rows: per-head scales spanning 1e-3..3e2, some all-zero heads, a few outlier channels.""" + g = torch.Generator(device="cuda").manual_seed(seed) + width = (hq + 2 * hkv) * HD + x = torch.randn(T, hq + 2 * hkv, HD, generator=g, device="cuda") + scale = torch.exp( + torch.empty(T, hq + 2 * hkv, 1, device="cuda").uniform_(-7.0, 5.7, generator=g) + ) + x = x * scale + x[:, :, :4] *= 20.0 # outlier channels + if T > 1: + x[1, 0] = 0.0 # an all-zero q head + x[T - 1, hq] = 0.0 # an all-zero k head + pos = torch.randint(0, MAX_POS, (T,), generator=g, device="cuda", dtype=torch.int32) + return x.reshape(T, width).bfloat16(), pos + + +def _stock(qkv, hq, hkv, qn, kn, pos): + qs, ks = hq * HD, hkv * HD + q = qn(qkv[:, :qs].reshape(-1, HD)).view(-1, qs) + k = kn(qkv[:, qs : qs + ks].reshape(-1, HD)).view(-1, ks) + fi_rope(pos, q, k, HD, _cs(), True) + return q, k + + +@pytest.mark.parametrize("hq,hkv", [(8, 1), (16, 1)], ids=["tp8", "tp4"]) +def test_bitwise_to_the_stock_lines(hq, hkv): + worst = 0 + for seed in range(16): + qn, kn = _norms(seed) + for T in (3, 4, 7, 8, 12, 24, 28, 32, 56): + qkv, pos = _qkv(T, hq, hkv, 100 * seed + T) + rq, rk = _stock(qkv, hq, hkv, qn, kn, pos) + got = qk_norm_rope(qkv, hq * HD, hkv * HD, HD, qn, kn, pos, _cs(), True) + assert got is not None + torch.cuda.synchronize() + bad = int((got[0] != rq).sum() + (got[1] != rk).sum()) + worst = max(worst, bad) + assert bad == 0, (seed, T, bad) + assert worst == 0 + + +def test_in_a_cuda_graph(): + hq, hkv = 8, 1 + qn, kn = _norms(0) + qkv, pos = _qkv(32, hq, hkv, 1) + qk_norm_rope(qkv, hq * HD, hkv * HD, HD, qn, kn, pos, _cs(), True) # compile outside capture + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + g = torch.cuda.CUDAGraph() + with torch.cuda.stream(s): + with torch.cuda.graph(g, stream=s): + got = qk_norm_rope(qkv, hq * HD, hkv * HD, HD, qn, kn, pos, _cs(), True) + torch.cuda.current_stream().wait_stream(s) + q2, p2 = _qkv(32, hq, hkv, 2) + qkv.copy_(q2), pos.copy_(p2) + g.replay() + rq, rk = _stock(qkv, hq, hkv, qn, kn, pos) + torch.cuda.synchronize() + assert torch.equal(got[0], rq) and torch.equal(got[1], rk) + + +def test_unsupported_inputs_keep_the_stock_path(): + hq, hkv = 8, 1 + qn, kn = _norms(0) + qkv, pos = _qkv(4, hq, hkv, 3) + assert ( + qk_norm_rope(qkv, hq * HD, hkv * HD, HD, qn, kn, pos, _cs(), False) is None + ) # interleaved RoPE + assert qk_norm_rope(qkv.float(), hq * HD, hkv * HD, HD, qn, kn, pos, _cs(), True) is None + assert qk_norm_rope(qkv, hq * HD, hkv * HD, HD, qn, kn, pos.long(), _cs(), True) is None + gem = RMSNorm(hidden_size=HD, eps=EPS, dtype=torch.bfloat16, device="cuda", use_gemma=True) + assert qk_norm_rope(qkv, hq * HD, hkv * HD, HD, gem, kn, pos, _cs(), True) is None + + +PS = 32 # the draft pool's page size + + +def _paged(B, K, hkv, pps, ctx, seed, idx_dtype=torch.int32): + """A random-filled HND cache [pages, 2, hkv, PS, HD], per-request block tables, and the append rows of a K-token + block per request at ctx[b] .. ctx[b] + K - 1 (as dflash_forward builds them).""" + g = torch.Generator(device="cuda").manual_seed(seed) + pages = B * pps + 3 + cache = torch.randn(pages, 2, hkv, PS, HD, generator=g, device="cuda").bfloat16() + perm = torch.randperm(pages, generator=g, device="cuda")[: B * pps] + kv_indices = perm.to(idx_dtype).contiguous() + kv_indptr = torch.arange(0, (B + 1) * pps, pps, device="cuda", dtype=torch.int32) + bidx = torch.arange(B, device="cuda", dtype=idx_dtype).repeat_interleave(K).contiguous() + apos = (torch.tensor(ctx, device="cuda", dtype=torch.int64).view(-1, 1) + + torch.arange(K, device="cuda")).reshape(-1).to(idx_dtype).contiguous() # fmt: skip + return cache, bidx, apos, kv_indices, kv_indptr + + +def _stock_append(qkv, hq, hkv, qn, kn, pos, cache, bidx, apos, kv_indices, kv_indptr): + import flashinfer + + q, k = _stock(qkv, hq, hkv, qn, kn, pos) + qs, ks = hq * HD, hkv * HD + v = qkv[:, qs + ks : qs + 2 * ks].reshape(-1, hkv, HD) + seq = torch.bincount(bidx.long(), minlength=kv_indptr.numel() - 1) + last = ((apos.view(seq.numel(), -1)[:, -1].long() % PS) + 1).to(torch.int32) + i32 = lambda t: t.to(torch.int32).contiguous() # noqa: E731 -- flashinfer takes int32 index tensors + flashinfer.page.append_paged_kv_cache(k.reshape(-1, hkv, HD).contiguous(), v, i32(bidx), i32(apos), cache, + i32(kv_indices), kv_indptr, last, kv_layout="HND") # fmt: skip + return q, k + + +@pytest.mark.parametrize("idx_dtype", [torch.int32, torch.int64], ids=["i32", "i64"]) +@pytest.mark.parametrize("hq,hkv", [(8, 1), (16, 1)], ids=["tp8", "tp4"]) +def test_append_pages_bitwise_to_flashinfer(hq, hkv, idx_dtype): + K, pps = 7, 4 + for seed in range(6): + qn, kn = _norms(seed) + for B, ctx in ( + (1, [29]), + (1, [0]), + (2, [30, 61]), + (3, [5, 95, 31]), + ): # blocks crossing pages, page 0, ... + T = B * K + qkv, pos = _qkv(T, hq, hkv, 500 * seed + 10 * B + ctx[0]) + cache, bidx, apos, kvi, kvp = _paged(B, K, hkv, pps, ctx, seed, idx_dtype) + ref_cache = cache.clone() + rq, rk = _stock_append(qkv, hq, hkv, qn, kn, pos, ref_cache, bidx, apos, kvi, kvp) + assert append_supported(qkv, hkv * HD, cache, bidx, apos, kvi, kvp) + got = qk_norm_rope_append(qkv, hq * HD, hkv * HD, HD, qn, kn, pos, _cs(), True, cache, bidx, apos, kvi, + kvp) # fmt: skip + assert got is not None + torch.cuda.synchronize() + assert torch.equal(got[0], rq) and torch.equal(got[1], rk), (seed, B, ctx) + assert torch.equal(cache.view(torch.int16), ref_cache.view(torch.int16)), (seed, B, ctx) + # Slot reuse: the same rows written again with new values (a rejected draft's slots are re-appended). + qkv2, _ = _qkv(T, hq, hkv, 777 + seed) + _stock_append(qkv2, hq, hkv, qn, kn, pos, ref_cache, bidx, apos, kvi, kvp) + qk_norm_rope_append( + qkv2, hq * HD, hkv * HD, HD, qn, kn, pos, _cs(), True, cache, bidx, apos, kvi, kvp + ) + torch.cuda.synchronize() + assert torch.equal(cache.view(torch.int16), ref_cache.view(torch.int16)), ( + "reuse", + seed, + B, + ctx, + ) + + +def test_append_in_a_cuda_graph(): + hq, hkv, K, B, pps = 8, 1, 7, 1, 4 + qn, kn = _norms(0) + qkv, pos = _qkv(B * K, hq, hkv, 5) + cache, bidx, apos, kvi, kvp = _paged(B, K, hkv, pps, [29], 0) + qk_norm_rope_append( + qkv, hq * HD, hkv * HD, HD, qn, kn, pos, _cs(), True, cache, bidx, apos, kvi, kvp + ) + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + g = torch.cuda.CUDAGraph() + with torch.cuda.stream(s): + with torch.cuda.graph(g, stream=s): + got = qk_norm_rope_append(qkv, hq * HD, hkv * HD, HD, qn, kn, pos, _cs(), True, cache, bidx, apos, kvi, + kvp) # fmt: skip + torch.cuda.current_stream().wait_stream(s) + for rep, c in enumerate((29, 60, 3, 95)): + q2, p2 = _qkv(B * K, hq, hkv, 40 + rep) + qkv.copy_(q2), pos.copy_(p2) + _, _, a2, _, _ = _paged(B, K, hkv, pps, [c], 0) + apos.copy_(a2) + ref_cache = cache.clone() + rq, rk = _stock_append(qkv, hq, hkv, qn, kn, pos, ref_cache, bidx, apos, kvi, kvp) + g.replay() + torch.cuda.synchronize() + assert torch.equal(got[0], rq) and torch.equal(got[1], rk), rep + assert torch.equal(cache.view(torch.int16), ref_cache.view(torch.int16)), rep + + +def test_append_unsupported_keeps_the_stock_path(): + hq, hkv, K = 8, 1, 7 + qn, kn = _norms(0) + qkv, pos = _qkv(K, hq, hkv, 9) + cache, bidx, apos, kvi, kvp = _paged(1, K, hkv, 4, [10], 0) + args = (qkv, hq * HD, hkv * HD, HD, qn, kn, pos, _cs(), True) + assert qk_norm_rope_append(*args, cache.float(), bidx, apos, kvi, kvp) is None # fp32 cache + assert ( + qk_norm_rope_append(*args, cache.transpose(2, 3), bidx, apos, kvi, kvp) is None + ) # NHD / strided + assert qk_norm_rope_append(*args, cache, bidx[:-1], apos, kvi, kvp) is None # rows mismatch + assert qk_norm_rope_append(*args, cache, bidx.float(), apos, kvi, kvp) is None + before = cache.clone() + assert qk_norm_rope_append(qkv, hq * HD, hkv * HD, HD, qn, kn, pos, _cs(), False, cache, bidx, apos, kvi, + kvp) is None # interleaved RoPE: nothing written # fmt: skip + assert torch.equal(cache.view(torch.int16), before.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_spec_strict_accept.py b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_spec_strict_accept.py new file mode 100644 index 000000000000..b725150e36c4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/spec_decode/test_modeling_v2_spec_strict_accept.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""`spec_decode/strict_accept`: both entries bitwise to their in-tree torch path over every cell of their `.CELLS`. + +One GPU; each cell is also run on a second random draw and inside a CUDA graph (captured after a warm call, replayed +on new inputs written into the captured buffers). +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.spec_decode.strict_accept import ( + SpecKvRewind, + SpecNextNewTokens, + SpecStrictAccept, + spec_kv_rewind, + spec_next_new_tokens, + spec_strict_accept, +) + +assert torch.cuda.is_available(), "strict_accept requires a CUDA device" + +VOCAB = 200064 + + +def _strict_inputs(spec, gen): + n_req, draft_len = spec["n_req"], spec["draft_len"] + t_dtype = spec.get("target_dtype", torch.int32) + k1 = draft_len + 1 + target = torch.randint(0, VOCAB, (n_req * k1,), generator=gen, device="cuda").to(t_dtype) + draft = target.view(n_req, k1)[:, :draft_len].clone().to(spec.get("draft_dtype", torch.int32)) + # A random prefix of each request's drafts matches: break the draft at a random position (or never). + cut = torch.randint(0, draft_len + 1, (n_req,), generator=gen, device="cuda") + pos = torch.arange(draft_len, device="cuda").unsqueeze(0) + breaks = pos >= cut.unsqueeze(1) + draft = torch.where(breaks & (pos == cut.unsqueeze(1)), (draft + 1) % VOCAB, draft) + # Past the first mismatch the drafts are random again (a later match must not count). + noise = torch.randint(0, VOCAB, draft.shape, generator=gen, device="cuda").to(draft.dtype) + draft = torch.where(breaks & (pos > cut.unsqueeze(1)), noise, draft) + if spec.get("pad_drafts"): + draft[::2, draft_len // 2 :] = -1 + if spec.get("extremes"): + draft[: n_req // 2] = target.view(n_req, k1)[: n_req // 2, :draft_len].to(draft.dtype) + draft[n_req // 2 :] = ( + target.view(n_req, k1)[n_req // 2 :, :draft_len].to(draft.dtype) + 1 + ) % VOCAB + return target, draft + + +def _next_new_inputs(spec, gen): + n_req, draft_len = spec["n_req"], spec["draft_len"] + width = draft_len + 1 + (4 if spec.get("wide") else 0) + rows = n_req + 3 + accepted = torch.randint( + 0, VOCAB, (rows, width), generator=gen, device="cuda", dtype=torch.int32 + ) + num = torch.randint(1, draft_len + 2, (n_req,), generator=gen, device="cuda", dtype=torch.int32) + batch_idx = torch.arange(n_req, device="cuda", dtype=torch.int32) + if spec.get("shuffle"): + batch_idx = torch.randperm(rows, generator=gen, device="cuda")[:n_req].to(torch.int32) + draft_dtype = spec.get("draft_dtype", torch.int32) + if spec.get("wide"): + # A strided view: rows of draft_len + 5 elements, the first draft_len used. + wide_rows = torch.randint(0, VOCAB, (n_req, draft_len + 5), generator=gen, device="cuda") + next_draft = wide_rows.to(draft_dtype)[:, :draft_len] + else: + next_draft = torch.randint(0, VOCAB, (n_req, draft_len), generator=gen, device="cuda").to( + draft_dtype + ) + return accepted, num, batch_idx, next_draft + + +def _assert_same(out, ref): + for o, r in zip( + out if isinstance(out, (list, tuple)) else [out], + ref if isinstance(ref, (list, tuple)) else [ref], + ): + assert o.dtype == r.dtype, (o.dtype, r.dtype) + assert o.shape == r.shape, (o.shape, r.shape) + assert torch.equal(o, r), (o, r) + + +@pytest.mark.parametrize("cell", SpecStrictAccept.CELLS, ids=lambda c: c.why[:60]) +def test_strict_accept_cells(cell): + gen = torch.Generator(device="cuda").manual_seed(7) + for _ in range(2): + target, draft = _strict_inputs(dict(cell.spec), gen) + _assert_same(spec_strict_accept(target, draft), spec_strict_accept.reference(target, draft)) + # Graph replay on new values written into the captured inputs. + target, draft = _strict_inputs(dict(cell.spec), gen) + spec_strict_accept(target, draft) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = spec_strict_accept(target, draft) + new_target, new_draft = _strict_inputs(dict(cell.spec), gen) + target.copy_(new_target) + draft.copy_(new_draft) + graph.replay() + torch.cuda.synchronize() + _assert_same(out, spec_strict_accept.reference(target, draft)) + + +@pytest.mark.parametrize("cell", SpecNextNewTokens.CELLS, ids=lambda c: c.why[:60]) +def test_next_new_tokens_cells(cell): + gen = torch.Generator(device="cuda").manual_seed(11) + for _ in range(2): + args = _next_new_inputs(dict(cell.spec), gen) + spec_next_new_tokens.is_valid(*args) + _assert_same(spec_next_new_tokens(*args), spec_next_new_tokens.reference(*args)) + args = _next_new_inputs(dict(cell.spec), gen) + spec_next_new_tokens(*args) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = spec_next_new_tokens(*args) + new_args = _next_new_inputs(dict(cell.spec), gen) + for dst, src in zip(args, new_args): + dst.copy_(src) + graph.replay() + torch.cuda.synchronize() + _assert_same(out, spec_next_new_tokens.reference(*args)) + + +def _kv_rewind_inputs(spec, gen): + nc, n = spec["nc"], spec["n"] + rows = nc + n + 2 # rows past the generation block must stay untouched + num = torch.randint(1, 5, (rows,), generator=gen, device="cuda", dtype=torch.int32) + kv = torch.randint(0, 1 << 20, (rows,), generator=gen, device="cuda").to( + spec.get("kv_dtype", torch.int32) + ) + rewind = torch.empty(64, dtype=torch.int32, device="cuda") + return num, kv, rewind + + +@pytest.mark.parametrize("cell", SpecKvRewind.CELLS, ids=lambda c: c.why[:60]) +def test_kv_rewind_cells(cell): + spec = dict(cell.spec) + nc, n = spec["nc"], spec["n"] + gen = torch.Generator(device="cuda").manual_seed(13) + for _ in range(2): + num, kv, rewind = _kv_rewind_inputs(spec, gen) + kv_ref = kv.clone() + out = spec_kv_rewind(num, kv, rewind, nc, n) + _assert_same(out, spec_kv_rewind.reference(num, kv_ref, rewind, nc, n)) + _assert_same(kv, kv_ref) + # Graph replay on new values written into the captured inputs. + num, kv, rewind = _kv_rewind_inputs(spec, gen) + spec_kv_rewind(num, kv.clone(), rewind, nc, n) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = spec_kv_rewind(num, kv, rewind, nc, n) + new_num, new_kv, _ = _kv_rewind_inputs(spec, gen) + num.copy_(new_num) + kv.copy_(new_kv) + kv_ref = kv.clone() + graph.replay() + torch.cuda.synchronize() + _assert_same(out, spec_kv_rewind.reference(num, kv_ref, rewind, nc, n)) + _assert_same(kv, kv_ref) diff --git a/tests/unittest/_torch/modeling_v2/test_m31_decode_scope_source.py b/tests/unittest/_torch/modeling_v2/test_m31_decode_scope_source.py new file mode 100644 index 000000000000..fb5698d43384 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_m31_decode_scope_source.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU checks of the production fused-dispatch predicates without CUDA imports.""" + +import ast +import types +import unittest +from pathlib import Path + +ROOT = ( + Path(__file__).resolve().parents[4] + / "tensorrt_llm" + / "_torch" + / "_experimental" + / "modeling_v2" +) + + +def dispatch(tp, rows, requests, capacity=64, ndim=1, embeddings=False): + path = ROOT / "models" / "minimax_m3" / f"minimax_m31_nvfp4__sm_103__tp{tp}" / "modeling.py" + tree = ast.parse(path.read_text()) + helper = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "is_decode_step" + ) + limit = next( + node + for node in tree.body + if isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == "FUSED_DECODE_MAX_TOKENS" + for target in node.targets + ) + ) + cls = next( + node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "DecodeTarget" + ) + forward = next( + node for node in cls.body if isinstance(node, ast.FunctionDef) and node.name == "forward" + ) + predicate = next( + node.value + for node in ast.walk(forward) + if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "kernels" for target in node.targets) + ) + env = { + "DECODE_MAX_TOKENS": capacity, + "input_ids": types.SimpleNamespace(ndim=ndim, numel=lambda: rows), + "inputs_embeds": object() if embeddings else None, + "attn_metadata": types.SimpleNamespace(num_seqs=requests), + } + exec(compile(ast.Module(body=[limit, helper], type_ignores=[]), str(path), "exec"), env) + return eval(compile(ast.Expression(predicate), str(path), "eval"), env) + + +class TestM31DecodeScope(unittest.TestCase): + def test_tp8_single_request_boundary(self): + for capacity in (48, 64): + for rows, expected in ( + (1, True), + (7, True), + (8, True), + (9, False), + (32, False), + (64, False), + ): + with self.subTest(capacity=capacity, rows=rows): + self.assertEqual(dispatch(8, rows, 1, capacity), expected) + + def test_tp8_multiple_requests_decline_even_when_rows_fit(self): + self.assertFalse(dispatch(8, 8, 2)) # K3/c2 must not enter the c1 fused path. + self.assertFalse(dispatch(8, 8, 8)) # Non-speculative c8 also has eight rows. + + def test_tp4_boundary(self): + for capacity in (48, 64): + for rows, requests, expected in ( + (8, 1, True), + (16, 2, True), + (32, 4, True), + (33, 5, False), + (64, 8, False), + ): + with self.subTest(capacity=capacity, rows=rows): + self.assertEqual(dispatch(4, rows, requests, capacity), expected) + + def test_embeddings_and_nonflat_ids_use_generic_path(self): + for tp in (4, 8): + with self.subTest(tp=tp): + self.assertFalse(dispatch(tp, 8, 1, ndim=2)) + self.assertFalse(dispatch(tp, 8, 1, embeddings=True)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/modeling_v2/test_m31_tp4_routing_source.py b/tests/unittest/_torch/modeling_v2/test_m31_tp4_routing_source.py new file mode 100644 index 000000000000..6ac180b441f6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_m31_tp4_routing_source.py @@ -0,0 +1,325 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU source-contract checks; no CUDA/model import or end-to-end gate implied.""" + +import ast +import enum +import importlib.util +import json +import os +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[4] +PACKAGE = "tensorrt_llm._torch._experimental.modeling_v2" +SOURCE = ROOT / "tensorrt_llm/_torch/_experimental/modeling_v2" + + +class Quant(enum.Enum): + FP8 = "FP8" + NVFP4 = "NVFP4" + MXFP8 = "MXFP8" + + +class NoTarget(Exception): + pass + + +def load_source(path, **values): + """Evaluate exact definitions with only top-level dependency imports supplied by fixtures.""" + tree = ast.parse(path.read_text()) + tree.body = [ + n + for n in tree.body + if not isinstance(n, (ast.Import, ast.ImportFrom)) + or (isinstance(n, ast.ImportFrom) and n.module == "__future__") + ] + namespace = dict(TYPE_CHECKING=False, enum=enum, os=os, json=json, Path=Path, **values) + exec(compile(tree, str(path), "exec"), namespace) + return namespace + + +def route(): + ns = load_source( + SOURCE / "models/minimax_m3/routing.py", + QuantAlgo=Quant, + NoTarget=NoTarget, + __package__=PACKAGE + ".models.minimax_m3", + ) + # The HF quant parser is a framework prerequisite; these tests supply its result. + ns["_quant_algo"] = lambda config: Quant.MXFP8 + return ns["route"] + + +def inputs(): + config = types.SimpleNamespace( + text_config=dict( + num_hidden_layers=60, + hidden_size=6144, + num_local_experts=128, + intermediate_size=3072, + num_attention_heads=64, + num_key_value_heads=4, + rope_theta=10_000_000, + sparse_attention_config=dict(use_sparse_attention=True, sparse_attention_freq=[1] * 60), + ), + quantization_config=dict(moe_quant_algo="NVFP4", moe_quant_format="nvfp4-pack-quantized"), + ) + mapping = types.SimpleNamespace( + world_size=4, + tp_size=4, + pp_size=1, + cp_size=1, + moe_tp_size=4, + moe_ep_size=1, + enable_attention_dp=False, + ) + return config, mapping + + +class TestM31Tp4RoutingSource(unittest.TestCase): + def test_declines_outside_scope(self): + cases = [("chip", (10, 0)), ("chip", (10, 7))] + cases += [ + (name, value) + for name, value in ( + ("world_size", 8), + ("tp_size", 8), + ("tp_size", 2), + ("pp_size", 2), + ("cp_size", 2), + ("moe_tp_size", 1), + ("moe_ep_size", 2), + ("enable_attention_dp", True), + ) + ] + cases += [ + (name, None) + for name in ("m3_rope", "m3_layers", "wrong_experts", "wrong_shape", "bf16_kv") + ] + for name, value in cases: + with self.subTest(name=name, value=value): + config, mapping = inputs() + chip, kv = (10, 3), Quant.NVFP4 + if name == "chip": + chip = value + elif hasattr(mapping, name): + setattr(mapping, name, value) + elif name == "m3_rope": + config.text_config["rope_theta"] = 5_000_000 + elif name == "m3_layers": + config.text_config["sparse_attention_config"]["sparse_attention_freq"] = [ + 0 + ] * 3 + [1] * 57 + elif name == "wrong_experts": + config.quantization_config["moe_quant_algo"] = "FP8" + elif name == "wrong_shape": + config.text_config["num_key_value_heads"] = 8 + else: + kv = None + with self.assertRaises(NoTarget): + route()(config, mapping, chip, None, kv) + + def test_tp4_routes_only_to_tp4(self): + config, mapping = inputs() + name = PACKAGE + ".models.minimax_m3.minimax_m31_nvfp4__sm_103__tp4.modeling" + module = types.ModuleType(name) + sentinel = type("Target", (), {}) + module.ModelingV2MiniMaxM31Nvfp4Sm103Tp4 = sentinel + with patch.dict(sys.modules, {name: module}): + self.assertIs(route()(config, mapping, (10, 3), None, Quant.NVFP4), sentinel) + + def test_router_arity_and_mode(self): + for arch, arity in [ + ("GptOssForCausalLM", 4), + ("DeepseekV3ForCausalLM", 4), + ("MiniMaxM3SparseForConditionalGeneration", 5), + ]: + with self.subTest(arch=arch): + ns = load_source( + SOURCE / "_router_index.py", logger=types.SimpleNamespace(info=lambda *a: None) + ) + calls = [] + target = type("Target", (), {}) + + def call(*args): + calls.append(args) + return target + + ns["routing_module"] = lambda name: types.SimpleNamespace(route=call) + ns["_device_compute_capability"] = lambda: (10, 3) + ns["_kv_cache_quant_algo"] = lambda *args: Quant.NVFP4 + config = types.SimpleNamespace( + pretrained_config=types.SimpleNamespace(architectures=[arch]), + mapping=object(), + spec_config=None, + ) + self.assertIs( + ns["modeling_v2_resolve"](config, types.SimpleNamespace(modeling_v2="require")), + target, + ) + self.assertEqual(len(calls[0]), arity) + if arity == 5: + self.assertIs(ns["resolved_mode"](target), ns["ModelingV2Mode"].REQUIRE) + self.assertEqual( + set(ns["MODELING_V2_ROUTERS"]), + { + "GptOssForCausalLM", + "DeepseekV3ForCausalLM", + "MiniMaxM3SparseForConditionalGeneration", + }, + ) + + def test_parent_import_and_late_native_registration(self): + """Import actual parent with no native ops, then load ops and repeat registration.""" + calls = [] + native = types.SimpleNamespace() + installed = {} + torch = types.ModuleType("torch") + torch.ops = types.SimpleNamespace(trtllm=native) + torch._C = types.SimpleNamespace( + _dispatch_has_kernel_for_dispatch_key=lambda name, key: installed.get(name) == "Meta" + ) + + def register_fake(name): + def decorate(fn): + self.assertNotIn(name, installed) + installed[name] = fn + calls.append(name) + return fn + + return decorate + + torch.library = types.SimpleNamespace(register_fake=register_fake) + registry = types.ModuleType("torch._library.simple_registry") + registry.singleton = types.SimpleNamespace( + find=lambda name: types.SimpleNamespace( + fake_impl=types.SimpleNamespace(kernel=installed.get(name)) + ) + ) + family_name = PACKAGE + ".models.minimax_m3" + catalog_name = PACKAGE + ".catalog.fakes.m31_fake_ops" + + def load_file(name, path, package=False): + spec = importlib.util.spec_from_file_location( + name, path, submodule_search_locations=[str(path.parent)] if package else None + ) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + with patch.dict(sys.modules, {"torch": torch, "torch._library.simple_registry": registry}): + for fullname in (family_name, catalog_name): + parts = fullname.split(".") + for length in range(1, len(parts)): + name = ".".join(parts[:length]) + package = types.ModuleType(name) + package.__path__ = [str(ROOT.joinpath(*parts[:length]))] + sys.modules[name] = package + parent = load_file(family_name, SOURCE / "models/minimax_m3/__init__.py", package=True) + self.assertEqual(calls, []) + self.assertFalse(hasattr(parent, "fake_ops")) + # A route rejected after the parent import can still choose the built-in model. + config, mapping = inputs() + with self.assertRaises(NoTarget): + route()(config, mapping, (10, 7), None, Quant.NVFP4) + load_file(catalog_name, SOURCE / "catalog/fakes/m31_fake_ops.py") + fake_module = load_file( + family_name + ".fake_ops", SOURCE / "models/minimax_m3/fake_ops.py" + ) + self.assertEqual(len(fake_module.FAKES), 7) + names = list(fake_module.FAKES) + for name in names: + setattr(native, name, object()) + installed["trtllm::" + names[0]] = "Meta" + previous_fake = object() + installed["trtllm::" + names[1]] = previous_fake + fake_module.register() + self.assertEqual(len(calls), 5) + fake_module.register() + self.assertEqual(len(calls), 5) + self.assertIs(installed["trtllm::" + names[1]], previous_fake) + self.assertEqual(installed["trtllm::" + names[0]], "Meta") + + def test_dormant_optins_removed_and_baseline_flags_retained(self): + root = SOURCE / "models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4" + rejected = { + "LC4_ONEDIV", + "DENSE_SWIGLU_CPP", + "R6", + "LDS1", + "S5", + "S2", + "S9B", + "FUSED_CAPTURE", + } + for path in root.rglob("*.py"): + for node in ast.walk(ast.parse(path.read_text())): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name): + if node.func.id in {"m31_enabled", "m31_disabled"} and node.args: + self.assertNotIn(getattr(node.args[0], "value", None), rejected, str(path)) + for diagnostic in ("UB_NOSELECT", "UB_DUPSELECT", "_ub_noselect_table"): + self.assertNotIn(diagnostic, (root / "modeling.py").read_text()) + s4 = (root / "dspark_s4.py").read_text() + for diagnostic in ("DSPARK_GATE_DUMP", "_s4_dump", "_s4_gate_dump", "_s4_logits_gate_dump"): + self.assertNotIn(diagnostic, s4) + model = ast.parse((root / "modeling.py").read_text()) + assigns = { + n.targets[0].id: ast.unparse(n.value) + for n in model.body + if isinstance(n, ast.Assign) and isinstance(n.targets[0], ast.Name) + } + self.assertTrue( + any( + isinstance(n, ast.AnnAssign) + and isinstance(n.target, ast.Name) + and n.target.id == "SWQ_ROWS" + for n in model.body + ) + ) + self.assertIn("AR_AUTOTUNE_ENABLED", assigns) + lc = ast.parse((root / "ops/attention/m3_sparse_decode_lc.py").read_text()) + flags = { + n.targets[0].id: ast.unparse(n.value) + for n in lc.body + if isinstance(n, ast.Assign) and isinstance(n.targets[0], ast.Name) + } + self.assertEqual(flags["LC4_BTSMEM"], "not m31_disabled('LC4_BTSMEM')") + self.assertEqual(flags["LC4_SYNC3"], "not m31_disabled('LC4_SYNC3')") + native_calls = [ + n + for n in ast.walk(lc) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and n.func.id == "lc4" + ] + self.assertEqual(len(native_calls), 1) + self.assertIs(native_calls[0].args[-1].value, False) + draft = (root / "dspark.py").read_text() + for marker in ("MXDRAFT = not m31_disabled", 'm31_disabled("S10")', 'm31_disabled("S10C")'): + self.assertIn(marker, draft) + self.assertFalse((root / "ops/gemm/m3_dense_swiglu.py").exists()) + + def test_pruned_activation_keeps_default_parameter_order(self): + path = SOURCE / "models/minimax_m3/minimax_m31_nvfp4__sm_103__tp4/modeling.py" + tree = ast.parse(path.read_text()) + fn = next( + n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "_swiglu_act" + ) + calls = [] + + def fallback(*args): + calls.append(args) + return "output" + + namespace = {"m3_silu_and_mul": fallback, "torch": types.SimpleNamespace(Tensor=object)} + exec(compile(ast.Module(body=[fn], type_ignores=[]), str(path), "exec"), namespace) + core = types.SimpleNamespace(swiglu_limit=7, swiglu_alpha=1.7, swiglu_beta=1.0) + self.assertEqual(namespace["_swiglu_act"](core, "gate_up"), "output") + self.assertEqual(calls, [("gate_up", 7, 1.7, 1.0)]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/modeling_v2/test_m31_tp8_routing_source.py b/tests/unittest/_torch/modeling_v2/test_m31_tp8_routing_source.py new file mode 100644 index 000000000000..57ca48949126 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_m31_tp8_routing_source.py @@ -0,0 +1,253 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU source-contract checks; no CUDA/model import or end-to-end gate implied.""" + +import ast +import enum +import importlib.util +import json +import os +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[4] +PACKAGE = "tensorrt_llm._torch._experimental.modeling_v2" +SOURCE = ROOT / "tensorrt_llm/_torch/_experimental/modeling_v2" + + +class Quant(enum.Enum): + FP8 = "FP8" + NVFP4 = "NVFP4" + MXFP8 = "MXFP8" + + +class NoTarget(Exception): + pass + + +def load_source(path, **values): + """Evaluate exact definitions with only top-level dependency imports supplied by fixtures.""" + tree = ast.parse(path.read_text()) + tree.body = [ + n + for n in tree.body + if not isinstance(n, (ast.Import, ast.ImportFrom)) + or (isinstance(n, ast.ImportFrom) and n.module == "__future__") + ] + namespace = dict(TYPE_CHECKING=False, enum=enum, os=os, json=json, Path=Path, **values) + exec(compile(tree, str(path), "exec"), namespace) + return namespace + + +def route(): + ns = load_source( + SOURCE / "models/minimax_m3/routing.py", + QuantAlgo=Quant, + NoTarget=NoTarget, + __package__=PACKAGE + ".models.minimax_m3", + ) + # The HF quant parser is a framework prerequisite; these tests supply its result. + ns["_quant_algo"] = lambda config: Quant.MXFP8 + return ns["route"] + + +def inputs(): + config = types.SimpleNamespace( + text_config=dict( + num_hidden_layers=60, + hidden_size=6144, + num_local_experts=128, + intermediate_size=3072, + num_attention_heads=64, + num_key_value_heads=4, + rope_theta=10_000_000, + sparse_attention_config=dict(use_sparse_attention=True, sparse_attention_freq=[1] * 60), + ), + quantization_config=dict(moe_quant_algo="NVFP4", moe_quant_format="nvfp4-pack-quantized"), + ) + mapping = types.SimpleNamespace( + world_size=8, + tp_size=8, + pp_size=1, + cp_size=1, + moe_tp_size=8, + moe_ep_size=1, + enable_attention_dp=False, + ) + return config, mapping + + +class TestM31Tp8RoutingSource(unittest.TestCase): + def test_declines_outside_scope(self): + cases = [("chip", (10, 0)), ("chip", (10, 7))] + cases += [ + (name, value) + for name, value in ( + ("world_size", 4), + ("tp_size", 4), + ("tp_size", 2), + ("pp_size", 2), + ("cp_size", 2), + ("moe_tp_size", 1), + ("moe_ep_size", 2), + ("enable_attention_dp", True), + ) + ] + cases += [ + (name, None) + for name in ("m3_rope", "m3_layers", "wrong_experts", "wrong_shape", "bf16_kv") + ] + for name, value in cases: + with self.subTest(name=name, value=value): + config, mapping = inputs() + chip, kv = (10, 3), Quant.NVFP4 + if name == "chip": + chip = value + elif hasattr(mapping, name): + setattr(mapping, name, value) + elif name == "m3_rope": + config.text_config["rope_theta"] = 5_000_000 + elif name == "m3_layers": + config.text_config["sparse_attention_config"]["sparse_attention_freq"] = [ + 0 + ] * 3 + [1] * 57 + elif name == "wrong_experts": + config.quantization_config["moe_quant_algo"] = "FP8" + elif name == "wrong_shape": + config.text_config["num_key_value_heads"] = 8 + else: + kv = None + with self.assertRaises(NoTarget): + route()(config, mapping, chip, None, kv) + + def test_tp8_routes_only_to_tp8(self): + config, mapping = inputs() + name = PACKAGE + ".models.minimax_m3.minimax_m31_nvfp4__sm_103__tp8.modeling" + module = types.ModuleType(name) + sentinel = type("Target", (), {}) + module.ModelingV2MiniMaxM31Nvfp4Sm103Tp8 = sentinel + fake_name = PACKAGE + ".models.minimax_m3.fake_ops" + fake_module = types.ModuleType(fake_name) + fake_module.register = lambda: None + with patch.dict(sys.modules, {name: module, fake_name: fake_module}): + self.assertIs(route()(config, mapping, (10, 3), None, Quant.NVFP4), sentinel) + + def test_router_arity_and_mode(self): + for arch, arity in [ + ("GptOssForCausalLM", 4), + ("DeepseekV3ForCausalLM", 4), + ("MiniMaxM3SparseForConditionalGeneration", 5), + ]: + with self.subTest(arch=arch): + ns = load_source( + SOURCE / "_router_index.py", logger=types.SimpleNamespace(info=lambda *a: None) + ) + calls = [] + target = type("Target", (), {}) + + def call(*args): + calls.append(args) + return target + + ns["routing_module"] = lambda name: types.SimpleNamespace(route=call) + ns["_device_compute_capability"] = lambda: (10, 3) + ns["_kv_cache_quant_algo"] = lambda *args: Quant.NVFP4 + config = types.SimpleNamespace( + pretrained_config=types.SimpleNamespace(architectures=[arch]), + mapping=object(), + spec_config=None, + ) + self.assertIs( + ns["modeling_v2_resolve"](config, types.SimpleNamespace(modeling_v2="require")), + target, + ) + self.assertEqual(len(calls[0]), arity) + if arity == 5: + self.assertIs(ns["resolved_mode"](target), ns["ModelingV2Mode"].REQUIRE) + self.assertEqual( + set(ns["MODELING_V2_ROUTERS"]), + { + "GptOssForCausalLM", + "DeepseekV3ForCausalLM", + "MiniMaxM3SparseForConditionalGeneration", + }, + ) + + def test_parent_import_and_late_native_registration(self): + """Import actual parent with no native ops, then load ops and repeat registration.""" + calls = [] + native = types.SimpleNamespace() + installed = {} + torch = types.ModuleType("torch") + torch.ops = types.SimpleNamespace(trtllm=native) + torch._C = types.SimpleNamespace( + _dispatch_has_kernel_for_dispatch_key=lambda name, key: installed.get(name) == "Meta" + ) + + def register_fake(name): + def decorate(fn): + self.assertNotIn(name, installed) + installed[name] = fn + calls.append(name) + return fn + + return decorate + + torch.library = types.SimpleNamespace(register_fake=register_fake) + registry = types.ModuleType("torch._library.simple_registry") + registry.singleton = types.SimpleNamespace( + find=lambda name: types.SimpleNamespace( + fake_impl=types.SimpleNamespace(kernel=installed.get(name)) + ) + ) + family_name = PACKAGE + ".models.minimax_m3" + catalog_name = PACKAGE + ".catalog.fakes.m31_fake_ops" + + def load_file(name, path, package=False): + spec = importlib.util.spec_from_file_location( + name, path, submodule_search_locations=[str(path.parent)] if package else None + ) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + with patch.dict(sys.modules, {"torch": torch, "torch._library.simple_registry": registry}): + for fullname in (family_name, catalog_name): + parts = fullname.split(".") + for length in range(1, len(parts)): + name = ".".join(parts[:length]) + package = types.ModuleType(name) + package.__path__ = [str(ROOT.joinpath(*parts[:length]))] + sys.modules[name] = package + load_file(catalog_name, SOURCE / "catalog/fakes/m31_fake_ops.py") + parent = load_file(family_name, SOURCE / "models/minimax_m3/__init__.py", package=True) + self.assertEqual(calls, []) + self.assertFalse(hasattr(parent, "fake_ops")) + fake_module = load_file( + family_name + ".fake_ops", SOURCE / "models/minimax_m3/fake_ops.py" + ) + self.assertEqual(len(fake_module.FAKES), 7) + # A route rejected after the parent import can still choose the built-in model. + config, mapping = inputs() + with self.assertRaises(NoTarget): + route()(config, mapping, (10, 7), None, Quant.NVFP4) + names = list(fake_module.FAKES) + for name in names: + setattr(native, name, object()) + installed["trtllm::" + names[0]] = "Meta" + previous_fake = object() + installed["trtllm::" + names[1]] = previous_fake + fake_module.register() + self.assertEqual(len(calls), 5) + fake_module.register() + self.assertEqual(len(calls), 5) + self.assertIs(installed["trtllm::" + names[1]], previous_fake) + self.assertEqual(installed["trtllm::" + names[0]], "Meta") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/speculative/test_rng_upload_pack.py b/tests/unittest/_torch/speculative/test_rng_upload_pack.py new file mode 100644 index 000000000000..4964734f7434 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_rng_upload_pack.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Philox seed/offset uploads of a sampled one-model step, staged into the step-input pack.""" + +import types +from typing import Optional + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import step_inputs +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack +from tensorrt_llm._torch.speculative.interface import SpecMetadata + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.skipif( + not (torch.cuda.is_available() and hasattr(torch.ops.trtllm, "step_inputs_scatter")), + reason="torch.ops.trtllm.step_inputs_scatter is not registered", + ), +] + +MAX_DRAFT_LEN = 3 +ROWS_PER_REQUEST = MAX_DRAFT_LEN + 1 + + +def _meta(pack_rng_uploads: bool) -> SpecMetadata: + meta = SpecMetadata( + max_num_requests=4, max_draft_len=MAX_DRAFT_LEN, max_total_draft_tokens=MAX_DRAFT_LEN + ) + meta.pack_rng_uploads = pack_rng_uploads + return meta + + +def _request(slot: int, seed: Optional[int]) -> types.SimpleNamespace: + return types.SimpleNamespace( + py_seq_slot=slot, + py_decoding_iter=0, + py_request_id=slot, + sampling_config=types.SimpleNamespace(seed=seed), + ) + + +def _populate(meta: SpecMetadata) -> None: + requests = [_request(2, seed=1234), _request(0, seed=None)] + normalized = [(0.7, 0, 0.9, 0.0, ROWS_PER_REQUEST) for _ in requests] + meta._populate_request_rng_state(requests, normalized) + + +def _rng_state(meta: SpecMetadata) -> dict[str, torch.Tensor]: + torch.cuda.synchronize() + return { + name: getattr(meta, name).clone() + for name in ("request_seeds", "request_offsets", "seeds", "offsets") + } + + +def test_packed_uploads_land_at_the_flush_with_the_plain_values(): + plain = _meta(pack_rng_uploads=False) + _populate(plain) + expected = _rng_state(plain) + + packed = _meta(pack_rng_uploads=True) + pack = StepInputPack(torch.device("cuda", torch.cuda.current_device())) + with step_inputs.active(pack): + step_inputs.begin() + _populate(packed) + before_flush = _rng_state(packed) + step_inputs.flush() + after_flush = _rng_state(packed) + + for name, tensor in expected.items(): + assert torch.equal(after_flush[name], tensor), name + # Staged, not copied: the buffers still hold their zero init until the flush. + assert not torch.equal(before_flush["seeds"], expected["seeds"]) + + +def test_without_an_open_pack_the_uploads_are_plain_copies(): + plain = _meta(pack_rng_uploads=False) + _populate(plain) + expected = _rng_state(plain) + + packed = _meta(pack_rng_uploads=True) + _populate(packed) + + for name, tensor in _rng_state(packed).items(): + assert torch.equal(tensor, expected[name]), name + + +def test_packing_is_off_by_default(): + assert not SpecMetadata( + max_num_requests=1, max_draft_len=MAX_DRAFT_LEN, max_total_draft_tokens=MAX_DRAFT_LEN + ).pack_rng_uploads + + +def test_under_capture_the_uploads_are_plain_copies(monkeypatch): + """step_inputs.copy_ never defers a copy while the stream captures (the captured copy reruns on every replay), + so a packed Philox upload issued during capture is a plain copy that lands without the flush.""" + plain = _meta(pack_rng_uploads=False) + _populate(plain) + expected = _rng_state(plain) + + packed = _meta(pack_rng_uploads=True) + pack = StepInputPack(torch.device("cuda", torch.cuda.current_device())) + with step_inputs.active(pack): + step_inputs.begin() + monkeypatch.setattr(step_inputs.torch.cuda, "is_current_stream_capturing", lambda: True) + _populate(packed) + monkeypatch.undo() + assert pack.h2d == [] and pack.staged == [] + landed = _rng_state(packed) + step_inputs.flush() + + for name, tensor in expected.items(): + assert torch.equal(landed[name], tensor), name + + +@pytest.mark.parametrize("pack_rng_uploads", [False, True]) +def test_an_all_greedy_batch_with_the_greedy_skip_uploads_nothing(pack_rng_uploads): + meta = _meta(pack_rng_uploads=pack_rng_uploads) + meta.skip_greedy_rng_upload = True + meta.is_all_greedy_sample = True + pack = StepInputPack(torch.device("cuda", torch.cuda.current_device())) + with step_inputs.active(pack): + step_inputs.begin() + _populate(meta) + assert pack.h2d == [] + step_inputs.flush() + + for name, tensor in _rng_state(meta).items(): + assert not bool(tensor.any()), name # still the zero init + + +def test_a_sampled_batch_with_the_greedy_skip_still_packs_its_uploads(): + plain = _meta(pack_rng_uploads=False) + _populate(plain) + expected = _rng_state(plain) + + packed = _meta(pack_rng_uploads=True) + packed.skip_greedy_rng_upload = True + packed.is_all_greedy_sample = False + pack = StepInputPack(torch.device("cuda", torch.cuda.current_device())) + with step_inputs.active(pack): + step_inputs.begin() + _populate(packed) + assert len(pack.h2d) == 4 # request_seeds, request_offsets, seeds, offsets + step_inputs.flush() + + for name, tensor in _rng_state(packed).items(): + assert torch.equal(tensor, expected[name]), name 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..a50e884a29d2 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py @@ -0,0 +1,648 @@ +# 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)) + + +def _host_outputs(state) -> dict[str, torch.Tensor]: + state.sampler_event.synchronize() + return dict( + new_tokens=state.host.new_tokens.clone(), + new_tokens_lens=state.host.new_tokens_lens.clone(), + next_draft_tokens=state.host.next_draft_tokens.clone(), + ) + + +def test_store_copies_stay_on_the_main_stream_by_default(): + sampler = _sampler() + scheduled, outputs = _step(MAX_DRAFT_LEN) + + state = sampler.sample_async(scheduled, outputs, []) + + assert not sampler.d2h_side_stream_copies + assert state.sampler_event.side_stream_event is None + + +@pytest.mark.parametrize("draft_width", [MAX_DRAFT_LEN, MAX_DRAFT_LEN - 1]) +def test_side_stream_store_copies_match_the_main_stream_copies(draft_width): + scheduled, outputs = _step(draft_width) + on_main = _host_outputs(_sampler().sample_async(scheduled, outputs, [])) + sampler = _sampler() + sampler.d2h_side_stream_copies = True + + state = sampler.sample_async(scheduled, outputs, []) + + assert state.sampler_event.side_stream_event is not None + on_side = _host_outputs(state) + for name, tensor in on_main.items(): + assert torch.equal(on_side[name], tensor), name + + +def test_a_fold_that_stored_the_rows_keeps_the_copies_on_the_main_stream(): + sampler = _sampler() + sampler.bind_spec_step_fold(_fold(rows_stored=True)) + sampler.d2h_side_stream_copies = True + scheduled, outputs = _step(MAX_DRAFT_LEN) + + state = sampler.sample_async(scheduled, outputs, []) + + assert state.sampler_event.side_stream_event is None + + +def test_a_fold_that_awaits_the_store_copies_moves_them_to_the_side_stream(): + """A fold with store_copies_done: the copies go on the side stream although its forward stored the rows, and + the event is recorded after them.""" + scheduled, outputs = _step(MAX_DRAFT_LEN) + fold = _fold(rows_stored=True) + fold.store_copies_done = torch.cuda.Event(external=True) + fold.store_copies_done.record() + sampler = _sampler() + sampler.bind_spec_step_fold(fold) + sampler.d2h_side_stream_copies = True + expected = _store_snapshot(sampler) + + state = sampler.sample_async(scheduled, outputs, []) + + assert state.sampler_event.side_stream_event is not None + state.sampler_event.synchronize() + fold.store_copies_done.synchronize() + host = _host_outputs(state) + for name in ("new_tokens", "new_tokens_lens", "next_draft_tokens"): + assert torch.equal(host[name], expected[name].cpu()), name + + +def test_a_fold_that_awaits_the_store_copies_keeps_them_on_the_main_stream_when_side_copies_are_off(): + scheduled, outputs = _step(MAX_DRAFT_LEN) + fold = _fold(rows_stored=True) + fold.store_copies_done = torch.cuda.Event(external=True) + sampler = _sampler() + sampler.bind_spec_step_fold(fold) + + state = sampler.sample_async(scheduled, outputs, []) + + assert state.sampler_event.side_stream_event is None + + +# HO-04 on GB: the folds the M3.1 sm_103 targets bind never store the rows in the forward (sampler_rows_stored is +# always False) and carry no store_copies_done. + + +class _ScatteringStepFold(_RecordingStepFold): + """The shape of the targets' M3SpecStepFold (Eagle3): the forward stores nothing, and ``scatter`` writes the + sampler store itself (as m3_spec_tail.slot_scatter does) and takes the step.""" + + def __init__(self) -> None: + super().__init__(rows_stored=False, scatter_taken=True) + + def scatter(self, seq_slots, new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens): + taken = super().scatter( + seq_slots, new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens + ) + store = self.store + count = seq_slots.shape[0] + + def padded(tensor: torch.Tensor, width: int) -> torch.Tensor: + tensor = tensor[:count, :width] + return torch.nn.functional.pad(tensor, (0, width - tensor.shape[1])) + + store.new_tokens.squeeze(-1).T.index_copy_( + 0, seq_slots, padded(new_tokens, store.new_tokens.shape[0]) + ) + store.next_new_tokens.squeeze(-1).T.index_copy_( + 0, seq_slots, padded(next_new_tokens, store.next_new_tokens.shape[0]) + ) + store.new_tokens_lens.index_copy_(0, seq_slots, new_tokens_lens[:count]) + store.next_draft_tokens.index_copy_( + 0, seq_slots, padded(next_draft_tokens, store.next_draft_tokens.shape[1]) + ) + return taken + + +class _AwaitedEvent: + """Stands in for the side-stream event and records that SamplerEvent.synchronize waited on it.""" + + def __init__(self, event: torch.cuda.Event) -> None: + self._event = event + self.waits = 0 + + def synchronize(self) -> None: + self.waits += 1 + self._event.synchronize() + + +@pytest.mark.parametrize("draft_width", [MAX_DRAFT_LEN, MAX_DRAFT_LEN - 1]) +def test_side_stream_copies_engage_under_the_gb_eagle3_fold(draft_width): + scheduled, outputs = _step(draft_width) + stock = _sampler() + stock.bind_spec_step_fold(_ScatteringStepFold()) + on_main = _host_outputs(stock.sample_async(scheduled, outputs, [])) + sampler = _sampler() + fold = _ScatteringStepFold() + sampler.bind_spec_step_fold(fold) + sampler.d2h_side_stream_copies = True + expected = _reference_store(_store_snapshot(sampler), outputs) + + state = sampler.sample_async(scheduled, outputs, []) + + assert not hasattr(fold, "store_copies_done") + assert fold.rows_stored_calls == [(0, SEQ_SLOTS)] and len(fold.scatter_calls) == 1 + side_event = state.sampler_event.side_stream_event + assert side_event is not None + awaited = _AwaitedEvent(side_event) + state.sampler_event.side_stream_event = awaited + host = _host_outputs(state) + assert awaited.waits == 1 + _assert_store_equals(sampler, expected) + for name, tensor in on_main.items(): + assert torch.equal(host[name], tensor), name + assert torch.equal(host[name], expected[name].cpu()), name + + +@pytest.mark.parametrize("draft_width", [MAX_DRAFT_LEN, MAX_DRAFT_LEN - 1]) +def test_side_stream_copies_engage_under_the_gb_dspark_fold(draft_width): + """The targets' M3DSparkStepFold (standalone DSpark): scatter declines every step, so the store is written only by + the sampler's own _store_outputs, and the copies go on the side stream.""" + scheduled, outputs = _step(draft_width) + sampler = _sampler() + fold = _fold(rows_stored=False, scatter_taken=False) + sampler.bind_spec_step_fold(fold) + sampler.d2h_side_stream_copies = True + before = _store_snapshot(sampler) + expected = _reference_store(before, outputs) + store_writes = [] + store_outputs = sampler._store_outputs + + def recording_store_outputs(*args): + store_writes.append(args[0]) + return store_outputs(*args) + + sampler._store_outputs = recording_store_outputs + + state = sampler.sample_async(scheduled, outputs, []) + + assert not hasattr(fold, "store_copies_done") + assert store_writes == [SEQ_SLOTS] + assert len(fold.scatter_calls) == 1 # offered, declined + assert state.sampler_event.side_stream_event is not None + host = _host_outputs(state) + _assert_store_equals(sampler, expected) + for name, tensor in host.items(): + assert torch.equal(tensor, expected[name].cpu()), name + + +def test_flags_off_issues_the_stock_copies_in_the_stock_order(): + """Flags off, sample_async issues exactly the stock calls: four _copy_to_host (new_tokens, new_tokens_lens, + next_draft_tokens, verify_lens) and one argument-less _record_sampler_event.""" + sampler = _sampler() + scheduled, outputs = _step(MAX_DRAFT_LEN) + rows = outputs["new_tokens_lens"].shape[0] + outputs["verify_lens"] = torch.full( + (rows,), MAX_DRAFT_LEN + 1, dtype=torch.int32, device="cuda" + ) + outputs["verify_lens_in_output_order"] = True + copies = [] + events = [] + copy_to_host = sampler._copy_to_host + record_sampler_event = sampler._record_sampler_event + + def recording_copy(src): + copies.append(src) + return copy_to_host(src) + + def recording_event(*args, **kwargs): + events.append((args, kwargs)) + return record_sampler_event(*args, **kwargs) + + sampler._copy_to_host = recording_copy + sampler._record_sampler_event = recording_event + + state = sampler.sample_async(scheduled, outputs, []) + + assert len(copies) == 4 + assert copies[0] is sampler.store.new_tokens + assert copies[1] is sampler.store.new_tokens_lens + assert copies[2] is sampler.store.next_draft_tokens + assert copies[3].data_ptr() == outputs["verify_lens"].data_ptr() + assert copies[3].shape[0] == len(SEQ_SLOTS) + assert events == [((), {})] + assert state.sampler_event.side_stream_event is None + + +def test_side_stream_copies_fall_back_to_the_main_stream_with_the_async_worker(): + sampler = _sampler() + sampler.d2h_side_stream_copies = True + scheduled, outputs = _step(MAX_DRAFT_LEN) + copies = [] + + def plain_copy(src): + # What _copy_to_host does without a worker thread; the worker itself is not started here. + copies.append(src) + dest = torch.empty_like(src, device="cpu", pin_memory=True) + dest.copy_(src, non_blocking=True) + return dest + + sampler._async_worker_active = lambda: True + sampler._copy_to_host = plain_copy + + state = sampler.sample_async(scheduled, outputs, []) + + assert len(copies) == 3 + assert state.sampler_event.side_stream_event is None + + +def test_a_sampler_built_without_init_reads_both_hooks_off(): + sampler = object.__new__(SpecSampler) + + assert sampler.d2h_side_stream_copies is False + assert sampler.static_slot_ids is False + assert sampler._slot_ids_cuda is None and sampler._slot_ids_host is None + + +# HO-12a: static slot ids. + + +def test_static_slot_ids_are_off_by_default(): + assert not _sampler().static_slot_ids + + +@pytest.mark.parametrize("draft_width", [MAX_DRAFT_LEN, MAX_DRAFT_LEN - 1]) +def test_static_slot_ids_store_the_same_rows(draft_width): + scheduled, outputs = _step(draft_width) + plain = _sampler() + expected = _reference_store(_store_snapshot(plain), outputs) + sampler = _sampler() + sampler.static_slot_ids = True + + sampler.sample_async(scheduled, outputs, []) + sampler.sample_async(scheduled, outputs, []) # second step reuses the uploaded ids + + _assert_store_equals(sampler, expected) + + +def test_static_slot_ids_follow_a_changed_batch(): + sampler = _sampler() + sampler.static_slot_ids = True + + first_view = sampler._device_slot_ids([4, 1]) + first = first_view.clone() + # Unchanged ids reuse the sampler's buffer without a new upload. The returned tensor is a view of that buffer, + # which the next call rewrites, so snapshot it before changing the batch. + reused_view = sampler._device_slot_ids([4, 1]) + assert reused_view.data_ptr() == first_view.data_ptr() + reused = reused_view.clone() + changed = sampler._device_slot_ids([2, 5, 0]).clone() + + torch.cuda.synchronize() + assert first.tolist() == [4, 1] + assert reused.tolist() == [4, 1] + assert changed.tolist() == [2, 5, 0] + + +@pytest.fixture +def slot_id_uploads(monkeypatch): + """Counts the copies into the sampler's slot-id buffer (bound with ``.watch(sampler)``).""" + + class _Uploads: + sampler: Optional[SpecSampler] = None + count = 0 + + def watch(self, sampler: SpecSampler) -> "_Uploads": + self.sampler = sampler + return self + + uploads = _Uploads() + copy_ = torch.Tensor.copy_ + + def counting_copy_(dst, src, *args, **kwargs): + buffer = uploads.sampler._slot_ids_cuda if uploads.sampler is not None else None + if ( + buffer is not None + and dst.is_cuda + and dst.untyped_storage().data_ptr() == buffer.untyped_storage().data_ptr() + ): + uploads.count += 1 + return copy_(dst, src, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "copy_", counting_copy_) + return uploads + + +def test_steady_slot_ids_upload_once_over_many_steps(slot_id_uploads): + sampler = _sampler() + sampler.static_slot_ids = True + uploads = slot_id_uploads.watch(sampler) + scheduled, outputs = _step(MAX_DRAFT_LEN) + + for _ in range(5): + sampler.sample_async(scheduled, outputs, []) + + assert uploads.count == 1 + torch.cuda.synchronize() + assert sampler._slot_ids_cuda[: len(SEQ_SLOTS)].tolist() == SEQ_SLOTS + + +def test_a_changed_slot_list_re_uploads(slot_id_uploads): + sampler = _sampler() + uploads = slot_id_uploads.watch(sampler) + + sampler._device_slot_ids([4, 1]) + sampler._device_slot_ids([4, 1]) + assert uploads.count == 1 + sampler._device_slot_ids([1, 4]) + assert uploads.count == 2 + sampler._device_slot_ids([1]) # a prefix of the buffer still differs from the last list + assert uploads.count == 3 + view = sampler._device_slot_ids([1]) + assert uploads.count == 3 + torch.cuda.synchronize() + assert view.tolist() == [1] + + +def test_a_longer_slot_list_reallocates_and_uploads(slot_id_uploads): + sampler = _sampler() + uploads = slot_id_uploads.watch(sampler) + + sampler._device_slot_ids([4, 1]) + small = sampler._slot_ids_cuda + assert small.numel() == 8 # max(n, 8) + long_ids = list(range(10)) + view = sampler._device_slot_ids(long_ids) + + assert sampler._slot_ids_cuda is not small and sampler._slot_ids_cuda.numel() == 10 + assert uploads.count == 2 + torch.cuda.synchronize() + assert view.tolist() == long_ids + # The same ids again: no upload into the new buffer. + sampler._device_slot_ids(long_ids) + assert uploads.count == 2 + + +def test_both_hooks_on_store_and_copy_what_the_stock_sampler_does_over_a_changing_batch(): + """The targets run HO-04 and HO-12a together: over steady, reordered and repeated steps the store and the host + copies are bitwise the stock ones.""" + stock = _sampler() + sampler = _sampler() + sampler.d2h_side_stream_copies = True + sampler.static_slot_ids = True + scheduled, outputs = _step(MAX_DRAFT_LEN) + for step in range(4): + if step == 2: # the batch's slot order changes + scheduled.generation_requests = scheduled.generation_requests[::-1] + step_outputs = {name: tensor + step for name, tensor in outputs.items()} + on_main = _host_outputs(stock.sample_async(scheduled, step_outputs, [])) + state = sampler.sample_async(scheduled, step_outputs, []) + assert state.sampler_event.side_stream_event is not None + on_side = _host_outputs(state) + for name, tensor in on_main.items(): + assert torch.equal(on_side[name], tensor), (step, name) + _assert_store_equals(sampler, _store_snapshot(stock)) diff --git a/tests/unittest/_torch/speculative/test_tail_bookkeep.py b/tests/unittest/_torch/speculative/test_tail_bookkeep.py new file mode 100644 index 000000000000..c8a98fecfed7 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_tail_bookkeep.py @@ -0,0 +1,590 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tail bookkeeping (RuntimeHooks.spec_tail_bookkeep -> SpecMetadata.tail_bookkeep): the shared items against the +in-tree path, bitwise. + +* B1 (``advanced_sample_draft_block``): at one decode request the draft block's sampling parameters are views of the + per-token buffers instead of four repeats of the per-request ones. +* S1 (``sample_draft_tokens``): a draft block with no d2t map returns its int32 tokens without the int64 round trip. +* A1 (``_sample_and_accept_draft_tokens_rejection``): a decode-only accept returns the acceptance kernel's outputs + instead of copying them into a ones-filled and an empty buffer. + +Each test runs the same step twice, on a metadata with the flag off (the in-tree path) and one with it on, and compares +every output (and every tensor the samplers were handed) bit for bit. The samplers are deterministic stand-ins that +record their inputs, so the comparison covers the parameter and Philox tensors, not only the tokens. + +The CPU tests fill the per-token buffers with a mirror of ``populate_sampling_params_for_one_model``'s layout and also +check that the flag off takes the in-tree lines. The CUDA tests fill them with the real populate, eagerly and inside +``torch.cuda.graph`` (captured on one request's values, replayed on another's). R1 and Z1 are the M3.1 target's +overrides and are tested with it (test_modeling_v2_m31_tailbk.py). +""" + +import types +from typing import Optional + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as runner_module +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.runtime_hooks import DEFAULT_RUNTIME_HOOKS, RuntimeHooks +from tensorrt_llm._torch.speculative import interface +from tensorrt_llm._torch.speculative.interface import ( + _RNG_SLOT_SPAN, + SpecMetadata, + SpeculativeDecodingMode, + SpecWorkerBase, + tail_bookkeep_draft_len, +) +from tensorrt_llm.llmapi.llm_args import AdvancedSamplingMode + +CUDA = torch.cuda.is_available() +needs_cuda = pytest.mark.skipif(not CUDA, reason="the real populate allocates CUDA buffers") + +V = 64 # vocab +MAX_REQUESTS = 8 +DRAFT_LENS = list(range(1, 8)) # K = 1..7 +# (temperature, top_k, top_p, min_p, seed), all sampled: unseeded, a seed at the int64 edge, a uint64 seed +# (reinterpreted as negative), seed 0, and top-k 2 with a large min-p (top-k 1 would be greedy). +REQUESTS = [ + (0.7, 0, 0.9, 0.0, None), + (1.3, 17, 0.5, 0.05, (1 << 63) - 1), + (0.25, 3, 1.0, 0.2, (1 << 64) - 1), + (1.0, 0, 0.95, 0.0, 0), + (2.0, 2, 0.1, 0.5, 1234), +] + + +class _Worker(SpecWorkerBase): + """SpecWorkerBase without its constructor (which needs flashinfer): the state the sampler tail reads.""" + + def _forward_impl(self, *args, **kwargs): # pragma: no cover - never run + raise NotImplementedError + + @property + def max_draft_len(self) -> int: + return max(DRAFT_LENS) + + +def _worker(d2t: Optional[torch.Tensor] = None, force_num_accepted_tokens: float = 0.0) -> _Worker: + worker = object.__new__(_Worker) + worker._d2t = d2t + worker.force_num_accepted_tokens = force_num_accepted_tokens + worker.mapping = None + return worker + + +def _request(params, slot: int = 3, request_id: int = 7): + temperature, top_k, top_p, min_p, seed = params + return types.SimpleNamespace( + sampling_config=types.SimpleNamespace( + temperature=temperature, top_k=top_k, top_p=top_p, min_p=min_p, seed=seed + ), + state=LlmRequestState.GENERATION_IN_PROGRESS, + py_seq_slot=slot, + py_request_id=request_id, + py_decoding_iter=0, + ) + + +def _meta(k: int, on: bool, mode=AdvancedSamplingMode.NO_TOPK) -> SpecMetadata: + """A one-model DSpark metadata with rejection sampling (the advanced draft path) at draft length ``k``.""" + meta = SpecMetadata( + max_num_requests=MAX_REQUESTS, + max_draft_len=k, + max_total_draft_tokens=k, + spec_dec_mode=SpeculativeDecodingMode.DSPARK, + ) + assert meta.spec_dec_mode.use_one_engine() + meta.runtime_draft_len = k + meta.advanced_sampling_mode = mode + meta.use_rejection_sampling = True + meta.vocab_size = V + meta.draft_vocab_size = V + meta.tail_bookkeep = on + return meta + + +def _fill_like_populate(meta: SpecMetadata, params, offset: int, device: str = "cpu") -> None: + """CPU mirror of populate_sampling_params_for_one_model + _populate_request_rng_state for one generation request + (the in-tree code allocates on CUDA): request_* hold its values, its K + 1 per-token rows hold the same values and + the offset window ``offset + r * SPAN``; every other row holds junk, so a wrong slice shows.""" + temperature, top_k, top_p, min_p, seed = params + k = meta.max_draft_len + rows = (k + 1) * MAX_REQUESTS + junk = torch.arange(rows, device=device) + + def per_token(value, dtype): + out = (junk + 1000).to(dtype) + out[: k + 1] = value + return out + + def per_request(value, dtype): + out = (junk[:MAX_REQUESTS] + 2000).to(dtype) + out[0] = value + return out + + seed = interface.DEFAULT_SAMPLING_SEED if seed is None else seed + seed = seed - (1 << 64) if seed >= (1 << 63) else seed + meta.temperatures = per_token(temperature, torch.float32) + meta.top_ks = per_token(top_k, torch.int32) + meta.top_ps = per_token(top_p, torch.float32) + meta.min_ps = per_token(min_p, torch.float32) + meta.request_temperatures = per_request(temperature, torch.float32) + meta.request_top_ks = per_request(top_k, torch.int32) + meta.request_top_ps = per_request(top_p, torch.float32) + meta.request_min_ps = per_request(min_p, torch.float32) + meta.request_seeds = per_request(seed, torch.int64) + meta.request_offsets = per_request(offset, torch.int64) + meta.seeds = per_token(seed, torch.int64) + meta.offsets = (junk + 3000).to(torch.int64) + meta.offsets[: k + 1] = torch.tensor([offset + r * _RNG_SLOT_SPAN for r in range(k + 1)]) + meta.is_all_greedy_sample = False + + +def _populate(meta: SpecMetadata, params, slot: int = 3, request_id: int = 7) -> None: + """The real populate for a one-request decode batch (CUDA).""" + meta.populate_sampling_params_for_one_model([_request(params, slot, request_id)]) + assert not meta.is_all_greedy_sample + + +class _Sampler: + """A deterministic stand-in for the advanced samplers: records (a copy of) every tensor it is handed and returns + the argmax as int32 tokens with softmax probs.""" + + def __init__(self): + self.seen = [] + + def __call__( + self, mode, logits, temperatures, top_ks, top_ps, min_ps, *, seed=None, offset=None + ): + self.seen.append( + tuple(t.clone() for t in (temperatures, top_ks, top_ps, min_ps, seed, offset)) + ) + return logits.argmax(-1).to(torch.int32), torch.softmax(logits, dim=-1) + + +def _fake_rejection(draft_probs, draft_token_ids, target_probs, seed, offset): + """A stand-in acceptance kernel: fresh int32 outputs that depend on every input, as the in-tree op's.""" + bonus = target_probs[:, -1].argmax(-1, keepdim=True).to(torch.int32) + accepted = torch.cat([draft_token_ids, bonus], dim=1) + num = ((seed ^ offset) % (draft_token_ids.shape[1] + 1) + 1).to(torch.int32) + num = num + (draft_probs.sum(dim=(1, 2)) < 0).to(torch.int32) # reads the draft probs (adds 0) + return accepted, num + + +def _patch_samplers(monkeypatch) -> _Sampler: + sampler = _Sampler() + monkeypatch.setattr(SpecWorkerBase, "_sample_from_logits_with_probs", staticmethod(sampler)) + monkeypatch.setattr( + SpecWorkerBase, + "_compute_probs_from_logits", + staticmethod(lambda mode, logits, t, k, p, m: torch.softmax(logits / t.unsqueeze(-1), -1)), + ) + monkeypatch.setattr(interface, "rejection_sampling_one_model", _fake_rejection) + return sampler + + +def _assert_same(a, b, what: str) -> None: + if isinstance(a, (tuple, list)): + assert isinstance(b, (tuple, list)) and len(a) == len(b), what + for i, (x, y) in enumerate(zip(a, b)): + _assert_same(x, y, f"{what}[{i}]") + return + if a is None or b is None: + assert a is None and b is None, what + return + assert a.dtype == b.dtype and a.shape == b.shape, ( + f"{what}: {a.dtype}{tuple(a.shape)} vs {b.dtype}{tuple(b.shape)}" + ) + if a.dtype == torch.float32: + a, b = a.view(torch.int32), b.view(torch.int32) # bits, so -0.0 / NaN payloads count too + assert torch.equal(a.cpu(), b.cpu()), what + + +# --------------------------------------------------------------------------------------------------------------------- +# The switch: off by default everywhere, copied to the spec metadata. +# --------------------------------------------------------------------------------------------------------------------- + + +@pytest.mark.cpu_only +def test_the_hook_and_the_metadata_flag_are_off_by_default(): + assert RuntimeHooks().spec_tail_bookkeep is False + assert DEFAULT_RUNTIME_HOOKS.spec_tail_bookkeep is False + assert SpecMetadata.tail_bookkeep is False + assert ( + SpecMetadata(max_num_requests=1, max_draft_len=1, max_total_draft_tokens=1).tail_bookkeep + is False + ) + assert not RuntimeHooks(spec_tail_bookkeep=True).skips_unchanged_metadata_uploads + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("on", [False, True]) +def test_setup_speculative_copies_the_hook_and_logs_nothing(monkeypatch, on): + created = types.SimpleNamespace() + lines = [] + monkeypatch.setattr(runner_module, "create_spec_metadata", lambda *args: created) + monkeypatch.setattr(runner_module.logger, "info", lines.append) + metadata = object.__new__(DecoderRunner) + metadata.runtime_hooks = RuntimeHooks(spec_tail_bookkeep=on) + metadata.spec_metadata = None + metadata._config = None + metadata.model = types.SimpleNamespace(config=None) + + assert metadata._set_up_spec_metadata(None) is created + assert created.tail_bookkeep is on + assert lines == [] # the arm line is the target's (it knows which items its worker takes) + + +# --------------------------------------------------------------------------------------------------------------------- +# The shared precondition (host state and shapes only). +# --------------------------------------------------------------------------------------------------------------------- + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("k", DRAFT_LENS) +def test_the_precondition_holds_for_one_decode_request(k): + meta = _meta(k, on=True) + _fill_like_populate(meta, REQUESTS[0], offset=0) + assert tail_bookkeep_draft_len(meta, 0, 1) == k + + +def _breaks(): + def flag_off(meta): + meta.tail_bookkeep = False + + def tree(meta): + meta.is_spec_dec_tree = True + + def dynamic_tree(meta): + meta.is_spec_dec_dynamic_tree = True + + def ragged(meta): + meta.verify_lens = torch.ones(1, dtype=torch.int32) + + def dynamic_draft_len(meta): + meta.runtime_draft_len = meta.max_draft_len - 1 + + def short_buffer(meta): + meta.seeds = meta.seeds[: meta.max_draft_len] + + def missing_buffer(meta): + meta.min_ps = None + + def wrong_dtype(meta): + meta.offsets = meta.offsets.to(torch.int32) + + def strided_buffer(meta): + meta.top_ps = torch.stack([meta.top_ps, meta.top_ps], dim=1)[:, 0] + + return [flag_off, tree, dynamic_tree, ragged, dynamic_draft_len, short_buffer, missing_buffer, wrong_dtype, + strided_buffer] # fmt: skip + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("brk", _breaks(), ids=lambda f: f.__name__) +def test_the_precondition_fails_closed(brk): + meta = _meta(3, on=True) + _fill_like_populate(meta, REQUESTS[0], offset=0) + brk(meta) + assert tail_bookkeep_draft_len(meta, 0, 1) == 0 + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("num_contexts,batch_size", [(0, 2), (1, 2), (0, 0), (1, 1), (0, None)]) +def test_the_precondition_needs_exactly_one_generation_request(num_contexts, batch_size): + meta = _meta(3, on=True) + _fill_like_populate(meta, REQUESTS[0], offset=0) + assert tail_bookkeep_draft_len(meta, num_contexts, batch_size) == 0 + + +# --------------------------------------------------------------------------------------------------------------------- +# B1 + S1: the draft block (sample_draft_tokens, block form). +# --------------------------------------------------------------------------------------------------------------------- + + +def _draft_block(meta, logits, d2t=None): + worker = _worker(d2t=d2t) + tokens = worker.sample_draft_tokens(logits, meta, 1, num_contexts=0) + return tokens + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("k", DRAFT_LENS) +@pytest.mark.parametrize("params", REQUESTS) +def test_draft_block_cpu_on_matches_off(monkeypatch, k, params): + sampler = _patch_samplers(monkeypatch) + logits = torch.randn(1, k, V, generator=torch.Generator().manual_seed(k)) + outs = {} + for on in (False, True): + meta = _meta(k, on) + _fill_like_populate(meta, params, offset=5 * k) + tokens = _draft_block(meta, logits) + outs[on] = (tokens, sampler.seen[-1]) + _assert_same(outs[False], outs[True], "tokens + sampler inputs") + assert outs[True][0].dtype == torch.int32 + + +@pytest.mark.cpu_only +def test_draft_block_cpu_flag_off_takes_the_in_tree_lines(monkeypatch): + """Off: the parameters are fresh repeats (not views of the per-token buffers) and the tokens a fresh int32 copy + (not the sampler's tensor).""" + k = 4 + seen_tokens = [] + + def sampler(mode, logits, temperatures, top_ks, top_ps, min_ps, *, seed=None, offset=None): + tokens = logits.argmax(-1).to(torch.int32) + seen_tokens.append(tokens) + sampler.params = (temperatures, top_ks, top_ps, min_ps) + return tokens, torch.softmax(logits, -1) + + monkeypatch.setattr(SpecWorkerBase, "_sample_from_logits_with_probs", staticmethod(sampler)) + logits = torch.randn(1, k, V) + for on in (False, True): + meta = _meta(k, on) + _fill_like_populate(meta, REQUESTS[1], offset=0) + tokens = _draft_block(meta, logits) + views = [p.untyped_storage().data_ptr() for p in sampler.params] + bufs = [meta.temperatures, meta.top_ks, meta.top_ps, meta.min_ps] + shared = [v == b.untyped_storage().data_ptr() for v, b in zip(views, bufs)] + aliases_sampler = ( + tokens.untyped_storage().data_ptr() == seen_tokens[-1].untyped_storage().data_ptr() + ) + assert shared == [on] * 4 + assert aliases_sampler is on + + +@pytest.mark.cpu_only +def test_s1_keeps_the_round_trip_with_a_d2t_map(monkeypatch): + _patch_samplers(monkeypatch) + k = 3 + d2t = torch.arange(V, dtype=torch.int64) % 5 + logits = torch.randn(1, k, V) + outs = [] + for on in (False, True): + meta = _meta(k, on) + _fill_like_populate(meta, REQUESTS[0], offset=0) + outs.append(_draft_block(meta, logits, d2t=d2t)) + _assert_same(outs[0], outs[1], "d2t-mapped tokens") + + +@pytest.mark.cpu_only +def test_b1_is_not_taken_for_a_padded_batch(monkeypatch): + """Two generation requests (a padded graph): the per-request repeats run with the flag on too.""" + sampler = _patch_samplers(monkeypatch) + k = 3 + meta = _meta(k, on=True) + _fill_like_populate(meta, REQUESTS[0], offset=0) + worker = _worker() + worker.advanced_sample_draft_block(torch.randn(2, k, V), meta, 0, 2) + temps = sampler.seen[-1][0] + assert torch.equal(temps, meta.request_temperatures[:2].repeat_interleave(k)) + + +# --------------------------------------------------------------------------------------------------------------------- +# A1: the rejection accept. +# --------------------------------------------------------------------------------------------------------------------- + + +def _accept(meta, logits, draft_tokens, draft_probs, num_contexts, batch_size, force=0.0): + worker = _worker(force_num_accepted_tokens=force) + return worker._sample_and_accept_draft_tokens_rejection( + logits, draft_tokens, draft_probs, num_contexts, batch_size, meta + ) + + +def _accept_inputs(k, num_gens, device="cpu", seed=0): + g = torch.Generator(device=device).manual_seed(seed) + logits = torch.randn(num_gens * (k + 1), V, generator=g, device=device) + draft_tokens = torch.randint(0, V, (num_gens, k), generator=g, device=device, dtype=torch.int32) + draft_probs = torch.softmax(torch.randn(num_gens, k, V, generator=g, device=device), -1) + return logits, draft_tokens, draft_probs + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("k", DRAFT_LENS) +@pytest.mark.parametrize("force", [0.0, 2.0]) +def test_accept_cpu_on_matches_off(monkeypatch, k, force): + _patch_samplers(monkeypatch) + logits, draft_tokens, draft_probs = _accept_inputs(k, 1, seed=k) + outs = [] + for on in (False, True): + meta = _meta(k, on) + _fill_like_populate(meta, REQUESTS[2], offset=96 * k) + outs.append(_accept(meta, logits, draft_tokens, draft_probs, 0, 1, force=force)) + _assert_same(outs[0], outs[1], "accepted tokens + counts") + + +@pytest.mark.cpu_only +def test_accept_flag_on_returns_the_kernel_outputs(monkeypatch): + returned = [] + + def rejection(*args, **kwargs): + returned.append(_fake_rejection(*args, **kwargs)) + return returned[-1] + + _patch_samplers(monkeypatch) + monkeypatch.setattr(interface, "rejection_sampling_one_model", rejection) + k = 3 + logits, draft_tokens, draft_probs = _accept_inputs(k, 1) + for on in (False, True): + meta = _meta(k, on) + _fill_like_populate(meta, REQUESTS[0], offset=0) + accepted, num = _accept(meta, logits, draft_tokens, draft_probs, 0, 1) + assert (accepted is returned[-1][0]) is on + assert (num is returned[-1][1]) is on + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "bad", + ["int64_tokens", "int64_counts", "strided_tokens"], +) +def test_accept_falls_back_when_the_kernel_outputs_differ_in_layout(monkeypatch, bad): + def rejection(*args, **kwargs): + accepted, num = _fake_rejection(*args, **kwargs) + if bad == "int64_tokens": + accepted = accepted.long() + elif bad == "int64_counts": + num = num.long() + else: + strided = torch.zeros(accepted.shape[0], 2 * accepted.shape[1], dtype=accepted.dtype) + strided[:, ::2] = accepted + accepted = strided[:, ::2] + return accepted, num + + _patch_samplers(monkeypatch) + monkeypatch.setattr(interface, "rejection_sampling_one_model", rejection) + k = 3 + logits, draft_tokens, draft_probs = _accept_inputs(k, 1) + outs = [] + for on in (False, True): + meta = _meta(k, on) + _fill_like_populate(meta, REQUESTS[0], offset=0) + outs.append(_accept(meta, logits, draft_tokens, draft_probs, 0, 1)) + _assert_same(outs[0], outs[1], "fallback") + assert outs[1][0].dtype == torch.int32 and outs[1][0].is_contiguous() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("num_gens", [2, 4]) +def test_accept_with_several_requests_matches_off(monkeypatch, num_gens): + """A1 is not limited to one request: any decode-only accept returns the kernel's outputs.""" + _patch_samplers(monkeypatch) + k = 3 + logits, draft_tokens, draft_probs = _accept_inputs(k, num_gens) + outs = [] + for on in (False, True): + meta = _meta(k, on) + _fill_like_populate(meta, REQUESTS[0], offset=0) + outs.append(_accept(meta, logits, draft_tokens, draft_probs, 0, num_gens)) + _assert_same(outs[0], outs[1], "accept") + + +@pytest.mark.cpu_only +def test_accept_with_a_context_row_takes_the_in_tree_lines(monkeypatch): + """num_contexts > 0: the outputs are assembled into fresh buffers whatever the flag.""" + _patch_samplers(monkeypatch) + monkeypatch.setattr( + _Worker, + "_sample_tokens_for_batch", + lambda self, logits, meta, n, b: logits.argmax(-1).to(torch.int32), + ) + returned = [] + + def rejection(*args, **kwargs): + returned.append(_fake_rejection(*args, **kwargs)) + return returned[-1] + + monkeypatch.setattr(interface, "rejection_sampling_one_model", rejection) + k = 3 + gen_logits, draft_tokens, draft_probs = _accept_inputs(k, 1) + logits = torch.cat([torch.randn(1, V), gen_logits]) + outs = [] + for on in (False, True): + meta = _meta(k, on) + _fill_like_populate(meta, REQUESTS[0], offset=0) + accepted, num = _accept(meta, logits, draft_tokens, draft_probs, 1, 2) + assert accepted is not returned[-1][0] and num is not returned[-1][1] + # The context row defines only its first token (the rest of its row is the in-tree torch.empty), so compare + # the defined entries: the context row's first token, the generation rows and every count. + outs.append((accepted[:1, :1], accepted[1:], num)) + _assert_same(outs[0], outs[1], "mixed batch") + + +# --------------------------------------------------------------------------------------------------------------------- +# CUDA: the real populate, eager and inside torch.cuda.graph. +# --------------------------------------------------------------------------------------------------------------------- + + +def _graph(fn): + """Capture ``fn`` (after a side-stream warmup, as torch.cuda.graphs asks); returns (graph, captured outputs).""" + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + fn() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = fn() + return graph, out + + +def _step(meta, worker, logits, accept_inputs, sampler): + """One sampled step's tail: the rejection accept, then the draft block; the outputs and the draft-prob row.""" + v_logits, draft_tokens, draft_probs = accept_inputs + accepted = worker._sample_and_accept_draft_tokens_rejection( + v_logits, draft_tokens, draft_probs, 0, 1, meta + ) + tokens = worker.sample_draft_tokens(logits, meta, 1, num_contexts=0) + return accepted, tokens, sampler.seen[-1] + + +@needs_cuda +@pytest.mark.parametrize("k", DRAFT_LENS) +@pytest.mark.parametrize("mode", [AdvancedSamplingMode.NO_TOPK, AdvancedSamplingMode.FULL]) +def test_cuda_eager_on_matches_off(monkeypatch, k, mode): + sampler = _patch_samplers(monkeypatch) + worker = _worker() + metas = {on: _meta(k, on, mode=mode) for on in (False, True)} + logits = torch.randn(1, k, V, device="cuda") + inputs = _accept_inputs(k, 1, device="cuda", seed=k) + for step, params in enumerate(REQUESTS): + outs = {} + for on, meta in metas.items(): + _populate(meta, params) + outs[on] = _step(meta, worker, logits + step, inputs, sampler) + outs[on] += (meta.draft_probs[3, :k].clone(),) + torch.cuda.synchronize() + _assert_same(outs[False], outs[True], f"step {step}") + + +@needs_cuda +@pytest.mark.parametrize("k", [1, 3, 7]) +def test_cuda_graph_on_matches_off(monkeypatch, k): + """The ON step captured on the first request's values, then replayed after each later populate (same buffers, + new values), against the in-tree step run eagerly on the same values.""" + sampler = _patch_samplers(monkeypatch) + worker = _worker() + on_meta, off_meta = _meta(k, True), _meta(k, False) + logits = torch.randn(1, k, V, device="cuda") + inputs = _accept_inputs(k, 1, device="cuda", seed=k) + _populate(on_meta, REQUESTS[0]) + _populate(off_meta, REQUESTS[0]) + graph, captured = _graph(lambda: _step(on_meta, worker, logits, inputs, sampler)) + for step, params in enumerate(REQUESTS[1:], start=1): + _populate(on_meta, params) + graph.replay() + _populate(off_meta, params) + expected = _step(off_meta, worker, logits, inputs, sampler) + torch.cuda.synchronize() + _assert_same(captured, expected, f"replay {step}") + _assert_same(on_meta.draft_probs[3, :k], off_meta.draft_probs[3, :k], f"draft probs {step}") + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"]))