diff --git a/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu b/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu index 102866de363f..6719dd3bf995 100644 --- a/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu +++ b/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu @@ -18,6 +18,7 @@ #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.h" +#include using namespace tensorrt_llm::common; TRTLLM_NAMESPACE_BEGIN @@ -194,6 +195,216 @@ void invokeRouterGemm(float* output, T const* mat_a, T const* mat_b, cudaStream_ &config, router_gemm_kernel, output, mat_a, mat_b)); } +// D += A(16x16, row) * B(16x8, col), bf16 inputs, fp32 accumulate. Fragment layout (PTX ISA, m16n8k16): +// a[0] = A[g][2q..2q+1] a[1] = A[g+8][2q..2q+1] a[2] = A[g][2q+8..2q+9] a[3] = A[g+8][2q+8..2q+9] +// b[0] = B[2q..2q+1][g] b[1] = B[2q+8..2q+9][g] c[0..1] = D[g][2q..2q+1] c[2..3] = D[g+8][2q..2q+1] +// with g = lane / 4, q = lane % 4. +__device__ __forceinline__ void mma_bf16_m16n8k16(float (&c)[4], uint32_t const (&a)[4], uint32_t const (&b)[2]) +{ + 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"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +// Tensor-core variant for large expert counts (Kimi K3: 896 experts). One warp computes a 16-token x 8-expert +// tile over a K-slice with mma.sync; the kNumWarps warps of a CTA split K and are reduced through shared memory +// in warp order (deterministic). Each lane loads 16 bytes (8 consecutive k) per row/column per 32-wide k-chunk; +// the k order inside a chunk is permuted identically for A and B (lane group q owns k = q*8..q*8+7, fed to the +// two MMAs as k-pairs {0,1},{2,3} and {4,5},{6,7}), which is legal because the dot product is order-invariant. +// Motivation (B200, M=8): router_gemm_kernel executes ~4M warp-instructions (M x 896 x 7168 FMAs + bf16->fp32 +// converts) and is issue-bound (~26% issue utilisation, time ~linear in M); one MMA replaces 64 warp-FFMAs. +template +__device__ __forceinline__ void routerGemmMma(float* out, __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b, + __nv_bfloat16* projection, __nv_bfloat16 const* downWeight) +{ + static_assert(kNumTokens >= 1 && kNumTokens <= 16, "one m16 tile in the token dimension"); + static_assert(kExpertsPerCta % 8 == 0, "expert columns per CTA must be whole n8 tiles"); + static_assert(kLatentDim == 0 || (kNumExperts % kExpertsPerCta == 0 && kLatentDim % kExpertsPerCta == 0), + "gate and latent output boundaries must align to CTA tiles"); + constexpr int kNumWarps = kBlockSize / 32; + constexpr int kTilesN = kExpertsPerCta / 8; + constexpr int kChunk = 32; // k elements per lane load (8 per lane group) = two m16n8k16 MMAs + static_assert(kHiddenDim % (kNumWarps * kChunk) == 0, "K must split evenly over warps and 32-wide chunks"); + constexpr int kChunksPerWarp = kHiddenDim / (kNumWarps * kChunk); + constexpr int kBatch = kChunksPerWarp < 8 ? kChunksPerWarp : 8; // chunks whose loads are kept in flight + static_assert(kChunksPerWarp % kBatch == 0, "chunks per warp must be a multiple of the load batch"); + constexpr bool kTwoRows = kNumTokens > 8; + + int const warpId = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + int const g = lane / 4; // token row (A) / expert column (B) owned by this lane in the fragments + int const q = lane % 4; // 8-wide k sub-chunk owned by this lane + int const n_base = blockIdx.x * kExpertsPerCta; + int const k_warp = warpId * (kChunksPerWarp * kChunk); + + __shared__ float sm_red[kNumWarps][kTilesN][32][4]; + + float acc[kTilesN][4] = {}; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); +#endif + + __nv_bfloat16 const* a_row0 = mat_a + g * kHiddenDim; + __nv_bfloat16 const* a_row1 = mat_a + (g + 8) * kHiddenDim; + __nv_bfloat16 const* weight = mat_b; + int weightColumn = n_base; + if constexpr (kLatentDim > 0) + { + if (n_base >= kNumExperts) + { + weight = downWeight; + weightColumn -= kNumExperts; + } + } + bool const row0_valid = g < kNumTokens; + bool const row1_valid = kTwoRows && (g + 8 < kNumTokens); + uint4 const zero4 = make_uint4(0u, 0u, 0u, 0u); + +#pragma unroll + for (int batch = 0; batch < kChunksPerWarp; batch += kBatch) + { + uint4 a0[kBatch]; + uint4 a1[kBatch]; + uint4 b[kBatch][kTilesN]; +#pragma unroll + for (int c = 0; c < kBatch; c++) + { + int const k = k_warp + (batch + c) * kChunk + q * 8; + a0[c] = row0_valid ? *reinterpret_cast(a_row0 + k) : zero4; + a1[c] = row1_valid ? *reinterpret_cast(a_row1 + k) : zero4; +#pragma unroll + for (int t = 0; t < kTilesN; t++) + { + b[c][t] = *reinterpret_cast( + weight + static_cast(weightColumn + t * 8 + g) * kHiddenDim + k); + } + } +#pragma unroll + for (int c = 0; c < kBatch; c++) + { + uint32_t const afrag0[4] = {a0[c].x, a1[c].x, a0[c].y, a1[c].y}; + uint32_t const afrag1[4] = {a0[c].z, a1[c].z, a0[c].w, a1[c].w}; +#pragma unroll + for (int t = 0; t < kTilesN; t++) + { + uint32_t const bfrag0[2] = {b[c][t].x, b[c][t].y}; + uint32_t const bfrag1[2] = {b[c][t].z, b[c][t].w}; + mma_bf16_m16n8k16(acc[t], afrag0, bfrag0); + mma_bf16_m16n8k16(acc[t], afrag1, bfrag1); + } + } + } + +#pragma unroll + for (int t = 0; t < kTilesN; t++) + { +#pragma unroll + for (int i = 0; i < 4; i++) + { + sm_red[warpId][t][lane][i] = acc[t][i]; + } + } + __syncthreads(); + + // One thread per (tile, lane slot, fragment element): sum the warp partials in warp order and scatter. + constexpr int kNumOutputs = kTilesN * 32 * 4; + static_assert(kNumOutputs <= kBlockSize, "need one thread per output element"); + if (threadIdx.x < kNumOutputs) + { + int const t = threadIdx.x / 128; + int const l = (threadIdx.x % 128) / 4; + int const i = threadIdx.x % 4; + float sum = 0.0f; +#pragma unroll + for (int w = 0; w < kNumWarps; w++) + { + sum += sm_red[w][t][l][i]; + } + int const row = (l / 4) + (i >= 2 ? 8 : 0); + int const col = n_base + t * 8 + (l % 4) * 2 + (i & 1); + if (row < kNumTokens) + { + if constexpr (kLatentDim > 0) + { + if (n_base >= kNumExperts) + { + projection[row * kLatentDim + col - kNumExperts] = __float2bfloat16_rn(sum); + } + else + { + out[row * kNumExperts + col] = sum; + } + } + else + { + out[row * kNumExperts + col] = sum; + } + } + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +template +__global__ __launch_bounds__(kBlockSize, 1) void router_gemm_mma_kernel( + float* out, __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b) +{ + routerGemmMma(out, mat_a, mat_b, nullptr, nullptr); +} + +template +__global__ __launch_bounds__(448, 1) void router_latent_gemm_mma_kernel(float* logits, __nv_bfloat16* projection, + __nv_bfloat16 const* input, __nv_bfloat16 const* gateWeight, __nv_bfloat16 const* downWeight) +{ + routerGemmMma(logits, input, gateWeight, projection, downWeight); +} + +template +void invokeRouterGemmMma(float* output, __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b, cudaStream_t stream) +{ + // 8 expert columns per CTA and 14 warps splitting K: for 896 x 7168 that is 112 CTAs (one wave on B200) + // each doing 16 k-chunks per warp. Measured B200 (graph replay, L2-warm) vs the cuBLAS fallback: + // M=1 2.79 vs 7.36 us, M=8 3.77 vs 7.62 us, M=16 4.70 vs 6.64 us. + constexpr int kExpertsPerCta = 8; + constexpr int kBlockSize = 448; + static_assert(kNumExperts % kExpertsPerCta == 0, "kNumExperts must be a multiple of kExpertsPerCta"); + cudaLaunchConfig_t config; + config.gridDim = kNumExperts / kExpertsPerCta; + config.blockDim = kBlockSize; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, + router_gemm_mma_kernel, output, mat_a, mat_b)); +} + +template +void invokeRouterLatentGemmMma(float* logits, __nv_bfloat16* projection, __nv_bfloat16 const* input, + __nv_bfloat16 const* gateWeight, __nv_bfloat16 const* downWeight, cudaStream_t stream) +{ + constexpr int kColumns = 896 + 3584; + constexpr int kColumnsPerCta = 8; + constexpr int kBlockSize = 448; + cudaLaunchConfig_t config{}; + config.gridDim = kColumns / kColumnsPerCta; + config.blockDim = kBlockSize; + config.stream = stream; + cudaLaunchAttribute attr{}; + attr.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr.val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = &attr; + TLLM_CUDA_CHECK(cudaLaunchKernelEx( + &config, router_latent_gemm_mma_kernel, logits, projection, input, gateWeight, downWeight)); +} + template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 1, 256, 7168>( float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); @@ -340,6 +551,49 @@ template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__n template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 16, 256, 4096>( float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); +// num_experts=896, hidden_dim=7168 instantiations (Kimi K3), tensor-core variant. +#define INSTANTIATE_ROUTER_GEMM_MMA_K3(kNumTokens) \ + template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemmMma( \ + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); +INSTANTIATE_ROUTER_GEMM_MMA_K3(1) +INSTANTIATE_ROUTER_GEMM_MMA_K3(2) +INSTANTIATE_ROUTER_GEMM_MMA_K3(3) +INSTANTIATE_ROUTER_GEMM_MMA_K3(4) +INSTANTIATE_ROUTER_GEMM_MMA_K3(5) +INSTANTIATE_ROUTER_GEMM_MMA_K3(6) +INSTANTIATE_ROUTER_GEMM_MMA_K3(7) +INSTANTIATE_ROUTER_GEMM_MMA_K3(8) +INSTANTIATE_ROUTER_GEMM_MMA_K3(9) +INSTANTIATE_ROUTER_GEMM_MMA_K3(10) +INSTANTIATE_ROUTER_GEMM_MMA_K3(11) +INSTANTIATE_ROUTER_GEMM_MMA_K3(12) +INSTANTIATE_ROUTER_GEMM_MMA_K3(13) +INSTANTIATE_ROUTER_GEMM_MMA_K3(14) +INSTANTIATE_ROUTER_GEMM_MMA_K3(15) +INSTANTIATE_ROUTER_GEMM_MMA_K3(16) +#undef INSTANTIATE_ROUTER_GEMM_MMA_K3 + +#define INSTANTIATE_ROUTER_LATENT_GEMM_MMA(kNumTokens) \ + template void invokeRouterLatentGemmMma( \ + float*, __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(1) +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(2) +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(3) +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(4) +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(5) +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(6) +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(7) +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(8) +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(9) +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(10) +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(11) +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(12) +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(13) +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(14) +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(15) +INSTANTIATE_ROUTER_LATENT_GEMM_MMA(16) +#undef INSTANTIATE_ROUTER_LATENT_GEMM_MMA + } // namespace kernels::dsv3MinLatencyKernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.h b/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.h index ffd77cf12a40..3aef4bdec6f9 100644 --- a/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.h +++ b/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -19,6 +19,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaUtils.h" #include +#include #include #include @@ -30,6 +31,16 @@ namespace kernels::dsv3MinLatencyKernels template void invokeRouterGemm(float* output, T const* mat_a, T const* mat_b, cudaStream_t stream); +// Tensor-core variant (mma.sync m16n8k16, bf16 in / fp32 out) for large expert counts such as Kimi K3's 896, +// where the scalar FFMA kernel above is instruction-bound. kNumTokens <= 16, kNumExperts % 8 == 0. +template +void invokeRouterGemmMma(float* output, __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b, cudaStream_t stream); + +// BF16 input [M,7168], gate [896,7168], down [3584,7168]; FP32 logits and BF16 projection, M <= 16. +template +void invokeRouterLatentGemmMma(float* logits, __nv_bfloat16* projection, __nv_bfloat16 const* input, + __nv_bfloat16 const* gateWeight, __nv_bfloat16 const* downWeight, cudaStream_t stream); + } // namespace kernels::dsv3MinLatencyKernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingCustomKernels.cuh b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingCustomKernels.cuh index d4c40f0aec35..87f053193ec7 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingCustomKernels.cuh +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingCustomKernels.cuh @@ -840,8 +840,13 @@ void launchCoopBlockKernel(Data const& data, uint32_t numThreadsHist, void* stre // //////////////////////////////////////////////////////////////////////////////////////////////////// +// launchDynBlockKernel launches min(max(tokens*32, MaxExperts), 1024) threads, i.e. 1024 for the +// 1024-expert tier now reachable from the post-topK path. Without a bound the 1024-expert +// instantiation compiled to 89 regs/thread -> 89*1024 > 64K -> cudaLaunchKernelEx "too many +// resources requested for launch". The <=576-expert tiers already use <=62 regs, so the bound only +// affects the large tiers. template -__global__ void routingIndicesDynBlockKernel(KernelParams params) +__global__ void __launch_bounds__(1024) routingIndicesDynBlockKernel(KernelParams params) { using OutputT = typename KernelParams::OutputT; using InputT = typename KernelParams::InputT; @@ -1734,6 +1739,20 @@ void launchClusterKernelBlockDim1024(Data const& data, void* stream) #if defined(TRTLLM_ROUTING_CUSTOM_ENTRY) void launchClusterKernel(Data const& data, void* stream) { + // Post-topK (mPtrScores == nullptr) only permutes: the cluster kernel is thread-per-expanded-index + // there, so the 512-thread variant is not bound by the warp-per-token limit below (capacity + // NumBlocksPerCluster * 512 = 4096 tokens). Measured on B200 (896 experts/top-16, E=1024 tier), + // same GPU: 512 threads (2 experts/thread) 4.3/4.6/4.8/5.1/5.4/6.1 us at 17/32/64/224/512/1024 + // tokens vs 256 threads (4 experts/thread) 5.2/5.7/6.2 us at 16/32/64 and 1024 threads + // (1 expert/thread) ~5.9 us at 224-512. At 4096 tokens the 512-thread variant degrades to ~9.8 us + // (16 expanded indices per thread), so beyond 1024 tokens keep the 1024-thread launch. + static constexpr int PostTopKClusterBlockDim512MaxNumTokens = 1024; + if (data.mPtrScores == nullptr && data.mNumTokens <= PostTopKClusterBlockDim512MaxNumTokens) + { + launchClusterKernelBlockDim512(data, stream); + return; + } + // Each warp owns one token, so the reduced-thread cluster variants have lower token capacity. // Use them only where the requested token count fits; otherwise keep the original 1024-thread launch. if (data.mNumTokens <= MaxNumTokensClusterScores256) diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingCustomPolicy.cuh b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingCustomPolicy.cuh index a75bab938b03..6dc9938e44cf 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingCustomPolicy.cuh +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingCustomPolicy.cuh @@ -587,6 +587,8 @@ struct PolicyTraits Tier<512, 16>, // Various 512-expert models with high topK Tier<512, 22>, // Nemotron Super V3 (512 experts, topK=22) Tier<576, 8>, // Customized model with 576 experts + Tier<1024, 32>, // Kimi K3 post-topK permutation (896 experts, topK 16): without this tier the + // >128-token cluster launch and the dyn-block dispatch fell to <2048,32> Tier<2048, 32> // Large-expert fallback >; }; diff --git a/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFromTopKIds.cu b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFromTopKIds.cu index 96d7c4f9ed9c..434ee5dad6a7 100644 --- a/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFromTopKIds.cu +++ b/cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFromTopKIds.cu @@ -74,8 +74,18 @@ void runPostTopKPipeline(DataType const& data, uint32_t /*numThreadsHist*/, void static int const smMajor = tensorrt_llm::common::getSMVersion() / 10; bool const useStaticBlock = data.mNumTokens <= routingCustom::BlockKernelMaxNumTokens; int32_t const dispatchedMaxExperts = routingCustom::queryDispatchedMaxExperts(customData); - bool const useDynBlock = !useStaticBlock && data.mNumTokens <= routingCustom::DynBlockKernelMaxNumTokens - && dispatchedMaxExperts <= routingCustom::DynBlockKernelMaxNumExperts; + // Post-topK only permutes (no in-kernel TopK), so the single-CTA dyn-block kernel is worth using + // for larger expert tiers than in the scores path: the 8-CTA cluster kernel pays a DSMEM gather + + // cluster-barrier latency floor regardless of token count. Measured on B200 (896 experts/top-16, + // E=1024 tier): dyn-block 3.1/3.3/4.2 us at 5/8/16 tokens vs the 256-thread cluster kernel + // 4.7/4.7/5.2 us; but the dyn-block per-expert loops run over all tokens (~0.11 us/token), while + // the post-topK 512-thread cluster variant (see launchClusterKernel) is ~4.3 us at 17 tokens and + // nearly flat -> crossover ~12 tokens (dyn-block 4.7 us vs cluster 4.3 us already at 16-17). + // smem = tokens * MaxNumExperts * 3 B (int8 k-idx + int16 offset) = 36 KB at 12 x 1024. + static constexpr int PostTopKDynBlockMaxNumTokens = 12; + static constexpr int PostTopKDynBlockMaxNumExperts = 1024; + bool const useDynBlock = !useStaticBlock && data.mNumTokens <= PostTopKDynBlockMaxNumTokens + && dispatchedMaxExperts <= PostTopKDynBlockMaxNumExperts; // runPostTopKPipeline only handles pre-computed topK (mPtrTopKIds or mPtrTopKPacked), // never raw scores. The cluster kernel's routingPermutation uses thread-per-expanded-index diff --git a/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu b/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu index e3bdd0d89645..78abbe667cfb 100644 --- a/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu +++ b/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu @@ -24,6 +24,7 @@ #include #include #include +#include namespace cg = cooperative_groups; using namespace tensorrt_llm::common; @@ -46,8 +47,55 @@ static __device__ inline float sigmoid_accurate(float x) return 0.5f * tanhf(0.5f * x) + 0.5f; } +// Warp-wide max of a 32-bit unsigned value: one redux.sync on SM100+, shuffle tree elsewhere. +static __device__ __forceinline__ uint32_t warpMaxU32(cg::thread_block_tile const& warp, uint32_t v) +{ + if constexpr (reduce_topk::kTLLM_GEN_HAS_FAST_REDUX) + { + uint32_t r; + asm("redux.sync.max.u32 %0, %1, 0xffffffff;\n" : "=r"(r) : "r"(v)); + return r; + } + else + { + return cg::reduce(warp, v, cg::greater{}); + } +} + +// One arg-max round over the warp's key slots (keys as produced by TopKRedType::makeCmpVal: +// high word = order-preserving value bits, low 16 bits = 65535 - expert index, so "larger key" means +// "larger value, then smaller index" — the same total order the iterated reduceTopK used). The winner +// is removed from the slot that held it (slot := 0, below every real key). Two 32-bit warp maxima +// replace the previous 64-bit shuffle reduction. +template +static __device__ __forceinline__ uint64_t warpArgMaxRound( + cg::thread_block_tile const& warp, uint64_t (&slot)[N]) +{ + uint64_t best = slot[0]; +#pragma unroll + for (int ii = 1; ii < N; ++ii) + { + best = slot[ii] > best ? slot[ii] : best; + } + uint32_t const bestHi = static_cast(best >> 32); + uint32_t const hi = warpMaxU32(warp, bestHi); + uint32_t const lo = warpMaxU32(warp, bestHi == hi ? static_cast(best & 0xFFFFu) : 0u); + uint64_t const winner = (static_cast(hi) << 32) | lo; +#pragma unroll + for (int ii = 0; ii < N; ++ii) + { + slot[ii] = slot[ii] == winner ? uint64_t{0} : slot[ii]; + } + return winner; +} + +// SmallBatch selects the ungrouped implementation: true = two-level parallel top-k (all warps, redux.sync) +// that minimizes per-CTA latency for decode batches; false = the original single-warp iterated +// reduceTopK, which issues fewer instructions and wins once thousands of CTAs saturate the GPU +// (prefill). Both produce bit-identical results; invokeNoAuxTc picks by token count. template + int MaxNumTopExperts = DefaultMaxNumTopExperts, int MaxNumTopGroups = DefaultMaxNumTopGroups, + bool SmallBatch = true> __device__ __forceinline__ void deepseek_v3_topk_block(InputT* scores, OutputT* topkValues, IdxT* topkIndices, BiasT* routingBias, int64_t const numTokens, int64_t const numGroup, int64_t const topkGroup, int64_t const topk, int64_t const numExperts, int64_t const numExpertsPerGroup, double const routedScalingFactor, int64_t tokenIdx) @@ -137,8 +185,9 @@ __device__ __forceinline__ void deepseek_v3_topk_block(InputT* scores, OutputT* } } } - else + else if constexpr (!SmallBatch) { + // Original ungrouped implementation (throughput regime): warp 0 runs the iterated reduceTopK. for (int e = threadIdx.x; e < numExperts; e += blockDim.x) { auto scoreIdx = tokenIdx * static_cast(numExperts) + e; @@ -180,10 +229,135 @@ __device__ __forceinline__ void deepseek_v3_topk_block(InputT* scores, OutputT* } } } + else + { + // Ungrouped (n_group == 1, e.g. Kimi K3: 896 experts / top-16): two-level parallel top-k. + // + // The previous implementation had warp 0 alone run Sort per lane followed by + // `topk` serial 64-bit shuffle reductions while the other warps idled, and its phase 1 issued + // the per-thread global loads one dependent iteration at a time -- ~6.8 us per CTA for the + // 1024-expert tier on B200; with one CTA per token that latency is fully exposed at small + // decode batches. Here: + // phase 1: every thread prefetches all its scores/bias (one round of load latency), then + // writes the sigmoid and the packed (value, index) key of reduce_topk (a strict total + // order: higher value first, then lower index); + // phase 2: each 128-expert batch is handled by one warp: `topk` arg-max rounds, each two + // 32-bit redux.sync maxima (value word, then index word) instead of a 64-bit shuffle + // tree -- all warps work in parallel; + // phase 3: warp 0 runs the same rounds over the NumBatches*topk candidates. + // Every round removes exactly the maximum key under the same total order the old iterated + // reduceTopK used, so the selected experts and their order are identical (bit-exact). + using RedType = reduce_topk::TopKRedType; + using KeyT = typename RedType::TypeCmp; + static_assert(std::is_same_v, "float keys are 64-bit"); + static constexpr int ChunksPerBatch = 4; + static constexpr int BatchSize = ChunksPerBatch * WARP_SIZE; // 128 experts per batch + static_assert(MaxNumExperts % BatchSize == 0, "MaxNumExperts must be a multiple of 128"); + static_assert(BatchSize >= MaxNumTopExperts, "a batch must hold at least topk experts"); + static constexpr int NumBatches = MaxNumExperts / BatchSize; + static constexpr int MaxCandPerLane = (NumBatches * MaxNumTopExperts + WARP_SIZE - 1) / WARP_SIZE; + // All launchers use blockDim >= 128 (128/256 threads, 448 in the K3 fused kernel). + static constexpr int MinBlockDim = 128; + static constexpr int MaxItersPhase1 = MaxNumExperts / MinBlockDim; + __shared__ KeyT __attribute((aligned(128))) smemKey[MaxNumExperts]; + __shared__ KeyT __attribute((aligned(128))) smemCand[NumBatches * MaxNumTopExperts]; + + int32_t const numWarps = blockDim.x / WARP_SIZE; + + // Phase 1: prefetch, then sigmoid + bias packed into a sortable key. Padding experts get -inf + // (never selected unless numExperts < topk, matching the previous behaviour). + float scoreArr[MaxItersPhase1]; + float biasArr[MaxItersPhase1]; +#pragma unroll + for (int ii = 0; ii < MaxItersPhase1; ++ii) + { + int const e = threadIdx.x + ii * blockDim.x; + if (e < numExperts) + { + scoreArr[ii] = static_cast(scores[tokenIdx * static_cast(numExperts) + e]); + biasArr[ii] = static_cast(routingBias[e]); + } + } +#pragma unroll + for (int ii = 0; ii < MaxItersPhase1; ++ii) + { + int const e = threadIdx.x + ii * blockDim.x; + if (e < MaxNumExperts) + { + float scoreBias = invalidScoreFloat; + float scoreSigmoid = 0.F; + if (e < numExperts) + { + scoreSigmoid = sigmoid_accurate(scoreArr[ii]); + scoreBias = scoreSigmoid + biasArr[ii]; + } + smemScoreSigmoid[e] = scoreSigmoid; + smemKey[e] = RedType::makeCmpVal(scoreBias, e); + } + } + __syncthreads(); + + // Phase 2: per-batch top-k candidates (one warp per batch, warps loop over batches). + for (int b = warpIdx; b < NumBatches; b += numWarps) + { + KeyT slot[ChunksPerBatch]; +#pragma unroll + for (int ii = 0; ii < ChunksPerBatch; ++ii) + { + slot[ii] = smemKey[b * BatchSize + ii * WARP_SIZE + laneIdx]; + } + for (int kk = 0; kk < topk; ++kk) + { + KeyT const winner = warpArgMaxRound(warp, slot); + if (laneIdx == kk) + { + smemCand[b * topk + kk] = winner; + } + } + } + __syncthreads(); + + // Phase 3 + finalize (warp 0): top-k over the NumBatches * topk candidates; lane kk keeps the + // kk-th largest expert, then the normalization is the same arithmetic as before. + if (warpIdx == 0) + { + int32_t const numCand = NumBatches * static_cast(topk); + KeyT slot[MaxCandPerLane]; +#pragma unroll + for (int ii = 0; ii < MaxCandPerLane; ++ii) + { + int const c = ii * WARP_SIZE + laneIdx; + slot[ii] = c < numCand ? smemCand[c] : KeyT{0}; + } + KeyT myOut = 0; + for (int kk = 0; kk < topk; ++kk) + { + KeyT const winner = warpArgMaxRound(warp, slot); + myOut = laneIdx == kk ? winner : myOut; + } + + int32_t expertIdx = MaxNumExperts - 1; + float scoreNorm = 0.F; + if (laneIdx < topk) + { + float unusedScore; + RedType::unpack(unusedScore, expertIdx, myOut); + scoreNorm = smemScoreSigmoid[expertIdx]; + } + auto redNorm = cg::reduce(warp, scoreNorm, cg::plus{}); + auto finalScore = static_cast(scoreNorm * routedScalingFactor / (redNorm + 1e-20)); + if (laneIdx < topk) + { + topkValues[laneIdx] = static_cast(finalScore); + topkIndices[laneIdx] = expertIdx; + } + } + } } template + int MaxNumTopExperts = DefaultMaxNumTopExperts, int MaxNumTopGroups = DefaultMaxNumTopGroups, + bool SmallBatch = true> __global__ void deepseek_v3_topk_kernel(InputT* scores, OutputT* topkValues, IdxT* topkIndices, BiasT* routingBias, int64_t const numTokens, int64_t const numGroup, int64_t const topkGroup, int64_t const topk, int64_t const numExperts, int64_t const numExpertsPerGroup, double const routedScalingFactor) @@ -198,8 +372,8 @@ __global__ void deepseek_v3_topk_kernel(InputT* scores, OutputT* topkValues, Idx return; } } - deepseek_v3_topk_block( - scores, topkValues, topkIndices, routingBias, numTokens, numGroup, topkGroup, topk, numExperts, + deepseek_v3_topk_block(scores, topkValues, topkIndices, routingBias, numTokens, numGroup, topkGroup, topk, numExperts, numExpertsPerGroup, routedScalingFactor, blockIdx.x); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); @@ -314,28 +488,48 @@ void invokeNoAuxTc(InputT* scores, BiasT* bias, OutputT* topk_values, IdxT* topk } else if (is_single_group) { + // Small batches (decode) use the latency-optimized two-level top-k; large batches keep the + // original single-warp implementation (fewer instructions once the GPU is saturated). + // Measured on B200 (148 SMs), 896 experts / top-16, kernel-only old -> new: 8 tok 6.84 -> 5.12 us, + // 128 tok 8.23 -> 5.70, 192-256 tok -3 %, 384 tok +9 %, 512 tok 10.9 -> 12.3, 8192 tok 94 -> 109. + // The break is where CTAs start sharing SMs, so the default crossover is 256 tokens. + // TLLM_NOAUX_TC_SMALL_BATCH_MAX_TOKENS overrides it (0 = always the original path). + static int const smallBatchMaxTokens = [] + { + char const* env = std::getenv("TLLM_NOAUX_TC_SMALL_BATCH_MAX_TOKENS"); + return env != nullptr ? std::atoi(env) : 256; + }(); + bool const smallBatch = num_tokens <= smallBatchMaxTokens; if (num_experts <= 128) { - kernel_instance - = &deepseek_v3_topk_kernel; + kernel_instance = smallBatch ? &deepseek_v3_topk_kernel + : &deepseek_v3_topk_kernel; num_threads = 128; } else if (num_experts <= 256) { - kernel_instance - = &deepseek_v3_topk_kernel; + kernel_instance = smallBatch ? &deepseek_v3_topk_kernel + : &deepseek_v3_topk_kernel; num_threads = 256; } else if (num_experts <= 512) { - kernel_instance - = &deepseek_v3_topk_kernel; + kernel_instance = smallBatch ? &deepseek_v3_topk_kernel + : &deepseek_v3_topk_kernel; num_threads = 256; } else { - kernel_instance - = &deepseek_v3_topk_kernel; + kernel_instance = smallBatch ? &deepseek_v3_topk_kernel + : &deepseek_v3_topk_kernel; num_threads = 256; } } diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index acde3a3d3cd2..d587ce255428 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -2024,6 +2024,7 @@ MnnvlFusionPattern getMnnvlFusionPattern(AllReduceFusionOp fusionOp) switch (fusionOp) { case AllReduceFusionOp::NONE: return MnnvlFusionPattern::kAllReduce; + case AllReduceFusionOp::RMS_NORM: return MnnvlFusionPattern::kARRMSNorm; case AllReduceFusionOp::RESIDUAL_RMS_NORM: return MnnvlFusionPattern::kARResidualRMSNorm; case AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_FP8: return MnnvlFusionPattern::kARResidualRMSNormFP8Quant; case AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_NVFP4: return MnnvlFusionPattern::kARResidualRMSNormFP4Quant; @@ -2052,7 +2053,7 @@ bool isMnnvlNvfp4FusionOp(AllReduceFusionOp fusionOp) bool hasMnnvlNormOutput(AllReduceFusionOp fusionOp) { - return fusionOp == AllReduceFusionOp::RESIDUAL_RMS_NORM + return fusionOp == AllReduceFusionOp::RMS_NORM || fusionOp == AllReduceFusionOp::RESIDUAL_RMS_NORM || fusionOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_OUT_QUANT_FP8 || fusionOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4; } @@ -2096,12 +2097,21 @@ std::vector mnnvlFusionAllReduce(torch::Tensor& input, torch::opt if (hasRmsNormFusion) { - TORCH_CHECK(residual_in.has_value() && gamma.has_value() && epsilon.has_value(), - "[mnnvlFusionAllReduce] residual_in, gamma, and epsilon must be provided for RMSNorm fusion"); - TORCH_CHECK(residual_in.value().is_contiguous(), "[mnnvlFusionAllReduce] residual_in must be contiguous"); + // kARRMSNorm has no residual operand (kHasResidual/kHasResidualOut + // are both false), so residual_in is optional for it alone. + bool const needsResidual = fusionOp != AllReduceFusionOp::RMS_NORM; + TORCH_CHECK(gamma.has_value() && epsilon.has_value(), + "[mnnvlFusionAllReduce] gamma and epsilon must be provided for RMSNorm fusion"); + TORCH_CHECK(!needsResidual || residual_in.has_value(), + "[mnnvlFusionAllReduce] residual_in must be provided for residual RMSNorm fusion"); + TORCH_CHECK(!residual_in.has_value() || residual_in.value().is_contiguous(), + "[mnnvlFusionAllReduce] residual_in must be contiguous"); TORCH_CHECK(gamma.value().is_contiguous(), "[mnnvlFusionAllReduce] gamma must be contiguous"); - residualOut = torch::empty_like(residual_in.value()); + if (needsResidual) + { + residualOut = torch::empty_like(residual_in.value()); + } if (hasMnnvlNormOutput(fusionOp)) { normOut = torch::empty_like(input); @@ -2169,12 +2179,12 @@ std::vector mnnvlFusionAllReduce(torch::Tensor& input, torch::opt { allreduce_params.output = normOut.defined() ? normOut.mutable_data_ptr() : nullptr; - allreduce_params.residualIn = residual_in.value().const_data_ptr(); + allreduce_params.residualIn = residual_in.has_value() ? residual_in.value().const_data_ptr() : nullptr; allreduce_params.gamma = gamma.value().const_data_ptr(); allreduce_params.epsilon = static_cast(epsilon.value()); allreduce_params.rmsNormFusion = true; - allreduce_params.residualOut = residualOut.mutable_data_ptr(); + allreduce_params.residualOut = residualOut.defined() ? residualOut.mutable_data_ptr() : nullptr; } if (isMnnvlQuantFusionOp(fusionOp)) @@ -2203,6 +2213,7 @@ std::vector mnnvlFusionAllReduce(torch::Tensor& input, torch::opt switch (fusionOp) { case AllReduceFusionOp::NONE: return {reduceOut}; + case AllReduceFusionOp::RMS_NORM: return {normOut}; case AllReduceFusionOp::RESIDUAL_RMS_NORM: return {normOut, residualOut}; case AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_FP8: return {quantOut, residualOut}; case AllReduceFusionOp::RESIDUAL_RMS_NORM_OUT_QUANT_FP8: return {normOut, quantOut, residualOut}; diff --git a/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp b/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp index d7b324598bd7..356d65816ef0 100644 --- a/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp +++ b/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp @@ -19,6 +19,9 @@ #include "tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.h" #include "tensorrt_llm/runtime/torchUtils.h" #include "tensorrt_llm/thop/cublasScaledMM.h" +#include +#include +#include namespace th = torch; namespace tl = tensorrt_llm; @@ -31,7 +34,24 @@ namespace torch_ext namespace { -template +// kUseMma selects the tensor-core kernel (invokeRouterGemmMma) instead of the scalar one. +template +void invokeRouterGemmForTokens( + float* output, __nv_bfloat16 const* input, __nv_bfloat16 const* weights, cudaStream_t stream) +{ + if constexpr (kUseMma) + { + tk::dsv3MinLatencyKernels::invokeRouterGemmMma( + output, input, weights, stream); + } + else + { + tk::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, kNumTokens, kNumExperts, kHiddenDim>( + output, input, weights, stream); + } +} + +template struct LoopUnroller { static void unroll( @@ -39,26 +59,25 @@ struct LoopUnroller { if (num_tokens == kBegin) { - tk::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, kBegin, kNumExperts, kHiddenDim>( - output, input, weights, stream); + invokeRouterGemmForTokens(output, input, weights, stream); } else { - LoopUnroller::unroll(num_tokens, output, input, weights, stream); + LoopUnroller::unroll( + num_tokens, output, input, weights, stream); } } }; -template -struct LoopUnroller +template +struct LoopUnroller { static void unroll( int num_tokens, float* output, __nv_bfloat16 const* input, __nv_bfloat16 const* weights, cudaStream_t stream) { if (num_tokens == kEnd) { - tk::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>( - output, input, weights, stream); + invokeRouterGemmForTokens(output, input, weights, stream); } else { @@ -66,8 +85,62 @@ struct LoopUnroller } } }; + +template +void invokeRouterLatentForTokens(int numTokens, float* logits, __nv_bfloat16* projection, __nv_bfloat16 const* input, + __nv_bfloat16 const* gateWeight, __nv_bfloat16 const* downWeight, cudaStream_t stream) +{ + if (numTokens == kNumTokens) + { + tk::dsv3MinLatencyKernels::invokeRouterLatentGemmMma( + logits, projection, input, gateWeight, downWeight, stream); + } + else if constexpr (kNumTokens < 16) + { + invokeRouterLatentForTokens( + numTokens, logits, projection, input, gateWeight, downWeight, stream); + } + else + { + TORCH_CHECK(false, "router + latent GEMM requires 1 <= M <= 16"); + } +} } // namespace +std::tuple dsv3_router_latent_gemm_op( + th::Tensor const& input, th::Tensor const& gateWeight, th::Tensor const& downWeight) +{ + constexpr int kHiddenDim = 7168; + constexpr int kNumExperts = 896; + constexpr int kLatentDim = 3584; + constexpr int kLoadAlignment = 16; + for (auto const& tensor : {input, gateWeight, downWeight}) + { + TORCH_CHECK(tensor.is_cuda() && tensor.device() == input.device(), "inputs must be on the same CUDA device"); + TORCH_CHECK(tensor.dim() == 2 && tensor.scalar_type() == th::kBFloat16, "inputs must be 2-D BF16 tensors"); + TORCH_CHECK(tensor.is_contiguous() && tensor.stride(0) == kHiddenDim && tensor.stride(1) == 1, + "inputs must have contiguous [N,7168] row-major storage"); + TORCH_CHECK(reinterpret_cast(tensor.data_ptr()) % kLoadAlignment == 0, + "inputs must have 16-byte-aligned storage"); + } + TORCH_CHECK(input.size(0) >= 1 && input.size(0) <= 16 && input.size(1) == kHiddenDim, + "input must have shape [M,7168], 1 <= M <= 16"); + TORCH_CHECK(gateWeight.size(0) == kNumExperts && gateWeight.size(1) == kHiddenDim, + "gate weight must have shape [896,7168]"); + TORCH_CHECK(downWeight.size(0) == kLatentDim && downWeight.size(1) == kHiddenDim, + "down weight must have shape [3584,7168]"); + + c10::cuda::CUDAGuard const deviceGuard(input.device()); + auto logits = th::empty({input.size(0), kNumExperts}, input.options().dtype(th::kFloat32)); + auto projection = th::empty({input.size(0), kLatentDim}, input.options()); + auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); + invokeRouterLatentForTokens(static_cast(input.size(0)), static_cast(logits.mutable_data_ptr()), + static_cast<__nv_bfloat16*>(projection.mutable_data_ptr()), static_cast<__nv_bfloat16 const*>(input.data_ptr()), + static_cast<__nv_bfloat16 const*>(gateWeight.data_ptr()), + static_cast<__nv_bfloat16 const*>(downWeight.data_ptr()), stream); + return {logits, projection}; +} + th::Tensor dsv3_router_gemm_op(th::Tensor const& mat_a, th::Tensor const& mat_b, std::optional const& bias, std::optional const& out_dtype) { @@ -77,7 +150,8 @@ th::Tensor dsv3_router_gemm_op(th::Tensor const& mat_a, th::Tensor const& mat_b, auto const out_dtype_ = out_dtype.value_or(mat_a.scalar_type()); auto const data_type = mat_a.scalar_type(); constexpr int kNumExperts = 256; - constexpr int kHiddenDim7168 = 7168; // DeepSeek-V3 / DeepSeek-V3.2 + constexpr int kNumExpertsKimiK3 = 896; + constexpr int kHiddenDim7168 = 7168; // DeepSeek-V3 / DeepSeek-V3.2 / Kimi K3 constexpr int kHiddenDim6144 = 6144; // GLM-5 constexpr int kHiddenDim4096 = 4096; // DeepSeek-V4 std::vector output_size = {mat_a.sizes()[0], mat_b.sizes()[1]}; @@ -86,11 +160,18 @@ th::Tensor dsv3_router_gemm_op(th::Tensor const& mat_a, th::Tensor const& mat_b, TORCH_CHECK(mat_a.strides()[1] == 1 && out.strides()[1] == 1); // Row-major TORCH_CHECK(mat_b.strides()[0] == 1); // Column-major auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device()); - bool const shape_ok - = (num_tokens >= 1 && num_tokens <= 16 && num_experts == kNumExperts && mat_b.sizes()[0] == hidden_dim - && data_type == torch::kBFloat16 && out_dtype_ == torch::kFloat32 && !bias.has_value()); + bool const base_ok = (num_tokens >= 1 && num_tokens <= 16 && mat_b.sizes()[0] == hidden_dim + && data_type == torch::kBFloat16 && out_dtype_ == torch::kFloat32 && !bias.has_value()); + bool const shape_ok = base_ok && num_experts == kNumExperts; - if (shape_ok && hidden_dim == kHiddenDim7168) + if (base_ok && num_experts == kNumExpertsKimiK3 && hidden_dim == kHiddenDim7168) + { + // Kimi K3: tensor-core kernel; the cuBLAS fallback below costs ~7.4 us at M=8 on B200 (split-K + reduce). + LoopUnroller<1, 16, kNumExpertsKimiK3, kHiddenDim7168, /*kUseMma=*/true>::unroll(num_tokens, + reinterpret_cast(out.mutable_data_ptr()), reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream); + } + else if (shape_ok && hidden_dim == kHiddenDim7168) { LoopUnroller<1, 16, kNumExperts, kHiddenDim7168>::unroll(num_tokens, reinterpret_cast(out.mutable_data_ptr()), reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), @@ -123,9 +204,13 @@ TRTLLM_NAMESPACE_END TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def("dsv3_router_gemm_op(Tensor mat_a, Tensor mat_b, Tensor? bias, ScalarType? out_dtype) -> (Tensor out)"); + m.def( + "dsv3_router_latent_gemm_op(Tensor input, Tensor gate_weight, Tensor down_weight) -> (Tensor logits, Tensor " + "projection)"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("dsv3_router_gemm_op", &tensorrt_llm::torch_ext::dsv3_router_gemm_op); + m.impl("dsv3_router_latent_gemm_op", &tensorrt_llm::torch_ext::dsv3_router_latent_gemm_op); } diff --git a/examples/disaggregated/slurm/benchmark/submit_dwdp.py b/examples/disaggregated/slurm/benchmark/submit_dwdp.py index 7423f3e1de1f..bd3a9154b087 100644 --- a/examples/disaggregated/slurm/benchmark/submit_dwdp.py +++ b/examples/disaggregated/slurm/benchmark/submit_dwdp.py @@ -25,6 +25,7 @@ import argparse import glob import json +import math import os import shutil import subprocess @@ -160,7 +161,22 @@ def submit_dwdp_job(config, log_dir, dry_run): 2 * ctx_world_size * gen_world_size if benchmark_config["mode"] == "e2e" else 0 ) - total_nodes = ctx_nodes + gen_nodes + # DWDP's context workers are TP=1 by construction, so the default + # "each worker owns whole nodes, no node sharing" rule spends one node per + # worker and leaves 3 of every 4 GPUs idle: a dwdp_size=16 group asks for + # 16 nodes instead of 4. With --segment that also means demanding a whole + # NVLink block, which on a busy cluster simply never schedules. + # ``hardware.compact_packing`` mirrors the option submit.py already has + # (see its total_gpus/ceil branch); allocate_gpus' assign_server_compact + # walks a per-GPU cursor, which for a divisible layout lands each rank on + # exactly the (node, gpu) that SLURM block distribution would give it, so + # the block-distribution launcher below stays valid. + compact_packing = bool(hw_config.get("compact_packing", False)) + if compact_packing: + total_gpus = ctx_world_size * ctx_num + gen_world_size * gen_num + total_nodes = math.ceil(total_gpus / gpus_per_node) + else: + total_nodes = ctx_nodes + gen_nodes total_tasks = total_nodes * gpus_per_node dwdp_size = worker_config.get("ctx", {}).get("dwdp_config", {}).get("dwdp_size", 1) @@ -220,6 +236,7 @@ def submit_dwdp_job(config, log_dir, dry_run): num_ctx_servers=ctx_num, gen_world_size=gen_world_size, ctx_world_size=ctx_world_size, + compact_packing=compact_packing, ) with open(os.path.join(log_dir, "allocations.json"), "w") as f: json.dump(allocations, f, indent=2) diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index b828abb30a85..71f400ec5348 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -527,6 +527,9 @@ The FMHA package is split by role: `TrtllmAttention` can pair it with a later causal-generation provider through `CombinedFmha`. - `fmha/cute_dsl_mla.py` implements the CuTe DSL MLA decode FMHA library. + With an FP8 KV cache it passes `mla_bmm1_scale` / `mla_bmm2_scale` to the + kernel as device tensors (element 0 of each), so the scales need no host + read and stay valid under CUDA graph replay. - `fmha/prims_ts_block_sparse.py` adapts generic block-sparse requests to the vendored PrimTS contiguous and paged wrappers. Paged generation passes a live, zero-copy 2D K-page-table view with its TRT-LLM padded row stride; it diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index 8ab481750855..625cc75c6b12 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -18,6 +18,8 @@ from typing import TYPE_CHECKING, Optional import torch +import triton +import triton.language as tl from tensorrt_llm._torch.attention.backends.interface import ( AttentionForwardArgs, @@ -37,12 +39,173 @@ TrtllmAttentionMetadata, ) -_LOG2_E = math.log2(math.e) + +# Per-request split-KV plan for the CuTe DSL MLA decode kernels. The kernel grid +# carries the tactic's split_kv slots per request; the plan decides how many a +# request uses this step from the batch's KV lengths (the kernel clamps it to +# the grid and the slots beyond exit before their mainloop). +_SPLIT_PLAN_MAX_REQUESTS = 1024 +# CuteDSLNVMlaDecodeBlackwellRunner._MAX_SPLIT_KV: the largest grid any tactic uses. +_SPLIT_PLAN_MAX_SPLIT = 32 +# mma_qk_tiler N of every tactic: the KV tile a split is made of. +_SPLIT_PLAN_KV_TILE = 128 +# Shortest piece worth its partial-output write and reduction read. +_SPLIT_PLAN_MIN_TILES = 4 +# Fixed cost of one piece (cluster launch, Q load, pipeline fill, partial store +# and its reduction) in KV tiles, fitted on SM107 forced-split sweeps (H96, +# 14K-256K KV, B1-B16): with it the wave-quantized makespan picks the measured +# optimum within 5%. +_SPLIT_PLAN_PIECE_COST_TILES = 60 + + +@triton.jit( + do_not_specialize=["num_requests", "rows_per_request", "wave_clusters", "max_split"], + do_not_specialize_on_alignment=[ + "num_requests", + "rows_per_request", + "wave_clusters", + "max_split", + ], +) +def _mla_split_kv_plan_kernel( + cache_seqs_ptr, + plan_ptr, + num_requests, + rows_per_request, + wave_clusters, + max_split, + KV_TILE: tl.constexpr, + MIN_TILES: tl.constexpr, + PIECE_COST: tl.constexpr, + MAX_PIECES: tl.constexpr, + BLOCK: tl.constexpr, +): + offsets = tl.arange(0, BLOCK) + mask = offsets < num_requests + kv_len = tl.load(cache_seqs_ptr + offsets, mask=mask, other=0) + tiles = tl.maximum((kv_len + KV_TILE - 1) // KV_TILE, 1) + tiles = tl.where(mask, tiles, 0) + longest = tl.max(tiles, axis=0) + # Cut the longest request into k equal pieces (every request then takes + # ceil(tiles / piece) pieces) and keep the k whose wave-quantized makespan, + # waves x (piece + fixed piece cost), is smallest; ties go to fewer pieces. + best_piece = longest + best_cost = tl.full([], 2**31 - 1, tl.int32) + for pieces in tl.static_range(1, MAX_PIECES + 1): + piece = tl.maximum((longest + pieces - 1) // pieces, MIN_TILES) + splits = tl.minimum(tl.maximum((tiles + piece - 1) // piece, 1), max_split) + clusters = tl.sum(tl.where(mask, splits, 0), axis=0) * rows_per_request + waves = (clusters + wave_clusters - 1) // wave_clusters + cost = waves * (piece + PIECE_COST) + better = cost < best_cost + best_cost = tl.where(better, cost, best_cost) + best_piece = tl.where(better, piece, best_piece) + plan = tl.minimum(tl.maximum((tiles + best_piece - 1) // best_piece, 1), max_split) + tl.store(plan_ptr + offsets, plan, mask=mask) + + +def plan_mla_split_kv( + kv_lens: list[int], + *, + rows_per_request: int, + wave_clusters: int, + max_split: int = _SPLIT_PLAN_MAX_SPLIT, + kv_tile: int = _SPLIT_PLAN_KV_TILE, + min_tiles: int = _SPLIT_PLAN_MIN_TILES, + piece_cost: int = _SPLIT_PLAN_PIECE_COST_TILES, +) -> list[int]: + """Host mirror of ``_mla_split_kv_plan_kernel``: per-request split counts. + + ``rows_per_request`` is the kernel's packed 128-row query tiles per + request, each a 2-CTA cluster per split; ``wave_clusters`` the clusters one + wave holds. + The longest request is cut into ``k`` equal pieces (every request takes + ``ceil(tiles / piece)`` of them) for the ``k`` with the smallest + wave-quantized makespan ``waves * (piece + piece_cost)``. + """ + num_requests = len(kv_lens) + if num_requests <= 0 or rows_per_request <= 0 or wave_clusters <= 0 or max_split <= 0: + raise ValueError("split plan needs requests, rows, wave clusters and a split cap") + tiles = [max(math.ceil(max(kv, 0) / kv_tile), 1) for kv in kv_lens] + longest = max(tiles) + best_piece, best_cost = longest, None + for pieces in range(1, max_split + 1): + piece = max(math.ceil(longest / pieces), min_tiles) + splits = [min(max(math.ceil(t / piece), 1), max_split) for t in tiles] + waves = math.ceil(sum(splits) * rows_per_request / wave_clusters) + cost = waves * (piece + piece_cost) + if best_cost is None or cost < best_cost: + best_piece, best_cost = piece, cost + return [min(max(math.ceil(t / best_piece), 1), max_split) for t in tiles] + + +def launch_mla_split_kv_plan( + cache_seqs: torch.Tensor, + plan: torch.Tensor, + *, + num_requests: int, + rows_per_request: int, + wave_clusters: int, + max_split: int = _SPLIT_PLAN_MAX_SPLIT, +) -> None: + """Write the per-request split counts for ``cache_seqs`` into ``plan`` (one launch).""" + if not 0 < num_requests <= _SPLIT_PLAN_MAX_REQUESTS: + raise ValueError(f"split plan supports 1..{_SPLIT_PLAN_MAX_REQUESTS} requests") + for name, tensor in (("cache_seqs", cache_seqs), ("plan", plan)): + if ( + tensor.dtype != torch.int32 + or tensor.ndim != 1 + or tensor.stride(0) != 1 + or tensor.numel() < num_requests + or not tensor.is_cuda + ): + raise ValueError(f"{name} must be a contiguous CUDA int32 vector of >= {num_requests}") + _mla_split_kv_plan_kernel[(1,)]( + cache_seqs, + plan, + num_requests, + max(rows_per_request, 1), + max(wave_clusters, 1), + max(max_split, 1), + KV_TILE=_SPLIT_PLAN_KV_TILE, + MIN_TILES=_SPLIT_PLAN_MIN_TILES, + PIECE_COST=_SPLIT_PLAN_PIECE_COST_TILES, + MAX_PIECES=_SPLIT_PLAN_MAX_SPLIT, + BLOCK=_SPLIT_PLAN_MAX_REQUESTS, + num_warps=4, + ) class CuteDslMlaFmha(PhasedFmha): """Blackwell CuTe DSL FMHA library for decode-only MLA.""" + @staticmethod + def _get_fp8_scales( + q: torch.Tensor, + fwd: AttentionForwardArgs, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return the FP8 BMM1/BMM2 scale device tensors the kernel reads. + + The kernel reads element 0 of each tensor on the device, so the scales + never need a host read: not on the first call and not under CUDA graph + capture. + """ + softmax_scale = fwd.mla_bmm1_scale + output_scale = fwd.mla_bmm2_scale + if softmax_scale is None or output_scale is None: + raise RuntimeError( + "CuTe DSL FP8 MLA decode requires mla_bmm1_scale and mla_bmm2_scale device tensors." + ) + for name, scale in ( + ("mla_bmm1_scale", softmax_scale), + ("mla_bmm2_scale", output_scale), + ): + if scale.dtype != torch.float32 or scale.numel() < 1: + raise RuntimeError(f"{name} must contain at least one float32 value.") + if scale.device != q.device: + raise RuntimeError(f"{name} must be on {q.device}, got {scale.device}.") + return softmax_scale, output_scale + @classmethod def _is_available(cls, attn: "TrtllmAttention") -> bool: if not IS_CUTLASS_DSL_AVAILABLE: @@ -179,7 +342,7 @@ def _kernel_can_implement( mma_qk_tiler_mn, mma_pv_tiler_mn, 1, # split_kv - True, # is_persistent + False, # is_persistent True, # is_var_seq False, # is_var_split_kv page_size, @@ -334,15 +497,30 @@ def _is_supported_with_reason( # failure. That makes this a correctness carve-out rather than a perf # tradeoff, so it is decided here instead of being folded into # _PERF_MIN_BATCH_FP8, which stays a pure measured-win table. - # Deliberately narrow: single-token H=96 is already a measured table - # entry and still goes through the gate, and non-helix multi-token - # H=96 keeps falling back exactly as it does today. Helix with - # seq_len_q > 1 implies _helix_spec_tokens_valid -- the helix block - # above returns False otherwise. + # Single-token H=96 is already a measured table entry and still goes + # through the gate. Helix with seq_len_q > 1 implies + # _helix_spec_tokens_valid -- the helix block above returns False + # otherwise. helix_h96_verify_group = ( attn.num_heads == 96 and seq_len_q > 1 and meta.helix_position_offsets is not None ) - if not helix_h96_verify_group: + # The same failure without Helix: an FP8-KV multi-token verify group + # (e.g. DSpark under attention DP, where every rank keeps all 96 heads) + # is not in the FP8 perf table, and on SM100/SM103 the next library is + # FlashInfer TRTLLM-Gen, which rejects these head counts. The kernel + # packs (token, head) rows into 128-row M tiles, so any H <= 128 with a + # linear-chain causal verify is valid (can_implement below still has the + # final say). SM107 is left out on purpose: TRTLLM-Gen FMHA is + # unavailable there and such a request already reaches the default + # TRTLLM path. + fp8_wide_head_verify_group = ( + 64 < attn.num_heads < 128 + and seq_len_q > 1 + and meta.helix_position_offsets is None + and self._get_kernel_dtype(attn, q) == torch.float8_e4m3fn + and get_sm_version() in (100, 103) + ) + if not (helix_h96_verify_group or fp8_wide_head_verify_group): # Perf gate (NOT a correctness limit). favorable, reason = self._is_perf_favorable( attn.num_heads, @@ -496,19 +674,11 @@ def _run_mla_decode( softmax_scale = float(1.0 / (math.sqrt(qk_nope_head_dim + d_rope) * attn.q_scaling)) output_scale = 1.0 if kernel_dtype == torch.float8_e4m3fn: - cached = getattr(self, "_cute_dsl_fp8_scale", None) - if cached is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "CuTe DSL MLA FMHA: fp8 decode scale was not cached for " - f"layer {attn.layer_idx} before CUDA graph capture." - ) - softmax_scale = float(params.fwd.mla_bmm1_scale[1].item()) / _LOG2_E - output_scale = float(params.fwd.mla_bmm2_scale[0].item()) - self._cute_dsl_fp8_scale = (softmax_scale, output_scale) - else: - softmax_scale, output_scale = cached + softmax_scale, output_scale = self._get_fp8_scales(q, params.fwd) + block_split_kvs = self._split_kv_plan( + meta, cache_seqs_base, batch_size, seq_len_q, num_heads + ) output_view = output.view(batch_size, seq_len_q, num_heads, d_latent) q_latent = q_view[..., :d_latent].permute(2, 3, 1, 0) @@ -540,8 +710,52 @@ def _run_mla_decode( if (meta.helix_position_offsets is not None and meta._helix_spec_tokens_valid) else None ), + block_split_kvs, + ) + + @staticmethod + def _split_kv_plan( + meta: "TrtllmAttentionMetadata", + cache_seqs: torch.Tensor, + batch_size: int, + seq_len_q: int, + num_heads: int, + ) -> Optional[torch.Tensor]: + """This step's per-request split-KV counts, planned once per step. + + The plan depends only on the step's KV lengths and the decode geometry, + so the first CuTe DSL MLA layer of a step launches it and the others + reuse it (``_invalidate_mla_scheduler_buffers`` resets the key when the + lengths change). None (uniform split) when it cannot be planned here. + """ + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + CuteDSLNVMlaDecodeBlackwellRunner, ) + if not 0 < batch_size <= _SPLIT_PLAN_MAX_REQUESTS: + return None + plan = getattr(meta, "_cute_dsl_mla_split_plan", None) + capacity = max(int(getattr(meta, "max_num_requests", 0) or 0), batch_size) + if plan is None or plan.numel() < batch_size or plan.device != cache_seqs.device: + if torch.cuda.is_current_stream_capturing(): + return None + plan = torch.empty(capacity, dtype=torch.int32, device=cache_seqs.device) + meta._cute_dsl_mla_split_plan = plan + runner = CuteDSLNVMlaDecodeBlackwellRunner + rows_per_request = runner.get_num_q_tiles(num_heads, seq_len_q) + wave_clusters = max(runner._get_max_active_blocks() // runner._CLUSTER_SHAPE_MNK[0], 1) + key = (batch_size, rows_per_request, wave_clusters, cache_seqs.data_ptr(), plan.data_ptr()) + if getattr(meta, "_cute_dsl_mla_split_plan_key", None) != key: + launch_mla_split_kv_plan( + cache_seqs, + plan, + num_requests=batch_size, + rows_per_request=rows_per_request, + wave_clusters=wave_clusters, + ) + meta._cute_dsl_mla_split_plan_key = key + return plan[:batch_size] + def run_mla_generation( self, params: FmhaParams, diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 048138a48582..a955ac3f560a 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -270,6 +270,14 @@ def effective_beam_width(self) -> int: _mla_ctx_cu_seqlens_valid: bool = field(default=False, init=False, repr=False) + # CuTe DSL MLA decode plans its per-request split-KV counts once per step + # (cute_dsl_mla.py); the plan buffer and the key that marks it current. + _cute_dsl_mla_split_plan: Optional[torch.Tensor] = field(init=False, + default=None, + repr=False) + _cute_dsl_mla_split_plan_key: Optional[tuple] = field(init=False, + default=None, + repr=False) # `DSAtrtllmAttentionMetadata` overrides this; the dense path keeps 0. num_sparse_topk: int = 0 @@ -692,6 +700,7 @@ def _invalidate_mla_scheduler_buffers(self) -> None: # buffers below must be rebuilt before the next MLA layer reads them. self._mla_scheduler_buffers_valid = False self._mla_ctx_cu_seqlens_valid = False + self._cute_dsl_mla_split_plan_key = None def restore_from_spec_dec(self) -> None: super().restore_from_spec_dec() diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index 6fac5ce27ff1..b475caa0e368 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -2059,25 +2059,41 @@ def setup_distributed_state(self, f"[AutoTuner] Whether using distributed tuning: {self._is_distributed()}" ) + def _tuning_group(self) -> Tuple[int, int]: + """Return ``(size, index)`` of this rank's tuning group. + + The group is the TP group. Helix CP ranks run the same SPMD forward + on the same token batch, so Helix adds the CP ranks too. Otherwise a + Helix generation worker with ``tp_size == 1`` would profile every + tactic on every CP rank. The index matches the TP-major, CP-minor + order of ``Distributed.tp_cp_allgather``, which merges the results. + """ + if self.mapping is None: + return 1, 0 + size = self.mapping.tp_size + index = self.mapping.tp_rank + if self.mapping.has_cp_helix(): + size *= self.mapping.cp_size + index = index * self.mapping.cp_size + self.mapping.cp_rank + return size, index + def _is_distributed(self) -> bool: """Check if we're in a distributed environment.""" - return self.mapping is not None and self.mapping.tp_size > 1 and self._dist is not None + return self._dist is not None and self._tuning_group()[0] > 1 def _maybe_parallelize_tactics( self, all_valid_tactics: List[Any], strategy: DistributedTuningStrategy) -> List[Any]: - """Parallelize tactics across all TP ranks if strategy is PARALLEL.""" + """Split tactics across the tuning group if strategy is PARALLEL.""" if strategy == DistributedTuningStrategy.PARALLEL and self._is_distributed( ): - # only distribute across TP ranks - # each TP rank will only tune the tactics that are assigned to it - tp_size = self.mapping.tp_size - tp_rank = self.mapping.tp_rank - valid_tactics = [] - for idx, tactic in enumerate(all_valid_tactics): - if idx % tp_size == tp_rank: - valid_tactics.append(tactic) - return valid_tactics + # Each rank of the tuning group only tunes its share of tactics; + # _merge_cache_data then keeps the fastest across the group. + group_size, group_index = self._tuning_group() + return [ + tactic for idx, tactic in enumerate(all_valid_tactics) + if idx % group_size == group_index + ] else: return all_valid_tactics diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index f500034c4df7..343d697d40a4 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -317,6 +317,26 @@ def _(mat_a, mat_b, bias, out_dtype): shape, dtype=out_dtype if out_dtype is not None else mat_a.dtype) return ret + @torch.library.register_fake("trtllm::dsv3_router_latent_gemm_op") + def _(input: torch.Tensor, gate_weight: torch.Tensor, + down_weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + for tensor in (input, gate_weight, down_weight): + torch._check(tensor.is_cuda and tensor.device == input.device) + torch._check(tensor.dim() == 2 and tensor.dtype == torch.bfloat16) + torch._check(tensor.is_contiguous()) + torch._check(tensor.stride(0) == 7168) + torch._check(tensor.stride(1) == 1) + torch._check(tensor.storage_offset() % 8 == 0) + torch._check(input.shape[0] >= 1) + torch._check(input.shape[0] <= 16) + torch._check(input.shape[1] == 7168) + torch._check(gate_weight.shape[0] == 896) + torch._check(gate_weight.shape[1] == 7168) + torch._check(down_weight.shape[0] == 3584) + torch._check(down_weight.shape[1] == 7168) + return (input.new_empty((input.shape[0], 896), dtype=torch.float32), + input.new_empty((input.shape[0], 3584), dtype=torch.bfloat16)) + @torch.library.register_fake("trtllm::dsv3_fused_a_gemm_op") def _(mat_a, mat_b, bias, out_dtype): shape = list(mat_a.shape) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 7d2b33f9e3de..a2a7e69c8dc8 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -11925,6 +11925,8 @@ def _( BlackwellMultiHeadLatentAttentionForwardFP8 from ..cute_dsl_kernels.blackwell.attention.mla.mla_decode_fp16 import \ BlackwellMultiHeadLatentAttentionForwardFP16 + from ..cute_dsl_kernels.blackwell.attention.mla.mla_helpers import \ + MAX_SPLITS as MLA_DECODE_MAX_SPLITS class CuteDSLNVMlaDecodeBlackwellRunner(TunableRunner): """Generic TunableRunner for the Blackwell CuTe DSL MLA decode kernels. @@ -11962,6 +11964,12 @@ class CuteDSLNVMlaDecodeBlackwellRunner(TunableRunner): _WORKSPACE_ALIGN = 128 _LSE_DTYPE_BYTES = 4 # float32 + _MMA_TILE_M = 128 + _MAX_SPLIT_KV = 32 + # CTA waves the split targets: with one wave, the longest request of a + # skewed decode batch sets the kernel time. + _SPLIT_KV_WAVES = 8 + def __init__( self, in_dtype, @@ -12021,25 +12029,25 @@ def _get_max_active_blocks(cls) -> int: cls._cute_dsl_max_active_blocks = cached return cached - @staticmethod - def get_default_split_kv(B: int, S: int, max_active_blocks: int) -> int: - # S is the number of packed 128-row query tiles per request. - max_split_kv = 32 - blocks_per_batch = max(1, max_active_blocks // B // (S * 2)) - split_kv = min(blocks_per_batch, max_split_kv) - return split_kv + @classmethod + def get_num_q_tiles(cls, num_heads: int, seq_len_q: int) -> int: + """Packed 128-row query tiles per request: the kernel packs the + ``seq_len_q * num_heads`` query rows into the M tile.""" + return ceil_div(num_heads * seq_len_q, cls._MMA_TILE_M) - @staticmethod - def get_default_is_persistent(B: int) -> bool: - if B >= 64: - return True - else: - return False + @classmethod + def get_default_split_kv(cls, B: int, S: int, + max_active_blocks: int) -> int: + """~_SPLIT_KV_WAVES waves of (2, B * S, split_kv) CTAs, as a power of + two; ``S`` is the number of packed query tiles per request.""" + target = cls._SPLIT_KV_WAVES * max_active_blocks / (2 * B * S) + return min(cls._MAX_SPLIT_KV, 1 << max(0, round(math.log2(target)))) @staticmethod def get_split_kv_candidates(B: int, S: int, max_active_blocks: int) -> List[int]: - # TODO: default split_kv is not always the best choice. We need to optimize it. + # One candidate: tuning profiles short KV, where split 1 wins, but the + # tactic is reused for long decode KV. return [ CuteDSLNVMlaDecodeBlackwellRunner.get_default_split_kv( B, S, max_active_blocks) @@ -12047,36 +12055,61 @@ def get_split_kv_candidates(B: int, S: int, @staticmethod def get_is_persistent_candidates() -> List[bool]: - return [True, False] + # The hardware scheduler balances skewed KV better than persistent + # round-robin; short-KV tuning would not see it. + return [False] + + @classmethod + def get_max_split_kv_rows(cls, H: int, S: int, + max_batch_size: int) -> int: + """Largest B * num_q_tiles(s) * split_kv over the batch buckets up + to ``max_batch_size`` and every s <= S, so the bound never shrinks + as S grows (the shared CUDA-graph workspace is sized once).""" + return cls._max_split_kv_rows(H, S, max_batch_size, + cls._get_max_active_blocks()) + + @classmethod + @functools.cache + def _max_split_kv_rows(cls, H: int, S: int, max_batch_size: int, + max_active_blocks: int) -> int: + rows = 0 + for s in range(1, S + 1): + num_q_tiles = cls.get_num_q_tiles(H, s) + bucket = 1 + while bucket <= max(1, max_batch_size): + split_kv = cls.get_default_split_kv(bucket, num_q_tiles, + max_active_blocks) + if split_kv > 1: + rows = max(rows, bucket * num_q_tiles * split_kv) + bucket *= 2 + return rows @classmethod def get_max_split_kv_workspace_size( cls, H: int, + S: int, D: int, + max_batch_size: int, acc_dtype: Type[cutlass.Numeric], ) -> int: """Raw bytes reserved for split-KV intermediates. - Batch-INDEPENDENT: CUDA graphs are captured per batch size in - descending order, and a later capture that needed a larger workspace - would resize the buffer, dangling the address baked into every - previously captured graph. + Fixed per (H, S, max_batch_size), never per batch: CUDA graphs are + captured per batch size in descending order, and a later capture + that needed a larger workspace would resize the buffer, dangling + the address baked into every previously captured graph. # cuda graph capture(B=8): eager warmup N times → capture graph_8 # cuda graph capture(B=4): eager warmup N times → capture graph_4 # ... # cuda graph replay # A later capture with a bigger workspace would resize it, so the - # bound covers every batch size up-front.""" - max_active_blocks = cls._get_max_active_blocks() - # Each query tile reserves 128 rows, including its masked tail. - # For split_kv > 1, the occupancy heuristic bounds - # B * num_q_tiles * split_kv by max_active_blocks / 2 at a tuning - # bucket. Runtime batches below the next bucket can use up to - # twice that space. split_kv == 1 does not use this workspace. - return (2 * max(H, 128) * (max_active_blocks // 2) * (D + 1) * - acc_dtype.width // 8) + # bound covers every batch size up-front; ``forward`` lowers + # split_kv for a batch between two buckets instead. Each query tile + # reserves _MMA_TILE_M rows, including its masked tail.""" + return (cls.get_max_split_kv_rows(H, S, max_batch_size) * + cls._MMA_TILE_M * (D + 1) * acc_dtype.width // 8) @classmethod def get_workspace_layout( @@ -12092,7 +12125,8 @@ def get_workspace_layout( lse_size = cls.get_lse_workspace_size(H, seq_len_q, max_batch_size) split_kv_offset = pad_up(lse_offset + lse_size, cls._WORKSPACE_ALIGN) - split_kv_size = cls.get_max_split_kv_workspace_size(H, D, acc_dtype) + split_kv_size = cls.get_max_split_kv_workspace_size( + H, seq_len_q, D, max_batch_size, acc_dtype) workspace_size = split_kv_offset + pad_up(split_kv_size, cls._WORKSPACE_ALIGN) return (lse_offset, lse_size, split_kv_offset, split_kv_size, @@ -12150,9 +12184,9 @@ def get_valid_tactics( ((128, 128), (128, 256)), ] max_active_blocks = self._get_max_active_blocks() - num_q_tiles = ceil_div(h * seq_len_q, 128) split_candidates = self.get_split_kv_candidates( - batch_size, num_q_tiles, max_active_blocks) + batch_size, self.get_num_q_tiles(h, seq_len_q), + max_active_blocks) persistent_candidates = self.get_is_persistent_candidates() valid = [] @@ -12329,24 +12363,18 @@ def default_tactic( ``batch_size`` is rounded down to its tuning bucket (``last_positive_power_of_2`` -- the same mapping the tuning - config uses) before deriving ``split_kv``: tuning profiles (and - therefore ``cute.compile``s) exactly the bucket-derived - ``split_kv`` variants, so a bucket-aligned fallback reuses an - already-compiled kernel where one exists instead of JIT-compiling - a fresh raw-batch ``split_kv`` variant in the serving loop. The - ``is_persistent`` choice is unaffected by the rounding (its - threshold is a power of two, so rounding down to a power of two - never crosses it), and both candidates are compiled during tuning - anyway.""" + config uses) before deriving ``split_kv``, so the fallback is the + tactic tuning would have stored for that bucket, and the one the + split-KV workspace is sized for.""" mma_qk_tiler_mn = (128, 128) mma_pv_tiler_mn = (128, 256) max_active_blocks = self._get_max_active_blocks() bucketed_batch_size = last_positive_power_of_2(batch_size) - num_q_tiles = ceil_div(self.num_heads * self.seq_len_q, 128) - split_kv = self.get_default_split_kv(bucketed_batch_size, - num_q_tiles, max_active_blocks) - is_persistent = self.get_default_is_persistent(bucketed_batch_size) - return (mma_qk_tiler_mn, mma_pv_tiler_mn, split_kv, is_persistent) + split_kv = self.get_default_split_kv( + bucketed_batch_size, + self.get_num_q_tiles(self.num_heads, self.seq_len_q), + max_active_blocks) + return (mma_qk_tiler_mn, mma_pv_tiler_mn, split_kv, False) def forward( self, @@ -12377,9 +12405,19 @@ def forward( softmax (max, sum) pair for Helix reduction. inputs[9] (kv_bounds): Optional contiguous int32 per-query local KV bounds of shape (B * S_q,). + inputs[10] (softmax_scale): FP8 runner only. float32 device + tensor; element 0 is the BMM1 (softmax) scale. + inputs[11] (output_scale): FP8 runner only. float32 device + tensor; element 0 is the BMM2 (output) scale. + inputs[12] (FP8) / inputs[10] (FP16) (block_split_kvs): + Optional int32 tensor of shape (B): the per-request + split-KV count planned on the device for this step's + KV lengths (clamped by the kernel to the tactic's + split_kv); absent = uniform split_kv. tactic: Tuple containing (mma_qk_tiler_mn, mma_pv_tiler_mn, split_kv, is_persistent). - **kwargs: Optional softmax_scale and output_scale values. + **kwargs: Optional scalar softmax_scale and output_scale values + for the FP16/BF16 runner. Returns: torch.Tensor: Output tensor of shape (H, D, S_q, B). The LSE @@ -12390,8 +12428,31 @@ def forward( # inputs[9] (optional): helix per-token attention bounds of shape # (B * S_q,), int32 — speculative verify groups only. kv_bounds = inputs[9] if len(inputs) > 9 else None - softmax_scale = float(kwargs.get("softmax_scale", 1.0)) - output_scale = float(kwargs.get("output_scale", 1.0)) + # FP8 KV keeps its scales on the device so neither the first call + # nor CUDA graph replay needs a host read of the scale tensors. + has_fp8_input = self.in_dtype == cutlass.Float8E4M3FN + fixed_inputs = 12 if has_fp8_input else 10 + block_split_kvs = (inputs[fixed_inputs] + if len(inputs) > fixed_inputs else None) + if has_fp8_input: + if len(inputs) not in (12, 13): + raise RuntimeError( + "CuteDSLNVMlaDecodeBlackwellRunner FP8 input requires " + "softmax_scale and output_scale tensors at inputs[10:12]." + ) + softmax_scale, output_scale = inputs[10:12] + for name, scale in (("softmax_scale", softmax_scale), + ("output_scale", output_scale)): + if scale.dtype != torch.float32 or scale.numel() < 1: + raise RuntimeError( + f"{name} must contain at least one float32 value.") + if scale.device != q_latent.device: + raise RuntimeError( + f"{name} must be on {q_latent.device}, got " + f"{scale.device}.") + else: + softmax_scale = float(kwargs.get("softmax_scale", 1.0)) + output_scale = float(kwargs.get("output_scale", 1.0)) if not (isinstance(tactic, tuple) and len(tactic) == 4): raise RuntimeError( @@ -12461,6 +12522,15 @@ def forward( # Kernel split-KV intermediates start AFTER the reserved LSE region. split_workspace = workspace_bytes[split_kv_offset:split_kv_offset + split_kv_size] + # The workspace fits each batch bucket's split; a batch between two + # buckets shrinks its split to fit instead of JIT-compiling a variant. + split_kv_rows = split_kv_size // (self._MMA_TILE_M * + (d_latent + 1) * + cutlass.Float32.width // 8) + fitting_split_kv = split_kv_rows // ( + batch_size * self.get_num_q_tiles(self.num_heads, seq_len_q)) + use_workspace = split_kv > 1 and fitting_split_kv >= 1 + split_kv = max(1, min(split_kv, fitting_split_kv)) if kv_bounds is not None and AutoTuner.get().is_tuning_mode: # Profiling rebuilds cache_seqs at bucketed sizes but input 9 @@ -12482,13 +12552,34 @@ def forward( f"{expected_bounds_shape}, got shape=" f"{tuple(kv_bounds.shape)}, dtype={kv_bounds.dtype}.") + if block_split_kvs is not None and ( + block_split_kvs.numel() != batch_size + or block_split_kvs.dtype != torch.int32 + or block_split_kvs.device != o.device + or not block_split_kvs.is_contiguous()): + if AutoTuner.get().is_tuning_mode: + # Profiling rebuilds cache_seqs at bucketed sizes; the + # step's plan does not apply to them (uniform split). + block_split_kvs = None + else: + raise RuntimeError( + "CuteDSLNVMlaDecodeBlackwellRunner requires contiguous " + f"int32 block_split_kvs of shape ({batch_size},) on the " + "output device.") + # The per-request plan only distributes the split-KV workspace, so + # it applies only when that workspace is used. + if not use_workspace: + block_split_kvs = None + is_var_split_kv = block_split_kvs is not None + # split_kv is a runtime Int32; only the workspace use is compiled in. cache_key = self.unique_id() + ( out_dtype, mma_qk_tiler_mn, mma_pv_tiler_mn, - split_kv, + use_workspace, is_persistent, kv_bounds is not None, + is_var_split_kv, ) if cache_key not in CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache: # A compile outside the tuning window stalls the serving loop @@ -12513,10 +12604,12 @@ def forward( self._SKIP_CORRECTION_THRESHOLD, is_persistent, self._IS_VAR_SEQ, - self._IS_VAR_SPLIT_KV, + is_var_split_kv, num_heads=self.num_heads, seq_len_q=seq_len_q, - reducer_max_splits=32, + # The variable split-KV reducer needs the generic capacity. + reducer_max_splits=(MLA_DECODE_MAX_SPLITS + if is_var_split_kv else 32), emit_softmax_stats=self.emit_softmax_stats, ) q_latent_ct = cute.runtime.from_dlpack( @@ -12544,7 +12637,6 @@ def forward( softmax_stats_kernel, assumed_align=16).mark_layout_dynamic( leading_dim=3) if softmax_stats_kernel is not None else None) - use_workspace = split_kv > 1 and split_workspace.numel() > 0 workspace_ct = (cute.runtime.from_dlpack( split_workspace, assumed_align=32).mark_layout_dynamic() if use_workspace else None) @@ -12553,8 +12645,17 @@ def forward( kv_bounds_ct = (cute.runtime.from_dlpack( kv_bounds, assumed_align=4).mark_layout_dynamic() if kv_bounds is not None else None) - # Variable split-KV (block_split_kvs) is not used on this path: - block_split_kvs_ct = None + block_split_kvs_ct = (cute.runtime.from_dlpack( + block_split_kvs, assumed_align=4).mark_layout_dynamic() + if is_var_split_kv else None) + if has_fp8_input: + softmax_scale_ct = cute.runtime.from_dlpack( + softmax_scale).mark_layout_dynamic() + output_scale_ct = cute.runtime.from_dlpack( + output_scale).mark_layout_dynamic() + else: + softmax_scale_ct = cutlass.Float32(softmax_scale) + output_scale_ct = cutlass.Float32(output_scale) compile_args = [ q_latent_ct, @@ -12576,8 +12677,8 @@ def forward( kv_bounds_ct, None, # cum_seq_lens_q: this path uses fixed query lengths block_split_kvs_ct, - cutlass.Float32(softmax_scale), - cutlass.Float32(output_scale), + softmax_scale_ct, + output_scale_ct, stream, ]) CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache[cache_key] = ( @@ -12609,13 +12710,12 @@ def forward( if self.emit_softmax_stats: runtime_args.append(softmax_stats_kernel) runtime_args.extend([ - split_workspace if - (split_kv > 1 and split_workspace.numel() > 0) else None, + split_workspace if use_workspace else None, split_kv, cache_seqs, kv_bounds, None, # cum_seq_lens_q: this path uses fixed query lengths - None, # block_split_kvs: var-split path unused (is_var_split_kv False) + block_split_kvs, softmax_scale, output_scale, stream, @@ -12640,19 +12740,24 @@ def cute_dsl_mla_decode_fp8_blackwell( num_heads: int, seq_len_q: int, page_size: int, - softmax_scale: float, - output_scale: float, + softmax_scale: torch.Tensor, + output_scale: torch.Tensor, # Keep the trailing arguments required in the custom-op schema. PyTorch # elides trailing default-valued arguments before its mutation fallback, # while mutates_args retains their positional indices. max_batch_size: int, softmax_stats: Optional[torch.Tensor], kv_bounds: Optional[torch.Tensor], + block_split_kvs: Optional[torch.Tensor] = None, ) -> None: """CuTe DSL FP8 MLA decode (Blackwell SM100/SM103/SM107). + softmax_scale / output_scale: float32 device tensors whose element 0 + is the BMM1 (softmax) and BMM2 (output) scale. kv_bounds: helix speculative verify groups -- per-token rank-local attention bounds of shape (B * seq_len_q,), int32. + block_split_kvs: optional per-request split-KV counts of shape (B,), + int32, planned on the device for the step's KV lengths. """ if (sm_version := get_sm_version()) not in (100, 103, 107): raise ValueError( @@ -12671,8 +12776,10 @@ def cute_dsl_mla_decode_fp8_blackwell( ) inputs = [ q_latent, q_rope, c_latent, c_rope, page_table, cache_seqs, o, - workspace, softmax_stats, kv_bounds + workspace, softmax_stats, kv_bounds, softmax_scale, output_scale ] + if block_split_kvs is not None: + inputs.append(block_split_kvs) tuner = AutoTuner.get() _, best_tactic = tuner.choose_one( "trtllm::cute_dsl_mla_decode_fp8_blackwell", @@ -12685,12 +12792,7 @@ def cute_dsl_mla_decode_fp8_blackwell( # is_persistent still come from a length-4 tactic. if not (isinstance(best_tactic, tuple) and len(best_tactic) == 4): best_tactic = runner.default_tactic(int(q_latent.shape[-1])) - runner( - inputs, - tactic=best_tactic, - softmax_scale=softmax_scale, - output_scale=output_scale, - ) + runner(inputs, tactic=best_tactic) @torch.library.register_fake("trtllm::cute_dsl_mla_decode_fp8_blackwell") def _( @@ -12705,11 +12807,12 @@ def _( num_heads: int, seq_len_q: int, page_size: int, - softmax_scale: float, - output_scale: float, + softmax_scale: torch.Tensor, + output_scale: torch.Tensor, max_batch_size: int, softmax_stats: Optional[torch.Tensor], kv_bounds: Optional[torch.Tensor], + block_split_kvs: Optional[torch.Tensor] = None, ) -> None: return None @@ -12736,11 +12839,14 @@ def cute_dsl_mla_decode_fp16_blackwell( max_batch_size: int, softmax_stats: Optional[torch.Tensor], kv_bounds: Optional[torch.Tensor], + block_split_kvs: Optional[torch.Tensor] = None, ) -> None: """CuTe DSL FP16/BF16 MLA decode (Blackwell SM100/SM103/SM107). kv_bounds: helix speculative verify groups — per-token rank-local attention bounds of shape (B * seq_len_q,), int32. + block_split_kvs: optional per-request split-KV counts of shape (B,), + int32, planned on the device for the step's KV lengths. """ if (sm_version := get_sm_version()) not in (100, 103, 107): raise ValueError( @@ -12778,6 +12884,8 @@ def cute_dsl_mla_decode_fp16_blackwell( q_latent, q_rope, c_latent, c_rope, page_table, cache_seqs, o, workspace, softmax_stats, kv_bounds ] + if block_split_kvs is not None: + inputs.append(block_split_kvs) tuner = AutoTuner.get() _, best_tactic = tuner.choose_one( "trtllm::cute_dsl_mla_decode_fp16_blackwell", @@ -12815,6 +12923,7 @@ def _( max_batch_size: int, softmax_stats: Optional[torch.Tensor], kv_bounds: Optional[torch.Tensor], + block_split_kvs: Optional[torch.Tensor] = None, ) -> None: return None @@ -14333,6 +14442,20 @@ class CuteDSLMXFP8RubinLinear(CuteDSLBlockScaledRubinLinear): use_cuda_graph=True, ) + class CuteDSLKimiK3MXFP8RubinLinear(CuteDSLMXFP8RubinLinear): + """MXFP8 runner with fine-grained tuning for K3's low-M range.""" + + tuning_config = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, _get_kimi_k3_mxfp8_tuning_buckets, + _kimi_k3_mxfp8_tuning_bucket), ), + constraint_specs=(ConstraintSpec(2, 0, + mxfp8_scale_infer_shape), ), + use_cold_l2_cache=True, + distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL, + use_cuda_graph=True, + ) + class CuteDSLNVFP4InplaceRubinLinear(CuteDSLNVFP4RubinLinear): kernel_cache = dict() tuning_config = TuningConfig( @@ -14356,8 +14479,12 @@ def cute_dsl_mxfp8_gemm_rubin( weight_scale: torch.Tensor, output_dtype: torch.dtype = torch.bfloat16, use_tvm_ffi: bool = True, + fine_grained_m: bool = False, ) -> torch.Tensor: - """Run the SM107 dense MXFP8 GEMM with K32 R128c4 UE8M0 scales.""" + """Run the SM107 dense MXFP8 GEMM with K32 R128c4 UE8M0 scales. + + ``fine_grained_m`` enables Kimi K3's 16-token low-M tuning buckets. + """ if output_dtype != torch.bfloat16: raise ValueError( f"CuteDSL MXFP8 only supports bfloat16 output, got " @@ -14373,8 +14500,10 @@ def cute_dsl_mxfp8_gemm_rubin( raise ValueError("CuteDSL MXFP8 scales must be UE8M0 uint8") alpha = _get_mxfp8_gemm_alpha(input.device) - runner = CuteDSLMXFP8RubinLinear(output_dtype=output_dtype, - use_tvm_ffi=use_tvm_ffi) + runner_cls = (CuteDSLKimiK3MXFP8RubinLinear + if fine_grained_m else CuteDSLMXFP8RubinLinear) + runner = runner_cls(output_dtype=output_dtype, + use_tvm_ffi=use_tvm_ffi) inputs = [input, weight, input_scale, weight_scale, alpha] _, best_tactic = AutoTuner.get().choose_one( "trtllm::cute_dsl_mxfp8_gemm_rubin", @@ -14392,6 +14521,7 @@ def _( weight_scale: torch.Tensor, output_dtype: torch.dtype = torch.bfloat16, use_tvm_ffi: bool = True, + fine_grained_m: bool = False, ) -> torch.Tensor: shape = list(mat_a.shape) shape[-1] = mat_b.shape[-2] @@ -18074,8 +18204,9 @@ class Sm107BlockScaledContiguousGroupedGemmFusedFc12Runner( TunableRunner): """Rubin runner for the fused FC1+FC2 (FC12) NVFP4 MoE kernel. - The fused kernel replaces the two-op FC1 (gather+GEMM+SwiGLU+quant) + The fused kernel replaces the two-op FC1 (gather+GEMM+gated-act+quant) and FC2 (GEMM+finalize) sequence with a single persistent kernel. + The gated epilogue is SwiGLU or SiTU (trace-time specialization). The only interface delta versus the existing CuteDSL backend is the three int32 atomic counters (fc1_ready / fc1_scheduler_counter / fc2_scheduler_counter), which are allocated and memset to zero here @@ -18088,16 +18219,20 @@ class Sm107BlockScaledContiguousGroupedGemmFusedFc12Runner( kernel_cache = dict() tuning_config_cache = dict() - def __init__(self, - num_experts: int, - top_k: int, - num_local_experts: int, - local_expert_offset: int, - tile_size: int, - scaling_vector_size: int = 16, - swiglu_limit: float = float("inf"), - ep_size: int = 1, - enable_alltoall: bool = False): + def __init__( + self, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + scaling_vector_size: int = 16, + swiglu_limit: float = float("inf"), + ep_size: int = 1, + enable_alltoall: bool = False, + activation_type: ActivationType = ActivationType.Swiglu, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None): super().__init__() self.num_experts = num_experts self.top_k = top_k @@ -18106,6 +18241,9 @@ def __init__(self, self.tile_size = tile_size self.scaling_vector_size = scaling_vector_size self.swiglu_limit = swiglu_limit + self.activation_type = ActivationType(int(activation_type)) + self.situ_beta = situ_beta + self.situ_linear_beta = situ_linear_beta # Used only by the in-op output memset (moved here so the memset # is the fused kernel's immediate stream predecessor). self.ep_size = ep_size @@ -18130,6 +18268,9 @@ def unique_id(self): self.tile_size, self.scaling_vector_size, self.swiglu_limit, + int(self.activation_type), + self.situ_beta, + self.situ_linear_beta, ) def get_valid_tactics( @@ -18299,16 +18440,18 @@ def forward(self, inputs: List[torch.Tensor], fc1_sfc = torch.empty(m * interm_size // sf_vec, dtype=fc1_sfa.dtype, device=fc1_sfa.device) - # Three atomic counters: allocate + memset-zero every launch. - fc1_ready = torch.zeros(num_tiles, - dtype=torch.int32, - device=fc1_a.device) - fc1_scheduler_counter = torch.zeros(1, - dtype=torch.int32, - device=fc1_a.device) - fc2_scheduler_counter = torch.zeros(1, - dtype=torch.int32, - device=fc1_a.device) + # Three atomic counters zeroed by one fill (they sit on the + # routing chain); each scheduler counter on its own 128 B line. + ints_per_line = 32 + fc1_at = ((num_tiles + ints_per_line - 1) // ints_per_line * + ints_per_line) + fc2_at = fc1_at + ints_per_line + counters = torch.zeros(fc2_at + ints_per_line, + dtype=torch.int32, + device=fc1_a.device) + fc1_ready = counters[:num_tiles] + fc1_scheduler_counter = counters[fc1_at:fc1_at + 1] + fc2_scheduler_counter = counters[fc2_at:fc2_at + 1] # Zero the scatter-add output right before the fused kernel so # the memset becomes the kernel's immediate stream predecessor @@ -18427,7 +18570,8 @@ def forward(self, inputs: List[torch.Tensor], cache_key = (self.scaling_vector_size, self.tile_size, self.top_k, mma_tiler, mma_inst_shape, cluster_shape_mn, max_active_clusters, - self.swiglu_limit) + self.swiglu_limit, int(self.activation_type), + self.situ_beta, self.situ_linear_beta) if cache_key not in self.__class__.kernel_cache: gemm = self.__class__.kernel_class( self.scaling_vector_size, @@ -18439,6 +18583,9 @@ def forward(self, inputs: List[torch.Tensor], use_pdl=True, swiglu_limit=self.swiglu_limit, scheduler="l2_atomic", + activation_type=self.activation_type, + situ_beta=self.situ_beta, + situ_linear_beta=self.situ_linear_beta, ) @cute.jit @@ -18674,6 +18821,9 @@ def _run_nvfp4_fc12_fused_rubin( ep_size: int, enable_alltoall: bool, tuner_key: str, + activation_type: ActivationType = ActivationType.Swiglu, + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, ) -> torch.Tensor: tuner = AutoTuner.get() runner = Sm107BlockScaledContiguousGroupedGemmFusedFc12Runner( @@ -18686,6 +18836,9 @@ def _run_nvfp4_fc12_fused_rubin( swiglu_limit=swiglu_limit, ep_size=ep_size, enable_alltoall=enable_alltoall, + activation_type=ActivationType(activation_type), + situ_beta=_canonicalize_situ_beta(situ_beta), + situ_linear_beta=_canonicalize_situ_beta(situ_linear_beta), ) # Input order matches Fc12FusedInputsHelper (FC1 prefix 0..9 mirrors # GatherGroupedGemmInputsHelper; FC2 tensors 10..14; expanded_idx 15 @@ -18722,7 +18875,9 @@ def _run_nvfp4_fc12_fused_rubin( "SymInt local_expert_offset, SymInt tile_size, float swiglu_limit, " "SymInt ep_size, bool enable_alltoall, " "SymInt scaling_vector_size=16, " - "str? precomputed_tactic=None) -> ()", + "str? precomputed_tactic=None, " + f"SymInt activation_type={int(ActivationType.Swiglu)}, " + "float situ_beta=-1.0, float situ_linear_beta=-1.0) -> ()", device_types="cuda") def cute_dsl_nvfp4_fc12_fused_rubin( input: torch.Tensor, @@ -18751,6 +18906,9 @@ def cute_dsl_nvfp4_fc12_fused_rubin( enable_alltoall: bool, scaling_vector_size: int = 16, precomputed_tactic: Optional[str] = None, + activation_type: int = int(ActivationType.Swiglu), + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, ) -> None: # In-place: finalize scatter-adds into ``output`` (mutates_args); # the op returns nothing so it does not alias its own input. @@ -18763,7 +18921,8 @@ def cute_dsl_nvfp4_fc12_fused_rubin( local_expert_offset, tile_size, scaling_vector_size, swiglu_limit, precomputed_tactic, expanded_idx_to_permuted_idx, ep_size, enable_alltoall, - "trtllm::cute_dsl_nvfp4_fc12_fused_rubin") + "trtllm::cute_dsl_nvfp4_fc12_fused_rubin", + ActivationType(activation_type), situ_beta, situ_linear_beta) @torch.library.register_fake("trtllm::cute_dsl_nvfp4_fc12_fused_rubin") def _( @@ -18793,5 +18952,8 @@ def _( enable_alltoall: bool, scaling_vector_size: int = 16, precomputed_tactic: Optional[str] = None, + activation_type: int = int(ActivationType.Swiglu), + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, ) -> None: return None diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py index 98e4a1d434fa..ece779ddcdb1 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py @@ -17,7 +17,7 @@ Wraps the source-integrated ``kda_decode_mtp_kernel`` (see ``cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py``) as the ``trtllm::kda_mtp_decode`` operator. One launch fuses, per generation -request: replay of previously-accepted draft tokens from the ``qkg/v/beta`` +request: replay of previously-accepted draft tokens from the ``kg/v/beta`` caches, causal conv + SiLU, Q/K L2 norm, beta sigmoid, lower-bound gate, and the KDA delta-rule recurrence over the ``1 + num_spec`` new tokens. The recurrent state and base conv windows are committed **in place** after the @@ -43,16 +43,16 @@ syncs. * The compile cache is keyed purely by dtype/shape/stride layouts and constexpr flags — no ``id()`` or ``data_ptr`` keys. -* Bias / output-norm / ``pad_slot_id`` arguments (unsupported by the - specialized kernel, previously validated-then-rejected) are dropped from - the signature. +* Bias / ``pad_slot_id`` arguments (unsupported by the specialized kernel, + previously validated-then-rejected) are dropped from the signature. Gated + output norm and MXFP8 quantization are optional packed-layout epilogues. Kernel shape contract: ``K == V == 128``, conv width ``W == 4``, -``HV == H``, TILE_V=64. ``num_spec`` is a compile-time constant per cache -allocation. Benchmark-tuned fast variants exist for ``N in (32, 128), H in -(2, 12, 32), num_spec == 2``; other shapes compile the general variant. +``HV == H``, one V pass (TILE_V=128, 512 threads). ``num_spec`` is a +compile-time constant per cache allocation. """ +import os from math import gcd from typing import Optional, Tuple @@ -66,17 +66,22 @@ import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute + import cutlass.utils.blackwell_helpers as sm100_utils + from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 from cutlass.cute.runtime import from_dlpack from ..cute_dsl_kernels.blackwell.kimi_k3_kda.kda_mtp_decode import ( NUM_THREADS, + SPLIT_MAX_TILES, + T_PAD, TILE_K, kda_decode_mtp_kernel, ) else: raise ImportError("Kimi K3 KDA MTP decode requires NVIDIA CUTLASS DSL") -_TILE_V = 64 +# One V pass: the kernel runs 16 warps x 8 rows over the full 128-row state. +_TILE_V = 128 if IS_CUTLASS_DSL_AVAILABLE: @@ -97,9 +102,13 @@ def _run_kda_decode_mtp( g: cute.Tensor, dt_bias: cute.Tensor, beta: cute.Tensor, + onorm_g: cute.Tensor, + onorm_weight: cute.Tensor, o: cute.Tensor, + output_scale: cute.Tensor, ht: cute.Tensor, - qkg_cache: cute.Tensor, + k_cache: cute.Tensor, + g_cache: cute.Tensor, v_cache: cute.Tensor, beta_cache: cute.Tensor, stage_timing: cute.Tensor, @@ -116,14 +125,19 @@ def _run_kda_decode_mtp( TILE_V: cutlass.Constexpr[int], KERNEL_WIDTH: cutlass.Constexpr[int], lower_bound: cutlass.Constexpr[float], + onorm_eps: cutlass.Constexpr[float], + scale_leading_dim: cutlass.Constexpr[int], USE_FLAT_LAYOUT: cutlass.Constexpr[bool], - USE_SETMAXREG: cutlass.Constexpr[bool], USE_REGULAR_METADATA: cutlass.Constexpr[bool], - USE_REG_Q_WEIGHTS: cutlass.Constexpr[bool], + USE_PACKED_TOKEN_LAYOUT: cutlass.Constexpr[bool], USE_ZERO_ACCEPTED: cutlass.Constexpr[bool], FUSE_PRECOMPUTE: cutlass.Constexpr[bool], RUNTIME_PRECOMPUTE_FLAG: cutlass.Constexpr[bool], + FUSE_OUTPUT_NORM: cutlass.Constexpr[bool], + QUANTIZE_OUTPUT: cutlass.Constexpr[bool], PROFILE_STAGES: cutlass.Constexpr[bool], + SPLIT_V: cutlass.Constexpr[int], + BF16_MMA: cutlass.Constexpr[bool], REPLAY_ONLY: cutlass.Constexpr[bool], stream: cuda.CUstream, ): @@ -132,7 +146,161 @@ def _run_kda_decode_mtp( else: t_max = 2 * NUM_SPEC + 1 smem_qk_layout = cute.make_layout((t_max, K), stride=(K, 1)) - kda_decode_mtp_kernel( + # BF16_MMA: bf16 tensor-core operands ([K~; Q~] only, K = 16 per instruction). + OP_DTYPE = cutlass.BFloat16 if BF16_MMA else cutlass.Float32 + N1 = 2 * T_PAD if BF16_MMA else 3 * T_PAD + # GEMM1: M = V rows of the state (TMEM), N = operand rows ([K~; Q~; K~lo] for tf32). + tiled_mma1 = sm100_utils.make_trivial_tiled_mma( + OP_DTYPE, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + (V, N1), + tcgen05.OperandSource.TMEM, + ) + smem_layout_b1 = sm100_utils.make_smem_layout_b(tiled_mma1, (V, N1, K), OP_DTYPE, 1) + # fp32 A-fragment layout of the (V x K) state in TMEM: the GEMM2 accumulator seed / state + # store views use it in every mode. + tiled_mma_f32a = sm100_utils.make_trivial_tiled_mma( + cutlass.Float32, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + (V, 3 * T_PAD), + tcgen05.OperandSource.TMEM, + ) + # GEMM1b: the S0_lo sweep only needs K~hi (N = T_PAD) -> a third of the N=48 work. + tiled_mma1k = sm100_utils.make_trivial_tiled_mma( + cutlass.Float32, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + (V, T_PAD), + tcgen05.OperandSource.TMEM, + ) + smem_layout_b1k = sm100_utils.make_smem_layout_b( + tiled_mma1k, (V, T_PAD, K), cutlass.Float32, 1 + ) + # TMA bulk store of the committed state: (V, K, tiles) view of ht, one 128B-swizzled + # (V/SPLIT_V x 32 fp32 | 64 bf16) box per 128 B-wide column slab. The view takes ht's own strides (the + # per-layer state cache is a strided slice of a larger buffer); the tile mode is the pool + # index (flat layout) or the (head, slot) pair. + if cutlass.const_expr(cute.rank(ht.shape) == 3): + mT_vkl = cute.make_tensor( + ht.iterator, + cute.make_layout( + (V, K, ht.shape[0]), + stride=(ht.layout.stride[1], ht.layout.stride[2], ht.layout.stride[0]), + ), + ) + else: + mT_vkl = cute.make_tensor( + ht.iterator, + cute.make_layout( + (V, K, (ht.shape[1], ht.shape[0])), + stride=( + ht.layout.stride[2], + ht.layout.stride[3], + (ht.layout.stride[1], ht.layout.stride[0]), + ), + ), + ) + # 128 B swizzled box rows: 32 fp32 or 64 bf16 state columns per box + box_cols = 128 // (h0.element_type.width // 8) + box_atom = sm100_utils.make_smem_layout_atom( + tcgen05.SmemLayoutAtomKind.K_SW128, h0.element_type + ) + box_layout = cute.tile_to_shape(box_atom, (V, box_cols), order=(0, 1)) + # committed-state store: each CTA writes its V/SPLIT_V rows through (V/SPLIT_V x box_cols) boxes + box_layout_t = cute.tile_to_shape(box_atom, (V // SPLIT_V, box_cols), order=(0, 1)) + tma_atom_t, mT_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), mT_vkl, box_layout_t, (V // SPLIT_V, box_cols) + ) + # TMA bulk load of the initial state S0 into the same box layout (issued at kernel start). + if cutlass.const_expr(cute.rank(h0.shape) == 3): + mS_vkl = cute.make_tensor( + h0.iterator, + cute.make_layout( + (V, K, h0.shape[0]), + stride=(h0.layout.stride[1], h0.layout.stride[2], h0.layout.stride[0]), + ), + ) + else: + mS_vkl = cute.make_tensor( + h0.iterator, + cute.make_layout( + (V, K, (h0.shape[1], h0.shape[0])), + stride=( + h0.layout.stride[2], + h0.layout.stride[3], + (h0.layout.stride[1], h0.layout.stride[0]), + ), + ), + ) + tma_atom_l, mS_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), mS_vkl, box_layout, (V, box_cols) + ) + # GEMM2: (V x T_PAD) x (T_PAD x K), both operands in smem, tf32; issued as two N = K/2 + # column halves so the state store can start on half 0. + tiled_mma2h = sm100_utils.make_trivial_tiled_mma( + OP_DTYPE, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + (V, K // 2), + tcgen05.OperandSource.SMEM, + ) + smem_layout_b2h = sm100_utils.make_smem_layout_b( + tiled_mma2h, (V, K // 2, T_PAD), OP_DTYPE, 1 + ) + smem_layout_a2 = sm100_utils.make_smem_layout_a( + tiled_mma2h, (V, K // 2, T_PAD), OP_DTYPE, 1 + ) + # GEMM2 half 0 also carries the output-combination GEMM: N = T_PAD (Bcoef) + K/2 (Kbar rows 0..63). + tiled_mma2q = sm100_utils.make_trivial_tiled_mma( + OP_DTYPE, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + (V, T_PAD + K // 2), + tcgen05.OperandSource.SMEM, + ) + smem_layout_b2q = sm100_utils.make_smem_layout_b( + tiled_mma2q, (V, T_PAD + K // 2, T_PAD), OP_DTYPE, 1 + ) + # GEMM3: T x T coupling matrices, A3 = [Kf; Qf] padded to 64 rows, B3 = Kb (16 rows). + tiled_mma3 = sm100_utils.make_trivial_tiled_mma( + OP_DTYPE, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + (4 * T_PAD, T_PAD), + tcgen05.OperandSource.SMEM, + ) + smem_layout_a3 = sm100_utils.make_smem_layout_a( + tiled_mma3, (4 * T_PAD, T_PAD, K), OP_DTYPE, 1 + ) + smem_layout_b3 = sm100_utils.make_smem_layout_b( + tiled_mma3, (4 * T_PAD, T_PAD, K), OP_DTYPE, 1 + ) + # Output combination D4[v, t] = sum_s V'[s, v] Bcoef[t, s]: computed as the first 16 columns + # of GEMM2 half 0; this tiled MMA only provides the accumulator layout of those columns. + tiled_mma4 = sm100_utils.make_trivial_tiled_mma( + cutlass.Float32, + OperandMajorMode.K, + OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + (V, T_PAD), + tcgen05.OperandSource.SMEM, + ) + launched = kda_decode_mtp_kernel( h0, x_q, x_k, @@ -147,9 +315,13 @@ def _run_kda_decode_mtp( g, dt_bias, beta, + onorm_g, + onorm_weight, o, + output_scale, ht, - qkg_cache, + k_cache, + g_cache, v_cache, beta_cache, smem_qk_layout, @@ -165,17 +337,49 @@ def _run_kda_decode_mtp( NUM_SPEC, KERNEL_WIDTH, lower_bound, + onorm_eps, + scale_leading_dim, USE_FLAT_LAYOUT, - USE_SETMAXREG, USE_REGULAR_METADATA, - USE_REG_Q_WEIGHTS, + USE_PACKED_TOKEN_LAYOUT, USE_ZERO_ACCEPTED, FUSE_PRECOMPUTE, RUNTIME_PRECOMPUTE_FLAG, + FUSE_OUTPUT_NORM, + QUANTIZE_OUTPUT, stage_timing, PROFILE_STAGES, + tiled_mma1, + tiled_mma3, + tiled_mma1k, + smem_layout_b1, + smem_layout_a2, + smem_layout_a3, + smem_layout_b3, + smem_layout_b1k, + tiled_mma4, + tma_atom_t, + mT_tma, + box_layout, + tma_atom_l, + mS_tma, + tiled_mma2h, + smem_layout_b2h, + box_layout_t, + tiled_mma2q, + smem_layout_b2q, + SPLIT_V, + BF16_MMA, + tiled_mma_f32a, REPLAY_ONLY, - ).launch(grid=(HV, N, 1), block=[NUM_THREADS, 1, 1], stream=stream) + ) + if cutlass.const_expr(SPLIT_V == 2): + # cluster pair: the two CTAs of a (head, request) must be co-scheduled (in-place state) + launched.launch( + grid=(HV, N, 2), block=[NUM_THREADS, 1, 1], cluster=(1, 1, 2), stream=stream + ) + else: + launched.launch(grid=(HV, N, 1), block=[NUM_THREADS, 1, 1], stream=stream) def _require_stride_layout( @@ -194,13 +398,16 @@ def _require_stride_layout( A_log, dt_bias, recurrent_state, - qkg_cache, + k_cache, + g_cache, v_cache, beta_cache, ssm_state_indices, cu_seqlens, num_accepted_tokens, out, + onorm_g, + onorm_weight, H, HV, K, @@ -208,6 +415,9 @@ def _require_stride_layout( W, num_spec, T_total, + packed_token_layout, + fuse_output_norm, + quantize_output, ): if x_q.ndim != 4 or x_k.ndim != 4 or x_v.ndim != 4: raise ValueError("Expected x_q/x_k/x_v to have shape [1, T, H, D].") @@ -219,8 +429,28 @@ def _require_stride_layout( raise ValueError(f"Expected g shape [1, {T_total}, {HV}, {K}].") if beta.ndim != 3 or beta.shape != (1, T_total, HV): raise ValueError(f"Expected beta shape [1, {T_total}, {HV}].") - if out.ndim != 4 or out.shape != (1, T_total, HV, V): - raise ValueError(f"Expected out shape [1, {T_total}, {HV}, {V}].") + if quantize_output: + if out.ndim != 2 or out.shape != (T_total, HV * V): + raise ValueError(f"Expected FP8 out shape [{T_total}, {HV * V}].") + if out.dtype is not torch.float8_e4m3fn: + raise TypeError("Expected FP8 out dtype torch.float8_e4m3fn.") + elif out.ndim != 4 or out.shape != (1, T_total, HV, V): + raise ValueError(f"Expected BF16 out shape [1, {T_total}, {HV}, {V}].") + if fuse_output_norm: + if not packed_token_layout: + raise ValueError("fuse_output_norm requires packed_token_layout=True.") + if onorm_g is None or onorm_weight is None: + raise ValueError("fuse_output_norm requires onorm_g and onorm_weight.") + if onorm_g.ndim != 4 or onorm_g.shape != (1, T_total, HV, V): + raise ValueError(f"Expected onorm_g shape [1, {T_total}, {HV}, {V}].") + if onorm_g.stride(-1) != 1 or onorm_g.stride(-2) != V: + raise ValueError("Expected onorm_g to have contiguous head/value blocks.") + if onorm_weight.ndim != 1 or onorm_weight.shape[0] != V: + raise ValueError(f"Expected onorm_weight shape [{V}].") + if onorm_weight.stride(0) != 1: + raise ValueError("Expected onorm_weight to be contiguous.") + if quantize_output and not fuse_output_norm: + raise ValueError("quantize_output requires fuse_output_norm=True.") last_dim_tensors = { "x_q": x_q, @@ -228,12 +458,14 @@ def _require_stride_layout( "x_v": x_v, "g": g, "beta": beta, - "out": out, "recurrent_state": recurrent_state, - "qkg_cache": qkg_cache, + "k_cache": k_cache, + "g_cache": g_cache, "v_cache": v_cache, "beta_cache": beta_cache, } + if not quantize_output: + last_dim_tensors["out"] = out for name, tensor in last_dim_tensors.items(): if tensor.stride(-1) != 1: raise ValueError(f"Expected {name} to be contiguous in its last dimension.") @@ -268,45 +500,45 @@ def _require_stride_layout( raise ValueError( f"Expected recurrent_state shape [pool, {HV}, {V}, {K}] (V-first pool layout)." ) - if recurrent_state.dtype != torch.float32: - # This kernel generation reads and writes the state pool as fp32. A - # bf16 pool (kv_cache_config.mamba_ssm_cache_dtype=bfloat16) is staged - # to fp32 for prefill and single-token decode, but the fused verify - # path updates the pool rows in place through ssm_state_indices and has - # no staging, so reject it here instead of reinterpreting the bytes. - # bf16 acceptance is a property of the tensor-core (tcgen05) kernel - # formulation, not a dtype switch: it derives the TMA box geometry from - # the element width (64 vs 32 state columns per 128 B row) and skips the - # hi/lo tf32 split GEMM that an fp32 state needs. This kernel generation - # is register-resident with no TMA or tcgen05 path to attach that to. - raise ValueError( - "kda_mtp_decode requires an fp32 recurrent_state pool; got " - f"{recurrent_state.dtype}. A bf16 KDA state pool is not supported " - "together with fused KDA MTP verify; set " - "kv_cache_config.mamba_ssm_cache_dtype=float32 (the Kimi K3 " - "default) to run speculative decoding." - ) - if qkg_cache.ndim != 4 or qkg_cache.shape[1:] != (num_spec, 3, H * K): - raise ValueError(f"Expected qkg_cache shape [pool, {num_spec}, 3, {H * K}].") + if recurrent_state.dtype not in (torch.float32, torch.bfloat16): + raise ValueError(f"recurrent_state must be fp32 or bf16, got {recurrent_state.dtype}") + if k_cache.ndim != 3 or k_cache.shape[1:] != (num_spec, H * K): + raise ValueError(f"Expected k_cache shape [pool, {num_spec}, {H * K}].") + if g_cache.ndim != 3 or g_cache.shape[1:] != (num_spec, H * K): + raise ValueError(f"Expected g_cache shape [pool, {num_spec}, {H * K}].") + if g_cache.dtype != torch.float32: + raise ValueError("g_cache (the replayed gates) must be fp32.") if v_cache.ndim != 3 or v_cache.shape[1:] != (num_spec, HV * V): raise ValueError(f"Expected v_cache shape [pool, {num_spec}, {HV * V}].") if beta_cache.ndim != 3 or beta_cache.shape[1:] != (num_spec, HV): raise ValueError(f"Expected beta_cache shape [pool, {num_spec}, {HV}].") if ( - qkg_cache.shape[0] < pool_size + k_cache.shape[0] < pool_size + or g_cache.shape[0] < pool_size or v_cache.shape[0] < pool_size or beta_cache.shape[0] < pool_size ): raise ValueError("Expected cache pool dimensions to cover recurrent_state rows.") - if ssm_state_indices.ndim != 1 or cu_seqlens.ndim != 1 or num_accepted_tokens.ndim != 1: - raise ValueError( - "Expected ssm_state_indices, cu_seqlens, and num_accepted_tokens to be 1D." - ) - if cu_seqlens.shape[0] != ssm_state_indices.shape[0] + 1: - raise ValueError("Expected cu_seqlens length to be N + 1.") - if num_accepted_tokens.shape[0] != ssm_state_indices.shape[0]: - raise ValueError("Expected num_accepted_tokens length to match N.") + if ssm_state_indices.ndim != 1 or num_accepted_tokens.ndim != 1: + raise ValueError("Expected ssm_state_indices and num_accepted_tokens to be 1D.") + num_requests = ssm_state_indices.shape[0] + if packed_token_layout: + if cu_seqlens is not None: + raise ValueError("packed_token_layout derives row offsets; cu_seqlens must be None.") + if T_total != num_requests * (num_spec + 1): + raise ValueError( + "packed_token_layout expects exactly num_spec + 1 new-token rows per request." + ) + if num_accepted_tokens.shape[0] < pool_size: + raise ValueError("Expected slot-indexed num_accepted_tokens to cover the state pool.") + else: + if cu_seqlens is None or cu_seqlens.ndim != 1: + raise ValueError("Expected cu_seqlens to be a 1D tensor.") + if cu_seqlens.shape[0] != num_requests + 1: + raise ValueError("Expected cu_seqlens length to be N + 1.") + if num_accepted_tokens.shape[0] != num_requests: + raise ValueError("Expected num_accepted_tokens length to match N.") def _fits_32bit_stride(tensor: torch.Tensor) -> bool: @@ -374,6 +606,19 @@ def _layout_key(tensor: torch.Tensor, dynamic_layout: bool = False, *, assumed_a _precompute_control_cache = {} +def _bf16_math_enabled() -> bool: + """Recurrence precision of the verify kernel, from ``TRTLLM_KDA_MTP_BF16_MATH``. + + Unset / 0 (default): 3xTF32 on every product that feeds the state. 1: bf16 tensor-core + operands (bf16 TMEM / smem tiles, K = 16 per instruction, fp32 accumulation -- the + prefill kernel's precision). + """ + value = os.environ.get("TRTLLM_KDA_MTP_BF16_MATH", "0") + if value not in ("0", "1"): + raise ValueError(f"TRTLLM_KDA_MTP_BF16_MATH must be 0 or 1, got {value!r}") + return value == "1" + + def _precompute_control_tensor(device: torch.device, enabled: bool) -> torch.Tensor: dev = torch.device(device) key = (dev.index if dev.index is not None else torch.cuda.current_device(), bool(enabled)) @@ -406,20 +651,6 @@ def _try_flatten_args( return True, h0, x_q_flat, x_k_flat, x_v_flat -def _is_benchmark_static_shape( - N: int, H: int, HV: int, K: int, V: int, W: int, num_spec: int -) -> bool: - return ( - K == 128 - and V == 128 - and W == 4 - and num_spec == 2 - and H == HV - and N in (32, 128) - and H in (2, 12, 32) - ) - - # Layout-and-constexpr-keyed compile cache. Request count and packed-token # length are dynamic; batches sharing the same kernel variant reuse one # artifact even when their launch grid and token-buffer extents differ. @@ -441,11 +672,12 @@ def kda_mtp_decode_impl( A_log: torch.Tensor, dt_bias: torch.Tensor, recurrent_state: torch.Tensor, - qkg_cache: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, v_cache: torch.Tensor, beta_cache: torch.Tensor, ssm_state_indices: torch.Tensor, - cu_seqlens: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], num_spec: int, num_accepted_tokens: torch.Tensor, lower_bound: float, @@ -453,39 +685,61 @@ def kda_mtp_decode_impl( out: Optional[torch.Tensor] = None, zero_accepted_hint: bool = False, regular_metadata_hint: bool = False, + packed_token_layout: bool = False, + onorm_g: Optional[torch.Tensor] = None, + onorm_weight: Optional[torch.Tensor] = None, + onorm_eps: float = 1e-5, + fuse_output_norm: bool = False, + quantize_output: bool = False, replay_only: bool = False, -) -> torch.Tensor: +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Launch the fused KDA MTP verify kernel. See the module docstring. Args (device tensors unless noted): x_q/x_k/x_v: post-projection, pre-conv token states ``[1, T_total, H, 128]`` bf16 — new tokens only, ``1 + - num_spec`` per request, packed per ``cu_seqlens``. + num_spec`` per request. Rows may be described by ``cu_seqlens`` + or packed uniformly by request. w_q/w_k/w_v: conv weights ``[H*128, W]`` fp32, width-contiguous. cs_q/cs_k/cs_v: extended conv caches ``[pool, H*128, >= W-1+M]`` fp32, dim-contiguous. Columns ``[0, W-1)`` are the committed window; tail columns hold raw pending-draft inputs. Mutated. g, beta: raw gate ``[1, T, H, 128]`` and beta ``[1, T, H]`` bf16. A_log, dt_bias: fp32 ``[H]`` / ``[H*128]``. - recurrent_state: pool ``[pool, H, V, K]`` fp32, **V-first** layout + recurrent_state: pool ``[pool, H, V, K]`` fp32 or bf16, **V-first** layout (matches the executor ssm pool and the single-token decode kernel). Committed in place. - qkg_cache/v_cache/beta_cache: replay caches ``[pool, M, 3, H*K]`` / - ``[pool, M, H*V]`` / ``[pool, M, H]`` fp32. Mutated. - ssm_state_indices / cu_seqlens / num_accepted_tokens: per-request - slot, token offsets ``[N+1]``, accepted-draft counts ``[N]``. + k_cache/g_cache/v_cache/beta_cache: replay caches ``[pool, M, H*K]`` (k, + fp32 or bf16) / ``[pool, M, H*K]`` (gate, fp32) / ``[pool, M, H*V]`` / + ``[pool, M, H]`` (fp32 or bf16). Mutated. + ssm_state_indices: per-request state-pool slots. On the legacy + metadata path, ``cu_seqlens`` supplies token offsets ``[N+1]`` + and accepted counts are request-indexed ``[N]``. On the packed + path, ``cu_seqlens`` is absent and accepted counts are pool-level. zero_accepted_hint: caller asserts every ``num_accepted_tokens`` is zero (compiles the smaller-smem no-replay variant). Wrong hints produce wrong results — pass True only when statically known. regular_metadata_hint: caller asserts ``cu_seqlens`` is the uniform ``arange * (2*num_spec+1)`` pattern and ``ssm_state_indices`` is ``arange(N)`` (benchmark identity layout). - replay_only: run only the replay steps and commit after them (see - ``kda_mtp_commit_pending_drafts``); the token inputs, weights - and ``out`` are not accessed. - - Returns the output ``[1, T_total, H, V]`` bf16 (rows for replayed - positions are zero; only new-token rows are written). + packed_token_layout: input projections contain exactly ``1 + + num_spec`` rows per request. ``cu_seqlens`` must be ``None`` and + ``num_accepted_tokens`` is indexed by state-pool slot. The + kernel derives every new-token row from the request index. + fuse_output_norm: on the packed path, apply gated RMSNorm in the CuTe + epilogue using section-strided ``onorm_g`` and ``onorm_weight``. + quantize_output: return the normalized activation directly as MXFP8 + E4M3 plus packed UE8M0 1x128 scales for a prequantized projection. + replay_only: run only the replay steps and commit the state after them + (see ``kda_mtp_commit_pending_drafts``); ``num_accepted_tokens`` + is indexed by request, and the token inputs, conv weights, gate + constants, conv caches, ``cu_seqlens`` and ``out`` are not + accessed. Takes the shifted-metadata layout without hints, output + norm or quantization. + + Returns either BF16 ``[1, T_total, H, V]`` or, with + ``quantize_output=True``, E4M3 ``[T_total, H*V]`` and packed UE8M0 scales + in the flat uint8 R128c4 layout consumed by the CuTe MXFP8 GEMM. """ _, T_total, _, D = x_q.shape H = A_log.shape[0] @@ -500,11 +754,56 @@ def kda_mtp_decode_impl( if scale is None: scale = K**-0.5 - N = cu_seqlens.shape[0] - 1 - if out is None: - out = torch.zeros(1, T_total, HV, V_dim, dtype=x_q.dtype, device=x_q.device) + if packed_token_layout and regular_metadata_hint: + raise ValueError("packed_token_layout and regular_metadata_hint are mutually exclusive.") + if replay_only and ( + packed_token_layout + or zero_accepted_hint + or regular_metadata_hint + or fuse_output_norm + or quantize_output + ): + raise ValueError( + "replay_only takes neither the packed layout, the metadata hints, " + "the output norm nor quantization" + ) + if packed_token_layout: + N = ssm_state_indices.shape[0] + elif cu_seqlens is None: + raise ValueError("cu_seqlens may be None only with packed_token_layout=True.") + else: + N = cu_seqlens.shape[0] - 1 + if quantize_output: + if out is not None: + raise ValueError("quantize_output does not accept a caller-provided out tensor.") + out = torch.empty( + (T_total, HV * V_dim), + dtype=torch.float8_e4m3fn, + device=x_q.device, + ) + scale_leading_dim = (T_total + 127) // 128 * 128 + output_scale = torch.empty( + (scale_leading_dim * HV * 4,), + dtype=torch.uint8, + device=x_q.device, + ) + elif out is None: + output_shape = (1, T_total, HV, V_dim) + if packed_token_layout: + # Packed outputs contain new-token rows only, and the kernel + # writes every head/value element. Avoid a redundant zero-fill. + out = torch.empty(output_shape, dtype=x_q.dtype, device=x_q.device) + else: + # Legacy layouts include replay-position holes that remain zero. + out = torch.zeros(output_shape, dtype=x_q.dtype, device=x_q.device) if num_accepted_tokens.dtype != torch.int32: num_accepted_tokens = num_accepted_tokens.to(torch.int32) + if not quantize_output: + # The kernel never reads output_scale without QUANTIZE_OUTPUT. Use the + # persistent 1-element control tensor as the placeholder: its static + # layout keeps the compile key independent of the request count N. + output_scale = _precompute_control_tensor(x_q.device, True) + scale_leading_dim = 1 _require_stride_layout( x_q=x_q, x_k=x_k, @@ -520,13 +819,16 @@ def kda_mtp_decode_impl( A_log=A_log, dt_bias=dt_bias, recurrent_state=recurrent_state, - qkg_cache=qkg_cache, + k_cache=k_cache, + g_cache=g_cache, v_cache=v_cache, beta_cache=beta_cache, ssm_state_indices=ssm_state_indices, cu_seqlens=cu_seqlens, num_accepted_tokens=num_accepted_tokens, out=out, + onorm_g=onorm_g, + onorm_weight=onorm_weight, H=H, HV=HV, K=K, @@ -534,10 +836,20 @@ def kda_mtp_decode_impl( W=W, num_spec=num_spec, T_total=T_total, + packed_token_layout=packed_token_layout, + fuse_output_norm=fuse_output_norm, + quantize_output=quantize_output, ) stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) precompute_control = _precompute_control_tensor(x_q.device, True) + # The packed variant proves at compile time that this tensor is never + # read. Reuse an existing metadata tensor as the CuTe placeholder. + cu_seqlens_arg = ssm_state_indices if cu_seqlens is None else cu_seqlens + # These tensors are compile-time dead in the raw-output variant. Reuse + # existing aligned arguments instead of allocating placeholders. + onorm_g_arg = onorm_g if onorm_g is not None else out + onorm_weight_arg = onorm_weight if onorm_weight is not None else A_log use_flat_layout, h0_arg, x_q_arg, x_k_arg, x_v_arg = _try_flatten_args( recurrent_state=recurrent_state, @@ -551,11 +863,6 @@ def kda_mtp_decode_impl( V=V_dim, ) pool_size = h0_arg.shape[0] - is_benchmark_static_shape = ( - _is_benchmark_static_shape(N, H, HV, K, V_dim, W, num_spec) and not replay_only - ) - use_setmaxreg = is_benchmark_static_shape - use_reg_q_weights = is_benchmark_static_shape use_regular_metadata = bool(regular_metadata_hint) # The kernel's USE_ZERO_ACCEPTED fast path unrolls exactly # 1 + NUM_SPEC == 3 new tokens, so it is only valid for num_spec == 2. @@ -563,8 +870,6 @@ def kda_mtp_decode_impl( # num_accepted_tokens is all zeros, so the generic path computes the # same result). use_zero_accepted = bool(zero_accepted_hint) and num_spec == 2 - if replay_only and (use_zero_accepted or use_regular_metadata): - raise ValueError("replay_only takes neither zero_accepted_hint nor regular_metadata_hint") # stage_timing is unused (PROFILE_STAGES=False); pass `out` as the # placeholder tensor argument like the drop's runner does. The alias is # only valid while profiling stays off: with PROFILE_STAGES=True the @@ -578,6 +883,9 @@ def kda_mtp_decode_impl( ) stage_timing_arg = out beta_cache_assumed_align = _beta_cache_assumed_align(beta_cache) + # Two CTAs per (head, request) while they fit in one wave (see SPLIT_MAX_TILES). + split_v = 2 if HV * N <= SPLIT_MAX_TILES else 1 + bf16_math = _bf16_math_enabled() key = ( x_q.dtype, @@ -604,18 +912,26 @@ def kda_mtp_decode_impl( _layout_key(g, dynamic_layout=True, assumed_align=16), _layout_key(dt_bias), _layout_key(beta, dynamic_layout=True, assumed_align=16), + _layout_key(onorm_g_arg, dynamic_layout=True, assumed_align=16), + _layout_key(onorm_weight_arg), _layout_key(out, dynamic_layout=True, assumed_align=16), - _layout_key(qkg_cache), + _layout_key(output_scale, assumed_align=16), + _layout_key(k_cache), + _layout_key(g_cache), _layout_key(v_cache), _layout_key(beta_cache, assumed_align=beta_cache_assumed_align), _layout_key(ssm_state_indices, dynamic_layout=True, assumed_align=4), - _layout_key(cu_seqlens, dynamic_layout=True, assumed_align=4), + _layout_key(cu_seqlens_arg, dynamic_layout=True, assumed_align=4), _layout_key(num_accepted_tokens, dynamic_layout=True, assumed_align=4), - use_setmaxreg, use_regular_metadata, - use_reg_q_weights, + bool(packed_token_layout), use_zero_accepted, - replay_only, + float(onorm_eps), + bool(fuse_output_norm), + bool(quantize_output), + split_v, + bf16_math, + bool(replay_only), ) if key not in _compiled_cache: @@ -623,34 +939,40 @@ def kda_mtp_decode_impl( f"kda_mtp_decode: compiling variant N={N} H={HV} T={T_total} " f"num_spec={num_spec} zero_accepted={use_zero_accepted} " f"regular_metadata={use_regular_metadata} " - f"static_shape={is_benchmark_static_shape} replay_only={replay_only}" + f"fuse_output_norm={fuse_output_norm} " + f"quantize_output={quantize_output} split_v={split_v} bf16_math={bf16_math} " + f"replay_only={replay_only}" ) _compiled_cache[key] = cute.compile( _run_kda_decode_mtp, - _from_dlpack_arg(h0_arg, assumed_align=16), + _from_dlpack_arg(h0_arg), _dlpack_arg(x_q_arg, assumed_align=16), _dlpack_arg(x_k_arg, assumed_align=16), _dlpack_arg(x_v_arg, assumed_align=16), - _from_dlpack_arg(w_q, assumed_align=16), - _from_dlpack_arg(w_k, assumed_align=16), - _from_dlpack_arg(w_v, assumed_align=16), - _from_dlpack_arg(cs_q, assumed_align=16), - _from_dlpack_arg(cs_k, assumed_align=16), - _from_dlpack_arg(cs_v, assumed_align=16), - _from_dlpack_arg(A_log, assumed_align=16), + _from_dlpack_arg(w_q), + _from_dlpack_arg(w_k), + _from_dlpack_arg(w_v), + _from_dlpack_arg(cs_q), + _from_dlpack_arg(cs_k), + _from_dlpack_arg(cs_v), + _from_dlpack_arg(A_log), _dlpack_arg(g, assumed_align=16), - _from_dlpack_arg(dt_bias, assumed_align=16), + _from_dlpack_arg(dt_bias), _dlpack_arg(beta, assumed_align=16), + _dlpack_arg(onorm_g_arg, assumed_align=16), + _from_dlpack_arg(onorm_weight_arg), _dlpack_arg(out, assumed_align=16), - _from_dlpack_arg(h0_arg, assumed_align=16), - _from_dlpack_arg(qkg_cache, assumed_align=16), - _from_dlpack_arg(v_cache, assumed_align=16), + _from_dlpack_arg(output_scale, assumed_align=16), + _from_dlpack_arg(h0_arg), + _from_dlpack_arg(k_cache), + _from_dlpack_arg(g_cache), + _from_dlpack_arg(v_cache), _from_dlpack_arg(beta_cache, assumed_align=beta_cache_assumed_align), _dlpack_arg(stage_timing_arg, assumed_align=16), _dlpack_arg(ssm_state_indices, assumed_align=4), - _dlpack_arg(cu_seqlens, assumed_align=4), + _dlpack_arg(cu_seqlens_arg, assumed_align=4), _dlpack_arg(num_accepted_tokens, assumed_align=4), - _from_dlpack_arg(precompute_control, assumed_align=16), + _from_dlpack_arg(precompute_control), scale=scale, HV=HV, K=K, @@ -660,18 +982,29 @@ def kda_mtp_decode_impl( TILE_V=_TILE_V, KERNEL_WIDTH=W, lower_bound=lower_bound, + onorm_eps=float(onorm_eps), + scale_leading_dim=scale_leading_dim, USE_FLAT_LAYOUT=use_flat_layout, - USE_SETMAXREG=use_setmaxreg, USE_REGULAR_METADATA=use_regular_metadata, - USE_REG_Q_WEIGHTS=use_reg_q_weights, + USE_PACKED_TOKEN_LAYOUT=bool(packed_token_layout), USE_ZERO_ACCEPTED=use_zero_accepted, FUSE_PRECOMPUTE=True, RUNTIME_PRECOMPUTE_FLAG=False, + FUSE_OUTPUT_NORM=bool(fuse_output_norm), + QUANTIZE_OUTPUT=bool(quantize_output), PROFILE_STAGES=profile_stages, - REPLAY_ONLY=replay_only, + SPLIT_V=split_v, + BF16_MMA=bf16_math, + REPLAY_ONLY=bool(replay_only), stream=stream, ) + # Runtime alignment policy: + # - 16 B: model inputs/outputs, parameters, recurrent/conv/kg/v pools, + # output-norm tensors, the persistent control scalar, and output_scale + # (separately allocated, or the control scalar as a placeholder). + # - beta_cache: derived from its actual per-layer physical span and dtype. + # - 4 B: int32 scheduler metadata, which may be an element-offset view. _compiled_cache[key]( _dlpack_arg(h0_arg, assumed_align=16), _dlpack_arg(x_q_arg, assumed_align=16), @@ -687,20 +1020,26 @@ def kda_mtp_decode_impl( _dlpack_arg(g, assumed_align=16), _dlpack_arg(dt_bias, assumed_align=16), _dlpack_arg(beta, assumed_align=16), + _dlpack_arg(onorm_g_arg, assumed_align=16), + _dlpack_arg(onorm_weight_arg, assumed_align=16), _dlpack_arg(out, assumed_align=16), + _dlpack_arg(output_scale, assumed_align=16), _dlpack_arg(h0_arg, assumed_align=16), - _dlpack_arg(qkg_cache, assumed_align=16), + _dlpack_arg(k_cache, assumed_align=16), + _dlpack_arg(g_cache, assumed_align=16), _dlpack_arg(v_cache, assumed_align=16), _dlpack_arg(beta_cache, assumed_align=beta_cache_assumed_align), _dlpack_arg(stage_timing_arg, assumed_align=16), _dlpack_arg(ssm_state_indices, assumed_align=4), - _dlpack_arg(cu_seqlens, assumed_align=4), + _dlpack_arg(cu_seqlens_arg, assumed_align=4), _dlpack_arg(num_accepted_tokens, assumed_align=4), _dlpack_arg(precompute_control, assumed_align=16), N, stream, ) + if quantize_output: + return out, output_scale return out @@ -729,7 +1068,8 @@ def unused(*shape, dtype=torch.bfloat16): "A_log": unused(H, dtype=torch.float32), "dt_bias": unused(H * K, dtype=torch.float32), "out": unused(1, N, H, V), - # One row per request keeps every request's [bos, eos) non-empty. + # One row per request keeps the shifted-metadata validation's [bos, eos) non-empty; + # the replay-only kernel itself does not read it. "cu_seqlens": torch.arange(N + 1, dtype=torch.int32, device=dev), } return _replay_only_inputs_cache[key] @@ -740,7 +1080,8 @@ def kda_mtp_commit_pending_drafts( cs_q: torch.Tensor, cs_k: torch.Tensor, cs_v: torch.Tensor, - qkg_cache: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, v_cache: torch.Tensor, beta_cache: torch.Tensor, ssm_state_indices: torch.Tensor, @@ -750,13 +1091,20 @@ def kda_mtp_commit_pending_drafts( ) -> None: """Commit the pending accepted drafts of the selected slots in place. - For request ``n`` this runs the first ``num_accepted_tokens[n]`` steps of - a ``kda_mtp_decode`` verify of slot ``ssm_state_indices[n]``, the replay - of its cached drafts, with the same kernel code, and then commits the - recurrent state and the base conv windows. The slot then holds what the - next verify would hold just before its golden token, with nothing left - pending; the caller sets the slot's accepted-draft count to 0. The - tensors are those of ``kda_mtp_decode``; the draft caches are only read. + For request ``n`` this runs the first ``num_accepted_tokens[n]`` rows of a + ``kda_mtp_decode`` verify chunk of slot ``ssm_state_indices[n]``, the + replay of its cached drafts, with the same kernel code, and commits the + recurrent state after them. The slot then holds the state the next verify + would reach just before its golden token, with nothing left pending; the + caller commits the conv window after the drafts and sets the slot's + accepted-draft count to 0. The tensors are those of ``kda_mtp_decode`` + (``k_cache`` / ``v_cache`` / ``beta_cache`` bf16 or fp32, ``g_cache`` + fp32); the conv and replay caches are only read. + + The tensor-core verify computes each chunk in one WY-form pass, so the + verify that follows this commit (a chunk of the new tokens only) rounds + differently from one that replays the drafts itself: the two agree to the + kernel's accuracy, not bit for bit. """ N = ssm_state_indices.shape[0] if N == 0: @@ -778,7 +1126,8 @@ def kda_mtp_commit_pending_drafts( A_log=inputs["A_log"], dt_bias=inputs["dt_bias"], recurrent_state=recurrent_state, - qkg_cache=qkg_cache, + k_cache=k_cache, + g_cache=g_cache, v_cache=v_cache, beta_cache=beta_cache, ssm_state_indices=ssm_state_indices, @@ -801,7 +1150,8 @@ def kda_mtp_commit_pending_drafts( "cs_k", "cs_v", "recurrent_state", - "qkg_cache", + "k_cache", + "g_cache", "v_cache", "beta_cache", ), @@ -822,17 +1172,23 @@ def kda_mtp_decode( A_log: torch.Tensor, dt_bias: torch.Tensor, recurrent_state: torch.Tensor, - qkg_cache: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, v_cache: torch.Tensor, beta_cache: torch.Tensor, ssm_state_indices: torch.Tensor, - cu_seqlens: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], num_spec: int, num_accepted_tokens: torch.Tensor, lower_bound: float, scale: Optional[float] = None, zero_accepted_hint: bool = False, regular_metadata_hint: bool = False, + packed_token_layout: bool = False, + onorm_g: Optional[torch.Tensor] = None, + onorm_weight: Optional[torch.Tensor] = None, + onorm_eps: float = 1e-5, + fuse_output_norm: bool = False, ) -> torch.Tensor: """Fused KDA multi-token verify with in-place state commit.""" return kda_mtp_decode_impl( @@ -850,7 +1206,8 @@ def kda_mtp_decode( A_log=A_log, dt_bias=dt_bias, recurrent_state=recurrent_state, - qkg_cache=qkg_cache, + k_cache=k_cache, + g_cache=g_cache, v_cache=v_cache, beta_cache=beta_cache, ssm_state_indices=ssm_state_indices, @@ -861,6 +1218,11 @@ def kda_mtp_decode( scale=scale, zero_accepted_hint=zero_accepted_hint, regular_metadata_hint=regular_metadata_hint, + packed_token_layout=packed_token_layout, + onorm_g=onorm_g, + onorm_weight=onorm_weight, + onorm_eps=onorm_eps, + fuse_output_norm=fuse_output_norm, ) @kda_mtp_decode.register_fake @@ -879,16 +1241,179 @@ def _( A_log: torch.Tensor, dt_bias: torch.Tensor, recurrent_state: torch.Tensor, - qkg_cache: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, v_cache: torch.Tensor, beta_cache: torch.Tensor, ssm_state_indices: torch.Tensor, - cu_seqlens: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], num_spec: int, num_accepted_tokens: torch.Tensor, lower_bound: float, scale: Optional[float] = None, zero_accepted_hint: bool = False, regular_metadata_hint: bool = False, + packed_token_layout: bool = False, + onorm_g: Optional[torch.Tensor] = None, + onorm_weight: Optional[torch.Tensor] = None, + onorm_eps: float = 1e-5, + fuse_output_norm: bool = False, ) -> torch.Tensor: + del onorm_g, onorm_weight, onorm_eps, fuse_output_norm return x_q.new_empty(x_v.shape) + + @torch.library.custom_op( + "trtllm::kda_mtp_decode_fp8", + mutates_args=( + "cs_q", + "cs_k", + "cs_v", + "recurrent_state", + "k_cache", + "g_cache", + "v_cache", + "beta_cache", + ), + device_types="cuda", + ) + def kda_mtp_decode_fp8( + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, + v_cache: torch.Tensor, + beta_cache: torch.Tensor, + ssm_state_indices: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], + num_spec: int, + num_accepted_tokens: torch.Tensor, + lower_bound: float, + scale: Optional[float] = None, + zero_accepted_hint: bool = False, + regular_metadata_hint: bool = False, + packed_token_layout: bool = False, + onorm_g: Optional[torch.Tensor] = None, + onorm_weight: Optional[torch.Tensor] = None, + onorm_eps: float = 1e-5, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Fused KDA verify, gated RMSNorm, and 1x128 MXFP8 quantization.""" + return kda_mtp_decode_impl( + x_q=x_q, + x_k=x_k, + x_v=x_v, + w_q=w_q, + w_k=w_k, + w_v=w_v, + cs_q=cs_q, + cs_k=cs_k, + cs_v=cs_v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + k_cache=k_cache, + g_cache=g_cache, + v_cache=v_cache, + beta_cache=beta_cache, + ssm_state_indices=ssm_state_indices, + cu_seqlens=cu_seqlens, + num_spec=num_spec, + num_accepted_tokens=num_accepted_tokens, + lower_bound=lower_bound, + scale=scale, + zero_accepted_hint=zero_accepted_hint, + regular_metadata_hint=regular_metadata_hint, + packed_token_layout=packed_token_layout, + onorm_g=onorm_g, + onorm_weight=onorm_weight, + onorm_eps=onorm_eps, + fuse_output_norm=True, + quantize_output=True, + ) + + @kda_mtp_decode_fp8.register_fake + def _( + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, + v_cache: torch.Tensor, + beta_cache: torch.Tensor, + ssm_state_indices: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], + num_spec: int, + num_accepted_tokens: torch.Tensor, + lower_bound: float, + scale: Optional[float] = None, + zero_accepted_hint: bool = False, + regular_metadata_hint: bool = False, + packed_token_layout: bool = False, + onorm_g: Optional[torch.Tensor] = None, + onorm_weight: Optional[torch.Tensor] = None, + onorm_eps: float = 1e-5, + ) -> Tuple[torch.Tensor, torch.Tensor]: + del ( + x_k, + w_q, + w_k, + w_v, + cs_q, + cs_k, + cs_v, + g, + beta, + A_log, + dt_bias, + recurrent_state, + k_cache, + g_cache, + v_cache, + beta_cache, + ssm_state_indices, + cu_seqlens, + num_spec, + num_accepted_tokens, + lower_bound, + scale, + zero_accepted_hint, + regular_metadata_hint, + packed_token_layout, + onorm_g, + onorm_weight, + onorm_eps, + ) + rows = x_q.shape[1] + heads = x_v.shape[2] + scale_leading_dim = (rows + 127) // 128 * 128 + output = x_q.new_empty((rows, heads * x_v.shape[3]), dtype=torch.float8_e4m3fn) + output_scale = torch.empty( + (scale_leading_dim * heads * 4,), + dtype=torch.uint8, + device=x_q.device, + ) + return output, output_scale diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py index 70365b95d87a..c29cba4ea6c4 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py @@ -1686,13 +1686,44 @@ def split_kv_kernel( mAccLSE=mAccLSE, mLSE=mLSE, ) - p_cor_consumer_state, mma_o_consumer_state = self.correction( - compute_common_params, - compute_epilogue_params, - k_tile_count=k_tile_count, - p_cor_consumer_state=p_cor_consumer_state, - mma_o_consumer_state=mma_o_consumer_state, - ) + if cutlass.const_expr(self.is_var_split_kv + and mAccO is not None): + # A request planned onto a single split stores its + # final output (and LSE / Helix stats) directly, as + # the split_kv == 1 build does; the reduction kernel + # skips its rows. Other requests write partials. + direct_common_params = SimpleNamespace( + **{ + **vars(compute_common_params), "mAccO": None + }) + direct_epilogue_params = SimpleNamespace( + **{ + **vars(compute_epilogue_params), "mAccLSE": None + }) + if local_split_kv == 1: + p_cor_consumer_state, mma_o_consumer_state = self.correction( + direct_common_params, + direct_epilogue_params, + k_tile_count=k_tile_count, + p_cor_consumer_state=p_cor_consumer_state, + mma_o_consumer_state=mma_o_consumer_state, + ) + else: + p_cor_consumer_state, mma_o_consumer_state = self.correction( + compute_common_params, + compute_epilogue_params, + k_tile_count=k_tile_count, + p_cor_consumer_state=p_cor_consumer_state, + mma_o_consumer_state=mma_o_consumer_state, + ) + else: + p_cor_consumer_state, mma_o_consumer_state = self.correction( + compute_common_params, + compute_epilogue_params, + k_tile_count=k_tile_count, + p_cor_consumer_state=p_cor_consumer_state, + mma_o_consumer_state=mma_o_consumer_state, + ) tile_sched.advance_to_next_work() work_tile = tile_sched.get_current_work() @@ -1794,99 +1825,111 @@ def reduction_kernel( q_tile_row = flat_q_row & 127 local_split_kv = split_kv if cutlass.const_expr(self.is_var_split_kv): - local_split_kv = block_split_kvs[bidz] + local_split_kv = cutlass.max( + cutlass.min(block_split_kvs[bidz], split_kv), 1) k_tile_total = cute.ceil_div(cache_seqs[bidz], self.mma_qk_tiler[1]) - k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) - if cutlass.const_expr(kv_bounds is not None): - k_tile_per_cta = cutlass.max(k_tile_per_cta, 1) + # A request without KV has no K tiles; keep the divisor nonzero. + k_tile_per_cta = cutlass.max( + cute.ceil_div(k_tile_total, local_split_kv), 1) local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) - gLSE = mAccLSE[q_tile_row, None, q_tile, bidz] - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - if warp_idx == 0: - # calculate the global lse and exp ^ (local_lse - global_lse) - lse_per_thread = cute.ceil_div(self.reducer_max_splits, - self.threads_per_warp) - - local_lse = cute.make_rmem_tensor( - cute.make_layout(lse_per_thread), self.lse_dtype) - lse_max = -self.lse_dtype.inf - # find the max lse - for i in cutlass.range_constexpr(lse_per_thread): - split_kv_idx = tidx + i * self.threads_per_warp - local_lse[i] = (gLSE[split_kv_idx] if cute.elem_less( - split_kv_idx, local_split_kv) else -self.lse_dtype.inf) - # reduce the local lse - lse_max = cute.arch.fmax(lse_max, local_lse[i]) - lse_max = cute.arch.warp_reduction_max(lse_max) - lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 - # calculate sum_lse - sum_lse = 0.0 - for i in cutlass.range_constexpr(lse_per_thread): - sum_lse += cute.math.exp2(local_lse[i] - lse_max, - fastmath=True) - sum_lse = cute.arch.warp_reduction_sum(sum_lse) - # calculate the global_lse - global_lse = (lse_max + cute.math.log2(sum_lse, fastmath=True) - if not sum_lse == self.lse_dtype(0.0) - or sum_lse != sum_lse else self.lse_dtype.inf) - if d_tile_idx == 0 and tidx == 0: - if cutlass.const_expr(self.is_var_q): - mLSE[head_idx, q_begin + bidy] = global_lse - if cutlass.const_expr(self.emit_softmax_stats): - if has_local_kv: - mSoftmaxStats[head_idx, q_begin + bidy, - 0] = (global_lse / LOG2_E) - mSoftmaxStats[head_idx, q_begin + bidy, 1] = 1.0 - else: - mSoftmaxStats[head_idx, q_begin + bidy, - 0] = -self.lse_dtype.inf - mSoftmaxStats[head_idx, q_begin + bidy, 1] = 0.0 - else: - mLSE[head_idx, bidy, bidz] = global_lse - if cutlass.const_expr(self.emit_softmax_stats): - if has_local_kv: - mSoftmaxStats[head_idx, bidy, bidz, - 0] = (global_lse / LOG2_E) - mSoftmaxStats[head_idx, bidy, bidz, 1] = 1.0 - else: - mSoftmaxStats[head_idx, bidy, bidz, - 0] = -self.lse_dtype.inf - mSoftmaxStats[head_idx, bidy, bidz, 1] = 0.0 - # store the scale to shared memory - for i in cutlass.range_constexpr(lse_per_thread): - split_kv_idx = tidx + i * self.threads_per_warp - if cute.elem_less(split_kv_idx, local_split_kv): - smem_lse_scale[split_kv_idx] = cute.math.exp2( - local_lse[i] - global_lse, fastmath=True) - - pipeline.sync(barrier_id=4) - - elements_per_thread = cute.ceil_div( - self.reducer_d_tile, - self.threads_per_warp * self.num_compute_warps, - ) - gAccO = mAccO[q_tile_row, None, None, q_tile, bidz] - rAccO = cute.make_rmem_tensor(cute.make_layout(elements_per_thread), - self.acc_dtype) - rO = cute.make_rmem_tensor(cute.make_layout(elements_per_thread), - self.o_dtype) - rAccO.fill(0.0) - for i in range(local_split_kv): + # A request planned onto a single split stored its output + # directly in the split kernel (variable split-KV only); K == 0 + # rows still take the canonical empty result from here. + reduce_row = cutlass.Boolean(True) + if cutlass.const_expr(self.is_var_split_kv): + reduce_row = local_split_kv != 1 + if reduce_row: + gLSE = mAccLSE[q_tile_row, None, q_tile, bidz] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 0: + # calculate the global lse and exp ^ (local_lse - global_lse) + lse_per_thread = cute.ceil_div(self.reducer_max_splits, + self.threads_per_warp) + + local_lse = cute.make_rmem_tensor( + cute.make_layout(lse_per_thread), self.lse_dtype) + lse_max = -self.lse_dtype.inf + # find the max lse + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + local_lse[i] = (gLSE[split_kv_idx] if cute.elem_less( + split_kv_idx, local_split_kv) else + -self.lse_dtype.inf) + # reduce the local lse + lse_max = cute.arch.fmax(lse_max, local_lse[i]) + lse_max = cute.arch.warp_reduction_max(lse_max) + lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 + # calculate sum_lse + sum_lse = 0.0 + for i in cutlass.range_constexpr(lse_per_thread): + sum_lse += cute.math.exp2(local_lse[i] - lse_max, + fastmath=True) + sum_lse = cute.arch.warp_reduction_sum(sum_lse) + # calculate the global_lse + global_lse = (lse_max + + cute.math.log2(sum_lse, fastmath=True) + if not sum_lse == self.lse_dtype(0.0) + or sum_lse != sum_lse else self.lse_dtype.inf) + if d_tile_idx == 0 and tidx == 0: + if cutlass.const_expr(self.is_var_q): + mLSE[head_idx, q_begin + bidy] = global_lse + if cutlass.const_expr(self.emit_softmax_stats): + if has_local_kv: + mSoftmaxStats[head_idx, q_begin + bidy, + 0] = (global_lse / LOG2_E) + mSoftmaxStats[head_idx, q_begin + bidy, + 1] = 1.0 + else: + mSoftmaxStats[head_idx, q_begin + bidy, + 0] = -self.lse_dtype.inf + mSoftmaxStats[head_idx, q_begin + bidy, + 1] = 0.0 + else: + mLSE[head_idx, bidy, bidz] = global_lse + if cutlass.const_expr(self.emit_softmax_stats): + if has_local_kv: + mSoftmaxStats[head_idx, bidy, bidz, + 0] = (global_lse / LOG2_E) + mSoftmaxStats[head_idx, bidy, bidz, 1] = 1.0 + else: + mSoftmaxStats[head_idx, bidy, bidz, + 0] = -self.lse_dtype.inf + mSoftmaxStats[head_idx, bidy, bidz, 1] = 0.0 + # store the scale to shared memory + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + if cute.elem_less(split_kv_idx, local_split_kv): + smem_lse_scale[split_kv_idx] = cute.math.exp2( + local_lse[i] - global_lse, fastmath=True) + + pipeline.sync(barrier_id=4) + + elements_per_thread = cute.ceil_div( + self.reducer_d_tile, + self.threads_per_warp * self.num_compute_warps, + ) + gAccO = mAccO[q_tile_row, None, None, q_tile, bidz] + rAccO = cute.make_rmem_tensor( + cute.make_layout(elements_per_thread), self.acc_dtype) + rO = cute.make_rmem_tensor( + cute.make_layout(elements_per_thread), self.o_dtype) + rAccO.fill(0.0) + for i in range(local_split_kv): + for j in cutlass.range_constexpr(elements_per_thread): + element_idx = ( + d_tile_idx * self.reducer_d_tile + tidx + + j * self.threads_per_warp * self.num_compute_warps) + rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i] + rO.store(rAccO.load().to(self.o_dtype)) for j in cutlass.range_constexpr(elements_per_thread): element_idx = ( d_tile_idx * self.reducer_d_tile + tidx + j * self.threads_per_warp * self.num_compute_warps) - rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i] - rO.store(rAccO.load().to(self.o_dtype)) - for j in cutlass.range_constexpr(elements_per_thread): - element_idx = ( - d_tile_idx * self.reducer_d_tile + tidx + - j * self.threads_per_warp * self.num_compute_warps) - if cutlass.const_expr(self.is_var_q): - mO[head_idx, element_idx, q_begin + bidy] = rO[j] - else: - mO[head_idx, element_idx, bidy, bidz] = rO[j] + if cutlass.const_expr(self.is_var_q): + mO[head_idx, element_idx, q_begin + bidy] = rO[j] + else: + mO[head_idx, element_idx, bidy, bidz] = rO[j] return @staticmethod @@ -1964,7 +2007,11 @@ def get_k_tile_count( """ K = cache_seqs[blk_coord[2]] if cutlass.const_expr(self.is_var_split_kv): - split_kv = block_split_kvs[blk_coord[2]] + # Per-request split planned on the device for the batch's KV + # lengths; clamped to the launched grid so the plan can be made + # for the largest split any tactic may use. + split_kv = cutlass.max( + cutlass.min(block_split_kvs[blk_coord[2]], split_kv), 1) k_tile_total = cute.ceil_div(K, self.mma_qk_tiler[1]) k_tile_per_cta = cute.ceil_div(k_tile_total, split_kv) @@ -1972,6 +2019,12 @@ def get_k_tile_count( k_tile_count = max( 0, min(k_tile_total, k_index + k_tile_per_cta) - k_index) + if cutlass.const_expr(self.is_var_split_kv): + # The split the request really uses after tile rounding (the + # reduction kernel derives the same value): one means this CTA + # stores the final output itself and the reduction skips the row. + split_kv = cute.ceil_div(k_tile_total, + cutlass.max(k_tile_per_cta, 1)) q_begin = cutlass.Int32(0) q_len = cutlass.Int32(self.seq_len_q) if cutlass.const_expr(self.is_var_q): diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py index cc6755dc3932..7a750bd38846 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py @@ -371,8 +371,8 @@ def __call__( kv_bounds: Optional[cute.Tensor], cum_seq_lens_q: Optional[cute.Tensor], block_split_kvs: Optional[cute.Tensor], - softmax_scale: cutlass.Float32, - output_scale: cutlass.Float32, + softmax_scale: cute.Tensor, + output_scale: cute.Tensor, stream: cuda.CUstream, ): self._run( @@ -412,8 +412,8 @@ def run_with_softmax_stats( kv_bounds: Optional[cute.Tensor], cum_seq_lens_q: Optional[cute.Tensor], block_split_kvs: Optional[cute.Tensor], - softmax_scale: cutlass.Float32, - output_scale: cutlass.Float32, + softmax_scale: cute.Tensor, + output_scale: cute.Tensor, stream: cuda.CUstream, ): self._run( @@ -453,8 +453,8 @@ def _run( kv_bounds: Optional[cute.Tensor], cum_seq_lens_q: Optional[cute.Tensor], block_split_kvs: Optional[cute.Tensor], - softmax_scale: cutlass.Float32, - output_scale: cutlass.Float32, + softmax_scale: cute.Tensor, + output_scale: cute.Tensor, stream: cuda.CUstream, ): """Execute the Multi-Head Latent Attention operation on the provided tensors. @@ -508,10 +508,10 @@ def _run( :type cum_seq_lens_q: cute.Tensor :param block_split_kvs: The block split KV tensor with shape [batch_size] :type block_split_kvs: cute.Tensor - :param softmax_scale: The scale factor for softmax - :type softmax_scale: cutlass.Float32 - :param output_scale: The scale factor for the output - :type output_scale: cutlass.Float32 + :param softmax_scale: Device tensor containing the scale factor for softmax + :type softmax_scale: cute.Tensor + :param output_scale: Device tensor containing the scale factor for the output + :type output_scale: cute.Tensor :param stream: The CUDA stream to execute the kernel on :type stream: cuda.CUstream @@ -1037,8 +1037,6 @@ class SplitKVKernelSharedStorage: # Tmem holding buffer tmem_holding_buf: cutlass.Int32 - softmax_scale_log2 = softmax_scale * LOG2_E - if cutlass.const_expr(self.emit_softmax_stats): split_kv_kernel = self.split_kv_kernel( qk_tiled_mma, @@ -1063,7 +1061,7 @@ class SplitKVKernelSharedStorage: kv_bounds, cum_seq_lens_q, block_split_kvs, - softmax_scale_log2, + softmax_scale, output_scale, q_latent_smem_layout_staged, q_rope_smem_layout_staged, @@ -1102,7 +1100,7 @@ class SplitKVKernelSharedStorage: kv_bounds, cum_seq_lens_q, block_split_kvs, - softmax_scale_log2, + softmax_scale, output_scale, q_latent_smem_layout_staged, q_rope_smem_layout_staged, @@ -1225,8 +1223,8 @@ def split_kv_kernel( kv_bounds: Optional[cute.Tensor], cum_seq_lens_q: Optional[cute.Tensor], block_split_kvs: cute.Tensor, - softmax_scale_log2: cutlass.Float32, - output_scale: cutlass.Float32, + softmax_scale: cute.Tensor, + output_scale: cute.Tensor, q_latent_smem_layout_staged: cute.ComposedLayout, q_rope_smem_layout_staged: cute.ComposedLayout, kc_latent_smem_layout_staged: cute.ComposedLayout, @@ -1292,10 +1290,10 @@ def split_kv_kernel( :type cache_seqs: cute.Tensor :param block_split_kvs: The per-block split_kv values tensor :type block_split_kvs: cute.Tensor - :param softmax_scale_log2: The log2 scale factor for softmax - :type softmax_scale_log2: cutlass.Float32 - :param output_scale: The scale factor for the output - :type output_scale: cutlass.Float32 + :param softmax_scale: Device tensor containing the scale factor for softmax + :type softmax_scale: cute.Tensor + :param output_scale: Device tensor containing the scale factor for the output + :type output_scale: cute.Tensor :param q_latent_smem_layout_staged: Shared memory layout for query tensor :type q_latent_smem_layout_staged: cute.ComposedLayout :param q_rope_smem_layout_staged: Shared memory layout for query rope tensor @@ -1324,6 +1322,8 @@ def split_kv_kernel( tidx, _, _ = cute.arch.thread_idx() bidx, _, _ = cute.arch.block_idx() + softmax_scale_log2 = softmax_scale[0] * LOG2_E + output_scale_value = output_scale[0] mma_tile_coord_v = bidx % cute.size(tiled_mma_qk.thr_id.shape) is_leader_cta = mma_tile_coord_v == 0 @@ -1751,8 +1751,39 @@ def split_kv_kernel( p_mma_pipeline=p_mma_pipeline, softmax_scale_log2=softmax_scale_log2, ) - mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state = ( - self.compute( + if cutlass.const_expr(self.is_var_split_kv + and mAccO is not None): + # A request planned onto a single split runs the + # non-split softmax, matching the direct-store + # correction it exchanges metadata with. + direct_common_params = SimpleNamespace( + **{ + **vars(compute_common_params), "mAccO": None + }) + if local_split_kv == 1: + mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state = self.compute( + direct_common_params, + compute_softmax_params, + k_index=k_index, + k_tile_count=k_tile_count, + mma_s_consumer_state=mma_s_consumer_state, + p_mma_producer_state=p_mma_producer_state, + p_cor_producer_state=p_cor_producer_state, + is_second_compute_warp=False, + ) + else: + mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state = self.compute( + compute_common_params, + compute_softmax_params, + k_index=k_index, + k_tile_count=k_tile_count, + mma_s_consumer_state=mma_s_consumer_state, + p_mma_producer_state=p_mma_producer_state, + p_cor_producer_state=p_cor_producer_state, + is_second_compute_warp=False, + ) + else: + mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state = self.compute( compute_common_params, compute_softmax_params, k_index=k_index, @@ -1761,7 +1792,7 @@ def split_kv_kernel( p_mma_producer_state=p_mma_producer_state, p_cor_producer_state=p_cor_producer_state, is_second_compute_warp=False, - )) + ) tile_sched.advance_to_next_work() work_tile = tile_sched.get_current_work() @@ -1833,16 +1864,48 @@ def split_kv_kernel( p_mma_pipeline=p_mma_pipeline, softmax_scale_log2=softmax_scale_log2, ) - mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state = self.compute( - compute_common_params, - compute_softmax_params, - k_index=k_index, - k_tile_count=k_tile_count, - mma_s_consumer_state=mma_s_consumer_state, - p_mma_producer_state=p_mma_producer_state, - p_cor_producer_state=p_cor_producer_state, - is_second_compute_warp=True, - ) + if cutlass.const_expr(self.is_var_split_kv + and mAccO is not None): + # A request planned onto a single split runs the + # non-split softmax, matching the direct-store + # correction it exchanges metadata with. + direct_common_params = SimpleNamespace( + **{ + **vars(compute_common_params), "mAccO": None + }) + if local_split_kv == 1: + mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state = self.compute( + direct_common_params, + compute_softmax_params, + k_index=k_index, + k_tile_count=k_tile_count, + mma_s_consumer_state=mma_s_consumer_state, + p_mma_producer_state=p_mma_producer_state, + p_cor_producer_state=p_cor_producer_state, + is_second_compute_warp=True, + ) + else: + mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state = self.compute( + compute_common_params, + compute_softmax_params, + k_index=k_index, + k_tile_count=k_tile_count, + mma_s_consumer_state=mma_s_consumer_state, + p_mma_producer_state=p_mma_producer_state, + p_cor_producer_state=p_cor_producer_state, + is_second_compute_warp=True, + ) + else: + mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state = self.compute( + compute_common_params, + compute_softmax_params, + k_index=k_index, + k_tile_count=k_tile_count, + mma_s_consumer_state=mma_s_consumer_state, + p_mma_producer_state=p_mma_producer_state, + p_cor_producer_state=p_cor_producer_state, + is_second_compute_warp=True, + ) tile_sched.advance_to_next_work() work_tile = tile_sched.get_current_work() # Neither compute group calls p_cor_pipeline.producer_tail(). @@ -1904,7 +1967,7 @@ def split_kv_kernel( ) if cutlass.const_expr(self.emit_softmax_stats): compute_epilogue_params = SimpleNamespace( - output_scale=output_scale, + output_scale=output_scale_value, softmax_scale_log2=softmax_scale_log2, mAccLSE=mAccLSE, mLSE=mLSE, @@ -1912,18 +1975,49 @@ def split_kv_kernel( ) else: compute_epilogue_params = SimpleNamespace( - output_scale=output_scale, + output_scale=output_scale_value, softmax_scale_log2=softmax_scale_log2, mAccLSE=mAccLSE, mLSE=mLSE, ) - p_cor_consumer_state, mma_o_consumer_state = self.correction( - compute_common_params, - compute_epilogue_params, - k_tile_count=k_tile_count, - p_cor_consumer_state=p_cor_consumer_state, - mma_o_consumer_state=mma_o_consumer_state, - ) + if cutlass.const_expr(self.is_var_split_kv + and mAccO is not None): + # A request planned onto a single split stores its + # final output (and LSE / Helix stats) directly, as + # the split_kv == 1 build does; the reduction kernel + # skips its rows. Other requests write partials. + direct_common_params = SimpleNamespace( + **{ + **vars(compute_common_params), "mAccO": None + }) + direct_epilogue_params = SimpleNamespace( + **{ + **vars(compute_epilogue_params), "mAccLSE": None + }) + if local_split_kv == 1: + p_cor_consumer_state, mma_o_consumer_state = self.correction( + direct_common_params, + direct_epilogue_params, + k_tile_count=k_tile_count, + p_cor_consumer_state=p_cor_consumer_state, + mma_o_consumer_state=mma_o_consumer_state, + ) + else: + p_cor_consumer_state, mma_o_consumer_state = self.correction( + compute_common_params, + compute_epilogue_params, + k_tile_count=k_tile_count, + p_cor_consumer_state=p_cor_consumer_state, + mma_o_consumer_state=mma_o_consumer_state, + ) + else: + p_cor_consumer_state, mma_o_consumer_state = self.correction( + compute_common_params, + compute_epilogue_params, + k_tile_count=k_tile_count, + p_cor_consumer_state=p_cor_consumer_state, + mma_o_consumer_state=mma_o_consumer_state, + ) tile_sched.advance_to_next_work() work_tile = tile_sched.get_current_work() return @@ -2004,105 +2098,117 @@ def reduction_kernel( q_tile_row = flat_q_row & 127 local_split_kv = split_kv if cutlass.const_expr(self.is_var_split_kv): - local_split_kv = block_split_kvs[bidz] + local_split_kv = cutlass.max( + cutlass.min(block_split_kvs[bidz], split_kv), 1) k_tile_total = cute.ceil_div(cache_seqs[bidz], self.mma_qk_tiler[1]) - k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) - if cutlass.const_expr(kv_bounds is not None): - k_tile_per_cta = cutlass.max(k_tile_per_cta, 1) + # A request without KV has no K tiles; keep the divisor nonzero. + k_tile_per_cta = cutlass.max( + cute.ceil_div(k_tile_total, local_split_kv), 1) local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) - gLSE = mAccLSE[q_tile_row, None, q_tile, bidz] - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - if warp_idx == 0: - # calculate the global lse and exp ^ (local_lse - global_lse) - lse_per_thread = cute.ceil_div(self.reducer_max_splits, - self.threads_per_warp) - - local_lse = cute.make_rmem_tensor( - cute.make_layout(lse_per_thread), self.lse_dtype) - lse_max = -self.lse_dtype.inf - # find the max lse - for i in cutlass.range_constexpr(lse_per_thread): - split_kv_idx = tidx + i * self.threads_per_warp - local_lse[i] = (gLSE[split_kv_idx] if cute.elem_less( - split_kv_idx, local_split_kv) else -self.lse_dtype.inf) - # reduce the local lse - lse_max = cute.arch.fmax(lse_max, local_lse[i]) - lse_max = cute.arch.warp_reduction_max(lse_max) - lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 - # calculate sum_lse - sum_lse = 0.0 - for i in cutlass.range_constexpr(lse_per_thread): - sum_lse += cute.math.exp2(local_lse[i] - lse_max, - fastmath=True) - sum_lse = cute.arch.warp_reduction_sum(sum_lse) - # calculate the global_lse - global_lse = (lse_max + cute.math.log2(sum_lse, fastmath=True) - if not sum_lse == self.lse_dtype(0.0) - or sum_lse != sum_lse else self.lse_dtype.inf) - if d_tile_idx == 0 and tidx == 0: - if cutlass.const_expr(self.is_var_q): - mLSE[head_idx, q_begin + bidy] = global_lse - if cutlass.const_expr(self.emit_softmax_stats): - if has_local_kv: - mSoftmaxStats[head_idx, q_begin + bidy, - 0] = (global_lse / LOG2_E) - mSoftmaxStats[head_idx, q_begin + bidy, 1] = 1.0 - else: - mSoftmaxStats[head_idx, q_begin + bidy, - 0] = -self.lse_dtype.inf - mSoftmaxStats[head_idx, q_begin + bidy, 1] = 0.0 - else: - mLSE[head_idx, bidy, bidz] = global_lse - if cutlass.const_expr(self.emit_softmax_stats): - if has_local_kv: - mSoftmaxStats[head_idx, bidy, bidz, - 0] = (global_lse / LOG2_E) - mSoftmaxStats[head_idx, bidy, bidz, 1] = 1.0 - else: - mSoftmaxStats[head_idx, bidy, bidz, - 0] = -self.lse_dtype.inf - mSoftmaxStats[head_idx, bidy, bidz, 1] = 0.0 - # store the scale to shared memory - for i in cutlass.range_constexpr(lse_per_thread): - split_kv_idx = tidx + i * self.threads_per_warp - if cute.elem_less(split_kv_idx, local_split_kv): - smem_lse_scale[split_kv_idx] = cute.math.exp2( - local_lse[i] - global_lse, fastmath=True) - - pipeline.sync(barrier_id=4) - - elements_per_thread = cute.ceil_div( - self.reducer_d_tile, - self.threads_per_warp * self.num_compute_warps, - ) - # FP8 register fragments need a full 32-bit pack. D-split reducers - # may own only one or two output bytes per thread; pad the register - # fragment while keeping global reads/stores limited to that band. - fragment_elements = elements_per_thread - if cutlass.const_expr(self.o_dtype == cutlass.Float8E4M3FN): - fragment_elements = max(4, elements_per_thread) - gAccO = mAccO[q_tile_row, None, None, q_tile, bidz] - rAccO = cute.make_rmem_tensor(cute.make_layout(fragment_elements), - self.acc_dtype) - rO = cute.make_rmem_tensor(cute.make_layout(fragment_elements), - self.o_dtype) - rAccO.fill(0.0) - for i in range(local_split_kv): + # A request planned onto a single split stored its output + # directly in the split kernel (variable split-KV only); K == 0 + # rows still take the canonical empty result from here. + reduce_row = cutlass.Boolean(True) + if cutlass.const_expr(self.is_var_split_kv): + reduce_row = local_split_kv != 1 + if reduce_row: + gLSE = mAccLSE[q_tile_row, None, q_tile, bidz] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 0: + # calculate the global lse and exp ^ (local_lse - global_lse) + lse_per_thread = cute.ceil_div(self.reducer_max_splits, + self.threads_per_warp) + + local_lse = cute.make_rmem_tensor( + cute.make_layout(lse_per_thread), self.lse_dtype) + lse_max = -self.lse_dtype.inf + # find the max lse + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + local_lse[i] = (gLSE[split_kv_idx] if cute.elem_less( + split_kv_idx, local_split_kv) else + -self.lse_dtype.inf) + # reduce the local lse + lse_max = cute.arch.fmax(lse_max, local_lse[i]) + lse_max = cute.arch.warp_reduction_max(lse_max) + lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 + # calculate sum_lse + sum_lse = 0.0 + for i in cutlass.range_constexpr(lse_per_thread): + sum_lse += cute.math.exp2(local_lse[i] - lse_max, + fastmath=True) + sum_lse = cute.arch.warp_reduction_sum(sum_lse) + # calculate the global_lse + global_lse = (lse_max + + cute.math.log2(sum_lse, fastmath=True) + if not sum_lse == self.lse_dtype(0.0) + or sum_lse != sum_lse else self.lse_dtype.inf) + if d_tile_idx == 0 and tidx == 0: + if cutlass.const_expr(self.is_var_q): + mLSE[head_idx, q_begin + bidy] = global_lse + if cutlass.const_expr(self.emit_softmax_stats): + if has_local_kv: + mSoftmaxStats[head_idx, q_begin + bidy, + 0] = (global_lse / LOG2_E) + mSoftmaxStats[head_idx, q_begin + bidy, + 1] = 1.0 + else: + mSoftmaxStats[head_idx, q_begin + bidy, + 0] = -self.lse_dtype.inf + mSoftmaxStats[head_idx, q_begin + bidy, + 1] = 0.0 + else: + mLSE[head_idx, bidy, bidz] = global_lse + if cutlass.const_expr(self.emit_softmax_stats): + if has_local_kv: + mSoftmaxStats[head_idx, bidy, bidz, + 0] = (global_lse / LOG2_E) + mSoftmaxStats[head_idx, bidy, bidz, 1] = 1.0 + else: + mSoftmaxStats[head_idx, bidy, bidz, + 0] = -self.lse_dtype.inf + mSoftmaxStats[head_idx, bidy, bidz, 1] = 0.0 + # store the scale to shared memory + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + if cute.elem_less(split_kv_idx, local_split_kv): + smem_lse_scale[split_kv_idx] = cute.math.exp2( + local_lse[i] - global_lse, fastmath=True) + + pipeline.sync(barrier_id=4) + + elements_per_thread = cute.ceil_div( + self.reducer_d_tile, + self.threads_per_warp * self.num_compute_warps, + ) + # FP8 register fragments need a full 32-bit pack. D-split reducers + # may own only one or two output bytes per thread; pad the register + # fragment while keeping global reads/stores limited to that band. + fragment_elements = elements_per_thread + if cutlass.const_expr(self.o_dtype == cutlass.Float8E4M3FN): + fragment_elements = max(4, elements_per_thread) + gAccO = mAccO[q_tile_row, None, None, q_tile, bidz] + rAccO = cute.make_rmem_tensor( + cute.make_layout(fragment_elements), self.acc_dtype) + rO = cute.make_rmem_tensor(cute.make_layout(fragment_elements), + self.o_dtype) + rAccO.fill(0.0) + for i in range(local_split_kv): + for j in cutlass.range_constexpr(elements_per_thread): + element_idx = ( + d_tile_idx * self.reducer_d_tile + tidx + + j * self.threads_per_warp * self.num_compute_warps) + rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i] + rO.store(rAccO.load().to(self.o_dtype)) for j in cutlass.range_constexpr(elements_per_thread): element_idx = ( d_tile_idx * self.reducer_d_tile + tidx + j * self.threads_per_warp * self.num_compute_warps) - rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i] - rO.store(rAccO.load().to(self.o_dtype)) - for j in cutlass.range_constexpr(elements_per_thread): - element_idx = ( - d_tile_idx * self.reducer_d_tile + tidx + - j * self.threads_per_warp * self.num_compute_warps) - if cutlass.const_expr(self.is_var_q): - mO[head_idx, element_idx, q_begin + bidy] = rO[j] - else: - mO[head_idx, element_idx, bidy, bidz] = rO[j] + if cutlass.const_expr(self.is_var_q): + mO[head_idx, element_idx, q_begin + bidy] = rO[j] + else: + mO[head_idx, element_idx, bidy, bidz] = rO[j] return @cute.jit @@ -2188,7 +2294,11 @@ def get_k_tile_count( """ K = cache_seqs[blk_coord[2]] if cutlass.const_expr(self.is_var_split_kv): - split_kv = block_split_kvs[blk_coord[2]] + # Per-request split planned on the device for the batch's KV + # lengths; clamped to the launched grid so the plan can be made + # for the largest split any tactic may use. + split_kv = cutlass.max( + cutlass.min(block_split_kvs[blk_coord[2]], split_kv), 1) k_tile_total = cute.ceil_div(K, self.mma_qk_tiler[1]) k_tile_per_cta = cute.ceil_div(k_tile_total, split_kv) @@ -2196,6 +2306,12 @@ def get_k_tile_count( k_tile_count = max( 0, min(k_tile_total, k_index + k_tile_per_cta) - k_index) + if cutlass.const_expr(self.is_var_split_kv): + # The split the request really uses after tile rounding (the + # reduction kernel derives the same value): one means this CTA + # stores the final output itself and the reduction skips the row. + split_kv = cute.ceil_div(k_tile_total, + cutlass.max(k_tile_per_cta, 1)) q_begin = cutlass.Int32(0) q_len = cutlass.Int32(self.seq_len_q) if cutlass.const_expr(self.is_var_q): @@ -5155,6 +5271,15 @@ def create_stats(): # Get the raw stream pointer as a CUstream stream = cuda.CUstream(torch_stream.cuda_stream) + # The kernel reads both scales from device memory (FP8 KV keeps them as + # device tensors so CUDA graph replay never needs a host read). + softmax_scale_t = cute.runtime.from_dlpack( + torch.tensor([softmax_scale], dtype=torch.float32, + device="cuda")).mark_layout_dynamic() + output_scale_t = cute.runtime.from_dlpack( + torch.tensor([output_scale], dtype=torch.float32, + device="cuda")).mark_layout_dynamic() + # compile mla kernel compiled_mla = cute.compile( mla.run_with_softmax_stats if emit_softmax_stats else mla, @@ -5172,8 +5297,8 @@ def create_stats(): kv_bounds_ct, cum_seq_lens_q, block_split_kvs, - softmax_scale, - output_scale, + softmax_scale_t, + output_scale_t, stream, options="--opt-level 2", ) @@ -5263,8 +5388,8 @@ def torch_reference_mla(q_latent, kv_bounds_ct, cum_seq_lens_q, block_split_kvs, - softmax_scale, - output_scale, + softmax_scale_t, + output_scale_t, stream, ) torch.cuda.synchronize() @@ -5408,8 +5533,8 @@ def generate_tensors(): kv_bounds_ct, cum_seq_lens_q, block_split_kvs, - softmax_scale, - output_scale, + softmax_scale_t, + output_scale_t, stream, ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py index 822b7f1e5097..9fac143f6fe6 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py @@ -15,36 +15,242 @@ """CuTe DSL device kernel for KDA multi-token speculative verify (conv-MTP). Source-integrated from the ``KDA_decode_mtp`` kernel drop ("readable KDA -conv-MTP"). Specialized for the benchmarked contract: conv enabled, no bias, -no output norm, lower_bound gate, Q/K L2 norm, beta sigmoid, TILE_V=64, -ILP=2, W=4, K == V == 128, HV == H. +conv-MTP") and restructured for the latency-bound production regime (one CTA +per (head, request), or a cluster pair of CTAs while the batch is small; TEP8 +launches only ``H * N = 24`` tiles on a 212-SM Rubin). Specialized for the +benchmarked contract: conv enabled, no bias, optional BF16 gated RMSNorm +output, lower_bound gate, Q/K L2 norm, beta sigmoid, W=4, K == V == 128, +HV == H, ``TILE_V == V`` (one V pass), 512 threads. -Per request ``n`` the kernel processes ``num_accepted_tokens[n] + 1 + +Per request ``n`` the kernel processes an accepted count plus ``1 + NUM_SPEC`` steps: it first *replays* the accepted draft tokens from the -``qkg/v/beta`` caches (raw conv inputs from the extended tail columns of +``k/g/v/beta`` caches (raw conv inputs from the extended tail columns of ``cs_q/cs_k/cs_v``), then processes the ``1 + NUM_SPEC`` new tokens. The recurrent state and the base conv windows are committed in place after the first new (golden) token; the new spec tokens are cached for the next round's replay. The pool invariant is therefore "state after the last golden token; accepted drafts pending in the replay caches". -Vendoring delta vs the drop: the ``v_row_a``/``v_row_b`` readout indices in -the output stage are hoisted above the ``if/elif`` on ``USE_ZERO_ACCEPTED`` -/ ``i_t >= commit_len``. In the drop they were first assigned inside the -*dynamic* ``elif`` branch, which the CuTe DSL cannot trace (names must -pre-exist before dynamic assignment), so every non-``USE_ZERO_ACCEPTED`` -compile — i.e. any replay round, and any non-benchmark shape — failed with -``DSLRuntimeError: v_row_a is None``. The hoist is semantics-preserving. +Structure (all stage boundaries are CTA barriers): + +1. **Load phase.** Every warp issues its global loads up front, so the whole + CTA waits for one DRAM round trip: the recurrent state ``h0`` through a TMA + bulk load into swizzled smem boxes (consumed by the tensor-core stage), the + raw conv inputs of every token (``x_q/x_k/x_v/g/beta`` for new tokens, the + replay caches and ``cs_*`` tail columns for accepted drafts) into smem + staging buffers, the conv weights, and the gated-RMSNorm operands for the + epilogue. +2. **Per-token-parallel precompute.** Warp ``j`` computes the Q/K conv + SiLU + + L2 norm of new token ``j`` and warp ``num_warps/2 + j`` its gate, beta and + V conv, reading the width-4 windows from the staged raw sequence (rows + ``0..W-2`` hold the committed window). Replay tokens were written to + ``sQ/sK/sG/sBeta/sVall`` directly by the load phase. Conv-window commits + and replay-cache writes are done by the warp that owns the token. +3. **Recurrence** in chunk (WY) form on the tcgen05 tensor cores, tf32 with + 3xTF32 (hi/lo splits) on every product that feeds the state: + + * GEMM1 ``[U0; O0] = S0 x [K~; Q~]`` -- S0 hi/lo staged into TMEM from the + smem boxes, ``K~ = k e^{G_t}``, ``Q~ = q e^{G_t}`` (``G_t`` = cumulative + gate of the chunk); + * GEMM3 ``L, B`` -- the T x T coupling matrices of the chunk (decays + factored around the chunk middle and clamped at ``e^87``); + * ``V' = (I + diag(beta) L)^-1 diag(beta) (V - U0)`` by forward + substitution (one warpgroup, one value row per thread); + * GEMM2 ``S_c = S0 e^{G_c} + V'^T Kbar`` -- accumulator seeded from the + registers holding S0, issued as two column halves so the first half can + start storing; the output combination ``O = O0 + B V'`` rides along as 16 + extra columns of the first half. +4. **Store / epilogue.** The committed state goes through swizzled smem boxes + and a TMA bulk tensor store; the gated RMSNorm and MXFP8 quantization read + their operands from smem and run on the warps that own no state rows while + the store drains. + +With ``SPLIT_V == 2`` two CTAs per (head, request) run as a cluster: both +recompute the k-side and the outputs, each seeds and stores half of the state +rows (the SM -> L2 drain of the 64 KB state is the tail of the kernel). The pair +must be co-scheduled because the state is updated in place. Two cluster +barrier phases order the pair: the first separates both CTAs' load-phase reads +of the conv windows and replay caches from their update, which only CTA 0 +stores; the second separates every CTA's S0 load from the pair's state stores. +The TMA views of the state follow the tensor's own strides (the runtime's +per-layer state cache is a strided slice of a larger buffer). The state may be +fp32 or bf16; a bf16 state is exact in tf32 (the ``S0_lo`` sweep is skipped) and the committed state +is rounded to nearest on the way out. + +Measured on Rubin (nsys, state L2-resident) against the previous CUDA-core +kernel at H=12/num_spec=7 with mixed acceptance: N=2 8.61 us -> 7.07 us, +N=16 11.78 us -> 9.06 us; H=96/N=16 91.1 us -> 74.2 us. The committed state +is within 2e-6 of the fp32 kernel (3xTF32), the FP8 output codes differ in +~0.2 % of the elements by one step; the conv pools and replay caches are +bit-exact. + +``REPLAY_ONLY`` compiles the replay steps alone, for committing a slot's pending +drafts before its pages leave the GPU (``kda_mtp_commit_pending_drafts``): the +chunk is the accepted drafts, the commit row the last of them, and only the +state is written. The verify that follows starts a new chunk at the golden +token, so it agrees with a verify that replays the drafts itself to the +kernel's accuracy, not bit for bit. + +``BF16_MMA`` (``TRTLLM_KDA_MTP_BF16_MATH=1``) trades that precision for the +bf16-operand / fp32-accumulate scheme of the prefill kernel: S0 packed as bf16 +in TMEM (64 columns), bf16 K-major smem tiles and K = 16 per instruction (18 +MMAs per CTA instead of 76). Gates, cumulative gates, exponents, reductions +and the state accumulator stay fp32; the committed-state error vs fp32 grows to +~2e-2, the conv pools and replay caches stay bit-exact between the two paths. """ import cutlass import cutlass.cute as cute from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.typing import Int64 from cutlass.cutlass_dsl import T, dsl_user_op -NUM_THREADS = 256 +# 16 warps: the full 128-row recurrent state is processed in one pass (eight +# rows per warp) and the per-token precompute runs one (token, role) per warp. +NUM_THREADS = 512 +# Two CTAs per (head, request) while HV * N <= SPLIT_MAX_TILES (up to 128 CTAs: one wave, the +# halved state store pays off); larger batches run one CTA per tile. +SPLIT_MAX_TILES = 64 TILE_K = 128 +# Chunk length padding (MMA N granularity). +T_PAD = 16 +# Largest draft count per verify: up to NUM_SPEC replayed drafts plus 1 + NUM_SPEC new tokens +# must fit the padded chunk (2 * 7 + 1 = 15 <= T_PAD). The executor falls back to the +# sequential verify for longer drafts. +MAX_NUM_SPEC = (T_PAD - 1) // 2 +# TMEM columns: D4 (16) + D2 halves (2 x 64) + S0 hi (128) + S0 lo (128) + D1 (48) + D3 (16) +# = 464 of the 512 allocated. +TMEM_COLS = 512 +# One 114 KB scratch region hosts, at different times, the TMA boxes of S0 / the committed +# state, the swizzled MMA operand tiles and (bf16 output path) the per-warp 32x32 transposes. +# Offsets in floats, all multiples of 256 (1024 B) for the 128B-swizzle atoms. +SCRATCH_FLOATS = 29184 +OFF_B4HI = 10240 # 16 x 16 Bcoef hi, immediately followed by Kbar hi -> N=80 GEMM2 half-0 tile +OFF_B4LO = 12544 # 16 x 16 Bcoef lo, immediately followed by Kbar lo +OFF_B1K = 27136 # 16 x 128 K~hi (duplicate, N=16 tile for the S0_lo sweep) +OFF_B1 = 0 # 48 x 128 [K~hi; Q~; K~lo] +OFF_A2HI = 6144 # 128 x 16 V'^T hi +OFF_A2LO = 8192 # 128 x 16 V'^T lo +OFF_B2HI = 10496 # 128 x 16 Kbar hi +OFF_B2LO = 12800 # 128 x 16 Kbar lo +OFF_A3 = 14848 # 64 x 128 [Kf; Qf hi; Qf lo; -] +OFF_B3HI = 23040 # 16 x 128 Kb hi +OFF_B3LO = 25088 # 16 x 128 Kb lo +TP_STRIDE = 17 # padded row stride of the transposed (v|k, s) staging buffers +# bf16 operand tiles (BF16_MMA): offsets in bf16 elements inside the same scratch region, all +# multiples of 512 (1024 B). [Bcoef; Kbar] is one contiguous SW32 tile (Bcoef rows 0..15, then the +# 128 Kbar rows); the GEMM2 half-1 B view starts at Kbar row 64. +OFFH_B1 = 0 # 32 x 128 [K~; Q~] +OFFH_A3 = 4096 # 64 x 128 [Kf; Qf; 0; -] +OFFH_B3 = 12288 # 16 x 128 Kb +OFFH_A2 = 14336 # 128 x 16 V'^T +OFFH_B4 = 16384 # 16 x 16 Bcoef, immediately followed by +OFFH_B2 = 16640 # 128 x 16 Kbar +XP_STRIDE = 33 # padded row stride of the 32x32 transpose blocks + + +def swz128(n, k, rows): + """Element offset of (row n, col k) in a K-major 128B-swizzled fp32 tile with + `rows` rows and 128 columns (canonical UMMA layout, see probe_swizzle.py).""" + off = n * 32 + (k % 32) + (k // 32) * (rows * 32) + return off ^ (((off >> 5) & 7) << 2) + + +def swz64_k16(n, k): + """Element offset in a K-major 64B-swizzled fp32 tile with 16 columns.""" + off = n * 16 + k + return off ^ (((off >> 5) & 3) << 2) + + +def sw128_bf16(n, k, rows): + """Element offset of (n, k) in a K-major 128B-swizzled bf16 tile with `rows` rows and K = 128: + two K blocks of 64 elements (rows x 128 B each), 16 B chunks XORed with n % 8.""" + kc = k % 64 + return (k // 64) * (rows * 64) + n * 64 + ((((kc >> 3) ^ (n & 7)) << 3) | (kc & 7)) + + +def sw32_k16_bf16(n, k): + """Element offset of (n, k) in a K-major 32B-swizzled bf16 tile with 16 columns.""" + off = n * 16 + k + return off ^ (((off >> 6) & 1) << 3) + + +@dsl_user_op +def cvt_bf16x2(hi: cutlass.Float32, lo: cutlass.Float32, *, loc=None, ip=None) -> cutlass.Float32: + """Pack two fp32 values into one bf16x2 word (round to nearest even, `lo` in the low half) + with a single cvt; returned as the fp32 bit pattern for the TMEM word stores.""" + w = 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), + ], + "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, + ) + ) + return w.bitcast(cutlass.Float32) + + +@dsl_user_op +def tcgen05_fence_before_thread_sync(*, loc=None, ip=None): + llvm.inline_asm( + None, + [], + "tcgen05.fence::before_thread_sync;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def tcgen05_fence_after_thread_sync(*, loc=None, ip=None): + llvm.inline_asm( + None, + [], + "tcgen05.fence::after_thread_sync;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def round_tf32(x: cutlass.Float32) -> cutlass.Float32: + """Round to the tf32 grid (rna on mantissa bit 13) so hi/lo splits are exact.""" + b = x.bitcast(cutlass.Int32) + b = (b + cutlass.Int32(0x1000)) & cutlass.Int32(-8192) + return b.bitcast(cutlass.Float32) + + +def round_bf16(x: cutlass.Float32) -> cutlass.Float32: + """Round to the bf16 grid (rna on mantissa bit 16, two integer ops like round_tf32); the + value stays an fp32 operand that is exact in tf32.""" + b = x.bitcast(cutlass.Int32) + b = (b + cutlass.Int32(0x8000)) & cutlass.Int32(-65536) + return b.bitcast(cutlass.Float32) + + +def round_hi(x: cutlass.Float32, BF16_MMA: bool) -> cutlass.Float32: + """MMA operand rounding: tf32 grid for the 3xTF32 path (so the hi/lo split is exact), bf16 + grid for the bf16 operand path (the conversion into the bf16 tiles is then exact).""" + if BF16_MMA: + return round_bf16(x) + return round_tf32(x) @dsl_user_op @@ -81,9 +287,13 @@ def kda_decode_mtp_kernel( g: cute.Tensor, dt_bias: cute.Tensor, beta: cute.Tensor, + onorm_g: cute.Tensor, + onorm_weight: cute.Tensor, o: cute.Tensor, + output_scale: cute.Tensor, ht: cute.Tensor, - qkg_cache: cute.Tensor, + k_cache: cute.Tensor, + g_cache: cute.Tensor, v_cache: cute.Tensor, beta_cache: cute.Tensor, smem_qk_layout: cute.Layout, @@ -99,82 +309,174 @@ def kda_decode_mtp_kernel( NUM_SPEC: cutlass.Constexpr[int], KERNEL_WIDTH: cutlass.Constexpr[int], lower_bound: cutlass.Constexpr[float], + onorm_eps: cutlass.Constexpr[float], + scale_leading_dim: cutlass.Constexpr[int], USE_FLAT_LAYOUT: cutlass.Constexpr[bool], - USE_SETMAXREG: cutlass.Constexpr[bool], USE_REGULAR_METADATA: cutlass.Constexpr[bool], - USE_REG_Q_WEIGHTS: cutlass.Constexpr[bool], + USE_PACKED_TOKEN_LAYOUT: cutlass.Constexpr[bool], USE_ZERO_ACCEPTED: cutlass.Constexpr[bool], FUSE_PRECOMPUTE: cutlass.Constexpr[bool], RUNTIME_PRECOMPUTE_FLAG: cutlass.Constexpr[bool], + FUSE_OUTPUT_NORM: cutlass.Constexpr[bool], + QUANTIZE_OUTPUT: cutlass.Constexpr[bool], stage_timing: cute.Tensor, PROFILE_STAGES: cutlass.Constexpr[bool], + tiled_mma1: cute.TiledMma, + tiled_mma3: cute.TiledMma, + tiled_mma1k: cute.TiledMma, + smem_layout_b1: cute.ComposedLayout, + smem_layout_a2: cute.ComposedLayout, + smem_layout_a3: cute.ComposedLayout, + smem_layout_b3: cute.ComposedLayout, + smem_layout_b1k: cute.ComposedLayout, + tiled_mma4: cute.TiledMma, + tma_atom_t: cute.CopyAtom, + mT_tma: cute.Tensor, + box_layout: cute.ComposedLayout, + tma_atom_l: cute.CopyAtom, + mS_tma: cute.Tensor, + tiled_mma2h: cute.TiledMma, + smem_layout_b2h: cute.ComposedLayout, + box_layout_t: cute.ComposedLayout, + tiled_mma2q: cute.TiledMma, + smem_layout_b2q: cute.ComposedLayout, + SPLIT_V: cutlass.Constexpr[int], + BF16_MMA: cutlass.Constexpr[bool], + tiled_mma_f32a: cute.TiledMma, REPLAY_ONLY: cutlass.Constexpr[bool], ): - """KDA MTP decode — SMEM pre-compute + register-resident state. - - With ``PROFILE_STAGES=True``, ``stage_timing`` must be an int64 tensor - with at least ``HV * N * 4`` elements, indexed as - ``(i_hv * grid_n + i_n) * 4``. With profiling off it is never accessed - and the host may pass any placeholder tensor. - - With ``REPLAY_ONLY=True`` the kernel runs only the replay steps: it - replays the ``num_accepted_tokens[n]`` pending drafts and commits the - recurrent state after them to ``ht``. It reads only ``h0``, the conv and - replay caches and the metadata, and writes only ``ht``; ``x_*``, ``w_*``, - ``g``, ``beta``, ``A_log``, ``dt_bias`` and ``o`` are not accessed. + """KDA MTP verify: smem precompute + tensor-core chunk recurrence. + + With ``PROFILE_STAGES=True``, ``stage_timing`` must be an int64 tensor with at + least ``SPLIT_V * HV * N * 64`` elements, indexed as + ``((vh * HV + i_hv) * grid_n + i_n) * 64``. With profiling off it is never + accessed and the host may pass any placeholder tensor. + + With ``REPLAY_ONLY=True`` the kernel runs only the replay steps: the chunk + is the ``num_accepted_tokens[n]`` pending drafts of request ``n`` (counts + indexed by request), and the recurrent state after the last of them is + committed to ``ht``; with no draft pending the state is stored unchanged. + It reads only ``h0``, the replay caches and the metadata, and writes only + ``ht``; the token inputs, conv weights, gate constants, the conv caches + ``cs_*``, ``cu_seqlens`` and ``o`` are not accessed, and nothing is cached + for a later round. The caller commits the conv window after the drafts. """ tidx, _, _ = cute.arch.thread_idx() in_warp_tid = tidx % 32 warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) - i_hv, i_n, _ = cute.arch.block_idx() + i_hv, i_n, vh = cute.arch.block_idx() + # Two CTAs per (head, request): CTA vh owns the state rows [vh * V/2, (vh + 1) * V/2). i_h = i_hv if cutlass.const_expr(PROFILE_STAGES): t_stage0 = read_globaltimer() # Pre-declare so the dynamic `run_precompute` branch below only # reassigns (first assignment inside a dynamic branch is untraceable; - # see the module docstring on v_row_a/v_row_b). + # the name has to exist before the branch). t_stage1 = Int64(0) - if cutlass.const_expr(USE_REGULAR_METADATA): + tp_a = Int64(0) + tp_b = Int64(0) + tp_c = Int64(0) + tp_d = Int64(0) + tp_e = Int64(0) + tp_f = Int64(0) + tp_g = Int64(0) + tp_h = Int64(0) + tp_i = Int64(0) + tp_j = Int64(0) + tp_k = Int64(0) + tp_l = Int64(0) + tp_m = Int64(0) + if cutlass.const_expr(REPLAY_ONLY): + # No token rows: the replayed drafts come from the slot's replay caches. + bos = cutlass.Int32(0) + eos = cutlass.Int32(0) + slot = ssm_state_indices[i_n] + elif cutlass.const_expr(USE_REGULAR_METADATA): bos = i_n * (2 * NUM_SPEC + 1) eos = bos + (2 * NUM_SPEC + 1) slot = i_n + elif cutlass.const_expr(USE_PACKED_TOKEN_LAYOUT): + # Runtime verification packs exactly 1 + NUM_SPEC new rows per + # request. Accepted drafts live in slot-indexed replay pools and do + # not occupy rows in the projection outputs. + bos = cutlass.Int32(0) + eos = cutlass.Int32(0) + slot = ssm_state_indices[i_n] else: bos = cu_seqlens[i_n] eos = cu_seqlens[i_n + 1] slot = ssm_state_indices[i_n] - # V2 can interleave many layers between state slots. Promote the index - # before multiplying by static strides: the element offset can exceed - # INT32_MAX even when each individual stride fits in int32. + # V2 state views may place many layer pages between adjacent slots. Widen + # the slot before indexing: slot * state stride can exceed INT32_MAX. slot = Int64(slot) h0_idx = slot * HV + i_hv hk_off = i_h * K hv_off = i_hv * V + # Recurrent-state element type (fp32 or bf16). The TMA boxes have 128 B swizzled rows, + # i.e. BOX_COLS = 32 fp32 or 64 bf16 columns; the state tile is NUM_SLABS such boxes. + STATE_DTYPE = h0.element_type + STATE_BF16 = STATE_DTYPE == cutlass.BFloat16 + # Recurrence precision. Default: 3xTF32 (hi/lo splits on every product that feeds the state). + # BF16_MMA: bf16 tensor-core operands -- S0 packed as bf16 in TMEM (64 columns), bf16 K-major + # smem tiles (half the bytes), K = 16 per MMA (a quarter of the instructions); the + # bf16-operand / fp32-accumulate scheme of the prefill kernel. Gates, cumulative gates, + # exponents, reductions and the state accumulator stay fp32 on both paths. + # S0_LO: the 3xTF32 path needs the lo part of a fp32 state (a bf16 state is exact in tf32). + S0_LO = (not STATE_BF16) and (not BF16_MMA) + N1 = 2 * T_PAD if BF16_MMA else 3 * T_PAD # GEMM1 N: [K~; Q~] (+ K~lo for the 3xTF32 sweep) + KB_MMA = 16 if BF16_MMA else 8 # K per tcgen05 instruction + # raw conv inputs of the committed window and of the pending drafts (bf16 in the runtime: the + # projections are bf16, so the copy is exact; fp32 caches are accepted as well) + CONV_DTYPE = cs_q.element_type + # replay caches: k / v / beta may be bf16 (one bf16 ulp on the replayed drafts), the gate + # cache stays fp32 (it feeds the cumulative-gate exponents) + K_CACHE_DTYPE = k_cache.element_type + V_CACHE_DTYPE = v_cache.element_type + BETA_CACHE_DTYPE = beta_cache.element_type + BOX_COLS = 128 // (STATE_DTYPE.width // 8) + NUM_SLABS = K // BOX_COLS if cutlass.const_expr(USE_ZERO_ACCEPTED): - commit_len = 0 + commit_len = cutlass.Int32(0) else: - commit_len = num_accepted_tokens[i_n] + if cutlass.const_expr(USE_PACKED_TOKEN_LAYOUT): + commit_len = num_accepted_tokens[slot] + else: + commit_len = num_accepted_tokens[i_n] # Only NUM_SPEC drafts can be pending from the previous round. Clamp # so a malformed count cannot drive T_loop past the t_max-sized SMEM # buffers or the num_spec extents of the replay caches. if commit_len > NUM_SPEC: commit_len = cutlass.Int32(NUM_SPEC) + if cutlass.const_expr(USE_PACKED_TOKEN_LAYOUT): + # Existing body addresses every new row as bos + i_t, where i_t is + # offset by commit_len after replay. Shift bos so those accesses map + # to request-major packed rows without materialized cu_seqlens. + bos = i_n * (1 + NUM_SPEC) - commit_len + eos = bos + commit_len + 1 + NUM_SPEC if cutlass.const_expr(USE_ZERO_ACCEPTED): - T_loop = 1 + NUM_SPEC + T_loop = cutlass.Int32(1 + NUM_SPEC) t_max = 1 + NUM_SPEC + elif cutlass.const_expr(REPLAY_ONLY): + T_loop = commit_len + t_max = 2 * NUM_SPEC + 1 else: - if cutlass.const_expr(REPLAY_ONLY): - T_loop = commit_len - else: - T_loop = commit_len + 1 + NUM_SPEC + T_loop = commit_len + 1 + NUM_SPEC t_max = 2 * NUM_SPEC + 1 + # Chunk row whose state is committed: the golden token (row commit_len), or with REPLAY_ONLY + # the last replayed draft (-1 when none is pending: the state is committed unchanged). + if cutlass.const_expr(REPLAY_ONLY): + commit_idx = commit_len - 1 + else: + commit_idx = commit_len vec_size = TILE_K // 32 num_v_tiles = V // TILE_V NUM_V_ROWS = TILE_V // (NUM_THREADS // 32) - if cutlass.const_expr(USE_REG_Q_WEIGHTS): - q_weight_elems = 0 - else: - q_weight_elems = KERNEL_WIDTH * K + # The kernel assumes one V pass with eight rows per warp. + assert num_v_tiles == 1 and NUM_V_ROWS == 8, ( + f"kda_decode_mtp_kernel expects TILE_V == V and 8 rows per warp, got TILE_V={TILE_V}, V={V}" + ) + q_weight_elems = KERNEL_WIDTH * K v_weight_elems = KERNEL_WIDTH * V k_weight_base = q_weight_elems v_weight_base = q_weight_elems + KERNEL_WIDTH * K @@ -184,9 +486,9 @@ def kda_decode_mtp_kernel( sK = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) sG = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max,)), 16) - # Reserve the 8-float scratch region formerly used by the removed - # output-norm reduction. Nothing reads or writes this allocation; it only - # preserves the offsets and bank mapping of the shared-memory buffers below. + # Preserve the offsets and bank mapping of the shared-memory buffers from + # the source kernel. The fused norm below reduces within one warp and does + # not need this legacy CTA-reduction scratch. smem.allocate_tensor(cutlass.Float32, cute.make_layout((8,)), 16) sVall = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max * V,)), 16) sConvW = smem.allocate_tensor( @@ -194,6 +496,195 @@ def kda_decode_mtp_kernel( cute.make_layout((conv_weight_elems,)), 16, ) + if cutlass.const_expr(FUSE_OUTPUT_NORM): + # Epilogue operands are staged in smem during the load phase. + sOnG = smem.allocate_tensor(cutlass.Float32, cute.make_layout(((1 + NUM_SPEC) * V,)), 16) + sOnW = smem.allocate_tensor(cutlass.Float32, cute.make_layout((V,)), 16) + # Staged per-token inputs (raw conv sequence with the committed window + # prepended in rows 0..W-2, raw gates, raw betas). + sXQ = smem.allocate_tensor(cutlass.Float32, cute.make_layout(((t_max + 3) * K,)), 16) + sXK = smem.allocate_tensor(cutlass.Float32, cute.make_layout(((t_max + 3) * K,)), 16) + sXV = smem.allocate_tensor(cutlass.Float32, cute.make_layout(((t_max + 3) * V,)), 16) + sGin = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max * K,)), 16) + sBetaRaw = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max,)), 16) + # Per-head gate constants, staged in the load phase together with the other inputs. + sDtBias = smem.allocate_tensor(cutlass.Float32, cute.make_layout((K,)), 16) + sAlog = smem.allocate_tensor(cutlass.Float32, cute.make_layout((32,)), 16) + # ---- tensor-core chunk recurrence buffers ---- + assert NUM_SPEC <= MAX_NUM_SPEC and t_max <= T_PAD, ( + f"chunk form pads the token axis to {T_PAD}: NUM_SPEC must be <= {MAX_NUM_SPEC}" + ) + tmem_holder = smem.allocate_array(cutlass.Int32, 1) + mbar1 = smem.allocate_array(cutlass.Int64, 1) + mbar2 = smem.allocate_array(cutlass.Int64, 1) + mbar_ld = smem.allocate_array(cutlass.Int64, 1) + mbar4 = smem.allocate_array(cutlass.Int64, 1) + mbar2b = smem.allocate_array(cutlass.Int64, 1) + sScratch = smem.allocate_tensor(cutlass.Float32, cute.make_layout((SCRATCH_FLOATS,)), 1024) + # Swizzled fp32 MMA operand tiles as views into the scratch region (tf32 modes; the layouts + # passed in are bf16 in mode 2, whose tiles are built below). + if cutlass.const_expr(not BF16_MMA): + sB1 = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B1, smem_layout_b1.inner, dtype=cutlass.Float32 + ), + smem_layout_b1.outer, + ) + sA2hi = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_A2HI, smem_layout_a2.inner, dtype=cutlass.Float32 + ), + smem_layout_a2.outer, + ) + sA2lo = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_A2LO, smem_layout_a2.inner, dtype=cutlass.Float32 + ), + smem_layout_a2.outer, + ) + sA3 = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_A3, smem_layout_a3.inner, dtype=cutlass.Float32 + ), + smem_layout_a3.outer, + ) + sB3hi = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B3HI, smem_layout_b3.inner, dtype=cutlass.Float32 + ), + smem_layout_b3.outer, + ) + sB3lo = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B3LO, smem_layout_b3.inner, dtype=cutlass.Float32 + ), + smem_layout_b3.outer, + ) + sB1k = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_B1K, smem_layout_b1k.inner, dtype=cutlass.Float32), + smem_layout_b1k.outer, + ) + sB1kf = cute.make_tensor(sScratch.iterator + OFF_B1K, cute.make_layout((T_PAD * 128,))) + # Flat (unswizzled-pointer) views for chunked element writes with swz128/swz64_k16 offsets. + sB1f = cute.make_tensor(sScratch.iterator + OFF_B1, cute.make_layout((48 * 128,))) + sA2hif = cute.make_tensor(sScratch.iterator + OFF_A2HI, cute.make_layout((128 * 16,))) + sA2lof = cute.make_tensor(sScratch.iterator + OFF_A2LO, cute.make_layout((128 * 16,))) + sB2hif = cute.make_tensor(sScratch.iterator + OFF_B2HI, cute.make_layout((128 * 16,))) + sB2lof = cute.make_tensor(sScratch.iterator + OFF_B2LO, cute.make_layout((128 * 16,))) + sB4hif = cute.make_tensor(sScratch.iterator + OFF_B4HI, cute.make_layout((T_PAD * T_PAD,))) + sB4lof = cute.make_tensor(sScratch.iterator + OFF_B4LO, cute.make_layout((T_PAD * T_PAD,))) + sA3f = cute.make_tensor(sScratch.iterator + OFF_A3, cute.make_layout((64 * 128,))) + sB3hif = cute.make_tensor(sScratch.iterator + OFF_B3HI, cute.make_layout((16 * 128,))) + sB3lof = cute.make_tensor(sScratch.iterator + OFF_B3LO, cute.make_layout((16 * 128,))) + # bf16 operand tiles (BF16_MMA): swizzled MMA views and flat views for the closed-form writes. + sScratchH = cute.make_tensor( + cute.recast_ptr(sScratch.iterator, dtype=cutlass.BFloat16), + cute.make_layout((2 * SCRATCH_FLOATS,)), + ) + if cutlass.const_expr(BF16_MMA): + sB1b = cute.make_tensor( + cute.recast_ptr( + sScratchH.iterator + OFFH_B1, smem_layout_b1.inner, dtype=cutlass.BFloat16 + ), + smem_layout_b1.outer, + ) + sA3b = cute.make_tensor( + cute.recast_ptr( + sScratchH.iterator + OFFH_A3, smem_layout_a3.inner, dtype=cutlass.BFloat16 + ), + smem_layout_a3.outer, + ) + sB3b = cute.make_tensor( + cute.recast_ptr( + sScratchH.iterator + OFFH_B3, smem_layout_b3.inner, dtype=cutlass.BFloat16 + ), + smem_layout_b3.outer, + ) + sA2b = cute.make_tensor( + cute.recast_ptr( + sScratchH.iterator + OFFH_A2, smem_layout_a2.inner, dtype=cutlass.BFloat16 + ), + smem_layout_a2.outer, + ) + sB1b_f = cute.make_tensor(sScratchH.iterator + OFFH_B1, cute.make_layout((32 * 128,))) + sA3b_f = cute.make_tensor(sScratchH.iterator + OFFH_A3, cute.make_layout((64 * 128,))) + sB3b_f = cute.make_tensor(sScratchH.iterator + OFFH_B3, cute.make_layout((16 * 128,))) + sA2b_f = cute.make_tensor(sScratchH.iterator + OFFH_A2, cute.make_layout((128 * 16,))) + sB4b_f = cute.make_tensor(sScratchH.iterator + OFFH_B4, cute.make_layout((16 * 16,))) + sB2b_f = cute.make_tensor(sScratchH.iterator + OFFH_B2, cute.make_layout((128 * 16,))) + # Per-warp 32 x 32 transpose blocks (bf16 output path: per-warp state store). + sXp = cute.make_tensor(sScratch.iterator, cute.make_layout((16 * 32 * XP_STRIDE,))) + # GEMM3 result [L; B hi; B lo; -] as a plain (64 x 16) row-major matrix. + sLB = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((4 * T_PAD, T_PAD), stride=(T_PAD, 1)), 16 + ) + sEc = smem.allocate_tensor(cutlass.Float32, cute.make_layout((K,)), 16) + # V'^T (v, s) and Kbar^T (k, s) staged with a padded row stride (conflict-free for both + # the row-per-thread producers and the 16-lanes-per-row tile writers). + sVpT = smem.allocate_tensor(cutlass.Float32, cute.make_layout((V * TP_STRIDE,)), 16) + sKbT = smem.allocate_tensor(cutlass.Float32, cute.make_layout((K * TP_STRIDE,)), 16) + # Cumulative gates G_t (per k) reuse the raw-gate staging buffer (no longer needed after the precompute). + sE = cute.make_tensor(sGin.iterator, cute.make_layout((t_max * K,))) + if warp_idx == 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_init(mbar1, 1) + cute.arch.mbarrier_init(mbar2, 1) + cute.arch.mbarrier_init(mbar4, 1) + cute.arch.mbarrier_init(mbar2b, 1) + cute.arch.mbarrier_init(mbar_ld, 1) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + # ---- TMA bulk load of S0 into the boxes; completes under the token staging / + # precompute and is consumed (mbar_ld) at the start of the chunk stage. + # Issued by the whole of one warp (a TMA load issued from a single thread never + # completes on this DSL build; the copy elects its issuing lane internally). + # Load-phase roles besides the per-warp token rows (warp w stages token w): the last K + # threads stage the committed q/k windows, the q/k conv weights and dt_bias, the V threads + # below them the committed v window, and the S0 TMA is issued by the last v-window warp + # (three window loads per lane). With 512 threads and K = V = 128 the staging lands on + # warps 8-15, of which 11-15 own no token row at production trip counts (t_max <= 15). + QK_STAGE_TID0 = NUM_THREADS - K + V_STAGE_TID0 = QK_STAGE_TID0 - V + TMA_WARP = QK_STAGE_TID0 // 32 - 1 + assert V_STAGE_TID0 >= 0 and V_STAGE_TID0 % 32 == 0 and QK_STAGE_TID0 % 32 == 0, ( + f"load-phase staging roles need whole warps for the last K + V threads, got " + f"NUM_THREADS={NUM_THREADS}, K={K}, V={V}" + ) + if warp_idx == TMA_WARP: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(mbar_ld, V * K * (STATE_DTYPE.width // 8)) + for sl_h in cutlass.range_constexpr(NUM_SLABS): + sTl_h = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + sl_h * (V * 32), box_layout.inner, dtype=STATE_DTYPE + ), + box_layout.outer, + ) + gS_h_all = cute.local_tile(mS_tma, (V, BOX_COLS, 1), (0, sl_h, None)) + if cutlass.const_expr(USE_FLAT_LAYOUT): + gS_h = gS_h_all[(None, None, 0, h0_idx)] + else: + gS_h = gS_h_all[(None, None, 0, (i_hv, slot))] + tSl_h, tGl_h = cpasync.tma_partition( + tma_atom_l, + 0, + cute.make_layout(1), + cute.group_modes(sTl_h, 0, 2), + cute.group_modes(gS_h, 0, 2), + ) + if warp_idx == TMA_WARP: + cute.copy(tma_atom_l, tGl_h, tSl_h, tma_bar_ptr=mbar_ld) + # TMEM allocation overlaps the in-flight loads; the pointer is retrieved after the + # precompute barrier (first TMEM use is the S0 staging in the chunk stage). + if warp_idx == 0: + cute.arch.alloc_tmem(TMEM_COLS, tmem_holder) + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + t_alloc = read_globaltimer() + _, grid_n4, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * grid_n4 + i_n) * 64 + 18] = t_alloc - t_stage0 + wg = warp_idx // 4 + wg_tidx = tidx % 128 r_q = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) r_k = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) r_decay = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) @@ -201,170 +692,176 @@ def kda_decode_mtp_kernel( r_state = cute.make_rmem_tensor( cute.make_layout((NUM_V_ROWS * vec_size,), stride=(1,)), cutlass.Float32 ) - if cutlass.const_expr(USE_REG_Q_WEIGHTS): - r_wq = cute.make_rmem_tensor( - cute.make_layout((KERNEL_WIDTH * vec_size,), stride=(1,)), cutlass.Float32 - ) r_exp_A = cutlass.Float32(0.0) - run_precompute = False - if cutlass.const_expr(USE_REGULAR_METADATA) or eos > bos: + # Issue the recurrent-state loads first so their DRAM latency overlaps the + # whole precompute stage (one V pass: TILE_V == V). + # S0 arrives through the TMA bulk load issued at kernel start (mbar_ld); it is read + # from the swizzled smem boxes at the start of the chunk stage. + # Warp-uniform request bounds: the TMA store inside this region needs a provably + # uniform branch condition (legacy shifted layout only; the other layouts are unconditional). + bos = cute.arch.make_warp_uniform(bos) + eos = cute.arch.make_warp_uniform(eos) + if ( + cutlass.const_expr(USE_REGULAR_METADATA or USE_PACKED_TOKEN_LAYOUT or REPLAY_ONLY) + or eos > bos + ): if cutlass.const_expr(FUSE_PRECOMPUTE or RUNTIME_PRECOMPUTE_FLAG): if cutlass.const_expr(RUNTIME_PRECOMPUTE_FLAG): run_precompute = precompute_control[0] != 0 else: run_precompute = True if run_precompute: - for i in range(vec_size): - k_idx = i * 32 + in_warp_tid - for w in range(KERNEL_WIDTH - 1): - r_state[w * vec_size + i] = cutlass.Float32(cs_q[slot, hk_off + k_idx, w]) - r_state[(KERNEL_WIDTH - 1) * vec_size + w * vec_size + i] = cutlass.Float32( - cs_k[slot, hk_off + k_idx, w] - ) + # The conv windows, conv weights and gate constants feed only the new tokens' + # precompute, which REPLAY_ONLY skips. if cutlass.const_expr(not REPLAY_ONLY): - for w in range(KERNEL_WIDTH): - if tidx < K: - if cutlass.const_expr(not USE_REG_Q_WEIGHTS): - sConvW[w * K + tidx] = cutlass.Float32(w_q[hk_off + tidx, w]) - sConvW[k_weight_base + w * K + tidx] = cutlass.Float32( - w_k[hk_off + tidx, w] + # The last K threads stage the committed q/k windows, the q/k conv weights and + # dt_bias (see the load-phase roles above). + ct = tidx - QK_STAGE_TID0 + if ct >= 0: + for w in cutlass.range(KERNEL_WIDTH - 1, unroll_full=True): + sXQ[w * K + ct] = cutlass.Float32(cs_q[slot, hk_off + ct, w]) + sXK[w * K + ct] = cutlass.Float32(cs_k[slot, hk_off + ct, w]) + for w in cutlass.range(KERNEL_WIDTH, unroll_full=True): + if ct >= 0: + sConvW[w * K + ct] = cutlass.Float32(w_q[hk_off + ct, w]) + sConvW[k_weight_base + w * K + ct] = cutlass.Float32( + w_k[hk_off + ct, w] ) - for ld in range(V * KERNEL_WIDTH // NUM_THREADS): + for ld in cutlass.range( + max(1, V * KERNEL_WIDTH // NUM_THREADS), unroll_full=True + ): flat = ld * NUM_THREADS + tidx - sConvW[v_weight_base + flat] = cutlass.Float32( - w_v[hv_off + flat % V, flat // V] - ) - if cutlass.const_expr(USE_REG_Q_WEIGHTS): - if warp_idx == 0: - for _w in range(KERNEL_WIDTH): - for _i in range(vec_size): - r_wq[_w * vec_size + _i] = cutlass.Float32( - w_q[hk_off + _i * 32 + in_warp_tid, _w] + if flat < V * KERNEL_WIDTH: + sConvW[v_weight_base + flat] = cutlass.Float32( + w_v[hv_off + flat % V, flat // V] + ) + if ct >= 0: + sDtBias[ct] = cutlass.Float32(dt_bias[i_h * K + ct]) + if tidx == 0: + sAlog[0] = cutlass.Float32(A_log[i_h]) + cv = tidx - V_STAGE_TID0 + if cv >= 0 and cv < V: + for _w in cutlass.range(KERNEL_WIDTH - 1, unroll_full=True): + sXV[_w * V + cv] = cutlass.Float32(cs_v[slot, hv_off + cv, _w]) + # Load phase: one round trip for every per-token input. Warp w + # stages tokens w, w + num_warps, ...; replay tokens go straight + # into sQ/sK/sG/sBeta/sVall, new tokens stage their raw inputs. + for _rep in cutlass.range( + (t_max + NUM_THREADS // 32 - 1) // (NUM_THREADS // 32), unroll_full=True + ): + _t = _rep * (NUM_THREADS // 32) + warp_idx + if _t < T_loop: + if _t < commit_len: + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + # replayed drafts produce no output: their q rows are unused + sQ[_t, k_idx] = cutlass.Float32(0.0) + sK[_t, k_idx] = cutlass.Float32(k_cache[slot, _t, hk_off + k_idx]) + sG[_t, k_idx] = cutlass.Float32(g_cache[slot, _t, hk_off + k_idx]) + # the drafts' raw conv inputs feed only the new tokens' windows + if cutlass.const_expr(not REPLAY_ONLY): + sXQ[(_t + 3) * K + k_idx] = cutlass.Float32( + cs_q[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + _t] ) - cute.arch.barrier() - if cutlass.const_expr(USE_SETMAXREG): - cute.arch.warpgroup_reg_dealloc(64) - if warp_idx < 3: - if cutlass.const_expr(not REPLAY_ONLY): - if warp_idx == 2: - r_exp_A = cute.math.exp(cutlass.Float32(A_log[i_h]), fastmath=True) - i_t = 0 - while i_t < T_loop: - if cutlass.const_expr(USE_ZERO_ACCEPTED): - replay_from_cache = False - elif cutlass.const_expr(REPLAY_ONLY): - replay_from_cache = True + sXK[(_t + 3) * K + k_idx] = cutlass.Float32( + cs_k[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + _t] + ) + sXV[(_t + 3) * V + k_idx] = cutlass.Float32( + cs_v[slot, hv_off + k_idx, KERNEL_WIDTH - 1 + _t] + ) + sVall[_t * V + k_idx] = cutlass.Float32( + v_cache[slot, _t, hv_off + k_idx] + ) + if in_warp_tid == 0: + sBeta[_t] = cutlass.Float32(beta_cache[slot, _t, i_hv]) else: - replay_from_cache = i_t < commit_len - if replay_from_cache: - if warp_idx == 0: - for i in range(vec_size): - k_idx = i * 32 + in_warp_tid - sQ[i_t, k_idx] = cutlass.Float32( - qkg_cache[slot, i_t, 0, hk_off + k_idx] + _tok = bos + _t + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + if cutlass.const_expr(USE_FLAT_LAYOUT): + sXQ[(_t + 3) * K + k_idx] = cutlass.Float32( + x_q[0, _tok, hk_off + k_idx] ) - for i in range(vec_size): - k_idx = i * 32 + in_warp_tid - r_xq_raw = cutlass.Float32( - cs_q[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + i_t] + sXK[(_t + 3) * K + k_idx] = cutlass.Float32( + x_k[0, _tok, hk_off + k_idx] ) - for w in range(KERNEL_WIDTH - 2): - r_state[w * vec_size + i] = r_state[(w + 1) * vec_size + i] - r_state[(KERNEL_WIDTH - 2) * vec_size + i] = r_xq_raw - elif warp_idx == 1: - for i in range(vec_size): - k_idx = i * 32 + in_warp_tid - sK[i_t, k_idx] = cutlass.Float32( - qkg_cache[slot, i_t, 1, hk_off + k_idx] + sXV[(_t + 3) * V + k_idx] = cutlass.Float32( + x_v[0, _tok, hv_off + k_idx] ) - if in_warp_tid == 0: - sBeta[i_t] = cutlass.Float32(beta_cache[slot, i_t, i_hv]) - for i in range(vec_size): - k_idx = i * 32 + in_warp_tid - r_xk_raw = cutlass.Float32( - cs_k[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + i_t] + else: + sXQ[(_t + 3) * K + k_idx] = cutlass.Float32( + x_q[0, _tok, i_h, k_idx] ) - for w in range(KERNEL_WIDTH - 2): - r_state[ - (KERNEL_WIDTH - 1) * vec_size + w * vec_size + i - ] = r_state[ - (KERNEL_WIDTH - 1) * vec_size + (w + 1) * vec_size + i - ] - r_state[ - (KERNEL_WIDTH - 1) * vec_size - + (KERNEL_WIDTH - 2) * vec_size - + i - ] = r_xk_raw - else: - for i in range(vec_size): - k_idx = i * 32 + in_warp_tid - r_gk_c2 = cutlass.Float32( - qkg_cache[slot, i_t, 2, hk_off + k_idx] + sXK[(_t + 3) * K + k_idx] = cutlass.Float32( + x_k[0, _tok, i_h, k_idx] ) - sG[i_t, k_idx] = cute.math.exp(r_gk_c2, fastmath=True) - else: - token = bos + i_t - if warp_idx == 0: - for i_pair in range(vec_size // 2): - i0 = i_pair * 2 - i1 = i_pair * 2 + 1 - k_idx0 = i0 * 32 + in_warp_tid - k_idx1 = i1 * 32 + in_warp_tid - r_conv_0 = 0.0 - r_conv_1 = 0.0 - if cutlass.const_expr(USE_REG_Q_WEIGHTS): - _cwq_0 = r_wq[0 * vec_size + i0] - _cwq_1 = r_wq[0 * vec_size + i1] - r_conv_0 += r_state[0 * vec_size + i0] * _cwq_0 - r_conv_1 += r_state[0 * vec_size + i1] * _cwq_1 - _cwq_0 = r_wq[1 * vec_size + i0] - _cwq_1 = r_wq[1 * vec_size + i1] - r_conv_0 += r_state[1 * vec_size + i0] * _cwq_0 - r_conv_1 += r_state[1 * vec_size + i1] * _cwq_1 - _cwq_0 = r_wq[2 * vec_size + i0] - _cwq_1 = r_wq[2 * vec_size + i1] - r_conv_0 += r_state[2 * vec_size + i0] * _cwq_0 - r_conv_1 += r_state[2 * vec_size + i1] * _cwq_1 - else: - _cwq_0 = cutlass.Float32(0.0) - _cwq_1 = cutlass.Float32(0.0) - for w in range(KERNEL_WIDTH - 1): - _cwq_0 = sConvW[w * K + i0 * 32 + in_warp_tid] - _cwq_1 = sConvW[w * K + i1 * 32 + in_warp_tid] - r_conv_0 += r_state[w * vec_size + i0] * _cwq_0 - r_conv_1 += r_state[w * vec_size + i1] * _cwq_1 - if cutlass.const_expr(USE_FLAT_LAYOUT): - r_xq_0 = cutlass.Float32(x_q[0, token, hk_off + k_idx0]) - r_xq_1 = cutlass.Float32(x_q[0, token, hk_off + k_idx1]) - else: - r_xq_0 = cutlass.Float32(x_q[0, token, i_h, k_idx0]) - r_xq_1 = cutlass.Float32(x_q[0, token, i_h, k_idx1]) - if cutlass.const_expr(USE_REG_Q_WEIGHTS): - _cwq_last_0 = r_wq[(KERNEL_WIDTH - 1) * vec_size + i0] - _cwq_last_1 = r_wq[(KERNEL_WIDTH - 1) * vec_size + i1] - else: - _cwq_last_0 = sConvW[ - (KERNEL_WIDTH - 1) * K + i0 * 32 + in_warp_tid - ] - _cwq_last_1 = sConvW[ - (KERNEL_WIDTH - 1) * K + i1 * 32 + in_warp_tid - ] - r_conv_0 += r_xq_0 * _cwq_last_0 - r_conv_1 += r_xq_1 * _cwq_last_1 + sXV[(_t + 3) * V + k_idx] = cutlass.Float32( + x_v[0, _tok, i_hv, k_idx] + ) + sGin[_t * K + k_idx] = cutlass.Float32(g[0, _tok, i_hv, k_idx]) + if in_warp_tid == 0: + sBetaRaw[_t] = cutlass.Float32(beta[0, _tok, i_hv]) + # Padding rows [T_loop, t_max) of the gates, values and betas are zeroed (under the + # load latency) so the chunk stage's r_Y / cumulative-gate / V' loops are branch-free. + if tidx < K: + for _z in cutlass.range(T_loop, t_max, 1): + sG[_z, tidx] = cutlass.Float32(0.0) + sVall[_z * V + tidx] = cutlass.Float32(0.0) + if tidx == 0: + sBeta[_z] = cutlass.Float32(0.0) + if cutlass.const_expr(FUSE_OUTPUT_NORM): + # Threads 224+ stage the gated-RMSNorm operands for the epilogue. + if tidx >= 224: + _pf_tid = tidx - 224 + for _e in cutlass.range( + ((1 + NUM_SPEC) * V + (NUM_THREADS - 224) - 1) // (NUM_THREADS - 224), + unroll_full=True, + ): + _flat = _e * (NUM_THREADS - 224) + _pf_tid + if _flat < (1 + NUM_SPEC) * V: + _tok = _flat // V + _vv = _flat % V + sOnG[_flat] = cutlass.Float32( + onorm_g[0, i_n * (1 + NUM_SPEC) + _tok, i_hv, _vv] + ) + for _e in cutlass.range( + (V + (NUM_THREADS - 224) - 1) // (NUM_THREADS - 224), unroll_full=True + ): + _flat = _e * (NUM_THREADS - 224) + _pf_tid + if _flat < V: + sOnW[_flat] = cutlass.Float32(onorm_weight[_flat]) + cute.arch.barrier() + # ---- per-token-parallel precompute ---- + # (REPLAY_ONLY has no new tokens: nothing to precompute, convolve or cache) + if cutlass.const_expr(not REPLAY_ONLY): + # Cluster phase 1 of 2: the precompute below overwrites the conv windows and + # replay caches that both CTAs of the pair read in the load phase. The CTA + # barrier above completed this CTA's reads; the cluster barrier makes the + # peer's complete too before CTA 0, the only writer, stores. Phase 2 is the + # S0 arrive / wait pair around the state store. + if cutlass.const_expr(SPLIT_V == 2): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + num_warps = NUM_THREADS // 32 + half = num_warps // 2 + role = warp_idx // half + r_exp_A = cute.math.exp(sAlog[0], fastmath=True) + for _rep in cutlass.range((1 + NUM_SPEC + half - 1) // half, unroll_full=True): + j = _rep * half + warp_idx % half + if j < 1 + NUM_SPEC: + t = commit_len + j + if role == 0: + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + r_conv_0 = cutlass.Float32(0.0) + for w in cutlass.range(KERNEL_WIDTH, unroll_full=True): + r_conv_0 += sXQ[(t + w) * K + k_idx] * sConvW[w * K + k_idx] e0 = cute.math.exp(-r_conv_0, fastmath=True) - e1 = cute.math.exp(-r_conv_1, fastmath=True) - sig_0 = cute.arch.rcp_approx(cutlass.Float32(1.0) + e0) - sig_1 = cute.arch.rcp_approx(cutlass.Float32(1.0) + e1) - r_q[i0] = r_conv_0 * sig_0 - r_q[i1] = r_conv_1 * sig_1 - r_state[0 * vec_size + i0] = r_state[1 * vec_size + i0] - r_state[0 * vec_size + i1] = r_state[1 * vec_size + i1] - r_state[1 * vec_size + i0] = r_state[2 * vec_size + i0] - r_state[1 * vec_size + i1] = r_state[2 * vec_size + i1] - r_state[2 * vec_size + i0] = r_xq_0 - r_state[2 * vec_size + i1] = r_xq_1 + r_q[i] = r_conv_0 * cute.arch.rcp_approx( + cutlass.Float32(1.0) + e0 + ) sum_q = 0.0 - for i in range(vec_size): + for i in cutlass.range(vec_size, unroll_full=True): sum_q += r_q[i] * r_q[i] for offset in [16, 8, 4, 2, 1]: sum_q += cute.arch.shuffle_sync_bfly( @@ -373,438 +870,1137 @@ def kda_decode_mtp_kernel( rnorm_q_scaled = ( cute.math.rsqrt(sum_q + 1e-06, fastmath=True) * scale ) - for i in range(vec_size): + for i in cutlass.range(vec_size, unroll_full=True): r_q[i] = r_q[i] * rnorm_q_scaled - for i in range(vec_size): + for i in cutlass.range(vec_size, unroll_full=True): k_idx = i * 32 + in_warp_tid - sQ[i_t, k_idx] = r_q[i] - elif warp_idx == 1: - r_b_raw = cutlass.Float32(0.0) - if in_warp_tid == 0: - r_b_raw = cutlass.Float32(beta[0, token, i_hv]) - for i in range(vec_size): + sQ[t, k_idx] = r_q[i] + for i in cutlass.range(vec_size, unroll_full=True): k_idx = i * 32 + in_warp_tid r_conv = ( - r_state[(KERNEL_WIDTH - 1) * vec_size + 0 * vec_size + i] - * sConvW[k_weight_base + 0 * K + i * 32 + in_warp_tid] - ) - r_conv += ( - r_state[(KERNEL_WIDTH - 1) * vec_size + 1 * vec_size + i] - * sConvW[k_weight_base + 1 * K + i * 32 + in_warp_tid] - ) - r_conv += ( - r_state[(KERNEL_WIDTH - 1) * vec_size + 2 * vec_size + i] - * sConvW[k_weight_base + 2 * K + i * 32 + in_warp_tid] - ) - if cutlass.const_expr(USE_FLAT_LAYOUT): - r_xk = cutlass.Float32(x_k[0, token, hk_off + k_idx]) - else: - r_xk = cutlass.Float32(x_k[0, token, i_h, k_idx]) - r_conv += ( - r_xk - * sConvW[ - k_weight_base - + (KERNEL_WIDTH - 1) * K - + i * 32 - + in_warp_tid - ] + sXK[t * K + k_idx] * sConvW[k_weight_base + 0 * K + k_idx] ) + for w in cutlass.range(1, KERNEL_WIDTH, unroll_full=True): + r_conv += ( + sXK[(t + w) * K + k_idx] + * sConvW[k_weight_base + w * K + k_idx] + ) r_conv = r_conv * cute.arch.rcp_approx( cutlass.Float32(1.0) + cute.math.exp(-r_conv, fastmath=True) ) r_k[i] = r_conv - r_state[(KERNEL_WIDTH - 1) * vec_size + 0 * vec_size + i] = ( - r_state[(KERNEL_WIDTH - 1) * vec_size + 1 * vec_size + i] - ) - r_state[(KERNEL_WIDTH - 1) * vec_size + 1 * vec_size + i] = ( - r_state[(KERNEL_WIDTH - 1) * vec_size + 2 * vec_size + i] - ) - r_state[(KERNEL_WIDTH - 1) * vec_size + 2 * vec_size + i] = r_xk sum_k = 0.0 - for i in range(vec_size): + for i in cutlass.range(vec_size, unroll_full=True): sum_k += r_k[i] * r_k[i] for offset in [16, 8, 4, 2, 1]: sum_k += cute.arch.shuffle_sync_bfly( sum_k, offset=offset, mask=-1, mask_and_clamp=31 ) rnorm_k = cute.math.rsqrt(sum_k + 1e-06, fastmath=True) - for i in range(vec_size): + for i in cutlass.range(vec_size, unroll_full=True): r_k[i] = r_k[i] * rnorm_k - for i in range(vec_size): + for i in cutlass.range(vec_size, unroll_full=True): k_idx = i * 32 + in_warp_tid - sK[i_t, k_idx] = r_k[i] - if in_warp_tid == 0: - sBeta[i_t] = cute.arch.rcp_approx( - cutlass.Float32(1.0) - + cute.math.exp(-r_b_raw, fastmath=True) - ) + sK[t, k_idx] = r_k[i] + # Both CTAs of a pair compute the same values; CTA 0 stores them. + if cutlass.const_expr(SPLIT_V == 1) or vh == 0: + if j == 0: + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + for w in cutlass.range( + KERNEL_WIDTH - 1, unroll_full=True + ): + cs_q[slot, hk_off + k_idx, w] = sXQ[ + (t + 1 + w) * K + k_idx + ].to(CONV_DTYPE) + cs_k[slot, hk_off + k_idx, w] = sXK[ + (t + 1 + w) * K + k_idx + ].to(CONV_DTYPE) + else: + cache_pos = j - 1 + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + k_cache[slot, cache_pos, hk_off + k_idx] = r_k[i].to( + K_CACHE_DTYPE + ) + cs_q[ + slot, hk_off + k_idx, KERNEL_WIDTH - 1 + cache_pos + ] = sXQ[(t + 3) * K + k_idx].to(CONV_DTYPE) + cs_k[ + slot, hk_off + k_idx, KERNEL_WIDTH - 1 + cache_pos + ] = sXK[(t + 3) * K + k_idx].to(CONV_DTYPE) else: - for i in range(vec_size): + for i in cutlass.range(vec_size, unroll_full=True): k_idx = i * 32 + in_warp_tid - r_g_raw = cutlass.Float32(g[0, token, i_hv, k_idx]) - r_g_raw = r_g_raw + cutlass.Float32(dt_bias[i_h * K + k_idx]) + r_g_raw = sGin[t * K + k_idx] + r_g_raw = r_g_raw + sDtBias[k_idx] exp_A_x = r_exp_A * r_g_raw sigmoid_val = cute.arch.rcp_approx( cutlass.Float32(1.0) + cute.math.exp(-exp_A_x, fastmath=True) ) r_gk = lower_bound * sigmoid_val - sG[i_t, k_idx] = cute.math.exp(r_gk, fastmath=True) + sG[t, k_idx] = r_gk r_decay[i] = r_gk - if i_t > commit_len: - cache_pos = i_t - commit_len - 1 - for i in range(vec_size): - k_idx = i * 32 + in_warp_tid - qkg_cache[slot, cache_pos, 2, hk_off + k_idx] = r_decay[i] - if i_t == commit_len: - if warp_idx == 0: - for i in range(vec_size): - k_idx = i * 32 + in_warp_tid - for w in range(KERNEL_WIDTH - 1): - cs_q[slot, hk_off + k_idx, w] = r_state[ - w * vec_size + i - ] - elif warp_idx == 1: - for i in range(vec_size): - k_idx = i * 32 + in_warp_tid - for w in range(KERNEL_WIDTH - 1): - cs_k[slot, hk_off + k_idx, w] = r_state[ - (KERNEL_WIDTH - 1) * vec_size + w * vec_size + i - ] - if i_t > commit_len: - cache_pos = i_t - commit_len - 1 - if warp_idx == 0: - for i in range(vec_size): - k_idx = i * 32 + in_warp_tid - qkg_cache[slot, cache_pos, 0, hk_off + k_idx] = sQ[ - i_t, k_idx - ] - for i in range(vec_size): - k_idx = i * 32 + in_warp_tid - cs_q[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + cache_pos] = ( - r_state[(KERNEL_WIDTH - 2) * vec_size + i] - ) - elif warp_idx == 1: - for i in range(vec_size): - k_idx = i * 32 + in_warp_tid - qkg_cache[slot, cache_pos, 1, hk_off + k_idx] = sK[ - i_t, k_idx - ] - if in_warp_tid == 0: - beta_cache[slot, cache_pos, i_hv] = sBeta[i_t] - for i in range(vec_size): - k_idx = i * 32 + in_warp_tid - cs_k[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + cache_pos] = ( - r_state[ - (KERNEL_WIDTH - 1) * vec_size - + (KERNEL_WIDTH - 2) * vec_size - + i - ] - ) - i_t = i_t + 1 - else: - _v_idx = tidx - 96 - if _v_idx < V: - _csv0 = cutlass.Float32(cs_v[slot, hv_off + _v_idx, 0]) - _csv1 = cutlass.Float32(cs_v[slot, hv_off + _v_idx, 1]) - _csv2 = cutlass.Float32(cs_v[slot, hv_off + _v_idx, 2]) - if cutlass.const_expr(USE_ZERO_ACCEPTED): - _wv0 = sConvW[v_weight_base + 0 * V + _v_idx] - _wv1 = sConvW[v_weight_base + 1 * V + _v_idx] - _wv2 = sConvW[v_weight_base + 2 * V + _v_idx] - _wv3 = sConvW[v_weight_base + (KERNEL_WIDTH - 1) * V + _v_idx] - if cutlass.const_expr(USE_FLAT_LAYOUT): - _xv0 = cutlass.Float32(x_v[0, bos + 0, hv_off + _v_idx]) - _xv1 = cutlass.Float32(x_v[0, bos + 1, hv_off + _v_idx]) - _xv2 = cutlass.Float32(x_v[0, bos + 2, hv_off + _v_idx]) - else: - _xv0 = cutlass.Float32(x_v[0, bos + 0, i_hv, _v_idx]) - _xv1 = cutlass.Float32(x_v[0, bos + 1, i_hv, _v_idx]) - _xv2 = cutlass.Float32(x_v[0, bos + 2, i_hv, _v_idx]) - _vconv0 = _csv0 * _wv0 - _vconv0 += _csv1 * _wv1 - _vconv0 += _csv2 * _wv2 - _vconv0 += _xv0 * _wv3 - _vconv0 = _vconv0 * cute.arch.rcp_approx( - cutlass.Float32(1.0) + cute.math.exp(-_vconv0, fastmath=True) - ) - sVall[0 * V + _v_idx] = _vconv0 - cs_v[slot, hv_off + _v_idx, 0] = _csv1 - cs_v[slot, hv_off + _v_idx, 1] = _csv2 - cs_v[slot, hv_off + _v_idx, 2] = _xv0 - _vconv1, _vconv2 = cute.arch.mul_packed_f32x2( - (_csv1, _csv2), (_wv0, _wv0) - ) - _vconv1, _vconv2 = cute.arch.fma_packed_f32x2( - (_csv2, _xv0), (_wv1, _wv1), (_vconv1, _vconv2) - ) - _vconv1, _vconv2 = cute.arch.fma_packed_f32x2( - (_xv0, _xv1), (_wv2, _wv2), (_vconv1, _vconv2) - ) - _vconv1, _vconv2 = cute.arch.fma_packed_f32x2( - (_xv1, _xv2), (_wv3, _wv3), (_vconv1, _vconv2) - ) - _vconv1 = _vconv1 * cute.arch.rcp_approx( - cutlass.Float32(1.0) + cute.math.exp(-_vconv1, fastmath=True) - ) - _vconv2 = _vconv2 * cute.arch.rcp_approx( - cutlass.Float32(1.0) + cute.math.exp(-_vconv2, fastmath=True) - ) - sVall[1 * V + _v_idx] = _vconv1 - v_cache[slot, 0, hv_off + _v_idx] = _vconv1 - cs_v[slot, hv_off + _v_idx, KERNEL_WIDTH - 1] = _xv1 - sVall[2 * V + _v_idx] = _vconv2 - v_cache[slot, 1, hv_off + _v_idx] = _vconv2 - cs_v[slot, hv_off + _v_idx, KERNEL_WIDTH] = _xv2 - else: - _i_t = 0 - while _i_t < T_loop: - if _i_t < commit_len: - sVall[_i_t * V + _v_idx] = cutlass.Float32( - v_cache[slot, _i_t, hv_off + _v_idx] - ) - _xv_replay = cutlass.Float32( - cs_v[slot, hv_off + _v_idx, KERNEL_WIDTH - 1 + _i_t] - ) - _csv0 = _csv1 - _csv1 = _csv2 - _csv2 = _xv_replay - else: - _token_v = bos + _i_t + r_beta_t = cute.arch.rcp_approx( + cutlass.Float32(1.0) + + cute.math.exp(-sBetaRaw[t], fastmath=True) + ) + if in_warp_tid == 0: + sBeta[t] = r_beta_t + for i in cutlass.range(vec_size, unroll_full=True): + v_idx = i * 32 + in_warp_tid _v_conv = 0.0 - _v_conv += _csv0 * sConvW[v_weight_base + 0 * V + _v_idx] - _v_conv += _csv1 * sConvW[v_weight_base + 1 * V + _v_idx] - _v_conv += _csv2 * sConvW[v_weight_base + 2 * V + _v_idx] - if cutlass.const_expr(USE_FLAT_LAYOUT): - _xv = cutlass.Float32(x_v[0, _token_v, hv_off + _v_idx]) - else: - _xv = cutlass.Float32(x_v[0, _token_v, i_hv, _v_idx]) - _v_conv += ( - _xv - * sConvW[v_weight_base + (KERNEL_WIDTH - 1) * V + _v_idx] - ) + for w in cutlass.range(KERNEL_WIDTH, unroll_full=True): + _v_conv += ( + sXV[(t + w) * V + v_idx] + * sConvW[v_weight_base + w * V + v_idx] + ) _v_conv = _v_conv * cute.arch.rcp_approx( cutlass.Float32(1.0) + cute.math.exp(-_v_conv, fastmath=True) ) - sVall[_i_t * V + _v_idx] = _v_conv - _csv0 = _csv1 - _csv1 = _csv2 - _csv2 = _xv - if _i_t == commit_len: - cs_v[slot, hv_off + _v_idx, 0] = _csv0 - cs_v[slot, hv_off + _v_idx, 1] = _csv1 - cs_v[slot, hv_off + _v_idx, 2] = _csv2 - if _i_t > commit_len: - _cp = _i_t - commit_len - 1 - v_cache[slot, _cp, hv_off + _v_idx] = _v_conv - cs_v[slot, hv_off + _v_idx, KERNEL_WIDTH - 1 + _cp] = _xv - _i_t = _i_t + 1 + sVall[t * V + v_idx] = _v_conv + r_bk[i] = _v_conv + if cutlass.const_expr(SPLIT_V == 1) or vh == 0: + if j == 0: + for i in cutlass.range(vec_size, unroll_full=True): + v_idx = i * 32 + in_warp_tid + for w in cutlass.range( + KERNEL_WIDTH - 1, unroll_full=True + ): + cs_v[slot, hv_off + v_idx, w] = sXV[ + (t + 1 + w) * V + v_idx + ].to(CONV_DTYPE) + else: + cache_pos = j - 1 + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + g_cache[slot, cache_pos, hk_off + k_idx] = r_decay[i] + v_cache[slot, cache_pos, hv_off + k_idx] = r_bk[i].to( + V_CACHE_DTYPE + ) + cs_v[ + slot, hv_off + k_idx, KERNEL_WIDTH - 1 + cache_pos + ] = sXV[(t + 3) * V + k_idx].to(CONV_DTYPE) + if in_warp_tid == 0: + beta_cache[slot, cache_pos, i_hv] = r_beta_t.to( + BETA_CACHE_DTYPE + ) if cutlass.const_expr(PROFILE_STAGES): cute.arch.barrier() t_stage1 = read_globaltimer() else: cute.arch.barrier() - if cutlass.const_expr(USE_SETMAXREG): - cute.arch.warpgroup_reg_dealloc(64) if cutlass.const_expr(PROFILE_STAGES): cute.arch.barrier() t_stage1 = read_globaltimer() else: cute.arch.barrier() - if cutlass.const_expr(USE_SETMAXREG): - cute.arch.warpgroup_reg_dealloc(64) if cutlass.const_expr(PROFILE_STAGES): cute.arch.barrier() t_stage1 = read_globaltimer() - for row in range(NUM_V_ROWS): - v_row = warp_idx * NUM_V_ROWS + row - for i in range(vec_size): - if cutlass.const_expr(USE_FLAT_LAYOUT): - r_state[row * vec_size + i] = cutlass.Float32( - h0[h0_idx, v_row, i * 32 + in_warp_tid] - ) + cute.arch.barrier() + # ================================================================= + # Chunk (WY-form) recurrence on tensor cores. + # ================================================================= + tmem_ptr = cute.arch.retrieve_tmem_ptr(cutlass.Float32, 16, tmem_holder) + pool = cutlass.utils.TmemBufferPool(tmem_ptr, TMEM_COLS) + tD4_frag = tiled_mma4.make_fragment_C(tiled_mma4.partition_shape_C((V, T_PAD))) + tD4 = pool.allocate_tensor(tD4_frag.layout, cutlass.Float32) # outputs (columns 0..15) + tD2h_frag = tiled_mma2h.make_fragment_C(tiled_mma2h.partition_shape_C((V, K // 2))) + tD2h0 = pool.allocate_tensor( + tD2h_frag.layout, cutlass.Float32 + ) # GEMM2 half 0: k columns 0..63 + tD2h1 = pool.allocate_tensor( + tD2h_frag.layout, cutlass.Float32 + ) # GEMM2 half 1: k columns 64..127 + tD2q_frag = tiled_mma2q.make_fragment_C(tiled_mma2q.partition_shape_C((V, T_PAD + K // 2))) + tD2q = cute.make_tensor(tD4.iterator, tD2q_frag.layout) # [D4 | ΔS half 0] accumulator view + tA_frag = tiled_mma_f32a.make_fragment_A(tiled_mma_f32a.partition_shape_A((V, K))) + if cutlass.const_expr(BF16_MMA): + # S0 as packed bf16 (even k in the low half of each 32-bit column): 64 columns, staged + # through an fp32 word view and consumed by the MMA through the bf16 A-fragment view. + tAw = pool.allocate_tensor( + cute.make_layout(((V, 8), 1, K // 16), stride=((65536, 1), 0, 8)), cutlass.Float32 + ) + tAb = cute.make_tensor( + cute.recast_ptr(tAw.iterator, dtype=cutlass.BFloat16), + tiled_mma1.make_fragment_A(tiled_mma1.partition_shape_A((V, K))).layout, + ) + tAhi = tAw + tAlo = tAw + else: + tAhi = pool.allocate_tensor(tA_frag.layout, cutlass.Float32) + tAlo = pool.allocate_tensor(tA_frag.layout, cutlass.Float32) + tD1_frag = tiled_mma1.make_fragment_C(tiled_mma1.partition_shape_C((V, N1))) + tD1 = pool.allocate_tensor(tD1_frag.layout, cutlass.Float32) + tD3_frag = tiled_mma3.make_fragment_C(tiled_mma3.partition_shape_C((4 * T_PAD, T_PAD))) + tD3 = pool.allocate_tensor(tD3_frag.layout, cutlass.Float32) + tD2 = cute.make_tensor( + tD2h0.iterator, tA_frag.layout + ) # A-view layout for st/ld (all 128 columns) + + st32 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x32), cutlass.Float32) + ld16 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + ld32 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), cutlass.Float32) + tiled_stA = tcgen05.make_tmem_copy(st32, tD2) + thr_stA = tiled_stA.get_slice(wg_tidx) + tD2_st = thr_stA.partition_D(tD2) + rA_shape = thr_stA.partition_S(tD2).shape + if cutlass.const_expr(BF16_MMA): + st16w = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + thr_stAw = tcgen05.make_tmem_copy(st16w, tAw).get_slice(wg_tidx) + tAw_st = thr_stAw.partition_D(tAw) + r_pack = cute.make_rmem_tensor(thr_stAw.partition_S(tAw).shape[:-1], cutlass.Float32) + r_pack_flat = cute.make_tensor(r_pack.iterator, cute.make_layout((16,))) + tAhi_st = tD2_st + tAlo_st = tD2_st + else: + tAhi_st = thr_stA.partition_D(tAhi) + tAlo_st = thr_stA.partition_D(tAlo) + tiled_ldD2 = tcgen05.make_tmem_copy(ld32, tD2) + thr_ldD2 = tiled_ldD2.get_slice(wg_tidx) + tD2_ld = thr_ldD2.partition_S(tD2) + rS = cute.make_rmem_tensor(thr_ldD2.partition_D(tD2).shape[:-1], cutlass.Float32) + rS_flat = cute.make_tensor(rS.iterator, cute.make_layout((32,))) + # TMA store partitions for the NUM_SLABS boxes of this CTA's committed-state tile. + # Created here, outside the dynamic request-bounds region (the exec-TMA lowering is + # only legal at this level); the copies are issued inside. + sBoxf = cute.make_tensor(sScratch.iterator, cute.make_layout((V * K,))) + sBoxh = cute.make_tensor( + cute.recast_ptr(sScratch.iterator, dtype=cutlass.BFloat16), cute.make_layout((V * K,)) + ) + tS_list = [] + tG_list = [] + for sl_h in cutlass.range_constexpr(NUM_SLABS): + sT_h = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + sl_h * ((V // SPLIT_V) * 32), + box_layout_t.inner, + dtype=STATE_DTYPE, + ), + box_layout_t.outer, + ) + gT_h_all = cute.local_tile(mT_tma, (V // SPLIT_V, BOX_COLS, 1), (vh, sl_h, None)) + if cutlass.const_expr(USE_FLAT_LAYOUT): + gT_h = gT_h_all[(None, None, 0, h0_idx)] + else: + gT_h = gT_h_all[(None, None, 0, (i_hv, slot))] + tS_h, tG_h = cpasync.tma_partition( + tma_atom_t, + 0, + cute.make_layout(1), + cute.group_modes(sT_h, 0, 2), + cute.group_modes(gT_h, 0, 2), + ) + tS_list.append(tS_h) + tG_list.append(tG_h) + r_hi = cute.make_rmem_tensor(rA_shape[:-1], cutlass.Float32) + r_lo = cute.make_rmem_tensor(rA_shape[:-1], cutlass.Float32) + r_hi_flat = cute.make_tensor(r_hi.iterator, cute.make_layout((32,))) + r_lo_flat = cute.make_tensor(r_lo.iterator, cute.make_layout((32,))) + + # ---- S0 row v = wg_tidx, columns [32 wg, +32), from the TMA-loaded swizzled box + # (16 B chunks: conflict-free LDS.128 under the 128B swizzle) + cute.arch.mbarrier_wait(mbar_ld, 0) + # in-place state update: the peer CTA must not store its half before our S0 load + # completed -> arrive here, wait right before the state store. + if cutlass.const_expr(SPLIT_V == 2): + cute.arch.cluster_arrive_relaxed() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq28 = read_globaltimer() + _, _gq28, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq28 + i_n) * 64 + 28] = _tq28 - t_stage0 + vx = 4 * (wg_tidx % 8) + if cutlass.const_expr(STATE_BF16): + # bf16 box: 64 columns per 128 B row, so this warpgroup's 32 columns are the 16 B + # chunks [4 (wg % 2), +4) of box wg // 2. A bf16 value is exact in tf32: no lo part, + # and the S0_lo sweep of GEMM1 is skipped. + for c in cutlass.range(4, unroll_full=True): + chunk = (4 * (wg % 2) + c) ^ (wg_tidx % 8) + off_l = (wg // 2) * (V * 64) + 64 * wg_tidx + 8 * chunk + for i in cutlass.range(8, unroll_full=True): + r_state[8 * c + i] = cutlass.Float32(sBoxh[off_l + i]) + if cutlass.const_expr(not BF16_MMA): + for i in cutlass.range(32, unroll_full=True): + r_hi_flat[i] = r_state[i] + for j in cutlass.range_constexpr(4): + if wg == j: + cute.copy(st32, r_hi, tAhi_st[(None, None, None, j)]) + else: + for c in cutlass.range(8, unroll_full=True): + off_l = wg * (V * 32) + 32 * wg_tidx + ((4 * c) ^ vx) + for i in cutlass.range(4, unroll_full=True): + r_state[4 * c + i] = sBoxf[off_l + i] + if cutlass.const_expr(not BF16_MMA): + for i in cutlass.range(32, unroll_full=True): + s_hi = round_hi(r_state[i], BF16_MMA) + r_hi_flat[i] = s_hi + if cutlass.const_expr(S0_LO): + r_lo_flat[i] = r_state[i] - s_hi + for j in cutlass.range_constexpr(4): + if wg == j: + cute.copy(st32, r_hi, tAhi_st[(None, None, None, j)]) + if cutlass.const_expr(S0_LO): + cute.copy(st32, r_lo, tAlo_st[(None, None, None, j)]) + if cutlass.const_expr(BF16_MMA): + # this warpgroup's 32 state columns -> 16 packed bf16x2 words -> TMEM word slab wg + for i in cutlass.range(16, unroll_full=True): + r_pack_flat[i] = cvt_bf16x2(r_state[2 * i + 1], r_state[2 * i]) + for j in cutlass.range_constexpr(4): + if wg == j: + cute.copy(st16w, r_pack, tAw_st[(None, None, None, j)]) + cute.arch.fence_view_async_tmem_store() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq29 = read_globaltimer() + _, _gq29, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq29 + i_n) * 64 + 29] = _tq29 - t_stage0 + + # ---- cumulative gates per k column (thread kq, quarter q owns tokens t = q mod 4): + # G_t -> sE, the T x T coupling factors, and the GEMM1 / GEMM3 operand rows written + # with closed-form 128B-swizzle offsets (lanes = consecutive k: conflict-free). + kq = tidx % K + quarter = tidx // K + t_mid = T_loop // 2 + r_G = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + g_run = cutlass.Float32(0.0) + for t in cutlass.range(T_PAD, unroll_full=True): + if t < t_max: + g_run = g_run + sG[t, kq] + r_G[t] = g_run + g_mid = cutlass.Float32(0.0) + for t in cutlass.range(T_PAD, unroll_full=True): + if t == t_mid: + g_mid = r_G[t] + # Factors of the T x T coupling: e^{G_t - G_mid} (t side) and e^{G_mid - G_s} (s side), + # clamped at e^87 so a chunk whose decay exceeds the fp32 range yields ~0 instead of + # inf/NaN (the affected products are < e^{-|g|} and negligible either way). + exp_clamp = cutlass.Float32(87.0) + # the scratch region (S0 boxes) is reused for the MMA operand tiles below + cute.arch.barrier() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq30 = read_globaltimer() + _, _gq30, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq30 + i_n) * 64 + 30] = _tq30 - t_stage0 + if cutlass.const_expr(REPLAY_ONLY): + # no draft pending: the commit row is -1 and the state is committed unchanged + # (e^{G_-1} = 1, and no row contributes to Kbar) + if (quarter == 0) & (commit_idx < 0): + sEc[kq] = cutlass.Float32(1.0) + ka = kq // 32 + kc = kq % 32 + b1_base = ka * (3 * T_PAD * 32) + a3_base = ka * (4 * T_PAD * 32) + b3_base = ka * (T_PAD * 32) + b1k_base = ka * (T_PAD * 32) + for t in cutlass.range(T_PAD, unroll_full=True): + if quarter == t % 4: + kt_hi = cutlass.Float32(0.0) + kt_lo = cutlass.Float32(0.0) + qt_r = cutlass.Float32(0.0) + kf = cutlass.Float32(0.0) + qf_hi = cutlass.Float32(0.0) + qf_lo = cutlass.Float32(0.0) + kb_hi = cutlass.Float32(0.0) + kb_lo = cutlass.Float32(0.0) + if t < T_loop: + d_t = r_G[t] - g_mid + e_t = cute.math.exp(-cute.arch.fmax(-d_t, -exp_clamp), fastmath=True) + e_bt = cute.math.exp(-cute.arch.fmax(d_t, -exp_clamp), fastmath=True) + eg_t = cute.math.exp(r_G[t], fastmath=True) + sE[t * K + kq] = r_G[t] + if t == commit_idx: + sEc[kq] = eg_t + k_t = sK[t, kq] + q_t = sQ[t, kq] + kt = k_t * eg_t + kt_hi = round_hi(kt, BF16_MMA) + qt_r = round_hi(q_t * eg_t, BF16_MMA) + kf = round_hi(k_t * e_t, BF16_MMA) + qf = q_t * e_t + qf_hi = round_hi(qf, BF16_MMA) + kb = k_t * e_bt + kb_hi = round_hi(kb, BF16_MMA) + if cutlass.const_expr(not BF16_MMA): + kt_lo = kt - kt_hi + qf_lo = qf - qf_hi + kb_lo = kb - kb_hi + if cutlass.const_expr(BF16_MMA): + sB1b_f[sw128_bf16(t, kq, 2 * T_PAD)] = cutlass.BFloat16(kt_hi) + sB1b_f[sw128_bf16(T_PAD + t, kq, 2 * T_PAD)] = cutlass.BFloat16(qt_r) + sA3b_f[sw128_bf16(t, kq, 4 * T_PAD)] = cutlass.BFloat16(kf) + sA3b_f[sw128_bf16(T_PAD + t, kq, 4 * T_PAD)] = cutlass.BFloat16(qf_hi) + # rows 32..47 feed the (unused, must-be-zero) B_lo block of D3 + sA3b_f[sw128_bf16(2 * T_PAD + t, kq, 4 * T_PAD)] = cutlass.BFloat16(0.0) + sB3b_f[sw128_bf16(t, kq, T_PAD)] = cutlass.BFloat16(kb_hi) else: - r_state[row * vec_size + i] = cutlass.Float32( - h0[slot, i_hv, v_row, i * 32 + in_warp_tid] + # rows t, 16+t, 32+t share t % 8 -> the same 16-byte-chunk XOR + xc = kc ^ (4 * (t % 8)) + sB1f[b1_base + 32 * t + xc] = kt_hi + if cutlass.const_expr(S0_LO): + sB1kf[b1k_base + 32 * t + xc] = kt_hi + sB1f[b1_base + 32 * (T_PAD + t) + xc] = qt_r + sB1f[b1_base + 32 * (2 * T_PAD + t) + xc] = kt_lo + sA3f[a3_base + 32 * t + xc] = kf + sA3f[a3_base + 32 * (T_PAD + t) + xc] = qf_hi + sA3f[a3_base + 32 * (2 * T_PAD + t) + xc] = qf_lo + sB3hif[b3_base + 32 * t + xc] = kb_hi + sB3lof[b3_base + 32 * t + xc] = kb_lo + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + if cutlass.const_expr(PROFILE_STAGES): + tp_a = read_globaltimer() + + # ---- operand fragments of GEMM1 (D1 = S0_hi B1^T + S0_lo K~hi^T) and GEMM3 + # (D3 = A3 Kb_hi^T + A3 Kb_lo^T); the MMAs are issued by warp 0 after the Kbar staging + if cutlass.const_expr(BF16_MMA): + tCrB1 = tiled_mma1.make_fragment_B(sB1b) + tCrA3 = tiled_mma3.make_fragment_A(sA3b) + tCrB3hi = tiled_mma3.make_fragment_B(sB3b) + tA1 = tAb + else: + sB1_tf = cute.make_tensor( + cute.recast_ptr(sB1.iterator, smem_layout_b1.inner, dtype=cutlass.TFloat32), + sB1.layout, + ) + tCrB1 = tiled_mma1.make_fragment_B(sB1_tf) + sB1k_tf = cute.make_tensor( + cute.recast_ptr(sB1k.iterator, smem_layout_b1k.inner, dtype=cutlass.TFloat32), + sB1k.layout, + ) + tCrB1k = tiled_mma1k.make_fragment_B(sB1k_tf) + tD1k = cute.make_tensor( + tD1.iterator, + tiled_mma1k.make_fragment_C(tiled_mma1k.partition_shape_C((V, T_PAD))).layout, + ) + tAhi_tf = cute.make_tensor( + cute.recast_ptr(tAhi.iterator, dtype=cutlass.TFloat32), tAhi.layout + ) + tAlo_tf = cute.make_tensor( + cute.recast_ptr(tAlo.iterator, dtype=cutlass.TFloat32), tAlo.layout + ) + sA3_tf = cute.make_tensor( + cute.recast_ptr(sA3.iterator, smem_layout_a3.inner, dtype=cutlass.TFloat32), + sA3.layout, + ) + sB3hi_tf = cute.make_tensor( + cute.recast_ptr(sB3hi.iterator, smem_layout_b3.inner, dtype=cutlass.TFloat32), + sB3hi.layout, + ) + sB3lo_tf = cute.make_tensor( + cute.recast_ptr(sB3lo.iterator, smem_layout_b3.inner, dtype=cutlass.TFloat32), + sB3lo.layout, + ) + tCrA3 = tiled_mma3.make_fragment_A(sA3_tf) + tCrB3hi = tiled_mma3.make_fragment_B(sB3hi_tf) + tCrB3lo = tiled_mma3.make_fragment_B(sB3lo_tf) + tA1 = tAhi_tf + # ---- GEMM2 accumulator seed (D2 = S0 e^{G_c}) and the Kbar tiles first; warp 0 then + # issues GEMM1/GEMM3 while the other warps build the B2 tiles. + # Kbar^T staging: thread (kq, quarter) -> sKbT[kq, s] for its 4 tokens (lanes = k: conflict-free) + for i in cutlass.range(4, unroll_full=True): + s_idx = quarter * 4 + i + b_val = cutlass.Float32(0.0) + if s_idx <= commit_idx: + b_val = sK[s_idx, kq] * cute.math.exp( + sE[commit_idx * K + kq] - sE[s_idx * K + kq], fastmath=True + ) + sKbT[kq * TP_STRIDE + s_idx] = b_val + if warp_idx == 0: + cute.arch.barrier_arrive(barrier_id=5, number_of_threads=NUM_THREADS) + tiled_mma3.set(tcgen05.Field.ACCUMULATE, False) + for kb_i in cutlass.range_constexpr(K // KB_MMA): + cute.gemm( + tiled_mma3, + tD3, + tCrA3[(None, None, kb_i, 0)], + tCrB3hi[(None, None, kb_i, 0)], + tD3, + ) + tiled_mma3.set(tcgen05.Field.ACCUMULATE, True) + if cutlass.const_expr(not BF16_MMA): + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma3, + tD3, + tCrA3[(None, None, kb_i, 0)], + tCrB3lo[(None, None, kb_i, 0)], + tD3, + ) + tiled_mma1.set(tcgen05.Field.ACCUMULATE, False) + for kb_i in cutlass.range_constexpr(K // KB_MMA): + cute.gemm( + tiled_mma1, tD1, tA1[(None, None, kb_i)], tCrB1[(None, None, kb_i, 0)], tD1 + ) + tiled_mma1.set(tcgen05.Field.ACCUMULATE, True) + if cutlass.const_expr(S0_LO): + tiled_mma1k.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma1k, + tD1k, + tAlo_tf[(None, None, kb_i)], + tCrB1k[(None, None, kb_i, 0)], + tD1k, ) - if cutlass.const_expr(USE_SETMAXREG): - cute.arch.warpgroup_reg_alloc(72) + with cute.arch.elect_one(): + tcgen05.commit(mbar1) + else: + cute.arch.barrier(barrier_id=5, number_of_threads=NUM_THREADS) + if cutlass.const_expr(PROFILE_STAGES): + tp_j = read_globaltimer() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 32: + _tq31 = read_globaltimer() + _, _gq31, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq31 + i_n) * 64 + 31] = _tq31 - t_stage0 + # ---- GEMM2 accumulator seed (overlaps the GEMM1/GEMM3 execution) + # D4 (output-combination) columns start from zero: GEMM2 half 0 accumulates into [D4 | ΔS] + st16 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + tiled_stD4 = tcgen05.make_tmem_copy(st16, tD4) + tD4_st = tiled_stD4.get_slice(wg_tidx).partition_D(tD4) + r_zero = cute.make_rmem_tensor( + tiled_stD4.get_slice(wg_tidx).partition_S(tD4).shape, cutlass.Float32 + ) + for i in cutlass.range(cute.size(r_zero), unroll_full=True): + r_zero[i] = cutlass.Float32(0.0) + if wg == 0: + cute.copy(st16, r_zero, tD4_st) + # only this CTA's state rows (TMEM lanes [vh * V/2, +V/2)) are seeded and stored + if cutlass.const_expr(SPLIT_V == 1) or ((warp_idx % 4) // 2 == vh): + for i in cutlass.range(32, unroll_full=True): + kcol = wg * 32 + i + r_hi_flat[i] = r_state[i] * sEc[kcol] + for j in cutlass.range_constexpr(4): + if wg == j: + cute.copy(st32, r_hi, tD2_st[(None, None, None, j)]) + cute.arch.fence_view_async_tmem_store() + # ---- B2 tiles (k rows x s cols, K-major SW64) from sKbT while the MMAs run: + # lanes 0..15 -> row 2w, lanes 16..31 -> row 2w+1, four 32-row rounds + s_col = in_warp_tid % T_PAD + if warp_idx > 0: + for r in cutlass.range(4, unroll_full=True): + k_row = 32 * r + 2 * warp_idx + in_warp_tid // T_PAD + b_val2 = sKbT[k_row * TP_STRIDE + s_col] + b_hi = round_hi(b_val2, BF16_MMA) + if cutlass.const_expr(BF16_MMA): + sB2b_f[sw32_k16_bf16(k_row, s_col)] = cutlass.BFloat16(b_hi) + else: + o_b2 = swz64_k16(k_row, s_col) + sB2hif[o_b2] = b_hi + sB2lof[o_b2] = b_val2 - b_hi + # warp 0 is issuing the MMAs: warps 1..4 cover its rows 32 r + {0, 1} + if warp_idx <= 4: + k_row0 = 32 * (warp_idx - 1) + in_warp_tid // T_PAD + b_val0 = sKbT[k_row0 * TP_STRIDE + s_col] + b_hi0 = round_hi(b_val0, BF16_MMA) + if cutlass.const_expr(BF16_MMA): + sB2b_f[sw32_k16_bf16(k_row0, s_col)] = cutlass.BFloat16(b_hi0) + else: + o_b20 = swz64_k16(k_row0, s_col) + sB2hif[o_b20] = b_hi0 + sB2lof[o_b20] = b_val0 - b_hi0 + + if cutlass.const_expr(PROFILE_STAGES): + tp_k = read_globaltimer() + # ---- wait for GEMM1/GEMM3 + cute.arch.mbarrier_wait(mbar1, 0) + tcgen05_fence_after_thread_sync() + if cutlass.const_expr(PROFILE_STAGES): + tp_b = read_globaltimer() + tiled_ld1 = tcgen05.make_tmem_copy(ld16, tD1) + thr_ld1 = tiled_ld1.get_slice(wg_tidx) + tD1_ld = thr_ld1.partition_S(tD1) + rD1 = cute.make_rmem_tensor(thr_ld1.partition_D(tD1).shape, cutlass.Float32) + cute.copy(ld16, tD1_ld, rD1) + cute.arch.fence_view_async_tmem_load() + rD1_flat = cute.make_tensor(rD1.iterator, cute.make_layout((N1,))) + v_row = wg_tidx + # warpgroup 0: D3 (64 x 16) -> smem [L; B_hi; B_lo; -] + ld3 = cute.make_copy_atom( + tcgen05.Ld16x128bOp(tcgen05.Repetition(T_PAD * 32 // 128)), cutlass.Float32 + ) + tiled_ld3 = tcgen05.make_tmem_copy(ld3, tD3) + thr_ld3 = tiled_ld3.get_slice(wg_tidx) + tD3_ld = thr_ld3.partition_S(tD3) + sLB3 = cute.make_tensor( + sLB.iterator, + cute.make_layout(((4 * T_PAD, T_PAD), 1, 1), stride=((T_PAD, 1), 0, 0)), + ) + tD3_sm = thr_ld3.partition_D(sLB3) + rD3 = cute.make_rmem_tensor(tD3_sm.shape, cutlass.Float32) + if wg == 0: + cute.copy(ld3, tD3_ld, rD3) + cute.arch.fence_view_async_tmem_load() + cute.autovec_copy(rD3, tD3_sm) cute.arch.barrier() - for i_v in range(num_v_tiles): - v_base = i_v * TILE_V - if i_v > 0: - for row in range(NUM_V_ROWS): - v_row = warp_idx * NUM_V_ROWS + row - for i in range(vec_size): - if cutlass.const_expr(USE_FLAT_LAYOUT): - r_state[row * vec_size + i] = cutlass.Float32( - h0[h0_idx, v_base + v_row, i * 32 + in_warp_tid] - ) + if cutlass.const_expr(PROFILE_STAGES): + tp_c = read_globaltimer() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _, _gq32, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq32 + i_n) * 64 + 32] = tp_c - t_stage0 + # ---- warpgroup 0: V' by forward substitution (row v per thread, four accumulators); + # warpgroups 1-2: output-combination operand Bcoef[t, s] = B_hi + B_lo (rows t < T_loop, + # s <= t; zero elsewhere), one element per thread, K-major SW64 16 x 16 tiles (hi / lo) + # that head the GEMM2 half-0 B tile + if wg == 0: + r_Y = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + for s in cutlass.range(T_PAD, unroll_full=True): + y = cutlass.Float32(0.0) + if s < t_max: + if cutlass.const_expr(BF16_MMA): + u0 = rD1_flat[s] + else: + u0 = rD1_flat[s] + rD1_flat[2 * T_PAD + s] + y = sBeta[s] * (sVall[s * V + v_row] - u0) + r_Y[s] = y + r_V = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + for t in cutlass.range(T_PAD, unroll_full=True): + vp = cutlass.Float32(0.0) + if t < t_max: + acc0 = cutlass.Float32(0.0) + acc1 = cutlass.Float32(0.0) + acc2 = cutlass.Float32(0.0) + acc3 = cutlass.Float32(0.0) + for s in range(t): + if s % 4 == 0: + acc0 = acc0 + sLB[t, s] * r_V[s] + elif s % 4 == 1: + acc1 = acc1 + sLB[t, s] * r_V[s] + elif s % 4 == 2: + acc2 = acc2 + sLB[t, s] * r_V[s] else: - r_state[row * vec_size + i] = cutlass.Float32( - h0[slot, i_hv, v_base + v_row, i * 32 + in_warp_tid] - ) - r_v_val = cutlass.Float32(0.0) - i_t = 0 - while i_t < T_loop: - if in_warp_tid < NUM_V_ROWS: - v_idx = v_base + warp_idx * NUM_V_ROWS + in_warp_tid - r_v_val = sVall[i_t * V + v_idx] - r_beta_val = sBeta[i_t] - for i_pair in range(vec_size // 2): - i0 = i_pair * 2 - i1 = i_pair * 2 + 1 - k_idx0 = i0 * 32 + in_warp_tid - k_idx1 = i1 * 32 + in_warp_tid - r_q[i0] = sQ[i_t, k_idx0] - r_q[i1] = sQ[i_t, k_idx1] - _k0 = sK[i_t, k_idx0] - _k1 = sK[i_t, k_idx1] - r_decay[i0] = sG[i_t, k_idx0] - r_decay[i1] = sG[i_t, k_idx1] - r_bk[i0], r_bk[i1] = cute.arch.mul_packed_f32x2( - (r_beta_val, r_beta_val), (_k0, _k1) + acc3 = acc3 + sLB[t, s] * r_V[s] + vp = r_Y[t] - sBeta[t] * ((acc0 + acc1) + (acc2 + acc3)) + r_V[t] = vp + for t in cutlass.range(T_PAD, unroll_full=True): + sVpT[v_row * TP_STRIDE + t] = r_V[t] + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq36 = read_globaltimer() + _, _gq36, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq36 + i_n) * 64 + 36] = _tq36 - t_stage0 + else: + if tidx < 128 + T_PAD * T_PAD: + t4 = (tidx - 128) // T_PAD + s4 = (tidx - 128) % T_PAD + bc = cutlass.Float32(0.0) + if (t4 < T_loop) & (s4 <= t4): + bc = sLB[T_PAD + t4, s4] + sLB[2 * T_PAD + t4, s4] + bc_hi = round_hi(bc, BF16_MMA) + if cutlass.const_expr(BF16_MMA): + sB4b_f[sw32_k16_bf16(t4, s4)] = cutlass.BFloat16(bc_hi) + else: + o_b4 = swz64_k16(t4, s4) + sB4hif[o_b4] = bc_hi + sB4lof[o_b4] = bc - bc_hi + cute.arch.barrier() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq37 = read_globaltimer() + _, _gq37, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq37 + i_n) * 64 + 37] = _tq37 - t_stage0 + # ---- A2 tiles (v rows x s cols) from sVpT, same 16-lanes-per-row mapping; s > c -> 0 + for r in cutlass.range(4, unroll_full=True): + v_row2 = 32 * r + 2 * warp_idx + in_warp_tid // T_PAD + a_val = sVpT[v_row2 * TP_STRIDE + s_col] + a_hi = round_hi(a_val, BF16_MMA) + if cutlass.const_expr(BF16_MMA): + sA2b_f[sw32_k16_bf16(v_row2, s_col)] = cutlass.BFloat16(a_hi) + else: + o_a2 = swz64_k16(v_row2, s_col) + sA2hif[o_a2] = a_hi + sA2lof[o_a2] = a_val - a_hi + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + if cutlass.const_expr(PROFILE_STAGES): + tp_d = read_globaltimer() + if cutlass.const_expr(BF16_MMA): + tCrA2hi = tiled_mma2h.make_fragment_A(sA2b) + tCrA2hi_q = tiled_mma2q.make_fragment_A(sA2b) + sB2b_h1 = cute.make_tensor( + cute.recast_ptr( + sScratchH.iterator + OFFH_B2 + 64 * T_PAD, + smem_layout_b2h.inner, + dtype=cutlass.BFloat16, + ), + smem_layout_b2h.outer, + ) + tCrB2hi_h1 = tiled_mma2h.make_fragment_B(sB2b_h1) + sB2qb = cute.make_tensor( + cute.recast_ptr( + sScratchH.iterator + OFFH_B4, smem_layout_b2q.inner, dtype=cutlass.BFloat16 + ), + smem_layout_b2q.outer, + ) + tCrB2q_hi = tiled_mma2q.make_fragment_B(sB2qb) + else: + sA2hi_tf = cute.make_tensor( + cute.recast_ptr(sA2hi.iterator, smem_layout_a2.inner, dtype=cutlass.TFloat32), + sA2hi.layout, + ) + sA2lo_tf = cute.make_tensor( + cute.recast_ptr(sA2lo.iterator, smem_layout_a2.inner, dtype=cutlass.TFloat32), + sA2lo.layout, + ) + tCrA2hi = tiled_mma2h.make_fragment_A(sA2hi_tf) + tCrA2lo = tiled_mma2h.make_fragment_A(sA2lo_tf) + # B2 halves: k rows 0..63 / 64..127 of the SW64 K-major tiles (a 4 KB shift keeps the swizzle) + sB2hi_h1 = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B2HI + 64 * T_PAD, + smem_layout_b2h.inner, + dtype=cutlass.Float32, + ), + smem_layout_b2h.outer, + ) + sB2lo_h1 = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B2LO + 64 * T_PAD, + smem_layout_b2h.inner, + dtype=cutlass.Float32, + ), + smem_layout_b2h.outer, + ) + tCrB2hi_h1 = tiled_mma2h.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sB2hi_h1.iterator, smem_layout_b2h.inner, dtype=cutlass.TFloat32 + ), + sB2hi_h1.layout, + ) + ) + tCrB2lo_h1 = tiled_mma2h.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sB2lo_h1.iterator, smem_layout_b2h.inner, dtype=cutlass.TFloat32 + ), + sB2lo_h1.layout, + ) + ) + sB2q_hi = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B4HI, smem_layout_b2q.inner, dtype=cutlass.Float32 + ), + smem_layout_b2q.outer, + ) + sB2q_lo = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B4LO, smem_layout_b2q.inner, dtype=cutlass.Float32 + ), + smem_layout_b2q.outer, + ) + tCrA2hi_q = tiled_mma2q.make_fragment_A(sA2hi_tf) + tCrA2lo_q = tiled_mma2q.make_fragment_A(sA2lo_tf) + tCrB2q_hi = tiled_mma2q.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sB2q_hi.iterator, smem_layout_b2q.inner, dtype=cutlass.TFloat32 + ), + sB2q_hi.layout, + ) + ) + tCrB2q_lo = tiled_mma2q.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sB2q_lo.iterator, smem_layout_b2q.inner, dtype=cutlass.TFloat32 + ), + sB2q_lo.layout, + ) + ) + if warp_idx == 0: + # GEMM2 half 0 fused with the output combination: [D4 | ΔS_0..63] += A2 x [Bcoef; Kbar_0..63] + tiled_mma2q.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(T_PAD // KB_MMA): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2hi_q[(None, None, kb_i, 0)], + tCrB2q_hi[(None, None, kb_i, 0)], + tD2q, + ) + if cutlass.const_expr(not BF16_MMA): + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2hi_q[(None, None, kb_i, 0)], + tCrB2q_lo[(None, None, kb_i, 0)], + tD2q, ) - r_k[i0], r_k[i1] = cute.arch.mul_packed_f32x2( - (r_decay[i0], r_decay[i1]), (_k0, _k1) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2lo_q[(None, None, kb_i, 0)], + tCrB2q_hi[(None, None, kb_i, 0)], + tD2q, ) - for row_pair in range(NUM_V_ROWS // 2): - ra = row_pair * 2 - rb = row_pair * 2 + 1 - r_va = cute.arch.shuffle_sync(r_v_val, ra) - r_vb = cute.arch.shuffle_sync(r_v_val, rb) - shk_a1 = 0.0 - shk_a2 = 0.0 - shk_b1 = 0.0 - shk_b2 = 0.0 - for _pi in range(vec_size // 2): - _p = _pi * 2 - shk_a1, shk_a2 = cute.arch.fma_packed_f32x2( - src_a=( - r_state[ra * vec_size + _p], - r_state[ra * vec_size + _p + 1], - ), - src_b=(r_k[_p], r_k[_p + 1]), - src_c=(shk_a1, shk_a2), - ) - shk_b1, shk_b2 = cute.arch.fma_packed_f32x2( - src_a=( - r_state[rb * vec_size + _p], - r_state[rb * vec_size + _p + 1], - ), - src_b=(r_k[_p], r_k[_p + 1]), - src_c=(shk_b1, shk_b2), - ) - shk_a = shk_a1 + shk_a2 - shk_b = shk_b1 + shk_b2 - for offset in [16, 8, 4, 2, 1]: - shk_a += cute.arch.shuffle_sync_bfly( - shk_a, offset=offset, mask=-1, mask_and_clamp=31 - ) - shk_b += cute.arch.shuffle_sync_bfly( - shk_b, offset=offset, mask=-1, mask_and_clamp=31 - ) - vn_a = r_va - shk_a - vn_b = r_vb - shk_b - shq_a1 = 0.0 - shq_a2 = 0.0 - shq_b1 = 0.0 - shq_b2 = 0.0 - for _pi in range(vec_size // 2): - _p = _pi * 2 - vnbk_a0, vnbk_a1 = cute.arch.mul_packed_f32x2( - (vn_a, vn_a), (r_bk[_p], r_bk[_p + 1]) + with cute.arch.elect_one(): + tcgen05.commit(mbar2) + tiled_mma2h.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(T_PAD // KB_MMA): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2hi[(None, None, kb_i, 0)], + tCrB2hi_h1[(None, None, kb_i, 0)], + tD2h1, + ) + if cutlass.const_expr(not BF16_MMA): + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2hi[(None, None, kb_i, 0)], + tCrB2lo_h1[(None, None, kb_i, 0)], + tD2h1, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2lo[(None, None, kb_i, 0)], + tCrB2hi_h1[(None, None, kb_i, 0)], + tD2h1, + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar2b) + if cutlass.const_expr(PROFILE_STAGES): + tp_e = read_globaltimer() + + # ---- outputs of the new tokens: o_t = O0_t + D4[v, t] (columns 0..15 of GEMM2 half 0), + # warpgroup w writes + # new tokens j = w mod 4 + cute.arch.mbarrier_wait(mbar2, 0) + tcgen05_fence_after_thread_sync() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq38 = read_globaltimer() + _, _gq38, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq38 + i_n) * 64 + 38] = _tq38 - t_stage0 + tiled_ld4 = tcgen05.make_tmem_copy(ld16, tD4) + thr_ld4 = tiled_ld4.get_slice(wg_tidx) + rD4 = cute.make_rmem_tensor(thr_ld4.partition_D(tD4).shape, cutlass.Float32) + cute.copy(ld16, thr_ld4.partition_S(tD4), rD4) + cute.arch.fence_view_async_tmem_load() + rD4_flat = cute.make_tensor(rD4.iterator, cute.make_layout((T_PAD,))) + # (REPLAY_ONLY: no new tokens, no outputs) + if cutlass.const_expr(not REPLAY_ONLY): + for j_new in cutlass.range(1 + NUM_SPEC, unroll_full=True): + if wg == j_new % 4: + t_new = commit_len + j_new + o_val = cutlass.Float32(0.0) + for tt in cutlass.range(T_PAD, unroll_full=True): + if tt == t_new: + o_val = rD1_flat[T_PAD + tt] + rD4_flat[tt] + if cutlass.const_expr(FUSE_OUTPUT_NORM): + sVall[j_new * V + v_row] = cutlass.Float32(cutlass.BFloat16(o_val)) + else: + o[0, bos + t_new, i_hv, v_row] = cutlass.BFloat16(o_val) + if cutlass.const_expr(PROFILE_STAGES): + tp_f = read_globaltimer() + # ---- committed-state store through TMA: D2 slab -> swizzled smem box (row per thread, + # 16 B chunks land in distinct banks under the 128B swizzle) -> one bulk tensor store + if cutlass.const_expr(SPLIT_V == 2): + cute.arch.cluster_wait() + if cutlass.const_expr(not FUSE_OUTPUT_NORM): + # bf16 path has no epilogue to overlap with: every warp stores its own slab + # (per-warp 32x32 smem transpose; the TMA store is used on the fused-norm path). + xp_base = warp_idx * 32 * XP_STRIDE + cute.arch.mbarrier_wait(mbar2, 0) + cute.arch.mbarrier_wait(mbar2b, 0) + tcgen05_fence_after_thread_sync() + for j in cutlass.range_constexpr(4): + if wg == j: + cute.copy(ld32, tD2_ld[(None, None, None, j)], rS) + cute.arch.fence_view_async_tmem_load() + cute.arch.barrier() + for i in cutlass.range(32, unroll_full=True): + sXp[xp_base + in_warp_tid * XP_STRIDE + i] = rS_flat[i] + cute.arch.sync_warp() + if cutlass.const_expr(SPLIT_V == 1) or ((warp_idx % 4) // 2 == vh): + for j in cutlass.range(32, unroll_full=True): + val = sXp[xp_base + j * XP_STRIDE + in_warp_tid] + if cutlass.const_expr(USE_FLAT_LAYOUT): + ht[h0_idx, (warp_idx % 4) * 32 + j, wg * 32 + in_warp_tid] = val.to( + STATE_DTYPE ) - vnbk_b0, vnbk_b1 = cute.arch.mul_packed_f32x2( - (vn_b, vn_b), (r_bk[_p], r_bk[_p + 1]) + else: + ht[slot, i_hv, (warp_idx % 4) * 32 + j, wg * 32 + in_warp_tid] = val.to( + STATE_DTYPE ) - r_state[ra * vec_size + _p], r_state[ra * vec_size + _p + 1] = ( - cute.arch.fma_packed_f32x2( - src_a=(r_decay[_p], r_decay[_p + 1]), - src_b=( - r_state[ra * vec_size + _p], - r_state[ra * vec_size + _p + 1], - ), - src_c=(vnbk_a0, vnbk_a1), + if cutlass.const_expr(FUSE_OUTPUT_NORM): + # All 16 warps produced disjoint value rows above. One CTA + # barrier makes the BF16-rounded rows visible before assigning + # one complete token to each warp. + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq39 = read_globaltimer() + _, _gq39, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq39 + i_n) * 64 + 39] = _tq39 - t_stage0 + cute.arch.barrier() + if cutlass.const_expr(PROFILE_STAGES): + tp_g = read_globaltimer() + # ---- committed state: the state-owning warps of warpgroups 0/1 stage k columns 0..63 + # (GEMM2 half 0, mbar2), those of warpgroups 2/3 columns 64..127 (half 1, mbar2b) + # into the swizzled smem boxes; the first staging warp of each pair issues the TMA + # stores once its half is visible (named barriers 6 / 7) and waits for the bulk + # groups at the very end. The warps without state rows run the norm epilogue. + if wg < 2: + cute.arch.mbarrier_wait(mbar2, 0) + else: + cute.arch.mbarrier_wait(mbar2b, 0) + tcgen05_fence_after_thread_sync() + if cutlass.const_expr(PROFILE_STAGES): + tp_l = read_globaltimer() + for sl in cutlass.range_constexpr(4): + if wg == sl: + if cutlass.const_expr(SPLIT_V == 1) or ((warp_idx % 4) // 2 == vh): + cute.copy(ld32, tD2_ld[(None, None, None, sl)], rS) + cute.arch.fence_view_async_tmem_load() + if cutlass.const_expr(STATE_BF16): + # bf16 box: chunks [4 (sl % 2), +4) of box sl // 2, round to nearest + for c in cutlass.range(4, unroll_full=True): + chunk = (4 * (sl % 2) + c) ^ (wg_tidx % 8) + off_b = ( + (sl // 2) * ((V // SPLIT_V) * 64) + + 64 * (wg_tidx - (V // SPLIT_V) * vh) + + 8 * chunk + ) + for i in cutlass.range(8, unroll_full=True): + sBoxh[off_b + i] = cutlass.BFloat16(rS_flat[8 * c + i]) + else: + for c in cutlass.range(8, unroll_full=True): + off_b = ( + sl * ((V // SPLIT_V) * 32) + + 32 * (wg_tidx - (V // SPLIT_V) * vh) + + ((4 * c) ^ vx) + ) + for i in cutlass.range(4, unroll_full=True): + sBoxf[off_b + i] = rS_flat[4 * c + i] + # last TMEM reads of the kernel: warps 1..15 signal barrier 8, warp 0 frees TMEM below + tcgen05_fence_before_thread_sync() + if warp_idx > 0: + cute.arch.barrier_arrive(barrier_id=8, number_of_threads=NUM_THREADS) + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq23 = read_globaltimer() + _, _gq23, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq23 + i_n) * 64 + 23] = _tq23 - t_stage0 + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 256: + _tq33 = read_globaltimer() + _, _gq33, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq33 + i_n) * 64 + 33] = _tq33 - t_stage0 + cute.arch.fence_view_async_shared() + if cutlass.const_expr(SPLIT_V == 1) or ((warp_idx % 4) // 2 == vh): + if wg < 2: + cute.arch.barrier(barrier_id=6, number_of_threads=256 // SPLIT_V) + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq24 = read_globaltimer() + _, _gq24, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq24 + i_n) * 64 + 24] = ( + _tq24 - t_stage0 ) - ) - r_state[rb * vec_size + _p], r_state[rb * vec_size + _p + 1] = ( - cute.arch.fma_packed_f32x2( - src_a=(r_decay[_p], r_decay[_p + 1]), - src_b=( - r_state[rb * vec_size + _p], - r_state[rb * vec_size + _p + 1], - ), - src_c=(vnbk_b0, vnbk_b1), + if tidx == (V // 2) * vh: + for s in cutlass.range_constexpr(NUM_SLABS // 2): + cute.copy(tma_atom_t, tS_list[s], tG_list[s]) + cute.arch.cp_async_bulk_commit_group() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + _tq25 = read_globaltimer() + _, _gq25, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq25 + i_n) * 64 + 25] = ( + _tq25 - t_stage0 ) + else: + cute.arch.barrier(barrier_id=7, number_of_threads=256 // SPLIT_V) + if tidx == 256 + (V // 2) * vh: + for s in cutlass.range_constexpr(NUM_SLABS // 2, NUM_SLABS): + cute.copy(tma_atom_t, tS_list[s], tG_list[s]) + cute.arch.cp_async_bulk_commit_group() + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 256: + _tq34 = read_globaltimer() + _, _gq34, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq34 + i_n) * 64 + 34] = ( + _tq34 - t_stage0 + ) + if warp_idx == 0: + # all TMEM reads are done (barrier 8): free the columns while the store drains + cute.arch.barrier(barrier_id=8, number_of_threads=NUM_THREADS) + tcgen05_fence_after_thread_sync() + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.dealloc_tmem(tmem_ptr, TMEM_COLS) + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 64: + _tq35 = read_globaltimer() + _, _gq35, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq35 + i_n) * 64 + 35] = _tq35 - t_stage0 + # norm epilogue (CTA 0) on the warps that own no state rows: (warp % 4) in {2, 3} + token_idx = (warp_idx // 4) * 2 + (warp_idx % 4 - 2) + if (vh == 0) & ((warp_idx % 4) // 2 == 1): + while token_idx < 1 + NUM_SPEC: + # gated values p_v = o_v * w_v * sigmoid(g_v); the two warp reductions (sum of + # squares, max |p_v|) are interleaved. bf16 rounding is monotonic, so + # amax(bf16(p_v * inv_rms)) = bf16(max|p_v| * inv_rms). + pre_values = cute.make_rmem_tensor( + cute.make_layout((V // 32,), stride=(1,)), cutlass.Float32 + ) + rms_partial = cutlass.Float32(0.0) + pre_max = cutlass.Float32(0.0) + for i in cutlass.range(V // 32, unroll_full=True): + v_idx = i * 32 + in_warp_tid + value = sVall[token_idx * V + v_idx] + gate = sOnG[token_idx * V + v_idx] + sigmoid_gate = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True) ) - shq_a1, shq_a2 = cute.arch.fma_packed_f32x2( - src_a=( - r_state[ra * vec_size + _p], - r_state[ra * vec_size + _p + 1], - ), - src_b=(r_q[_p], r_q[_p + 1]), - src_c=(shq_a1, shq_a2), + p_v = value * sOnW[v_idx] * sigmoid_gate + pre_values[i] = p_v + rms_partial += value * value + pre_max = cute.arch.fmax(pre_max, cute.arch.fmax(p_v, -p_v)) + for offset in [16, 8, 4, 2, 1]: + rms_partial += cute.arch.shuffle_sync_bfly( + rms_partial, offset=offset, mask=-1, mask_and_clamp=31 ) - shq_b1, shq_b2 = cute.arch.fma_packed_f32x2( - src_a=( - r_state[rb * vec_size + _p], - r_state[rb * vec_size + _p + 1], + pre_max = cute.arch.fmax( + pre_max, + cute.arch.shuffle_sync_bfly( + pre_max, offset=offset, mask=-1, mask_and_clamp=31 ), - src_b=(r_q[_p], r_q[_p + 1]), - src_c=(shq_b1, shq_b2), - ) - shq_a = shq_a1 + shq_a2 - shq_b = shq_b1 + shq_b2 - for offset in [16, 8, 4, 2, 1]: - shq_a += cute.arch.shuffle_sync_bfly( - shq_a, offset=offset, mask=-1, mask_and_clamp=31 ) - shq_b += cute.arch.shuffle_sync_bfly( - shq_b, offset=offset, mask=-1, mask_and_clamp=31 + inv_rms = cute.math.rsqrt( + rms_partial / cutlass.Float32(V) + cutlass.Float32(onorm_eps), + fastmath=True, + ) + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 64: + _tq19 = read_globaltimer() + _, _gq19, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq19 + i_n) * 64 + 19] = ( + _tq19 - t_stage0 + ) + output_row = i_n * (1 + NUM_SPEC) + token_idx + normalized_values = cute.make_rmem_tensor( + cute.make_layout((V // 32,), stride=(1,)), cutlass.Float32 + ) + for i in cutlass.range(V // 32, unroll_full=True): + normalized = pre_values[i] * inv_rms + if cutlass.const_expr(QUANTIZE_OUTPUT): + # Preserve the old standalone RMSNorm's BF16 output + # boundary in registers. The value goes straight to + # MXFP8 below; it is never written to or reloaded from + # HBM as BF16. + normalized = cutlass.Float32(cutlass.BFloat16(normalized)) + normalized_values[i] = normalized + if cutlass.const_expr(QUANTIZE_OUTPUT): + output_amax = cutlass.Float32(cutlass.BFloat16(pre_max * inv_rms)) + output_amax = cute.arch.fmax(output_amax, cutlass.Float32(1.0e-10)) + output_sf = (output_amax * cutlass.Float32(1.0 / 448.0)).to( + cutlass.Float8E8M0FNU ) - if in_warp_tid == 0: - v_row_a = warp_idx * NUM_V_ROWS + ra - v_row_b = warp_idx * NUM_V_ROWS + rb - if cutlass.const_expr(USE_ZERO_ACCEPTED): - o[0, bos + i_t, i_hv, v_base + v_row_a] = cutlass.BFloat16(shq_a) - o[0, bos + i_t, i_hv, v_base + v_row_b] = cutlass.BFloat16(shq_b) - elif i_t >= commit_len: - o[0, bos + i_t, i_hv, v_base + v_row_a] = cutlass.BFloat16(shq_a) - o[0, bos + i_t, i_hv, v_base + v_row_b] = cutlass.BFloat16(shq_b) - if i_t == commit_len: - for row in range(NUM_V_ROWS): - v_row = warp_idx * NUM_V_ROWS + row - for i in range(vec_size): - if cutlass.const_expr(USE_FLAT_LAYOUT): - ht[h0_idx, v_base + v_row, i * 32 + in_warp_tid] = r_state[ - row * vec_size + i - ] - else: - ht[slot, i_hv, v_base + v_row, i * 32 + in_warp_tid] = r_state[ - row * vec_size + i - ] - i_t = i_t + 1 - if cutlass.const_expr(REPLAY_ONLY): - for row in range(NUM_V_ROWS): - v_row = warp_idx * NUM_V_ROWS + row - for i in range(vec_size): - if cutlass.const_expr(USE_FLAT_LAYOUT): - ht[h0_idx, v_base + v_row, i * 32 + in_warp_tid] = r_state[ - row * vec_size + i - ] - else: - ht[slot, i_hv, v_base + v_row, i * 32 + in_warp_tid] = r_state[ - row * vec_size + i - ] + quant_scale = cute.arch.rcp_approx(cutlass.Float32(output_sf)) + for i in cutlass.range(V // 32, unroll_full=True): + v_idx = i * 32 + in_warp_tid + o[output_row, i_hv * V + v_idx] = cutlass.Float8E4M3FN( + normalized_values[i] * quant_scale + ) + if in_warp_tid == 0: + scale_bytes = cute.make_tensor( + cute.recast_ptr( + output_scale.iterator, + dtype=cutlass.Float8E8M0FNU, + ), + cute.make_layout(output_scale.shape), + ) + scale_offset = ( + (output_row // 128) * HV * 512 + + i_hv * 512 + + (output_row % 32) * 16 + + ((output_row % 128) // 32) * 4 + ) + for replica in cutlass.range_constexpr(4): + scale_bytes[scale_offset + replica] = output_sf + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 64: + _tq20 = read_globaltimer() + _, _gq20, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq20 + i_n) * 64 + 20] = ( + _tq20 - t_stage0 + ) + else: + for i in cutlass.range(V // 32, unroll_full=True): + v_idx = i * 32 + in_warp_tid + o[0, output_row, i_hv, v_idx] = cutlass.BFloat16(normalized_values[i]) + token_idx = token_idx + NUM_THREADS // 32 + if (tidx == (V // 2) * vh) | (tidx == 256 + (V // 2) * vh): + cute.arch.cp_async_bulk_wait_group(0, read=True) + if cutlass.const_expr(PROFILE_STAGES): + tp_m = read_globaltimer() + if tidx == 0: + _, grid_n2, _ = cute.arch.grid_dim() + tb2 = ((vh * HV + i_hv) * grid_n2 + i_n) * 64 + stage_timing[tb2 + 15] = tp_l - t_stage0 + stage_timing[tb2 + 16] = tp_m - t_stage0 + if tidx == 256: + _, _gq26, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq26 + i_n) * 64 + 26] = tp_l - t_stage0 + if cutlass.const_expr(PROFILE_STAGES): + tp_h = read_globaltimer() + tp_i = tp_h + if tidx == 32: + _, _gq21, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq21 + i_n) * 64 + 21] = tp_h - t_stage0 + if tidx == 224: + _, _gq22, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq22 + i_n) * 64 + 22] = tp_h - t_stage0 if cutlass.const_expr(PROFILE_STAGES): cute.arch.barrier() t_stage2 = read_globaltimer() + if tidx == 256: + _, _gq27, _ = cute.arch.grid_dim() + stage_timing[((vh * HV + i_hv) * _gq27 + i_n) * 64 + 27] = t_stage2 - t_stage0 if tidx == 0: _, grid_n, _ = cute.arch.grid_dim() - timing_base = (i_hv * grid_n + i_n) * 4 + timing_base = ((vh * HV + i_hv) * grid_n + i_n) * 64 stage_timing[timing_base + 0] = t_stage1 - t_stage0 stage_timing[timing_base + 1] = t_stage2 - t_stage1 stage_timing[timing_base + 2] = t_stage2 - t_stage0 stage_timing[timing_base + 3] = t_stage0 + stage_timing[timing_base + 4] = tp_a - t_stage0 + stage_timing[timing_base + 5] = tp_b - t_stage0 + stage_timing[timing_base + 6] = tp_c - t_stage0 + stage_timing[timing_base + 7] = tp_d - t_stage0 + stage_timing[timing_base + 8] = tp_e - t_stage0 + stage_timing[timing_base + 9] = tp_f - t_stage0 + stage_timing[timing_base + 10] = tp_g - t_stage0 + stage_timing[timing_base + 11] = tp_h - t_stage0 + stage_timing[timing_base + 12] = tp_i - t_stage0 + stage_timing[timing_base + 13] = tp_j - t_stage0 + stage_timing[timing_base + 14] = tp_k - t_stage0 + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + if cutlass.const_expr(not FUSE_OUTPUT_NORM): + # (the fused-norm path released TMEM right after the state staging) + if warp_idx == 0: + cute.arch.relinquish_tmem_alloc_permit() + tmem_ptr_d = cute.arch.retrieve_tmem_ptr(cutlass.Float32, 16, tmem_holder) + cute.arch.dealloc_tmem(tmem_ptr_d, TMEM_COLS) + if cutlass.const_expr(PROFILE_STAGES): + if tidx == 0: + t_exit = read_globaltimer() + _, grid_n3, _ = cute.arch.grid_dim() + tb3 = ((vh * HV + i_hv) * grid_n3 + i_n) * 64 + stage_timing[tb3 + 17] = t_exit - t_stage0 diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_scheduler.py b/tensorrt_llm/_torch/cute_dsl_kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_scheduler.py index 275247681c18..75ef3eccaabd 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_scheduler.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_scheduler.py @@ -670,8 +670,6 @@ def _wait_for_fc1_prologue_claims( ) ) spin_wait(fc1_counter_pointer, lambda value: value >= claim_target) - if stream_ends_before_target: - fc1_exhausted = Boolean(True) return fc1_exhausted @cute.jit diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_fused_fc12.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_fused_fc12.py index 7681fbde79bc..66f43e28ed48 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_fused_fc12.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_fused_fc12.py @@ -29,6 +29,8 @@ from dataclasses import dataclass from typing import NamedTuple, Optional, Tuple, Type, Union +from ....utils import ActivationType + import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute @@ -1076,6 +1078,22 @@ def silu_f32( return a * sigmoid_f32(a, fastmath=fastmath) +_FC12_ACTIVATION_TYPES = ( + ActivationType.Swiglu, + ActivationType.SiTu, +) + + +def validate_fc12_activation_type(activation_type) -> ActivationType: + """Fused FC12 is gated-only: SwiGLU or SiTU. Relu2 would keep full N.""" + activation_type = ActivationType(int(activation_type)) + if activation_type not in _FC12_ACTIVATION_TYPES: + raise ValueError( + f"Fused FC12 supports SwiGLU and SiTU only, got {activation_type}" + ) + return activation_type + + class S2TCopyBundle(NamedTuple): """Bundle of tiled copy and partitioned tensors for smem-to-tmem copies.""" @@ -1127,13 +1145,14 @@ def sm100_tcgen05_st_32x32b_x4( Compute: fc1_acc = fc1_alpha * (SFA * A[token_ids]) * (SFB * B) - C = up * silu(gate) # SwiGLU on interleaved acc + C = gated act on gate16/up16 acc (SwiGLU or SiTU; both halve N) fc2_acc = fc2_alpha * (SFC * C) * (FC2_SFB * FC2_B) + optional NVFP4 quantization (generates SFC) when c_dtype == Float4E2M1FN. Shapes: A is M×K×1; B is N×K×L (L = num experts), with FC1 N ordered as [gate16, up16] groups. B and SFB must use the same N order; -C is M×(N/2)×1 in logical channel order. SFA/SFB layouts follow +C is M×(N/2)×1 in logical channel order (N halved by the gated epilogue). +SFA/SFB layouts follow BlockScaledBasicChunk. ``permuted_idx_to_expanded_idx`` is shared by both phases: FC1 divides by ``topk`` to recover the receive row, while FC2 also uses the remainder to select the route scale. @@ -1169,8 +1188,9 @@ class Sm107BlockScaledContiguousGroupedGemmFusedFc12Kernel: - A/SFA load: four gather warps issue CpAsync128.CG into a shared fc1_a_pipeline. SFA is consumed directly by a Cp128x128b UTCCP in the MMA warp using the 128dp_Unique layout. - - SwiGLU epilogue: C = up * silu(gate), where up/gate come from the - gate16/up16 FC1 weight layout → output N is halved. + - Gated epilogue: SwiGLU (``up * silu(gate)``) or SiTU (Kimi K3); + up/gate come from the gate16/up16 FC1 weight layout → output N is + halved. Selected by ``activation_type`` at trace time. - Optional NVFP4 quant: when c_dtype == Float4E2M1FN, the epilogue also generates SFC and quantizes the output. @@ -1189,6 +1209,11 @@ class Sm107BlockScaledContiguousGroupedGemmFusedFc12Kernel: :param vectorized_f32: Use vectorized f32x2 ops in epilogue. :param topk: Experts selected per token. :param swiglu_limit: DS-v4 SwiGLU clamp limit; ``+inf`` disables clamp. + Mutually exclusive with ``ActivationType.SiTu``. + :param activation_type: ``Swiglu`` (default) or ``SiTu``. Folded at trace + time; a different value compiles a different kernel. + :param situ_beta: Gate-side SiTU constant; required for ``SiTu`` only. + :param situ_linear_beta: Linear-side SiTU constant; required for ``SiTu`` only. :param scheduler: Persistent work-ID scheduler: ``static`` or ``l2_atomic``. FC1 B and SFB must already use gate16/up16 order (tile N 128/256). FC1 output remains in logical channel order; FC2 weight layout is unchanged. @@ -1205,6 +1230,9 @@ def __init__( use_pdl: bool = True, swiglu_limit: cutlass.Float32 = float("inf"), scheduler: str = "static", + activation_type: ActivationType = ActivationType.Swiglu, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, ): self.sf_vec_size = sf_vec_size self.topk = topk @@ -1226,6 +1254,31 @@ def __init__( swiglu_limit = float("inf") self.swiglu_limit = swiglu_limit self.has_swiglu_limit = swiglu_limit != float("inf") + self.activation_type = validate_fc12_activation_type(activation_type) + if self.activation_type == ActivationType.SiTu: + if situ_beta is None or situ_linear_beta is None: + raise ValueError( + "ActivationType.SiTu requires both situ_beta and " + f"situ_linear_beta, got {situ_beta} and {situ_linear_beta}." + ) + if situ_beta <= 0 or situ_linear_beta <= 0: + raise ValueError( + f"SiTU betas must be positive, got {situ_beta} and {situ_linear_beta}." + ) + if self.has_swiglu_limit: + raise ValueError( + "Fused FC12 SiTU does not support a SwiGLU clamp; " + "drop swiglu_limit for SiTU checkpoints." + ) + elif situ_beta is not None or situ_linear_beta is not None: + raise ValueError( + "situ_beta / situ_linear_beta require " + f"ActivationType.SiTu, got {self.activation_type.name}." + ) + self.situ_beta = None if situ_beta is None else float(situ_beta) + self.situ_linear_beta = ( + None if situ_linear_beta is None else float(situ_linear_beta) + ) self.use_2cta_instrs = mma_inst_shape[0] == 256 self.cta_group = ( @@ -1968,7 +2021,7 @@ def __call__( This method performs FC1 layer computation: 1. GEMM: acc = fc1_alpha * (SFA * A[token_ids]) * (SFB * B) - 2. SwiGLU: C = up * silu(gate), decoded from gate16/up16 accumulator groups + 2. Gated act: SwiGLU or SiTU, decoded from gate16/up16 accumulator groups (N/2) 3. Optional Quant: When c_dtype is Float4E2M1FN, generates SFC and quantizes output Data loading: @@ -1988,7 +2041,7 @@ def __call__( shared memory when use_2cta_instrs is True - TMA warp: Load B and SFB with multicast - MMA warp: Perform matrix multiply-accumulate - - Epilogue warps: Apply SwiGLU activation, optional quantization, and store results + - Epilogue warps: Apply gated activation, optional quantization, and store results :param fc1_a: FC1 input A (MxKx1), gathered with the shared route mapping :type fc1_a: cute.Tensor @@ -2615,6 +2668,173 @@ class SharedStorageCpasync2cta: ) return + @cute.jit + def _apply_swiglu_epilogue( + self, + acc_vec_up: cute.Tensor, + acc_vec_gate: cute.Tensor, + alpha_val, + tCompute: cute.Tensor, + ): + """SwiGLU: ``tCompute[i] = (alpha * up[i]) * silu(alpha * gate[i])``. + + Optional DS-v4 clamp via ``has_swiglu_limit`` is folded at trace time. + """ + if cutlass.const_expr(self.vectorized_f32): + LOG2_E = cutlass.Float32(1.4426950408889634) + for i in cutlass.range_constexpr(0, cute.size(acc_vec_up.shape), 2): + acc_vec_up_alpha = cute.arch.mul_packed_f32x2( + (acc_vec_up[i], acc_vec_up[i + 1]), + (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)), + ) + acc_vec_gate_alpha = cute.arch.mul_packed_f32x2( + (acc_vec_gate[i], acc_vec_gate[i + 1]), + (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)), + ) + if cutlass.const_expr(self.has_swiglu_limit): + acc_vec_gate_alpha = ( + fmin(acc_vec_gate_alpha[0], self.swiglu_limit), + fmin(acc_vec_gate_alpha[1], self.swiglu_limit), + ) + acc_vec_up_alpha = ( + fclip_xorsign(acc_vec_up_alpha[0], self.swiglu_limit), + fclip_xorsign(acc_vec_up_alpha[1], self.swiglu_limit), + ) + tCompute_log2e = cute.arch.mul_packed_f32x2( + (acc_vec_gate_alpha[0], acc_vec_gate_alpha[1]), + (-LOG2_E, -LOG2_E), + ) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.add_packed_f32x2( + ( + cute.math.exp2(tCompute_log2e[0], fastmath=True), + cute.math.exp2(tCompute_log2e[1], fastmath=True), + ), + (1.0, 1.0), + ) + tCompute[i] = cute.arch.rcp_approx(tCompute[i]) + tCompute[i + 1] = cute.arch.rcp_approx(tCompute[i + 1]) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (acc_vec_gate_alpha[0], acc_vec_gate_alpha[1]), + ) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (acc_vec_up_alpha[0], acc_vec_up_alpha[1]), + ) + else: + for i in cutlass.range_constexpr(cute.size(acc_vec_up.shape)): + acc_vec_up_alpha = acc_vec_up[i] * cutlass.Float32(alpha_val) + acc_vec_gate_alpha = acc_vec_gate[i] * cutlass.Float32(alpha_val) + if cutlass.const_expr(self.has_swiglu_limit): + acc_vec_gate_alpha = fmin(acc_vec_gate_alpha, self.swiglu_limit) + acc_vec_up_alpha = fclip_xorsign( + acc_vec_up_alpha, self.swiglu_limit + ) + tCompute[i] = acc_vec_up_alpha * silu_f32( + acc_vec_gate_alpha, fastmath=True + ) + + @cute.jit + def _apply_situ_epilogue( + self, + acc_vec_up: cute.Tensor, + acc_vec_gate: cute.Tensor, + alpha_val, + tCompute: cute.Tensor, + ): + """SiTU (Kimi K3), matching ``SituAndMul`` / the two-op Rubin kernel:: + + g = alpha * gate, u = alpha * up + situ_gate = beta * tanh(g / beta) * sigmoid(g) + situ_up = linear_beta * tanh(u / linear_beta) + tCompute = situ_gate * situ_up + + Packed path uses ``tanh(z) = 2 * sigmoid(2z) - 1``. Betas fold at + trace time; a different pair compiles a different kernel. + """ + beta = self.situ_beta + linear_beta = self.situ_linear_beta + if cutlass.const_expr(self.vectorized_f32): + LOG2_E = cutlass.Float32(1.4426950408889634) + neg_log2e_pair = (-LOG2_E, -LOG2_E) + one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) + + inv_2beta = cutlass.Float32(2.0 / beta) + two_beta = cutlass.Float32(2.0 * beta) + neg_beta = cutlass.Float32(-beta) + inv_2lbeta = cutlass.Float32(2.0 / linear_beta) + two_lbeta = cutlass.Float32(2.0 * linear_beta) + neg_lbeta = cutlass.Float32(-linear_beta) + + def _sigmoid(p0, p1): + neg = cute.arch.mul_packed_f32x2((p0, p1), neg_log2e_pair) + e = ( + cute.math.exp2(neg[0], fastmath=True), + cute.math.exp2(neg[1], fastmath=True), + ) + d = cute.arch.add_packed_f32x2(e, one_pair) + return (cute.arch.rcp_approx(d[0]), cute.arch.rcp_approx(d[1])) + + alpha_pair = (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)) + for i in cutlass.range_constexpr(0, cute.size(acc_vec_up.shape), 2): + g = cute.arch.mul_packed_f32x2( + (acc_vec_gate[i], acc_vec_gate[i + 1]), alpha_pair + ) + u = cute.arch.mul_packed_f32x2( + (acc_vec_up[i], acc_vec_up[i + 1]), alpha_pair + ) + + sigmoid_g = _sigmoid(g[0], g[1]) + + gs = _sigmoid( + *cute.arch.mul_packed_f32x2(g, (inv_2beta, inv_2beta)) + ) + tanh_g = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(gs, (two_beta, two_beta)), + (neg_beta, neg_beta), + ) + + us = _sigmoid( + *cute.arch.mul_packed_f32x2(u, (inv_2lbeta, inv_2lbeta)) + ) + tanh_u = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(us, (two_lbeta, two_lbeta)), + (neg_lbeta, neg_lbeta), + ) + + situ_gate = cute.arch.mul_packed_f32x2(tanh_g, sigmoid_g) + out_pair = cute.arch.mul_packed_f32x2(situ_gate, tanh_u) + tCompute[i] = out_pair[0] + tCompute[i + 1] = out_pair[1] + else: + inv_2beta = cutlass.Float32(2.0 / beta) + two_beta = cutlass.Float32(2.0 * beta) + beta_f32 = cutlass.Float32(beta) + inv_2lbeta = cutlass.Float32(2.0 / linear_beta) + two_lbeta = cutlass.Float32(2.0 * linear_beta) + lbeta_f32 = cutlass.Float32(linear_beta) + for i in cutlass.range_constexpr(cute.size(acc_vec_up.shape)): + g = acc_vec_gate[i] * cutlass.Float32(alpha_val) + u = acc_vec_up[i] * cutlass.Float32(alpha_val) + tanh_g = ( + two_beta * sigmoid_f32(g * inv_2beta, fastmath=True) + - beta_f32 + ) + tanh_u = ( + two_lbeta * sigmoid_f32(u * inv_2lbeta, fastmath=True) + - lbeta_f32 + ) + tCompute[i] = (tanh_g * sigmoid_f32(g, fastmath=True)) * tanh_u + # GPU device kernel @cute.kernel def kernel( @@ -4792,7 +5012,8 @@ def kernel( acc_pipeline.consumer_wait(acc_consumer_state) # Each N64 output subtile consumes two consecutive N64 - # accumulator chunks, each [gate16, up16, gate16, up16]. + # accumulator chunks, each [gate16, up16, gate16, up16]; + # the decoded up/gate feed SwiGLU or SiTU. # tFC1TR_tAcc: (T2R, T2R_M, T2R_N, EPI_M, EPI_N, # STAGE), sliced on STAGE. bSG_gC: ((ATOM_V, REST_V), # EPI_M, EPI_N, loopM, loopN, loopL). @@ -4841,103 +5062,22 @@ def kernel( acc_vec_up = tFC1TR_rAcc_up.load() acc_vec_gate = tFC1TR_rAcc_gate.load() - # SwiGLU: output = up * silu(gate), silu(x) = x * sigmoid(x). + # Gated epilogue on the decoded up/gate; C has N/2 + # cols. SwiGLU vs SiTU is a const_expr specialization + # (trace-time). tCompute = cute.make_rmem_tensor( acc_vec_gate.shape, self.acc_dtype ) - if cutlass.const_expr(self.vectorized_f32): - # SwiGLU Packed Version: uses f32x2 packed operations for better performance - # Computes: output = (alpha * up) * silu(alpha * gate) - # where silu(x) = x * sigmoid(x) = x / (1 + exp(-x)) - LOG2_E = cutlass.Float32(1.4426950408889634) - for i in cutlass.range_constexpr( - 0, cute.size(tFC1TR_rAcc_up), 2 - ): - acc_vec_up_alpha = cute.arch.mul_packed_f32x2( - (acc_vec_up[i], acc_vec_up[i + 1]), - ( - cutlass.Float32(alpha_val), - cutlass.Float32(alpha_val), - ), - ) - acc_vec_gate_alpha = cute.arch.mul_packed_f32x2( - (acc_vec_gate[i], acc_vec_gate[i + 1]), - ( - cutlass.Float32(alpha_val), - cutlass.Float32(alpha_val), - ), - ) - if cutlass.const_expr(self.has_swiglu_limit): - acc_vec_gate_alpha = ( - fmin( - acc_vec_gate_alpha[0], self.swiglu_limit - ), - fmin( - acc_vec_gate_alpha[1], self.swiglu_limit - ), - ) - acc_vec_up_alpha = ( - fclip_xorsign( - acc_vec_up_alpha[0], self.swiglu_limit - ), - fclip_xorsign( - acc_vec_up_alpha[1], self.swiglu_limit - ), - ) - tCompute_log2e = cute.arch.mul_packed_f32x2( - (acc_vec_gate_alpha[0], acc_vec_gate_alpha[1]), - (-LOG2_E, -LOG2_E), - ) - ( - tCompute[i], - tCompute[i + 1], - ) = cute.arch.add_packed_f32x2( - ( - cute.math.exp2( - tCompute_log2e[0], fastmath=True - ), - cute.math.exp2( - tCompute_log2e[1], fastmath=True - ), - ), - (1.0, 1.0), - ) - tCompute[i] = cute.arch.rcp_approx(tCompute[i]) - tCompute[i + 1] = cute.arch.rcp_approx(tCompute[i + 1]) - ( - tCompute[i], - tCompute[i + 1], - ) = cute.arch.mul_packed_f32x2( - (tCompute[i], tCompute[i + 1]), - (acc_vec_gate_alpha[0], acc_vec_gate_alpha[1]), - ) - ( - tCompute[i], - tCompute[i + 1], - ) = cute.arch.mul_packed_f32x2( - (tCompute[i], tCompute[i + 1]), - (acc_vec_up_alpha[0], acc_vec_up_alpha[1]), - ) + if cutlass.const_expr( + self.activation_type == ActivationType.SiTu + ): + self._apply_situ_epilogue( + acc_vec_up, acc_vec_gate, alpha_val, tCompute + ) else: - # SwiGLU Unpacked Version: scalar operations - # Computes: output = (alpha * up) * silu(alpha * gate) - for i in cutlass.range_constexpr(cute.size(tFC1TR_rAcc_up)): - acc_vec_up_alpha = acc_vec_up[i] * cutlass.Float32( - alpha_val - ) - acc_vec_gate_alpha = acc_vec_gate[i] * cutlass.Float32( - alpha_val - ) - if cutlass.const_expr(self.has_swiglu_limit): - acc_vec_gate_alpha = fmin( - acc_vec_gate_alpha, self.swiglu_limit - ) - acc_vec_up_alpha = fclip_xorsign( - acc_vec_up_alpha, self.swiglu_limit - ) - tCompute[i] = acc_vec_up_alpha * silu_f32( - acc_vec_gate_alpha, fastmath=True - ) + self._apply_swiglu_epilogue( + acc_vec_up, acc_vec_gate, alpha_val, tCompute + ) if cutlass.const_expr(self.fc1_generate_sfc): # Float4E2M1FN quantization: per-vector absmax → diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 87506e0070b3..bee09228302b 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -797,6 +797,7 @@ class MNNVLAllReduce(nn.Module): threading.Lock] = threading.Lock() SUPPORTED_FUSION_OPS: frozenset[AllReduceFusionOp] = frozenset({ + AllReduceFusionOp.RMS_NORM, AllReduceFusionOp.RESIDUAL_RMS_NORM, AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8, AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4, @@ -996,7 +997,11 @@ def forward( all_reduce_params.scale, # scale int(fusion_op), ) - return tuple(outputs) if is_fusion else outputs[0] + # RMS_NORM returns norm_out alone; match AllReduce.forward's + # "unwrap a single-element result" contract so callers see a tensor. + if is_fusion and len(outputs) > 1: + return tuple(outputs) + return outputs[0] class AllReduce(nn.Module): diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index ccb4609d09c8..ef6f20365727 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2844,10 +2844,13 @@ class MLADSparkForCausalLM(_DSparkHeadMixin, DFlashForCausalLM): # head, head_dim 576) would fail both. The backend NAME still selects, via # _mla_block_decode_variant below. _uses_worker_attention_backend = False - # Selectable but not the default: needs a cute_dsl_mla_decode_fp16_blackwell - # taking per-token kv_bounds, absent on this branch. Worth 35.7% off the - # drafter region (947.50 -> 608.95 us/iter, jobs 3006025/3006026). - _default_attention_backend = "TRTLLM" + # CUTEDSL by default: one cute-dsl pass over context and block through the + # per-token kv_bounds of cute_dsl_mla_decode_fp16_blackwell, with no + # upper-triangle fixup -- 35.7% off the drafter region (947.50 -> 608.95 + # us/iter) at unchanged acceptance and accuracy. It does not degrade (see + # _auto_fallback_attention_backend): a build without the kernel raises in + # __init__ with the reason from cute_dsl_mla_decode_unavailability_reason. + _default_attention_backend = "CUTEDSL" _supported_attention_backends = ("VANILLA", "TRTLLM", "CUTEDSL") # No AUTO degrade. This drafter needs a 288 GB Blackwell part to hold the # target at all, so a build that cannot run its kernels cannot run the @@ -3033,8 +3036,7 @@ def _assert_cute_dsl_can_implement(self, batch, block_size, page_size, dtype) -> f"attention_backend='CUTEDSL' cannot serve this shape: batch={batch}, " f"block_size={block_size}, page_size={page_size}, " f"num_heads={self._num_heads}, dtype={dtype}. Use " - f"attention_backend='TRTLLM', or set kv_cache_config.tokens_per_block " - f"to a divisor of 128 other than 1." + f"attention_backend='TRTLLM'." ) self._cute_dsl_checked_shape = key diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 07d232f62afe..3806e7148029 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -106,19 +106,19 @@ from safetensors import safe_open from torch import nn -from ..._utils import is_sm_100f +from ..._utils import get_sm_version, is_sm_100f from ...logger import logger from ...mapping import Mapping from ...models.modeling_utils import QuantAlgo, QuantConfig from ..attention.backends import AttentionMetadata -from ..distributed import AllReduce, AllReduceParams +from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams from ..model_config import ModelConfig from ..modules.gated_mlp import GatedMLP from ..modules.kimi_kda import KimiKDALinearAttention from ..modules.kimi_kda.kimi_k3_mamba_metadata import KimiK3MambaMetadata from ..modules.linear import Linear as TrtllmLinear from ..modules.linear import TensorParallelMode, load_weight_shard -from ..modules.multi_stream_utils import maybe_execute_in_parallel +from ..modules.multi_stream_utils import do_multi_stream, maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm from ..modules.situ import SituAndMul from ..moe.fused_moe import ConfigurableMoE, SiTuActivation, TRTLLMGenFusedMoE, create_moe @@ -804,12 +804,27 @@ def _helix_cp_v_b_shard( # --------------------------------------------------------------------------- +def _rubin_mxfp8_is_available() -> bool: + """Whether this rank can run the SM107 native block-scaled MXFP8 GEMM. + + Recomputed rather than cached: the SM version cannot change within a run, + so weight preparation and forward always agree. + """ + from ..cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE + + return get_sm_version() == 107 and IS_CUTLASS_DSL_RUBIN_AVAILABLE + + class _Fp8BlockScaleWeightReadLinear(nn.Module): """Bias-free FP8 block-scale GEMM for checkpoint attention and converted MLPs. Attention loads checkpoint E4M3 codes and 128x128 scales directly. The optional - MLP path quantizes BF16 weights once at load. Both prepare UE8M0 scales - for DeepGEMM and quantize BF16 activations inside the GEMM. + MLP path quantizes BF16 weights once at load. On SM100/SM103 the UE8M0 + scales are packed for DeepGEMM, which quantizes BF16 activations inside the + GEMM. On Rubin (SM107) the scales are packed into the CuTe DSL K32 R128c4 + layout and the activation is quantized by the specialized CuTe quantizer + before the native MXFP8 GEMM, which runs with K3's fine-grained low-M + tuning buckets. """ def __init__( @@ -836,9 +851,14 @@ def has_w4a8_nvfp4_fp8(self) -> bool: """Match the ``Linear`` interface consumed by ``GatedMLP``.""" return False + @property + def supports_prequantized_input(self) -> bool: + """Whether a producer may hand this GEMM an FP8 activation in CuTe layout.""" + return _rubin_mxfp8_is_available() + @staticmethod def quantize_weight(weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - """BF16 ``[out, in]`` weight -> (FP8 weight, deep_gemm-ready scale). + """BF16 ``[out, in]`` weight -> (FP8 weight, backend-ready scale). Both dims must be multiples of 128. Because the 128x128 block scale is computed per block, concatenating several such weights along ``out`` @@ -858,14 +878,22 @@ def _prepare_weights( ) -> Tuple[torch.Tensor, torch.Tensor]: from ...quantization.utils.fp8_utils import ( resmooth_to_fp8_e8m0, + transform_k128_scales_to_cutedsl_mxfp8_layout, transform_sf_into_required_layout, ) - # fp8_swap_ab_gemm with disable_ue8m0_cast=True consumes packed, - # TMA-aligned UE8M0 scales rather than the checkpoint's FP32 grid. + # Both GEMMs consume UE8M0 scales rather than the checkpoint's FP32 grid. weight_fp8, weight_scale = resmooth_to_fp8_e8m0( weight_fp8.contiguous(), weight_scale.contiguous().float() ) + if _rubin_mxfp8_is_available(): + # cute_dsl_mxfp8_gemm_rubin reads K32 R128c4 UE8M0 scales. + weight_scale = transform_k128_scales_to_cutedsl_mxfp8_layout( + weight_scale, mn=weight_fp8.shape[0], k=weight_fp8.shape[1] + ) + return weight_fp8, weight_scale + # fp8_swap_ab_gemm with disable_ue8m0_cast=True consumes packed, + # TMA-aligned UE8M0 scales. weight_scale = transform_sf_into_required_layout( weight_scale, mn=weight_fp8.shape[0], @@ -911,14 +939,54 @@ def forward( layer_idx: Optional[int] = None, ) -> torch.Tensor: out_shape = (*x.shape[:-1], self.out_features) - out = torch.ops.trtllm.fp8_swap_ab_gemm( - x.reshape(-1, x.shape[-1]), + x_flat = x.reshape(-1, x.shape[-1]) + if _rubin_mxfp8_is_available(): + act_fp8, act_scale = torch.ops.trtllm.fp8_quantize_1x128_cutedsl_ue8m0(x_flat) + out = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin( + act_fp8, + self.weight, + act_scale, + self.weight_scale, + output_dtype=torch.bfloat16, + fine_grained_m=True, + ) + else: + out = torch.ops.trtllm.fp8_swap_ab_gemm( + x_flat, + self.weight, + self.weight_scale, + output_dtype=x.dtype, + disable_ue8m0_cast=True, + ) + return out.reshape(out_shape) + + def forward_prequantized( + self, + activation: torch.Tensor, + activation_scale: torch.Tensor, + ) -> torch.Tensor: + """Project an FP8 activation a fused producer already quantized. + + Args: + activation: ``[..., in_features]`` ``float8_e4m3fn`` activation. + activation_scale: ``uint8`` UE8M0 scales in the CuTe K32 R128c4 + layout that ``fp8_quantize_1x128_cutedsl_ue8m0`` produces. + """ + if activation.dtype is not torch.float8_e4m3fn: + raise TypeError("prequantized Kimi K3 activation must be float8_e4m3fn") + if not self.supports_prequantized_input: + raise RuntimeError("prequantized Kimi K3 activation requires the Rubin CuTe GEMM") + if activation_scale.dtype is not torch.uint8: + raise TypeError("prequantized Kimi K3 activation scale must be uint8 R128c4") + out = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin( + activation.reshape(-1, activation.shape[-1]), self.weight, + activation_scale, self.weight_scale, - output_dtype=x.dtype, - disable_ue8m0_cast=True, + output_dtype=torch.bfloat16, + fine_grained_m=True, ) - return out.reshape(out_shape) + return out.reshape(activation.shape[:-1] + (self.out_features,)) def _swap_linear_to_fp8_weight_read( @@ -1192,7 +1260,7 @@ def __init__( # CUTLASS is absent on purpose: it is the fallback target, so # "degraded to CUTLASS" is not a thing that can happen to it. allow_backend_degradation=routed_moe_model_config.moe_backend - not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", "CUTEDSL"), + not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", "CUTEDSL", "CUTEDSL_FC12"), ) self._check_trtllm_situ_quant( routed_moe_model_config.moe_backend, routed_quant_config.quant_algo @@ -1230,6 +1298,16 @@ def __init__( # Direct MoE-TP leaves both branches as partials for one concatenated # all-reduce. use_shared_tp = not attention_dp and model_config.mapping.tp_size > 1 + # Fused-comm MegaMoE runs its own cross-rank exchange inside the + # kernel. Order the shared-expert all-reduce after the routed MoE + # instead of running the two collectives concurrently (startup + # deadlocks). CuteDSL is the validated scope, not a correctness + # dependency: DeepGEMM rejects TP without attention DP. + self._defer_shared_allreduce = ( + use_shared_tp + and routed_moe_model_config.moe_backend == "MEGAMOE_CUTEDSL" + and self.routed_experts.backend.scheduler_kind == MoESchedulerKind.FUSED_COMM + ) self._reduce_routed_output = ( use_shared_tp and self.routed_experts.backend.scheduler_kind != MoESchedulerKind.FUSED_COMM @@ -1262,6 +1340,17 @@ def __init__( self.shared_expert_stream = aux_stream_dict[AuxStreamType.MoeShared] self.moe_main_event = torch.cuda.Event() self.moe_shared_event = torch.cuda.Event() + self.moe_routed_event = torch.cuda.Event() if self._defer_shared_allreduce else None + # CuteDSL already quantizes the MoE input on the MoeOutputMemset stream; + # the latent down-projection runs there too, beside the routing chain. + self.latent_down_stream = ( + aux_stream_dict.get(AuxStreamType.MoeOutputMemset) + if routed_moe_model_config.moe_backend in ("CUTEDSL", "CUTEDSL_FC12") + else None + ) + has_latent_stream = self.latent_down_stream is not None + self.moe_latent_fork_event = torch.cuda.Event() if has_latent_stream else None + self.moe_latent_ready_event = torch.cuda.Event() if has_latent_stream else None self.routed_expert_down_proj = nn.Linear( cfg.hidden_size, self.moe_hidden_size, bias=False, dtype=dtype ) @@ -1282,6 +1371,45 @@ def _routed_projection(hidden_states: torch.Tensor, projection: nn.Module) -> to hidden_states, projection.weight.t(), None, None ) + def _can_fuse_router_down(self, hidden_states: torch.Tensor) -> bool: + """Whether one fused GEMM can produce the router logits and the latent + down-projection (trtllm::dsv3_router_latent_gemm_op). + + Scoped to the validated configuration (TP MEGAMOE_CUTEDSL with fused + communication on SM107); the fusion itself does not depend on it. + """ + if ( + torch.compiler.is_compiling() + or _K3_DISABLE_MIN_LATENCY_LATENT_PROJ + or not hidden_states.is_cuda + or self.routed_experts.use_dp + or self.routed_experts.parallel_size <= 1 + or self.routed_experts.moe_backend != "MEGAMOE_CUTEDSL" + or self.routed_experts.backend.scheduler_kind != MoESchedulerKind.FUSED_COMM + or get_sm_version() != 107 + ): + return False + down = self.routed_expert_down_proj + if not isinstance(down, nn.Linear) or down.bias is not None: + return False + if ( + hidden_states.ndim != 2 + or not 1 <= hidden_states.shape[0] <= 16 + or hidden_states.shape[1] != 7168 + or self.gate.weight.shape != (896, 7168) + or down.weight.shape != (3584, 7168) + ): + return False + # Inspect the loaded weights: FP8 conversion can replace the down module. + return all( + tensor.dtype == torch.bfloat16 + and tensor.is_cuda + and tensor.device == hidden_states.device + and tensor.stride() == (7168, 1) + and tensor.data_ptr() % 16 == 0 + for tensor in (hidden_states, self.gate.weight, down.weight) + ) + @staticmethod def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: """Resolve the routed-expert ``(moe_tp, moe_ep)`` split. @@ -1400,18 +1528,20 @@ def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: config. Default split is EP-only; see ``_select_moe_tp_ep``.""" # Every backend here declares ``ActivationType.SiTu`` in its # ``activation_support``; the list is not a preference order. CUTEDSL - # joined once its act-fusion kernel grew the SiTU epilogue. + # joined once its act-fusion kernel grew the SiTU epilogue, and + # CUTEDSL_FC12 (Rubin/SM107) once its fused FC1+FC2 kernel did. supported_backends = { "CUTLASS", "TRTLLM", "CUTEDSL", + "CUTEDSL_FC12", "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", } if model_config.moe_backend not in supported_backends: raise ValueError( "Kimi K3 SiTU routed experts only support the CUTLASS, TRTLLM, " - "CUTEDSL, MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL backends; " + "CUTEDSL, CUTEDSL_FC12, MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL backends; " f"got {model_config.moe_backend!r}." ) if model_config.moe_load_balancer is not None: @@ -1421,7 +1551,26 @@ def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: ) mapping = model_config.mapping if getattr(mapping, "_dwdp_size", 0) > 1: - raise NotImplementedError("Kimi K3 packed-checkpoint streaming does not support DWDP.") + # DWDP does not hook the load path at all: setup_dwdp() scans the + # model *after* every weight is loaded, reads w3_w1_weight / + # w2_weight and their block scales straight off the backend + # (dwdp/setup.py::collect_moe_params), copies them into fabric + # memory, frees the originals, and rebinds param.data to a + # composite VA view. K3's packed-checkpoint streaming produces + # exactly those tensors, so the two are compatible -- but only on + # the backends that declare supports_dwdp, which is the CuteDSL + # family (fused_moe_cute_dsl.py) and NVFP4 only. + # + # MEGAMOE_CUTEDSL is excluded deliberately rather than by + # oversight: its process_weights_after_loading packs into the mega + # buffers and shrinks the raw source params back to placeholders + # (see maybe_finalize_layer), which has never been checked against + # DWDP's copy-then-free of those same originals. + if model_config.moe_backend not in ("CUTEDSL", "CUTEDSL_FC12"): + raise NotImplementedError( + "Kimi K3 with DWDP requires moe_config.backend CUTEDSL or " + f"CUTEDSL_FC12; got {model_config.moe_backend!r}." + ) moe_tp, moe_ep = KimiK3MoERuntime._select_moe_tp_ep(mapping) if moe_tp < 1 or moe_ep < 1 or moe_tp * moe_ep != mapping.tp_size: @@ -1483,8 +1632,38 @@ def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: def forward(self, hidden_states: torch.Tensor, all_rank_num_tokens=None) -> torch.Tensor: """``hidden_states``: ``[num_tokens, hidden_size]`` bf16.""" identity = hidden_states - router_logits = self.gate.compute_logits(hidden_states) + projected_input = None + early_routed_in = None + if self._can_fuse_router_down(hidden_states): + router_logits, projected_input = torch.ops.trtllm.dsv3_router_latent_gemm_op( + hidden_states, self.gate.weight, self.routed_expert_down_proj.weight + ) + else: + if ( + self.latent_down_stream is not None + and do_multi_stream() + and not torch.compiler.is_compiling() + ): + # Forked ahead of the router GEMM: the latent chain needs only + # hidden_states. + self.moe_latent_fork_event.record() + with torch.cuda.stream(self.latent_down_stream): + self.moe_latent_fork_event.wait() + early_routed_in = self._routed_projection( + hidden_states, self.routed_expert_down_proj + ) + self.moe_latent_ready_event.record() + # No record_stream (under graph capture it holds every layer's + # blocks until capture ends): both tensors are read on main only + # after the join, and later work on this stream forks from main. + router_logits = self.gate.compute_logits(hidden_states) moe_all_reduce = self.routed_experts.all_reduce if self._reduce_routed_output else None + overlap_shared_allreduce = ( + self._defer_shared_allreduce + and do_multi_stream() + and self.shared_expert_stream is not None + and not torch.compiler.is_compiling() + ) def _routed_output(): # Latent down/up projections via the min-latency fused GEMM op: @@ -1497,36 +1676,75 @@ def _routed_output(): # replaced the projection module, call it directly: its weight is # an e4m3 buffer the bf16 dsv3 op must not read, and its forward # is already a single fused GEMM (fp8_swap_ab_gemm). - routed_in = self._routed_projection(hidden_states, self.routed_expert_down_proj) + moe_kwargs = {} + if projected_input is not None: + routed_in = projected_input + elif early_routed_in is not None: + routed_in = early_routed_in + # ConfigurableMoE joins the event before its first read of x. + moe_kwargs = {"input_ready_event": self.moe_latent_ready_event} + else: + routed_in = self._routed_projection(hidden_states, self.routed_expert_down_proj) y = self.routed_experts( routed_in, router_logits, all_rank_num_tokens=all_rank_num_tokens, + **moe_kwargs, ) + if overlap_shared_allreduce: + self.moe_routed_event.record() if self._reduce_routed_output: return y # Communication-backed paths return a complete routed result. y = self.routed_expert_norm(y) return self._routed_projection(y, self.routed_expert_up_proj) + def _shared_output(): + shared_out = self.shared_experts( + identity, + final_all_reduce_params=( + AllReduceParams(enable_allreduce=False) + if self._defer_shared_allreduce + else None + ), + ) + if overlap_shared_allreduce: + # Wait for all routed chunks, but overlap the norm/up tail. + self.moe_routed_event.wait() + shared_out = self.shared_experts.down_proj.all_reduce(shared_out) + return shared_out + # Shared experts depend only on the block input, so overlap their GEMMs # with the routed dispatch/expert/combine chain. Multi-stream engages # only under CUDA graphs; otherwise both branches run in order on the # default stream. The shared GatedMLP includes its output all-reduce on - # the auxiliary stream. The join below must precede the routed - # all-reduce: concurrent collectives on different streams can corrupt - # SYMM_MEM all-reduce state. + # the auxiliary stream (deferred past the routed MoE for fused-comm + # MegaMoE). The final join includes the shared all-reduce, so it + # precedes the routed all-reduce: concurrent collectives on different + # streams can corrupt SYMM_MEM all-reduce state. routed_out, shared_out = maybe_execute_in_parallel( _routed_output, - lambda: self.shared_experts(identity), + _shared_output, self.moe_main_event, self.moe_shared_event, self.shared_expert_stream, disable_on_compile=True, ) + if self._defer_shared_allreduce and not overlap_shared_allreduce: + shared_out = self.shared_experts.down_proj.all_reduce(shared_out) if self._reduce_routed_output: - routed_latent = moe_all_reduce(routed_out) - routed_latent = self.routed_expert_norm(routed_latent) + # RMSNorm folded into the routed all-reduce (RMS_NORM fusion): the + # MNNVL kernel normalizes straight out of the lamport buffer, so + # the reduced row never round-trips through HBM for a norm kernel. + # Strategies without the fusion fall back inside AllReduce. + routed_latent = moe_all_reduce( + routed_out, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RMS_NORM, + norm_weight=self.routed_expert_norm.weight, + eps=float(self.routed_expert_norm.variance_epsilon), + ), + ) routed_out = self._routed_projection(routed_latent, self.routed_expert_up_proj) return routed_out + shared_out @@ -3118,6 +3336,11 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: projection = getattr(attention, name) projection._fused_projection = (weakref.ref(fused), offset) offset += projection.out_features + # Repoint to the fused rows now rather than in post_load: + # deferring kept every layer's standalone codes alongside + # its fused copy, ~22.6 GiB at peak across K3's 69 KDA + # layers, which overflows a DEP8 context rank. + projection.transform_weights() logger.info( f"Kimi K3: loaded {num_params} parameters and expert slices of {num_moe_layers} layers" ) diff --git a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py index 56c1179eb870..151e6b4a799e 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py @@ -64,8 +64,8 @@ def get_kda_sm_version() -> int: def is_kda_optimized_supported() -> bool: - """The optimized prefill/decode kernels are Blackwell sm_100 only.""" - return get_kda_sm_version() in (100, 103) + """The optimized prefill/decode kernels need Blackwell sm_100 or Rubin sm_107.""" + return get_kda_sm_version() in (100, 103, 107) @triton.jit(do_not_specialize=["num_tokens", "feature_stride"]) @@ -471,7 +471,8 @@ def commit_kda_pending_drafts( conv_q: torch.Tensor, conv_k: torch.Tensor, conv_v: torch.Tensor, - qkg_cache: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, v_cache: torch.Tensor, beta_cache: torch.Tensor, state_indices: torch.Tensor, @@ -483,14 +484,17 @@ def commit_kda_pending_drafts( ``num_accepted_tokens[n]`` is the number of drafts of slot ``state_indices[n]`` the last verify accepted, which the next verify would take up first. Afterwards ``recurrent_state`` and ``conv_pool`` - hold the state and the conv window that verify would start from, bit for - bit, and the caller clears the slot's count: + hold the state and the conv window that verify would start from (the + window bit for bit), and the caller clears the slot's count: * with per-token states (``state_tok``), the drafts are replayed from the per-draft records the Kimi K3 verify kernels keep there (``replay_kda_token_records``); - * otherwise they are replayed from the replay caches by the fused verify - kernel's own replay steps (``kda_mtp_commit_pending_drafts``). + * otherwise they are replayed from the replay caches (``k_cache`` / + ``g_cache`` / ``v_cache`` / ``beta_cache``) by the fused verify kernel's + own replay steps (``kda_mtp_commit_pending_drafts``). The tensor-core + verify computes a chunk in one pass, so this commit rounds the state + differently from a verify that replays the drafts itself. """ num_spec = conv_q.shape[2] - conv_pool.shape[2] if state_tok is not None: @@ -501,7 +505,8 @@ def commit_kda_pending_drafts( conv_q, conv_k, conv_v, - qkg_cache, + k_cache, + g_cache, v_cache, beta_cache, state_indices, @@ -594,6 +599,37 @@ def is_kda_mtp_verify_available() -> bool: return is_kda_optimized_supported() and is_intree_mtp_available() +def kda_mtp_max_num_spec() -> int: + """Largest draft count the fused multi-token verify kernel supports.""" + from tensorrt_llm._torch.cute_dsl_kernels.blackwell.kimi_k3_kda.kda_mtp_decode import ( + MAX_NUM_SPEC, + ) + + return MAX_NUM_SPEC + + +def kda_mtp_replay_num_spec(num_spec: int) -> Optional[int]: + """Draft count to size the KDA replay caches for, or ``None``. + + ``None`` keeps the legacy intermediate verification buffers, so the + module runs the sequential FLA verify: either the fused kernel cannot + run here, or ``num_spec`` exceeds what it supports. + """ + if not is_kda_mtp_verify_available(): + return None + max_num_spec = kda_mtp_max_num_spec() + if num_spec > max_num_spec: + from tensorrt_llm.logger import logger + + logger.warning_once( + f"KDA fused multi-token verify supports at most {max_num_spec} draft " + f"tokens, got {num_spec}; using the sequential FLA verify.", + key="kda_mtp_num_spec_fallback", + ) + return None + return num_spec + + # --------------------------------------------------------------------------- # Dispatch API used by the KimiKDALinearAttention module. # --------------------------------------------------------------------------- @@ -653,10 +689,15 @@ def __init__( f"verify={self.verify_kernel_path}" ) - def mtp_verify(self, **kwargs) -> torch.Tensor: + def mtp_verify( + self, *, mxfp8_output: bool = False, **kwargs + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Run the fused KDA multi-token verify kernel. - Thin passthrough to ``trtllm::kda_mtp_decode`` (see + Thin passthrough to ``trtllm::kda_mtp_decode``, or with + ``mxfp8_output`` to ``trtllm::kda_mtp_decode_fp8``, which always + fuses the output norm and returns the E4M3 core and its R128c4 + scales (see ``custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py`` for the full argument and state-management contract). Only defined on the optimized path; the FLA fallback is the module's sequential @@ -667,7 +708,11 @@ def mtp_verify(self, **kwargs) -> torch.Tensor: "mtp_verify called on non-optimized path; use the module's " "sequential FLA verify fallback instead." ) - _load_mtp_module() # registers trtllm::kda_mtp_decode + _load_mtp_module() # registers both the BF16 and the MXFP8-output ops + if mxfp8_output: + if not kwargs.pop("fuse_output_norm", True): + raise ValueError("the MXFP8-output KDA MTP op always fuses the output norm") + return torch.ops.trtllm.kda_mtp_decode_fp8(**kwargs) return torch.ops.trtllm.kda_mtp_decode(**kwargs) def can_use_indexed_prefill( diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py index 624fa5baddfd..30231196346d 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py @@ -20,7 +20,7 @@ import math import os import weakref -from typing import List, Optional, Tuple +from typing import List, NamedTuple, Optional, Tuple, Union import torch from fla.modules import ShortConvolution @@ -50,6 +50,18 @@ _KDA_BFA_MULTISTREAM_MAX_ROWS = 128 +class _KdaPrequantizedCore(NamedTuple): + """MXFP8 KDA core emitted by the fused verify epilogue. + + ``activation`` is the ``[tokens, H * V]`` E4M3 post-o_norm core and + ``scale`` its UE8M0 1x128 scales in the flat uint8 R128c4 layout that + ``o_proj.forward_prequantized`` consumes. + """ + + activation: torch.Tensor + scale: torch.Tensor + + def _meta_safe_cast_dtype(module: nn.Module, dtype: torch.dtype) -> None: """Cast floating parameters while preserving meta-device construction. @@ -442,12 +454,14 @@ def _forward_impl( hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: + ) -> Union[torch.Tensor, _KdaPrequantizedCore]: """Run metadata-dependent KDA prefill/decode/verify dispatch. ``output`` is the BCG post-o_norm, pre-o_proj core buffer. When it is supplied, subpaths fill and return its corresponding slices; - otherwise this method returns an allocated core tensor. + otherwise this method returns an allocated core tensor. A + verify-only batch without ``output`` may instead return the fused + verify epilogue's MXFP8 core (see ``forward_verify_fused``). """ mamba_metadata = attn_metadata.mamba_metadata num_prefills = attn_metadata.num_contexts @@ -515,6 +529,9 @@ def _forward_impl( ssm_pool, generation_state_indices, output=(output[num_ctx_tokens:num_tokens] if output is not None else None), + # The MXFP8 core cannot be concatenated with a prefill + # core or stored into the BCG buffer. + allow_prequantized_output=output is None and num_prefills == 0, ) if output is None: cores.append(verify_core) @@ -522,9 +539,12 @@ def _forward_impl( return output return cores[0] if len(cores) == 1 else torch.cat(cores, dim=0) - def _project_output(self, core: torch.Tensor) -> torch.Tensor: + def _project_output(self, core: Union[torch.Tensor, _KdaPrequantizedCore]) -> torch.Tensor: """Project the post-o_norm KDA core and reduce TP partials.""" - out = self.o_proj(core.reshape(-1, self.proj_size)) + if isinstance(core, _KdaPrequantizedCore): + out = self.o_proj.forward_prequantized(core.activation, core.scale) + else: + out = self.o_proj(core.reshape(-1, self.proj_size)) if self._o_allreduce is not None: # Head-sharded TP: every rank ran its head shard on the same # local batch; sum the row-sharded o_proj partials. @@ -1084,7 +1104,8 @@ def forward_verify( ssm_pool, slot_indices, output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: + allow_prequantized_output: bool = False, + ) -> Union[torch.Tensor, _KdaPrequantizedCore]: """Speculative verification: advance each request ``num_steps`` tokens (1 golden + ``num_steps - 1`` padded drafts). @@ -1109,7 +1130,13 @@ def forward_verify( "were not allocated so there is no fallback" ) return self.forward_verify_fused( - x2d, num_steps, layer_cache, ssm_pool, slot_indices, output=output + x2d, + num_steps, + layer_cache, + ssm_pool, + slot_indices, + output=output, + allow_prequantized_output=allow_prequantized_output, ) return self.forward_verify_sequential( x2d, @@ -1122,7 +1149,7 @@ def forward_verify( ) def _project_verify_inputs( - self, x: torch.Tensor, num_rows: int + self, x: torch.Tensor, num_rows: int, preserve_section_strides: bool = False ) -> Optional[ tuple[ torch.Tensor, @@ -1133,7 +1160,13 @@ def _project_verify_inputs( Optional[torch.Tensor], ] ]: - """Project fused QKVG and [f_a | b] inputs for target verification.""" + """Project fused QKVG and [f_a | b] inputs for target verification. + + With ``preserve_section_strides`` the q/k/v/gate sections stay + strided views of the fused QKVG output (and beta of the fused + ``[f_a | b]`` output); the packed fused verify kernel reads them in + place instead of paying one copy per section. + """ qkvg_weight = self._qkvg_proj_weight fused_qkvg = self.qkvg_proj if qkvg_weight is None and fused_qkvg is None: @@ -1165,17 +1198,21 @@ def _project_qkvg() -> torch.Tensor: forget_gate = self.f_b_proj(self.f_a_proj(x)) d = self.proj_size - q_proj, k_proj, v_proj = (part.contiguous() for part in qkvg[..., : 3 * d].split(d, dim=-1)) + q_proj, k_proj, v_proj = qkvg[..., : 3 * d].split(d, dim=-1) qkvg_split_sizes = self.qkvg_split_sizes has_onorm_gate = self.use_full_rank_gate and ( qkvg_weight is not None or (qkvg_split_sizes is not None and len(qkvg_split_sizes) == 4) ) if has_onorm_gate: - onorm_g = qkvg[..., 3 * d : 4 * d].contiguous() + onorm_g = qkvg[..., 3 * d : 4 * d] elif onorm_lowrank is not None: - onorm_g = onorm_lowrank.contiguous() + onorm_g = onorm_lowrank else: onorm_g = None + if not preserve_section_strides: + q_proj, k_proj, v_proj = q_proj.contiguous(), k_proj.contiguous(), v_proj.contiguous() + if onorm_g is not None: + onorm_g = onorm_g.contiguous() return q_proj, k_proj, v_proj, forget_gate, beta, onorm_g def forward_verify_fused( @@ -1186,15 +1223,27 @@ def forward_verify_fused( ssm_pool, slot_indices, output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: + allow_prequantized_output: bool = False, + ) -> Union[torch.Tensor, _KdaPrequantizedCore]: """Fused multi-token verify via ``trtllm::kda_mtp_decode``. - Token layout: the kernel indexes each request's new tokens at - ``cu_seqlens[n] + num_accepted[n] + i``. The runtime packs the - ``num_steps`` new tokens per request contiguously, so we pass - ``cu_seqlens[n] = n * num_steps - num_accepted[n]`` — the shift - lands the kernel's reads/writes exactly on the packed rows. A - negative entry for request 0 is fine: ``bos`` is only ever used + Fused epilogue path (the output gate is available and the output-norm + weights are finalized): the runtime packs ``num_steps`` new rows + contiguously per request, so the kernel derives the row offsets from + the request index and reads each request's pending accepted count + from the slot-indexed replay pool (``packed_token_layout``). It also + applies the gated RMSNorm in its epilogue (``fuse_output_norm``), so + it returns the post-o_norm core directly. When + ``allow_prequantized_output`` is set and ``o_proj`` accepts a + prequantized activation (Rubin CuTe MXFP8 GEMM), the epilogue also + quantizes the core to MXFP8 and ``_project_output`` hands it to + ``o_proj.forward_prequantized``. + + Fallback path: the kernel indexes each request's new tokens at + ``cu_seqlens[n] + num_accepted[n] + i``, so we pass + ``cu_seqlens[n] = n * num_steps - num_accepted[n]``, which lands the + kernel's reads and writes on the packed rows; the output gate runs in + Python. A negative entry for request 0 is fine: ``bos`` is only used additively with a token offset ``>= num_accepted``. """ num_generations = x2d.shape[0] // num_steps @@ -1204,7 +1253,10 @@ def forward_verify_fused( x = x2d.view(num_generations, num_steps, -1) # [B, T, hidden] T_total = num_generations * num_steps - projections = self._project_verify_inputs(x, T_total) + can_fuse_output_norm = self._onorm_w_f32 is not None + projections = self._project_verify_inputs( + x, T_total, preserve_section_strides=can_fuse_output_norm + ) if projections is None: q_proj = self.q_proj(x) k_proj = self.k_proj(x) @@ -1214,26 +1266,21 @@ def forward_verify_fused( onorm_g = None else: q_proj, k_proj, v_proj, forget_gate, beta_proj, onorm_g = projections + fuse_output_norm = can_fuse_output_norm and onorm_g is not None + if not fuse_output_norm: + # The legacy layout has only ever seen dense projections. + q_proj, k_proj, v_proj = q_proj.contiguous(), k_proj.contiguous(), v_proj.contiguous() + beta_proj = beta_proj.contiguous() x_q = q_proj.view(1, T_total, H, K) x_k = k_proj.view(1, T_total, H, K) x_v = v_proj.view(1, T_total, H, self.head_dim) # Raw gate / beta: the kernel applies dt_bias, A_log, the # lower-bound sigmoid gate, and the beta sigmoid itself. g = forget_gate.view(1, T_total, H, K) - beta = beta_proj.contiguous().view(1, T_total, H) + beta = beta_proj.view(1, T_total, H) w_q, w_k, w_v = self._get_mtp_conv_weights() - lower_bound = self.gate_lower_bound - - pending = layer_cache.prev_num_accepted_tokens[ - slot_indices - ] # accepted drafts of the previous round, per req - cu_seqlens = torch.arange( - 0, (num_generations + 1) * num_steps, num_steps, dtype=torch.int32, device=x2d.device - ) - cu_seqlens[:num_generations].sub_(pending) - - out = self._dispatch.mtp_verify( + verify_kwargs = dict( x_q=x_q, x_k=x_k, x_v=x_v, @@ -1250,16 +1297,52 @@ def forward_verify_fused( A_log=self.A_log.detach(), dt_bias=self.dt_bias.detach(), recurrent_state=ssm_pool, - qkg_cache=layer_cache.kda_qkg_cache, + k_cache=layer_cache.kda_k_cache, + g_cache=layer_cache.kda_g_cache, v_cache=layer_cache.kda_v_cache, beta_cache=layer_cache.kda_beta_cache, ssm_state_indices=slot_indices, - cu_seqlens=cu_seqlens, num_spec=num_spec, - num_accepted_tokens=pending, - lower_bound=lower_bound, + lower_bound=self.gate_lower_bound, scale=self.head_k_dim**-0.5, ) + + if fuse_output_norm: + mxfp8_output = ( + allow_prequantized_output + and output is None + and bool(getattr(self.o_proj, "supports_prequantized_input", False)) + and callable(getattr(self.o_proj, "forward_prequantized", None)) + ) + out = self._dispatch.mtp_verify( + **verify_kwargs, + cu_seqlens=None, + num_accepted_tokens=layer_cache.prev_num_accepted_tokens, + packed_token_layout=True, + onorm_g=onorm_g.view(1, T_total, H, self.head_dim), + onorm_weight=self._onorm_w_f32, + onorm_eps=self.o_norm.eps, + fuse_output_norm=True, + mxfp8_output=mxfp8_output, + ) + if mxfp8_output: + return _KdaPrequantizedCore(*out) + # Can be removed once the op writes the core in place. + return self._store_core(out, output) + + pending = layer_cache.prev_num_accepted_tokens[ + slot_indices + ] # accepted drafts of the previous round, per req + cu_seqlens = torch.arange( + 0, (num_generations + 1) * num_steps, num_steps, dtype=torch.int32, device=x2d.device + ) + cu_seqlens[:num_generations].sub_(pending) + + out = self._dispatch.mtp_verify( + **verify_kwargs, + cu_seqlens=cu_seqlens, + num_accepted_tokens=pending, + ) o = out.view(num_generations, num_steps, H, self.head_dim) core = self._output_gate(x, o, onorm_g) # Can be removed once attention writes the core in place. diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index 99891bd3d3a4..c09a514fdf11 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -227,7 +227,7 @@ Still on old path (standalone, with embedded communication): | `fused_moe_densegemm.py` | `TrtllmCutedslDenseGemmNvfp4Impl` (aliased as `DenseGEMMFusedMoE`) | SM100/SM103 | NVFP4 min-latency; CuTe DSL dense GEMM packs all experts into one matrix (vs Cutlass per-expert scatter), efficient for small token counts | `EXTERNAL_COMM` | | `fused_moe_cute_dsl.py` | `CuteDslFusedMoE` | SM100/SM103/SM107 | High throughput NVFP4, generally faster than Cutlass | `EXTERNAL_COMM` | | `fused_moe_cute_dsl_b12x.py` | `CuteDslB12xFusedMoE` | SM120/SM121 | NVFP4 hybrid CUTLASS-prefill / FlashInfer NVFP4 MoE decode — best perf on RTX PRO 6000 (SM120) and DGX Spark (SM121); select via the `CUTEDSL` backend path (it heads that family's candidate list, so it wins on SM120/121 when flashinfer is present and yields to `CuteDslFusedMoE` otherwise); single-GPU-shaped topology only — it rejects both `ep_size > 1` and attention-DP, because it has no dispatch/combine kernel and has never been exercised behind a DP allgather | `EXTERNAL_COMM` | -| `fused_moe_cute_dsl_fc12.py` | `TrtllmCutedslFusedFc12Nvfp4Impl` (`trtllm.cutedsl.fused_fc12.nvfp4`, aliased as `CuteDslFc12FusedMoE`) | SM107 (Rubin) | NVFP4 fused FC1+FC2: a `MoEImplBase` leaf sharing `CuteDslFusedMoE`'s NVFP4 weight layout and outer autotune runner, driving the single persistent `cute_dsl_nvfp4_fc12_fused_rubin` kernel (keeps the FC1->FC2 intermediate on-chip, removing the global round-trip and one launch); reachable through the `CUTEDSL_FC12` backend family or a pinned `impl_id`; requires CuTe DSL Rubin support (`MoEDep.CUTEDSL_RUBIN` in the collected environment); routing tile 128 (1-CTA) or 256 (2-CTA, cluster (2,1)), chosen by the autotuner per shape; non-uGPU | `EXTERNAL_COMM` | +| `fused_moe_cute_dsl_fc12.py` | `TrtllmCutedslFusedFc12Nvfp4Impl` (`trtllm.cutedsl.fused_fc12.nvfp4`, aliased as `CuteDslFc12FusedMoE`) | SM107 (Rubin) | NVFP4 fused FC1+FC2: a `MoEImplBase` leaf sharing `CuteDslFusedMoE`'s NVFP4 weight layout and outer autotune runner, driving the single persistent `cute_dsl_nvfp4_fc12_fused_rubin` kernel, whose gated epilogue is a trace-time SwiGLU or SiTU specialization (keeps the FC1->FC2 intermediate on-chip, removing the global round-trip and one launch); reachable through the `CUTEDSL_FC12` backend family or a pinned `impl_id`; requires CuTe DSL Rubin support (`MoEDep.CUTEDSL_RUBIN` in the collected environment); routing tile 128 (1-CTA) or 256 (2-CTA, cluster (2,1)), chosen by the autotuner per shape; non-uGPU | `EXTERNAL_COMM` | | `mega_moe/mega_moe_deepgemm.py` | `DeepgemmCudaW4a8Mxfp4Mxfp8Impl` (aliased as `MegaMoEDeepGemm`) | SM100/SM103/SM107 | W4A8_MXFP4_MXFP8 via DeepGEMM `fp8_fp4_mega_moe` fused dispatch+GEMM+act+GEMM+combine kernel; requires `hidden_size % 512 == 0` | `FUSED_COMM` | | `mega_moe/mega_moe_cute_dsl.py` | `TrtllmCutedslMegaMoeNvfp4Impl` (aliased as `MegaMoECuteDsl`) | SM100/SM103/SM107 | NVFP4 fused dispatch+FC1+act+FC2+combine; internally selects one of its per-architecture kernel variants. Requires CUDA 13 Cutlass DSL and a symmetric-memory provider. Uses uniform activation constants and applies routing weights before FC2 quantization for DeepSeek-V4, after FC2 for other models. | `FUSED_COMM` | | `marlin/` (+ the `fused_moe_marlin.py` alias module) | the abstract `MarlinFusedMoEBase` and its two leaves, `MarlinCudaNvfp4Impl` and `MarlinCudaW4a16Nvfp4Impl` | SM89-SM99 | W4A16 NVFP4 on Ada/Hopper (BF16 activations + FP4 weights, fused single-launch `marlin_nvfp4_moe_gemm` kernel); supports attention-DP + EP via external comm (scheduler precomputes routing; dispatch payload is plain BF16, no activation scales); non-NVFP4 layers (e.g. unquantized MTP draft layers) degrade to Cutlass in `resolve_moe_impl`, recorded in the layer's `MoEResolutionReport`; no dynamic EPLB | `EXTERNAL_COMM` | diff --git a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py index cbb7c5965b6e..0909c8e7f528 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py @@ -652,6 +652,7 @@ def forward_impl( all_rank_num_tokens: Optional[List[int]] = None, use_dp_padding: Optional[bool] = None, lora_params: Optional[Dict] = None, + input_ready_event: Optional[torch.cuda.Event] = None, **kwargs, ) -> torch.Tensor: """Forward entry point. @@ -665,6 +666,9 @@ def forward_impl( 4. Advances the EPLB ``repeat_idx``. DP-padding handling and chunking live in the scheduler. + + ``input_ready_event``: ``x`` is written on another stream; the join is + deferred to the first kernel that reads ``x``, so routing overlaps it. """ input_ids = kwargs.get("input_ids") routing_aux = kwargs.get("routing_aux") @@ -695,6 +699,7 @@ def forward_impl( input_ids=input_ids, lora_params=lora_params, routing_aux=routing_aux, + input_ready_event=input_ready_event, ) # DWDP: record compute and trigger next prefetch (per-layer, not per-chunk). diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py index d568e7e792a1..45fb7da17d83 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py @@ -55,6 +55,10 @@ NVFP4CuteDslFusedMoEMethod) from .routing import BaseMoeRoutingMethod +# Kill switch for CuteDslFusedMoE.quantize_input_async (fork the NVFP4 input +# quantization onto the MoE aux stream so it overlaps top-k + moe_sort). +ASYNC_INPUT_QUANT = os.environ.get("TLLM_MOE_ASYNC_INPUT_QUANT", "1") == "1" + # These runners are defined inside cute_dsl_custom_ops' ``if # IS_CUTLASS_DSL_AVAILABLE:`` block, which has no else-branch, so importing them # unconditionally would break every importer of this file -- and create_moe @@ -342,7 +346,8 @@ def __init__(self, scaling_vector_size: int = 16, use_direct_expert_metadata: bool = False, use_locality_domain: bool = False, - workload_identity: Optional[Tuple] = None): + workload_identity: Optional[Tuple] = None, + prime_inner_tactics: bool = False): super().__init__() self.forward_impl = forward_impl self.num_experts = num_experts @@ -357,7 +362,14 @@ def __init__(self, assert output_dtype == torch.bfloat16 self.output_dtype = output_dtype self.scaling_vector_size = scaling_vector_size + # Extra terms that distinguish two otherwise identically-shaped + # workloads (locality-domain topology, activation type, ...). Cache key + # only -- it must not imply anything about ``forward_impl``'s signature. self.workload_identity = workload_identity + # Locality domain only: ``forward_impl`` nests its own AutoTuner calls + # and accepts ``overlap_moe_output_memset``. Single-op backends must + # leave this off. + self.prime_inner_tactics = prime_inner_tactics def unique_id(self): identity = ( @@ -443,7 +455,7 @@ def forward(self, tactic: Optional[int], do_preparation: bool = False) -> torch.Tensor: if do_preparation: - if self.workload_identity is not None: + if self.prime_inner_tactics: # Inner FC tuning cannot run from inside the CUDA graph used to # profile an outer tile. Prime every tile's FC1/FC2 cache for # this optimization profile before outer profiling starts. @@ -558,7 +570,8 @@ def __init__(self, local_expert_offset: int, enable_alltoall: bool = False, output_dtype: torch.dtype = torch.bfloat16, - workload_identity: Optional[Tuple] = None): + workload_identity: Optional[Tuple] = None, + prime_inner_tactics: bool = False): super().__init__() self.forward_impl = forward_impl self.num_experts = num_experts @@ -567,7 +580,10 @@ def __init__(self, self.local_expert_offset = local_expert_offset self.enable_alltoall = enable_alltoall self.output_dtype = output_dtype + # See CuteDslFusedMoENvfp4Runner: cache key only, and locality-domain + # only. self.workload_identity = workload_identity + self.prime_inner_tactics = prime_inner_tactics def unique_id(self): identity = ( @@ -622,7 +638,7 @@ def forward(self, tactic: Optional[int], do_preparation: bool = False) -> torch.Tensor: if do_preparation: - if self.workload_identity is not None: + if self.prime_inner_tactics: # See the NVFP4 runner: nested FC tuning must complete before # the outer tile is profiled under CUDA graph capture. for tile_size in self._tile_sizes(): @@ -837,13 +853,9 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: MoERejectReason.DEP_MISSING, "NVFP4 CuteDSL MoE on SM107 requires Rubin support in CuTe DSL" ) - # Keep SiTU enablement scoped to the Blackwell path. Rubin - # integration needs separate end-to-end validation. - if p.activation == "SiTu" and sm_version == 107: - return _reject( - MoERejectReason.ACTIVATION_UNSUPPORTED, - "CuteDSL SiTU is enabled only on SM100/SM103; " - "SM107 integration is not enabled") + # SiTU is served on SM107 too: the Rubin act-fusion kernel carries + # the same trace-time SiTU epilogue as the Blackwell one, and Kimi + # K3 NVFP4 on CUTEDSL is what the Rubin K3 deployments run. # process_weights_after_loading() unswizzles the FC1 block scales, # which asserts 128-row tiles; without this gate an unaligned shard # dies mid weight load with a bare swizzle error. @@ -919,6 +931,10 @@ def __init__( key: torch.cuda.Event() for key in [EventType.Main, EventType.MoeOutputMemset] } + # Async input quantization (see quantize_input_async / + # _join_pending_x_ready). + self._x_quant_event = torch.cuda.Event() + self._pending_x_ready_event: Optional[torch.cuda.Event] = None self._weights_created = False @@ -1055,6 +1071,54 @@ def quantize_input(self, x_sf = x_sf.view(x_row, scale_cols) return x, x_sf + def quantize_input_async( + self, + x: torch.Tensor, + input_ready_event: Optional[torch.cuda.Event] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], + Optional[torch.cuda.Event]]: + """Quantize ``x`` on the MoE aux stream so it overlaps top-k and ``moe_sort``. + + At small decode batches the three per-layer prologue kernels (top-k, + input quantization, routing/permutation) are all latency-bound and + serialize on the main stream (~25 us/layer for 8 tokens on GB300). + Quantization only depends on ``x``, so the scheduler forks it here + before routing; the consumer joins via the returned event right before + the first kernel that reads ``x``/``x_sf`` (``run_moe_nvfp4_impl``, + after ``moe_sort``). Falls back to the synchronous ``quantize_input`` + (event ``None``) when disabled, not NVFP4, already quantized, or no aux + stream is available. The stream running the quantization joins + ``input_ready_event`` (``x`` written on another stream). + """ + if (not ASYNC_INPUT_QUANT or not self.has_nvfp4 + or isinstance(x, Fp4QuantizedTensor) + or not self._has_moe_output_memset_aux_stream()): + if input_ready_event is not None: + input_ready_event.wait() + x_q, x_sf = self.quantize_input(x, post_quant_comm=False) + return x_q, x_sf, None + main_stream = torch.cuda.current_stream() + aux_stream = self.aux_stream_dict[AuxStreamType.MoeOutputMemset] + self.event_dict[EventType.Main].record() + x.record_stream(aux_stream) + with torch.cuda.stream(aux_stream): + self.event_dict[EventType.Main].wait() + if input_ready_event is not None: + input_ready_event.wait() + x_q, x_sf = self.quantize_input(x, post_quant_comm=False) + self._x_quant_event.record() + x_q.record_stream(main_stream) + if x_sf is not None: + x_sf.record_stream(main_stream) + return x_q, x_sf, self._x_quant_event + + def _join_pending_x_ready(self) -> None: + """Wait for the async input quantization (once per forward).""" + event = getattr(self, "_pending_x_ready_event", None) + if event is not None: + event.wait() + self._pending_x_ready_event = None + def run_moe_nvfp4( self, x: torch.Tensor, @@ -1067,6 +1131,7 @@ def run_moe_nvfp4( recv_expert_count: Optional[torch.Tensor] = None, deep_ep_expert_capacity: Optional[int] = None, use_deep_ep_direct_metadata: bool = False, + x_ready_event: Optional[torch.cuda.Event] = None, ) -> torch.Tensor: """NVFP4 MoE computation. @@ -1079,6 +1144,10 @@ def run_moe_nvfp4( weight_view: Bundled weight tensors. Must not be None. use_deep_ep_direct_metadata: Use adapter-free, count-native DeepEP metadata. The scheduler sets this only for the supported path. + x_ready_event: Set when ``x``/``x_sf`` come from + ``quantize_input_async``; joined inside ``run_moe_nvfp4_impl`` + right before the FC1 gather GEMM (the locality-domain path and + empty micro-batches join here). """ assert self.has_nvfp4 assert weight_view is not None @@ -1121,6 +1190,16 @@ def run_moe_nvfp4( self.hidden_size) assert moe_output.dtype == output_dtype + # The async input quantization forked onto the aux stream must be + # joined on every path (CUDA graph capture rejects unjoined work). + # Only the single-tensor impl overlaps it with moe_sort. + if use_locality_domain or token_selected_experts.size(0) == 0: + if x_ready_event is not None: + x_ready_event.wait() + self._pending_x_ready_event = None + else: + self._pending_x_ready_event = x_ready_event + # Empty micro-batches are valid at the backend boundary. Avoid # entering autotuning because its synthetic grouped-GEMM inputs # require at least one output row. @@ -1167,6 +1246,7 @@ def run_moe_nvfp4( workload_identity=workload_identity, use_direct_expert_metadata=use_direct_expert_metadata, use_locality_domain=use_locality_domain, + prime_inner_tactics=use_locality_domain, ) if use_direct_expert_metadata: @@ -1324,6 +1404,10 @@ def run_moe_nvfp4_impl( if self.activation_type == ActivationType.Swiglu: gather_act_kwargs["swiglu_limit_scalar"] = self.act_clamp + # Join the async input quantization (quantize_input_async): x/x_sf were + # produced on the aux stream and are first read here, after moe_sort + # ran on the main stream in parallel with the quantization. + self._join_pending_x_ready() x, x_sf = gather_act_op(**gather_act_kwargs) if self.use_fused_finalize: @@ -1478,6 +1562,7 @@ def run_moe_bf16( enable_alltoall=enable_alltoall, output_dtype=output_dtype, workload_identity=workload_identity, + prime_inner_tactics=use_locality_domain, ) inputs = [x, token_selected_experts, token_final_scales, moe_output] @@ -2016,6 +2101,7 @@ def run_moe( recv_expert_count=plan.recv_expert_count, deep_ep_expert_capacity=plan.deep_ep_expert_capacity, use_deep_ep_direct_metadata=plan.use_deep_ep_direct_metadata, + x_ready_event=ctx.x_ready_event, ) elif self.has_deepseek_fp8_block_scales: result = self.run_moe_fp8_block_scales( diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_fc12.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_fc12.py index 5969ad27d6c3..2472709f992a 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_fc12.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_fc12.py @@ -15,9 +15,10 @@ """TrtllmCutedslFusedFc12Nvfp4Impl: ``trtllm.cutedsl.fused_fc12.nvfp4``. FC1+FC2-fused CuteDSL NVFP4 MoE backend for Rubin (SM107). One persistent kernel -(``trtllm::cute_dsl_nvfp4_fc12_fused_rubin``) does gather + FC1 GEMM + SwiGLU + -requant + FC2 GEMM + finalize (scatter-add into ``moe_output``), so the FC1->FC2 -intermediate never round-trips through global memory. +(``trtllm::cute_dsl_nvfp4_fc12_fused_rubin``) does gather + FC1 GEMM + gated +epilogue (SwiGLU, or SiTU for Kimi K3) + requant + FC2 GEMM + finalize +(scatter-add into ``moe_output``), so the FC1->FC2 intermediate never +round-trips through global memory. A leaf in the ``MoEImplBase`` sense: it declares its ``descriptor``, registers itself, and owns the four abstract methods. The NVFP4 weight layout, quant method @@ -106,7 +107,7 @@ class TrtllmCutedslFusedFc12Nvfp4Impl(MoEImplBase): supports_apply_router_weight_on_input=True, ), input_requirement=MoEInputRequirement(routing_scales_dtype=torch.float32), - doc="Rubin (SM107) NVFP4 fused FC1 + SwiGLU + FC2 + finalize persistent " + doc="Rubin (SM107) NVFP4 fused FC1 + SwiGLU/SiTU + FC2 + finalize persistent " "CuTe DSL grouped GEMM.", ) # Taken off the descriptor rather than restated: the scheduler reads these @@ -116,11 +117,14 @@ class TrtllmCutedslFusedFc12Nvfp4Impl(MoEImplBase): capabilities = descriptor.capabilities input_requirement = descriptor.input_requirement - # The fused FC12 kernel has no activation selector: FC1 weights are the - # gate/up pair and the epilogue is SwiGLU (+ clamp). Relu2 is not gated and - # is not supported, so the resolver turns such layers down here. + # FC1 weights are the gate/up pair and the epilogue is a trace-time + # specialization of SwiGLU (+ clamp) or SiTU, whose two soft-caps arrive as + # the uniform-scalar alpha/beta pair (``SiTuActivation.constants()``). Relu2 + # is not gated and is not supported, so the resolver turns such layers down + # here. activation_support = MoEActivationSupport( - kinds=frozenset({ActivationType.Swiglu}), + kinds=frozenset({ActivationType.Swiglu, ActivationType.SiTu}), + alpha_beta=ActivationParamShape.UNIFORM_SCALAR, limit=ActivationParamShape.UNIFORM_SCALAR, limit_when_absent=float("inf"), ) @@ -380,6 +384,9 @@ def run_moe_nvfp4( local_expert_offset=weight_view.slot_start, enable_finalize_fusion=self.use_fused_finalize, enable_alltoall=enable_alltoall, + # The inner kernel cache keys on the epilogue; the outer AutoTuner + # unique_id must too, or SwiGLU and SiTU layers share a tactic. + workload_identity=self._epilogue_identity(), ) inputs = [x, token_selected_experts, token_final_scales, x_sf, moe_output, weight_view] _, best_tactic = tuner.choose_one( @@ -390,6 +397,19 @@ def run_moe_nvfp4( ) return runner(inputs, tactic=best_tactic) + def _situ_betas(self) -> tuple[float, float]: + """``(situ_beta, situ_linear_beta)`` for the fused op; -1 disables them. + + ``SiTuActivation.constants()`` lands gate_softcap in ``act_alpha`` and + linear_softcap in ``act_beta``, the same mapping CuteDslFusedMoE uses. + """ + if self.activation_type != ActivationType.SiTu: + return -1.0, -1.0 + return float(self.act_alpha), float(self.act_beta) + + def _epilogue_identity(self) -> tuple[int, float, float]: + return (int(self.activation_type), *self._situ_betas()) + def run_moe_nvfp4_impl( self, x: torch.Tensor, @@ -417,6 +437,7 @@ def run_moe_nvfp4_impl( ) effective_top_k = token_selected_experts.size(1) esp = weight_view.expert_size_per_partition + situ_beta, situ_linear_beta = self._situ_betas() slot_start = weight_view.slot_start ( @@ -436,8 +457,9 @@ def run_moe_nvfp4_impl( tile_tokens_dim=tile_size, ) - # One fused op: gather + FC1 GEMM + SwiGLU + requant + FC2 GEMM + - # finalize (scatter-add into moe_output = a2a combine workspace). + # One fused op: gather + FC1 GEMM + gated act (SwiGLU or SiTU) + requant + # + FC2 GEMM + finalize (scatter-add into moe_output = a2a combine + # workspace). # fc1_alpha/fc2_alpha map 1:1 to the two-op path's per-expert global # scales (the kernel takes split alphas). The three atomic counters are # allocated + memset inside the op runner. The moe_output zeroing memset @@ -471,6 +493,9 @@ def run_moe_nvfp4_impl( ep_size=self.mapping.moe_ep_size, enable_alltoall=enable_alltoall, scaling_vector_size=16, + activation_type=int(self.activation_type), + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, ) return moe_output diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py index 65d34e5e22ac..683302725064 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py @@ -609,6 +609,10 @@ class MoERunContext: # produced by quantize_input x: torch.Tensor # activations [num_tokens, hidden_size] x_sf: Optional[torch.Tensor] # scale factors, when the input is quantized + # set when x/x_sf were produced on an auxiliary stream (async input + # quantization); the impl must wait on it before the first kernel that + # reads x/x_sf. None when quantization ran on the main stream. + x_ready_event: Optional[torch.cuda.Event] = None # produced by the outer forward output_dtype: Optional[torch.dtype] = None do_finalize: bool = True diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py index 8ad1a76da6a2..52504f7b6980 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py @@ -69,6 +69,16 @@ from .configurable_moe import ConfigurableMoE +def _wait_input_ready(event: Optional[torch.cuda.Event]) -> None: + """Join a pending ``input_ready_event`` before the first read of ``x``. + + Returns ``None`` so callers can clear their handle in one statement. + """ + if event is not None: + event.wait() + return None + + class MoEScheduler(ABC): """Forward-execution strategy for ConfigurableMoE. @@ -93,6 +103,7 @@ def forward( input_ids: Optional[torch.Tensor] = None, lora_params: Optional[Dict] = None, routing_aux: Optional[torch.Tensor] = None, + input_ready_event: Optional[torch.cuda.Event] = None, ) -> torch.Tensor: ... @@ -135,7 +146,11 @@ def forward( input_ids: Optional[torch.Tensor] = None, lora_params: Optional[Dict] = None, routing_aux: Optional[torch.Tensor] = None, + input_ready_event: Optional[torch.cuda.Event] = None, ) -> torch.Tensor: + """``input_ready_event``: ``x`` is written on another stream. The join + is deferred to the first kernel that reads ``x`` so routing overlaps it. + """ moe = self.moe # ========== Step 1: Handle padding ========== @@ -181,6 +196,7 @@ def forward( all_rank_num_tokens_padded = [all_rank_max_num_tokens] * len(all_rank_num_tokens) local_n = x.shape[0] if local_n < all_rank_max_num_tokens: + input_ready_event = _wait_input_ready(input_ready_event) pad = all_rank_max_num_tokens - local_n x = torch.cat([x, x.new_zeros((pad, x.shape[1]))], dim=0) router_logits = torch.cat( @@ -207,8 +223,10 @@ def forward( input_ids, lora_params=lora_params, routing_aux=routing_aux, + input_ready_event=input_ready_event, ) else: + input_ready_event = _wait_input_ready(input_ready_event) outputs = self._forward_multiple_chunks( x, router_logits, @@ -326,6 +344,7 @@ def _forward_single_chunk( input_ids: Optional[torch.Tensor] = None, lora_params: Optional[Dict] = None, routing_aux: Optional[torch.Tensor] = None, + input_ready_event: Optional[torch.cuda.Event] = None, ) -> torch.Tensor: moe = self.moe is_first_call = moe.repeat_idx == 0 @@ -346,6 +365,7 @@ def _forward_single_chunk( input_ids=input_ids, lora_params=lora_params, routing_aux=routing_aux, + input_ready_event=input_ready_event, ) def _forward_chunk_impl( @@ -363,6 +383,7 @@ def _forward_chunk_impl( lora_params: Optional[Dict] = None, row_offset: Optional[int] = 0, routing_aux: Optional[torch.Tensor] = None, + input_ready_event: Optional[torch.cuda.Event] = None, ) -> torch.Tensor: """Unified per-chunk execution flow for all external-comm backends. @@ -386,6 +407,25 @@ def _forward_chunk_impl( # ========== Step 1: EPLB - Start wait GPU stage ========== moe._load_balancer_start_wait_gpu_stage(is_first_call) + # Async input quantization (backends providing quantize_input_async, + # currently CuteDSL NVFP4): quantize x on the MoE aux stream now so it + # overlaps routing (top-k) and the backend's moe_sort on the main + # stream -- at decode batch 8 these three latency-bound kernels + # otherwise serialize to ~25 us per layer. Only when nothing modifies + # x before quantization (no apply_router_weight_on_input) and no comm + # strategy decides the quantize/dispatch order. + async_quant = None + x_ready_event = None + if ( + moe.comm is None + and not moe.apply_router_weight_on_input + and not isinstance(x, Fp4QuantizedTensor) + and hasattr(moe.backend, "quantize_input_async") + ): + # The quantization stream joins input_ready_event itself. + async_quant = moe.backend.quantize_input_async(x, input_ready_event) + input_ready_event = None + # ========== Step 2: Apply routing ========== # External dispatch (Step 5) sends per-token expert/scale payloads, so # routing must be precomputed whenever a comm strategy is active — even @@ -411,6 +451,8 @@ def _forward_chunk_impl( and not moe.apply_router_weight_on_input and routing_aux is None ): + # A fused route+quant reads x. + input_ready_event = _wait_input_ready(input_ready_event) fused_result = moe.backend.try_fused_route_quant(x, router_logits) else: fused_result = None @@ -461,6 +503,7 @@ def _forward_chunk_impl( # apply_router_weight_on_input: fuse top-k weight onto x if moe.apply_router_weight_on_input: + input_ready_event = _wait_input_ready(input_ready_event) assert x.dtype != torch.float8_e4m3fn, ( "Current workaround for apply_router_weight_on_input does not support fp8 input" ) @@ -536,6 +579,7 @@ def _forward_chunk_impl( should_update_eplb_after_dispatch = True # ========== Step 5: Quantization + dispatch (pre/post-quant adaptive ordering) ========== + input_ready_event = _wait_input_ready(input_ready_event) use_deep_ep_direct_metadata = False if moe.comm is not None: # Debug: optional dummy AllReduce to break load-balancing artifacts @@ -591,8 +635,10 @@ def _forward_chunk_impl( if not uses_internal_dispatch_quantization: x, x_sf = moe.backend.quantize_input(x, post_quant_comm=False) else: - # No comm: just quantize - if not used_fused_route_quant: + # No comm: just quantize (or pick up the async result forked above) + if async_quant is not None: + x, x_sf, x_ready_event = async_quant + elif not used_fused_route_quant: x, x_sf = moe.backend.quantize_input(x, post_quant_comm=False) # ========== Step 6: MoE computation ========== @@ -608,6 +654,7 @@ def _forward_chunk_impl( all_rank_num_tokens=all_rank_num_tokens, lora_params=lora_params, use_deep_ep_direct_metadata=use_deep_ep_direct_metadata, + x_ready_event=x_ready_event, ) final_hidden_states = moe.backend.run_moe( ctx, @@ -980,6 +1027,7 @@ def _build_run_context( all_rank_num_tokens: Optional[List[int]], lora_params: Optional[Dict], use_deep_ep_direct_metadata: bool = False, + x_ready_event: Optional[torch.cuda.Event] = None, ) -> MoERunContext: """The single ``run_moe`` argument set, identical for every backend. @@ -993,6 +1041,7 @@ def _build_run_context( token_final_scales=token_final_scales, x=x, x_sf=x_sf, + x_ready_event=x_ready_event, output_dtype=output_dtype, do_finalize=do_finalize, lora_params=lora_params if moe.backend.capabilities.supports_moe_lora else None, @@ -1091,6 +1140,7 @@ def forward( input_ids: Optional[torch.Tensor], lora_params: Optional[Dict] = None, routing_aux: Optional[torch.Tensor] = None, + input_ready_event: Optional[torch.cuda.Event] = None, ) -> torch.Tensor: """Sequential multi-chunk path for MegaMoE-style backends. @@ -1101,6 +1151,7 @@ def forward( kernel for the cross-rank barrier. """ del use_dp_padding # MegaMoE has no host-side cross-rank shape alignment. + _wait_input_ready(input_ready_event) # Fused-comm (MegaMoE) backends cannot carry LoRA adapters; routed-expert # MoE LoRA is supported only on the CUTLASS backend. Reject rather than diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index bcfb7529311b..4e8e4003f0ae 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -86,6 +86,15 @@ GB = 1 << 30 +# Page size of a standalone drafter's separate KV pool (DFlash, standalone +# DSpark). Deliberately not kv_cache_config.tokens_per_block: that describes +# the target pool, and target attention backends and deployment configs set it +# to values the drafter's block-decode kernels reject (FP4 MLA forces 128, which +# the trtllm-gen MLA decode does not take). The drafter adopts whatever page +# size its pool is carved with, so the pool picks one every drafter kernel +# accepts. Matches the KvCacheConfig.tokens_per_block default. +_STANDALONE_DRAFT_TOKENS_PER_BLOCK = 32 + def ceil_div(a: int, b: int) -> int: return (a + b - 1) // b @@ -1043,10 +1052,13 @@ def _per_manager_cache_cost(self, *, is_draft: bool = False, mapping=None, + tokens_per_block: Optional[int] = None, include_request_cost: bool = False, **extra_kwargs) -> CacheCost: kv_cache_config = (kv_cache_config if kv_cache_config is not None else self._kv_cache_config) + if tokens_per_block is None: + tokens_per_block = self._tokens_per_block if not is_draft: derived_windows = _derive_v2_layer_type_attention_windows( kv_cache_config, manager_cls, model_config, self._max_seq_len) @@ -1061,7 +1073,7 @@ def _per_manager_cache_cost(self, manager_cls.get_cache_size_per_token( model_config, mapping if mapping is not None else self._mapping, - tokens_per_block=self._tokens_per_block, + tokens_per_block=tokens_per_block, max_seq_len=self._max_seq_len, max_batch_size=self._max_batch_size, max_num_tokens=self._max_num_tokens, @@ -1077,6 +1089,7 @@ def _per_manager_cache_cost(self, model_config, kv_cache_config, mapping=mapping, + tokens_per_block=tokens_per_block, is_draft=is_draft, # Zero is a valid weight: a manager whose entire retention # is intercept-funded spends nothing per request from the @@ -1093,6 +1106,7 @@ def _per_manager_cache_cost(self, model_config, kv_cache_config, mapping=mapping, + tokens_per_block=tokens_per_block, is_draft=is_draft, **extra_kwargs) if quota_per_request is not None: @@ -1107,6 +1121,7 @@ def _per_manager_request_estimate(self, kv_cache_config: KvCacheConfig, *, mapping=None, + tokens_per_block: Optional[int] = None, is_draft: bool = False, allow_zero: bool = False, **extra_kwargs) -> Optional[int]: @@ -1122,7 +1137,8 @@ def _per_manager_request_estimate(self, estimate = estimate_method( model_config, mapping if mapping is not None else self._mapping, - tokens_per_block=self._tokens_per_block, + tokens_per_block=(tokens_per_block if tokens_per_block is not None + else self._tokens_per_block), max_seq_len=self._max_seq_len, max_batch_size=self._max_batch_size, max_num_tokens=self._max_num_tokens, @@ -1227,6 +1243,9 @@ def scaled(cost: CacheCost) -> CacheCost: effective_draft_config = self._get_draft_kv_model_config() draft_kv_cache_config = self._get_one_model_draft_kv_cache_config( kv_cache_config, self._max_seq_len) + # Charge the page size the pool is allocated with, not the + # target's: per-request quotas round up to whole pages. + draft_tokens_per_block = self._get_draft_tokens_per_block() # Resolve draft manager class from draft config — may differ # from target (e.g. hybrid target + plain transformer draft). draft_kv_cache_manager_cls = get_kv_cache_manager_cls( @@ -1245,6 +1264,7 @@ def scaled(cost: CacheCost) -> CacheCost: effective_draft_config, draft_kv_cache_config, mapping=draft_mapping, + tokens_per_block=draft_tokens_per_block, is_draft=True, include_request_cost=include_request_cost)) elif self._mapping.is_last_pp_rank(): @@ -1255,6 +1275,7 @@ def scaled(cost: CacheCost) -> CacheCost: effective_draft_config, draft_kv_cache_config, mapping=draft_mapping, + tokens_per_block=draft_tokens_per_block, num_layers=self._get_num_draft_layers(), is_draft=True, include_request_cost=include_request_cost)) @@ -2143,6 +2164,28 @@ def _get_draft_kv_model_config(self) -> ModelConfig: effective_draft_config.quant_config = neutral_quant return effective_draft_config + def _get_draft_tokens_per_block(self) -> int: + """Page size of the separate one-model draft KV pool. + + A standalone drafter (DFlash, standalone DSpark) gets its own page + size, ``_STANDALONE_DRAFT_TOKENS_PER_BLOCK``, NOT self._tokens_per_block: + the target's value can be one the drafter's kernels reject (FP4 MLA + forces 128; trtllm-gen MLA decode takes only 32 or 64). Those drafters + read their pool geometry from the manager, so nothing else has to agree. + + Every other separate draft pool (MTP, Eagle3, ...) keeps the target's + page size: its draft layers run through the target's attention + metadata, which is not re-planned for a different page size. + + Both the budget split (_get_draft_cache_cost) and the allocation + (_create_one_model_draft_kv_cache_manager) must go through here, or + the split charges one page granularity and the pool uses another. + """ + spec_dec_mode = self._speculative_config.spec_dec_mode + if spec_dec_mode.is_dflash() or spec_dec_mode.is_dspark(): + return _STANDALONE_DRAFT_TOKENS_PER_BLOCK + return self._tokens_per_block + def _get_num_draft_layers(self) -> int: """Return the actual number of draft KV cache layers. @@ -2256,6 +2299,13 @@ def _create_one_model_draft_kv_cache_manager( draft_mapping = self._mapping if draft_mapping.has_cp_helix(): draft_mapping = draft_mapping.repurpose_helix_cp_to_tp() + draft_tokens_per_block = self._get_draft_tokens_per_block() + if draft_tokens_per_block != self._tokens_per_block: + # Logged because nothing fails when the two differ: only the draft + # pool's block count, block-table width and fragmentation change. + logger.info( + f"Separate draft KV cache page size {draft_tokens_per_block} " + f"(target pool uses {self._tokens_per_block}).") return _create_kv_cache_manager( model_engine=None, max_cuda_graph_batch_size=self._model_engine. @@ -2263,7 +2313,7 @@ def _create_one_model_draft_kv_cache_manager( kv_cache_manager_cls=draft_kv_cache_manager_cls, mapping=draft_mapping, kv_cache_config=draft_kv_config, - tokens_per_block=self._tokens_per_block, + tokens_per_block=draft_tokens_per_block, max_seq_len=max_seq_len, max_batch_size=self._max_batch_size, spec_config=self._speculative_config, @@ -2550,6 +2600,40 @@ def _split_kv_cache_budget_for_draft( return target_kv_cache_config, draft_kv_cache_config target_budget, draft_budget = shares + # The proportional split assumes the draft pool holds as many tokens as + # the target pool. A disaggregated generation worker never runs the + # prompt forward, so an external drafter there only stores the tokens + # generated on this worker; TLLM_DRAFT_KV_GPU_BUDGET_GIB caps the + # draft's GPU share and hands the rest to the target pool. + draft_cap_gib = os.environ.get("TLLM_DRAFT_KV_GPU_BUDGET_GIB") + if draft_cap_gib and budget_attr == "max_gpu_total_bytes": + try: + cap_gib = float(draft_cap_gib) + except ValueError as e: + raise ValueError( + "TLLM_DRAFT_KV_GPU_BUDGET_GIB must be a number of GiB, got " + f"{draft_cap_gib!r}") from e + # Reject NaN too: it fails every comparison below. + if not cap_gib > 0: + raise ValueError( + "TLLM_DRAFT_KV_GPU_BUDGET_GIB must be positive, got " + f"{draft_cap_gib!r}") + draft_cap = draft_kv.intercept + int(cap_gib * GB) + if draft_kv.quota_per_request is not None: + # The allocator raises the draft quota to one request's floor + # regardless of the split (see _compute_draft_budget_shares), + # so a lower cap would silently overcommit the target's share. + draft_cap = max(draft_cap, + draft_kv.intercept + draft_kv.quota_per_request) + if draft_cap < draft_budget: + logger.info( + f"Capping the draft KV cache GPU budget at " + f"{draft_cap / GB:.2f} GiB (TLLM_DRAFT_KV_GPU_BUDGET_GIB=" + f"{draft_cap_gib}); the proportional share was " + f"{draft_budget / GB:.2f} GiB.") + draft_budget = draft_cap + target_budget = total_budget - draft_budget + logger.info( f"Splitting KV cache {budget_attr}: total={total_budget / GB:.2f} GiB, " f"target={target_budget / GB:.2f} GiB ({target_kv}), " @@ -3357,11 +3441,11 @@ def _create_kv_cache_manager( MambaHybridCacheManagerV2) if (spec_config is not None and issubclass(kv_cache_manager_cls, kda_replay_manager_types)): - from ..modules.kimi_kda._kda_kernels import \ - is_kda_mtp_verify_available - if is_kda_mtp_verify_available(): - kda_extra_kwargs["kda_replay_num_spec"] = ( - spec_config.tokens_per_gen_step - 1) + from ..modules.kimi_kda._kda_kernels import kda_mtp_replay_num_spec + kda_replay_num_spec = kda_mtp_replay_num_spec( + spec_config.tokens_per_gen_step - 1) + if kda_replay_num_spec is not None: + kda_extra_kwargs["kda_replay_num_spec"] = kda_replay_num_spec if is_glm5_next: # The manager places an indexer buffer on every attention layer. from ..attention.backends.sparse.glm_kpool import \ diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index a9b2b16bfbe3..8878090ebe18 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -339,8 +339,8 @@ def resolve_auto_ssm_cache_dtype(config, fallback): kv_cache_config.mamba_ssm_cache_dtype; prefill, decode and sequential verify then stage the addressed rows through an fp32 copy and round the committed state back to bf16. The fused MTP verify kernel - (``trtllm::kda_mtp_decode``) has no such staging and rejects a non-fp32 - pool outright, so bf16 and fused KDA MTP verify are mutually exclusive. + (``trtllm::kda_mtp_decode``) reads and writes a bf16 pool directly and + rounds the committed state to nearest. A checkpoint-declared mamba_ssm_cache_dtype is not applied to Kimi K3 (the released checkpoints do not carry the field); it is logged when it would have diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index c74a0a8d9e7e..aeb0c9f806c6 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -835,9 +835,19 @@ def _warmup_scheduled(self, resource_manager: ResourceManager, kv_cache_manager) gc.collect() torch.cuda.empty_cache() - # Helix CP is decode-only and runs into issues with the - # autotuner warmup's context requests. - if self._context_warmups_supported and not self.mapping.has_cp_helix(): + # Helix CP is decode-only, so it tunes on generation requests. + if self._context_warmups_supported and self.mapping.has_cp_helix(): + with ( + self._warmup_timer.phase( + "autotuner", metrics=self._metrics, metric_name="autotuner_warmup_seconds" + ), + self._maybe_bypass_torch_compile(bypass=True), + ): + self._run_autotuner_warmup(resource_manager, generation_only=True) + log_mem_snapshot("warmup/after_autotuner") + gc.collect() + torch.cuda.empty_cache() + elif self._context_warmups_supported: with ( self._warmup_timer.phase( "autotuner", metrics=self._metrics, metric_name="autotuner_warmup_seconds" @@ -1497,8 +1507,14 @@ def _release_megamoe_profiling_scratch(): if release_megamoe_scratch is not None: release_megamoe_scratch() - def _run_autotuner_warmup(self, resource_manager: ResourceManager) -> None: - """Runs forward passes to populate the autotuner cache.""" + def _run_autotuner_warmup( + self, resource_manager: ResourceManager, generation_only: bool = False + ) -> None: + """Runs forward passes to populate the autotuner cache. + + ``generation_only`` tunes on one batch of ``max_batch_size`` generation + requests instead of a context batch, for Helix CP generation workers. + """ from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( IS_FLASHINFER_MXFP8_CUTE_DSL_AVAILABLE, MXFP8GemmRunner, @@ -1592,14 +1608,22 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager) -> None: enable_native_mxfp8_autotuner = bool(native_mxfp8_methods) kv_cache_manager = resource_manager.get_resource_manager(self.kv_cache_manager_key) - token_num_upper_bound = min( - self._config.max_num_tokens, - self._config.max_batch_size * (self._config.max_seq_len - 1), - ) - curr_max_num_tokens = kv_cache_manager.get_num_available_tokens( - token_num_upper_bound=token_num_upper_bound, - max_num_draft_tokens=self._config.original_max_draft_len, - ) + if generation_only: + tokens_per_gen_request = 1 + self._config.max_total_draft_tokens + num_gen_requests = min( + self._config.max_batch_size, + self._config.max_num_tokens // tokens_per_gen_request, + ) + curr_max_num_tokens = num_gen_requests * tokens_per_gen_request + else: + token_num_upper_bound = min( + self._config.max_num_tokens, + self._config.max_batch_size * (self._config.max_seq_len - 1), + ) + curr_max_num_tokens = kv_cache_manager.get_num_available_tokens( + token_num_upper_bound=token_num_upper_bound, + max_num_draft_tokens=self._config.original_max_draft_len, + ) # The warmup runs again after the KV cache is rebuilt; when the token # shape is unchanged and no MXFP8 tuning is pending, the profiling @@ -1629,10 +1653,13 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager) -> None: f"FlashInfer MXFP8={enable_flashinfer_mxfp8_autotuner})..." ) - warmup_configs = [(curr_max_num_tokens, 0)] - if self.guided_decoder is None and not self.mapping.has_pp(): - # Add generation request to warmup the autotuner cache. - warmup_configs.append((1 + self._config.max_total_draft_tokens, 1)) + if generation_only: + warmup_configs = [(curr_max_num_tokens, num_gen_requests)] + else: + warmup_configs = [(curr_max_num_tokens, 0)] + if self.guided_decoder is None and not self.mapping.has_pp(): + # Add generation request to warmup the autotuner cache. + warmup_configs.append((1 + self._config.max_total_draft_tokens, 1)) def run_autotuner_pass(autotune_context: Any, synchronize_trtllm_cache: bool) -> bool: """Run one isolated tuning pass with fresh synthetic batches.""" diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 0108de49cb14..207817a9ec05 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -97,19 +97,19 @@ def _allocate_kda_beta_cache( shape: Tuple[int, ...], device: Optional[torch.device], ) -> torch.Tensor: - """Allocate a logically shaped beta cache with aligned physical rows.""" + """Allocate a logically shaped bf16 beta cache with aligned physical rows.""" logical_num_heads = shape[-1] heads_per_alignment = (_KDA_BETA_CACHE_ALIGNMENT_BYTES // - torch.float32.itemsize) + torch.bfloat16.itemsize) padded_num_heads = ((logical_num_heads + heads_per_alignment - 1) // heads_per_alignment) * heads_per_alignment # CuTe requires every nested cache view to be 16-byte aligned. Keep the - # logical head count while padding the physical stride (six fp32 heads are - # 24 bytes, so an unpadded row makes alternating views only 8-byte aligned). + # logical head count while padding the physical stride (six bf16 heads are + # 12 bytes, so an unpadded row makes alternating views only 4-byte aligned). return torch.zeros( *shape[:-1], padded_num_heads, - dtype=torch.float32, + dtype=torch.bfloat16, device=device, )[..., :logical_num_heads] @@ -248,6 +248,19 @@ def _allocate_mamba_seed_buffer(cache_size: int, rank_offset: int, return torch.tensor(slot_seeds, dtype=torch.int64, device=device) +def _host_slots_to_device(slots: List[int], + device: torch.device) -> torch.Tensor: + """Copy host slot ids to ``device`` without a host sync. + + A pageable copy would wait for the queued step. Each call stages through a + fresh pinned tensor, so no later host write can race the copy. + """ + return torch.tensor(slots, + dtype=torch.long, + pin_memory=device.type == "cuda" + and prefer_pinned()).to(device, non_blocking=True) + + def _mamba_rank_offset(mapping: Mapping) -> int: """Distinct seed offset per (tp_rank, pp_rank) so independent ranks don't draw identical streams when not coordinated.""" @@ -426,17 +439,19 @@ class SpeculativeState(State): # the LAST GOLDEN token; the accepted drafts recorded in # prev_num_accepted_tokens are pending in these caches and are # replayed by the kernel at the start of the next verify round. - # Extended conv caches [layers, slots, dim, (W-1) + num_spec] fp32, - # dim-contiguous (stride(dim) == 1): columns [0, W-1) are the + # Extended conv caches [layers, slots, dim, (W-1) + num_spec] bf16 + # (raw projection inputs, so bf16 is exact), dim-contiguous (stride(dim) == 1): columns [0, W-1) are the # committed raw-input window, tail columns the pending drafts' raw # inputs. kda_conv_q: torch.Tensor | None = None kda_conv_k: torch.Tensor | None = None kda_conv_v: torch.Tensor | None = None - # Post-processed per-draft quantities for replay: - # [layers, slots, num_spec, 3, H*K] (q/k/gate), [.., num_spec, H*V], - # [.., num_spec, H] — all fp32. - kda_qkg_cache: torch.Tensor | None = None + # Post-processed per-draft quantities for replay: k + # [layers, slots, num_spec, H*K] bf16, gate [.., num_spec, H*K] fp32 + # (feeds the cumulative-gate exponents), v [.., num_spec, H*V] bf16, + # beta [.., num_spec, H] bf16. q is recomputed from the conv caches. + kda_k_cache: torch.Tensor | None = None + kda_g_cache: torch.Tensor | None = None kda_v_cache: torch.Tensor | None = None kda_beta_cache: torch.Tensor | None = None @@ -459,7 +474,7 @@ class SpeculativeState(State): def has_kda_replay_caches(self) -> bool: """True when the KDA fused-verify replay caches were allocated (the fused ``trtllm::kda_mtp_decode`` verify path).""" - return self.kda_qkg_cache is not None + return self.kda_k_cache is not None def commit_conv_window(self, slot_indices: torch.Tensor, conv_pool: torch.Tensor) -> None: @@ -638,11 +653,13 @@ def __init__( extended_s = w_kernel - 1 + M def _dim_contiguous_conv_cache(): + # raw bf16 projection inputs: bf16 is exact and halves + # the per-draft replay traffic of the verify kernel return torch.zeros(num_local_layers, max_batch_size, extended_s, section_dim, - dtype=torch.float32, + dtype=torch.bfloat16, device=device).transpose(-1, -2) spec_kwargs['prev_num_accepted_tokens'] = torch.zeros( @@ -650,25 +667,31 @@ def _dim_contiguous_conv_cache(): spec_kwargs['kda_conv_q'] = _dim_contiguous_conv_cache() spec_kwargs['kda_conv_k'] = _dim_contiguous_conv_cache() spec_kwargs['kda_conv_v'] = _dim_contiguous_conv_cache() - spec_kwargs['kda_qkg_cache'] = torch.zeros(num_local_layers, - max_batch_size, - M, - 3, - section_dim, - dtype=torch.float32, - device=device) - spec_kwargs['kda_v_cache'] = torch.zeros(num_local_layers, + spec_kwargs['kda_k_cache'] = torch.zeros(num_local_layers, + max_batch_size, + M, + section_dim, + dtype=torch.bfloat16, + device=device) + spec_kwargs['kda_g_cache'] = torch.zeros(num_local_layers, max_batch_size, M, section_dim, dtype=torch.float32, device=device) + spec_kwargs['kda_v_cache'] = torch.zeros(num_local_layers, + max_batch_size, + M, + section_dim, + dtype=torch.bfloat16, + device=device) spec_kwargs['kda_beta_cache'] = _allocate_kda_beta_cache( (num_local_layers, max_batch_size, M, nheads), device) ssm_spec_cache = [ spec_kwargs['kda_conv_q'], spec_kwargs['kda_conv_k'], - spec_kwargs['kda_conv_v'], spec_kwargs['kda_qkg_cache'], - spec_kwargs['kda_v_cache'], spec_kwargs['kda_beta_cache'] + spec_kwargs['kda_conv_v'], spec_kwargs['kda_k_cache'], + spec_kwargs['kda_g_cache'], spec_kwargs['kda_v_cache'], + spec_kwargs['kda_beta_cache'] ] spec_path_label = "kda-replay" elif self._use_replay_state_update: @@ -896,7 +919,7 @@ def seed_kda_replay_caches_for_disagg_gen(self, device=cache.device) seeded[:, :, :, :committed] = section.to(cache.dtype) cache.index_copy_(1, idx, seeded) - for buf in (self.mamba_cache.kda_qkg_cache, + for buf in (self.mamba_cache.kda_k_cache, self.mamba_cache.kda_g_cache, self.mamba_cache.kda_v_cache, self.mamba_cache.kda_beta_cache): buf.index_fill_(1, idx, 0) @@ -1127,7 +1150,8 @@ def _drop(tensor): kda_conv_q=_drop(self.mamba_cache.kda_conv_q), kda_conv_k=_drop(self.mamba_cache.kda_conv_k), kda_conv_v=_drop(self.mamba_cache.kda_conv_v), - kda_qkg_cache=_drop(self.mamba_cache.kda_qkg_cache), + kda_k_cache=_drop(self.mamba_cache.kda_k_cache), + kda_g_cache=_drop(self.mamba_cache.kda_g_cache), kda_v_cache=_drop(self.mamba_cache.kda_v_cache), kda_beta_cache=_drop(self.mamba_cache.kda_beta_cache), ) @@ -3470,7 +3494,8 @@ def _allocate_pool_replay_buffers( self.kda_conv_q = None self.kda_conv_k = None self.kda_conv_v = None - self.kda_qkg_cache = None + self.kda_k_cache = None + self.kda_g_cache = None self.kda_v_cache = None self.kda_beta_cache = None @@ -3490,12 +3515,14 @@ def _allocate_pool_replay_buffers( extended_window = committed_window + num_spec def allocate_dim_contiguous_conv_cache() -> torch.Tensor: + # raw bf16 projection inputs: bf16 is exact and halves the + # per-draft replay traffic of the verify kernel return torch.zeros( self.local_num_mamba_layers, cache_size, extended_window, section_dim, - dtype=torch.float32, + dtype=torch.bfloat16, device=device, ).transpose(-1, -2) @@ -3505,11 +3532,21 @@ def allocate_dim_contiguous_conv_cache() -> torch.Tensor: self.kda_conv_q = allocate_dim_contiguous_conv_cache() self.kda_conv_k = allocate_dim_contiguous_conv_cache() self.kda_conv_v = allocate_dim_contiguous_conv_cache() - self.kda_qkg_cache = torch.zeros( + # k / v / beta replay entries in bf16 (halves the per-draft replay + # traffic of the verify kernel); the gate cache stays fp32 because it + # feeds the cumulative-gate exponents of the replayed chunk. + self.kda_k_cache = torch.zeros( + self.local_num_mamba_layers, + cache_size, + num_spec, + section_dim, + dtype=torch.bfloat16, + device=device, + ) + self.kda_g_cache = torch.zeros( self.local_num_mamba_layers, cache_size, num_spec, - 3, section_dim, dtype=torch.float32, device=device, @@ -3519,7 +3556,7 @@ def allocate_dim_contiguous_conv_cache() -> torch.Tensor: cache_size, num_spec, section_dim, - dtype=torch.float32, + dtype=torch.bfloat16, device=device, ) self.kda_beta_cache = _allocate_kda_beta_cache( @@ -4287,8 +4324,10 @@ def _reset_context_mamba_slots(self, num_contexts: int) -> None: if num_contexts == 0 or not self.use_kda_replay_update: return assert self.prev_num_accepted_tokens is not None - context_slots = self.cuda_state_indices[:num_contexts].long() - self.prev_num_accepted_tokens[context_slots] = 0 + context_slots = _host_slots_to_device( + self._host_state_indices[:num_contexts].tolist(), + self.prev_num_accepted_tokens.device) + self.prev_num_accepted_tokens.index_fill_(0, context_slots, 0) def _reset_kda_replay_slots( self, @@ -4297,11 +4336,17 @@ def _reset_kda_replay_slots( """Reset pending KDA drafts for newly initialized logical slots.""" assert self.use_kda_replay_update assert self.prev_num_accepted_tokens is not None - slots = torch.as_tensor(slots, - dtype=torch.long, - device=self.prev_num_accepted_tokens.device) + device = self.prev_num_accepted_tokens.device + if isinstance(slots, torch.Tensor) and slots.device.type != "cpu": + slots = slots.to(device=device, dtype=torch.long, non_blocking=True) + else: + if isinstance(slots, torch.Tensor): + slots = slots.tolist() + if not slots: + return + slots = _host_slots_to_device(slots, device) if slots.numel() > 0: - self.prev_num_accepted_tokens[slots] = 0 + self.prev_num_accepted_tokens.index_fill_(0, slots, 0) def _kda_replay_state_buffers(self) -> List[Tuple[torch.Tensor, int]]: """The per-slot KDA replay state kept outside the V2 pools, as @@ -4309,8 +4354,8 @@ def _kda_replay_state_buffers(self) -> List[Tuple[torch.Tensor, int]]: assert self.prev_num_accepted_tokens is not None buffers = [(self.prev_num_accepted_tokens, 0)] for replay_buffer in (self.kda_conv_q, self.kda_conv_k, self.kda_conv_v, - self.kda_qkg_cache, self.kda_v_cache, - self.kda_beta_cache): + self.kda_k_cache, self.kda_g_cache, + self.kda_v_cache, self.kda_beta_cache): assert replay_buffer is not None buffers.append((replay_buffer, 1)) return buffers @@ -4326,12 +4371,9 @@ def _relocate_kda_replay_slots(self, old_slots: List[int], return device = self.prev_num_accepted_tokens.device - source_slots = torch.tensor([old for old, _ in moves], - dtype=torch.long, - device=device) - destination_slots = torch.tensor([new for _, new in moves], - dtype=torch.long, - device=device) + source_slots = _host_slots_to_device([old for old, _ in moves], device) + destination_slots = _host_slots_to_device([new for _, new in moves], + device) for buffer, dim in self._kda_replay_state_buffers(): buffer.index_copy_(dim, destination_slots, buffer.index_select(dim, source_slots)) @@ -4389,8 +4431,9 @@ def _commit_kda_pending_drafts(self, slot: int) -> None: """Commit the slot's pending accepted drafts and its conv window into its V2 SSM and conv states, in place, and clear its draft count. - Afterwards the V2 states hold exactly what the next verify would have - started from (``commit_kda_pending_drafts``). + Afterwards the V2 states hold what the next verify would have started + from: the conv window exactly, the SSM state to the verify kernel's + accuracy (``commit_kda_pending_drafts``). """ from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import \ commit_kda_pending_drafts @@ -4413,7 +4456,8 @@ def _commit_kda_pending_drafts(self, slot: int) -> None: self.kda_conv_q[layer_offset], self.kda_conv_k[layer_offset], self.kda_conv_v[layer_offset], - self.kda_qkg_cache[layer_offset], + self.kda_k_cache[layer_offset], + self.kda_g_cache[layer_offset], self.kda_v_cache[layer_offset], self.kda_beta_cache[layer_offset], slots, @@ -4602,10 +4646,8 @@ def seed_kda_replay_caches_for_disagg_gen(self, return assert self.prev_num_accepted_tokens is not None - device = self.prev_num_accepted_tokens.device - state_indices = torch.tensor(sorted(set(slots)), - dtype=torch.long, - device=device) + state_indices = _host_slots_to_device( + sorted(set(slots)), self.prev_num_accepted_tokens.device) self._seed_kda_replay_slots(state_indices) def _seed_kda_replay_slots(self, state_indices: torch.Tensor) -> None: @@ -4630,11 +4672,11 @@ def _seed_kda_replay_slots(self, state_indices: torch.Tensor) -> None: replay_layer[:, :, :committed_window].index_copy_( 0, state_indices, section.to(replay_layer.dtype)) - for replay_buffer in (self.kda_qkg_cache, self.kda_v_cache, - self.kda_beta_cache): + for replay_buffer in (self.kda_k_cache, self.kda_g_cache, + self.kda_v_cache, self.kda_beta_cache): assert replay_buffer is not None replay_buffer.index_fill_(1, state_indices, 0) - self.prev_num_accepted_tokens[state_indices] = 0 + self.prev_num_accepted_tokens.index_fill_(0, state_indices, 0) def mamba_layer_cache( self, layer_idx: int @@ -4648,7 +4690,8 @@ def mamba_layer_cache( assert self.kda_conv_q is not None assert self.kda_conv_k is not None assert self.kda_conv_v is not None - assert self.kda_qkg_cache is not None + assert self.kda_k_cache is not None + assert self.kda_g_cache is not None assert self.kda_v_cache is not None assert self.kda_beta_cache is not None spec_kwargs = { @@ -4656,7 +4699,8 @@ def mamba_layer_cache( "kda_conv_q": self.kda_conv_q[layer_offset], "kda_conv_k": self.kda_conv_k[layer_offset], "kda_conv_v": self.kda_conv_v[layer_offset], - "kda_qkg_cache": self.kda_qkg_cache[layer_offset], + "kda_k_cache": self.kda_k_cache[layer_offset], + "kda_g_cache": self.kda_g_cache[layer_offset], "kda_v_cache": self.kda_v_cache[layer_offset], "kda_beta_cache": self.kda_beta_cache[layer_offset], } @@ -4942,7 +4986,8 @@ def shutdown(self): self.kda_conv_q = None self.kda_conv_k = None self.kda_conv_v = None - self.kda_qkg_cache = None + self.kda_k_cache = None + self.kda_g_cache = None self.kda_v_cache = None self.kda_beta_cache = None self._kda_replay_suspended.clear() diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 3b7e55c8fe33..ebcf794bfad6 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -1534,9 +1534,9 @@ def _try_schedule_generation( def _try_allocate_generation(self, req: LlmRequest) -> bool: """Atomically admit one generation step in target and draft pools. - The draft pool mirrors the target pool (same requests, same block - boundaries, same ``tokens_per_block``), so a draft-side failure just - means this request does not fit. Roll the target growth back and report + The draft pool mirrors the target pool's requests and token counts (a + standalone drafter's pool may use its own ``tokens_per_block``), so a + draft-side failure just means this request does not fit. Roll the target growth back and report a plain allocation failure, letting the caller run the same evict / recompute-pause / self-suspend ladder it uses for target-pool pressure. diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index ea0ea7287599..08eb2c94d863 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3646,10 +3646,11 @@ class DSparkDecodingConfig(DecodingBaseConfig): "namespace). Ignored by the embedded DeepSeek-V4-Pro draft. Independent " "of the backend that builds the drafter's own attention modules.\n\n" "AUTO resolves per drafter family and is right unless you are pinning a " - "kernel: a GQA backbone degrades when its kernel is missing, an MLA one " - "raises. TRTLLM needs FlashInfer and SM100/SM103. CUTEDSL is MLA-only " - "and needs a cute-dsl MLA decode taking per-token kv_bounds that is not " - "upstream yet. Which kernel each name selects: MLADSparkForCausalLM.") + "kernel: TRTLLM for a GQA backbone, which degrades when its kernel is " + "missing, and CUTEDSL for an MLA one, which raises. TRTLLM needs " + "FlashInfer and SM100/SM103. CUTEDSL is MLA-only: one cute-dsl MLA " + "decode pass over context and block, on SM100/SM103/SM107. Which " + "kernel each name selects: MLADSparkForCausalLM.") @model_validator(mode="after") def set_max_total_draft_tokens(self): diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 1211e43a6390..ce332dfae73a 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -975,7 +975,7 @@ }, "attention_backend": { "default": "AUTO", - "description": "Block-decode attention backend for a standalone DSpark drafter (one shipped as its own checkpoint, not inside the target's mtp.* namespace). Ignored by the embedded DeepSeek-V4-Pro draft. Independent of the backend that builds the drafter's own attention modules.\n\nAUTO resolves per drafter family and is right unless you are pinning a kernel: a GQA backbone degrades when its kernel is missing, an MLA one raises. TRTLLM needs FlashInfer and SM100/SM103. CUTEDSL is MLA-only and needs a cute-dsl MLA decode taking per-token kv_bounds that is not upstream yet. Which kernel each name selects: MLADSparkForCausalLM.", + "description": "Block-decode attention backend for a standalone DSpark drafter (one shipped as its own checkpoint, not inside the target's mtp.* namespace). Ignored by the embedded DeepSeek-V4-Pro draft. Independent of the backend that builds the drafter's own attention modules.\n\nAUTO resolves per drafter family and is right unless you are pinning a kernel: TRTLLM for a GQA backbone, which degrades when its kernel is missing, and CUTEDSL for an MLA one, which raises. TRTLLM needs FlashInfer and SM100/SM103. CUTEDSL is MLA-only: one cute-dsl MLA decode pass over context and block, on SM100/SM103/SM107. Which kernel each name selects: MLADSparkForCausalLM.", "enum": [ "AUTO", "VANILLA", @@ -2798,7 +2798,7 @@ }, "enable_kv_pool_rebalance": { "default": false, - "description": "Opt in to the KVCacheManagerV2 auto-tuner (``adjust()``) for rebalancing pool-group ratios between iterations. When True the PyExecutor calls ``adjust()`` opportunistically; the auto-tuner itself remains gated by V2's internal 2000-sample / 120s cooldown. When False (default) the rebalance hook is skipped entirely and pool ratios remain at their warmup-derived values. Beta: enable at your own risk. Only used when using KV cache manager v2 (experimental). This option is incompatible with dtype='fp8_ds_mla'.", + "description": "Opt in to the KVCacheManagerV2 auto-tuner (``adjust()``) for rebalancing pool-group ratios between iterations. When True the PyExecutor calls ``adjust()`` opportunistically; the auto-tuner itself remains gated by the ``kv_pool_rebalance_config`` thresholds (2000 samples / 120s by default). When False (default) the rebalance hook is skipped entirely and pool ratios remain at their warmup-derived values. Beta: enable at your own risk. Only used when using KV cache manager v2 (experimental). This option is incompatible with dtype='fp8_ds_mla'.", "status": "prototype", "title": "Enable Kv Pool Rebalance", "type": "boolean" @@ -2897,6 +2897,11 @@ "description": "Streaming (push-based) KV cache event publishing (KV cache manager V2 only). When set, each rank publishes its own events directly (e.g. over ZeroMQ) instead of the buffered event_buffer_max_size gather/poll path.", "status": "prototype" }, + "kv_pool_rebalance_config": { + "$ref": "#/$defs/KvPoolRebalanceConfig", + "description": "Tuning for the KVCacheManagerV2 pool rebalance auto-tuner. Only used when ``enable_kv_pool_rebalance`` is True.", + "status": "prototype" + }, "mamba_ssm_cache_dtype": { "default": "auto", "description": "The data type to use for the Mamba SSM cache. If set to 'auto', the data type will be inferred from the model config.", @@ -3159,6 +3164,63 @@ "title": "KvCacheConnectorConfig", "type": "object" }, + "KvPoolRebalanceConfig": { + "additionalProperties": false, + "description": "Tuning for the KV cache manager v2 pool rebalance auto-tuner.\n\nOnly used when ``kv_cache_config.enable_kv_pool_rebalance`` is True.", + "properties": { + "check_interval": { + "default": 10, + "description": "Number of PyExecutor iterations between checks of whether a pool rebalance should be attempted.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Check Interval", + "type": "integer" + }, + "cooldown_secs": { + "default": 120.0, + "description": "Minimum time, in seconds, between successive rebalance adjustments.", + "minimum": 0, + "status": "prototype", + "title": "Cooldown Secs", + "type": "number" + }, + "min_sampled_kv_caches": { + "default": 2000, + "description": "Minimum number of sampled (closed) KvCaches observed before the auto-tuner will consider making an adjustment.", + "minimum": 0, + "status": "prototype", + "title": "Min Sampled Kv Caches", + "type": "integer" + }, + "moving_average_decay": { + "default": 0.9999, + "description": "Decay factor for the exponential moving averages (reused length, capacity, history length) used to compute target pool ratios. Higher values react more slowly to change.", + "exclusiveMaximum": 1, + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Moving Average Decay", + "type": "number" + }, + "ratio_tolerance": { + "default": 0.25, + "description": "How far a pool group's current ratio may drift from its target before a rebalance is triggered, as a fraction of the smaller of the two. The comparison is symmetric, so 0.25 means a rebalance is triggered once one of the two exceeds the other by more than 25%.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Ratio Tolerance", + "type": "number" + }, + "target_ratio_update_interval": { + "default": 100, + "description": "Number of newly sampled KvCaches between recomputations of the target pool-group ratios.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Target Ratio Update Interval", + "type": "integer" + } + }, + "title": "KvPoolRebalanceConfig", + "type": "object" + }, "LayerwiseBenchmarksConfig": { "additionalProperties": false, "description": "Configuration for layer-wise benchmarks calibration.", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 15599ad03a1f..e34b7f2adaa7 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -975,7 +975,7 @@ }, "attention_backend": { "default": "AUTO", - "description": "Block-decode attention backend for a standalone DSpark drafter (one shipped as its own checkpoint, not inside the target's mtp.* namespace). Ignored by the embedded DeepSeek-V4-Pro draft. Independent of the backend that builds the drafter's own attention modules.\n\nAUTO resolves per drafter family and is right unless you are pinning a kernel: a GQA backbone degrades when its kernel is missing, an MLA one raises. TRTLLM needs FlashInfer and SM100/SM103. CUTEDSL is MLA-only and needs a cute-dsl MLA decode taking per-token kv_bounds that is not upstream yet. Which kernel each name selects: MLADSparkForCausalLM.", + "description": "Block-decode attention backend for a standalone DSpark drafter (one shipped as its own checkpoint, not inside the target's mtp.* namespace). Ignored by the embedded DeepSeek-V4-Pro draft. Independent of the backend that builds the drafter's own attention modules.\n\nAUTO resolves per drafter family and is right unless you are pinning a kernel: TRTLLM for a GQA backbone, which degrades when its kernel is missing, and CUTEDSL for an MLA one, which raises. TRTLLM needs FlashInfer and SM100/SM103. CUTEDSL is MLA-only: one cute-dsl MLA decode pass over context and block, on SM100/SM103/SM107. Which kernel each name selects: MLADSparkForCausalLM.", "enum": [ "AUTO", "VANILLA", @@ -4348,7 +4348,7 @@ }, "enable_kv_pool_rebalance": { "default": false, - "description": "Opt in to the KVCacheManagerV2 auto-tuner (``adjust()``) for rebalancing pool-group ratios between iterations. When True the PyExecutor calls ``adjust()`` opportunistically; the auto-tuner itself remains gated by V2's internal 2000-sample / 120s cooldown. When False (default) the rebalance hook is skipped entirely and pool ratios remain at their warmup-derived values. Beta: enable at your own risk. Only used when using KV cache manager v2 (experimental). This option is incompatible with dtype='fp8_ds_mla'.", + "description": "Opt in to the KVCacheManagerV2 auto-tuner (``adjust()``) for rebalancing pool-group ratios between iterations. When True the PyExecutor calls ``adjust()`` opportunistically; the auto-tuner itself remains gated by the ``kv_pool_rebalance_config`` thresholds (2000 samples / 120s by default). When False (default) the rebalance hook is skipped entirely and pool ratios remain at their warmup-derived values. Beta: enable at your own risk. Only used when using KV cache manager v2 (experimental). This option is incompatible with dtype='fp8_ds_mla'.", "status": "prototype", "title": "Enable Kv Pool Rebalance", "type": "boolean" @@ -4447,6 +4447,11 @@ "description": "Streaming (push-based) KV cache event publishing (KV cache manager V2 only). When set, each rank publishes its own events directly (e.g. over ZeroMQ) instead of the buffered event_buffer_max_size gather/poll path.", "status": "prototype" }, + "kv_pool_rebalance_config": { + "$ref": "#/$defs/KvPoolRebalanceConfig", + "description": "Tuning for the KVCacheManagerV2 pool rebalance auto-tuner. Only used when ``enable_kv_pool_rebalance`` is True.", + "status": "prototype" + }, "mamba_ssm_cache_dtype": { "default": "auto", "description": "The data type to use for the Mamba SSM cache. If set to 'auto', the data type will be inferred from the model config.", @@ -4709,6 +4714,63 @@ "title": "KvCacheConnectorConfig", "type": "object" }, + "KvPoolRebalanceConfig": { + "additionalProperties": false, + "description": "Tuning for the KV cache manager v2 pool rebalance auto-tuner.\n\nOnly used when ``kv_cache_config.enable_kv_pool_rebalance`` is True.", + "properties": { + "check_interval": { + "default": 10, + "description": "Number of PyExecutor iterations between checks of whether a pool rebalance should be attempted.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Check Interval", + "type": "integer" + }, + "cooldown_secs": { + "default": 120.0, + "description": "Minimum time, in seconds, between successive rebalance adjustments.", + "minimum": 0, + "status": "prototype", + "title": "Cooldown Secs", + "type": "number" + }, + "min_sampled_kv_caches": { + "default": 2000, + "description": "Minimum number of sampled (closed) KvCaches observed before the auto-tuner will consider making an adjustment.", + "minimum": 0, + "status": "prototype", + "title": "Min Sampled Kv Caches", + "type": "integer" + }, + "moving_average_decay": { + "default": 0.9999, + "description": "Decay factor for the exponential moving averages (reused length, capacity, history length) used to compute target pool ratios. Higher values react more slowly to change.", + "exclusiveMaximum": 1, + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Moving Average Decay", + "type": "number" + }, + "ratio_tolerance": { + "default": 0.25, + "description": "How far a pool group's current ratio may drift from its target before a rebalance is triggered, as a fraction of the smaller of the two. The comparison is symmetric, so 0.25 means a rebalance is triggered once one of the two exceeds the other by more than 25%.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Ratio Tolerance", + "type": "number" + }, + "target_ratio_update_interval": { + "default": 100, + "description": "Number of newly sampled KvCaches between recomputations of the target pool-group ratios.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Target Ratio Update Interval", + "type": "integer" + } + }, + "title": "KvPoolRebalanceConfig", + "type": "object" + }, "LayerwiseBenchmarksConfig": { "additionalProperties": false, "description": "Configuration for layer-wise benchmarks calibration.", diff --git a/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json index 6cab6b77e810..cffa1974bcdc 100644 --- a/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json @@ -397,6 +397,24 @@ "title": "Ring Size", "type": "integer" }, + "tp_layout": { + "anyOf": [ + { + "enum": [ + "replicated", + "token_sharded" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Activation layout between TP projections: 'replicated' (all-reduce TP; the default when unset) or 'token_sharded' (each TP rank holds 1/tp_size of the tokens; each all-reduce becomes a reduce-scatter plus an all-gather). 'token_sharded' requires tp_size > 1 and no ulysses/ring/attn2d.", + "status": "prototype", + "title": "Tp Layout" + }, "tp_size": { "default": 1, "description": "Tensor parallel group size. Heads are sharded across tp_size GPUs.", diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py index 438e181797f7..91288f6fef6d 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py @@ -972,11 +972,12 @@ def create_kv_cache_manager( # Multi-token verify reads per-slot replay caches when the fused # kernel is available, else the legacy per-step buffers. from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import ( - is_kda_mtp_verify_available, + kda_mtp_replay_num_spec, ) - if is_kda_mtp_verify_available(): - kimi_extra_kwargs["kda_replay_num_spec"] = spec_config.tokens_per_gen_step - 1 + kda_replay_num_spec = kda_mtp_replay_num_spec(spec_config.tokens_per_gen_step - 1) + if kda_replay_num_spec is not None: + kimi_extra_kwargs["kda_replay_num_spec"] = kda_replay_num_spec kv_cache_manager = kv_cache_manager_cls( # mamba (KDA) cache parameters mamba_params.state_size, diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 60afaa3d2be2..206d12745133 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -319,6 +319,11 @@ moonshotai/Kimi-K3: - quant_algo: MIXED_PRECISION spec_dec_algo: DSpark accuracy: 96.5 + # Same requant, non-speculative, on the aggregate DEP8 recipe with + # TRTLLM-Gen MoE (TestKimiK3::test_gsm8k_nvfp4_dep8). Measured 96.361 on all + # 1,319 samples on SM107; 96.0 keeps margin for run-to-run variation. + - quant_algo: MIXED_PRECISION + accuracy: 96.0 nvidia/Nemotron-MOE: - accuracy: 88.249 - quant_algo: FP8 diff --git a/tests/integration/defs/accuracy/test_kimi3.py b/tests/integration/defs/accuracy/test_kimi3.py index ad94695f3563..48a1d05030c6 100644 --- a/tests/integration/defs/accuracy/test_kimi3.py +++ b/tests/integration/defs/accuracy/test_kimi3.py @@ -207,6 +207,65 @@ def test_gpqa_diamond_w4a16_mxfp4(self) -> None: extra_evaluator_kwargs=dict(apply_chat_template=True), ) + # The NVFP4 requant: the Blackwell/Rubin deployment precision, served + # here by the aggregate DEP8 recipe rather than the DEP16 one above. + NVFP4_MODEL_PATH = f"{llm_models_root()}/Kimi-K3-NVFP4" + + @skip_pre_blackwell + @pytest.mark.skip_less_mpi_world_size(8) + # Attention DP replicates every non-expert weight on each rank, so DEP8 + # holds far more weight per GPU than TEP8 does; only GB300-class memory + # leaves room for the KV cache. + @pytest.mark.skip_less_device_memory(250000) + def test_gsm8k_nvfp4_dep8(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Full GSM8K on the NVFP4 checkpoint with the aggregate DEP8 recipe. + + TP=8 / EP=8 with attention DP, TRTLLM-Gen MoE (fused NVFP4 SiTu + cubins), CUDA graphs and chunked prefill. + """ + # Attention DP also replicates the shared-expert and latent MoE MLPs. + # Kept in BF16 (the default) they cost about 15.7 GiB more per rank + # and leave DEP8 no KV-cache room, so read them as FP8 block-scale + # weights; the GSM8K reference was measured with this cast. The + # variable is read in the worker ranks while the model is built, so it + # travels as an LLM env override. LLM also sets it in this process and + # does not restore it, hence the monkeypatch teardown. + fp8_moe_mlp_env = {"KIMI_K3_FP8_WEIGHT_READ_MOE_MLP": "1"} + for key, value in fp8_moe_mlp_env.items(): + monkeypatch.setenv(key, value) + with LLM( + self.NVFP4_MODEL_PATH, + env_overrides=fp8_moe_mlp_env, + tensor_parallel_size=8, + moe_expert_parallel_size=8, + enable_attention_dp=True, + max_batch_size=8, + max_num_tokens=8192, + max_seq_len=8192, + trust_remote_code=True, + disable_overlap_scheduler=False, + enable_chunked_prefill=True, + cuda_graph_config=CudaGraphConfig(enable_padding=True, max_batch_size=8), + moe_config=MoeConfig( + backend="TRTLLM", + max_num_tokens=131072, + use_low_precision_moe_combine=True, + ), + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + # Attention DP replicates KV per rank. A smaller fraction lets + # clamp_max_seq_len_for_mem drop max_seq_len below a GSM8K + # prompt, which then overruns max_blocks_per_seq mid-run. + free_gpu_memory_fraction=0.4, + tokens_per_block=64, + ), + ) as llm: + # The requant ships hf_quant_config.json, so the matcher sees + # MIXED_PRECISION -- only the experts are 4-bit. + assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + def _assert_checkpoint_routing(self) -> None: config = load_pretrained_config(self.MODEL_PATH, trust_remote_code=True) assert isinstance(config, KimiK3Config) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 101ebbd27a06..27c832b0fee4 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -63,6 +63,7 @@ accuracy/test_glm53_flash.py::TestGLM53FlashFP8::test_attention_dp_mtp_and_fp8_k accuracy/test_glm53_flash.py::TestGLM53FlashFP8::test_block_reuse[tp_size=4-ep_size=4] accuracy/test_glm53_flash.py::TestGLM53FlashFP8::test_mmmu[tp_size=4-ep_size=4] accuracy/test_glm53_flash.py::TestGLM53FlashFP8::test_mtp[tp_size=4-ep_size=4] +accuracy/test_kimi3.py::TestKimiK3::test_gsm8k_nvfp4_dep8 accuracy/test_kimi3.py::TestKimiK3DSpark::test_gsm8k_tep8 accuracy/test_kv_pool_rebalance_accuracy.py::TestKvPoolRebalanceAccuracy::test_rebalance_matches_baseline[no_overlap] accuracy/test_kv_pool_rebalance_accuracy.py::TestKvPoolRebalanceAccuracy::test_rebalance_matches_baseline[overlap] diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py index 7fb79f29925e..3bc87dfe6f48 100644 --- a/tests/unittest/_torch/attention/test_attention_mla.py +++ b/tests/unittest/_torch/attention/test_attention_mla.py @@ -734,18 +734,14 @@ def run_once(batch_size: int = 64) -> None: kernel_keys = list(CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache) assert kernel_keys, "tuning-mode pass compiled no CuTe DSL MLA kernels" # Kernel-cache key layout: unique_id + (out_dtype, mma_qk, mma_pv, - # split_kv, is_persistent, has_kv_bounds). The trailing has_kv_bounds - # separates the helix per-token-bounds kernel variant from the plain one - # and is not a tuner tactic, so index the two tactic elements from the - # end past it. Both tactic elements chosen by the tuner must have been - # exercised during profiling. - persistent_variants = {key[-2] for key in kernel_keys} + # use_workspace, is_persistent, has_kv_bounds, is_var_split_kv). split_kv + # is a runtime argument and not part of the key. The runner offers only + # the non-persistent schedule, so no persistent variant may be compiled. + persistent_variants = {key[-3] for key in kernel_keys} assert persistent_variants == { - True, False - }, (f"expected both is_persistent tactic candidates to be profiled, " + False + }, (f"expected only the non-persistent tactic to be profiled, " f"got {persistent_variants}") - split_kv_variants = {key[-3] for key in kernel_keys} - assert split_kv_variants, "no split_kv tactic variant was profiled" # Serving-mode pass: tuned tactics must be reused as-is -- any new # kernel_cache entry means a runtime cute.compile happened post-tuning. @@ -797,19 +793,30 @@ def spying_default_tactic(self, batch_size: int): # The assertions above can hold even for an unbucketed fallback when this # GPU's occupancy makes split_kv(65) == split_kv(64), so also check the # bucketing itself with the occupancy ceiling pinned to a value where the - # raw batch and its bucket disagree: 256 // 65 // 2 == 1, while bucket 64 - # gives 256 // 64 // 2 == 2. + # raw batch and its bucket disagree: with 91 blocks, batch 65 gives + # split_kv 4 while bucket 64 gives 8. + pinned_max_active_blocks = 91 + num_q_tiles = runner.get_num_q_tiles(runner.num_heads, runner.seq_len_q) + bucket_split_kv = runner.get_default_split_kv(64, num_q_tiles, + pinned_max_active_blocks) + raw_split_kv = runner.get_default_split_kv(65, num_q_tiles, + pinned_max_active_blocks) + assert bucket_split_kv != raw_split_kv, ( + f"pinned max_active_blocks {pinned_max_active_blocks} no longer " + f"separates batch 65 from bucket 64 (split_kv {raw_split_kv} vs " + f"{bucket_split_kv}); pick another value") with mock.patch.object(CuteDSLNVMlaDecodeBlackwellRunner, "_cute_dsl_max_active_blocks", - 256, + pinned_max_active_blocks, create=True): pinned_tactic = runner.default_tactic(65) assert pinned_tactic == runner.default_tactic(64), ( f"default_tactic no longer buckets the batch: got {pinned_tactic} " f"for batch 65 vs {runner.default_tactic(64)} for batch 64") - assert pinned_tactic[2] == 2, ( - f"expected bucket-64 split_kv 2 with max_active_blocks pinned to " - f"256, got {pinned_tactic[2]}") + assert pinned_tactic[2] == bucket_split_kv, ( + f"expected bucket-64 split_kv {bucket_split_kv} with " + f"max_active_blocks pinned to {pinned_max_active_blocks}, " + f"got {pinned_tactic[2]}") def _run_test_for_backend(backend_name, diff --git a/tests/unittest/_torch/attention/test_cute_dsl_mla.py b/tests/unittest/_torch/attention/test_cute_dsl_mla.py new file mode 100644 index 000000000000..8e0b2cfaf916 --- /dev/null +++ b/tests/unittest/_torch/attention/test_cute_dsl_mla.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host-side contract tests for the CuTe DSL MLA FMHA library (FP8 scales, split-KV sizing).""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.fmha.cute_dsl_mla import CuteDslMlaFmha +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + + +def test_fp8_mla_scales_remain_tensors() -> None: + """FP8 scales reach the kernel as the same device tensors, never as host floats.""" + q = torch.empty(1) + bmm1_scale = torch.tensor([0.25, 0.25 * 1.442695], dtype=torch.float32) + bmm2_scale = torch.tensor([4.0], dtype=torch.float32) + + actual_bmm1_scale, actual_bmm2_scale = CuteDslMlaFmha._get_fp8_scales( + q, + SimpleNamespace(mla_bmm1_scale=bmm1_scale, mla_bmm2_scale=bmm2_scale), + ) + + assert actual_bmm1_scale is bmm1_scale + assert actual_bmm2_scale is bmm2_scale + + +@pytest.mark.parametrize( + "bmm1_scale, bmm2_scale, match", + [ + (None, torch.tensor([1.0]), "requires mla_bmm1_scale and mla_bmm2_scale"), + (torch.tensor([1.0], dtype=torch.float16), torch.tensor([1.0]), "float32"), + (torch.tensor([1.0]), torch.empty(0), "float32"), + ], +) +def test_fp8_mla_scales_are_validated(bmm1_scale, bmm2_scale, match) -> None: + with pytest.raises(RuntimeError, match=match): + CuteDslMlaFmha._get_fp8_scales( + torch.empty(1), + SimpleNamespace(mla_bmm1_scale=bmm1_scale, mla_bmm2_scale=bmm2_scale), + ) + + +@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="nvidia-cutlass-dsl is not installed") +def test_fp8_mla_custom_op_accepts_tensor_scales() -> None: + # Importing the module registers the CuTe DSL custom ops. + import tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops # noqa: F401 + + schema = torch.ops.trtllm.cute_dsl_mla_decode_fp8_blackwell.default._schema + argument_types = {argument.name: str(argument.type) for argument in schema.arguments} + + assert argument_types["softmax_scale"] == "Tensor" + assert argument_types["output_scale"] == "Tensor" + + +def _mla_decode_runner_cls(): + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLNVMlaDecodeBlackwellRunner + + return CuteDSLNVMlaDecodeBlackwellRunner + + +@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="nvidia-cutlass-dsl is not installed") +def test_mla_decode_split_kv_is_a_bounded_power_of_two() -> None: + runner_cls = _mla_decode_runner_cls() + max_active_blocks = 296 + previous = None + for batch_size in (1, 2, 4, 8, 16, 32, 64, 128, 256): + split_kv = runner_cls.get_default_split_kv(batch_size, 1, max_active_blocks) + assert 1 <= split_kv <= runner_cls._MAX_SPLIT_KV + assert split_kv & (split_kv - 1) == 0 + # More requests already fill more CTAs, so the split never grows. + if previous is not None: + assert split_kv <= previous + previous = split_kv + # Persistent round-robin is never a candidate: skewed decode KV balances + # better on the hardware scheduler. + assert runner_cls.get_is_persistent_candidates() == [False] + + +@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="nvidia-cutlass-dsl is not installed") +@pytest.mark.parametrize("num_heads", [8, 128]) +def test_mla_decode_split_kv_workspace_rows_cover_every_bucket(num_heads: int) -> None: + """The CUDA-graph workspace is sized once, so its row bound must cover every + batch bucket and query length below the configured maxima, and must not + shrink when either maximum grows.""" + runner_cls = _mla_decode_runner_cls() + max_active_blocks = 296 + max_seq_len_q, max_batch_size = 4, 64 + + rows = runner_cls._max_split_kv_rows( + num_heads, max_seq_len_q, max_batch_size, max_active_blocks + ) + for seq_len_q in range(1, max_seq_len_q + 1): + num_q_tiles = runner_cls.get_num_q_tiles(num_heads, seq_len_q) + assert 1 <= num_q_tiles <= seq_len_q * max(1, num_heads // runner_cls._MMA_TILE_M) + bucket = 1 + while bucket <= max_batch_size: + split_kv = runner_cls.get_default_split_kv(bucket, num_q_tiles, max_active_blocks) + if split_kv > 1: + assert bucket * num_q_tiles * split_kv <= rows + bucket *= 2 + + assert ( + runner_cls._max_split_kv_rows( + num_heads, max_seq_len_q + 1, max_batch_size, max_active_blocks + ) + >= rows + ) + assert ( + runner_cls._max_split_kv_rows( + num_heads, max_seq_len_q, 2 * max_batch_size, max_active_blocks + ) + >= rows + ) + + +class _TrtllmGenStub: + """Stands in for FlashInfer TRTLLM-Gen, the next MLA library on SM100/SM103.""" + + def is_supported(self, *args, **kwargs) -> bool: + return True + + +def _select_mla_generation_fmha( + monkeypatch, *, sm: int, num_heads: int, seq_len_q: int, fp8_kv: bool, helix: bool +): + """Run FmhaManager._select_mla over [CuteDslMlaFmha, TRTLLM-Gen stub] on CPU.""" + from tensorrt_llm._torch.attention.backends.fmha import cute_dsl_mla + from tensorrt_llm._torch.attention.backends.fmha.manager import FmhaManager + from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, + ) + from tensorrt_llm._torch.autotuner import AutoTuner + + monkeypatch.setattr(cute_dsl_mla, "get_sm_version", lambda: sm) + # AutoTuner.get() builds the singleton from CUDA device properties. + monkeypatch.setattr( + AutoTuner, "get", staticmethod(lambda: SimpleNamespace(is_tuning_mode=False)) + ) + # The kernel's can_implement needs nvidia-cutlass-dsl; this test covers only + # the host-side routing in front of it. + monkeypatch.setattr( + CuteDslMlaFmha, "_kernel_can_implement", staticmethod(lambda *args, **kwargs: (True, "")) + ) + + batch_size = 4 + num_tokens = batch_size * seq_len_q + attn = SimpleNamespace( + num_heads=num_heads, + has_fp8_kv_cache=fp8_kv, + kv_lora_rank=512, + qk_rope_head_dim=64, + layer_idx=0, + ) + meta = SimpleNamespace( + num_contexts=0, + num_generations=batch_size, + beam_width=1, + is_spec_dec_tree=False, + is_spec_dec_dynamic_tree=False, + helix_position_offsets=torch.zeros(num_tokens, dtype=torch.int32) if helix else None, + _helix_spec_tokens_valid=helix, + kv_cache_manager=SimpleNamespace( + get_buffers=lambda layer_idx: torch.empty(0, dtype=torch.bfloat16) + ), + tokens_per_block=64, + ) + fwd = AttentionForwardArgs( + attention_input_type=AttentionInputType.generation_only, + output=torch.empty(num_tokens, num_heads, 512, dtype=torch.bfloat16), + softmax_stats_tensor=( + torch.empty(num_tokens, num_heads, 2, dtype=torch.float32) if helix else None + ), + ) + q = torch.empty(num_tokens, num_heads, 576, dtype=torch.bfloat16) + + cute = CuteDslMlaFmha.__new__(CuteDslMlaFmha) + cute._attn_ref = lambda: attn + trtllm_gen = _TrtllmGenStub() + manager = FmhaManager.__new__(FmhaManager) + manager.fmha_libs = [cute, trtllm_gen] + selected = manager._select_mla(q, None, None, meta, fwd) + return "cute_dsl_mla" if selected is cute else "trtllm_gen" + + +@pytest.mark.parametrize( + "sm, num_heads, seq_len_q, fp8_kv, helix, expected", + [ + # Non-helix FP8 verify group at H=96 (DSpark under attention DP): the + # TRTLLM-Gen MLA decode rejects 64 < H < 128, so it must not be chosen. + (103, 96, 8, True, False, "cute_dsl_mla"), + (100, 96, 8, True, False, "cute_dsl_mla"), + (103, 96, 2, True, False, "cute_dsl_mla"), + # SM107 has no TRTLLM-Gen FMHA; the request keeps falling through. + (107, 96, 8, True, False, "trtllm_gen"), + # Helix H=96 verify groups are unchanged. + (103, 96, 4, True, True, "cute_dsl_mla"), + (107, 96, 4, True, True, "cute_dsl_mla"), + # Single-token H=96 is unchanged (measured perf-table entry). + (103, 96, 1, True, False, "cute_dsl_mla"), + # Head counts TRTLLM-Gen serves keep the perf gate. + (103, 128, 8, True, False, "trtllm_gen"), + (103, 64, 8, True, False, "trtllm_gen"), + # BF16 KV is unchanged (H=96 is already admitted for any seq_len_q). + (103, 96, 8, False, False, "cute_dsl_mla"), + ], +) +def test_mla_generation_routes_wide_head_fp8_verify_away_from_trtllm_gen( + monkeypatch, sm, num_heads, seq_len_q, fp8_kv, helix, expected +) -> None: + assert ( + _select_mla_generation_fmha( + monkeypatch, + sm=sm, + num_heads=num_heads, + seq_len_q=seq_len_q, + fp8_kv=fp8_kv, + helix=helix, + ) + == expected + ) diff --git a/tests/unittest/_torch/attention/test_cute_dsl_mla_helix.py b/tests/unittest/_torch/attention/test_cute_dsl_mla_helix.py index 4345eb3110de..b8884fc40a45 100644 --- a/tests/unittest/_torch/attention/test_cute_dsl_mla_helix.py +++ b/tests/unittest/_torch/attention/test_cute_dsl_mla_helix.py @@ -19,6 +19,24 @@ ) +def _scale_inputs(input_dtype: torch.dtype, softmax_scale: float, device) -> list: + """FP8 runners read both scales from device tensors at inputs[10:12].""" + if input_dtype != torch.float8_e4m3fn: + return [] + return [ + None, # kv_bounds + torch.tensor([softmax_scale], dtype=torch.float32, device=device), + torch.tensor([1.0], dtype=torch.float32, device=device), + ] + + +def _scale_kwargs(input_dtype: torch.dtype, softmax_scale: float) -> dict: + """FP16/BF16 runners keep their scales as host scalars.""" + if input_dtype == torch.float8_e4m3fn: + return {} + return {"softmax_scale": softmax_scale, "output_scale": 1.0} + + @pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float8_e4m3fn]) @pytest.mark.parametrize("split_kv", [1, 4]) def test_cute_dsl_mla_helix_stats_and_empty_local_kv( @@ -98,21 +116,21 @@ def test_cute_dsl_mla_helix_stats_and_empty_local_kv( num_heads, seq_len_q, latent_dim, batch_size, cutlass.Float32 ) workspace = torch.empty(workspace_size, device=device, dtype=torch.uint8) + inputs = [ + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + cache_seqs, + output, + workspace, + softmax_stats, + ] runner.forward( - [ - q_latent, - q_rope, - c_latent, - c_rope, - page_table, - cache_seqs, - output, - workspace, - softmax_stats, - ], + inputs + _scale_inputs(input_dtype, softmax_scale, device), tactic=((128, 128), (128, 256), split_kv, False), - softmax_scale=softmax_scale, - output_scale=1.0, + **_scale_kwargs(input_dtype, softmax_scale), ) # The existing Helix sanitizer makes a rank with no local pages the identity @@ -146,3 +164,276 @@ def test_cute_dsl_mla_helix_stats_and_empty_local_kv( actual_log_partition = softmax_stats[1, :, 0] + softmax_stats[1, :, 1].log() expected_log_partition = torch.logsumexp(scores, dim=-1) torch.testing.assert_close(actual_log_partition, expected_log_partition, rtol=1e-5, atol=1e-5) + + +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("num_heads", [96, 12]) +def test_cute_dsl_mla_helix_per_token_bounds(input_dtype: torch.dtype, num_heads: int) -> None: + """Verify-group semantics: kv_bounds replaces the causal bound per token. + + num_heads=96 is the helix production shape (fold ratio 1); num_heads=12 + folds seq_len_q into the M tile (fold ratio 4), exercising the + folded-grid token indexing in both the masking and the reduction + stats gating. + """ + import cutlass + + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLNVMlaDecodeBlackwellRunner + + torch.manual_seed(23) + device = torch.device("cuda") + batch_size, seq_len_q = 2, 4 + latent_dim, rope_dim, page_size, kv_len = 512, 64, 64, 256 + blocks_per_sequence = kv_len // page_size + num_pages = batch_size * blocks_per_sequence + softmax_scale = 1.0 / math.sqrt(latent_dim + rope_dim) + + q_storage = ( + torch.randn( + batch_size, + seq_len_q, + num_heads, + latent_dim + rope_dim, + device=device, + dtype=torch.bfloat16, + ) + * 0.1 + ).to(input_dtype) + cache_storage = ( + torch.randn( + num_pages, + page_size, + latent_dim + rope_dim, + device=device, + dtype=torch.bfloat16, + ) + * 0.1 + ).to(input_dtype) + q_latent = q_storage[..., :latent_dim].permute(2, 3, 1, 0) + q_rope = q_storage[..., latent_dim:].permute(2, 3, 1, 0) + c_latent = cache_storage[..., :latent_dim].permute(1, 2, 0) + c_rope = cache_storage[..., latent_dim:].permute(1, 2, 0) + page_table = ( + torch.arange(num_pages, dtype=torch.int32, device=device) + .view(batch_size, blocks_per_sequence) + .transpose(0, 1) + ) + cache_seqs = torch.tensor([kv_len, kv_len], dtype=torch.int32, device=device) + # Per-token rank-local bounds: sequence 0 permutes non-causal values in + # [K - S_q, K] (the table must win over the causal formula); sequence 1's + # leading token owns no local KV (stats must gate it out per token). + kv_bounds = torch.tensor( + [253, 252, 255, 254, 0, 256, 253, 256], dtype=torch.int32, device=device + ) + output_storage = torch.empty( + batch_size, + seq_len_q, + num_heads, + latent_dim, + device=device, + dtype=torch.bfloat16, + ) + output = output_storage.permute(2, 3, 1, 0) + softmax_stats = torch.empty( + batch_size * seq_len_q, + num_heads, + 2, + device=device, + dtype=torch.float32, + ) + + cutlass_dtype = cutlass.Float8E4M3FN if input_dtype == torch.float8_e4m3fn else cutlass.BFloat16 + runner = CuteDSLNVMlaDecodeBlackwellRunner( + in_dtype=cutlass_dtype, + num_heads=num_heads, + seq_len_q=seq_len_q, + page_size=page_size, + max_batch_size=batch_size, + emit_softmax_stats=True, + ) + workspace_size = runner.get_max_padded_workspace_size( + num_heads, seq_len_q, latent_dim, batch_size, cutlass.Float32 + ) + workspace = torch.empty(workspace_size, device=device, dtype=torch.uint8) + inputs = [ + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + cache_seqs, + output, + workspace, + softmax_stats, + kv_bounds, + ] + scale_kwargs = { + "softmax_scale": softmax_scale, + "output_scale": 1.0, + } + if input_dtype == torch.float8_e4m3fn: + inputs.extend( + [ + torch.tensor([softmax_scale], device=device), + torch.tensor([1.0], device=device), + ] + ) + scale_kwargs = {} + runner.forward( + inputs, + tactic=((128, 128), (128, 256), 4, False), + **scale_kwargs, + ) + + output_atol = 5e-2 if input_dtype == torch.float8_e4m3fn else 3e-2 + for b in range(batch_size): + page_ids = page_table[:, b].long() + key_latent = c_latent[:, :, page_ids].permute(2, 0, 1).reshape(kv_len, latent_dim).float() + key_rope = c_rope[:, :, page_ids].permute(2, 0, 1).reshape(kv_len, rope_dim).float() + for t in range(seq_len_q): + bound = int(kv_bounds[b * seq_len_q + t]) + stats_row = softmax_stats[b * seq_len_q + t] + if bound == 0: + # The reduction gates zero-KV tokens out of the CP merge with + # the canonical empty pair; the output row is unspecified. + assert torch.isneginf(stats_row[:, 0]).all() + torch.testing.assert_close(stats_row[:, 1], torch.zeros_like(stats_row[:, 1])) + continue + query_latent = q_latent[:, :, t, b].float() + query_rope = q_rope[:, :, t, b].float() + scores = ( + query_latent @ key_latent[:bound].transpose(0, 1) + + query_rope @ key_rope[:bound].transpose(0, 1) + ) * softmax_scale + expected_output = torch.softmax(scores, dim=-1) @ key_latent[:bound] + torch.testing.assert_close( + output_storage[b, t].float(), + expected_output, + rtol=output_atol, + atol=output_atol, + ) + actual_log_partition = stats_row[:, 0] + stats_row[:, 1].log() + expected_log_partition = torch.logsumexp(scores, dim=-1) + torch.testing.assert_close( + actual_log_partition, expected_log_partition, rtol=1e-4, atol=1e-4 + ) + + +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +def test_cute_dsl_mla_helix_bounds_tile_boundary(input_dtype: torch.dtype) -> None: + """Make the masked-span widening load-bearing: with S_q=2 and K=129 + (two 128-wide QK tiles), the causal span masks only the last tile while + the minimum per-token bound (K - S_q = 127) lands in the first, so a + port that drops the widening returns unmasked probability mass.""" + import cutlass + + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLNVMlaDecodeBlackwellRunner + + torch.manual_seed(29) + device = torch.device("cuda") + batch_size, seq_len_q, num_heads = 1, 2, 96 + latent_dim, rope_dim, page_size, kv_len = 512, 64, 64, 129 + blocks_per_sequence = (kv_len + page_size - 1) // page_size + num_pages = batch_size * blocks_per_sequence + softmax_scale = 1.0 / math.sqrt(latent_dim + rope_dim) + + q_storage = ( + torch.randn( + batch_size, + seq_len_q, + num_heads, + latent_dim + rope_dim, + device=device, + dtype=torch.bfloat16, + ) + * 0.1 + ).to(input_dtype) + cache_storage = ( + torch.randn( + num_pages, + page_size, + latent_dim + rope_dim, + device=device, + dtype=torch.bfloat16, + ) + * 0.1 + ).to(input_dtype) + q_latent = q_storage[..., :latent_dim].permute(2, 3, 1, 0) + q_rope = q_storage[..., latent_dim:].permute(2, 3, 1, 0) + c_latent = cache_storage[..., :latent_dim].permute(1, 2, 0) + c_rope = cache_storage[..., latent_dim:].permute(1, 2, 0) + page_table = ( + torch.arange(num_pages, dtype=torch.int32, device=device) + .view(batch_size, blocks_per_sequence) + .transpose(0, 1) + ) + cache_seqs = torch.tensor([kv_len], dtype=torch.int32, device=device) + kv_bounds = torch.tensor([127, 129], dtype=torch.int32, device=device) + output_storage = torch.empty( + batch_size, seq_len_q, num_heads, latent_dim, device=device, dtype=torch.bfloat16 + ) + output = output_storage.permute(2, 3, 1, 0) + softmax_stats = torch.empty( + batch_size * seq_len_q, num_heads, 2, device=device, dtype=torch.float32 + ) + + cutlass_dtype = cutlass.Float8E4M3FN if input_dtype == torch.float8_e4m3fn else cutlass.BFloat16 + runner = CuteDSLNVMlaDecodeBlackwellRunner( + in_dtype=cutlass_dtype, + num_heads=num_heads, + seq_len_q=seq_len_q, + page_size=page_size, + max_batch_size=batch_size, + emit_softmax_stats=True, + ) + workspace_size = runner.get_max_padded_workspace_size( + num_heads, seq_len_q, latent_dim, batch_size, cutlass.Float32 + ) + workspace = torch.empty(workspace_size, device=device, dtype=torch.uint8) + inputs = [ + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + cache_seqs, + output, + workspace, + softmax_stats, + kv_bounds, + ] + scale_kwargs = {"softmax_scale": softmax_scale, "output_scale": 1.0} + if input_dtype == torch.float8_e4m3fn: + inputs.extend( + [ + torch.tensor([softmax_scale], device=device), + torch.tensor([1.0], device=device), + ] + ) + scale_kwargs = {} + runner.forward( + inputs, + tactic=((128, 128), (128, 256), 1, False), + **scale_kwargs, + ) + + page_ids = page_table[:, 0].long() + key_latent = ( + c_latent[:, :, page_ids].permute(2, 0, 1).reshape(num_pages * page_size, latent_dim).float() + ) + key_rope = ( + c_rope[:, :, page_ids].permute(2, 0, 1).reshape(num_pages * page_size, rope_dim).float() + ) + for t in range(seq_len_q): + bound = int(kv_bounds[t]) + query_latent = q_latent[:, :, t, 0].float() + query_rope = q_rope[:, :, t, 0].float() + scores = ( + query_latent @ key_latent[:bound].transpose(0, 1) + + query_rope @ key_rope[:bound].transpose(0, 1) + ) * softmax_scale + actual_log_partition = softmax_stats[t, :, 0] + softmax_stats[t, :, 1].log() + expected_log_partition = torch.logsumexp(scores, dim=-1) + torch.testing.assert_close( + actual_log_partition, expected_log_partition, rtol=1e-4, atol=1e-4 + ) diff --git a/tests/unittest/_torch/attention/test_cute_dsl_mla_split_plan.py b/tests/unittest/_torch/attention/test_cute_dsl_mla_split_plan.py new file mode 100644 index 000000000000..566fe85dd9e4 --- /dev/null +++ b/tests/unittest/_torch/attention/test_cute_dsl_mla_split_plan.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request split-KV plan of the CuTe DSL MLA decode: policy, kernel vs mirror, kernel output.""" + +import math +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or not IS_CUTLASS_DSL_AVAILABLE + or get_sm_version() not in (100, 103, 107), + reason="CuTe DSL MLA decode runs on SM100/SM103/SM107 only (see CuteDslMlaFmha)", +) + +# One SM107 wave: 212 SMs, one 2-CTA cluster per SM pair -> 106 clusters. +WAVE = 106 + + +def _plan_mod(): + import tensorrt_llm._torch.attention.backends.fmha.cute_dsl_mla as mod + + return mod + + +@pytest.mark.parametrize( + "kv_lens,rows,expected", + [ + # helix16 bs8 MTP7 (8 rows per request): 64 clusters fill the wave at one task each. + ([14208] * 8, 8, [1] * 8), + # helix16 bs4: 32 rows -> three pieces per request (96 clusters, one wave). + ([14208] * 4, 8, [3] * 4), + # helix16 bs2: six pieces of 19 tiles. + ([14208] * 2, 8, [6] * 2), + # helix16 bs16: already over one wave, a second piece would cost a third wave. + ([14208] * 16, 8, [1] * 16), + # One long request alone (H128, single token): the whole grid. + ([65536], 1, [32]), + # 64K x 8 (64 rows): three pieces fill two waves better than one 512-tile pass. + ([65536] * 8, 8, [3] * 8), + # Short requests never get pieces below the minimum tile count. + ([512] * 4, 8, [1] * 4), + # DEP skew: the 256K request takes eight 256-tile pieces, the 32K ones stay whole. + ([32768] * 15 + [262144], 8, [1] * 15 + [8]), + ], +) +def test_plan_mla_split_kv_policy(kv_lens: list[int], rows: int, expected: list[int]) -> None: + plan = _plan_mod().plan_mla_split_kv(kv_lens, rows_per_request=rows, wave_clusters=WAVE) + assert plan == expected + + +def test_plan_kernel_matches_host_mirror() -> None: + mod = _plan_mod() + generator = torch.Generator().manual_seed(3) + for rows in (1, 4, 8): + for batch in (1, 3, 8, 16, 37, 128): + kv = torch.randint(1, 300_000, (batch,), generator=generator) + kv[0] = 262144 # a long one + kv_lens = kv.tolist() + expected = mod.plan_mla_split_kv(kv_lens, rows_per_request=rows, wave_clusters=WAVE) + cache_seqs = torch.tensor(kv_lens, dtype=torch.int32, device="cuda") + plan = torch.full((batch + 3,), -5, dtype=torch.int32, device="cuda") + mod.launch_mla_split_kv_plan( + cache_seqs, + plan, + num_requests=batch, + rows_per_request=rows, + wave_clusters=WAVE, + ) + torch.cuda.synchronize() + assert plan[:batch].tolist() == expected, (rows, batch, kv_lens[:4]) + assert (plan[batch:] == -5).all() + + +def test_split_kv_plan_is_launched_once_per_step(monkeypatch) -> None: + mod = _plan_mod() + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLNVMlaDecodeBlackwellRunner + + monkeypatch.setattr( + CuteDSLNVMlaDecodeBlackwellRunner, "_get_max_active_blocks", classmethod(lambda cls: 212) + ) + launches = [] + monkeypatch.setattr( + mod, "launch_mla_split_kv_plan", lambda *args, **kwargs: launches.append(kwargs) + ) + meta = SimpleNamespace(max_num_requests=16, _cute_dsl_mla_split_plan_key=None) + cache_seqs = torch.full((8,), 14208, dtype=torch.int32, device="cuda") + first = mod.CuteDslMlaFmha._split_kv_plan(meta, cache_seqs, 8, 8, 96) + assert first.numel() == 8 and meta._cute_dsl_mla_split_plan.numel() == 16 + # 8 tokens x 96 heads pack into six 128-row query tiles per request. + assert launches == [dict(num_requests=8, rows_per_request=6, wave_clusters=106)] + # The other layers of the step reuse the plan. + second = mod.CuteDslMlaFmha._split_kv_plan(meta, cache_seqs, 8, 8, 96) + assert second.data_ptr() == first.data_ptr() and len(launches) == 1 + # A new step (metadata invalidation) plans again. + meta._cute_dsl_mla_split_plan_key = None + mod.CuteDslMlaFmha._split_kv_plan(meta, cache_seqs, 8, 8, 96) + assert len(launches) == 2 + # Beyond the plan's request capacity: uniform split. + assert mod.CuteDslMlaFmha._split_kv_plan(meta, cache_seqs, 2000, 8, 96) is None + + +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("empty_request", [False, True]) +def test_runner_variable_split_matches_uniform_split( + input_dtype: torch.dtype, empty_request: bool +) -> None: + """Per-request splits under a 16-slot grid reproduce the uniform-split output and Helix stats. + + Covers requests planned onto one split (stored directly by the split + kernel), onto several, onto more splits than KV tiles, and a request + without local KV (the reduction's canonical empty result). + """ + import cutlass + + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLNVMlaDecodeBlackwellRunner + + torch.manual_seed(5) + device = torch.device("cuda") + batch_size, seq_len_q, num_heads = 4, 2, 96 + latent_dim, rope_dim, page_size = 512, 64, 64 + kv_lens = [4096, 0 if empty_request else 300, 9000, 1500] + blocks_per_sequence = math.ceil(max(kv_lens) / page_size) + num_pages = batch_size * blocks_per_sequence + softmax_scale = 1.0 / math.sqrt(latent_dim + rope_dim) + q_storage = ( + torch.randn(batch_size, seq_len_q, num_heads, latent_dim + rope_dim, device=device) * 0.1 + ).to(input_dtype) + cache_storage = ( + torch.randn(num_pages, page_size, latent_dim + rope_dim, device=device) * 0.1 + ).to(input_dtype) + q_latent = q_storage[..., :latent_dim].permute(2, 3, 1, 0) + q_rope = q_storage[..., latent_dim:].permute(2, 3, 1, 0) + c_latent = cache_storage[..., :latent_dim].permute(1, 2, 0) + c_rope = cache_storage[..., latent_dim:].permute(1, 2, 0) + page_table = ( + torch.arange(num_pages, dtype=torch.int32, device=device) + .view(batch_size, blocks_per_sequence) + .transpose(0, 1) + ) + cache_seqs = torch.tensor(kv_lens, dtype=torch.int32, device=device) + kv_bounds = torch.tensor( + [max(k - (seq_len_q - 1 - q), 0) for k in kv_lens for q in range(seq_len_q)], + dtype=torch.int32, + device=device, + ) + cutlass_dtype = cutlass.Float8E4M3FN if input_dtype == torch.float8_e4m3fn else cutlass.BFloat16 + runner = CuteDSLNVMlaDecodeBlackwellRunner( + in_dtype=cutlass_dtype, + num_heads=num_heads, + seq_len_q=seq_len_q, + page_size=page_size, + max_batch_size=batch_size, + emit_softmax_stats=True, + ) + workspace = torch.empty( + runner.get_max_padded_workspace_size( + num_heads, seq_len_q, latent_dim, batch_size, cutlass.Float32 + ), + device=device, + dtype=torch.uint8, + ) + + def run(block_split_kvs, split_kv): + output_storage = torch.zeros( + batch_size, seq_len_q, num_heads, latent_dim, device=device, dtype=torch.bfloat16 + ) + stats = torch.zeros(batch_size * seq_len_q, num_heads, 2, device=device) + inputs = [ + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + cache_seqs, + output_storage.permute(2, 3, 1, 0), + workspace, + stats, + kv_bounds, + ] + kwargs = {"softmax_scale": softmax_scale, "output_scale": 1.0} + if input_dtype == torch.float8_e4m3fn: + inputs.extend( + [torch.tensor([softmax_scale], device=device), torch.tensor([1.0], device=device)] + ) + kwargs = {} + if block_split_kvs is not None: + inputs.append(torch.tensor(block_split_kvs, dtype=torch.int32, device=device)) + runner.forward(inputs, tactic=((128, 128), (128, 256), split_kv, False), **kwargs) + torch.cuda.synchronize() + return output_storage.float(), stats + + ref_out, ref_stats = run(None, 1) + for plan, split_kv in ( + ([3, 1, 7, 2], 16), + ([16, 16, 16, 16], 16), + ([40, 1, 2, 1], 8), + ([1, 1, 1, 1], 16), + ): + out, stats = run(plan, split_kv) + tol = 2e-2 if input_dtype == torch.float8_e4m3fn else 1e-2 + torch.testing.assert_close(out, ref_out, rtol=tol, atol=tol) + log_partition = stats[:, :, 0] + stats[:, :, 1].log() + ref_log_partition = ref_stats[:, :, 0] + ref_stats[:, :, 1].log() + finite = torch.isfinite(ref_log_partition) + assert torch.equal(torch.isfinite(log_partition), finite) + torch.testing.assert_close( + log_partition[finite], ref_log_partition[finite], rtol=1e-4, atol=1e-4 + ) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index 3dea34c17649..9ba390d40484 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py @@ -21,7 +21,11 @@ import numpy as np import pytest -from tensorrt_llm._torch.pyexecutor._util import CacheCost, KvCacheCreator +from tensorrt_llm._torch.pyexecutor._util import ( + _STANDALONE_DRAFT_TOKENS_PER_BLOCK, + CacheCost, + KvCacheCreator, +) from tensorrt_llm._torch.pyexecutor.config_utils import uses_vswa_kv_cache_layout from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode @@ -1199,6 +1203,81 @@ def test_non_external_drafter_is_untouched(self, mocker): assert c._get_draft_kv_model_config() is draft_model_config +class TestStandaloneDrafterPageSize: + """A standalone drafter's separate pool has its own page size. + + The target's tokens_per_block is whatever its attention backend or the + deployment config chose (FP4 MLA forces 128), and the drafter's block-decode + kernels need not accept it: trtllm-gen MLA decode takes 32 or 64 only. The + pool is built with the drafter's page size, and the budget split must charge + the same one, since per-request quotas round up to whole pages. + """ + + TARGET_TOKENS_PER_BLOCK = 128 + + def _creator(self, mocker): + from tensorrt_llm.models.modeling_utils import QuantConfig + + c, draft_model_config, _ = TestExternalDrafterKvDtype()._creator(mocker, QuantConfig()) + c._tokens_per_block = self.TARGET_TOKENS_PER_BLOCK + # Read by the draft manager-class resolution on both the cost and the + # allocation path; the resolution itself is not under test here. + c._cache_transceiver_config = None + c._max_beam_width = 1 + c._validate_or_fallback_kv_cache_manager_v2 = Mock(side_effect=lambda cls, *_: cls) + return c, draft_model_config + + def test_standalone_drafter_page_size_is_independent_of_the_target(self, mocker): + c, _ = self._creator(mocker) + assert c._get_draft_tokens_per_block() == _STANDALONE_DRAFT_TOKENS_PER_BLOCK + assert c._get_draft_tokens_per_block() != c._tokens_per_block + + def test_draft_layers_of_the_target_keep_the_target_page_size(self, mocker): + """MTP/Eagle3 draft layers run on the target's attention metadata.""" + c, _ = self._creator(mocker) + c._speculative_config = MTPDecodingConfig(num_nextn_predict_layers=1, max_draft_len=4) + assert c._get_draft_tokens_per_block() == self.TARGET_TOKENS_PER_BLOCK + + def test_cost_charges_the_draft_page_size(self, mocker): + c, _ = self._creator(mocker) + size_spy = mocker.spy(KVCacheManagerV2, "get_cache_size_per_token") + bytes_spy = mocker.spy(KVCacheManagerV2, "get_cache_bytes_per_request") + quota_spy = mocker.spy(KVCacheManagerV2, "get_cache_quota_per_request") + + c._get_draft_cache_cost( + c._kv_cache_config, + use_separate_draft_kv_cache=True, + include_request_cost=True, + ) + + for spy in (size_spy, bytes_spy, quota_spy): + assert spy.call_count == 1, spy + assert spy.call_args.kwargs["tokens_per_block"] == ( + _STANDALONE_DRAFT_TOKENS_PER_BLOCK + ), spy + + def test_allocation_uses_the_page_size_the_cost_charged(self, mocker): + c, draft_model_config = self._creator(mocker) + draft_model_config.sparse_attention_config = None + c._model_engine._max_cuda_graph_batch_size = 1 + c._skip_est = False + c._enable_kv_cache_stats = Mock(return_value=False) + c._execution_stream = None + c._kv_connector_manager = None + c._disable_overlap_scheduler = False + c._joint_kv_cache_reuse = False + create = mocker.patch("tensorrt_llm._torch.pyexecutor._util._create_kv_cache_manager") + size_spy = mocker.spy(KVCacheManagerV2, "get_cache_size_per_token") + + c._create_one_model_draft_kv_cache_manager(c._max_seq_len) + c._get_kv_size_per_token() + + assert create.call_count == 1 + allocated = create.call_args.kwargs["tokens_per_block"] + assert allocated == _STANDALONE_DRAFT_TOKENS_PER_BLOCK + assert size_spy.call_args.kwargs["tokens_per_block"] == allocated + + class TestMambaEffectiveTpSize: """The sharding rule shared by the mamba pool allocator and the budget. @@ -1619,6 +1698,67 @@ def test_floor_exceeding_slope_budget_is_infeasible(self): with pytest.raises(ValueError, match="insufficient"): c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") + def test_draft_gpu_budget_cap_moves_the_excess_to_the_target(self, monkeypatch): + """TLLM_DRAFT_KV_GPU_BUDGET_GIB caps the draft's GPU share (above its + intercept) and hands the rest of the budget to the target.""" + monkeypatch.setenv("TLLM_DRAFT_KV_GPU_BUDGET_GIB", "1.5") + c = self._creator( + 10 * GB, + target=CacheCost(slope=80, intercept=GB, bytes_per_request=GB), + draft=CacheCost(slope=20, intercept=GB // 2, bytes_per_request=GB), + ) + + target_config, draft_config = c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") + + assert draft_config.max_gpu_total_bytes == GB // 2 + 3 * GB // 2 + assert target_config.max_gpu_total_bytes == 10 * GB - 2 * GB + + def test_draft_gpu_budget_cap_keeps_the_allocator_floor(self, monkeypatch): + """The cap never goes below the draft's one-request allocator quota.""" + monkeypatch.setenv("TLLM_DRAFT_KV_GPU_BUDGET_GIB", "0.5") + c = self._creator( + 10 * GB, + target=CacheCost(slope=80, bytes_per_request=GB), + draft=CacheCost(slope=20, bytes_per_request=GB, quota_per_request=2 * GB), + ) + + target_config, draft_config = c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") + + assert draft_config.max_gpu_total_bytes == 2 * GB + assert target_config.max_gpu_total_bytes == 8 * GB + + @pytest.mark.parametrize( + "value, match", + [ + ("abc", "must be a number"), + ("0", "must be positive"), + ("-1", "must be positive"), + ("nan", "must be positive"), + ], + ) + def test_draft_gpu_budget_cap_rejects_invalid_values(self, monkeypatch, value, match): + """A malformed or non-positive cap fails with the variable's name instead + of leaving the draft a zero or negative slope budget.""" + monkeypatch.setenv("TLLM_DRAFT_KV_GPU_BUDGET_GIB", value) + c = self._creator( + 10 * GB, + target=CacheCost(slope=80, bytes_per_request=GB), + draft=CacheCost(slope=20, bytes_per_request=GB), + ) + + with pytest.raises(ValueError, match=f"TLLM_DRAFT_KV_GPU_BUDGET_GIB {match}"): + c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") + + def test_draft_gpu_budget_cap_ignores_offload_budgets(self, monkeypatch): + """The cap applies to the GPU budget only.""" + monkeypatch.setenv("TLLM_DRAFT_KV_GPU_BUDGET_GIB", "0") + c = _make_creator(max_gpu_total_bytes=1_000, host_cache_size=1_000) + + target_config, draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") + + assert draft_config.host_cache_size == 200 + assert target_config.host_cache_size == 800 + def test_fully_covered_target_weight_hands_slope_budget_to_draft(self): """A target whose entire retention is intercept-funded reports a zero slope-funded weight: the draft takes the remainder, still capped at diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index f76bae05c042..0387894bf82c 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -385,6 +385,37 @@ def test_kimi_explicit_v2_manager_enables_kda_replay( assert kwargs["conv_state_layout"] == "q_k_v" +@pytest.mark.parametrize("max_draft_len, replay_enabled", [(7, True), (8, False)]) +def test_kimi_kda_replay_respects_fused_verify_draft_limit( + monkeypatch: pytest.MonkeyPatch, max_draft_len: int, replay_enabled: bool +) -> None: + """Draft lengths beyond the fused KDA verify kernel's limit keep the + legacy intermediate buffers (sequential FLA verify) instead of + allocating replay caches the kernel would reject.""" + kda_kernels = "tensorrt_llm._torch.modules.kimi_kda._kda_kernels" + monkeypatch.setattr(f"{kda_kernels}.is_kda_mtp_verify_available", lambda: True) + monkeypatch.setattr(f"{kda_kernels}.kda_mtp_max_num_spec", lambda: 7) + spec_config = MTPDecodingConfig(max_draft_len=max_draft_len) + + _, kwargs = _capture_kimi_v2_manager_ctor(monkeypatch, spec_config) + + if replay_enabled: + assert kwargs["kda_replay_num_spec"] == max_draft_len + else: + assert "kda_replay_num_spec" not in kwargs + + +def test_kda_mtp_max_num_spec_fits_the_kernel_chunk() -> None: + pytest.importorskip("cutlass") + from tensorrt_llm._torch.cute_dsl_kernels.blackwell.kimi_k3_kda.kda_mtp_decode import T_PAD + from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import kda_mtp_max_num_spec + + max_num_spec = kda_mtp_max_num_spec() + # A verify chunk holds up to max_num_spec replayed drafts plus 1 + max_num_spec new tokens. + assert max_num_spec == 7 + assert 2 * max_num_spec + 1 <= T_PAD + + def test_kimi_explicit_v2_manager_uses_qkv_convolution_layout( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -3649,11 +3680,12 @@ def test_v2_kda_replay_allocates_logical_slot_caches(): assert layer_cache.kda_conv_q.shape == (cache_size, 48, 6) assert layer_cache.kda_conv_k.shape == (cache_size, 48, 6) assert layer_cache.kda_conv_v.shape == (cache_size, 48, 6) - assert layer_cache.kda_conv_q.dtype is torch.float32 + assert layer_cache.kda_conv_q.dtype is torch.bfloat16 assert layer_cache.kda_conv_q.stride(-2) == 1 - assert layer_cache.kda_qkg_cache.shape == (cache_size, 2, 3, 48) + assert layer_cache.kda_k_cache.shape == (cache_size, 2, 48) + assert layer_cache.kda_g_cache.shape == (cache_size, 2, 48) assert layer_cache.kda_v_cache.shape == (cache_size, 2, 48) - # Six fp32 heads are 24 bytes, so the row is padded to 32 (stride 8) + # Six bf16 heads are 12 bytes, so the row is padded to 16 (stride 8) # to keep every nested per-draft view 16-byte aligned for CuTe. assert layer_cache.kda_beta_cache.shape == (cache_size, 2, 6) assert layer_cache.kda_beta_cache.stride(-2) == 8 @@ -3665,7 +3697,10 @@ def test_v2_kda_replay_allocates_logical_slot_caches(): % _KDA_BETA_CACHE_ALIGNMENT_BYTES == 0 ) - assert layer_cache.kda_qkg_cache.dtype is torch.float32 + assert layer_cache.kda_k_cache.dtype is torch.bfloat16 + assert layer_cache.kda_g_cache.dtype is torch.float32 + assert layer_cache.kda_v_cache.dtype is torch.bfloat16 + assert layer_cache.kda_beta_cache.dtype is torch.bfloat16 assert layer_cache.prev_num_accepted_tokens.data_ptr() == ( mgr.prev_num_accepted_tokens.data_ptr() ) @@ -3766,6 +3801,7 @@ def test_v2_kda_replay_resets_context_slots(): mgr._use_kda_replay_update = True mgr._use_replay_state_update = False mgr.cuda_state_indices = torch.tensor([2, 1], dtype=torch.int32) + mgr._host_state_indices = torch.tensor([2, 1], dtype=torch.int32) mgr.prev_num_accepted_tokens = torch.tensor([5, 6, 7], dtype=torch.int32) mgr.mamba_ssm_rand_seed = None @@ -3813,7 +3849,8 @@ def test_v2_kda_replay_host_drafter_records_active_requests(monkeypatch): "kda_conv_q", "kda_conv_k", "kda_conv_v", - "kda_qkg_cache", + "kda_k_cache", + "kda_g_cache", "kda_v_cache", "kda_beta_cache", ) @@ -3918,7 +3955,7 @@ def test_v2_kda_replay_suspended_request_resumes_from_its_v2_states(resume_slot) cache = getattr(mgr, name)[0, resume_slot] assert torch.equal(cache[:, :3], window[2 * section : 2 * section + 2]) assert not cache[:, 3:].any() - for name in ("kda_qkg_cache", "kda_v_cache", "kda_beta_cache"): + for name in ("kda_k_cache", "kda_g_cache", "kda_v_cache", "kda_beta_cache"): assert not getattr(mgr, name)[0, resume_slot].any() if resume_slot != 0: assert _kda_replay_slot(mgr, 0) == (3, [7.0] * len(_KDA_REPLAY_BUFFERS)) @@ -4091,7 +4128,8 @@ def inputs(*shape): A_log=weights["A_log"], dt_bias=weights["dt_bias"], recurrent_state=mgr.all_ssm_states[layer_offset], - qkg_cache=mgr.kda_qkg_cache[layer_offset], + k_cache=mgr.kda_k_cache[layer_offset], + g_cache=mgr.kda_g_cache[layer_offset], v_cache=mgr.kda_v_cache[layer_offset], beta_cache=mgr.kda_beta_cache[layer_offset], ssm_state_indices=slots, @@ -4174,7 +4212,9 @@ def _run_kda_verify_rounds(accepted_per_round, suspend_after=None, slot_taken=Fa @pytest.mark.parametrize("slot_taken", [True, False], ids=["slot_taken", "slot_kept"]) def test_v2_kda_replay_suspend_resume_matches_an_unsuspended_run(slot_taken): """A request is suspended with two accepted drafts pending and resumes; every later verify gives the outputs and - leaves the states of the same request never suspended, bit for bit, so its tokens cannot change. + leaves the states of the same request never suspended. The verifies up to the suspend match bit for bit. The + suspend commits the pending drafts as a chunk of their own, and the tensor-core verify computes each chunk in one + WY-form pass, so the later ones match to the kernel's accuracy rather than bit for bit. slot_taken: while the request is suspended, new caches fill the SSM pool, so the V2 runtime moves its SSM page to the host tier and gives its slot to one of them, which overwrites the slot's states; the request resumes in @@ -4184,14 +4224,20 @@ def test_v2_kda_replay_suspend_resume_matches_an_unsuspended_run(slot_taken): if not is_kda_mtp_verify_available(): pytest.skip("needs the fused KDA verify kernel (sm_100 / sm_103 with the CuTe DSL)") accepted_per_round = [2, 1, 2, 0, 2, 1] + suspend_after = 2 expected_outputs, expected_states = _run_kda_verify_rounds(accepted_per_round) outputs, states = _run_kda_verify_rounds( - accepted_per_round, suspend_after=2, slot_taken=slot_taken + accepted_per_round, suspend_after=suspend_after, slot_taken=slot_taken ) for round_idx, (actual, expected) in enumerate(zip(outputs, expected_outputs)): - assert torch.equal(actual, expected), f"verify round {round_idx}" - assert torch.equal(states, expected_states) + if round_idx <= suspend_after: + assert torch.equal(actual, expected), f"verify round {round_idx}" + else: + torch.testing.assert_close( + actual, expected, atol=2e-2, rtol=2e-2, msg=f"verify round {round_idx}" + ) + torch.testing.assert_close(states, expected_states, atol=1e-3, rtol=1e-3) @skip_no_cuda @@ -4212,7 +4258,8 @@ def test_v2_kda_replay_relocates_live_slot_history(): mgr.kda_conv_q, mgr.kda_conv_k, mgr.kda_conv_v, - mgr.kda_qkg_cache, + mgr.kda_k_cache, + mgr.kda_g_cache, mgr.kda_v_cache, mgr.kda_beta_cache, ) @@ -4266,7 +4313,8 @@ def close(self) -> None: mgr.kda_conv_q, mgr.kda_conv_k, mgr.kda_conv_v, - mgr.kda_qkg_cache, + mgr.kda_k_cache, + mgr.kda_g_cache, mgr.kda_v_cache, mgr.kda_beta_cache, ) @@ -4309,7 +4357,8 @@ def dim_contiguous_conv_cache() -> torch.Tensor: mgr.kda_conv_q = dim_contiguous_conv_cache() mgr.kda_conv_k = dim_contiguous_conv_cache() mgr.kda_conv_v = dim_contiguous_conv_cache() - mgr.kda_qkg_cache = torch.full((2, 4, 2, 3, 2), 7.0) + mgr.kda_k_cache = torch.full((2, 4, 2, 2), 7.0) + mgr.kda_g_cache = torch.full((2, 4, 2, 2), 7.0) mgr.kda_v_cache = torch.full((2, 4, 2, 2), 7.0) mgr.kda_beta_cache = torch.full((2, 4, 2, 1), 7.0) @@ -4327,7 +4376,8 @@ def dim_contiguous_conv_cache() -> torch.Tensor: conv_state[slot, start : start + 2, :], ) assert torch.count_nonzero(replay_buffer[layer_offset, slot, :, 4:]) == 0 - assert torch.count_nonzero(mgr.kda_qkg_cache[layer_offset, slot]) == 0 + assert torch.count_nonzero(mgr.kda_k_cache[layer_offset, slot]) == 0 + assert torch.count_nonzero(mgr.kda_g_cache[layer_offset, slot]) == 0 assert torch.count_nonzero(mgr.kda_v_cache[layer_offset, slot]) == 0 assert torch.count_nonzero(mgr.kda_beta_cache[layer_offset, slot]) == 0 assert mgr.prev_num_accepted_tokens[slot] == 0 @@ -4341,8 +4391,9 @@ def test_v2_kda_replay_seeds_bf16_conv_state_from_disagg_transfer(): ``extract_mamba_kv_cache_params`` forces the KDA delta-rule (ssm) pool to fp32 but leaves the short-convolution pool at the model dtype, so what the context rank serializes and the generation rank restores for the conv slot - is bf16 while the SA replay caches the generation rank rebuilds from it are - fp32. Seeding must widen the restored rows losslessly, keep the logical + is bf16, the same dtype as the replay conv caches the generation rank + rebuilds from it (k / v / beta replay caches are bf16 too, the gate cache + fp32). Seeding must copy the restored rows bit for bit, keep the logical shape, leave the restored bf16 pool itself untouched so the decode path reuses exactly what was transferred, and clear only the draft scratch of the seeded slots. @@ -4360,15 +4411,18 @@ def test_v2_kda_replay_seeds_bf16_conv_state_from_disagg_transfer(): mgr.all_conv_states = [(torch.randn(4, 6, 4) * 0.1).to(torch.bfloat16) for _ in range(2)] restored = [conv.clone() for conv in mgr.all_conv_states] + # Same dtypes as _allocate_pool_replay_buffers: bf16 conv / k / v / beta + # replay caches, fp32 gate cache. def dim_contiguous_conv_cache() -> torch.Tensor: - return torch.full((2, 4, 6, 2), 7.0).transpose(-1, -2) + return torch.full((2, 4, 6, 2), 7.0, dtype=torch.bfloat16).transpose(-1, -2) mgr.kda_conv_q = dim_contiguous_conv_cache() mgr.kda_conv_k = dim_contiguous_conv_cache() mgr.kda_conv_v = dim_contiguous_conv_cache() - mgr.kda_qkg_cache = torch.full((2, 4, 2, 3, 2), 7.0) - mgr.kda_v_cache = torch.full((2, 4, 2, 2), 7.0) - mgr.kda_beta_cache = torch.full((2, 4, 2, 1), 7.0) + mgr.kda_k_cache = torch.full((2, 4, 2, 2), 7.0, dtype=torch.bfloat16) + mgr.kda_g_cache = torch.full((2, 4, 2, 2), 7.0, dtype=torch.float32) + mgr.kda_v_cache = torch.full((2, 4, 2, 2), 7.0, dtype=torch.bfloat16) + mgr.kda_beta_cache = torch.full((2, 4, 2, 1), 7.0, dtype=torch.bfloat16) mgr.seed_kda_replay_caches_for_disagg_gen([101, 202]) @@ -4380,27 +4434,29 @@ def dim_contiguous_conv_cache() -> torch.Tensor: (mgr.kda_conv_k, 2), (mgr.kda_conv_v, 4), ): - # fp32 replay cache, logical shape unchanged by the widening. - assert replay_buffer.dtype is torch.float32 + # bf16 replay cache, logical shape unchanged by the seeding. + assert replay_buffer.dtype is torch.bfloat16 seeded = replay_buffer[layer_offset, slot, :, :4] assert seeded.shape == (2, 4) - # bf16 -> fp32 is lossless, so the restored values must survive - # the widening bit for bit. + # Same dtype on both sides: the restored values are copied bit + # for bit. torch.testing.assert_close( seeded, - conv_state[slot, start : start + 2, :].float(), + conv_state[slot, start : start + 2, :], rtol=0, atol=0, ) assert torch.count_nonzero(replay_buffer[layer_offset, slot, :, 4:]) == 0 - assert torch.count_nonzero(mgr.kda_qkg_cache[layer_offset, slot]) == 0 + assert torch.count_nonzero(mgr.kda_k_cache[layer_offset, slot]) == 0 + assert torch.count_nonzero(mgr.kda_g_cache[layer_offset, slot]) == 0 assert torch.count_nonzero(mgr.kda_v_cache[layer_offset, slot]) == 0 assert torch.count_nonzero(mgr.kda_beta_cache[layer_offset, slot]) == 0 assert mgr.prev_num_accepted_tokens[slot] == 0 # Slots no request restored keep their pre-seed contents. for slot in (0, 2): assert (mgr.kda_conv_q[layer_offset, slot] == 7.0).all() - assert (mgr.kda_qkg_cache[layer_offset, slot] == 7.0).all() + assert (mgr.kda_k_cache[layer_offset, slot] == 7.0).all() + assert (mgr.kda_g_cache[layer_offset, slot] == 7.0).all() # The restored bf16 pool is read-only here: plain decode (speculation # off) reuses exactly the state the transfer wrote. torch.testing.assert_close(conv_state, restored[layer_offset], rtol=0, atol=0) @@ -5065,3 +5121,48 @@ def test_v2_stack_state_views_rejects_non_affine_layout(): assert stack([pool[0], pool[1], pool[3]]) is None assert stack([pool[0], torch.zeros(5, 8)]) is None assert stack([]) is None + + +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=skip_no_cuda)]) +def test_v2_kda_admission_slot_updates_without_sync(device): + """Disagg GEN admission resets, relocates and seeds KDA replay slots; on CUDA none of it may + synchronize the stream (the previous step is still queued).""" + num_slots, num_layers, num_spec, dim, window = 8, 2, 2, 2, 3 + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr._use_kda_replay_update = True + mgr._use_replay_state_update = False + mgr.mamba_ssm_rand_seed = None + mgr.local_num_mamba_layers = num_layers + mgr.prev_num_accepted_tokens = torch.arange(1, num_slots + 1, dtype=torch.int32, device=device) + # Stale device indices: the reset must take its slots from _host_state_indices. + mgr.cuda_state_indices = torch.full((num_slots,), 7, dtype=torch.int32, device=device) + mgr._host_state_indices = torch.tensor([5, 1, 0, 0, 0, 0, 0, 0], dtype=torch.int32) + mgr._request_id_to_state_index = {10: 4, 11: 2} + mgr._kda_replay_unseeded = set() + mgr.conv_state_shape = [3 * dim, window] + mgr.conv_section_dims = [dim] * 3 + mgr.all_conv_states = [ + torch.randn(num_slots, 3 * dim, window, device=device) for _ in range(num_layers) + ] + mgr.kda_conv_q, mgr.kda_conv_k, mgr.kda_conv_v = ( + torch.randn(num_layers, num_slots, dim, window + num_spec, device=device) for _ in range(3) + ) + mgr.kda_k_cache, mgr.kda_g_cache, mgr.kda_v_cache, mgr.kda_beta_cache = ( + torch.randn(num_layers, num_slots, num_spec, dim, device=device) for _ in range(4) + ) + + if device == "cuda": + torch.cuda.synchronize() + previous_mode = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode("error") + try: + mgr._reset_context_mamba_slots(2) + mgr._reset_kda_replay_slots([3]) + mgr._relocate_kda_replay_slots([1], [6]) + mgr.seed_kda_replay_caches_for_disagg_gen([10, 11]) + finally: + if device == "cuda": + torch.cuda.set_sync_debug_mode(previous_mode) + + # Reset 5, 1 and 3; slot 6 takes slot 1's (reset) count; seeding resets 4 and 2; 7 untouched. + assert mgr.prev_num_accepted_tokens.tolist() == [1, 0, 0, 0, 0, 0, 0, 8] diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py index 83d37b1c22af..33feb9c13fa3 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -831,27 +831,30 @@ def test_warmup_phases_pass_configured_speculation_state(phase: str) -> None: class _Tracker: - """Records method-call order via mock side_effects.""" + """Records method-call order (and each call's arguments) via mock side_effects.""" def __init__(self): self.calls = [] + self.call_args = {} def __call__(self, name): def _wrapped(*args, **kwargs): self.calls.append(name) + self.call_args.setdefault(name, []).append((args, kwargs)) return _wrapped def _run_warmup_tracked( - model_engine, resource_manager, *, force_helix_cp=False, capture_logs=False + model_engine, resource_manager, *, force_helix_cp=False, capture_logs=False, tracker=None ): """Patch warmup stages and cleanup operations, then run warmup. - Optionally force helix CP and capture logs. Returns - (call_order_list, log_records_or_None). + Optionally force helix CP and capture logs. Pass ``tracker`` to inspect the + recorded call arguments. Returns (call_order_list, log_records_or_None). """ - tracker = _Tracker() + if tracker is None: + tracker = _Tracker() # Engines left by earlier tests are cyclic; finalize them before tracking. gc.collect() # Observe this module's explicit cleanup without counting unrelated @@ -1235,18 +1238,99 @@ def test_empty_cache_count_under_default(self): f"Expected exactly 2 empty_cache calls; got order={calls}", ) - def test_step_b_cleanup_skipped_with_helix_cp(self): - """With Helix CP, can_run_general_warmup is False AND step (b) is - gated off -> no empty_cache calls inside warmup().""" + def test_helix_cp_runs_only_the_generation_autotuner_warmup(self): + """With Helix CP, can_run_general_warmup is False and the context + warmups (general, Mamba hybrid) are skipped; the autotuner warmup still + runs, on generation requests, followed by one empty_cache.""" model_engine, resource_manager = _build_engine_and_resource_manager() - calls, _ = _run_warmup_tracked(model_engine, resource_manager, force_helix_cp=True) - self.assertNotIn("autotuner", calls, f"Helix CP should skip autotuner; got {calls}") + tracker = _Tracker() + calls, _ = _run_warmup_tracked( + model_engine, resource_manager, force_helix_cp=True, tracker=tracker + ) + self.assertEqual( + tracker.call_args["autotuner"], + [((resource_manager,), {"generation_only": True})], + "Helix CP must tune on generation requests only", + ) + self.assertNotIn("general_warmup", calls, f"Helix CP should skip general; got {calls}") + self.assertIn("autotuner", calls, f"Helix CP should run the autotuner; got {calls}") self.assertEqual( calls.count("empty_cache"), - 0, - f"Helix CP should skip all warmup cleanup; got {calls}", + 1, + f"Helix CP should clean up once, after the autotuner; got {calls}", + ) + self.assertEqual(calls[calls.index("autotuner") + 1], "empty_cache") + + @pytest.mark.cpu_only + def test_generation_only_autotuner_warmup_tunes_one_full_generation_batch(self): + """Helix CP tunes on max_batch_size generation requests (bounded by + max_num_tokens), without the context batch.""" + tuner = SimpleNamespace( + setup_distributed_state=Mock(), + cache_pp_recv=Mock(), + cache_pp_send=Mock(), + clean_pp_flag=Mock(), + profiling_cache={}, + print_profiling_cache=Mock(), + ) + get_num_available_tokens = Mock(return_value=4096) + engine = SimpleNamespace( + _warmup_timer=_WarmupTimer(rank=0), + _config=SimpleNamespace( + enable_autotuner=True, + torch_compile_enabled=False, + max_num_tokens=100, + max_batch_size=32, + max_seq_len=1024, + original_max_draft_len=3, + max_total_draft_tokens=3, + is_spec_decode=True, + max_draft_len=3, + spec_config=None, + ), + cuda_graph_runner=SimpleNamespace(enabled=False), + model=SimpleNamespace(modules=lambda: []), + kv_cache_manager_key="kv_cache", + mapping=SimpleNamespace(tp_size=1, has_pp=lambda: False), + dist=object(), + guided_decoder=None, + no_cuda_graph=lambda: contextlib.nullcontext(), + _create_warmup_request=Mock(return_value=object()), + _release_batch_context=Mock(return_value=contextlib.nullcontext(object())), + _should_run_warmup_batch=Mock(return_value=True), + _release_megamoe_profiling_scratch=Mock(), + _forward_warmup=Mock(), + _completed_autotuner_warmup_num_tokens=None, + _agree_warmup_flag=lambda flag: flag, + ) + kv_cache_manager = SimpleNamespace(get_num_available_tokens=get_num_available_tokens) + resource_manager = SimpleNamespace( + get_resource_manager=lambda key: kv_cache_manager if key == "kv_cache" else None ) + with ( + patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner.AutoTuner.get", + return_value=tuner, + ), + patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner.autotune", + return_value=contextlib.nullcontext(), + ), + patch("torch.cuda.synchronize"), + patch("torch.cuda.empty_cache"), + patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner.clear_memory_buffers" + ), + ): + DecoderRunner._run_autotuner_warmup(engine, resource_manager, generation_only=True) + + # 100 tokens fit 25 requests of 1 + 3 tokens, below max_batch_size. + engine._create_warmup_request.assert_called_once_with(resource_manager, 100, 25) + engine._forward_warmup.assert_called_once() + get_num_available_tokens.assert_not_called() + self.assertEqual(engine._completed_autotuner_warmup_num_tokens, 100) + @pytest.mark.cpu_only def test_flashinfer_mxfp8_respects_disabled_global_autotuner(self): """The global autotuner switch also disables automatic FlashInfer tuning.""" diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index 8db7b62229ba..7b7d44422ff6 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/tests/unittest/_torch/misc/test_autotuner.py @@ -1745,6 +1745,52 @@ def test_post_tune_merge_tactics_single_rank_noop(): assert tuner.profiling_cache.cache == original +@pytest.mark.parametrize("tp_size,cp_size", [(1, 4), (2, 2)]) +def test_parallel_tactics_split_across_helix_cp_ranks(tp_size, cp_size): + """Helix CP ranks share the PARALLEL tactic sweep with the TP ranks.""" + from tensorrt_llm.mapping import CpType + + world_size = tp_size * cp_size + tactics = list(range(11)) + shares = [] + for rank in range(world_size): + tuner = autotuner.AutoTuner() + mapping = Mapping(world_size=world_size, + rank=rank, + tp_size=tp_size, + cp_size=cp_size, + cp_config={"cp_type": CpType.HELIX}) + tuner.setup_distributed_state(mapping, dist=object()) + assert tuner._is_distributed() + assert tuner._tuning_group() == (world_size, rank) + shares.append( + tuner._maybe_parallelize_tactics( + tactics, DistributedTuningStrategy.PARALLEL)) + assert tuner._maybe_parallelize_tactics( + tactics, DistributedTuningStrategy.MERGE) == tactics + + # Every tactic is tuned on exactly one rank. + assert sorted(t for share in shares for t in share) == tactics + + +def test_parallel_tactics_split_skips_non_helix_cp_ranks(): + """Non-Helix CP keeps the TP-only tuning group.""" + from tensorrt_llm.mapping import CpType + + tuner = autotuner.AutoTuner() + mapping = Mapping(world_size=4, + rank=3, + tp_size=1, + cp_size=4, + cp_config={"cp_type": CpType.ULYSSES}) + tuner.setup_distributed_state(mapping, dist=object()) + assert tuner._tuning_group() == (1, 0) + assert not tuner._is_distributed() + tactics = list(range(5)) + assert tuner._maybe_parallelize_tactics( + tactics, DistributedTuningStrategy.PARALLEL) == tactics + + def test_profiling_cache_enum_tactic_roundtrip(tmp_path): # An enum tactic must survive save -> load: repr("") isn't # ast.literal_eval-parsable, so the cache serializes tactic.value instead. diff --git a/tests/unittest/_torch/modeling/test_kimi_linear_checkpoint.py b/tests/unittest/_torch/modeling/test_kimi_linear_checkpoint.py index 077d1425121b..f031f418a53d 100644 --- a/tests/unittest/_torch/modeling/test_kimi_linear_checkpoint.py +++ b/tests/unittest/_torch/modeling/test_kimi_linear_checkpoint.py @@ -381,10 +381,6 @@ def test_kda_load_preserves_checkpoint_shards_and_fused_outputs(tp_size, tp_rank assert isinstance(module, _Fp8BlockScaleWeightReadLinear) expected = checkpoint[f"model.layers.0.self_attn.{name}.weight"] expected = expected.chunk(tp_size, dim=1 if name == "o_proj" else 0)[tp_rank] - assert torch.equal( - module.weight[: module.out_features].view(torch.uint8), expected.view(torch.uint8) - ) - assert module.weight.dtype == torch.float8_e4m3fn scales = checkpoint[f"model.layers.0.self_attn.{name}.weight_scale"][:, 0, :, 0] axis = 1 if name == "o_proj" else 0 width = module.in_features if axis == 1 else module.out_features @@ -392,6 +388,14 @@ def test_kda_load_preserves_checkpoint_shards_and_fused_outputs(tp_size, tp_rank indices = [slice(None), slice(None)] indices[axis] = slice(start // 128, (end + 127) // 128) scales = scales[tuple(indices)] + if module._fused_projection is not None: + # q/k/v/g are fused into qkvg_proj and prepared at load, so their + # checkpoint shard is held in its GEMM-ready (UE8M0) form. + expected, scales = _Fp8BlockScaleWeightReadLinear._prepare_weights(expected, scales) + assert torch.equal( + module.weight[: module.out_features].view(torch.uint8), expected.view(torch.uint8) + ) + assert module.weight.dtype == torch.float8_e4m3fn torch.testing.assert_close(module.weight_scale, scales, rtol=0, atol=0) for module in model.modules(): if isinstance(module, (Linear, _Fp8BlockScaleWeightReadLinear)): diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_mixed_batch_core_layout.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_mixed_batch_core_layout.py index b228d9b44f53..8d7270cdff3a 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kda_mixed_batch_core_layout.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_mixed_batch_core_layout.py @@ -44,7 +44,8 @@ class _LayerCache: def __init__(self, conv: torch.Tensor, temporal: torch.Tensor) -> None: self.conv = conv self.temporal = temporal - self.kda_qkg_cache = None + self.kda_k_cache = None + self.kda_g_cache = None self.has_kda_replay_caches = False diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py index c909afad0c83..9f2f2351a736 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py @@ -34,7 +34,7 @@ the drop's benchmark-tuned set {2, 12, 32}; the vendored v_row hoist fix makes them compile, and this test validates them numerically. -Requires: 1 GPU (sm100 for the CuTe kernel), fla-core, nvidia-cutlass-dsl, +Requires: 1 GPU (sm100/sm103/sm107), fla-core, nvidia-cutlass-dsl, cuda-bindings. Skips cleanly otherwise. """ @@ -57,12 +57,12 @@ def _is_blackwell(): if not torch.cuda.is_available(): return False prop = torch.cuda.get_device_properties(0) - return prop.major * 10 + prop.minor in (100, 103) + return prop.major * 10 + prop.minor in (100, 103, 107) pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"), - pytest.mark.skipif(not _is_blackwell(), reason="needs sm100/sm103"), + pytest.mark.skipif(not _is_blackwell(), reason="needs sm100/sm103/sm107"), pytest.mark.skipif(not _HAVE_DEPS, reason=f"deps: {_DEP_ERR}"), ] @@ -103,17 +103,20 @@ def make_conv_data(B, H, K=128, V=128, M=2, W=4, lower_bound=-5.0, scale=None, s A_log = torch.randn(H, dtype=torch.float32, device=device) dt_bias = torch.randn(dim, dtype=torch.float32, device=device) # dim-contiguous extended conv caches: allocate [B, S, dim], transpose. - cs_q = torch.randn(B, S, dim, dtype=torch.float32, device=device).transpose(1, 2) - cs_k = torch.randn(B, S, dim, dtype=torch.float32, device=device).transpose(1, 2) - cs_v = torch.randn(B, S, dim, dtype=torch.float32, device=device).transpose(1, 2) + # raw conv inputs are cached in bf16 like the projections that produce them + cs_q = torch.randn(B, S, dim, dtype=torch.bfloat16, device=device).transpose(1, 2) + cs_k = torch.randn(B, S, dim, dtype=torch.bfloat16, device=device).transpose(1, 2) + cs_v = torch.randn(B, S, dim, dtype=torch.bfloat16, device=device).transpose(1, 2) initial_state_kfirst = torch.randn(B, H, K, V, dtype=torch.float32, device=device) initial_state_cute = initial_state_kfirst.permute(0, 1, 3, 2).contiguous() cu_seqlens = torch.arange(0, B * T + 1, T, dtype=torch.int32, device=device) ssm_state_indices = torch.arange(B, dtype=torch.int32, device=device) num_accepted_tokens = torch.zeros(B, dtype=torch.int32, device=device) - qkg_cache = torch.zeros(B, M, 3, dim, dtype=torch.float32, device=device) - v_cache = torch.zeros(B, M, H * V, dtype=torch.float32, device=device) - beta_cache = torch.zeros(B, M, H, dtype=torch.float32, device=device) + # replay caches in the runtime's dtypes: k / v / beta bf16, gate fp32 + k_cache = torch.zeros(B, M, dim, dtype=torch.bfloat16, device=device) + g_cache = torch.zeros(B, M, dim, dtype=torch.float32, device=device) + v_cache = torch.zeros(B, M, H * V, dtype=torch.bfloat16, device=device) + beta_cache = torch.zeros(B, M, H, dtype=torch.bfloat16, device=device) return { "x_q": x_q, "x_k": x_k, @@ -130,7 +133,8 @@ def make_conv_data(B, H, K=128, V=128, M=2, W=4, lower_bound=-5.0, scale=None, s "cs_v": cs_v, "initial_state_kfirst": initial_state_kfirst, "initial_state_cute": initial_state_cute, - "qkg_cache": qkg_cache, + "k_cache": k_cache, + "g_cache": g_cache, "v_cache": v_cache, "beta_cache": beta_cache, "cu_seqlens": cu_seqlens, @@ -168,9 +172,16 @@ def cpu_reference(data): cs_q = data["cs_q"].float().cpu().clone() cs_k = data["cs_k"].float().cpu().clone() cs_v = data["cs_v"].float().cpu().clone() - qkg_cache = data["qkg_cache"].float().cpu().clone() + # fp32 working copies; new entries are rounded to the caches' dtypes like the kernel does + k_cache = data["k_cache"].float().cpu().clone() + g_cache = data["g_cache"].float().cpu().clone() v_cache = data["v_cache"].float().cpu().clone() beta_cache = data["beta_cache"].float().cpu().clone() + k_dtype, v_dtype, beta_dtype = ( + data["k_cache"].dtype, + data["v_cache"].dtype, + data["beta_cache"].dtype, + ) ht = data["initial_state_kfirst"].float().cpu().clone() ht_commit = ht.clone() out = torch.zeros(1, B * T, H, V, dtype=torch.float32) @@ -189,9 +200,10 @@ def cpu_reference(data): h_state = ht[slot, h].clone() for i_t in range(T_loop): if i_t < commit_len: - q_t = qkg_cache[slot, i_t, 0, hk : hk + K] - k_t = qkg_cache[slot, i_t, 1, hk : hk + K] - gk_t = qkg_cache[slot, i_t, 2, hk : hk + K] + # replayed drafts produce no output; their q is not cached + q_t = torch.zeros(K, dtype=torch.float32) + k_t = k_cache[slot, i_t, hk : hk + K] + gk_t = g_cache[slot, i_t, hk : hk + K] v_t = v_cache[slot, i_t, hv : hv + V] beta_t = beta_cache[slot, i_t, h] xq_raw = cs_q[slot, hk : hk + K, W - 1 + i_t] @@ -239,11 +251,10 @@ def cpu_reference(data): cs_v[slot, hv : hv + V, : W - 1] = hist_v if i_t > commit_len: cache_pos = i_t - commit_len - 1 - qkg_cache[slot, cache_pos, 0, hk : hk + K] = q_t - qkg_cache[slot, cache_pos, 1, hk : hk + K] = k_t - qkg_cache[slot, cache_pos, 2, hk : hk + K] = gk_t - v_cache[slot, cache_pos, hv : hv + V] = v_t - beta_cache[slot, cache_pos, h] = beta_t + k_cache[slot, cache_pos, hk : hk + K] = k_t.to(k_dtype).float() + g_cache[slot, cache_pos, hk : hk + K] = gk_t + v_cache[slot, cache_pos, hv : hv + V] = v_t.to(v_dtype).float() + beta_cache[slot, cache_pos, h] = beta_t.to(beta_dtype).float() cs_q[slot, hk : hk + K, W - 1 + cache_pos] = xq_raw cs_k[slot, hk : hk + K, W - 1 + cache_pos] = xk_raw cs_v[slot, hv : hv + V, W - 1 + cache_pos] = xv_raw @@ -251,7 +262,8 @@ def cpu_reference(data): return { "out": out, "recurrent_state": ht_commit, - "qkg_cache": qkg_cache, + "k_cache": k_cache, + "g_cache": g_cache, "v_cache": v_cache, "beta_cache": beta_cache, "cs_q": cs_q, @@ -260,11 +272,25 @@ def cpu_reference(data): } -def cute_run(data, zero_accepted_hint=False, beta_cache_override=None): - """Run the in-tree op on cloned caches; return the drop-format dict.""" +def cute_run( + data, + zero_accepted_hint=False, + packed_token_layout=False, + out=None, + fuse_output_norm=False, + quantize_output=False, + beta_cache_override=None, + state_override=None, +): + """Run the op, optionally substituting a pre-populated beta-cache view or a + caller-allocated recurrent-state tensor (e.g. a strided slice of a larger cache).""" import tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops # noqa: F401 - state = data["initial_state_cute"].clone() + if state_override is None: + state = data["initial_state_cute"].clone() + else: + state = state_override + state.copy_(data["initial_state_cute"]) cs = {} for name in ("cs_q", "cs_k", "cs_v"): src = data[name] @@ -273,10 +299,11 @@ def cute_run(data, zero_accepted_hint=False, beta_cache_override=None): ).transpose(1, 2) dst.copy_(src) cs[name] = dst - qkg_cache = data["qkg_cache"].clone() + k_cache = data["k_cache"].clone() + g_cache = data["g_cache"].clone() v_cache = data["v_cache"].clone() beta_cache = data["beta_cache"].clone() if beta_cache_override is None else beta_cache_override - out = torch.ops.trtllm.kda_mtp_decode( + op_kwargs = dict( x_q=data["x_q"], x_k=data["x_k"], x_v=data["x_v"], @@ -291,7 +318,8 @@ def cute_run(data, zero_accepted_hint=False, beta_cache_override=None): A_log=data["A_log"], dt_bias=data["dt_bias"], recurrent_state=state, - qkg_cache=qkg_cache, + k_cache=k_cache, + g_cache=g_cache, v_cache=v_cache, beta_cache=beta_cache, ssm_state_indices=data["ssm_state_indices"], @@ -301,13 +329,30 @@ def cute_run(data, zero_accepted_hint=False, beta_cache_override=None): lower_bound=data["lower_bound"], scale=data["scale"], zero_accepted_hint=zero_accepted_hint, + packed_token_layout=packed_token_layout, + onorm_g=data.get("onorm_g"), + onorm_weight=data.get("onorm_weight"), + onorm_eps=data.get("onorm_eps", 1e-5), + fuse_output_norm=fuse_output_norm, ) + output_scale = None + if quantize_output: + op_kwargs.pop("fuse_output_norm") + out, output_scale = torch.ops.trtllm.kda_mtp_decode_fp8(**op_kwargs) + elif out is None: + out = torch.ops.trtllm.kda_mtp_decode(**op_kwargs) + else: + from tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops import kda_mtp_decode_impl + + out = kda_mtp_decode_impl(out=out, **op_kwargs) return { "out": out, + "output_scale": output_scale, # committed state back in K-first layout for CPU-golden comparison "recurrent_state": state.permute(0, 1, 3, 2).contiguous(), "state_v_first": state, - "qkg_cache": qkg_cache, + "k_cache": k_cache, + "g_cache": g_cache, "v_cache": v_cache, "beta_cache": beta_cache, "cs_q": cs["cs_q"], @@ -316,6 +361,58 @@ def cute_run(data, zero_accepted_hint=False, beta_cache_override=None): } +def _pack_new_token_projection_views(data): + """Build production-like section views over compact fused GEMM outputs.""" + B, H, K, M, source_T = ( + data["B"], + data["H"], + data["K"], + data["M"], + data["T"], + ) + accepted = data["num_accepted_tokens"].cpu() + rows = torch.cat( + [ + torch.arange(n * source_T + int(accepted[n]), n * source_T + int(accepted[n]) + M + 1) + for n in range(B) + ] + ).cuda() + num_rows = B * (M + 1) + dim = H * K + + # q/k/v are adjacent sections of [q | k | v | onorm_g]. Their token + # stride is therefore 4 * dim, not dim. + qkvg = torch.empty(1, num_rows, 4 * dim, dtype=torch.bfloat16, device="cuda") + qkvg[..., :dim].copy_(data["x_q"][0, rows].reshape(1, num_rows, dim)) + qkvg[..., dim : 2 * dim].copy_(data["x_k"][0, rows].reshape(1, num_rows, dim)) + qkvg[..., 2 * dim : 3 * dim].copy_(data["x_v"][0, rows].reshape(1, num_rows, dim)) + qkvg[..., 3 * dim : 4 * dim].copy_(data["g"][0, rows].reshape(1, num_rows, dim)) + + # beta is the H-wide tail of the padded [f_a | beta] projection. + bfa_width = ((K + H + 7) // 8) * 8 + bfa = torch.empty(1, num_rows, bfa_width, dtype=torch.bfloat16, device="cuda") + bfa[..., K : K + H].copy_(data["beta"][0, rows].reshape(1, num_rows, H)) + + packed = dict(data) + packed.update( + x_q=qkvg[..., :dim].view(1, num_rows, H, K), + x_k=qkvg[..., dim : 2 * dim].view(1, num_rows, H, K), + x_v=qkvg[..., 2 * dim : 3 * dim].view(1, num_rows, H, K), + onorm_g=qkvg[..., 3 * dim : 4 * dim].view(1, num_rows, H, K), + g=data["g"][:, rows].contiguous(), + beta=bfa[..., K : K + H], + cu_seqlens=None, + T=M + 1, + T_total=num_rows, + ) + assert packed["x_q"].stride(1) == 4 * dim + assert packed["x_k"].stride(1) == 4 * dim + assert packed["x_v"].stride(1) == 4 * dim + assert packed["onorm_g"].stride(1) == 4 * dim + assert packed["beta"].stride(1) == bfa_width + return packed, rows + + def _fla_sequential_reference(data, num_accepted): """Per-request sequential conv+SiLU (fp32 torch) + fused_recurrent_kda. @@ -329,7 +426,7 @@ def _fla_sequential_reference(data, num_accepted): from fla.ops.kda import fused_recurrent_kda B, H, K, V, W = data["B"], data["H"], data["K"], data["V"], data["W"] - T = data["T"] + M, T = data["M"], data["T"] dim = H * K dev = data["x_q"].device out = torch.zeros(1, B * T, H, V, dtype=torch.float32, device=dev) @@ -355,7 +452,7 @@ def _fla_sequential_reference(data, num_accepted): xk = data["cs_k"][n, :, W - 1 + i_t].float() xv = data["cs_v"][n, :, W - 1 + i_t].float() tok = None - g_t = data["qkg_cache"][n, i_t, 2].float() + g_t = data["g_cache"][n, i_t].float() beta_t = data["beta_cache"][n, i_t].float() replay = True else: @@ -377,7 +474,7 @@ def conv_step(hist_s, x_raw, w): if replay: # Replayed tokens use the cached post-processed k/g/v/beta # exactly as the kernel does (delta rule applied directly). - k_t = data["qkg_cache"][n, i_t, 1].float().view(H, K) + k_t = data["k_cache"][n, i_t].float().view(H, K) gk_t = g_t.view(H, K) v_t = data["v_cache"][n, i_t].float().view(H, V) st = state[0] @@ -449,8 +546,8 @@ def test_round1_zero_accepted(B, H): _assert_close( "state(cute vs cpu)", cute_out["recurrent_state"], cpu["recurrent_state"], atol=1e-4 ) - for name in ("qkg_cache", "v_cache", "beta_cache", "cs_q", "cs_k", "cs_v"): - _assert_close(f"{name}(cute vs cpu)", cute_out[name], cpu[name], atol=2e-2) + for name in ("k_cache", "g_cache", "v_cache", "beta_cache", "cs_q", "cs_k", "cs_v"): + _assert_close(f"{name}(cute vs cpu)", cute_out[name], cpu[name], atol=2e-2, rtol=4e-3) # Kernel vs the FLA sequential path (the in-tree fused verify math). _assert_close( @@ -465,11 +562,16 @@ def test_round1_zero_accepted(B, H): ) -@pytest.mark.parametrize("B,H", [(128, 2), (32, 12), (32, 6)], ids=lambda v: str(v)) -def test_round2_replay(B, H): +@pytest.mark.parametrize( + "B,H,M", + [(128, 2, 2), (32, 12, 2), (32, 6, 2), (2, 12, 7), (16, 12, 7)], + ids=lambda v: str(v), +) +def test_round2_replay(B, H, M): """Replay round: mixed num_accepted per request, chained from a CPU- golden round 1. Validates the kernel's cache-replay state math (the - path the drop's own self-check never exercised).""" + path the drop's own self-check never exercised). ``M=7`` covers the + production MTP7 verify (T_loop up to 15) at the TEP8 shape.""" data = make_conv_data(B, H, M=M, seed=7) T = data["T"] @@ -479,8 +581,8 @@ def test_round2_replay(B, H): # Round 2 inputs: fresh tokens, caches/state/conv from round 1. data2 = make_conv_data(B, H, M=M, seed=8) - for name in ("qkg_cache", "v_cache", "beta_cache"): - data2[name] = r1[name].cuda().contiguous() + for name in ("k_cache", "g_cache", "v_cache", "beta_cache"): + data2[name] = r1[name].cuda().to(data2[name].dtype).contiguous() for name in ("cs_q", "cs_k", "cs_v"): # Preserve the contract's dim-contiguous (transposed) layout. src = r1[name].cuda() @@ -515,6 +617,284 @@ def test_round2_replay(B, H): ) +def _round2_inputs(B, H, M): + """Round-2 inputs chained from a CPU-golden round 1 (caches, conv windows, committed state) + with mixed acceptance 0..M cycling across requests. Returns (data2, accept).""" + data = make_conv_data(B, H, M=M, seed=7) + r1 = cpu_reference(data) + data2 = make_conv_data(B, H, M=M, seed=8) + for name in ("k_cache", "g_cache", "v_cache", "beta_cache"): + data2[name] = r1[name].cuda().to(data2[name].dtype).contiguous() + for name in ("cs_q", "cs_k", "cs_v"): + src = r1[name].cuda() + dst = torch.empty( + src.shape[0], src.shape[2], src.shape[1], dtype=src.dtype, device=src.device + ).transpose(1, 2) + dst.copy_(src) + data2[name] = dst + data2["initial_state_kfirst"] = r1["recurrent_state"].cuda().contiguous() + data2["initial_state_cute"] = r1["recurrent_state"].cuda().permute(0, 1, 3, 2).contiguous() + accept = torch.arange(B, dtype=torch.int32) % (M + 1) + data2["num_accepted_tokens"] = accept.cuda() + return data2, accept + + +@pytest.mark.parametrize("B,H,M", [(2, 12, 7), (16, 12, 7)], ids=lambda v: str(v)) +def test_round2_replay_bf16_math(B, H, M, monkeypatch): + """bf16 tensor-core operand path (``TRTLLM_KDA_MTP_BF16_MATH=1``) on the MTP7 replay round. + The conv pools and replay caches do not depend on the recurrence precision and must match + the 3xTF32 kernel bit for bit; the committed state and the outputs are held to the tolerance + of the bf16 tensor-core references (FLA, prefill).""" + data2, accept = _round2_inputs(B, H, M) + T = data2["T"] + monkeypatch.setenv("TRTLLM_KDA_MTP_BF16_MATH", "0") + ref = cute_run(data2) + monkeypatch.setenv("TRTLLM_KDA_MTP_BF16_MATH", "1") + fast = cute_run(data2) + cpu2 = {k: v.cuda() for k, v in cpu_reference(data2).items()} + rows = torch.cat( + [torch.arange(n * T + int(accept[n]), n * T + int(accept[n]) + 1 + M) for n in range(B)] + ).cuda() + + for name in ("cs_q", "cs_k", "cs_v"): + _assert_close(f"{name}(fast vs 3xTF32)", fast[name], ref[name], atol=0.0) + for name in ("k_cache", "g_cache", "v_cache", "beta_cache"): + _assert_close(f"{name}(fast vs 3xTF32)", fast[name], ref[name], atol=0.0) + state_err = (fast["recurrent_state"] - cpu2["recurrent_state"]).abs().max().item() + ref_err = (ref["recurrent_state"] - cpu2["recurrent_state"]).abs().max().item() + out_err = (fast["out"][0, rows].float() - cpu2["out"][0, rows].float()).abs().max().item() + print( + f"[bf16_math B={B} H={H} M={M}] state max err vs fp32 golden " + f"{state_err:.3e} (3xTF32 kernel {ref_err:.3e}), out max err {out_err:.3e}" + ) + # measured max |err| vs the fp32 golden at (16,12,7): bf16 1.6e-2 (3xTF32: 4e-5) + tol = 3e-2 + _assert_close( + "state(fast vs cpu)", fast["recurrent_state"], cpu2["recurrent_state"], atol=tol, rtol=tol + ) + _assert_close( + "out(fast vs cpu)", fast["out"][0, rows], cpu2["out"][0, rows], atol=5e-2, rtol=5e-2 + ) + + +@pytest.mark.parametrize("B,H,M", [(2, 12, 7), (16, 12, 7)], ids=lambda v: str(v)) +def test_round2_replay_strided_state_cache(B, H, M): + """Replay round with the recurrent state living in a strided slice of a larger buffer + (slot stride 3 * H * V * K), like the runtime's per-layer view of the SSM cache. Covers the + kernel's own global-memory views (TMA state load/store), which only a non-contiguous state + exercises; ``B=2`` runs the cluster-pair form, ``B=16`` the single-CTA form.""" + data = make_conv_data(B, H, M=M, seed=7) + T = data["T"] + r1 = cpu_reference(data) + data2 = make_conv_data(B, H, M=M, seed=8) + for name in ("k_cache", "g_cache", "v_cache", "beta_cache"): + data2[name] = r1[name].cuda().to(data2[name].dtype).contiguous() + for name in ("cs_q", "cs_k", "cs_v"): + src = r1[name].cuda() + dst = torch.empty( + src.shape[0], src.shape[2], src.shape[1], dtype=src.dtype, device=src.device + ).transpose(1, 2) + dst.copy_(src) + data2[name] = dst + data2["initial_state_kfirst"] = r1["recurrent_state"].cuda().contiguous() + data2["initial_state_cute"] = r1["recurrent_state"].cuda().permute(0, 1, 3, 2).contiguous() + accept = torch.arange(B, dtype=torch.int32) % (M + 1) + data2["num_accepted_tokens"] = accept.cuda() + + state_shape = tuple(data2["initial_state_cute"].shape) # (B, H, V, K) + big = torch.zeros((B, 3) + state_shape[1:], dtype=torch.float32, device="cuda") + strided_state = big[:, 1] + assert not strided_state.is_contiguous() + + cpu2 = {k: v.cuda() for k, v in cpu_reference(data2).items()} + cute2 = cute_run(data2, state_override=strided_state) + + rows = torch.cat( + [torch.arange(n * T + int(accept[n]), n * T + int(accept[n]) + 1 + M) for n in range(B)] + ).cuda() + _assert_close("out2(cute vs cpu)", cute2["out"][0, rows], cpu2["out"][0, rows], atol=2e-2) + _assert_close( + "state2(cute vs cpu)", cute2["recurrent_state"], cpu2["recurrent_state"], atol=1e-4 + ) + # the neighbouring slices of the buffer must be untouched + assert torch.count_nonzero(big[:, 0]) == 0 and torch.count_nonzero(big[:, 2]) == 0 + + +@pytest.mark.parametrize("B,H,M", [(2, 12, 7), (16, 12, 7)], ids=lambda v: str(v)) +def test_round2_replay_bf16_state(B, H, M): + """Replay round with a bf16 recurrent-state pool: the kernel reads the bf16 state (exact in + tf32) and rounds the committed state to nearest bf16. The CPU golden starts from the same + bf16-rounded state; its fp32 result is compared against the kernel's bf16 output with a + half-ulp (2^-9) bf16 tolerance. ``B=2`` runs the cluster-pair form, ``B=16`` the single-CTA form.""" + data = make_conv_data(B, H, M=M, seed=7) + T = data["T"] + r1 = cpu_reference(data) + data2 = make_conv_data(B, H, M=M, seed=8) + for name in ("k_cache", "g_cache", "v_cache", "beta_cache"): + data2[name] = r1[name].cuda().to(data2[name].dtype).contiguous() + for name in ("cs_q", "cs_k", "cs_v"): + src = r1[name].cuda() + dst = torch.empty( + src.shape[0], src.shape[2], src.shape[1], dtype=src.dtype, device=src.device + ).transpose(1, 2) + dst.copy_(src) + data2[name] = dst + state_bf16 = r1["recurrent_state"].cuda().permute(0, 1, 3, 2).to(torch.bfloat16).contiguous() + data2["initial_state_cute"] = state_bf16 + data2["initial_state_kfirst"] = state_bf16.permute(0, 1, 3, 2).float().contiguous() + accept = torch.arange(B, dtype=torch.int32) % (M + 1) + data2["num_accepted_tokens"] = accept.cuda() + + cpu2 = {k: v.cuda() for k, v in cpu_reference(data2).items()} + cute2 = cute_run(data2) + assert cute2["state_v_first"].dtype == torch.bfloat16 + + rows = torch.cat( + [torch.arange(n * T + int(accept[n]), n * T + int(accept[n]) + 1 + M) for n in range(B)] + ).cuda() + _assert_close("out2(cute vs cpu)", cute2["out"][0, rows], cpu2["out"][0, rows], atol=2e-2) + _assert_close( + "state2(cute bf16 vs cpu fp32)", + cute2["recurrent_state"].float(), + cpu2["recurrent_state"], + atol=2e-3, + rtol=4e-3, + ) + + +@pytest.mark.parametrize("B,H,num_spec", [(4, 6, 2), (2, 12, 7)]) +def test_packed_projection_views_match_shifted_metadata_replay(B, H, num_spec): + """Compact row-strided projection views match the legacy shifted layout.""" + data = make_conv_data(B=B, H=H, M=num_spec, seed=10) + request_accepted = torch.arange(B, dtype=torch.int32, device="cuda") + request_accepted[-1] = num_spec + data["num_accepted_tokens"] = request_accepted + packed, legacy_rows = _pack_new_token_projection_views(data) + + # The production count is pool-level, while request rows may map to any + # state slot. Move every request's initial pools to a non-identity slot + # and scatter its accepted count to the same slot. + slots = torch.arange(B - 1, -1, -1, dtype=torch.int32, device="cuda") + packed["ssm_state_indices"] = slots + for name in ( + "initial_state_kfirst", + "initial_state_cute", + "k_cache", + "g_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): + request_order = packed[name] + slot_order = torch.empty_like(request_order) + slot_order[slots.long()] = request_order + packed[name] = slot_order + pool_accepted = torch.zeros(B, dtype=torch.int32, device="cuda") + pool_accepted[slots.long()] = request_accepted + packed["num_accepted_tokens"] = pool_accepted + + legacy = cute_run(data) + actual = cute_run(packed, packed_token_layout=True) + poison = torch.full( + (1, B * (num_spec + 1), H, data["V"]), + float("nan"), + dtype=torch.bfloat16, + device="cuda", + ) + poisoned = cute_run(packed, packed_token_layout=True, out=poison) + assert poisoned["out"].data_ptr() == poison.data_ptr() + assert torch.isfinite(poisoned["out"]).all() + _assert_close("poisoned packed out", poisoned["out"], actual["out"], atol=1e-5) + _assert_close( + "packed out", + actual["out"].reshape(-1, data["H"], data["V"]), + legacy["out"][0, legacy_rows], + atol=1e-5, + ) + + # Stage3a oracle: fuse only gated RMSNorm into the packed CuTe consumer. + # The raw recurrence output above is BF16, matching the explicit BF16 + # round-trip boundary retained in the fused epilogue. + from tensorrt_llm._torch.modules.mamba.layernorm_gated import rms_norm_gated_token_major + + fused_data = dict(packed) + fused_data["onorm_weight"] = torch.linspace( + 0.75, 1.25, data["V"], dtype=torch.float32, device="cuda" + ) + fused_data["onorm_eps"] = 1e-5 + expected_norm = rms_norm_gated_token_major( + actual["out"].reshape(-1, data["V"]), + packed["onorm_g"].reshape(-1, data["H"], data["V"]), + fused_data["onorm_weight"], + fused_data["onorm_eps"], + gate_activation="sigmoid", + ) + fused_norm = cute_run( + fused_data, + packed_token_layout=True, + fuse_output_norm=True, + ) + assert torch.isfinite(fused_norm["out"]).all() + fused_norm_diff = ( + fused_norm["out"].float().reshape_as(expected_norm) - expected_norm.float() + ).abs() + fused_norm_rel_l2 = ( + fused_norm_diff.norm() / expected_norm.float().norm().clamp_min(1e-12) + ).item() + print( + f" fused_norm B={B} H={H} M={num_spec}: " + f"max_abs={fused_norm_diff.max().item():.3e} rel_l2={fused_norm_rel_l2:.3e}" + ) + _assert_close( + "fused gated RMSNorm", + fused_norm["out"].reshape(-1, data["V"]), + expected_norm, + atol=2e-2, + rtol=1e-2, + ) + fused_fp8 = cute_run( + fused_data, + packed_token_layout=True, + quantize_output=True, + ) + expected_rows = B * (num_spec + 1) + assert fused_fp8["out"].dtype is torch.float8_e4m3fn + assert fused_fp8["out"].shape == (expected_rows, H * data["V"]) + assert torch.isfinite(fused_fp8["out"].float()).all() + assert fused_fp8["output_scale"].dtype is torch.uint8 + assert fused_fp8["output_scale"].shape == ((expected_rows + 127) // 128 * 128 * H * 4,) + expected_fp8, expected_scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0( + expected_norm.reshape(expected_rows, H * data["V"]), + ) + assert torch.equal(fused_fp8["out"], expected_fp8) + rows = torch.arange(expected_rows, device="cuda", dtype=torch.int64).view(-1, 1) + scale_columns = torch.arange(H * 4, device="cuda", dtype=torch.int64).view(1, -1) + offsets = ( + ((rows // 128 * H + scale_columns // 4) * 32 + rows % 32) * 4 + (rows % 128) // 32 + ) * 4 + scale_columns % 4 + assert torch.equal( + fused_fp8["output_scale"][offsets.flatten()], + expected_scale[offsets.flatten()], + ) + for name in ( + "recurrent_state", + "k_cache", + "g_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): + _assert_close( + f"{name}(packed vs shifted)", actual[name][slots.long()], legacy[name], atol=1e-5 + ) + _assert_close(f"{name}(norm vs raw)", fused_norm[name], actual[name], atol=1e-5) + _assert_close(f"{name}(fp8 vs raw)", fused_fp8[name], actual[name], atol=1e-5) + + @pytest.mark.parametrize("B,H", [(32, 12)], ids=lambda v: str(v)) def test_zero_accepted_hint_variant(B, H): """The zero_accepted_hint fast variant matches the general variant.""" @@ -524,7 +904,8 @@ def test_zero_accepted_hint_variant(B, H): for name in ( "out", "recurrent_state", - "qkg_cache", + "k_cache", + "g_cache", "v_cache", "beta_cache", "cs_q", @@ -534,15 +915,39 @@ def test_zero_accepted_hint_variant(B, H): _assert_close(f"{name}(fast vs general)", fast[name], general[name], atol=1e-5) -@pytest.mark.parametrize("field", ["ssm_state_indices", "cu_seqlens", "num_accepted_tokens"]) -def test_misaligned_scalar_metadata_after_aligned_warmup(field): - """Every scalar CuTe metadata argument accepts an offset int32 view. +def test_misaligned_state_indices_after_aligned_warmup(): + """Mixed context/generation metadata remains valid after op warmup.""" + data = make_conv_data(B=1, H=6, M=M, seed=13) + expected = cute_run(data) - Mixed context/generation batches hand the op a view into the shared buffer - that starts after one to three context rows: contiguous, int32, and not - 16-byte aligned. The op states 4-byte alignment for these, so the CuTe - bridge no longer has to be told a guarantee the caller cannot make. - """ + index_storage = torch.empty(2, dtype=torch.int32, device="cuda") + index_storage[0] = -1 + index_storage[1:].copy_(data["ssm_state_indices"]) + misaligned_indices = index_storage[1:] + assert misaligned_indices.is_contiguous() + assert misaligned_indices.data_ptr() % 16 != 0 + + misaligned_data = dict(data) + misaligned_data["ssm_state_indices"] = misaligned_indices + actual = cute_run(misaligned_data) + + for name in ( + "out", + "recurrent_state", + "k_cache", + "g_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): + _assert_close(f"{name}(misaligned vs aligned)", actual[name], expected[name], atol=1e-5) + + +@pytest.mark.parametrize("field", ("cu_seqlens", "num_accepted_tokens")) +def test_misaligned_scalar_metadata_after_aligned_warmup(field): + """Every scalar CuTe metadata argument accepts an offset int32 view.""" data = make_conv_data(B=1, H=6, M=M, seed=17) expected = cute_run(data) @@ -558,7 +963,17 @@ def test_misaligned_scalar_metadata_after_aligned_warmup(field): misaligned_data[field] = misaligned actual = cute_run(misaligned_data) - for name in ("out", "recurrent_state", "qkg_cache", "v_cache", "beta_cache"): + for name in ( + "out", + "recurrent_state", + "k_cache", + "g_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): _assert_close( f"{name}({field} misaligned vs aligned)", actual[name], expected[name], atol=1e-5 ) @@ -646,7 +1061,8 @@ def test_beta_cache_sibling_layer_after_aligned_warmup( for name in ( "out", "recurrent_state", - "qkg_cache", + "k_cache", + "g_cache", "v_cache", "beta_cache", "cs_q", @@ -680,7 +1096,7 @@ def _verify_after_round_one(B, H, num_spec, seed): data[name].copy_(data[name].bfloat16().float()) first = cute_run(data) data = make_conv_data(B, H, M=num_spec, seed=seed + 1) - for name in ("qkg_cache", "v_cache", "beta_cache", "cs_q", "cs_k", "cs_v"): + for name in ("k_cache", "g_cache", "v_cache", "beta_cache", "cs_q", "cs_k", "cs_v"): data[name] = first[name] data["initial_state_cute"] = first["state_v_first"] data["num_accepted_tokens"] = torch.arange(B, dtype=torch.int32, device="cuda") % (num_spec + 1) @@ -690,15 +1106,17 @@ def _verify_after_round_one(B, H, num_spec, seed): @pytest.mark.parametrize("B,H,num_spec", [(32, 12, 2), (32, 6, 2), (4, 6, 7)], ids=lambda v: str(v)) def test_commit_pending_drafts_matches_the_next_verify(B, H, num_spec): """Committing the pending drafts first leaves the next verify's results - unchanged, bit for bit. + unchanged. This is what the cache manager does across a suspend. ``kda_mtp_commit_pending_drafts`` folds the accepted drafts into the state, ``store_kda_replay_conv_window`` puts the conv window after them in the bf16 conv pool, and the resume seeds the replay caches from that pool - with nothing pending. The verify that follows must give exactly the - outputs, state, caches and windows of the verify that replays the drafts - itself.""" + with nothing pending. The verify that follows must give the caches and + windows of the verify that replays the drafts itself bit for bit, and its + outputs and state to the kernel's accuracy: the tensor-core verify computes + a chunk in one WY-form pass, so splitting the replayed drafts and the new + tokens into two chunks rounds differently.""" from tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops import ( kda_mtp_commit_pending_drafts, ) @@ -717,7 +1135,8 @@ def test_commit_pending_drafts_matches_the_next_verify(B, H, num_spec): kda_mtp_commit_pending_drafts( state, *cs, - data["qkg_cache"].clone(), + data["k_cache"].clone(), + data["g_cache"].clone(), data["v_cache"].clone(), data["beta_cache"].clone(), slots, @@ -742,7 +1161,8 @@ def test_commit_pending_drafts_matches_the_next_verify(B, H, num_spec): cs_q=cs[0], cs_k=cs[1], cs_v=cs[2], - qkg_cache=torch.zeros_like(data["qkg_cache"]), + k_cache=torch.zeros_like(data["k_cache"]), + g_cache=torch.zeros_like(data["g_cache"]), v_cache=torch.zeros_like(data["v_cache"]), beta_cache=torch.zeros_like(data["beta_cache"]), num_accepted_tokens=torch.zeros_like(accepted), @@ -752,17 +1172,86 @@ def test_commit_pending_drafts_matches_the_next_verify(B, H, num_spec): ) resumed = cute_run(resumed) - for name in ( - "out", - "state_v_first", - "qkg_cache", - "v_cache", - "beta_cache", - "cs_q", - "cs_k", - "cs_v", - ): + for name in ("k_cache", "g_cache", "v_cache", "beta_cache", "cs_q", "cs_k", "cs_v"): assert torch.equal(resumed[name], replayed[name]), name + _assert_close("out(commit vs replay)", resumed["out"], replayed["out"], atol=2e-2) + _assert_close( + "state(commit vs replay)", + resumed["state_v_first"], + replayed["state_v_first"], + atol=2e-4, + rtol=1e-4, + ) + + +def _replay_drafts_reference(data, slots, accepted): + """fp32 CPU replay of each selected slot's first ``accepted`` cached drafts onto its state, + the replay steps of ``cpu_reference``. Returns the K-first states of every slot.""" + H, K, V = data["H"], data["K"], data["V"] + state = data["initial_state_kfirst"].float().cpu().clone() + caches = { + name: data[name].float().cpu() for name in ("k_cache", "g_cache", "v_cache", "beta_cache") + } + for slot, count in zip(slots.tolist(), accepted.tolist()): + for h in range(H): + hk, hv = h * K, h * V + h_state = state[slot, h] + for t in range(count): + k_t = caches["k_cache"][slot, t, hk : hk + K] + decay = torch.exp(caches["g_cache"][slot, t, hk : hk + K]) + v_t = caches["v_cache"][slot, t, hv : hv + V] + beta_t = caches["beta_cache"][slot, t, h] + h_state = h_state * decay.unsqueeze(1) + v_new = (v_t - (h_state * k_t.unsqueeze(1)).sum(dim=0)) * beta_t + h_state = h_state + k_t.unsqueeze(1) * v_new.unsqueeze(0) + state[slot, h] = h_state + return state + + +@pytest.mark.parametrize( + "B,H,M", [(32, 12, 2), (32, 6, 2), (2, 12, 7), (16, 12, 7)], ids=lambda v: str(v) +) +def test_commit_pending_drafts_matches_the_replay_reference(B, H, M): + """The replay-only variant of the verify kernel commits, for each selected slot, the state + after its accepted drafts: the replay steps of a full verify, checked against the fp32 CPU + replay. It only reads the conv and replay caches, leaves a slot with nothing pending and the + unselected slots bit for bit, and covers both the cluster-pair and the one-CTA launch.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops import ( + kda_mtp_commit_pending_drafts, + ) + + data, accept = _round2_inputs(B, H, M) + # Every other slot, in reverse order, with its own accepted count (0 included). + slots = torch.arange(B - 1, -1, -2, dtype=torch.int32) + accepted = accept[slots.long()] + expected = _replay_drafts_reference(data, slots, accepted).cuda() + + state = data["initial_state_cute"].clone() + cs = [_dim_contiguous_clone(data[name]) for name in ("cs_q", "cs_k", "cs_v")] + caches = [data[name].clone() for name in ("k_cache", "g_cache", "v_cache", "beta_cache")] + kda_mtp_commit_pending_drafts( + state, + *cs, + *caches, + slots.cuda(), + accepted.cuda(), + M, + conv_width=data["W"], + ) + + committed = state.permute(0, 1, 3, 2) + _assert_close( + "state(commit vs cpu replay)", committed[slots.long()], expected[slots.long()], atol=1e-4 + ) + for slot, count in zip(slots.tolist(), accepted.tolist()): + if count == 0: + assert torch.equal(state[slot], data["initial_state_cute"][slot]), slot + others = [slot for slot in range(B) if slot not in slots.tolist()] + assert torch.equal(state[others], data["initial_state_cute"][others]) + for name, cache in zip(("cs_q", "cs_k", "cs_v"), cs): + assert torch.equal(cache, data[name]), name + for name, cache in zip(("k_cache", "g_cache", "v_cache", "beta_cache"), caches): + assert torch.equal(cache, data[name]), name def test_commit_pending_drafts_replays_the_per_token_records(): @@ -803,7 +1292,17 @@ def small(*shape, scale): unread = torch.empty(0, device="cuda") commit_kda_pending_drafts( - state, conv_pool, *cs, unread, unread, unread, slots, accepted, state_tok=state_tok + state, + conv_pool, + *cs, + # k / g / v / beta replay caches: unread with per-token states. + unread, + unread, + unread, + unread, + slots, + accepted, + state_tok=state_tok, ) for slot, count in zip(slots.tolist(), accepted.tolist()): diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_bf16_state_pool.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_bf16_state_pool.py index 9ec6908b2567..d0a27a505632 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_bf16_state_pool.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_bf16_state_pool.py @@ -49,10 +49,10 @@ class _LayerCache: """Only what the plain-decode path reads; no fused-verify replay caches. ``has_kda_replay_caches`` mirrors the manager's property, which is - ``kda_qkg_cache is not None``; the mixer reads it unguarded. + ``kda_k_cache is not None``; the mixer reads it unguarded. """ - kda_qkg_cache = None + kda_k_cache = None has_kda_replay_caches = False diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py index 7c4078286f36..1fa34bd0b26a 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py @@ -28,9 +28,8 @@ Identical hidden states are fed to both worlds; the fused world additionally uses fused QKV(G) and full- or low-rank gate projections with multi-stream overlap. With mixed per-request acceptance between rounds, matching round-2 outputs proves the -projection fusion and replay bookkeeping (shifted ``cu_seqlens`` layout, -conv-window seeding, pending-count plumbing) reproduce the promoted-state -semantics. +projection fusion, packed row/slot bookkeeping, conv-window seeding and the +in-kernel gated RMSNorm epilogue reproduce the promoted-state semantics. Requires 1 Blackwell GPU, fla-core, nvidia-cutlass-dsl. Skips otherwise. """ @@ -60,12 +59,12 @@ def _is_blackwell(): if not torch.cuda.is_available(): return False prop = torch.cuda.get_device_properties(0) - return prop.major * 10 + prop.minor in (100, 103) + return prop.major * 10 + prop.minor in (100, 103, 107) pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"), - pytest.mark.skipif(not _is_blackwell(), reason="needs sm100/sm103"), + pytest.mark.skipif(not _is_blackwell(), reason="needs sm100/sm103/sm107"), pytest.mark.skipif(not _HAVE_DEPS, reason=f"deps: {_DEP_ERR}"), ] @@ -163,9 +162,10 @@ def _conv_cache(section): kda_conv_q=_conv_cache(0), kda_conv_k=_conv_cache(1), kda_conv_v=_conv_cache(2), - kda_qkg_cache=torch.zeros(B, M, 3, d, device="cuda", dtype=torch.float32), - kda_v_cache=torch.zeros(B, M, d, device="cuda", dtype=torch.float32), - kda_beta_cache=torch.zeros(B, M, num_heads, device="cuda", dtype=torch.float32), + kda_k_cache=torch.zeros(B, M, d, device="cuda", dtype=torch.bfloat16), + kda_g_cache=torch.zeros(B, M, d, device="cuda", dtype=torch.float32), + kda_v_cache=torch.zeros(B, M, d, device="cuda", dtype=torch.bfloat16), + kda_beta_cache=torch.zeros(B, M, num_heads, device="cuda", dtype=torch.bfloat16), prev_num_accepted_tokens=torch.zeros(B, dtype=torch.int32, device="cuda"), has_kda_replay_caches=True, intermediate_conv_window=None, @@ -176,7 +176,7 @@ def _conv_cache(section): def _make_seq_layer_cache(B, *, num_heads=H): d = num_heads * K return SimpleNamespace( - kda_qkg_cache=None, + kda_k_cache=None, has_kda_replay_caches=False, intermediate_conv_window=torch.zeros( B, M + 1, 3 * d, W - 1, device="cuda", dtype=torch.bfloat16 @@ -193,6 +193,37 @@ def _promote_sequential(layer_cache, conv_pool, ssm_pool, accept): ssm_pool.copy_(layer_cache.intermediate_ssm[rows, accept]) +def _capture_mtp_verify(runtime): + """Record the kwargs of every fused verify launch, then run it.""" + calls = [] + mtp_verify = runtime._dispatch.mtp_verify + + def _capturing_mtp_verify(**kwargs): + calls.append(kwargs) + return mtp_verify(**kwargs) + + runtime._dispatch.mtp_verify = _capturing_mtp_verify + return calls + + +def _assert_fused_epilogue_call(call, runtime, use_full_rank_gate): + """The fused verify ran with the in-kernel output norm on packed rows.""" + assert call["fuse_output_norm"] + assert call["packed_token_layout"] + assert call["cu_seqlens"] is None + assert not call["mxfp8_output"] + onorm_g = call["onorm_g"] + assert onorm_g.stride(-1) == 1 and onorm_g.stride(-2) == K + if use_full_rank_gate: + # The gate and q/k/v stay strided views of the fused QKVG output. + row_stride = 4 * runtime.proj_size + assert onorm_g.stride(1) == row_stride + for name in ("x_q", "x_k", "x_v"): + assert call[name].stride(1) == row_stride + assert call[name].stride(-1) == 1 + assert call["beta"].stride(-1) == 1 + + def _rep(name, a, b): a, b = a.float(), b.float() cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item() @@ -247,6 +278,8 @@ def test_fused_vs_sequential_two_rounds( else: assert rt_fused._qkvg_proj_weight is not None assert rt_fused._bfa_proj_weight is not None + assert rt_fused._onorm_w_f32 is not None + mtp_calls = _capture_mtp_verify(rt_fused) slot_indices = torch.arange(B, dtype=torch.int32, device="cuda") conv_pool_seq, ssm_pool_seq = _make_pools(B, seed=2, num_heads=num_heads) @@ -288,6 +321,7 @@ def tokens(scale=0.5): x1, T, cache_fused, conv_pool_fused, ssm_pool_fused, slot_indices ) ) + _assert_fused_epilogue_call(mtp_calls[-1], rt_fused, use_full_rank_gate) print("round 1:") ok &= _rep("out", out1_fused, out1_seq) @@ -315,6 +349,8 @@ def tokens(scale=0.5): output=core2_fused, ) assert result2_fused is core2_fused + assert len(mtp_calls) == 2 + _assert_fused_epilogue_call(mtp_calls[-1], rt_fused, use_full_rank_gate) out2_fused = rt_fused._project_output(core2_fused) print("round 2 (mixed replay):") ok &= _rep("out", out2_fused, out2_seq) @@ -328,3 +364,117 @@ def tokens(scale=0.5): ok &= _rep("committed ssm", ssm_pool_fused, ssm_pool_seq) assert ok + + +def _run_one_verify_round(runtime, x, batch, seed, **verify_kwargs): + conv_pool, ssm_pool = _make_pools(batch, seed=seed) + layer_cache = _make_fused_layer_cache(batch, conv_pool) + slot_indices = torch.arange(batch, dtype=torch.int32, device="cuda") + return runtime.forward_verify( + x, M + 1, layer_cache, conv_pool, ssm_pool, slot_indices, **verify_kwargs + ) + + +@torch.no_grad() +def test_fused_verify_falls_back_without_output_norm_weights(): + """Without finalized o_norm weights the verify keeps the Python output gate.""" + batch = 2 + runtime = _make_runtime(seed=5) + runtime.finalize_decode_weights() + x = (torch.randn(batch * (M + 1), HIDDEN, device="cuda") * 0.5).to(torch.bfloat16) + fused_out = runtime._project_output(_run_one_verify_round(runtime, x, batch, seed=6)) + + runtime._onorm_w_f32 = None + calls = _capture_mtp_verify(runtime) + fallback_out = runtime._project_output(_run_one_verify_round(runtime, x, batch, seed=6)) + + assert len(calls) == 1 + assert "fuse_output_norm" not in calls[0] + assert calls[0]["cu_seqlens"] is not None + assert _rep("fallback vs fused epilogue", fallback_out, fused_out) + + +class _PrequantizedProjectionStub(torch.nn.Module): + """o_proj stand-in that records the prequantized pair it receives.""" + + supports_prequantized_input = True + + def __init__(self, out_features): + super().__init__() + self.out_features = out_features + self.calls = [] + + def forward(self, x): + return x.new_zeros(x.shape[0], self.out_features) + + def forward_prequantized(self, activation, activation_scale): + self.calls.append((activation, activation_scale)) + return torch.zeros( + activation.shape[0], self.out_features, dtype=torch.bfloat16, device="cuda" + ) + + +@torch.no_grad() +@pytest.mark.parametrize("allow, with_output", [(True, False), (False, False), (True, True)]) +def test_fused_verify_mxfp8_output_gating(allow, with_output): + """MXFP8 output only when allowed, unbuffered and accepted by o_proj.""" + batch = 2 + rows = batch * (M + 1) + runtime = _make_runtime(seed=9) + runtime.finalize_decode_weights() + stub = _PrequantizedProjectionStub(HIDDEN) + runtime.o_proj = stub + calls = [] + + def _fake_mtp_verify(**kwargs): + calls.append(kwargs) + if kwargs["mxfp8_output"]: + activation = torch.zeros(rows, H * K, device="cuda").to(torch.float8_e4m3fn) + return activation, torch.zeros(128 * H * 4, dtype=torch.uint8, device="cuda") + return torch.zeros(1, rows, H, K, dtype=torch.bfloat16, device="cuda") + + runtime._dispatch.mtp_verify = _fake_mtp_verify + x = torch.zeros(rows, HIDDEN, dtype=torch.bfloat16, device="cuda") + output = x.new_empty(rows, H, K) if with_output else None + core = _run_one_verify_round( + runtime, x, batch, seed=10, output=output, allow_prequantized_output=allow + ) + out = runtime._project_output(core) + + expect_mxfp8 = allow and not with_output + assert calls[-1]["fuse_output_norm"] + assert calls[-1]["mxfp8_output"] is expect_mxfp8 + assert len(stub.calls) == int(expect_mxfp8) + assert out.shape == (rows, HIDDEN) + if with_output: + assert core is output + + +@torch.no_grad() +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability(0) != (10, 7), + reason="the MXFP8 verify epilogue feeds the Rubin (SM107) CuTe MXFP8 GEMM", +) +def test_fused_verify_mxfp8_output_matches_bf16_projection(): + """SM107: epilogue-quantized o_proj input matches the BF16 core path.""" + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE + from tensorrt_llm._torch.models.modeling_kimi_linear import _Fp8BlockScaleWeightReadLinear + + if not IS_CUTLASS_DSL_RUBIN_AVAILABLE: + pytest.skip("needs the Rubin CuTe DSL") + batch = 4 + runtime = _make_runtime(seed=11) + runtime.finalize_decode_weights() + runtime.o_proj = _Fp8BlockScaleWeightReadLinear.from_linear(runtime.o_proj) + assert runtime.o_proj.supports_prequantized_input + calls = _capture_mtp_verify(runtime) + x = (torch.randn(batch * (M + 1), HIDDEN, device="cuda") * 0.5).to(torch.bfloat16) + + bf16_out = runtime._project_output(_run_one_verify_round(runtime, x, batch, seed=12)) + mxfp8_core = _run_one_verify_round(runtime, x, batch, seed=12, allow_prequantized_output=True) + mxfp8_out = runtime._project_output(mxfp8_core) + + assert [call["mxfp8_output"] for call in calls] == [False, True] + assert mxfp8_core.activation.dtype is torch.float8_e4m3fn + assert mxfp8_core.scale.dtype is torch.uint8 + assert _rep("mxfp8 vs bf16 o_proj input", mxfp8_out, bf16_out) diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py index 167200bf9d37..48e4c29b8ca5 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py @@ -113,6 +113,7 @@ def forward_verify( ssm_pool: torch.Tensor, slot_indices: torch.Tensor, output: torch.Tensor | None = None, + allow_prequantized_output: bool = False, ) -> torch.Tensor: captured["slot_indices"] = slot_indices return torch.zeros_like(hidden_states) diff --git a/tests/unittest/_torch/moe/test_kimi_k3_mlp.py b/tests/unittest/_torch/moe/test_kimi_k3_mlp.py index 7a4de96dad62..83d216c933b4 100644 --- a/tests/unittest/_torch/moe/test_kimi_k3_mlp.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_mlp.py @@ -21,6 +21,8 @@ import os import sys +from collections.abc import Iterator +from contextlib import contextmanager from types import SimpleNamespace import cloudpickle @@ -204,13 +206,21 @@ def test_gated_mlp_supports_fused_situ(situ_beta, situ_linear_beta): @pytest.mark.parametrize( - "attention_dp,tp_size,rank,expected_shared_tp,expected_shared_rank", + "attention_dp,tp_size,rank,expected_shared_tp,expected_shared_rank,moe_backend,scheduler_kind", [ - (True, 8, 7, 1, 0), - (False, 1, 0, 1, 0), - (False, 8, 7, 8, 7), + (True, 8, 7, 1, 0, "TRTLLM", "EXTERNAL_COMM"), + (False, 1, 0, 1, 0, "TRTLLM", "EXTERNAL_COMM"), + (False, 8, 7, 8, 7, "TRTLLM", "EXTERNAL_COMM"), + (True, 8, 7, 1, 0, "MEGAMOE_CUTEDSL", "FUSED_COMM"), + (False, 8, 7, 8, 7, "MEGAMOE_CUTEDSL", "FUSED_COMM"), + ], + ids=[ + "attention_dp", + "single_rank", + "direct_tp", + "megamoe_attention_dp", + "megamoe_direct_tp", ], - ids=["attention_dp", "single_rank", "direct_tp"], ) def test_kimi_k3_shared_expert_parallel_construction( monkeypatch, @@ -219,6 +229,8 @@ def test_kimi_k3_shared_expert_parallel_construction( rank, expected_shared_tp, expected_shared_rank, + moe_backend, + scheduler_kind, ): """Shared experts are replicated or sharded for the selected parallel mode.""" from tensorrt_llm._torch import distributed @@ -241,7 +253,7 @@ def __init__(self): # the runtime reads it unguarded; the stand-in must too. self.backend = SimpleNamespace( initial_local_expert_ids=[0, 1, 2, 3], - scheduler_kind=MoESchedulerKind.EXTERNAL_COMM, + scheduler_kind=MoESchedulerKind[scheduler_kind], ) self.comm = None self.layer_load_balancer = None @@ -268,7 +280,7 @@ def __init__(self): model_config = ModelConfig( mapping=mapping, quant_config=QuantConfig(), - moe_backend="TRTLLM", + moe_backend=moe_backend, ) config = _runtime_config() aux_stream_dict = _make_aux_stream_dict() @@ -280,6 +292,16 @@ def __init__(self): ) shared = runtime.shared_experts + assert runtime._defer_shared_allreduce is ( + not attention_dp + and tp_size > 1 + and moe_backend == "MEGAMOE_CUTEDSL" + and scheduler_kind == "FUSED_COMM" + ) + assert (runtime.moe_routed_event is not None) is runtime._defer_shared_allreduce + if runtime._defer_shared_allreduce: + assert runtime.moe_routed_event is not runtime.moe_main_event + assert runtime.moe_routed_event is not runtime.moe_shared_event assert create_moe_kwargs["aux_stream_dict"] is aux_stream_dict assert AuxStreamType.MoeBalancer in create_moe_kwargs["aux_stream_dict"] assert runtime.shared_expert_stream is aux_stream_dict[AuxStreamType.MoeShared] @@ -302,6 +324,246 @@ def __init__(self): assert shared.down_proj.weight.shape == (config.hidden_size, local_intermediate) +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "defer,multi_stream,is_compiling,has_aux_stream", + [ + (False, True, False, True), + (True, True, False, True), + (True, False, False, True), + ], + ids=["ordinary_shared", "parallel", "serial"], +) +def test_kimi_k3_shared_allreduce_ordering( + monkeypatch: pytest.MonkeyPatch, + defer: bool, + multi_stream: bool, + is_compiling: bool, + has_aux_stream: bool, +) -> None: + """Check shared-allreduce dependencies and result propagation without CUDA work.""" + from tensorrt_llm._torch.distributed import AllReduceParams + from tensorrt_llm._torch.models import modeling_kimi_linear + from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream + + calls = [] + current_stream = "main" + aux_stream = object() if has_aux_stream else None + parallel = multi_stream and has_aux_stream and not is_compiling + overlap = defer and parallel + + class _Event: + def __init__(self, name: str, enabled: bool) -> None: + self.name = name + self.enabled = enabled + self.recorded = False + + def record(self) -> None: + assert self.enabled + assert not self.recorded + self.recorded = True + calls.append((f"{self.name}_record", current_stream)) + + def wait(self) -> None: + assert self.enabled + assert self.recorded + self.recorded = False + calls.append((f"{self.name}_wait", current_stream)) + + @contextmanager + def use_stream(stream: object) -> Iterator[None]: + nonlocal current_stream + assert parallel and stream is aux_stream + previous_stream = current_stream + current_stream = "aux" + try: + yield + finally: + current_stream = previous_stream + + class _Projection(nn.Module): + def __init__(self, name: str) -> None: + super().__init__() + self.name = name + + def forward(self, value: torch.Tensor) -> torch.Tensor: + calls.append((self.name, current_stream)) + return value + + def routed_experts(value: torch.Tensor, logits: torch.Tensor, **kwargs) -> torch.Tensor: + calls.append(("routed_moe", current_stream)) + return value * 2 + + def shared_allreduce(value: torch.Tensor) -> torch.Tensor: + calls.append(("shared_allreduce", current_stream)) + return value + 5 + + class _Shared(nn.Module): + def __init__(self) -> None: + super().__init__() + self.down_proj = SimpleNamespace(all_reduce=shared_allreduce) + + def forward( + self, value: torch.Tensor, final_all_reduce_params: AllReduceParams | None = None + ) -> torch.Tensor: + calls.append(("shared_local", current_stream)) + partial = value * 3 + if defer: + assert final_all_reduce_params.enable_allreduce is False + return partial + assert final_all_reduce_params is None + return self.down_proj.all_reduce(partial) + + runtime = modeling_kimi_linear.KimiK3MoERuntime.__new__(modeling_kimi_linear.KimiK3MoERuntime) + nn.Module.__init__(runtime) + runtime.gate = SimpleNamespace(compute_logits=lambda value: value) + runtime.routed_experts = routed_experts + runtime.routed_expert_down_proj = _Projection("routed_down") + runtime.routed_expert_norm = _Projection("routed_norm") + runtime.routed_expert_up_proj = _Projection("routed_up") + runtime.shared_experts = _Shared() + runtime._reduce_routed_output = False + runtime._defer_shared_allreduce = defer + runtime.shared_expert_stream = aux_stream + runtime.moe_main_event = _Event("fork", parallel) + runtime.moe_shared_event = _Event("shared_done", parallel) + midpoint_event = _Event("midpoint", overlap) if defer else None + runtime.moe_routed_event = midpoint_event + runtime.latent_down_stream = None + runtime.moe_latent_fork_event = None + runtime.moe_latent_ready_event = None + + monkeypatch.setattr(torch.cuda, "stream", use_stream) + # maybe_execute_in_parallel records the aux outputs on the calling stream + # after the join; keep that query and lookup off the CUDA runtime too. + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: current_stream) + monkeypatch.setattr(torch.cuda, "Event", lambda: pytest.fail("Events must be initialized once")) + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: is_compiling) + + shared_stream = "aux" if parallel else "main" + + # Reusing the same owner must record fresh dependencies without replacing events. + for _ in range(2): + calls.clear() + with with_multi_stream(multi_stream): + output = runtime(torch.ones(2, 4)) + assert runtime.moe_routed_event is midpoint_event + assert sum(name == "shared_allreduce" for name, _ in calls) == 1 + allreduce = calls.index(("shared_allreduce", shared_stream)) + assert calls.index(("shared_local", shared_stream)) < allreduce + if overlap: + # Fence all routed chunks, but leave the norm/up tail free to overlap. + recorded = calls.index(("midpoint_record", "main")) + assert calls.index(("routed_moe", "main")) < recorded + assert recorded < calls.index(("routed_norm", "main")) + assert recorded < calls.index(("midpoint_wait", "aux")) < allreduce + if parallel: + # The final join must cover the shared all-reduce. + completed = calls.index(("shared_done_record", "aux")) + assert allreduce < completed < calls.index(("shared_done_wait", "main")) + elif defer: + assert calls.index(("routed_moe", "main")) < allreduce + torch.testing.assert_close(output, torch.full((2, 4), 10.0)) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("multi_stream", [True, False], ids=["forked", "serial"]) +def test_kimi_k3_latent_down_projection_overlaps_routing( + monkeypatch: pytest.MonkeyPatch, multi_stream: bool +) -> None: + """The latent down-projection forks onto the CuteDSL quantization stream + before the router GEMM, and the MoE receives the event to join.""" + from tensorrt_llm._torch.models import modeling_kimi_linear + from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream + + calls = [] + current_stream = "main" + latent_stream = object() + + class _Event: + def __init__(self, name: str) -> None: + self.name = name + + def record(self) -> None: + calls.append((f"{self.name}_record", current_stream)) + + def wait(self) -> None: + calls.append((f"{self.name}_wait", current_stream)) + + @contextmanager + def use_stream(stream: object) -> Iterator[None]: + nonlocal current_stream + assert stream is latent_stream + previous_stream = current_stream + current_stream = "latent" + try: + yield + finally: + current_stream = previous_stream + + class _Projection(nn.Module): + def __init__(self, name: str) -> None: + super().__init__() + self.name = name + + def forward(self, value: torch.Tensor) -> torch.Tensor: + calls.append((self.name, current_stream)) + return value + 1 + + received = {} + + def routed_experts(value: torch.Tensor, logits: torch.Tensor, **kwargs) -> torch.Tensor: + calls.append(("routed_moe", current_stream)) + received.update(kwargs) + return value + + def compute_logits(value: torch.Tensor) -> torch.Tensor: + calls.append(("router_gemm", current_stream)) + return value + + runtime = modeling_kimi_linear.KimiK3MoERuntime.__new__(modeling_kimi_linear.KimiK3MoERuntime) + nn.Module.__init__(runtime) + runtime.gate = SimpleNamespace(compute_logits=compute_logits) + runtime.routed_experts = routed_experts + runtime.routed_expert_down_proj = _Projection("routed_down") + runtime.routed_expert_norm = _Projection("routed_norm") + runtime.routed_expert_up_proj = _Projection("routed_up") + runtime.shared_experts = lambda value, final_all_reduce_params=None: value * 0 + runtime._reduce_routed_output = False + runtime._defer_shared_allreduce = False + runtime.shared_expert_stream = None + runtime.moe_main_event = _Event("fork") + runtime.moe_shared_event = _Event("shared_done") + runtime.moe_routed_event = None + runtime.latent_down_stream = latent_stream + runtime.moe_latent_fork_event = _Event("latent_fork") + runtime.moe_latent_ready_event = _Event("latent_ready") + + monkeypatch.setattr(torch.cuda, "stream", use_stream) + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False) + + with with_multi_stream(multi_stream): + output = runtime(torch.zeros(2, 4)) + + # down (+1), norm (+1), up (+1); the shared branch contributes zero. + torch.testing.assert_close(output, torch.full((2, 4), 3.0)) + if multi_stream: + assert calls.index(("latent_fork_record", "main")) < calls.index( + ("latent_fork_wait", "latent") + ) + assert calls.index(("latent_fork_wait", "latent")) < calls.index(("routed_down", "latent")) + assert calls.index(("routed_down", "latent")) < calls.index( + ("latent_ready_record", "latent") + ) + # The router GEMM is issued on main after the fork, not after the join. + assert calls.index(("latent_fork_record", "main")) < calls.index(("router_gemm", "main")) + assert received["input_ready_event"] is runtime.moe_latent_ready_event + else: + assert ("routed_down", "main") in calls + assert "input_ready_event" not in received + + def _make_kimi_k3_moe_weights(config): generator = torch.Generator().manual_seed(2026) @@ -540,6 +802,62 @@ def _run_kimi_k3_moe_multi_rank_worker(attention_dp): os.environ["TRTLLM_FORCE_COMM_METHOD"] = previous_comm_method +@requires_cuda +@pytest.mark.parametrize( + "case, num_tokens, expected", + [ + ("eligible", 1, True), + ("eligible", 16, True), + ("eligible", 17, False), + ("non_linear_down", 1, False), + ("other_backend", 1, False), + ], +) +def test_kimi_k3_router_down_fusion_gate( + monkeypatch: pytest.MonkeyPatch, case: str, num_tokens: int, expected: bool +) -> None: + """``_can_fuse_router_down`` admits only the validated router + latent + down-projection fusion: 1..16 tokens, a bias-free BF16 ``nn.Linear`` down + projection (FP8 conversion replaces it), and MEGAMOE_CUTEDSL with fused + communication on SM107.""" + from tensorrt_llm._torch.models import modeling_kimi_linear + + monkeypatch.setattr(modeling_kimi_linear, "get_sm_version", lambda: 107) + monkeypatch.setattr(modeling_kimi_linear, "_K3_DISABLE_MIN_LATENCY_LATENT_PROJ", False) + device = torch.device("cuda") + hidden_size, latent_size, num_experts = 7168, 3584, 896 + + class _Fp8ConvertedDown(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter( + torch.empty(latent_size, hidden_size, device=device, dtype=torch.bfloat16), + requires_grad=False, + ) + self.bias = None + + runtime = modeling_kimi_linear.KimiK3MoERuntime.__new__(modeling_kimi_linear.KimiK3MoERuntime) + nn.Module.__init__(runtime) + runtime.gate = SimpleNamespace( + weight=torch.empty(num_experts, hidden_size, device=device, dtype=torch.bfloat16) + ) + if case == "non_linear_down": + runtime.routed_expert_down_proj = _Fp8ConvertedDown() + else: + runtime.routed_expert_down_proj = nn.Linear( + hidden_size, latent_size, bias=False, device=device, dtype=torch.bfloat16 + ) + runtime.routed_experts = SimpleNamespace( + use_dp=False, + parallel_size=8, + moe_backend="CUTEDSL" if case == "other_backend" else "MEGAMOE_CUTEDSL", + backend=SimpleNamespace(scheduler_kind=modeling_kimi_linear.MoESchedulerKind.FUSED_COMM), + ) + hidden_states = torch.empty(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) + + assert runtime._can_fuse_router_down(hidden_states) is expected + + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs 2 GPUs to run this test") @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) diff --git a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py index 384474fd1448..e737d86b502e 100644 --- a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py @@ -1388,7 +1388,7 @@ def test_tp16_nvfp4_padded_loaders_preserve_rank_ownership(): #: importing ``cute_dsl_utils`` pulls in the DSL package, which appends its own #: directory to ``sys.path``, and this repository fails the whole pytest #: session when a test file does that at collection time. -_NVFP4_SITU_BACKENDS = ["CUTLASS", "TRTLLM", "CUTEDSL"] +_NVFP4_SITU_BACKENDS = ["CUTLASS", "TRTLLM", "CUTEDSL", "CUTEDSL_FC12"] def _skip_if_backend_unavailable(moe_backend): @@ -1399,20 +1399,30 @@ def _skip_if_backend_unavailable(moe_backend): early puts the wheel's package directory on ``sys.path`` for every other test file in the session. """ + if moe_backend == "CUTEDSL_FC12": + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE + + # The fused FC1+FC2 kernel is Rubin-only; can_implement() turns every + # other SM down, so there is no end-to-end path to exercise. + if get_sm_version() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE: + pytest.skip("CUTEDSL_FC12 needs SM107 and CuTe DSL Rubin support") + return if moe_backend != "CUTEDSL": return from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE if not IS_CUTLASS_DSL_AVAILABLE: pytest.skip("CuteDSL MoE requires the CuTe DSL wheel") - # nvfp4_moe_supported admits every SM >= 100, but this integration enables - # SiTU only on Blackwell. Match can_implement() rather than exercising an - # end-to-end path that has not been enabled on SM107. - if get_sm_version() not in (100, 103): - pytest.skip( - f"CuteDSL SiTU MoE needs the Blackwell act-fusion kernel " - f"(SM100/SM103), got SM{get_sm_version()}" - ) + # nvfp4_moe_supported admits every SM >= 100; the CuteDSL act-fusion + # kernels with a SiTU epilogue exist for SM100/SM103 and SM107 (Rubin). + sm_version = get_sm_version() + if sm_version not in (100, 103, 107): + pytest.skip(f"CuteDSL SiTU MoE needs SM100/SM103/SM107, got SM{sm_version}") + if sm_version == 107: + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE + + if not IS_CUTLASS_DSL_RUBIN_AVAILABLE: + pytest.skip("CuteDSL SiTU MoE on SM107 needs CuTe DSL Rubin support") def _make_nvfp4_expert_bank(num_experts, intermediate, hidden, seed=907): @@ -2202,6 +2212,72 @@ def fixed_tactic(custom_op, runners, tuning_config, inputs, **kwargs): assert not torch.allclose(first, output), f"soft-caps {betas} reused the first output" +@nvfp4_moe_supported +def test_cutedsl_fc12_survives_autotune_warmup(): + """Drive FC12 through a cold AutoTuner warmup. + + ``do_preparation`` only runs from inside ``AutoTuner.choose_one``'s + profiling path. Other FC12 tests call ``moe.forward`` outside any + ``autotune()`` context, so that branch never runs there. ``clear_cache()`` + is load-bearing: in tuning mode a cache hit also returns early. + """ + from tensorrt_llm._torch.autotuner import AutoTuner, AutoTunerStatistics, autotune + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE + + if get_sm_version() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE: + pytest.skip("CUTEDSL_FC12 needs SM107 and CuTe DSL Rubin support") + + custom_op = "CuteDslFc12FusedMoE::run_moe_nvfp4" + num_experts, hidden, inter = _TP_EXPERTS, _TP_HIDDEN, _TP_INTERMEDIATE + gate = _make_test_gate(num_experts=num_experts) + + torch.manual_seed(91) + x = torch.randn(8, hidden, dtype=torch.bfloat16, device="cuda") * 0.5 + act_scale = float(x.abs().max().float() / (448 * 6)) + w1 = [ + torch.randn(inter, hidden, dtype=torch.bfloat16, device="cuda") * 0.05 + for _ in range(num_experts) + ] + w3 = [ + torch.randn(inter, hidden, dtype=torch.bfloat16, device="cuda") * 0.05 + for _ in range(num_experts) + ] + w2 = [ + torch.randn(hidden, inter, dtype=torch.bfloat16, device="cuda") * 0.05 + for _ in range(num_experts) + ] + bank = [_quantize_expert_to_nvfp4(w1[e], w2[e], w3[e], act_scale) for e in range(num_experts)] + + moe = _make_nvfp4_moe(gate, num_experts=num_experts, moe_backend="CUTEDSL_FC12") + _stream_nvfp4_bank(moe, bank) + moe.backend._weights_transformed = False + moe.post_load_weights() + + router_logits = gate.compute_logits(x) + + autotuner = AutoTuner.get() + saved = (autotuner.warmup, autotuner.repeat, autotuner.stream_delay_micro_secs) + autotuner.clear_cache() + # The singleton outlives the test; fresh statistics keep the positive + # control below from being satisfied by an earlier test. + autotuner.stats = AutoTunerStatistics() + autotuner.warmup, autotuner.repeat, autotuner.stream_delay_micro_secs = 0, 1, 10 + try: + # No cache_path: a shared on-disk cache would turn a cold miss into a + # hit and skip the branch under test. + with torch.inference_mode(), autotune(): + tuned = moe.forward(x, router_logits, all_rank_num_tokens=None) + finally: + autotuner.warmup, autotuner.repeat, autotuner.stream_delay_micro_secs = saved + + assert tuned.shape == x.shape + assert torch.isfinite(tuned.float()).all() + # Positive control: without it the test passes when nothing is profiled. + profiled = autotuner.stats.tuned_op_profiled_configs.get(custom_op, 0) + assert profiled > 0, f"{custom_op} was never profiled; the warmup branch did not run" + assert not autotuner.stats.failed_profiling_count.get(custom_op, set()) + + def test_fp8_block_scaled_dequantization(): """FP8_PB_WO attention weights must be dequantized, never reinterpreted. diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index adbe7eaf2f0b..ed26f89c4268 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -20,6 +20,7 @@ import logging import os import sys +from collections.abc import Callable from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace @@ -73,6 +74,10 @@ CuteDslFusedMoENvfp4Runner, ) from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE +from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl_fc12 import ( + CuteDslFc12FusedMoE, + CuteDslFc12FusedMoENvfp4Runner, +) from tensorrt_llm._torch.moe.fused_moe.fused_moe_cutlass import CutlassFusedMoE from tensorrt_llm._torch.moe.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE from tensorrt_llm._torch.moe.fused_moe.fused_moe_marlin import MarlinFusedMoE @@ -2441,6 +2446,172 @@ def test_megamoe_cutedsl_minimax_m3_swiglu_bias_numerics() -> None: check_accuracy(output, ref_output, rtol=0.1, atol=0.1, percent=0.94) +@pytest.fixture(scope="module") +def _run_megamoe_prologue_helper() -> Callable[[int, bool], torch.Tensor]: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7): + pytest.skip("The scheduler helper regression requires an SM107 CUDA GPU") + + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + + if not IS_CUTLASS_DSL_AVAILABLE: + pytest.skip("cutlass-dsl is not available") + + import cuda.bindings.driver as cuda + import cutlass.cute as cute + import cutlass.utils as utils + from cutlass.cute.runtime import from_dlpack + from cutlass.cutlass_dsl import Boolean, Int32 + from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + + from tensorrt_llm._torch.cute_dsl_kernels.cutedsl_megamoe.api import ImplDesc, ProblemDesc + from tensorrt_llm._torch.cute_dsl_kernels.cutedsl_megamoe.helpers.device_workspace import ( + DeviceWorkspace, + ) + from tensorrt_llm._torch.cute_dsl_kernels.cutedsl_megamoe.helpers.smem_workspace import ( + SmemWorkspace, + ) + from tensorrt_llm._torch.cute_dsl_kernels.cutedsl_megamoe.kernel_src.schedulers.fc12_mapping import ( + BlockPhase, + map_phase_interleaved_fc12_work_id, + resolve_phase_interleaved_fc1_claim_target, + ) + from tensorrt_llm._torch.cute_dsl_kernels.cutedsl_megamoe.kernel_src.schedulers.fc12_scheduler import ( + PhaseInterleavedFc12Scheduler, + ) + + class PrologueProbe: + def __init__(self) -> None: + self.scheduler = PhaseInterleavedFc12Scheduler( + ProblemDesc(dict(expert_count=1, intermediate_gateup_size=128, hidden_size=128)), + ImplDesc( + dict( + num_scheduler_consumer_threads=32, + mma_tiler_mnk=(128, 128, 64), + cluster_shape_mn=(1, 1), + use_2cta_instrs=False, + hint=1, + token_padding_block=128, + sf_padding_block=128, + work_id_mode="atomic_counter", + is_swap_ab=False, + launch_cluster_count=3, + ) + ), + ) + assert self.scheduler.minimum_global_fc1_claims == 3 + self.smem = SmemWorkspace() + self.device = DeviceWorkspace() + self.scheduler.register_smem_regions(self.smem) + self.scheduler.register_device_workspace(self.device) + self.smem.finalize() + self.device.finalize() + alignment = self.smem.base_alignment + self.smem_launch_bytes = ( + (self.smem.total_bytes + alignment - 1) // alignment * alignment + ) + assert self.smem_launch_bytes >= self.smem.storage_class().size_in_bytes() + + @cute.jit + def __call__( + self, + counts: cute.Tensor, + workspace: cute.Tensor, + exhausted: cute.Tensor, + output: cute.Tensor, + stream: cuda.CUstream, + ) -> None: + self.kernel(counts, workspace, exhausted, output).launch( + grid=(1, 1, 1), block=(32, 1, 1), smem=self.smem_launch_bytes, stream=stream + ) + + @cute.kernel + def kernel( + self, + counts: cute.Tensor, + workspace: cute.Tensor, + exhausted: cute.Tensor, + output: cute.Tensor, + ) -> None: + storage = utils.SmemAllocator().allocate(self.smem.storage_class()) + self.device.assign_device_members(workspace.iterator) + self.scheduler.assign_device_members( + expert_token_sizes=counts, + expert_token_prefix_sum=None, + actual_expert_shape=None, + block_idx=cute.arch.block_idx(), + smem_workspace=self.smem, + smem_base=storage.buffer.data_ptr(), + device_workspace=self.device, + ) + pipeline_init_arrive(cluster_shape_mn=(1, 1), is_relaxed=True) + pipeline_init_wait(cluster_shape_mn=(1, 1)) + worker = self.scheduler._work_id_worker + state = self.scheduler._task_mapping_state + _, fc1_valid, state = map_phase_interleaved_fc12_work_id( + Int32(0), Int32(BlockPhase.Linear1), (Int32(0), Int32(0), Int32(0)), state + ) + _, fc2_valid, state = map_phase_interleaved_fc12_work_id( + Int32(0), Int32(BlockPhase.Linear2), (Int32(0), Int32(0), Int32(0)), state + ) + target, short_stream = resolve_phase_interleaved_fc1_claim_target( + Int32(self.scheduler.minimum_global_fc1_claims), state + ) + input_exhausted = Boolean(exhausted[0]) + pointer = worker.get_atomic_counter_pointer(Int32(0)) + result = self.scheduler._wait_for_fc1_prologue_claims(input_exhausted, worker, state) + tid, _, _ = cute.arch.thread_idx() + output[tid, 0] = Int32(result) + output[tid, 1] = cute.arch.load(pointer, Int32, sem="acquire", scope="gpu") + output[tid, 2] = target + output[tid, 3] = Int32(short_stream) + output[tid, 4] = Int32(fc1_valid) + output[tid, 5] = Int32(fc2_valid) + + probe = PrologueProbe() + counts = torch.tensor([1], device="cuda", dtype=torch.int32) + workspace = torch.zeros(probe.device.total_bytes("local"), device="cuda", dtype=torch.uint8) + exhausted = torch.zeros(1, device="cuda", dtype=torch.int32) + output = torch.full((32, 6), -1, device="cuda", dtype=torch.int32) + tensors = tuple( + from_dlpack(t, assumed_align=16) for t in (counts, workspace, exhausted, output) + ) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + compiled = cute.compile(probe, *tensors, stream) + counter_index = probe.device.offset(probe.scheduler._work_id_worker.work_id_counter_region) // 4 + + def run(counter: int, input_exhausted: bool) -> torch.Tensor: + workspace.zero_() + workspace.view(torch.int32)[counter_index] = counter + workspace.view(torch.int32)[counter_index + 1] = 1 + exhausted.fill_(int(input_exhausted)) + output.fill_(-1) + compiled(*tensors, stream) + torch.cuda.synchronize() + return output.cpu() + + return run + + +@pytest.mark.parametrize( + "counter,input_exhausted", + [(1, False), (3, False), (2, True)], + ids=["slow-peek", "fast-peek", "already-exhausted"], +) +def test_fc1_prologue_wait_preserves_exhaustion( + _run_megamoe_prologue_helper: Callable[[int, bool], torch.Tensor], + counter: int, + input_exhausted: bool, +) -> None: + # One observer warp represents a three-cluster snapshot with one real FC1 + # tile. Counter 1 precedes peer arrivals; counter 3 includes two invalid peer + # claims. Already-exhausted controls include an invalid local claim. + actual = _run_megamoe_prologue_helper(counter, input_exhausted) + expected = torch.tensor( + [int(input_exhausted), counter, 1, 1, 1, 1], dtype=torch.int32 + ).expand_as(actual) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + def run_backend_moe( backend: MoE, backend_type: MoeBackendType, @@ -3897,8 +4068,8 @@ def test_nvfp4_fc1_row_alignment_gate( @pytest.mark.parametrize( "backend_cls", - [CutlassFusedMoE, CuteDslFusedMoE], - ids=["cutlass", "cutedsl"], + [CutlassFusedMoE, CuteDslFusedMoE, CuteDslFc12FusedMoE], + ids=["cutlass", "cutedsl", "cutedsl_fc12"], ) def test_situ_survives_resolution_not_just_construction(backend_cls): """A SiTU layer must be admitted by the *resolver*, not only build. @@ -3919,6 +4090,61 @@ def test_situ_survives_resolution_not_just_construction(backend_cls): assert rejection is None, f"{backend_cls.__name__} refuses SiTU at resolution: {rejection}" +@pytest.mark.parametrize("sm", [100, 103, 107]) +def test_cutedsl_nvfp4_situ_admitted_on_blackwell_and_rubin(sm): + """CuteDSL NVFP4 SiTU is eligible on SM100/SM103 and on Rubin (SM107).""" + deps = (MoEDep.CUTEDSL_RUBIN,) if sm == 107 else () + deployment = MoEDeployment( + ep_size=1, + tp_size=1, + parallel_size=1, + use_dp=False, + num_slots=256, + env=MoEEnvironment(sm=sm, available_deps=deps), + ) + verdict = CuteDslFusedMoE.can_implement(_nvfp4_problem(2048, "SiTu"), deployment) + assert verdict.reject_reason is not MoERejectReason.ACTIVATION_UNSUPPORTED, verdict.detail + + +def test_cutedsl_fc12_refuses_relu2_at_resolution(): + """Relu2 is not gated; the fused FC12 epilogue would compile it as SwiGLU.""" + rejection = _reject_unsupported_activation(CuteDslFc12FusedMoE, _nvfp4_problem(2048, "Relu2")) + assert rejection is not None + assert rejection.reason is MoERejectReason.ACTIVATION_UNSUPPORTED + + +def test_cutedsl_fc12_outer_runner_unique_id_includes_epilogue(): + """The outer AutoTuner key must not share a tile tactic across epilogues.""" + + def _forward(*args, **kwargs): + raise AssertionError("not launched") + + def _runner(identity): + return CuteDslFc12FusedMoENvfp4Runner( + forward_impl=_forward, + num_experts=8, + top_k=2, + num_local_experts=8, + local_expert_offset=0, + enable_finalize_fusion=True, + enable_alltoall=False, + workload_identity=identity, + ) + + swiglu = _runner((int(ActivationType.Swiglu), -1.0, -1.0)) + situ = _runner((int(ActivationType.SiTu), 4.0, 25.0)) + assert swiglu.unique_id() != situ.unique_id() + + # A non-None workload_identity is a cache key only. It must not open the + # locality-domain inner-tactic priming branch, which calls forward_impl with + # overlap_moe_output_memset -- a kwarg the fused single-op impl does not + # take. _forward raises if reached; no GPU needed, the failure is at kwarg + # binding. + assert situ.prime_inner_tactics is False + inputs = [torch.empty(0) for _ in range(5)] + assert situ(inputs, tactic=-1, do_preparation=True) is inputs[4] + + def test_unresolvable_layer_error_carries_rejection_details(): """describe() prints reason codes only, so impl_class_for has to add the details -- without them a shape rejection reaches the operator as a bare diff --git a/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py b/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py index 9ec0e5c10073..bec466f6423c 100644 --- a/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py +++ b/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py @@ -487,6 +487,59 @@ def mnnvl_checkpoint_rejects_wrong_membership(world_size: int, os.environ.pop(name, None) +@torch.inference_mode() +def mnnvl_rms_norm_fusion_forward( + input: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + dtype: torch.dtype, + tensor_parallel_size: int, + tensor_parallel_rank: int, + reference_norm: torch.Tensor, +): + input = input.cuda() + norm_weight = norm_weight.cuda() + reference_norm = reference_norm.cuda() + + os.environ["TLLM_TEST_MNNVL"] = "1" + MPI.COMM_WORLD.barrier() + + allreduce = AllReduce( + mapping=Mapping( + world_size=tensor_parallel_size, + tp_size=tensor_parallel_size, + rank=tensor_parallel_rank, + ), + strategy=AllReduceStrategy.MNNVL, + dtype=dtype, + ) + assert allreduce.mnnvl_allreduce is not None + params = AllReduceParams( + fusion_op=AllReduceFusionOp.RMS_NORM, + norm_weight=norm_weight, + eps=eps, + ) + # Call the MNNVL module directly: AllReduce.forward would silently fall + # back to another strategy if MNNVL declined the fusion. + output = allreduce.mnnvl_allreduce(input, all_reduce_params=params) + assert isinstance(output, torch.Tensor), type(output) + torch.testing.assert_close(output, reference_norm, rtol=0.05, atol=0.15) + + +def run_rms_norm_single_rank(tensor_parallel_size, input, norm_weight, eps, + dtype, reference_norm): + rank = tensorrt_llm.mpi_rank() + torch.cuda.set_device(rank) + try: + mnnvl_rms_norm_fusion_forward(input, norm_weight, eps, dtype, + tensor_parallel_size, rank, + reference_norm) + except Exception: + traceback.print_exc() + raise + return True + + @torch.inference_mode() def mnnvl_quant_fusion_forward( input: torch.Tensor, @@ -887,6 +940,41 @@ def test_mnnvl_quant_fusion(seq_len, hidden_size, dtype, fusion_op, assert r is True +@pytest.mark.skipif(torch.cuda.device_count() < 2, + reason="needs 2 GPUs to run this test") +@pytest.mark.parametrize("seq_len, hidden_size", ( + pytest.param(16, 3584, id="oneshot_m16_n3584"), + pytest.param(2048, 3584, id="twoshot_m2048_n3584"), +)) +@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) +def test_mnnvl_rms_norm_fusion_without_residual(seq_len, hidden_size, + mpi_pool_executor): + """AllReduce + RMSNorm with no residual operand (Kimi K3 routed latent).""" + torch.manual_seed(42) + tensor_parallel_size = mpi_pool_executor.num_workers + dtype = torch.bfloat16 + eps = 1e-5 + + x = torch.randn((tensor_parallel_size, seq_len, hidden_size), dtype=dtype) + norm_weight = torch.randn((hidden_size, ), dtype=dtype) + reference_norm = rms_norm( + torch.sum(x, dim=0).to(torch.float32), norm_weight, eps).to(dtype) + + results = mpi_pool_executor.map( + run_rms_norm_single_rank, + *zip(*[( + tensor_parallel_size, + x[i, :, :], + norm_weight, + eps, + dtype, + reference_norm, + ) for i in range(tensor_parallel_size)]), + ) + for r in results: + assert r is True + + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs 2 GPUs to run this test") @pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py index ff4d75f354dc..a8f8eddae83c 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py @@ -1727,8 +1727,14 @@ def test_mla_dspark_auto_does_not_degrade_to_the_eager_reference(monkeypatch): ``_auto_fallback_attention_backend = None``. """ from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models import modeling_dspark from tensorrt_llm._torch.models.modeling_dspark import MLADSparkForCausalLM + # AUTO's pick is CUTEDSL; make it unavailable, and TRTLLM too, so a + # degrade to either kernel backend would also be caught. + monkeypatch.setattr( + modeling_dspark, "cute_dsl_mla_decode_unavailability_reason", lambda: "probe says no" + ) monkeypatch.setattr(MLADSparkForCausalLM, "_mla_decode_op", staticmethod(lambda: None)) model_config = ModelConfig(pretrained_config=_tiny_mla_config(), attn_backend="VANILLA") with pytest.raises(ValueError, match="does not degrade"): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py index 930c1e78a5a4..76d91360159b 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py @@ -462,7 +462,8 @@ class _FakeSpecState: cache.kda_conv_q = torch.full((L, SLOTS, D, committed + M), 7.0) cache.kda_conv_k = torch.full((L, SLOTS, D, committed + M), 7.0) cache.kda_conv_v = torch.full((L, SLOTS, D, committed + M), 7.0) - cache.kda_qkg_cache = torch.full((L, SLOTS, M, 3, D), 7.0) + cache.kda_k_cache = torch.full((L, SLOTS, M, D), 7.0) + cache.kda_g_cache = torch.full((L, SLOTS, M, D), 7.0) cache.kda_v_cache = torch.full((L, SLOTS, M, D), 7.0) cache.kda_beta_cache = torch.full((L, SLOTS, M, NH), 7.0) cache.prev_num_accepted_tokens = torch.full((SLOTS,), 5, dtype=torch.int32) @@ -490,7 +491,8 @@ def test_seeds_committed_window_and_clears_scratch(self): kda[:, slot, :, :committed], conv_before[:, slot, lo:hi].to(kda.dtype) ) assert (kda[:, slot, :, committed:] == 0).all() - assert (cache.kda_qkg_cache[:, slot] == 0).all() + assert (cache.kda_k_cache[:, slot] == 0).all() + assert (cache.kda_g_cache[:, slot] == 0).all() assert (cache.kda_v_cache[:, slot] == 0).all() assert (cache.kda_beta_cache[:, slot] == 0).all() assert cache.prev_num_accepted_tokens[slot] == 0 @@ -499,7 +501,8 @@ def test_seeds_committed_window_and_clears_scratch(self): # Untouched slots keep their contents. for slot in (0, 2): assert (cache.kda_conv_q[:, slot] == 7.0).all() - assert (cache.kda_qkg_cache[:, slot] == 7.0).all() + assert (cache.kda_k_cache[:, slot] == 7.0).all() + assert (cache.kda_g_cache[:, slot] == 7.0).all() assert cache.prev_num_accepted_tokens[slot] == 5 def test_noop_without_kda_replay_or_unknown_ids(self): diff --git a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py index 706d55ae10c7..7bb17f69a1f7 100644 --- a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py +++ b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py @@ -2747,6 +2747,7 @@ def forward_impl(*args, **kwargs): num_local_experts=1, local_expert_offset=0, workload_identity=("locality_domain",), + prime_inner_tactics=True, ) inputs = [torch.empty(0) for _ in range(num_inputs)] @@ -2760,6 +2761,229 @@ def forward_impl(*args, **kwargs): assert kwargs["overlap_moe_output_memset"] is False +@pytest.mark.parametrize( + "module_name,runner_name,num_inputs,output_idx", + [ + ("fused_moe_cute_dsl", "CuteDslFusedMoENvfp4Runner", 5, 4), + ("fused_moe_cute_dsl", "CuteDslFusedMoEBF16Runner", 4, 3), + # The subclass production instantiates for CUTEDSL_FC12. It overrides + # _tile_sizes, so a green result on the parent is not evidence about it. + ("fused_moe_cute_dsl_fc12", "CuteDslFc12FusedMoENvfp4Runner", 5, 4), + ], +) +def test_workload_identity_alone_does_not_prime_inner_tactics( + monkeypatch, + module_name: str, + runner_name: str, + num_inputs: int, + output_idx: int, +): + """workload_identity is a cache key, not a "forward_impl nests tuning" flag. + + Non-locality-domain backends (e.g. CuteDslFc12FusedMoE, which keys the outer + tactic on the activation epilogue) pass a workload_identity, but their + ``run_moe_*_impl`` has no ``overlap_moe_output_memset`` parameter; priming + them raises TypeError during autotuner warmup. + """ + import importlib + + from tensorrt_llm._torch.moe.fused_moe import fused_moe_cute_dsl + + monkeypatch.setattr(fused_moe_cute_dsl, "get_sm_version", lambda: 107) + preparation_calls = [] + + def forward_impl(*args, enable_alltoall=False, tile_size=128): + # Mirrors the single-op signature: no ``overlap_moe_output_memset``, + # so any priming call is a TypeError. + preparation_calls.append((args, tile_size)) + return args[output_idx] + + module = importlib.import_module(f"tensorrt_llm._torch.moe.fused_moe.{module_name}") + runner_cls = getattr(module, runner_name) + runner = runner_cls( + forward_impl, + num_experts=1, + top_k=1, + num_local_experts=1, + local_expert_offset=0, + workload_identity=(int(ActivationType.SiTu), 1.0, 0.5), + ) + assert runner.prime_inner_tactics is False + inputs = [torch.empty(0) for _ in range(num_inputs)] + + result = runner(inputs, tactic=-1, do_preparation=True) + + assert result is inputs[output_idx] + assert preparation_calls == [] + + +@pytest.mark.parametrize("reason", ["disabled", "not_nvfp4", "no_aux_stream"]) +def test_quantize_input_async_falls_back_to_sync(monkeypatch, reason: str): + """Without an aux stream, NVFP4, or with the kill switch off, the async + quantization is the synchronous one and hands back no event to join.""" + from tensorrt_llm._torch.moe.fused_moe import fused_moe_cute_dsl + + monkeypatch.setattr(fused_moe_cute_dsl, "ASYNC_INPUT_QUANT", reason != "disabled") + calls = [] + backend = SimpleNamespace( + has_nvfp4=reason != "not_nvfp4", + quantize_input=lambda x, post_quant_comm: calls.append(post_quant_comm) or (x, None), + _has_moe_output_memset_aux_stream=lambda: reason != "no_aux_stream", + ) + x = torch.empty(2, 4) + + x_q, x_sf, event = fused_moe_cute_dsl.CuteDslFusedMoE.quantize_input_async(backend, x) + + assert x_q is x and x_sf is None and event is None + assert calls == [False] + + +def test_join_pending_x_ready_waits_once(): + from tensorrt_llm._torch.moe.fused_moe import fused_moe_cute_dsl + + waits = [] + backend = SimpleNamespace(_pending_x_ready_event=SimpleNamespace(wait=lambda: waits.append(1))) + join = fused_moe_cute_dsl.CuteDslFusedMoE._join_pending_x_ready + + join(backend) + join(backend) + + assert waits == [1] + assert backend._pending_x_ready_event is None + + +@pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() < 100, + reason="NVFP4 quantization needs SM100+", +) +def test_quantize_input_async_matches_sync(): + """The aux-stream quantization produces the same bytes as the main-stream one.""" + from tensorrt_llm._torch.moe.fused_moe import fused_moe_cute_dsl + from tensorrt_llm._torch.utils import AuxStreamType, EventType + + hidden = 256 + backend = SimpleNamespace( + has_nvfp4=True, + hidden_size=hidden, + scaling_vector_size=16, + fc31_input_scale=torch.tensor([1.0], device="cuda"), + aux_stream_dict={AuxStreamType.MoeOutputMemset: torch.cuda.Stream()}, + event_dict={ + EventType.Main: torch.cuda.Event(), + EventType.MoeOutputMemset: torch.cuda.Event(), + }, + _x_quant_event=torch.cuda.Event(), + ) + cls = fused_moe_cute_dsl.CuteDslFusedMoE + backend.quantize_input = cls.quantize_input.__get__(backend) + backend._has_moe_output_memset_aux_stream = cls._has_moe_output_memset_aux_stream.__get__( + backend + ) + x = torch.randn(8, hidden, dtype=torch.bfloat16, device="cuda") + + ref_q, ref_sf = backend.quantize_input(x, post_quant_comm=False) + x_q, x_sf, event = cls.quantize_input_async(backend, x) + assert event is backend._x_quant_event + event.wait() + + assert torch.equal(x_q.view(torch.uint8), ref_q.view(torch.uint8)) + assert torch.equal(x_sf.view(torch.uint8), ref_sf.view(torch.uint8)) + + +def test_fc12_warmup_preparation_does_not_call_the_fused_impl(monkeypatch): + """Drive the real FC12 runner through do_preparation, on CPU. + + The runner is built through ``CuteDslFc12FusedMoE.run_moe_nvfp4`` so the + workload_identity and the forward_impl are production's. Priming the + fused impl would raise ``TypeError: ... unexpected keyword argument + 'overlap_moe_output_memset'`` at kwarg binding, before any CUDA work. + """ + import inspect + + from tensorrt_llm._torch.moe.fused_moe import fused_moe_cute_dsl_fc12 as fc12 + + captured = {} + + class _SpyRunner(fc12.CuteDslFc12FusedMoENvfp4Runner): + def __init__(self, **kwargs): + captured.update(kwargs) + captured["runner"] = self + super().__init__(**kwargs) + + def forward(self, inputs, tactic, do_preparation=False): + if do_preparation: + return super().forward(inputs, tactic=tactic, do_preparation=True) + # The steady-state launch needs a GPU and is not on trial. + return inputs[4] + + monkeypatch.setattr(fc12, "CuteDslFc12FusedMoENvfp4Runner", _SpyRunner) + monkeypatch.setattr( + fc12.AutoTuner, + "get", + staticmethod( + lambda: SimpleNamespace(choose_one=lambda op, runners, cfg, inputs: (runners[0], 128)) + ), + ) + + num_tokens, hidden = 4, 16 + impl_calls = [] + backend = SimpleNamespace( + has_nvfp4=True, + num_slots=8, + hidden_size=hidden, + use_fused_finalize=True, + activation_type=ActivationType.SiTu, + act_alpha=4.0, + act_beta=25.0, + ) + for name in ("_situ_betas", "_epilogue_identity"): + setattr(backend, name, getattr(fc12.CuteDslFc12FusedMoE, name).__get__(backend)) + # Bound, like production's ``self.run_moe_nvfp4_impl``. + real_impl = fc12.CuteDslFc12FusedMoE.run_moe_nvfp4_impl.__get__(backend) + + def tracking_impl(*args, **kwargs): + impl_calls.append(kwargs) + return real_impl(*args, **kwargs) + + backend.run_moe_nvfp4_impl = tracking_impl + weight_view = SimpleNamespace(expert_size_per_partition=8, slot_start=0) + + inputs = dict( + x=torch.empty(num_tokens, hidden, dtype=torch.bfloat16), + token_selected_experts=torch.zeros(num_tokens, 2, dtype=torch.int32), + token_final_scales=torch.ones(num_tokens, 2), + x_sf=torch.empty(0), + moe_output=torch.empty(num_tokens, hidden, dtype=torch.bfloat16), + weight_view=weight_view, + ) + out = fc12.CuteDslFc12FusedMoE.run_moe_nvfp4(backend, **inputs) + assert out.shape == (num_tokens, hidden) + + # Positive control: the activation-keyed identity is non-None. + assert captured["workload_identity"] == (int(ActivationType.SiTu), 4.0, 25.0) + + runner = captured["runner"] + assert runner.prime_inner_tactics is False + ordered = [ + inputs[k] + for k in ( + "x", + "token_selected_experts", + "token_final_scales", + "x_sf", + "moe_output", + "weight_view", + ) + ] + assert runner(ordered, tactic=-1, do_preparation=True) is ordered[4] + assert impl_calls == [], f"warmup priming reached the fused impl with {impl_calls}" + # And the reason it must not: the impl cannot take that kwarg. + assert ( + "overlap_moe_output_memset" + not in inspect.signature(fc12.CuteDslFc12FusedMoE.run_moe_nvfp4_impl).parameters + ) + + def _cute_dsl_eligibility(**deployment_kwargs): """``can_implement`` verdict for NVFP4 on an SM107 machine with Rubin DSL.""" from tensorrt_llm._torch.moe.fused_moe.impl_contract import ( diff --git a/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py b/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py index 1e03da59024a..fe286df2d7e6 100644 --- a/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py +++ b/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py @@ -24,7 +24,9 @@ def router_gemm_ref(input, weight, bias, dtype): @pytest.mark.parametrize( "num_tokens", [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]) -@pytest.mark.parametrize("num_experts", [256]) +# 256 experts: DeepSeek-V3/V4, GLM-5 (scalar FFMA kernel, all three hidden sizes). +# 896 experts: Kimi K3 (tensor-core kernel for hidden 7168; cuBLAS fallback otherwise). +@pytest.mark.parametrize("num_experts", [256, 896]) @pytest.mark.parametrize("hidden_size", [7168, 6144, 4096]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) def test_router_gemm_run(num_tokens, num_experts, hidden_size, dtype): @@ -39,4 +41,75 @@ def test_router_gemm_run(num_tokens, num_experts, hidden_size, dtype): torch.float32) logtis_ref = router_gemm_ref(input.float(), weight.t().float(), bias, torch.float32) - assert torch.allclose(logits, logtis_ref, rtol=5e-2) + # Logits have std ~ sqrt(hidden_size) ~ 85; fp32 accumulation-order noise (~1e-3) on + # near-zero logits fails a pure relative tolerance, hence the absolute term. + assert torch.allclose(logits, logtis_ref, rtol=5e-2, atol=1e-2) + # The router output only matters through the top-k selection that follows it. + top_k = 16 if num_experts == 896 else 8 + assert torch.equal( + torch.topk(logits, top_k, dim=-1).indices.sort(dim=-1).values, + torch.topk(logtis_ref, top_k, dim=-1).indices.sort(dim=-1).values) + + +@pytest.mark.skipif(not torch.cuda.is_available() + or torch.cuda.get_device_capability() != (10, 7), + reason="requires SM107") +@pytest.mark.parametrize("num_tokens", [1, 8, 16]) +def test_router_latent_gemm_matches_separate_ops(num_tokens: int) -> None: + from types import SimpleNamespace + + from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoEGate + + torch.manual_seed(24) + torch.cuda.manual_seed(24) + config = SimpleNamespace(hidden_size=7168, + num_experts=896, + num_experts_per_token=16, + routed_scaling_factor=2.5, + moe_router_activation_func="sigmoid", + moe_renormalize=True) + gate = KimiK3MoEGate(config, + logits_gemm_dtype=torch.bfloat16, + device=torch.device("cuda")) + down = torch.nn.Linear(config.hidden_size, + 3584, + bias=False, + device="cuda", + dtype=torch.bfloat16) + hidden = torch.randn(num_tokens, + config.hidden_size, + dtype=torch.bfloat16, + device="cuda") + with torch.no_grad(): + gate.weight.normal_() + down.weight.normal_() + expected_logits = gate.compute_logits(hidden) + expected_projection = down(hidden) + precise_projection = (hidden.double() @ down.weight.double().t()).to( + torch.bfloat16) + logits, projection = torch.ops.trtllm.dsv3_router_latent_gemm_op( + hidden, gate.weight, down.weight) + + torch.testing.assert_close(logits, expected_logits, rtol=0, atol=0) + assert torch.isfinite(logits).all() + assert projection.shape == expected_projection.shape + assert projection.dtype == expected_projection.dtype == torch.bfloat16 + # Unit-normal K=7168 dot products need an FP32 accumulation allowance near + # zero; one BF16 step alone can reject even the correctly rounded result. + accumulation_atol = 1e-3 + for actual, reference in ((projection, expected_projection), + (projection, precise_projection), + (expected_projection, precise_projection)): + lower = torch.nextafter(reference, + torch.full_like(reference, -torch.inf)).float() + upper = torch.nextafter(reference, + torch.full_like(reference, torch.inf)).float() + within_tolerance = (torch.isfinite(actual) + & torch.isfinite(reference) + & (actual.float() >= lower - accumulation_atol) + & (actual.float() <= upper + accumulation_atol)) + mismatches = (~within_tolerance).sum().item() + max_abs = (actual.float() - reference.float()).abs().max().item() + assert mismatches == 0, ( + f"{mismatches}/{actual.numel()} projection values exceed one BF16 " + f"step + {accumulation_atol} or are nonfinite; max_abs={max_abs}") diff --git a/tests/unittest/_torch/thop/parallel/test_kimi_k3_fp8_weight_read_linear.py b/tests/unittest/_torch/thop/parallel/test_kimi_k3_fp8_weight_read_linear.py index d36fe773f004..8e42a5df2adc 100644 --- a/tests/unittest/_torch/thop/parallel/test_kimi_k3_fp8_weight_read_linear.py +++ b/tests/unittest/_torch/thop/parallel/test_kimi_k3_fp8_weight_read_linear.py @@ -140,6 +140,28 @@ def test_single_cute_scale_is_built_on_rubin(): assert not hasattr(mod, "gemm_alpha") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@RUBIN +def test_checkpoint_pair_carries_cute_scale(): + """A fused checkpoint FP8 pair must be prepared into the CuTe scale layout.""" + in_features, parts = 1024, [512, 512] + torch.random.manual_seed(6) + weights, scales = [], [] + for p in parts: + wp = torch.randn((p, in_features), device="cuda", dtype=torch.bfloat16) / in_features**0.5 + q, s = per_block_cast_to_fp8(wp) + weights.append(q) + scales.append(s.float()) + + weight, weight_scale = K3Fp8Linear._prepare_weights( + torch.cat(weights, dim=0), torch.cat(scales, dim=0) + ) + mod = K3Fp8Linear(weight, weight_scale, sum(parts)) + assert mod.weight_scale.numel() > 0 + assert mod.weight_scale.dtype is torch.uint8 + assert not hasattr(mod, "gemm_alpha") + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") @RUBIN def test_weight_preparation_returns_only_cute_pair():