diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 85ee5ccf8970..f3a8adfe64fd 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -800,10 +800,13 @@ size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int3 auto const batch_size = static_cast(max_num_seq); auto const kv_seq_length = (isCrossAttention() ? cross_kv_length : input_seq_length); - // The unfused-MHA buffers below must upper-bound the enqueueContext carve, which sizes them by - // batch_size * input_seq_length (not num_tokens): with padding removal the actual token count can be - // smaller than batch_size * max(context q length), so sizing by max_num_tokens underestimates. - size_t const attention_mask_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_length * kv_seq_length; + // Unfused context attention operates on padded [batch, sequence] tensors, + // even when the input QKV is packed. Size those buffers from the padded + // token counts exactly as enqueueContext does; max_num_tokens remains the + // packed count used by the fused paths below. + size_t const padded_num_tokens = batch_size * static_cast(input_seq_length); + size_t const padded_kv_tokens = batch_size * static_cast(kv_seq_length); + size_t const attention_mask_size = mEnableContextFMHA ? 0 : size * padded_num_tokens * kv_seq_length; size_t const cu_seqlens_size = sizeof(int) * (batch_size + 1); size_t const rotary_inv_freq_size = sizeof(float) * batch_size * mRotaryEmbeddingDim / 2; @@ -823,7 +826,7 @@ size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int3 size_t const v_buf_2_size = mEnableContextFMHA ? 0 : size * batch_size * kv_seq_length * local_hidden_units_kv; size_t const qk_buf_size = mEnableContextFMHA ? 0 : size * batch_size * mNumHeads * input_seq_length * kv_seq_length; - size_t const qkv_buf_2_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_length * local_hidden_units_qo; + size_t const qkv_buf_2_size = mEnableContextFMHA ? 0 : size * padded_num_tokens * local_hidden_units_qo; size_t const qk_buf_float_size = mEnableContextFMHA ? 0 : sizeof(float) * batch_size * mNumHeads * input_seq_length * kv_seq_length; int dim_q_per_head = (mMLAParams.qk_rope_head_dim + mMLAParams.qk_nope_head_dim); @@ -899,8 +902,8 @@ size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int3 ? sizeof(float) * tc::divUp(local_hidden_units_kv, std::max(1, mSageAttnNumEltsPerBlkV)) : 0; - size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * batch_size * input_seq_length; - size_t const encoder_padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * batch_size * cross_kv_length; + size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * padded_num_tokens; + size_t const encoder_padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * padded_kv_tokens; // Each token holds (batch_idx, token_idx_in_seq) int2. size_t const tokens_info_size = sizeof(int2) * max_num_tokens; size_t const fmha_scheduler_counter = mEnableContextFMHA ? sizeof(uint32_t) : 0; @@ -1886,6 +1889,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea preprocessingParams.is_last_chunk = !mAttentionChunkSize.has_value() || (params.input_seq_length == params.max_past_kv_length); + if (!(mIsMLAEnabled && params.mla_param != nullptr && params.mla_param->q_rope_applied)) { std::string const beforeRopeStr = "ctx attention before RoPE at layer " + std::to_string(mLayerIdx); TLLM_CHECK_DEBUG_WITH_INFO(tensorrt_llm::runtime::utils::tensorHasInvalid(params.num_tokens, @@ -1991,6 +1995,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea invokeQKVPreprocessing(preprocessingParams, stream); } sync_check_cuda_error(stream); + if (!(mIsMLAEnabled && params.mla_param != nullptr && params.mla_param->q_rope_applied)) { std::string const afterRopeStr = "ctx attention after RoPE at layer " + std::to_string(mLayerIdx); TLLM_CHECK_DEBUG_WITH_INFO(tensorrt_llm::runtime::utils::tensorHasInvalid(params.num_tokens, diff --git a/cpp/tensorrt_llm/common/envUtils.h b/cpp/tensorrt_llm/common/envUtils.h index e6d354d05745..a81e1b362f82 100644 --- a/cpp/tensorrt_llm/common/envUtils.h +++ b/cpp/tensorrt_llm/common/envUtils.h @@ -72,10 +72,10 @@ bool getEnvUseFineGrainedSync(); void setFineGrainedSyncDisabledOverride(bool disabled); template -inline void launchWithPdlWhenEnabled(char const* name, KernelFn kernelFn, dim3 grid, dim3 block, size_t dynamicShmSize, - cudaStream_t stream, Args&&... args) +inline void launchWithPdl(char const* name, bool enablePdl, KernelFn kernelFn, dim3 grid, dim3 block, + size_t dynamicShmSize, cudaStream_t stream, Args&&... args) { - TLLM_LOG_DEBUG("Enable PDL in %s", name); + TLLM_LOG_DEBUG("PDL in %s: %s", name, enablePdl ? "enabled" : "disabled"); cudaLaunchConfig_t kernelConfig; kernelConfig.gridDim = grid; kernelConfig.blockDim = block; @@ -84,13 +84,20 @@ inline void launchWithPdlWhenEnabled(char const* name, KernelFn kernelFn, dim3 g cudaLaunchAttribute attrs[1]; attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + attrs[0].val.programmaticStreamSerializationAllowed = enablePdl; kernelConfig.attrs = attrs; kernelConfig.numAttrs = 1; TLLM_CUDA_CHECK(cudaLaunchKernelEx(&kernelConfig, kernelFn, std::forward(args)...)); } +template +inline void launchWithPdlWhenEnabled(char const* name, KernelFn kernelFn, dim3 grid, dim3 block, size_t dynamicShmSize, + cudaStream_t stream, Args&&... args) +{ + launchWithPdl(name, getEnvEnablePDL(), kernelFn, grid, block, dynamicShmSize, stream, std::forward(args)...); +} + bool getEnvUseUCXKvCache(); bool getEnvUseMPIKvCache(); diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh index 7817012f2a97..b997621a8fcc 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh +++ b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh @@ -46,6 +46,7 @@ // empty_input[N_INPUT] pmap -> TMA (input slot empty) // full_cast[2] pmap -> MMA (A ready in TMEM) // empty_cast[2] MMA -> pmap (TMEM slot empty) +// full_mix[1] TMA -> pmap (post_mix+comb_mix arrived) // tmem_full[1] MMA -> epilogue #pragma once @@ -184,6 +185,36 @@ __device__ __forceinline__ void stsm_x4_b16_rout(void* smem_dst, uint32_t a, uin "r"(a), "r"(b), "r"(c), "r"(d)); } +// Blackwell can issue two independent FP32 FMAs with one packed instruction. +// The pmap hot loop always updates adjacent bf16 values with the same mixing +// coefficient, so keep the pair packed instead of lowering it to two FFMA. +__device__ __forceinline__ float2 fma_f32x2(float2 const& a, float2 const& b, float2 const& c) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (__CUDA_ARCH__ < 1100) + float2 result; + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(reinterpret_cast(result)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b)), + "l"(reinterpret_cast(c))); + return result; +#else + return make_float2(fmaf(a.x, b.x, c.x), fmaf(a.y, b.y, c.y)); +#endif +} + +__device__ __forceinline__ float2 mul_f32x2(float2 const& a, float2 const& b) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (__CUDA_ARCH__ < 1100) + float2 result; + asm volatile("mul.f32x2 %0, %1, %2;" + : "=l"(reinterpret_cast(result)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b))); + return result; +#else + return make_float2(a.x * b.x, a.y * b.y); +#endif +} + template @@ -191,9 +222,10 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 const uint32_t shape_m, const __grid_constant__ cute::TmaDescriptor tensor_map_residual, const __grid_constant__ cute::TmaDescriptor tensor_map_x, const __grid_constant__ cute::TmaDescriptor tensor_map_b, const __grid_constant__ cute::TmaDescriptor tensor_map_residual_out, - float* __restrict__ D, // [M, SHAPE_N] (caller memsets to 0) - float const* __restrict__ post_mix, float const* __restrict__ comb_mix, float* __restrict__ sqr_sum) -{ // [M] (caller memsets to 0) + const __grid_constant__ cute::TmaDescriptor tensor_map_post, + const __grid_constant__ cute::TmaDescriptor tensor_map_comb, float const* __restrict__ post_mix, + float const* __restrict__ comb_mix, float* __restrict__ D, float* __restrict__ sqr_sum) +{ // D [M, SHAPE_N], sqr_sum [M] (caller memsets both to 0) #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) and (__CUDA_ARCH__ < 1100)) or defined(__CLION_IDE__) using Barrier = cutlass::arch::ClusterTransactionBarrier; @@ -243,6 +275,8 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 cute::prefetch_tma_descriptor(&tensor_map_x); cute::prefetch_tma_descriptor(&tensor_map_b); cute::prefetch_tma_descriptor(&tensor_map_residual_out); + cute::prefetch_tma_descriptor(&tensor_map_post); + cute::prefetch_tma_descriptor(&tensor_map_comb); } // SMEM layout: [cd, B stages, res stages, x stages, post, comb, rc (HC_MULT slices)] @@ -275,9 +309,10 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 [=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + i; }); auto empty_cast = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + kNumCastStages + i; }); - auto tmem_full_barrier = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages; + auto full_mix = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages; + auto tmem_full_barrier = full_mix + 1; - cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages + 1) * sizeof(Barrier); + cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages + 2) * sizeof(Barrier); auto tmem_ptr_in_smem = reinterpret_cast(cursor); if (warp_idx == 1 and cute::elect_one_sync()) @@ -300,6 +335,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 full_cast[i]->init(kNumPmapThreads); empty_cast[i]->init(1); } + full_mix->init(1); tmem_full_barrier->init(1); cutlass::arch::fence_barrier_init(); } @@ -313,8 +349,8 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 const uint32_t m_block_idx = block_idx / kNumSplits; const uint32_t k_split_idx = block_idx % kNumSplits; const uint32_t m_offset = m_block_idx * BLOCK_M; - // Give the first H_TILES_EXTRA splits one extra tile. Even splits fold back - // to the original constants. + // Give the first H_TILES_EXTRA splits one additional tile. This admits + // Rubin exact-wave KS=106/53 shapes for H=7168 (112 H tiles). uint32_t h_tile_start; uint32_t h_tiles_this_split; if constexpr (H_TILES_EXTRA == 0) @@ -327,16 +363,24 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 h_tile_start = k_split_idx * H_TILES_BASE + cute::min(k_split_idx, H_TILES_EXTRA); h_tiles_this_split = H_TILES_BASE + static_cast(k_split_idx < H_TILES_EXTRA); } - const uint32_t num_total_stages = h_tiles_this_split * HC_MULT; + uint32_t const num_total_stages = h_tiles_this_split * HC_MULT; - // Prologue removed: pmap threads load their own post_mix/comb_mix rows - // directly into registers below, so the MMA/TMA warps never wait on those - // global loads and the TMA pipeline starts immediately. + // Pmap threads load their own post_mix/comb_mix rows directly into registers, + // so the MMA/TMA warps can start filling the TMA pipeline immediately. if (warp_idx < kNumMMAThreads / 32) { // ----- TMA warp (warp 0) ----- if (warp_idx == 0 and cute::elect_one_sync()) { + // Load the two row-wise pmap coefficient tiles asynchronously. + // Only the pmap warp group consumes them, so it waits on full_mix + // independently while this warp continues filling input/B stages. + deep_gemm::tma::copy( + &tensor_map_post, full_mix, smem_post, /*inner_idx=*/0, m_offset); + deep_gemm::tma::copy( + &tensor_map_comb, full_mix, smem_comb, /*inner_idx=*/0, m_offset); + full_mix->arrive_and_expect_tx(SMEM_POST_SIZE + SMEM_COMB_SIZE); + uint32_t b_stage = 0; uint32_t i_stage = 0; uint32_t s = 0; @@ -471,6 +515,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 else { // ----- Pmap warp group (warps 4..7, 128 threads) ----- + full_mix->wait(0); const uint32_t sub_warp_idx = warp_idx - kNumMMAThreads / 32; const uint32_t upper_row = sub_warp_idx * 16 + lane_idx / 4; const uint32_t lower_row = upper_row + 8; @@ -595,17 +640,43 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 #pragma unroll for (uint32_t hc = 0; hc < HC_MULT; ++hc) { - float2 nu0 = mul_f32x2(pm_u[hc], xf[0][i + 0]); - float2 nu1 = mul_f32x2(pm_u[hc], xf[0][i + 1]); - float2 nl0 = mul_f32x2(pm_l[hc], xf[1][i + 0]); - float2 nl1 = mul_f32x2(pm_l[hc], xf[1][i + 1]); + float2 nu0, nu1, nl0, nl1; + if constexpr (kNumSplits >= 53) + { + // Packed coefficients are faster for Rubin exact-wave small-M splits. + float2 const coefficientU = make_float2(pm_u[hc], pm_u[hc]); + float2 const coefficientL = make_float2(pm_l[hc], pm_l[hc]); + nu0 = mul_f32x2(coefficientU, xf[0][i + 0]); + nu1 = mul_f32x2(coefficientU, xf[0][i + 1]); + nl0 = mul_f32x2(coefficientL, xf[1][i + 0]); + nl1 = mul_f32x2(coefficientL, xf[1][i + 1]); + } + else + { + nu0 = mul_f32x2(pm_u[hc], xf[0][i + 0]); + nu1 = mul_f32x2(pm_u[hc], xf[0][i + 1]); + nl0 = mul_f32x2(pm_l[hc], xf[1][i + 0]); + nl1 = mul_f32x2(pm_l[hc], xf[1][i + 1]); + } #pragma unroll for (uint32_t j = 0; j < HC_MULT; ++j) { - nu0 = fma_f32x2(cm_u[j][hc], r_u[j][0], nu0); - nu1 = fma_f32x2(cm_u[j][hc], r_u[j][1], nu1); - nl0 = fma_f32x2(cm_l[j][hc], r_l[j][0], nl0); - nl1 = fma_f32x2(cm_l[j][hc], r_l[j][1], nl1); + if constexpr (kNumSplits >= 53) + { + float2 const coefficientU = make_float2(cm_u[j][hc], cm_u[j][hc]); + float2 const coefficientL = make_float2(cm_l[j][hc], cm_l[j][hc]); + nu0 = fma_f32x2(coefficientU, r_u[j][0], nu0); + nu1 = fma_f32x2(coefficientU, r_u[j][1], nu1); + nl0 = fma_f32x2(coefficientL, r_l[j][0], nl0); + nl1 = fma_f32x2(coefficientL, r_l[j][1], nl1); + } + else + { + nu0 = fma_f32x2(cm_u[j][hc], r_u[j][0], nu0); + nu1 = fma_f32x2(cm_u[j][hc], r_u[j][1], nu1); + nl0 = fma_f32x2(cm_l[j][hc], r_l[j][0], nl0); + nl1 = fma_f32x2(cm_l[j][hc], r_l[j][1], nl1); + } } nv_bfloat162 b_u0 = __float22bfloat162_rn(nu0); nv_bfloat162 b_u1 = __float22bfloat162_rn(nu1); @@ -738,17 +809,17 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) const __grid_constant__ cute::TmaDescriptor tensor_map_x, // x_prev, bf16 const __grid_constant__ cute::TmaDescriptor tensor_map_b, // W_T, tf32 const __grid_constant__ cute::TmaDescriptor tensor_map_residual_out, // residual_cur, bf16 (TMA store) + const __grid_constant__ cute::TmaDescriptor tensor_map_post, // post_mix_prev, fp32 + const __grid_constant__ cute::TmaDescriptor tensor_map_comb, // comb_mix_prev, fp32 __nv_bfloat16 const* __restrict__ residual_cur_ptr, // same buffer as TMA target __nv_bfloat16* __restrict__ layer_input_out, // [M, HIDDEN] bf16 - float* __restrict__ D, // [M, SHAPE_N] fp32 (y_acc, caller zeros) - float* __restrict__ sqr_sum, // [M] fp32 (r_acc, caller zeros) - int* __restrict__ done_counter, // [ceil(M/BLOCK_M)] int (caller zeros) - float const* __restrict__ post_mix_prev, // [M, HC_MULT] - float const* __restrict__ comb_mix_prev, // [M, HC_MULT, HC_MULT] - float const* __restrict__ hc_scale, // [3] - float const* __restrict__ hc_base, // [HC_MULT*(2+HC_MULT)] - float* __restrict__ post_mix_out, // [M, HC_MULT] - float* __restrict__ comb_mix_out, // [M, HC_MULT, HC_MULT] + float* __restrict__ D, // [M, SHAPE_N] fp32 (y_acc, caller zeros) + float* __restrict__ sqr_sum, // [M] fp32 (r_acc, caller zeros) + int* __restrict__ done_counter, // [ceil(M/BLOCK_M)] int (caller zeros) + float const* __restrict__ hc_scale, // [3] + float const* __restrict__ hc_base, // [HC_MULT*(2+HC_MULT)] + float* __restrict__ post_mix_out, // [M, HC_MULT] + float* __restrict__ comb_mix_out, // [M, HC_MULT, HC_MULT] // When kFuseNorm: layer_input_out receives the RMSNorm-normalized // values: out[t,h] = bf16(li[t,h] * rsqrt(mean(li²)+norm_eps) * w[h]). // norm_weight must be bf16 [HIDDEN]; norm_eps is the RMSNorm epsilon. @@ -808,6 +879,8 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) cute::prefetch_tma_descriptor(&tensor_map_x); cute::prefetch_tma_descriptor(&tensor_map_b); cute::prefetch_tma_descriptor(&tensor_map_residual_out); + cute::prefetch_tma_descriptor(&tensor_map_post); + cute::prefetch_tma_descriptor(&tensor_map_comb); } // SMEM layout: [cd, B stages, res stages, x stages, post, comb, rc (HC_MULT slices)] @@ -840,9 +913,10 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) [=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + i; }); auto empty_cast = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + kNumCastStages + i; }); - auto tmem_full_barrier = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages; + auto full_mix = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages; + auto tmem_full_barrier = full_mix + 1; - cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages + 1) * sizeof(Barrier); + cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages + 2) * sizeof(Barrier); auto tmem_ptr_in_smem = reinterpret_cast(cursor); if (warp_idx == 1 and cute::elect_one_sync()) @@ -865,6 +939,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) full_cast[i]->init(kNumPmapThreads); empty_cast[i]->init(1); } + full_mix->init(1); tmem_full_barrier->init(1); cutlass::arch::fence_barrier_init(); } @@ -878,8 +953,9 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) const uint32_t m_block_idx = block_idx / kNumSplits; const uint32_t k_split_idx = block_idx % kNumSplits; const uint32_t m_offset = m_block_idx * BLOCK_M; - // Give the first H_TILES_EXTRA splits one extra tile. Even splits fold back - // to the original constants. + // Distribute remainder tiles over the first splits. Even-split instances + // fold to the original constants; KS=106/53 provide exact 212-CTA waves + // for H=7168 at M=128/256 on Rubin. uint32_t h_tile_start; uint32_t h_tiles_this_split; if constexpr (H_TILES_EXTRA == 0) @@ -894,44 +970,19 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) } const uint32_t num_total_stages = h_tiles_this_split * HC_MULT; - // Prologue: pmap warp group loads post_mix_prev, comb_mix_prev into SMEM - if (warp_idx >= kNumMMAThreads / 32) - { - const uint32_t pmap_tid = threadIdx.x - kNumMMAThreads; -#pragma unroll - for (uint32_t t = 0; t < 2; ++t) - { - uint32_t idx = pmap_tid + t * kNumPmapThreads; - if (idx < BLOCK_M * HC_MULT) - { - uint32_t m = idx / HC_MULT; - uint32_t hc = idx % HC_MULT; - uint32_t gmem_m = m_offset + m; - float v = (gmem_m < shape_m) ? post_mix_prev[gmem_m * HC_MULT + hc] : 0.f; - smem_post[idx] = v; - } - } -#pragma unroll - for (uint32_t t = 0; t < 8; ++t) - { - uint32_t idx = pmap_tid + t * kNumPmapThreads; - if (idx < BLOCK_M * HC_MULT * HC_MULT) - { - uint32_t m = idx / (HC_MULT * HC_MULT); - uint32_t jk = idx % (HC_MULT * HC_MULT); - uint32_t gmem_m = m_offset + m; - float v = (gmem_m < shape_m) ? comb_mix_prev[gmem_m * HC_MULT * HC_MULT + jk] : 0.f; - smem_comb[idx] = v; - } - } - } - __syncthreads(); - if (warp_idx < kNumMMAThreads / 32) { // ----- TMA warp (warp 0) ----- if (warp_idx == 0 and cute::elect_one_sync()) { + // Fetch coefficient tiles asynchronously while this warp starts + // filling the regular input/B pipeline. Only pmap waits on full_mix. + deep_gemm::tma::copy( + &tensor_map_post, full_mix, smem_post, /*inner_idx=*/0, m_offset); + deep_gemm::tma::copy( + &tensor_map_comb, full_mix, smem_comb, /*inner_idx=*/0, m_offset); + full_mix->arrive_and_expect_tx(SMEM_POST_SIZE + SMEM_COMB_SIZE); + uint32_t b_stage = 0; uint32_t i_stage = 0; uint32_t s = 0; @@ -1066,6 +1117,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) else { // ----- Pmap warp group (warps 4..7, 128 threads) ----- + full_mix->wait(0); const uint32_t sub_warp_idx = warp_idx - kNumMMAThreads / 32; const uint32_t upper_row = sub_warp_idx * 16 + lane_idx / 4; const uint32_t lower_row = upper_row + 8; @@ -1350,67 +1402,47 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) if (tok >= shape_m) continue; - // Lanes 0..HC_MULT-1 compute rmsnorm / sigmoid / sinkhorn; pre_mix is - // held in `pre_mix_local` on lanes 0..HC_MULT-1 and later broadcast to - // all 32 lanes via __shfl_sync. All warps in a team redundantly run - // these ~tens of FLOPs (cheap) to avoid a cross-warp SMEM sync; only - // warp_in_team==0 writes comb_mix_out / post_mix_out to GMEM. + // Every team warp needs pre_mix locally for the HIDDEN-stride loop. + // post_mix and iterative Sinkhorn are token-only results, so only the + // first warp computes them. KS=112/56/28 thereby removes 8x/4x/2x + // redundant Sinkhorn work without adding a cross-warp synchronization. + // Lanes 0..HC_MULT-1 hold the four rows and broadcast pre_mix below. float pre_mix_local = 0.f; if (lane_bf < HC_MULT) { float const r_val = sqr_sum[tok]; - float y_local[HC_MULT3]; float const* y_row = D + static_cast(tok) * SHAPE_N; -#pragma unroll - for (uint32_t c = 0; c < HC_MULT3; ++c) - y_local[c] = y_row[c]; - float const rstd = rsqrtf(r_val / static_cast(HC_MULT * HIDDEN) + rms_eps); float const s0 = hc_scale[0]; - float const s1 = hc_scale[1]; - float const s2 = hc_scale[2]; - float v = y_local[lane_bf] * rstd * s0 + hc_base[lane_bf]; + float v = y_row[lane_bf] * rstd * s0 + hc_base[lane_bf]; pre_mix_local = 1.0f / (1.0f + __expf(-v)) + hc_pre_eps; - v = y_local[HC_MULT + lane_bf] * rstd * s1 + hc_base[HC_MULT + lane_bf]; - float post_val = 1.0f / (1.0f + __expf(-v)) * hc_post_mult_value; if (warp_in_team == 0) { + float const s1 = hc_scale[1]; + float const s2 = hc_scale[2]; + v = y_row[HC_MULT + lane_bf] * rstd * s1 + hc_base[HC_MULT + lane_bf]; + float const post_val = 1.0f / (1.0f + __expf(-v)) * hc_post_mult_value; post_mix_out[tok * HC_MULT + lane_bf] = post_val; - } - float cm_vals[HC_MULT]; + float cm_vals[HC_MULT]; #pragma unroll - for (uint32_t k = 0; k < HC_MULT; ++k) - cm_vals[k] = y_local[2 * HC_MULT + lane_bf * HC_MULT + k] * rstd * s2 - + hc_base[2 * HC_MULT + lane_bf * HC_MULT + k]; + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_vals[k] = y_row[2 * HC_MULT + lane_bf * HC_MULT + k] * rstd * s2 + + hc_base[2 * HC_MULT + lane_bf * HC_MULT + k]; - constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; - float const rowMax = fmaxf(fmaxf(cm_vals[0], cm_vals[1]), fmaxf(cm_vals[2], cm_vals[3])); -#pragma unroll - for (uint32_t k = 0; k < HC_MULT; ++k) - cm_vals[k] = __expf(cm_vals[k] - rowMax); - // Reciprocal-multiply for sinkhorn: 1 fdiv + 4 fmul instead of 4 - // fdivs per row-normalize. Equivalent under fp32 round-off. - float inv_rs = 1.0f / (cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3]); -#pragma unroll - for (uint32_t k = 0; k < HC_MULT; ++k) - cm_vals[k] = cm_vals[k] * inv_rs + hc_sinkhorn_eps; + constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; + float const rowMax = fmaxf(fmaxf(cm_vals[0], cm_vals[1]), fmaxf(cm_vals[2], cm_vals[3])); #pragma unroll - for (uint32_t k = 0; k < HC_MULT; ++k) - { - float cs = cm_vals[k]; - cs += __shfl_xor_sync(LANE_MASK, cs, 1); - cs += __shfl_xor_sync(LANE_MASK, cs, 2); - cm_vals[k] *= 1.0f / (cs + hc_sinkhorn_eps); - } - for (uint32_t it = 1; it < sinkhorn_repeat; ++it) - { - inv_rs = 1.0f / (cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3] + hc_sinkhorn_eps); + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_vals[k] = __expf(cm_vals[k] - rowMax); + // Reciprocal-multiply for sinkhorn: 1 fdiv + 4 fmul instead of 4 + // fdivs per row-normalize. Equivalent under fp32 round-off. + float inv_rs = 1.0f / (cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3]); #pragma unroll for (uint32_t k = 0; k < HC_MULT; ++k) - cm_vals[k] *= inv_rs; + cm_vals[k] = cm_vals[k] * inv_rs + hc_sinkhorn_eps; #pragma unroll for (uint32_t k = 0; k < HC_MULT; ++k) { @@ -1419,9 +1451,21 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) cs += __shfl_xor_sync(LANE_MASK, cs, 2); cm_vals[k] *= 1.0f / (cs + hc_sinkhorn_eps); } - } - if (warp_in_team == 0) - { + for (uint32_t it = 1; it < sinkhorn_repeat; ++it) + { + inv_rs = 1.0f / (cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3] + hc_sinkhorn_eps); +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_vals[k] *= inv_rs; +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + { + float cs = cm_vals[k]; + cs += __shfl_xor_sync(LANE_MASK, cs, 1); + cs += __shfl_xor_sync(LANE_MASK, cs, 2); + cm_vals[k] *= 1.0f / (cs + hc_sinkhorn_eps); + } + } float* cm_out_ptr = comb_mix_out + tok * HC_MULT2; #pragma unroll for (uint32_t k = 0; k < HC_MULT; ++k) @@ -1439,10 +1483,6 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) // When WARPS_PER_TOK>1, warp_in_team 0..WARPS_PER_TOK-1 together cover // HIDDEN in strides of WARPS_PER_TOK * 32 * 8. When WARPS_PER_TOK==1, // each warp sweeps HIDDEN alone (same as the original single-warp case). - - // First read of these addresses, ordered by the acquire fence above, so - // __ldg stays for residual read throughput. The layer_input reload below - // is the same-thread store-to-load case and cannot. __nv_bfloat16 const* rbase = residual_cur_ptr + static_cast(tok) * HC_MULT * HIDDEN; __nv_bfloat16* obase = layer_input_out + static_cast(tok) * HIDDEN; @@ -1542,7 +1582,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) // Reduce: intra-warp __shfl_xor; cross-warp via SMEM + per-team // named PTX barrier when WARPS_PER_TOK > 1 (KS ≥ 16 instances). // Inactive teams `continue`'d above and never reach this barrier. - // Pass 2: volatile-reload layer_input_out, multiply by + // Pass 2: re-LDG layer_input_out from L2, multiply by // rsqrt * norm_weight, STG normalized bf16 back to the same // address. Avoids the FMA recompute that doubling pass 1 would // require (Path D Phase 4 is already FMA-heavy). @@ -1654,10 +1694,8 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) float const rsqrt_val = rsqrtf(sum_sq_local / static_cast(HIDDEN) + norm_eps); - // Pass 2: reload the un-normalized layer_input this thread wrote, - // LDG norm_weight, normalize, STG back. No FMA recompute. Each - // thread reloads only its own stores, so no CTA fence is needed -- - // just a load that bypasses the read-only cache. + // Pass 2: re-LDG the un-normalized layer_input we just wrote + // (L2-hot), LDG norm_weight, normalize, STG back. No FMA recompute. __nv_bfloat16 const* nbase = norm_weight; #pragma unroll for (uint32_t h = h_start; h < H_VEC_END; h += H_STRIDE) diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu index 69bd1ec3b9c2..c098a4203b0e 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu @@ -128,10 +128,10 @@ static bool isSupportedFhcHiddenRuntime(int hidden_size) return hidden_size == static_cast(FHC_HIDDEN_FLASH) || hidden_size == static_cast(FHC_HIDDEN_PRO); } -// Validate the tcgen05 MMA fused-HC compile-time shape contract. Hidden must -// be divisible into BLOCK_K tiles, and the hidden dimension must be a multiple -// of BF16_VEC_LI (per-thread vector load granularity in the Phase 4 -// layer_input loop). The (Hidden % team-stride) +// Validate the tcgen05 all-in-one fused-HC compile-time shape contract. Hidden +// must be divisible into BLOCK_K tiles, KS must evenly divide those tiles, and +// the hidden dimension must be a multiple of BF16_VEC_LI (per-thread vector +// load granularity in the Phase 4 layer_input loop). The (Hidden % team-stride) // alignment is no longer required: the layer_input loop has a scalar-vec tail // that handles the residue after the vectorized main loop. Keep this in sync // with the Python tactic filter (_fused_hc_mma_ks_supported in mhc_cuda.py). @@ -144,13 +144,28 @@ static constexpr bool isSupportedFhcMmaKS() constexpr uint32_t hTilesPerHc = Hidden / FHC_BLOCK_K; constexpr uint32_t bf16VecLi = 8; - // The kernels distribute remainder tiles over their first splits, so KS need - // not divide hTilesPerHc. Keep the uneven surface to the two measured H=7168 - // shapes rather than instantiating every split count. constexpr bool evenSplit = hTilesPerHc % KS == 0; - constexpr bool singleWaveSplit = Hidden == FHC_HIDDEN_PRO && (KS == 53 || KS == 106); + constexpr bool rubinExactSplit = Hidden == FHC_HIDDEN_PRO && (KS == 53 || KS == 106); - return Hidden % FHC_BLOCK_K == 0 && KS <= hTilesPerHc && (evenSplit || singleWaveSplit) && Hidden % bf16VecLi == 0; + return Hidden % FHC_BLOCK_K == 0 && KS <= hTilesPerHc && (evenSplit || rubinExactSplit) && Hidden % bf16VecLi == 0; +} + +// The half-MMA kernel can distribute remainder tiles over its first splits. +// Keep the uneven support surface limited to the two measured exact-wave +// H=7168 shapes for 212-SM Rubin. +template +static constexpr bool isSupportedFhcHalfMmaKS() +{ + static_assert(isSupportedFhcHidden(), "Unsupported fused-HC hidden size"); + static_assert(KS > 0, "kNumSplits must be positive"); + + constexpr uint32_t hTilesPerHc = Hidden / FHC_BLOCK_K; + constexpr uint32_t bf16VecLi = 8; + + constexpr bool evenSplit = hTilesPerHc % KS == 0; + constexpr bool rubinExactSplit = Hidden == FHC_HIDDEN_PRO && (KS == 53 || KS == 106); + + return Hidden % FHC_BLOCK_K == 0 && KS <= hTilesPerHc && (evenSplit || rubinExactSplit) && Hidden % bf16VecLi == 0; } static CUtensorMap makeTma2D(void* base, CUtensorMapDataType dtype, uint64_t gmemInner, uint64_t gmemOuter, @@ -178,9 +193,10 @@ static CUtensorMap makeTma2D(void* base, CUtensorMapDataType dtype, uint64_t gme // ---- TMA descriptor cache -------------------------------------------------- // // cuTensorMapEncodeTiled is a host-side call that takes ~1-2 µs per descriptor. -// Each fused_hc launch builds 4 descriptors (residual_in, x_in, W, -// residual_cur), so the per-call descriptor build is 4-8 µs — 25-50% of total -// wall time at small M (M ≤ 64). +// On a cache miss, the half-MMA path builds 6 descriptors (residual_in, x_in, +// W, residual_cur, post_mix, comb_mix) and the all-in-one path builds 4. The +// per-call encode cost is material at small M, so hits stay entirely host-side +// and each launcher resolves the current device only once. // // Cache scope: per-host-thread (`thread_local`). Same host thread launching to // multiple CUDA streams shares one cache (descriptor content depends only on @@ -201,7 +217,7 @@ static CUtensorMap makeTma2D(void* base, CUtensorMapDataType dtype, uint64_t gme // fresh `base` pointers as public outputs are allocated, so the unbounded // version would grow across shape transitions. 128 entries × ~256 B = ~32 KB // per host thread — fits in L1, sized to cover the working set of any single -// model (~4-8 distinct shapes × 4 descriptors each = O(20) live, with +// model (~4-8 distinct shapes × up to 6 descriptors each = O(30) live, with // headroom for shape transitions). namespace { @@ -254,11 +270,10 @@ struct TmaDescCache }; CUtensorMap getCachedTma2D(void* base, CUtensorMapDataType dtype, uint64_t gmemInner, uint64_t gmemOuter, - uint32_t smemInner, uint32_t smemOuter, uint64_t gmemOuterStrideBytes, uint32_t swizzleBytes, uint32_t elemBytes) + uint32_t smemInner, uint32_t smemOuter, uint64_t gmemOuterStrideBytes, uint32_t swizzleBytes, uint32_t elemBytes, + int device_id) { static thread_local TmaDescCache cache; - int device_id = 0; - cudaGetDevice(&device_id); TmaDescKey const key{base, gmemInner, gmemOuter, smemInner, smemOuter, gmemOuterStrideBytes, swizzleBytes, elemBytes, dtype, device_id}; auto it = cache.index.find(key); @@ -294,20 +309,20 @@ static constexpr uint32_t fhcSmemSize() constexpr uint32_t SMEM_COMB = FHC_BLOCK_M * FHC_HC_MULT * FHC_HC_MULT * sizeof(float); constexpr uint32_t SMEM_RC = FHC_HC_MULT * FHC_BLOCK_M * FHC_BLOCK_K * sizeof(__nv_bfloat16); constexpr uint32_t kNumCast = 4; - constexpr uint32_t barriers = 2 * FHC_N_B_STAGES + 2 * FHC_N_INPUT_STG + 2 * kNumCast + 1; + constexpr uint32_t barriers = 2 * FHC_N_B_STAGES + 2 * FHC_N_INPUT_STG + 2 * kNumCast + 2; // 4 bytes for the tmem ptr word + 32 bytes padding for alignment headroom. return SMEM_CD + FHC_N_B_STAGES * SMEM_B + FHC_N_INPUT_STG * (SMEM_RES_ISTG + SMEM_X_ISTG) + SMEM_POST + SMEM_COMB + SMEM_RC + barriers * 8 + 4 + 32; } -using FusedRoutFn = void (*)( - uint32_t, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, float*, float const*, float const*, float*); +using FusedRoutFn = void (*)(uint32_t, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, + float const*, float const*, float*, float*); template static FusedRoutFn fhcInstance() { static_assert(isSupportedFhcHidden(), "Unsupported fused-HC hidden size"); - static_assert(isSupportedFhcMmaKS(), "Unsupported fused-HC MMA kNumSplits for hidden size"); + static_assert(isSupportedFhcHalfMmaKS(), "Unsupported fused-HC half-MMA kNumSplits for hidden size"); return &fused_mhc::fused_tf32_pmap_gemm_rout_atomic_impl; @@ -316,7 +331,7 @@ static FusedRoutFn fhcInstance() template static FusedRoutFn fhcInstanceIfSupported() { - if constexpr (isSupportedFhcMmaKS()) + if constexpr (isSupportedFhcHalfMmaKS()) { return fhcInstance(); } @@ -401,6 +416,8 @@ static void mhcFusedHcLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloat16 cons constexpr uint32_t SHAPE_K = FHC_HC_MULT * Hidden; uint32_t const m_u = static_cast(M); + int device_id = 0; + TLLM_CUDA_CHECK(cudaGetDevice(&device_id)); uint32_t const ks = (num_k_splits > 0) ? static_cast(num_k_splits) : pickKSplits(M); int const bs = (bigfuse_block_size > 0) ? bigfuse_block_size : selectBigFuseBS(M); @@ -416,19 +433,27 @@ static void mhcFusedHcLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloat16 cons // ---- Build TMA descriptors (cached by ptr+shape) ---- CUtensorMap desc_res = getCachedTma2D(const_cast<__nv_bfloat16*>(residual_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, SHAPE_K, m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast(SHAPE_K) * sizeof(__nv_bfloat16), - /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + /*swizzleBytes=*/128, sizeof(__nv_bfloat16), device_id); CUtensorMap desc_x = getCachedTma2D(const_cast<__nv_bfloat16*>(x_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, Hidden, m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast(Hidden) * sizeof(__nv_bfloat16), - /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + /*swizzleBytes=*/128, sizeof(__nv_bfloat16), device_id); CUtensorMap desc_b = getCachedTma2D(const_cast(w_t), CU_TENSOR_MAP_DATA_TYPE_TFLOAT32, SHAPE_K, FHC_SHAPE_N, FHC_BLOCK_K, FHC_BLOCK_N, static_cast(SHAPE_K) * sizeof(float), - /*swizzleBytes=*/128, sizeof(float)); + /*swizzleBytes=*/128, sizeof(float), device_id); CUtensorMap desc_res_out = getCachedTma2D(residual_cur, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, SHAPE_K, m_u, FHC_BLOCK_K, /*smemOuter=*/16, static_cast(SHAPE_K) * sizeof(__nv_bfloat16), - /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + /*swizzleBytes=*/128, sizeof(__nv_bfloat16), device_id); + + CUtensorMap desc_post = getCachedTma2D(const_cast(post_mix_prev), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, + FHC_HC_MULT, m_u, FHC_HC_MULT, FHC_BLOCK_M, static_cast(FHC_HC_MULT) * sizeof(float), + /*swizzleBytes=*/0, sizeof(float), device_id); + + CUtensorMap desc_comb = getCachedTma2D(const_cast(comb_mix_prev), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, + FHC_HC_MULT * FHC_HC_MULT, m_u, FHC_HC_MULT * FHC_HC_MULT, FHC_BLOCK_M, + static_cast(FHC_HC_MULT * FHC_HC_MULT) * sizeof(float), /*swizzleBytes=*/0, sizeof(float), device_id); // ---- Step 1: fused post-mapping + TF32 GEMM + sqrsum + residual_out ---- constexpr uint32_t fused_smem = fhcSmemSize(); @@ -439,8 +464,8 @@ static void mhcFusedHcLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloat16 cons uint32_t const m_tiles = (m_u + FHC_BLOCK_M - 1) / FHC_BLOCK_M; dim3 const grid(m_tiles * ks); dim3 const block(FHC_NUM_MMA_TH + FHC_NUM_PMAP_TH); - fa<<>>( - m_u, desc_res, desc_x, desc_b, desc_res_out, y_acc_workspace, post_mix_prev, comb_mix_prev, r_acc_workspace); + fa<<>>(m_u, desc_res, desc_x, desc_b, desc_res_out, desc_post, desc_comb, + post_mix_prev, comb_mix_prev, y_acc_workspace, r_acc_workspace); // ---- Step 2: big-fuse postlogue (RMS + sigmoid + Sinkhorn + pre-apply) ---- // Delegate to mhcBigFuseLaunch (defined in mhcKernels.cu) to avoid @@ -584,8 +609,8 @@ static constexpr uint32_t fhcAllInOneSmemSize() return fhcSmemSize(); } -using FusedAllInOneFn = void (*)(uint32_t, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, __nv_bfloat16 const*, - __nv_bfloat16*, float*, float*, int*, float const*, float const*, float const*, float const*, float*, float*, +using FusedAllInOneFn = void (*)(uint32_t, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, + __nv_bfloat16 const*, __nv_bfloat16*, float*, float*, int*, float const*, float const*, float*, float*, __nv_bfloat16 const*, float, float, float, float, float, uint32_t); template @@ -657,6 +682,8 @@ static void mhcFusedHcAllInOneLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloa constexpr uint32_t SHAPE_K = FHC_HC_MULT * Hidden; uint32_t const m_u = static_cast(M); + int device_id = 0; + TLLM_CUDA_CHECK(cudaGetDevice(&device_id)); uint32_t const ks = (num_k_splits > 0) ? static_cast(num_k_splits) : 1u; uint32_t const m_tiles = (m_u + FHC_BLOCK_M - 1) / FHC_BLOCK_M; @@ -673,19 +700,27 @@ static void mhcFusedHcAllInOneLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloa // ---- Build TMA descriptors (cached by ptr+shape) ---- CUtensorMap desc_res = getCachedTma2D(const_cast<__nv_bfloat16*>(residual_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, SHAPE_K, m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast(SHAPE_K) * sizeof(__nv_bfloat16), - /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + /*swizzleBytes=*/128, sizeof(__nv_bfloat16), device_id); CUtensorMap desc_x = getCachedTma2D(const_cast<__nv_bfloat16*>(x_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, Hidden, m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast(Hidden) * sizeof(__nv_bfloat16), - /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + /*swizzleBytes=*/128, sizeof(__nv_bfloat16), device_id); CUtensorMap desc_b = getCachedTma2D(const_cast(w_t), CU_TENSOR_MAP_DATA_TYPE_TFLOAT32, SHAPE_K, FHC_SHAPE_N, FHC_BLOCK_K, FHC_BLOCK_N, static_cast(SHAPE_K) * sizeof(float), - /*swizzleBytes=*/128, sizeof(float)); + /*swizzleBytes=*/128, sizeof(float), device_id); CUtensorMap desc_res_out = getCachedTma2D(residual_cur, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, SHAPE_K, m_u, FHC_BLOCK_K, /*smemOuter=*/16, static_cast(SHAPE_K) * sizeof(__nv_bfloat16), - /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + /*swizzleBytes=*/128, sizeof(__nv_bfloat16), device_id); + + CUtensorMap desc_post = getCachedTma2D(const_cast(post_mix_prev), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, + FHC_HC_MULT, m_u, FHC_HC_MULT, FHC_BLOCK_M, static_cast(FHC_HC_MULT) * sizeof(float), + /*swizzleBytes=*/0, sizeof(float), device_id); + + CUtensorMap desc_comb = getCachedTma2D(const_cast(comb_mix_prev), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, + FHC_HC_MULT * FHC_HC_MULT, m_u, FHC_HC_MULT * FHC_HC_MULT, FHC_BLOCK_M, + static_cast(FHC_HC_MULT * FHC_HC_MULT) * sizeof(float), /*swizzleBytes=*/0, sizeof(float), device_id); // ---- Launch the single all-in-one kernel ---- // Dispatch on `norm_weight != nullptr` to a kFuseNorm=true instance that @@ -699,8 +734,8 @@ static void mhcFusedHcAllInOneLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloa dim3 const grid(m_tiles * ks); dim3 const block(FHC_NUM_MMA_TH + FHC_NUM_PMAP_TH); - fa<<>>(m_u, desc_res, desc_x, desc_b, desc_res_out, residual_cur, layer_input_cur, - y_acc_workspace, r_acc_workspace, done_counter_workspace, post_mix_prev, comb_mix_prev, hc_scale, hc_base, + fa<<>>(m_u, desc_res, desc_x, desc_b, desc_res_out, desc_post, desc_comb, + residual_cur, layer_input_cur, y_acc_workspace, r_acc_workspace, done_counter_workspace, hc_scale, hc_base, post_mix_cur, comb_mix_cur, norm_weight, norm_eps, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, static_cast(sinkhorn_repeat)); } diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu index 48d2a9e5e3a7..2c3129fa2a7a 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu @@ -27,6 +27,22 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels::mhc { +// Blackwell can issue two independent FP32 FMAs with one packed instruction. +// Phase 2 always updates adjacent bf16 values with the same pre-mix scalar. +__device__ __forceinline__ float2 mhcFmaF32x2(float2 const& a, float2 const& b, float2 const& c) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (__CUDA_ARCH__ < 1100) + float2 result; + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(reinterpret_cast(result)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b)), + "l"(reinterpret_cast(c))); + return result; +#else + return make_float2(fmaf(a.x, b.x, c.x), fmaf(a.y, b.y, c.y)); +#endif +} + // =================================================================== // Kernel 1: big_fuse — one CTA per token // @@ -37,11 +53,11 @@ namespace kernels::mhc // Phase 1a (warp 0, lanes 0-3): RMS norm + sigmoid → s_pre_mix, post_mix // ── __syncthreads ── // Phase 1b (warp 0, lanes 0-3) ‖ Phase 2 (remaining warps) — overlapped -// 1b: parallel Sinkhorn (4 lanes, __shfl_xor col normalize) → comb_mix +// 1b: one 4-element Sinkhorn row per lane → comb_mix // 2: stream residual × pre_mix → layer_input // =================================================================== -template +template __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __restrict__ y_acc, float const* __restrict__ r_acc, __nv_bfloat16 const* __restrict__ residual, float const* __restrict__ hc_scale, float const* __restrict__ hc_base, float* __restrict__ post_mix, float* __restrict__ comb_mix, @@ -67,40 +83,73 @@ __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __re cudaGridDependencySynchronize(); #endif + // Small-M/BS512 specialization: issue one bulk asynchronous copy for the + // token's contiguous [4, hidden] residual before Phase 1a, then overlap it + // with RMS/sigmoid/Sinkhorn. Consumers wait on the mbarrier independently, + // so warp 0 can keep running Sinkhorn while the copy completes. + extern __shared__ __align__(16) unsigned char s_tma_residual_raw[]; + __shared__ alignas(8) uint64_t s_tma_bar; + if constexpr (kUseTma) + { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + uint32_t const tma_bytes = HC_MULT * hidden_size * static_cast(sizeof(__nv_bfloat16)); + if (tid == 0) + { + uint32_t const bar_addr = static_cast(__cvta_generic_to_shared(&s_tma_bar)); + uint32_t const dst_addr = static_cast(__cvta_generic_to_shared(s_tma_residual_raw)); + __nv_bfloat16 const* src = residual + static_cast(token) * HC_MULT * hidden_size; + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" : : "r"(bar_addr) : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" + : + : "r"(bar_addr), "r"(tma_bytes) + : "memory"); + asm volatile("cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" + : + : "r"(dst_addr), "l"(reinterpret_cast(src)), "r"(tma_bytes), "r"(bar_addr) + : "memory"); + } +#endif + } + __shared__ float s_pre_mix[HC_MULT]; float cm[HC_MULT]; // ---- Phase 1a (warp 0, lanes 0..3): split-K reduce → RMS norm → sigmoid ---- - // Each lane handles one of the 4 hc_mult slots. - // Produces: s_pre_mix[4] (shared), post_mix[4] (global), cm[4×4] (registers). + // Keep one full comb row per lane: its four independent values provide ILP + // across Sinkhorn's reciprocal/shuffle dependency chains. Load only the + // six y values this lane actually consumes instead of materializing all 24. if (warp_id == 0 && lane < HC_MULT) { float r_val; - float y_local[HC_MULT3]; + float y_pre = 0.0f; + float y_post = 0.0f; +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + cm[k] = 0.0f; if constexpr (NUM_SPLITS == 1) { r_val = r_acc[token]; float const* y_row = y_acc + token * HC_MULT3; + y_pre = y_row[lane]; + y_post = y_row[HC_MULT + lane]; #pragma unroll - for (int c = 0; c < HC_MULT3; c++) - y_local[c] = y_row[c]; + for (int k = 0; k < HC_MULT; k++) + cm[k] = y_row[2 * HC_MULT + lane * HC_MULT + k]; } else { - // Reduce across split-K partials r_val = 0.0f; -#pragma unroll - for (int c = 0; c < HC_MULT3; c++) - y_local[c] = 0.0f; for (int s = 0; s < NUM_SPLITS; s++) { r_val += r_acc[s * M + token]; float const* y_row = y_acc + (static_cast(s) * M + token) * HC_MULT3; + y_pre += y_row[lane]; + y_post += y_row[HC_MULT + lane]; #pragma unroll - for (int c = 0; c < HC_MULT3; c++) - y_local[c] += y_row[c]; + for (int k = 0; k < HC_MULT; k++) + cm[k] += y_row[2 * HC_MULT + lane * HC_MULT + k]; } } @@ -108,76 +157,60 @@ __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __re float const rstd = rsqrtf(r_val / static_cast(K) + rms_eps); float const s0 = hc_scale[0], s1 = hc_scale[1], s2 = hc_scale[2]; - // y_local layout: [pre_mix(4) | post_mix(4) | comb_mix(4×4)] - // pre_mix: sigmoid(norm * scale0 + base) + eps → shared for Phase 2 - float v = y_local[lane] * rstd * s0 + hc_base[lane]; + float v = y_pre * rstd * s0 + hc_base[lane]; s_pre_mix[lane] = 1.0f / (1.0f + expf(-v)) + hc_pre_eps; - // post_mix: sigmoid(norm * scale1 + base) * mult → global - v = y_local[HC_MULT + lane] * rstd * s1 + hc_base[HC_MULT + lane]; + v = y_post * rstd * s1 + hc_base[HC_MULT + lane]; post_mix[token * HC_MULT + lane] = 1.0f / (1.0f + expf(-v)) * hc_post_mult_value; - // comb_mix init: norm * scale2 + base → cm[4] per lane (one row of 4×4 matrix) #pragma unroll for (int k = 0; k < HC_MULT; k++) - cm[k] = y_local[2 * HC_MULT + lane * HC_MULT + k] * rstd * s2 + hc_base[2 * HC_MULT + lane * HC_MULT + k]; + cm[k] = cm[k] * rstd * s2 + hc_base[2 * HC_MULT + lane * HC_MULT + k]; } __syncthreads(); // ---- Phase 1b (warp 0, lanes 0..3): Sinkhorn normalization ---- - // Each lane holds one row of the 4×4 comb matrix. - // Row normalize via local sum, column normalize via __shfl_xor across 4 lanes. if (warp_id == 0 && lane < HC_MULT) { - constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; // 0xf for HC_MULT=4 + constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; // Softmax rows: subtract the row max to avoid inf / inf when comb logits // are large. - float const rowMax = fmaxf(fmaxf(cm[0], cm[1]), fmaxf(cm[2], cm[3])); + float const row_max = fmaxf(fmaxf(cm[0], cm[1]), fmaxf(cm[2], cm[3])); #pragma unroll for (int k = 0; k < HC_MULT; k++) - cm[k] = expf(cm[k] - rowMax); - // Replace per-element fdiv with one reciprocal + 4 fmul. fp32 fdiv on - // B200 is multi-cycle while fmul retires at peak rate; sinkhorn's - // O(HC_MULT * sinkhorn_repeat) divisions per token (160 at sinkhorn=20) - // dominate the bigfuse epilogue cost on this 4-lane warp. Math is - // identical modulo last-bit round-off, which sinkhorn iteration - // absorbs. - float inv_rs = 1.0f / (cm[0] + cm[1] + cm[2] + cm[3]); + cm[k] = expf(cm[k] - row_max); + float inv_row_sum = 1.0f / (cm[0] + cm[1] + cm[2] + cm[3]); #pragma unroll for (int k = 0; k < HC_MULT; k++) - cm[k] = cm[k] * inv_rs + hc_sinkhorn_eps; - - // Column normalize: sum across lanes (rows) via butterfly shuffle + cm[k] = cm[k] * inv_row_sum + hc_sinkhorn_eps; #pragma unroll for (int k = 0; k < HC_MULT; k++) { - float cs = cm[k]; - cs += __shfl_xor_sync(LANE_MASK, cs, 1); - cs += __shfl_xor_sync(LANE_MASK, cs, 2); - cm[k] *= 1.0f / (cs + hc_sinkhorn_eps); + float col_sum = cm[k]; + col_sum += __shfl_xor_sync(LANE_MASK, col_sum, 1); + col_sum += __shfl_xor_sync(LANE_MASK, col_sum, 2); + cm[k] *= 1.0f / (col_sum + hc_sinkhorn_eps); } // Remaining Sinkhorn iterations: alternate row / column normalize for (int it = 1; it < sinkhorn_repeat; it++) { - inv_rs = 1.0f / (cm[0] + cm[1] + cm[2] + cm[3] + hc_sinkhorn_eps); + inv_row_sum = 1.0f / (cm[0] + cm[1] + cm[2] + cm[3] + hc_sinkhorn_eps); #pragma unroll for (int k = 0; k < HC_MULT; k++) - cm[k] *= inv_rs; - + cm[k] *= inv_row_sum; #pragma unroll for (int k = 0; k < HC_MULT; k++) { - float cs = cm[k]; - cs += __shfl_xor_sync(LANE_MASK, cs, 1); - cs += __shfl_xor_sync(LANE_MASK, cs, 2); - cm[k] *= 1.0f / (cs + hc_sinkhorn_eps); + float col_sum = cm[k]; + col_sum += __shfl_xor_sync(LANE_MASK, col_sum, 1); + col_sum += __shfl_xor_sync(LANE_MASK, col_sum, 2); + cm[k] *= 1.0f / (col_sum + hc_sinkhorn_eps); } } - // Write 4×4 comb_mix to global (lane = row index, k = col index) float* cm_out = comb_mix + token * HC_MULT2; #pragma unroll for (int k = 0; k < HC_MULT; k++) @@ -206,7 +239,30 @@ __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __re for (int j = 0; j < HC_MULT; j++) pm[j] = s_pre_mix[j]; - __nv_bfloat16 const* rbase = residual + static_cast(token) * HC_MULT * hidden_size; + __nv_bfloat16 const* rbase; + if constexpr (kUseTma) + { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + uint32_t const bar_addr = static_cast(__cvta_generic_to_shared(&s_tma_bar)); + uint32_t complete; + do + { + asm volatile( + "{ .reg .pred P; mbarrier.try_wait.parity.shared::cta.b64 P, [%1], %2;" + " selp.b32 %0, 1, 0, P; }" + : "=r"(complete) + : "r"(bar_addr), "r"(0u) + : "memory"); + } while (!complete); + rbase = reinterpret_cast<__nv_bfloat16 const*>(s_tma_residual_raw); +#else + rbase = residual + static_cast(token) * HC_MULT * hidden_size; +#endif + } + else + { + rbase = residual + static_cast(token) * HC_MULT * hidden_size; + } __nv_bfloat16* obase = layer_input + static_cast(token) * hidden_size; int const p2_tid = tid - WARP_SIZE; @@ -215,19 +271,19 @@ __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __re float sum_sq_local = 0.f; for (int h = p2_tid * BF16_VEC; h < hidden_size; h += p2_threads * BF16_VEC) { - float acc[BF16_VEC] = {}; + float2 acc[BF16_VEC / 2] = {}; #pragma unroll for (int j = 0; j < HC_MULT; j++) { uint4 raw = *reinterpret_cast(&rbase[j * hidden_size + h]); __nv_bfloat162 const* pairs = reinterpret_cast<__nv_bfloat162 const*>(&raw); + float2 const coefficient = make_float2(pm[j], pm[j]); #pragma unroll for (int v = 0; v < BF16_VEC / 2; v++) { - float2 f = __bfloat1622float2(pairs[v]); - acc[2 * v + 0] += pm[j] * f.x; - acc[2 * v + 1] += pm[j] * f.y; + float2 const f = __bfloat1622float2(pairs[v]); + acc[v] = mhcFmaF32x2(coefficient, f, acc[v]); } } @@ -235,7 +291,7 @@ __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __re __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); #pragma unroll for (int v = 0; v < BF16_VEC / 2; v++) - opairs[v] = __float22bfloat162_rn(make_float2(acc[2 * v], acc[2 * v + 1])); + opairs[v] = __float22bfloat162_rn(acc[v]); *reinterpret_cast(&obase[h]) = out_raw; if constexpr (kFuseNorm) @@ -327,11 +383,19 @@ __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __re // turns those otherwise-ignored qualifiers into part of the function type. Clang then rejects the instantiation; // GCC does not diagnose it. #define INST_BIGFUSE(NS, BS) \ - template __global__ void mhcBigFuseKernel(float const* __restrict__, \ - float const* __restrict__, __nv_bfloat16 const* __restrict__, float const* __restrict__, \ - float const* __restrict__, float* __restrict__, float* __restrict__, __nv_bfloat16* __restrict__, int, int, \ - int, float, float, float, float, int, __nv_bfloat16 const* __restrict__, float); \ - template __global__ void mhcBigFuseKernel(float const* __restrict__, \ + template __global__ void mhcBigFuseKernel( \ + float const* __restrict__, float const* __restrict__, __nv_bfloat16 const* __restrict__, \ + float const* __restrict__, float const* __restrict__, float* __restrict__, float* __restrict__, \ + __nv_bfloat16* __restrict__, int, int, int, float, float, float, float, int, \ + __nv_bfloat16 const* __restrict__, float); \ + template __global__ void mhcBigFuseKernel( \ + float const* __restrict__, float const* __restrict__, __nv_bfloat16 const* __restrict__, \ + float const* __restrict__, float const* __restrict__, float* __restrict__, float* __restrict__, \ + __nv_bfloat16* __restrict__, int, int, int, float, float, float, float, int, \ + __nv_bfloat16 const* __restrict__, float); + +#define INST_BIGFUSE_TMA(FN) \ + template __global__ void mhcBigFuseKernel<1, 512, /*kFuseNorm=*/FN, /*kUseTma=*/true>(float const* __restrict__, \ float const* __restrict__, __nv_bfloat16 const* __restrict__, float const* __restrict__, \ float const* __restrict__, float* __restrict__, float* __restrict__, __nv_bfloat16* __restrict__, int, int, \ int, float, float, float, float, int, __nv_bfloat16 const* __restrict__, float); @@ -351,6 +415,9 @@ INST_BIGFUSE(8, 512) INST_BIGFUSE(16, 128) INST_BIGFUSE(16, 256) INST_BIGFUSE(16, 512) +INST_BIGFUSE_TMA(false) +INST_BIGFUSE_TMA(true) +#undef INST_BIGFUSE_TMA #undef INST_BIGFUSE // =================================================================== @@ -791,22 +858,44 @@ static void mhcBigFuseDispatch(float const* y_acc, float const* r_acc, __nv_bflo { dim3 grid(static_cast(M)); -#define LAUNCH_BF(BS) \ - tensorrt_llm::common::launchWithPdlWhenEnabled("mhcBigFuseKernel", mhcBigFuseKernel, \ - grid, dim3(BS), 0, stream, y_acc, r_acc, residual, hc_scale, hc_base, post_mix, comb_mix, layer_input, M, K, \ - hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, norm_weight, norm_eps) +#define LAUNCH_BF(BS, USE_TMA, SMEM_BYTES) \ + tensorrt_llm::common::launchWithPdlWhenEnabled("mhcBigFuseKernel", \ + mhcBigFuseKernel, grid, dim3(BS), SMEM_BYTES, stream, y_acc, r_acc, \ + residual, hc_scale, hc_base, post_mix, comb_mix, layer_input, M, K, hidden_size, rms_eps, hc_pre_eps, \ + hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, norm_weight, norm_eps) if (block_size >= 512) { - LAUNCH_BF(512); + if constexpr (NUM_SPLITS == 1) + { + // BS512 is the measured winner for the small-M one-wave regime. Bulk + // TMA wins only through M=64; its shared-memory footprint regresses + // M>=128, where regular LDG has enough latency hiding. + if (M <= 64) + { + size_t const tma_smem_bytes = static_cast(4) * hidden_size * sizeof(__nv_bfloat16); + TLLM_CUDA_CHECK(cudaFuncSetAttribute( + reinterpret_cast(mhcBigFuseKernel), + cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(tma_smem_bytes))); + LAUNCH_BF(512, true, tma_smem_bytes); + } + else + { + LAUNCH_BF(512, false, 0); + } + } + else + { + LAUNCH_BF(512, false, 0); + } } else if (block_size >= 256) { - LAUNCH_BF(256); + LAUNCH_BF(256, false, 0); } else { - LAUNCH_BF(128); + LAUNCH_BF(128, false, 0); } #undef LAUNCH_BF } diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh index 48a1290df564..1ce437a24549 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh @@ -798,9 +798,9 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1 // Phase 4: inline bigFuse for the TM tokens in this batch. // Layout: FULL_N = HC_MULT*(2+HC_MULT) = 24 - // y_local[0..HC_MULT) → s_pre_mix (sigmoid gate) - // y_local[HC_MULT..2*HC_MULT) → post_mix_out (sigmoid*hc_post_mult) - // y_local[2*HC_MULT..FULL_N) → comb_mix_out (Sinkhorn) + // y_acc[0..HC_MULT) → s_pre_mix (sigmoid gate) + // y_acc[HC_MULT..2*HC_MULT) → post_mix_out (sigmoid*hc_post_mult) + // y_acc[2*HC_MULT..FULL_N) → comb_mix_out (Sinkhorn) __shared__ float s_pre_mix[TM][HC_MULT]; #pragma unroll @@ -812,30 +812,26 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1 continue; } int const tok = base_tok + t; - float cm_vals[HC_MULT]; if (warp_id == 0 && lane < HC_MULT) { + float cm_vals[HC_MULT]; float const r_val = r_acc[tok]; - float y_local[HC_MULT3]; float const* y_row = y_acc + static_cast(tok) * FULL_N; -#pragma unroll - for (int c = 0; c < HC_MULT3; c++) - y_local[c] = y_row[c]; float const rstd = rsqrtf(r_val / static_cast(K) + rms_eps); float const s0 = hc_scale[0], s1 = hc_scale[1], s2 = hc_scale[2]; - float v = y_local[lane] * rstd * s0 + hc_base[lane]; + float v = y_row[lane] * rstd * s0 + hc_base[lane]; s_pre_mix[t][lane] = 1.0f / (1.0f + expf(-v)) + hc_pre_eps; - v = y_local[HC_MULT + lane] * rstd * s1 + hc_base[HC_MULT + lane]; + v = y_row[HC_MULT + lane] * rstd * s1 + hc_base[HC_MULT + lane]; post_mix_out[tok * HC_MULT + lane] = 1.0f / (1.0f + expf(-v)) * hc_post_mult_value; #pragma unroll for (int k = 0; k < HC_MULT; k++) cm_vals[k] - = y_local[2 * HC_MULT + lane * HC_MULT + k] * rstd * s2 + hc_base[2 * HC_MULT + lane * HC_MULT + k]; + = y_row[2 * HC_MULT + lane * HC_MULT + k] * rstd * s2 + hc_base[2 * HC_MULT + lane * HC_MULT + k]; constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; float const rowMax = fmaxf(fmaxf(cm_vals[0], cm_vals[1]), fmaxf(cm_vals[2], cm_vals[3])); @@ -877,19 +873,18 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1 } __syncthreads(); - // layer_input: warp>0 threads process hidden in parallel. + // layer_input: all eight warps process hidden in parallel. // // When kFuseNorm is true, accumulate per-thread sum_sq while writing // un-normalized layer_input bf16; after a __syncthreads, all threads // run pass 2: re-LDG from L2 (hot from pass 1's STGs), multiply by // rsqrt * norm_weight, STG normalized bf16. Saves one HBM read+write // pair vs the separate flashinfer.rmsnorm kernel. - constexpr int kFmaBigFuseWarps = BLOCK_SIZE / WARP_SIZE - 1; + constexpr int kFmaBigFuseWarps = BLOCK_SIZE / WARP_SIZE; __shared__ float s_sumsq_li[kFmaBigFuseWarps]; __shared__ float s_rsqrt_li; constexpr int BF16_VEC_LI = 8; __nv_bfloat16* obase = layer_input_out + static_cast(tok) * hidden_size; - if (warp_id > 0) { float pm[HC_MULT]; #pragma unroll @@ -897,8 +892,8 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1 pm[j] = s_pre_mix[t][j]; __nv_bfloat16 const* rbase = residual_out + static_cast(tok) * HC_MULT * hidden_size; - int const p2_tid = tid - WARP_SIZE; - constexpr int p2_threads = BLOCK_SIZE - WARP_SIZE; + int const p2_tid = tid; + constexpr int p2_threads = BLOCK_SIZE; float sum_sq_local = 0.f; for (int h = p2_tid * BF16_VEC_LI; h < hidden_size; h += p2_threads * BF16_VEC_LI) @@ -942,7 +937,7 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1 sum_sq_local += __shfl_xor_sync(0xffffffff, sum_sq_local, 2); sum_sq_local += __shfl_xor_sync(0xffffffff, sum_sq_local, 1); if ((tid & 31) == 0) - s_sumsq_li[warp_id - 1] = sum_sq_local; + s_sumsq_li[warp_id] = sum_sq_local; } } if constexpr (kFuseNorm) @@ -957,10 +952,9 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1 s_rsqrt_li = rsqrtf(total / static_cast(hidden_size) + norm_eps); } __syncthreads(); - if (warp_id > 0) { - int const p2_tid = tid - WARP_SIZE; - constexpr int p2_threads = BLOCK_SIZE - WARP_SIZE; + int const p2_tid = tid; + constexpr int p2_threads = BLOCK_SIZE; float const rsqrt_val = s_rsqrt_li; for (int h = p2_tid * BF16_VEC_LI; h < hidden_size; h += p2_threads * BF16_VEC_LI) { diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.cu b/cpp/tensorrt_llm/kernels/mlaKernels.cu index 422569c3b3b2..1d984adeb49a 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.cu +++ b/cpp/tensorrt_llm/kernels/mlaKernels.cu @@ -406,15 +406,16 @@ inline __device__ void dequantCopy( // `kOutputFp8Q`: when true, write the rotated Q rope segment directly to // `quant_q_buf` as FP8 (scaled by `*quant_scale_qkv`) and skip the bf16 STG to -// `q_ptr`. Companion: `deepseek_v4_q_norm_fused_fp8` pre-fills the nope segment -// of `quant_q_buf`, so the standalone quantizeCopyInputToFp8Kernel can be -// dropped. `quant_q_buf`/`quant_scale_qkv`/bmm_scale outputs are unused when -// `kOutputFp8Q == false`. +// `q_ptr`. `deepseek_v4_q_norm_fused_fp8` can pre-fill the nope segment for +// this kernel to append RoPE, while the q_b fusion can pre-fill the complete +// rotated Q and set `q_rope_applied` to preserve it. Both paths drop the +// standalone quantizeCopyInputToFp8Kernel. `quant_q_buf`/`quant_scale_qkv` and +// bmm-scale outputs are unused when `kOutputFp8Q == false`. template __global__ void applyMLARopeAndAssignQKVKernelOptContext(T* q_ptr, T* q_pe, T* k_ptr, T const* fuse_buf, KVCacheBuffer kv_cache, int q_pe_ld, int q_pe_stride, float2 const* cos_sin_cache, size_t head_num, int head_size, int c_k, int* cu_q_seqlens, int32_t const* kv_cache_lengths, uint32_t max_input_seq_len, KvCacheDataType cache_type, - float const* quant_scale_kv, int32_t const* helix_position_offsets, bool absorption_mode, + float const* quant_scale_kv, int32_t const* helix_position_offsets, bool absorption_mode, bool q_rope_applied, __nv_fp8_e4m3* quant_q_buf = nullptr, float const* quant_scale_qkv = nullptr, float* bmm1_scale_out = nullptr, float* bmm2_scale_out = nullptr, float const* dequant_scale_q = nullptr, float const* dequant_scale_kv = nullptr, float const* quant_scale_o = nullptr, float host_bmm1_scale = 1.0f) @@ -501,19 +502,24 @@ __global__ void applyMLARopeAndAssignQKVKernelOptContext(T* q_ptr, T* q_pe, T* k float2 const* rotary_coef_cache_buffer = cos_sin_cache + static_cast(ROPE_DIM) * position_id + (head_dim_idx / 2); - VecT q, k; + // When q_rope_applied is true, Q is already complete in the FP8 + // output buffer. Zero-initialize the unused Q fragment so the + // shared GPT-J helper can still rotate K without reading q_pe. + VecT q{}, k; auto const src_k_global_offset = static_cast(global_token_idx) * (c_k + ROPE_DIM) + c_k; - auto src_q_global_offset = static_cast(global_token_idx) * head_num * (head_size + ROPE_DIM) - + (head_size + ROPE_DIM) * head_idx + head_size; - // In the absorption mode, we load pe from q_pe instead of q_ptr. - T* q_pe_input = q_ptr; - if (absorption_mode) + if (!q_rope_applied) { - q_pe_input = q_pe; - src_q_global_offset = static_cast(global_token_idx) * q_pe_stride + q_pe_ld * head_idx; + auto src_q_global_offset = static_cast(global_token_idx) * head_num * (head_size + ROPE_DIM) + + (head_size + ROPE_DIM) * head_idx + head_size; + // In absorption mode, load the positional segment from q_pe. + T* q_pe_input = q_ptr; + if (absorption_mode) + { + q_pe_input = q_pe; + src_q_global_offset = static_cast(global_token_idx) * q_pe_stride + q_pe_ld * head_idx; + } + q = *reinterpret_cast(&q_pe_input[src_q_global_offset + head_dim_idx]); } - - q = *reinterpret_cast(&q_pe_input[src_q_global_offset + head_dim_idx]); k = *reinterpret_cast(&fuse_buf[src_k_global_offset + head_dim_idx]); // Pack two elements into one for gptj rotary embedding. @@ -548,14 +554,17 @@ __global__ void applyMLARopeAndAssignQKVKernelOptContext(T* q_ptr, T* q_pe, T* k + head_idx * (nope_head_size_q + ROPE_DIM) + nope_head_size_q + head_dim_idx; auto const dst_k_idx = static_cast(global_token_idx) * head_num * (head_size + ROPE_DIM) + head_idx * (head_size + ROPE_DIM) + head_size + head_dim_idx; - if constexpr (kOutputFp8Q) + if (!q_rope_applied) { - quantCopy( - quant_q_buf + dst_q_idx, reinterpret_cast(&q), quant_scale_qkv_val); - } - else - { - reinterpret_cast(q_ptr)[dst_q_idx / ELTS_PER_VEC] = q; + if constexpr (kOutputFp8Q) + { + quantCopy( + quant_q_buf + dst_q_idx, reinterpret_cast(&q), quant_scale_qkv_val); + } + else + { + reinterpret_cast(q_ptr)[dst_q_idx / ELTS_PER_VEC] = q; + } } // Only write to k_pe to k_buf in the non-absorption mode. if (!absorption_mode) @@ -626,7 +635,8 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_pe, int* seqKVOffsets, int q_pe_ld, int q_pe_stride, KvCacheDataType cache_type, float* bmm1_scale, float* bmm2_scale, float const* quant_scale_o, float const* quant_scale_q, float const* quant_scale_kv, float const* dequant_scale_q, float const* dequant_scale_kv, float host_bmm1_scale, int32_t const* helix_position_offsets, - bool const* helix_is_inactive_rank, bool precomputed_cu_seqlens = false, bool precomputed_fmha_scheduler = false) + bool const* helix_is_inactive_rank, bool precomputed_cu_seqlens = false, bool precomputed_fmha_scheduler = false, + bool q_rope_applied = false) { // Constants. using VecT = typename VecType::Type; @@ -645,7 +655,9 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_pe, constexpr auto TOTAL_VEC_PER_HEAD = VECS_PER_HEAD + K_VECS_PER_HEAD; // Block/Head idx. - size_t const head_idx = blockIdx.y; + // Remap a compact pre-rotated-Q launch directly onto the K-RoPE and + // KV-copy work, skipping the Q-RoPE and Q-nope quantization blocks. + size_t const head_idx = q_rope_applied ? blockIdx.y + head_num : blockIdx.y; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); #endif @@ -806,7 +818,7 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_pe, } else if (head_idx <= head_num + 8) { - int block_dim = gridDim.y - head_num - 1; + constexpr int block_dim = 8; int block_id = head_idx - head_num - 1; size_t const head_dim_vec_idx = (threadIdx.x % K_VECS_PER_HEAD); size_t const head_dim_idx = head_dim_vec_idx * ELTS_PER_VEC; @@ -1720,13 +1732,13 @@ void invokeMLARopeContext(MlaParams& params, KVCacheBuffer kv_cache_buffer, c { if (useFusedFp8Q) { - applyMLARopeAndAssignQKVKernelOptContext - <<>>(params.q_buf, params.q_pe, params.k_buf, params.latent_cache, - kv_cache_buffer, params.q_pe_ld, params.q_pe_stride, params.cos_sin_cache, params.head_num, - head_size, params.meta.kv_lora_rank, params.cu_q_seqlens, params.cache_seq_lens, - params.max_input_seq_len, params.cache_type, params.quant_scale_kv, params.helix_position_offsets, - params.absorption_mode, quant_q_fp8, params.quant_scale_qkv, params.bmm1_scale, params.bmm2_scale, - params.dequant_scale_q, params.dequant_scale_kv, params.quant_scale_o, params.host_bmm1_scale); + applyMLARopeAndAssignQKVKernelOptContext<<>>( + params.q_buf, params.q_pe, params.k_buf, params.latent_cache, kv_cache_buffer, params.q_pe_ld, + params.q_pe_stride, params.cos_sin_cache, params.head_num, head_size, params.meta.kv_lora_rank, + params.cu_q_seqlens, params.cache_seq_lens, params.max_input_seq_len, params.cache_type, + params.quant_scale_kv, params.helix_position_offsets, params.absorption_mode, params.q_rope_applied, + quant_q_fp8, params.quant_scale_qkv, params.bmm1_scale, params.bmm2_scale, params.dequant_scale_q, + params.dequant_scale_kv, params.quant_scale_o, params.host_bmm1_scale); } else { @@ -1734,7 +1746,7 @@ void invokeMLARopeContext(MlaParams& params, KVCacheBuffer kv_cache_buffer, c params.q_buf, params.q_pe, params.k_buf, params.latent_cache, kv_cache_buffer, params.q_pe_ld, params.q_pe_stride, params.cos_sin_cache, params.head_num, head_size, params.meta.kv_lora_rank, params.cu_q_seqlens, params.cache_seq_lens, params.max_input_seq_len, params.cache_type, - params.quant_scale_kv, params.helix_position_offsets, params.absorption_mode); + params.quant_scale_kv, params.helix_position_offsets, params.absorption_mode, params.q_rope_applied); } } else @@ -1781,13 +1793,13 @@ void invokeMLARopeContext(MlaParams& params, KVCacheBuffer kv_cache_buffer, c if (useFusedFp8Q) { - applyMLARopeAndAssignQKVKernelOptContext - <<>>(params.q_buf, params.q_pe, params.k_buf, params.latent_cache, - kv_cache_buffer, params.q_pe_ld, params.q_pe_stride, params.cos_sin_cache, params.head_num, - head_size, params.meta.kv_lora_rank, params.cu_q_seqlens, params.cache_seq_lens, - params.max_input_seq_len, params.cache_type, params.quant_scale_kv, params.helix_position_offsets, - params.absorption_mode, quant_q_fp8, params.quant_scale_qkv, params.bmm1_scale, params.bmm2_scale, - params.dequant_scale_q, params.dequant_scale_kv, params.quant_scale_o, params.host_bmm1_scale); + applyMLARopeAndAssignQKVKernelOptContext<<>>( + params.q_buf, params.q_pe, params.k_buf, params.latent_cache, kv_cache_buffer, params.q_pe_ld, + params.q_pe_stride, params.cos_sin_cache, params.head_num, head_size, params.meta.kv_lora_rank, + params.cu_q_seqlens, params.cache_seq_lens, params.max_input_seq_len, params.cache_type, + params.quant_scale_kv, params.helix_position_offsets, params.absorption_mode, params.q_rope_applied, + quant_q_fp8, params.quant_scale_qkv, params.bmm1_scale, params.bmm2_scale, params.dequant_scale_q, + params.dequant_scale_kv, params.quant_scale_o, params.host_bmm1_scale); } else { @@ -1795,7 +1807,7 @@ void invokeMLARopeContext(MlaParams& params, KVCacheBuffer kv_cache_buffer, c params.q_buf, params.q_pe, params.k_buf, params.latent_cache, kv_cache_buffer, params.q_pe_ld, params.q_pe_stride, params.cos_sin_cache, params.head_num, head_size, params.meta.kv_lora_rank, params.cu_q_seqlens, params.cache_seq_lens, params.max_input_seq_len, params.cache_type, - params.quant_scale_kv, params.helix_position_offsets, params.absorption_mode); + params.quant_scale_kv, params.helix_position_offsets, params.absorption_mode, params.q_rope_applied); } } } @@ -1906,8 +1918,11 @@ void invokeMLARopeGeneration(MlaParams& params, KVCacheBuffer kv_cache_buffer bool const useFusedFp8Q = params.fuse_q_fp8_in_rope && params.cache_type == KvCacheDataType::FP8 && params.quant_q_buf != nullptr && params.quant_scale_qkv != nullptr; - dim3 grid(int(tensorrt_llm::common::divUp(params.acc_q_len, 32)), params.head_num + 1 + 8); - if ((params.cache_type == KvCacheDataType::FP8 && !useFusedFp8Q) || params.cache_type == KvCacheDataType::NVFP4) + dim3 grid(int(tensorrt_llm::common::divUp(params.acc_q_len, 32)), + params.q_rope_applied ? 1 + 8 : params.head_num + 1 + 8); + if (!params.q_rope_applied + && ((params.cache_type == KvCacheDataType::FP8 && !useFusedFp8Q) + || params.cache_type == KvCacheDataType::NVFP4)) grid.y += params.head_num * 8; TLLM_CHECK_WITH_INFO(params.acc_q_len % params.batch_size == 0, "MLA can only support input sequences with the same sequence length."); @@ -1946,7 +1961,8 @@ void invokeMLARopeGeneration(MlaParams& params, KVCacheBuffer kv_cache_buffer params.cache_seq_lens, params.cu_kv_seqlens, params.q_pe_ld, params.q_pe_stride, params.cache_type, params.bmm1_scale, params.bmm2_scale, params.quant_scale_o, quant_scale_q_eff, params.quant_scale_kv, params.dequant_scale_q, params.dequant_scale_kv, params.host_bmm1_scale, params.helix_position_offsets, - params.helix_is_inactive_rank, params.precomputed_cu_seqlens, params.precomputed_fmha_scheduler); + params.helix_is_inactive_rank, params.precomputed_cu_seqlens, params.precomputed_fmha_scheduler, + params.q_rope_applied); } template diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.h b/cpp/tensorrt_llm/kernels/mlaKernels.h index 6a71f4369ca3..a0125d236a89 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.h +++ b/cpp/tensorrt_llm/kernels/mlaKernels.h @@ -145,6 +145,9 @@ struct MlaParams // `latent_cache` row stride in elements; the fused path passes a slice of // kv_a_proj, so rows are wider than packed. 0 means packed. int latent_row_stride = 0; + // The caller has already applied Q RoPE and written the complete FP8 Q to + // `quant_q_buf`. Context preprocessing must still rotate/cache K. + bool q_rope_applied = false; // DSv4 fused inverse-RoPE + FP8 quant epilogue parameters. Dsv4EpilogueFusionParams dsv4_epilogue_fusion; diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 6c2e03c6415e..2e9472a064fc 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -1151,7 +1151,7 @@ class TllmGenFmhaKernel options.mIsCustomSpecDecodingGen = !isContext && params.mMaxSeqLenQ > 1 && params.mIsSpecDecTree; options.mIsCausalSpecDecodingGen = !isContext && params.mMaxSeqLenQ > 1 && !params.mIsSpecDecTree; options.mNumSpecDecodingTokens = !isContext && params.mMaxSeqLenQ > 1 ? params.mMaxSeqLenQ : 0; - // Carry static tree length into FMHA kernel selection. + // Carry the tree length into FMHA kernel selection. options.mSpecDecodingTargetMaxGenLen = params.mSpecDecodingTargetMaxGenLen; options.mIsTrtllmLayout = true; diff --git a/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp b/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp index 64552c0e973c..b6ef2e13cbf5 100644 --- a/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp +++ b/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp @@ -90,6 +90,7 @@ struct MlaRopeGenArgs // `kv_only` launches the KV half, `kv_done_elsewhere` the Q half. bool kv_only; bool kv_done_elsewhere; + bool q_rope_applied; }; template @@ -133,6 +134,7 @@ void invokeMLARopeGenerationHelper(T const* latent_cache_ptr, T* q_pe_ptr, T* fu mla_params.host_bmm1_scale = args.host_bmm1_scale; mla_params.helix_position_offsets = args.helix_position_offsets_ptr; mla_params.helix_is_inactive_rank = args.helix_is_inactive_rank_ptr; + mla_params.q_rope_applied = args.q_rope_applied; mla_params.precomputed_cu_seqlens = args.precomputed_cu_seqlens; mla_params.precomputed_fmha_scheduler = args.precomputed_fmha_scheduler; @@ -175,7 +177,7 @@ void MLARopeGeneration(std::optional fused_q, // [tokens, num_hea int64_t qk_nope_head_dim, int64_t qk_rope_head_dim, int64_t v_head_dim, bool rope_append, std::optional kv_norm_weight, double const kv_norm_eps, bool const precomputed_cu_seqlens, bool const precomputed_fmha_scheduler, bool const kv_only, bool const kv_done_elsewhere, - std::optional quant_scale_qkv) + std::optional quant_scale_qkv, bool q_rope_applied) { // `kv_only` runs before q_pe exists, so the Q tensors are absent. TORCH_CHECK(kv_only || (fused_q.has_value() && q_pe.has_value()), @@ -186,8 +188,8 @@ void MLARopeGeneration(std::optional fused_q, // [tokens, num_hea TLLM_CHECK_WITH_INFO( head_size == kv_lora_rank + qk_rope_head_dim, "head_size must = kv_lora_rank + qk_rope_head_dim"); TLLM_CHECK_WITH_INFO(num_kv_heads == 1, "num_kv_heads must = 1"); - TLLM_CHECK_WITH_INFO(residual_dim == 0 || residual_dim == qk_rope_head_dim, - "MLA KV residual_dim must be 0 or qk_rope_head_dim (%ld), got %ld", qk_rope_head_dim, residual_dim); + TLLM_CHECK_WITH_INFO(residual_dim >= 0 && residual_dim <= qk_rope_head_dim && residual_dim % 16 == 0, + "MLA KV residual_dim must be a multiple of 16 in [0, qk_rope_head_dim], got %ld", residual_dim); TORCH_CHECK(helix_tensor_params.size() == 2, "Expecting 2 tensors for helix_tensor_params: helix_position_offsets and helix_is_inactive_rank."); @@ -330,7 +332,7 @@ void MLARopeGeneration(std::optional fused_q, // [tokens, num_hea quant_scale_o_ptr, kv_scale_orig_quant_ptr, kv_scale_quant_orig_ptr, kv_cache_scale_orig_quant_ptr, host_bmm1_scale, helix_position_offsets_ptr, helix_is_inactive_rank_ptr, kv_norm_weight_ptr, static_cast(kv_norm_eps), latent_row_stride, precomputed_cu_seqlens, precomputed_fmha_scheduler, kv_only, - kv_done_elsewhere}; + kv_done_elsewhere, q_rope_applied}; void* q_pe_ptr = kv_only ? nullptr : q_pe->data_ptr(); void* fused_q_ptr = kv_only ? nullptr : fused_q->data_ptr(); @@ -414,6 +416,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) ", bool kv_only=False" ", bool kv_done_elsewhere=False" ", Tensor? quant_scale_qkv=None" + ", bool q_rope_applied=False" ") -> ()"); } diff --git a/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp b/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp index 8f955399a60e..b0f2dc0dcd1a 100644 --- a/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp +++ b/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp @@ -15,6 +15,7 @@ */ #include "tensorrt_llm/common/attentionWorkspace.h" +#include "tensorrt_llm/common/attentionOp.h" #include "tensorrt_llm/common/workspace.h" @@ -56,8 +57,73 @@ void expectNextSlice(char const* name, Slice const& slice, size_t size, size_t& expectedOffset += tc::alignSize(size, kAlignment); } +constexpr int32_t kBatchSize = 2; +constexpr int32_t kInputSequenceLength = 11; +constexpr int32_t kCrossKvLength = 7; +constexpr int32_t kPackedTokenCount = 14; +constexpr int32_t kHeadSize = 8; + +void configureUnfusedAttention(tcop::AttentionOp& op, bool crossAttention) +{ + op.mNumHeads = 1; + op.mNumKVHeads = 1; + op.mHeadSize = kHeadSize; + op.mNumAttnHeads = 1; + op.mNumAttnKVHeads = 1; + op.mEnableContextFMHA = false; + op.mCrossAttention = crossAttention; +} + +size_t expectedUnfusedContextWorkspace(bool crossAttention) +{ + constexpr size_t kElementSize = sizeof(half); + size_t const batchSize = kBatchSize; + size_t const inputSequenceLength = kInputSequenceLength; + size_t const kvSequenceLength = crossAttention ? kCrossKvLength : kInputSequenceLength; + size_t const paddedTokenCount = batchSize * inputSequenceLength; + size_t const paddedKvTokenCount = batchSize * kvSequenceLength; + + tcop::AttentionContextWorkspaceSizes sizes{}; + sizes.attentionMask = kElementSize * paddedTokenCount * kvSequenceLength; + sizes.cuQSeqlens = sizeof(int) * (batchSize + 1); + sizes.cuKvSeqlens = sizes.cuQSeqlens; + sizes.cuMaskRows = sizes.cuQSeqlens; + sizes.qBuf = kElementSize * paddedTokenCount * kHeadSize; + sizes.kBuf = kElementSize * paddedKvTokenCount * kHeadSize; + sizes.vBuf = sizes.kBuf; + sizes.qkBuf = kElementSize * batchSize * inputSequenceLength * kvSequenceLength; + sizes.qkvBuf = kElementSize * paddedTokenCount * kHeadSize; + sizes.qkFloatBuf = sizeof(float) * batchSize * inputSequenceLength * kvSequenceLength; + sizes.paddingOffset = sizeof(int) * paddedTokenCount; + sizes.encoderPaddingOffset = sizeof(int) * paddedKvTokenCount; + sizes.tokensInfo = sizeof(int2) * kPackedTokenCount; + return tcop::AttentionWorkspaceManager::buildContextLayout(sizes).totalSize; +} + +size_t getUnfusedContextWorkspace(tcop::AttentionOp const& op) +{ + return op.getWorkspaceSizeForContext( + tensorrt_llm::DataType::kHALF, kBatchSize, kInputSequenceLength, kCrossKvLength, kPackedTokenCount); +} + } // namespace +TEST(AttentionWorkspaceManagerTest, RaggedUnfusedSelfAttentionUsesPaddedTokenCounts) +{ + tcop::AttentionOp op; + configureUnfusedAttention(op, false); + + EXPECT_EQ(getUnfusedContextWorkspace(op), expectedUnfusedContextWorkspace(false)); +} + +TEST(AttentionWorkspaceManagerTest, RaggedUnfusedCrossAttentionUsesPaddedTokenCounts) +{ + tcop::AttentionOp op; + configureUnfusedAttention(op, true); + + EXPECT_EQ(getUnfusedContextWorkspace(op), expectedUnfusedContextWorkspace(true)); +} + TEST(AttentionWorkspaceManagerTest, ContextLayoutMatchesAttentionOpOrdering) { tcop::AttentionContextWorkspaceSizes sizes{}; diff --git a/tensorrt_llm/_ipc_utils.py b/tensorrt_llm/_ipc_utils.py index 4d5ecefc9a9b..3178597b5087 100644 --- a/tensorrt_llm/_ipc_utils.py +++ b/tensorrt_llm/_ipc_utils.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import array +import ctypes import struct import sys from typing import List, Tuple @@ -26,6 +27,9 @@ from .logger import logger from .mapping import Mapping +# Must match CUDA_IPC_HANDLE_SIZE / CU_IPC_HANDLE_SIZE. +_IPC_MEM_HANDLE_SIZE = 64 + def _raise_if_error(error: cudart.cudaError_t | cuda.CUresult): if isinstance(error, cudart.cudaError_t): @@ -36,6 +40,33 @@ def _raise_if_error(error: cudart.cudaError_t | cuda.CUresult): raise RuntimeError(f"CUDA Driver API error: {repr(error)}") +def _ipc_mem_handle_to_bytes(handle) -> bytes: + """Serialize a cudaIpcMemHandle across cuda-python / cuda-bindings versions.""" + reserved = getattr(handle, "reserved", None) + if reserved is not None: + return bytes(reserved) + # Some cuda-bindings builds omit the .reserved attribute on cudaIpcMemHandle_t. + return ctypes.string_at(handle.getPtr(), _IPC_MEM_HANDLE_SIZE) + + +def _ipc_mem_handle_from_bytes(data: bytes): + """Rebuild a cudaIpcMemHandle from the opaque 64-byte payload.""" + handle = cudart.cudaIpcMemHandle_t() + if hasattr(handle, "reserved"): + try: + handle.reserved = data + except TypeError: + # Older bindings expect list[int] rather than bytes. + handle.reserved = list(data) + else: + if len(data) != _IPC_MEM_HANDLE_SIZE: + raise ValueError( + f"Invalid CUDA IPC mem handle size: {len(data)} (expected {_IPC_MEM_HANDLE_SIZE})" + ) + ctypes.memmove(handle.getPtr(), data, _IPC_MEM_HANDLE_SIZE) + return handle + + def can_access_peer(mapping: Mapping) -> bool: src_node = mapping.local_rank @@ -117,13 +148,8 @@ def align_size(size, alignment): _raise_if_error(cudart.cudaMemset(local_ptr, 0, aligned_size)[0]) error, local_handle = cudart.cudaIpcGetMemHandle(local_ptr) _raise_if_error(error) - handles_reserved = dist.tp_allgather(local_handle.reserved) - - handles = [] - for reserved in handles_reserved: - handle = cudart.cudaIpcMemHandle_t() - handle.reserved = reserved - handles.append(handle) + handles_reserved = dist.tp_allgather(_ipc_mem_handle_to_bytes(local_handle)) + handles = [_ipc_mem_handle_from_bytes(reserved) for reserved in handles_reserved] peer_ptrs = [] for node, handle in enumerate(handles): diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index c5dfe5585fbc..4b2769391270 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -142,13 +142,14 @@ statically checkable without runtime signature inspection. Ordinary sparse variants use `attention_output_hidden_size` and the shared output allocation. DeepSeek-V4's fused epilogue instead uses the optional -output-preparation hook to create one token-major O-LoRA output tensor. Its -context- and generation-phase helpers allocate the private FP8 attention and -scale buffers, then write the O-LoRA result into the corresponding token range. -The shared MLA custom-op contract exposes exactly one mutable output tensor; -`_create_outputs()` keeps that tensor in a single-entry list through forward -and output projection. Phase-specific scratch buffers remain inside the -DeepSeek-V4 algorithm module and do not widen the generic hook facade. +output-preparation hook to create one token-major O-LoRA-sized output tensor. +Its context- and generation-phase helpers allocate the private FP8 attention +and scale buffers, run both O-LoRA projections, and write the final hidden +states into the leading columns of that tensor. The shared MLA custom-op +contract exposes exactly one mutable output tensor; `_create_outputs()` keeps +that tensor in a single-entry list through forward and output projection. +Phase-specific scratch buffers remain inside the DeepSeek-V4 algorithm module +and do not widen the generic hook facade. Sparse prediction inputs stay out of shared MLA APIs. Algorithm modules wrap their module-to-backend inputs in a `SparseBackendForwardArgs` subclass and diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/indexer.py index 4ede4dfac0b6..995fc1c96ac8 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/indexer.py @@ -4,6 +4,7 @@ from __future__ import annotations +import os from typing import TYPE_CHECKING, Optional, Tuple import torch @@ -129,11 +130,10 @@ def _apply_q_rope(self, q: torch.Tensor, position_ids: torch.Tensor) -> torch.Te ) return q - def _project_and_quantize_q( - self, qr: torch.Tensor, position_ids: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Project and quantize Q, using the fused MXFP4 path when supported.""" - use_fused_project_mxfp4 = ( + def _is_fused_project_mxfp4_enabled(self, input_dtype: torch.dtype) -> bool: + if os.environ.get("TRTLLM_DISABLE_DSA_FUSED_INDEXER_Q", "0") == "1": + return False + return ( self.indexer_cache_dtype == KVCacheDtype.MXFP4_BLOCKWISE and not HAS_FAST_HADAMARD and not self.rotary_emb.is_neox @@ -145,9 +145,15 @@ def _project_and_quantize_q( torch.ops.trtllm, "cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell", ) - and qr.dtype == torch.bfloat16 + and input_dtype == torch.bfloat16 and is_sm_100f() ) + + def _project_and_quantize_q( + self, qr: torch.Tensor, position_ids: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Project and quantize Q, using the fused MXFP4 path when supported.""" + use_fused_project_mxfp4 = self._is_fused_project_mxfp4_enabled(qr.dtype) if use_fused_project_mxfp4: q_fp4, q_scale = torch.ops.trtllm.cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell( qr, @@ -226,6 +232,7 @@ def precompute_aux( self, hidden_states: torch.Tensor, metadata: DeepseekV4TrtllmAttentionMetadata, + start_event: Optional[torch.cuda.Event] = None, ) -> Optional[Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]]: """Pre-launch the qr-independent half of the indexer prepare phase. @@ -236,14 +243,17 @@ def precompute_aux( prepare path skip its own aux-stream launch and consume these results directly. - Returns ``None`` when multi-stream mode is off (caller should fall - back to the normal ``forward()`` call without ``pre_aux``). + Returns ``None`` when multi-stream mode is off or fused Indexer-Q + requires serial prepare (caller should fall back to the normal + ``forward()`` call without ``pre_aux``). """ if not (do_multi_stream() and self.aux_stream is not None): return None - self.indexer_start_event.record() + if start_event is None: + start_event = self.indexer_start_event + start_event.record() with torch.cuda.stream(self.aux_stream): - self.indexer_start_event.wait() + start_event.wait() weights = self.weights_proj(hidden_states) self.weights_proj_event.record() k_fp8, k_scale = self.compressor(hidden_states, metadata) @@ -308,16 +318,16 @@ def _run_overlapped_indexer_prepare( self.k_cache_update_event.record() else: weights, k_fp8, k_scale = pre_aux - # pre_aux tensors were allocated on aux_stream; record on the - # consuming stream so the caching allocator can't recycle them mid-use. - cur_stream = torch.cuda.current_stream() - weights.record_stream(cur_stream) - if k_fp8 is not None: - k_fp8.record_stream(cur_stream) - if k_scale is not None: - k_scale.record_stream(cur_stream) q_fp8, q_scale = self._project_and_quantize_q(qr, position_ids) + # Aux-stream tensors are consumed on the current stream after waits. + cur_stream = torch.cuda.current_stream() + weights.record_stream(cur_stream) + if k_fp8 is not None: + k_fp8.record_stream(cur_stream) + if k_scale is not None: + k_scale.record_stream(cur_stream) + self.weights_proj_event.wait() weights = self._apply_weight_scale(weights, q_scale) @@ -363,7 +373,12 @@ def forward( Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]] ] = None, ): - if do_multi_stream() and self.aux_stream is not None: + use_overlapped_prepare = ( + do_multi_stream() + and self.aux_stream is not None + and (pre_aux is not None or not self._is_fused_project_mxfp4_enabled(qr.dtype)) + ) + if use_overlapped_prepare: q_fp8, q_scale, k_fp8, k_scale, weights = self._run_overlapped_indexer_prepare( qr, hidden_states, @@ -372,7 +387,7 @@ def forward( pre_aux=pre_aux, ) else: - assert pre_aux is None, "pre_aux requires multi-stream mode" + assert pre_aux is None, "pre_aux requires the overlapped indexer prepare path" q_fp8, q_scale, k_fp8, k_scale, weights = self._run_serial_indexer_prepare( qr, hidden_states, metadata, position_ids ) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/params.py b/tensorrt_llm/_torch/attention/backends/sparse/params.py index 8e6829f8c879..e21af93568c0 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/params.py @@ -14,11 +14,34 @@ # limitations under the License. """Shared sparse attention parameter types.""" +import os from dataclasses import dataclass from typing import Literal, Optional import torch +_INDEXER_MQA_LOGITS_DEFAULT_ELEM_BUDGET = 1 << 31 +_INDEXER_MQA_LOGITS_BYTES_PER_ELEMENT = 4 + + +def get_indexer_mqa_logits_elem_budget() -> int: + """Return the per-call Indexer MQA-logits cap used by the runtime.""" + return int( + os.environ.get( + "TLLM_INDEXER_MQA_LOGITS_ELEM_BUDGET", _INDEXER_MQA_LOGITS_DEFAULT_ELEM_BUDGET + ) + ) + + +def get_indexer_mqa_logits_workspace_bytes( + max_num_tokens: Optional[int] = None, max_seq_len: Optional[int] = None +) -> int: + """Return the reachable maximum bytes for one FP32 MQA-logits tile.""" + elem_budget = get_indexer_mqa_logits_elem_budget() + if max_num_tokens is not None and max_seq_len is not None: + elem_budget = min(elem_budget, max_num_tokens * max_seq_len) + return elem_budget * _INDEXER_MQA_LOGITS_BYTES_PER_ELEMENT + class SparseParams: """Base parameters for a sparse attention backend.""" diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index d41ae7c5b066..a2b10d9dd80f 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1399,6 +1399,7 @@ def _( host_kv_cache_pool_mapping: Optional[torch.Tensor], kv_scale_orig_quant: Optional[torch.Tensor], kv_scale_quant_orig: Optional[torch.Tensor], + kv_cache_scale_orig_quant: Optional[torch.Tensor], out_scale: Optional[torch.Tensor], block_ids_per_seq: Optional[torch.Tensor], helix_tensor_params: List[Optional[torch.Tensor]], @@ -1407,6 +1408,7 @@ def _( num_heads: int, num_kv_heads: int, head_size: int, + residual_dim: int, tokens_per_block: int, attention_window_size: int, beam_width: int, @@ -1425,6 +1427,9 @@ def _( kv_only: bool = False, kv_done_elsewhere: bool = False, quant_scale_qkv: Optional[torch.Tensor] = None, + # Declared by the schema in dsv3RopeOp.cpp; meta dispatch passes it + # positionally, so the fake has to accept it. + q_rope_applied: bool = False, ) -> None: # This is a fake implementation for shape inference # The actual operation modifies fused_q and q_pe in-place 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 0e02a8759ab4..00127e89a966 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -42,6 +42,68 @@ # Torch schema parsing rejects ``inf`` as a default value. SWIGLU_LIMIT_SCALAR_DISABLED = -1.0 +# The torch.library schema needs a concrete float, so "unset" is a sentinel +# rather than ``None``. SiTU betas are required to be positive, so any +# non-positive value is unambiguously "not provided". +SITU_BETA_DISABLED = -1.0 + + +def _canonicalize_situ_beta(situ_beta: float) -> Optional[float]: + return None if situ_beta <= 0 else float(situ_beta) + + +_KIMI_K3_MXFP8_TUNING_BUCKETS = (1, 2, 4, 8, *range(16, 193, 16)) + +# Both dense MXFP8 wrappers pass a scalar one to the CuTe kernel. Keep one +# read-only tensor per CUDA device so steady-state calls, including CUDA graph +# replay, do not launch a scalar fill kernel. +_MXFP8_GEMM_ALPHA_CACHE: "dict[torch.device, torch.Tensor]" = {} + + +def _get_mxfp8_gemm_alpha(device: torch.device) -> torch.Tensor: + """Return the cached FP32 scalar one for ``device``. + + Allocation is deliberately rejected during CUDA graph capture. Callers + must run one eager warmup on each device, which is already required for + GEMM autotuning, before capturing the steady-state path. + """ + device = torch.device(device) + if device.type != "cuda": + raise ValueError( + f"MXFP8 GEMM alpha requires a CUDA device, got {device}.") + if device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + + alpha = _MXFP8_GEMM_ALPHA_CACHE.get(device) + if alpha is None: + with torch.cuda.device(device): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "MXFP8 GEMM alpha cache must be initialized before CUDA graph " + f"capture on {device}; run one eager GEMM warmup first.") + alpha = torch.ones((), dtype=torch.float32, device=device) + _MXFP8_GEMM_ALPHA_CACHE[device] = alpha + return alpha + + +def _get_kimi_k3_mxfp8_tuning_buckets(max_num_tokens: int) -> Tuple[int, ...]: + """Generate every K3 bucket used through ``max_num_tokens``.""" + max_bucket = _kimi_k3_mxfp8_tuning_bucket(max_num_tokens) + low_m = [m for m in _KIMI_K3_MXFP8_TUNING_BUCKETS if m <= max_bucket] + high_m = [ + m for m in get_last_power_of_2_num_tokens_buckets(max_bucket) if m > 192 + ] + return tuple((*low_m, *high_m)) + + +def _kimi_k3_mxfp8_tuning_bucket(num_tokens: int) -> int: + """Use lower power-of-two buckets for small M, then upper-bound buckets.""" + if num_tokens <= 32: + return last_positive_power_of_2(num_tokens) + if num_tokens <= 192: + return next(m for m in _KIMI_K3_MXFP8_TUNING_BUCKETS if num_tokens <= m) + return next_positive_power_of_2(num_tokens) + def _with_input_cuda_device(function): """Run a custom-op implementation under its input tensor's CUDA device.""" @@ -3516,17 +3578,24 @@ def __init__(self, scaling_vector_size: int = 16, activation_type: ActivationType = ActivationType.Swiglu, swiglu_limit_scalar: float = float("inf"), - use_expert_counts: bool = False): + use_expert_counts: bool = False, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None): """Initialize the runner. Args: activation_type: ``ActivationType`` for the fused epilogue. Only - ``Swiglu`` (gated) and ``Relu2`` (non-gated) are supported. + ``Swiglu`` (gated), ``Relu2`` (non-gated) and ``SiTu`` + (gated) are supported. swiglu_limit_scalar: Uniform clamp limit for SwiGLU. ``+inf`` disables clamp. + situ_beta: Gate-side SiTU constant; required for ``SiTu`` only. + situ_linear_beta: Linear-side SiTU constant; required for ``SiTu`` only. """ super().__init__() self.activation_type = validate_activation_type(activation_type) self.is_gated = is_gated_activation(self.activation_type) + self.situ_beta = situ_beta + self.situ_linear_beta = situ_linear_beta self.num_experts = num_experts self.top_k = top_k self.num_local_experts = num_local_experts @@ -3564,6 +3633,8 @@ def unique_id(self): self.activation_type, self.swiglu_limit_scalar, self.use_expert_counts, + self.situ_beta, + self.situ_linear_beta, ) def get_valid_tactics( @@ -3823,10 +3894,13 @@ def forward(self, inputs: List, assert mma_tiler_mn[ 0] == self.tile_size, f"Tactic ({tactic}) is incompatible with tile size ({self.tile_size})" + # The SiTU betas are folded into the kernel at trace time, so they + # are part of the compiled-kernel identity, not just runtime args. cache_key = (self.scaling_vector_size, self.tile_size, self.top_k, mma_tiler_mn, cluster_shape_mn, raster_along_m, self.activation_type, self.swiglu_limit_scalar, - self.use_expert_counts, self.num_local_experts) + self.use_expert_counts, self.num_local_experts, + self.situ_beta, self.situ_linear_beta) if cache_key not in self.__class__.kernel_cache: gemm = self.__class__.kernel_class( @@ -3840,6 +3914,8 @@ def forward(self, inputs: List, swiglu_limit=self.swiglu_limit_scalar, use_expert_counts=self.use_expert_counts, num_local_experts=self.num_local_experts, + situ_beta=self.situ_beta, + situ_linear_beta=self.situ_linear_beta, ) hardware_info = cutlass.utils.HardwareInfo() max_active_clusters = hardware_info.get_max_active_clusters( @@ -3929,12 +4005,17 @@ def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell( swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED, expert_counts: Optional[torch.Tensor] = None, expert_capacity: int = 0, + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, ) -> Tuple[torch.Tensor, torch.Tensor]: """CuteDSL-based NVFP4 gather grouped GEMM with activation fusion. - Supports ``ActivationType.Swiglu`` (gated) and ``ActivationType.Relu2`` - (non-gated) epilogues; other ``ActivationType`` values raise an - assertion in the runner. + Supports ``ActivationType.Swiglu`` (gated), ``ActivationType.Relu2`` + (non-gated) and ``ActivationType.SiTu`` (gated) epilogues; other + ``ActivationType`` values raise an assertion in the runner. + + ``situ_beta``/``situ_linear_beta`` are only meaningful for + ``ActivationType.SiTu``; values <= 0 disable them. """ tuner = AutoTuner.get() swiglu_limit_scalar = _canonicalize_swiglu_limit_scalar( @@ -3960,7 +4041,9 @@ def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell( scaling_vector_size, activation_type=ActivationType(activation_type), swiglu_limit_scalar=swiglu_limit_scalar, - use_expert_counts=expert_counts is not None) + use_expert_counts=expert_counts is not None, + situ_beta=_canonicalize_situ_beta(situ_beta), + situ_linear_beta=_canonicalize_situ_beta(situ_linear_beta)) inputs = [ input, weight, input_scale, weight_scale, alpha, tile_idx_to_group_idx, tile_idx_to_mn_limit, @@ -3999,6 +4082,8 @@ def _fake_single_b( swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED, expert_counts: Optional[torch.Tensor] = None, expert_capacity: int = 0, + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, ) -> Tuple[torch.Tensor, torch.Tensor]: if expert_counts is not None: helper = GroupedGemmInputsHelper(num_experts, top_k, @@ -13291,7 +13376,7 @@ def cute_dsl_mxfp8_gemm_rubin( if input_scale.dtype != torch.uint8 or weight_scale.dtype != torch.uint8: raise ValueError("CuteDSL MXFP8 scales must be UE8M0 uint8") - alpha = torch.ones((), dtype=torch.float32, device=input.device) + alpha = _get_mxfp8_gemm_alpha(input.device) runner = CuteDSLMXFP8RubinLinear(output_dtype=output_dtype, use_tvm_ffi=use_tvm_ffi) inputs = [input, weight, input_scale, weight_scale, alpha] @@ -14083,6 +14168,8 @@ def _( from ..cute_dsl_kernels.rubin.moe.rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion import \ Sm107BlockScaledContiguousGatherGroupedGemmActFusionKernel + from ..cute_dsl_kernels.rubin.moe.rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion import \ + validate_activation_type as validate_rubin_activation_type class Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner( TunableRunner): @@ -14090,7 +14177,7 @@ class Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner( SM107 counterpart to Blackwell's ``Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner``. - Supports SwiGLU and Relu2. + Supports SwiGLU, SiTU, and Relu2. Key differences from Blackwell: - Uses LDGSTS (cp.async) for A/SFA loading instead of TMA - Supports B-reuse pattern (mma_tiler_m = 2 * mma_inst_shape_m) @@ -14108,7 +14195,9 @@ def __init__( local_expert_offset: int, tile_size: int, scaling_vector_size: int = 16, - activation_type: ActivationType = ActivationType.Swiglu): + 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 @@ -14116,12 +14205,10 @@ def __init__( self.local_expert_offset = local_expert_offset self.tile_size = tile_size self.scaling_vector_size = scaling_vector_size - self.activation_type = ActivationType(int(activation_type)) - if self.activation_type not in (ActivationType.Swiglu, - ActivationType.Relu2): - raise ValueError( - f"Rubin NVFP4 CuteDSL FC1 does not support " - f"{self.activation_type.name}") + self.activation_type = validate_rubin_activation_type( + activation_type) + self.situ_beta = situ_beta + self.situ_linear_beta = situ_linear_beta self.is_gated = is_gated_activation(self.activation_type) if (sm_version := get_sm_version()) != 107: @@ -14143,6 +14230,8 @@ def unique_id(self): self.tile_size, self.scaling_vector_size, int(self.activation_type), + self.situ_beta, + self.situ_linear_beta, ) def get_valid_tactics( @@ -14445,7 +14534,8 @@ def forward(self, inputs: List[torch.Tensor], self.top_k, mma_tiler, mma_inst_shape, cluster_shape_mn, raster_along_m, locality_domain_half_gemm, a_path, - int(self.activation_type), max_active_clusters) + int(self.activation_type), self.situ_beta, + self.situ_linear_beta, max_active_clusters) if cache_key not in self.__class__.kernel_cache: gemm = self.__class__.kernel_class( sf_vec_size=self.scaling_vector_size, @@ -14458,6 +14548,8 @@ def forward(self, inputs: List[torch.Tensor], locality_domain_half_gemm=locality_domain_half_gemm, a_path=a_path, activation_type=self.activation_type, + situ_beta=self.situ_beta, + situ_linear_beta=self.situ_linear_beta, ) compiled_gemm = cute.compile( gemm.wrapper, @@ -14534,6 +14626,8 @@ def _run_nvfp4_gather_grouped_gemm_act_fusion_rubin( scaling_vector_size: int, partition_id: int, activation_type: ActivationType, + situ_beta: float, + situ_linear_beta: float, precomputed_tactic: Optional[str], tuner_key: str, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: @@ -14560,6 +14654,8 @@ def _run_nvfp4_gather_grouped_gemm_act_fusion_rubin( tile_size, scaling_vector_size, activation_type=activation_type, + situ_beta=_canonicalize_situ_beta(situ_beta), + situ_linear_beta=_canonicalize_situ_beta(situ_linear_beta), ) inputs = [ input, weight, input_scale, weight_scale, alpha, @@ -14600,6 +14696,7 @@ def _run_nvfp4_gather_grouped_gemm_act_fusion_rubin( "Tensor(a16!)? output_tensor, Tensor(a17!)? output_sf_tensor, " "SymInt scaling_vector_size=16, SymInt partition_id=-1, " f"SymInt activation_type={int(ActivationType.Swiglu)}, " + "float situ_beta=-1.0, float situ_linear_beta=-1.0, " "str? precomputed_tactic=None) -> (Tensor?, Tensor?)", device_types="cuda") def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( @@ -14623,6 +14720,8 @@ def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( scaling_vector_size: int = 16, partition_id: int = -1, activation_type: int = int(ActivationType.Swiglu), + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, precomputed_tactic: Optional[str] = None, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: return _run_nvfp4_gather_grouped_gemm_act_fusion_rubin( @@ -14632,7 +14731,8 @@ def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( num_experts, top_k, num_local_experts, local_expert_offset, tile_size, output_tensor, output_sf_tensor, scaling_vector_size, partition_id, - ActivationType(activation_type), precomputed_tactic, + ActivationType(activation_type), situ_beta, situ_linear_beta, + precomputed_tactic, "trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin") @torch.library.register_fake( @@ -14658,6 +14758,8 @@ def _( scaling_vector_size: int = 16, partition_id: int = -1, activation_type: int = int(ActivationType.Swiglu), + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, precomputed_tactic: Optional[str] = None, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: m = permuted_idx_to_expanded_idx.size(0) @@ -14689,30 +14791,33 @@ def _( "SymInt num_local_experts, SymInt local_expert_offset, " "SymInt tile_size, Tensor(a!) output_tensor, " "Tensor(b!) output_sf_tensor, SymInt scaling_vector_size=16, " - f"SymInt activation_type={int(ActivationType.Swiglu)}) -> ()", + 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_gather_grouped_gemm_act_fusion_locality_domain_inplace_rubin( - input: torch.Tensor, - weight_0: torch.Tensor, - weight_1: torch.Tensor, - input_scale: torch.Tensor, - weight_scale_0: torch.Tensor, - weight_scale_1: torch.Tensor, - alpha: torch.Tensor, - tile_idx_to_group_idx: torch.Tensor, - tile_idx_to_mn_limit: torch.Tensor, - permuted_idx_to_expanded_idx: torch.Tensor, - num_non_exiting_tiles: torch.Tensor, - global_sf: torch.Tensor, - num_experts: int, - top_k: int, - num_local_experts: int, - local_expert_offset: int, - tile_size: int, - output_tensor: torch.Tensor, - output_sf_tensor: torch.Tensor, - scaling_vector_size: int = 16, - activation_type: int = int(ActivationType.Swiglu), + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + input_scale: torch.Tensor, + weight_scale_0: torch.Tensor, + weight_scale_1: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_tensor: torch.Tensor, + output_sf_tensor: torch.Tensor, + scaling_vector_size: int = 16, + activation_type: int = int(ActivationType.Swiglu), + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, ) -> None: """Tune and launch both Rubin locality domain NVFP4 MoE FC1 partitions. @@ -14754,6 +14859,8 @@ def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_locality_domain_inplace_rubin( tile_size, scaling_vector_size, activation_type=ActivationType(activation_type), + situ_beta=_canonicalize_situ_beta(situ_beta), + situ_linear_beta=_canonicalize_situ_beta(situ_linear_beta), )) inputs = [ input, @@ -14799,6 +14906,8 @@ def launch_partition( scaling_vector_size=scaling_vector_size, partition_id=partition_id, activation_type=activation_type, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, precomputed_tactic=repr(tactic), ) @@ -14817,29 +14926,886 @@ def launch_partition( "trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_locality_domain_inplace_rubin" ) def _( + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + input_scale: torch.Tensor, + weight_scale_0: torch.Tensor, + weight_scale_1: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_tensor: torch.Tensor, + output_sf_tensor: torch.Tensor, + scaling_vector_size: int = 16, + activation_type: int = int(ActivationType.Swiglu), + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, + ) -> None: + return None + + # ---------------------------------------------------------------- + # Rubin NVFP4 Fused FC12 (FC1 gather+gated-act + FC2 finalize in ONE kernel) + # ---------------------------------------------------------------- + # Compat shim: the delivered fused kernel accesses ``cutlass.memory.*`` + # (SmemAllocator / TmemAllocator / get_smem_capacity_in_bytes) and + # ``cutlass.tensor_utils.LayoutEnum`` as attributes of the top-level + # ``cutlass`` module. In the pinned nvidia-cutlass-dsl-internal + # (0.3.0+...c907734) these are real submodules (e.g. the dspark kernel + # does ``from cutlass import memory``) that are simply not auto-exposed + # until explicitly imported. Import them so the attribute access + # resolves; fall back to aliasing ``cutlass.utils`` (which carries the + # same symbols in this build) if a submodule is genuinely absent. This + # keeps the vendored kernel byte-identical instead of editing its API + # references. + import cutlass.utils as _cutlass_utils_compat + try: + import cutlass.memory # noqa: F401 real submodule in pinned build + except ImportError: + if not hasattr(cutlass, "memory"): + cutlass.memory = _cutlass_utils_compat + try: + import cutlass.tensor_utils # noqa: F401 + except ImportError: + if not hasattr(cutlass, "tensor_utils"): + cutlass.tensor_utils = _cutlass_utils_compat + + # The fused FC12 kernel module ships with the CuteDSL MoE backend + # change. Keep the op registration optional so that importing this + # module still succeeds on an internal CuTe DSL build that does not + # yet carry the kernel; the MoE backend checks for the op before use. + try: + from ..cute_dsl_kernels.rubin.moe.rubin_contiguous_grouped_blockscaled_gemm_fused_fc12 import \ + Sm107BlockScaledContiguousGroupedGemmFusedFc12Kernel + except ImportError: + Sm107BlockScaledContiguousGroupedGemmFusedFc12Kernel = None + logger.debug( + "Fused FC12 CuteDSL kernel is unavailable; " + "trtllm::cute_dsl_nvfp4_fc12_fused_rubin is not registered.") + + if Sm107BlockScaledContiguousGroupedGemmFusedFc12Kernel is not None: + + class Fc12FusedInputsHelper(GatherGroupedGemmInputsHelper): + """Autotuning helper for the fused FC12 op. + + Inputs 0..9 keep the FC1 gather layout so the parent moe_sort + regeneration is reused verbatim; inputs 10..14 carry the FC2 + tensors, of which fc2_c and fc2_routing_scales depend on the + token count and must be resized to match the regenerated tiles + during profiling. + """ + IDX_FC2_B = 10 + IDX_FC2_SFB = 11 + IDX_FC2_ALPHA = 12 + IDX_FC2_C = 13 + IDX_FC2_ROUTING = 14 + # expanded_idx_to_permuted_idx (dim0 = num_tokens); consumed only by + # the in-op output memset. Appended after the FC2 tensors so the + # existing 0..14 positions (and their constraints) stay put. + IDX_EXPANDED_IDX = 15 + + def inputs_pre_hook( + self, inputs: List[torch.Tensor]) -> List[torch.Tensor]: + # Reuse the parent's moe_sort regeneration on the FC1 prefix + # (feed two None placeholders so it skips the + # locality-domain output resize). + base = list(super().inputs_pre_hook( + list(inputs[:10]) + [None, None])) + fc1_prefix = base[:10] + permuted = fc1_prefix[self.IDX_PERMUTED_IDX_TO_EXPANDED_IDX] + num_tokens = self.infer_num_tokens(permuted.size(0)) + fc2_b = inputs[self.IDX_FC2_B] + fc2_sfb = inputs[self.IDX_FC2_SFB] + fc2_alpha = inputs[self.IDX_FC2_ALPHA] + fc2_c = inputs[self.IDX_FC2_C] + fc2_routing = inputs[self.IDX_FC2_ROUTING] + new_fc2_c = fc2_c.new_empty((num_tokens, fc2_c.size(1))) + new_routing = fc2_routing.new_empty( + (num_tokens, fc2_routing.size(1))) + # Resize expanded_idx to the regenerated token count so the + # unpacked forward sees a consistent 16-tensor list. Its + # contents are unused during tuning (the in-op memset falls back + # to an index-free full zero while AutoTuner.is_tuning_mode, so + # the uninitialised index values are never dereferenced). + expanded_idx = inputs[self.IDX_EXPANDED_IDX] + new_expanded = expanded_idx.new_empty( + (num_tokens, expanded_idx.size(1))) + return (*fc1_prefix, fc2_b, fc2_sfb, fc2_alpha, new_fc2_c, + new_routing, new_expanded) + + class Sm107BlockScaledContiguousGroupedGemmFusedFc12Runner( + TunableRunner): + """Rubin runner for the fused FC1+FC2 (FC12) NVFP4 MoE kernel. + + 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 + on every launch. v1 exposes the GEN-phase geometry only: + mma_tiler 128x{128,256}, cluster (1,1), scheduler="l2_atomic". + 2CTA/CTX geometries are added once CTX perf tuning lands. + """ + kernel_class = Sm107BlockScaledContiguousGroupedGemmFusedFc12Kernel + 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, + 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 + self.num_local_experts = num_local_experts + self.local_expert_offset = local_expert_offset + self.tile_size = tile_size + self.scaling_vector_size = scaling_vector_size + self.swiglu_limit = swiglu_limit + self.activation_type = ActivationType(int(activation_type)) + # The fused kernel derives interm_size as fc1_n // 2, so a + # non-gated activation would silently halve the wrong + # dimension and produce wrong output rather than failing. + if self.activation_type not in (ActivationType.Swiglu, + ActivationType.SiTu): + raise ValueError( + f"{self.__class__.kernel_class.__name__} supports " + f"gated activations (Swiglu, SiTu) only, but got " + f"{self.activation_type.name}") + 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 + self.enable_alltoall = enable_alltoall + if (sm_version := get_sm_version()) != 107: + raise ValueError( + f"{self.__class__.kernel_class.__name__} supports SM 107 " + f"(Rubin) only, but got SM {sm_version}") + # v1 fused kernel only supports the 128-wide routing tile. + if self.tile_size not in (128, 256): + raise ValueError( + f"{self.__class__.kernel_class.__name__} supports " + f"tile_size 128 (1-CTA) or 256 (2-CTA) only, but got " + f"{self.tile_size}") + + def unique_id(self): + return ( + self.num_experts, + self.top_k, + self.num_local_experts, + self.local_expert_offset, + 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( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + **kwargs, + ) -> List[Tuple]: + (fc1_a, fc1_b, fc1_sfa, fc1_sfb, fc1_alpha, + tile_idx_to_group_idx, tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, num_non_exiting_tiles, + fc1_norm_const, fc2_b, *_) = inputs + m = permuted_idx_to_expanded_idx.size(0) + k = fc1_a.size(1) * 2 + l, fc1_n = fc1_b.size(0), fc1_b.size(1) # noqa: E741 + fc2_n, fc2_k = fc2_b.size(1), fc2_b.size(2) * 2 + + # Fixed K for FP4: mma_tiler_k=256, mma_inst_k=128. + mma_tiler_k = 256 + mma_inst_k = 128 + # Mirror the CuteDSL grouped-GEMM runners: the MMA M-tile equals + # the routing tile_size and the cluster M = tile_size // 128, so + # tile_size=128 -> 1-CTA cluster (1,1), tile_size=256 -> 2-CTA + # cluster (2,1). moe_sort tiles the tokens by tile_size, so the + # kernel's M-tile always matches the routing tile. mma_n is free + # {128, 256}. + mma_tiler_m = self.tile_size + mma_inst_m = self.tile_size + cluster_shape_mn = (self.tile_size // 128, 1) + mma_n_candidates = [128, 256] + + valid_tactics = [] + for mma_n in mma_n_candidates: + # No "cluster CTAs > tiles" guard here (unlike the Sm100 + # runners): this kernel's per-CTA M-tile is 128, so one + # 256-row logical tile is exactly one (2,1) cluster and a + # single-tile problem is valid for 2-CTA. + # The fused N-tile must divide both FC1 and FC2 output N. + if fc1_n % mma_n != 0 or fc2_n % mma_n != 0: + continue + + mma_tiler = (mma_tiler_m, mma_n, mma_tiler_k) + mma_inst_shape = (mma_inst_m, mma_n, mma_inst_k) + + if self.__class__.kernel_class.can_implement( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + sf_dtype=cutlass.Float8E4M3FN, + sf_vec_size=self.scaling_vector_size, + fc1_c_dtype=cutlass.Float4E2M1FN, + fc2_c_dtype=cutlass.BFloat16, + mma_inst_shape=mma_inst_shape, + mma_tiler=mma_tiler, + cluster_shape_mn=cluster_shape_mn, + fc1_gemm_shape=(m, fc1_n, k, l), + fc2_gemm_shape=(m, fc2_n, fc2_k, l), + a_major="k", + b_major="k", + fc1_c_major="n", + fc2_c_major="n", + ): + valid_tactics.append( + (mma_tiler, mma_inst_shape, cluster_shape_mn)) + + logger.debug( + f"CuteDSL Rubin FusedFC12: Found {len(valid_tactics)} valid " + f"tactics for M={m}, FC1_N={fc1_n}, FC2_N={fc2_n}, K={k}, " + f"L={l}") + return valid_tactics + + def get_tuning_config(self) -> TuningConfig: + key = self.unique_id() + if key not in self.__class__.tuning_config_cache: + helper = Fc12FusedInputsHelper(self.num_experts, + self.top_k, + self.num_local_experts, + self.local_expert_offset, + self.tile_size) + self.__class__.tuning_config_cache[key] = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + GatherGroupedGemmInputsHelper.IDX_SHAPE_INFER, + 0, helper.gen_tuning_buckets, + helper.map_to_tuning_buckets), ), + constraint_specs=( + ConstraintSpec(0, 0, + helper.infer_shape_num_tokens), + ConstraintSpec(2, 0, + helper.infer_shape_num_tokens), + ConstraintSpec( + 5, 0, helper.infer_shape_max_num_tiles), + ConstraintSpec( + 6, 0, helper.infer_shape_max_num_tiles), + # fc2_c/fc2_routing_scales dim0 = num_tokens (follows M bucket like inputs 0/2) + ConstraintSpec(13, 0, + helper.infer_shape_num_tokens), + ConstraintSpec(14, 0, + helper.infer_shape_num_tokens), + # expanded_idx_to_permuted_idx dim0 = num_tokens too + ConstraintSpec(15, 0, + helper.infer_shape_num_tokens), + ), + inputs_pre_hook=helper.inputs_pre_hook, + ) + return self.__class__.tuning_config_cache[key] + + def forward(self, inputs: List[torch.Tensor], + tactic: Optional[tuple], **kwargs) -> torch.Tensor: + (fc1_a, fc1_b, fc1_sfa, fc1_sfb, fc1_alpha, + tile_idx_to_group_idx, tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, num_non_exiting_tiles, + fc1_norm_const, fc2_b, fc2_sfb, fc2_alpha, fc2_c, + fc2_routing_scales, expanded_idx_to_permuted_idx) = inputs + + assert fc1_a.dtype == torch.float4_e2m1fn_x2 + assert fc1_b.dtype == torch.float4_e2m1fn_x2 + assert fc2_b.dtype == torch.float4_e2m1fn_x2 + assert fc1_sfa.dtype == torch.uint8 + assert fc1_sfb.dtype == torch.uint8 + assert fc2_sfb.dtype == torch.uint8 + assert fc1_alpha.dtype == torch.float32 + assert fc2_alpha.dtype == torch.float32 + assert fc1_norm_const.dtype == torch.float32 + assert fc2_c.dtype == torch.bfloat16 + + sf_vec = self.scaling_vector_size + orig_m, k = fc1_a.size(0), fc1_a.size(1) * 2 + m = permuted_idx_to_expanded_idx.size(0) + l, fc1_n = fc1_b.size(0), fc1_b.size(1) # noqa: E741 + interm_size = fc1_n // 2 # SwiGLU (gated) halves N + k // sf_vec + fc2_n, fc2_k = fc2_b.size(1), fc2_b.size(2) * 2 + fc2_k // sf_vec + num_tokens = fc2_c.size(0) + num_tiles = m // self.tile_size + assert m % self.tile_size == 0 + assert fc2_k == interm_size, ( + f"FC2 K ({fc2_k}) must equal FC1 intermediate ({interm_size})" + ) + + if isinstance(tactic, tuple): + mma_tiler, mma_inst_shape, cluster_shape_mn = tactic + else: + # Fallback geometry must satisfy the kernel validator + # (is_valid_mma_tiler_and_cluster_shape): inst_m == + # tile_m and cluster_m == inst_m // 128. Mirrors + # _get_sm107_nvfp4_default_mma_config used by the sibling + # runners; the old fixed cluster (1,1) was illegal for + # tile_size=256 (a 2-CTA MMA in a 1-CTA cluster). + mma_inst_m = min(self.tile_size, 256) + mma_tiler = (self.tile_size, 128, 256) + mma_inst_shape = (mma_inst_m, 128, 128) + cluster_shape_mn = (mma_inst_m // 128, 1) + # The MMA M-tile and cluster M must match the routing tile this + # runner was built for; a mismatch (e.g. a tactic captured under + # another tile_size) would index routing metadata wrongly and + # can read uninitialised permuted-index padding. + assert ( + mma_tiler[0] == self.tile_size + and cluster_shape_mn[0] == self.tile_size // 128), ( + f"FC12 tactic/tile mismatch: mma_tiler={mma_tiler} " + f"mma_inst_shape={mma_inst_shape} " + f"cluster_shape_mn={cluster_shape_mn} " + f"tile_size={self.tile_size}") + + # FC1 intermediate output (kept on-chip by the kernel; passed as + # a scratch tensor) + its dynamic block scale. + fc1_c = torch.empty(m, + interm_size // 2, + dtype=fc1_a.dtype, + device=fc1_a.device) + 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) + + # Zero the scatter-add output right before the fused kernel so + # the memset becomes the kernel's immediate stream predecessor + # (PDL prologue can then overlap a longer predecessor than the + # 1-int counter). The zeroing must run on every launch: autotune() + # keeps is_tuning_mode set for its whole context, including the + # final real launch, so gating the memset on it would leave real + # outputs unzeroed. While tuning, the profiling inputs carry an + # uninitialised expanded_idx (see Fc12FusedInputsHelper + # .inputs_pre_hook), so use an index-free full zero there instead + # of the sparse, index-driven memset; it is semantically equivalent + # (a superset of the rows the finalize scatter-adds into) and still + # charges a memset to the profiled time. + if AutoTuner.get().is_tuning_mode: + fc2_c.zero_() + else: + torch.ops.trtllm.moe_output_memset_inplace( + input=fc2_c, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx= + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx= + permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + tile_tokens_dim=self.tile_size, + top_k=self.top_k, + ep_size=self.ep_size, + enable_alltoall=self.enable_alltoall, + ) + + fc1_a_ptr = make_ptr(cutlass.Float4E2M1FN, + fc1_a.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32) + fc1_b_ptr = make_ptr(cutlass.Float4E2M1FN, + fc1_b.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32) + fc1_c_ptr = make_ptr(cutlass.Float4E2M1FN, + fc1_c.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32) + fc1_sfa_ptr = make_ptr(cutlass.Float8E4M3FN, + fc1_sfa.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + fc1_sfb_ptr = make_ptr(cutlass.Float8E4M3FN, + fc1_sfb.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + fc1_sfc_ptr = make_ptr(cutlass.Float8E4M3FN, + fc1_sfc.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + fc1_norm_const_ptr = make_ptr(cutlass.Float32, + fc1_norm_const.data_ptr(), + cute.AddressSpace.gmem) + fc1_alpha_ptr = make_ptr(cutlass.Float32, + fc1_alpha.data_ptr(), + cute.AddressSpace.gmem) + fc2_alpha_ptr = make_ptr(cutlass.Float32, + fc2_alpha.data_ptr(), + cute.AddressSpace.gmem) + tile_idx_to_group_idx_ptr = make_ptr( + cutlass.Int32, tile_idx_to_group_idx.data_ptr(), + cute.AddressSpace.gmem) + tile_idx_to_mn_limit_ptr = make_ptr( + cutlass.Int32, tile_idx_to_mn_limit.data_ptr(), + cute.AddressSpace.gmem) + permuted_idx_to_expanded_idx_ptr = make_ptr( + cutlass.Int32, permuted_idx_to_expanded_idx.data_ptr(), + cute.AddressSpace.gmem) + num_non_exiting_tiles_ptr = make_ptr( + cutlass.Int32, num_non_exiting_tiles.data_ptr(), + cute.AddressSpace.gmem) + fc1_ready_ptr = make_ptr(cutlass.Int32, + fc1_ready.data_ptr(), + cute.AddressSpace.gmem) + fc1_scheduler_counter_ptr = make_ptr( + cutlass.Int32, fc1_scheduler_counter.data_ptr(), + cute.AddressSpace.gmem) + fc2_scheduler_counter_ptr = make_ptr( + cutlass.Int32, fc2_scheduler_counter.data_ptr(), + cute.AddressSpace.gmem) + fc2_b_ptr = make_ptr(cutlass.Float4E2M1FN, + fc2_b.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32) + fc2_sfb_ptr = make_ptr(cutlass.Float8E4M3FN, + fc2_sfb.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + fc2_c_ptr = make_ptr(cutlass.BFloat16, + fc2_c.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + fc2_routing_scales_ptr = make_ptr( + cutlass.Float32, fc2_routing_scales.data_ptr(), + cute.AddressSpace.gmem) + + torch_stream = torch.cuda.current_stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + max_active_clusters = get_max_activate_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1]) + + # The fused kernel exposes only ``__call__`` (cute.Tensor args), + # no ptr-based wrapper. ``make_ordered_layout`` needs an MLIR + # context, so build the cute tensors INSIDE a @cute.jit wrapper + # (like the standalone FC1/FC2 kernels' ``wrapper`` methods do + # internally) that captures ``gemm`` and takes ptrs + Int64 + # shapes. Shapes are dynamic so one compile serves all shapes. + sf_vec_cx = self.scaling_vector_size + tile_size_cx = self.tile_size + top_k_cx = self.top_k + 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, 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, + mma_inst_shape, + mma_tiler, + cluster_shape_mn, + True, # vectorized_f32 + topk=self.top_k, + 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 + def _fc12_wrapper( + a_ptr, + b_ptr, + c_ptr, + sfa_ptr, + sfb_ptr, + sfc_ptr, + norm_const_ptr, + tile_grp_ptr, + tile_mn_ptr, + num_non_exiting_ptr, + fc1_alpha_ptr, + fc2_alpha_ptr, + ready_ptr, + fc1_sched_ptr, + fc2_sched_ptr, + w2_ptr, + out_ptr, + w2_sf_ptr, + permuted_ptr, + routing_ptr, + orig_m: cutlass.Int64, + m: cutlass.Int64, + fc1_n: cutlass.Int64, + k: cutlass.Int64, + l: cutlass.Int64, # noqa: E741 + fc2_n: cutlass.Int64, + fc2_k: cutlass.Int64, + num_tokens: cutlass.Int64, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + interm_size = fc1_n // 2 # SwiGLU (gated) halves N + scale_k = k // sf_vec_cx + fc2_scale_k = fc2_k // sf_vec_cx + num_tiles = m // tile_size_cx + a = cute.make_tensor( + a_ptr, + layout=cute.make_ordered_layout( + (orig_m, k, 1), order=(1, 0, 2))) + b = cute.make_tensor( + b_ptr, + layout=cute.make_ordered_layout( + (fc1_n, k, l), order=(1, 0, 2))) + c = cute.make_tensor(c_ptr, + layout=cute.make_layout( + (m, interm_size, 1), + stride=(interm_size, 1, + m * interm_size))) + sfa = cute.make_tensor( + sfa_ptr, + layout=cute.make_ordered_layout( + (orig_m, scale_k, 1), order=(1, 0, 2))) + sfb = cute.make_tensor( + sfb_ptr, + layout=cute.make_ordered_layout( + (32, 4, fc1_n // 128, 4, scale_k // 4, l), + order=(2, 1, 4, 0, 3, 5))) + sfc = cute.make_tensor( + sfc_ptr, + layout=cute.make_ordered_layout( + (32, 4, m // 128, 4, interm_size // + (sf_vec_cx * 4), l), + order=(2, 1, 4, 0, 3, 5))) + norm_const = cute.make_tensor( + norm_const_ptr, layout=cute.make_layout((1, ))) + fc1_alpha = cute.make_tensor( + fc1_alpha_ptr, layout=cute.make_layout((l, ))) + fc2_alpha = cute.make_tensor( + fc2_alpha_ptr, layout=cute.make_layout((l, ))) + tile_grp = cute.make_tensor(tile_grp_ptr, + layout=cute.make_layout( + (num_tiles, ))) + tile_mn = cute.make_tensor(tile_mn_ptr, + layout=cute.make_layout( + (num_tiles, ))) + num_non_exiting = cute.make_tensor( + num_non_exiting_ptr, + layout=cute.make_layout((1, ))) + ready = cute.make_tensor(ready_ptr, + layout=cute.make_layout( + (num_tiles, ))) + fc1_sched = cute.make_tensor( + fc1_sched_ptr, layout=cute.make_layout((1, ))) + fc2_sched = cute.make_tensor( + fc2_sched_ptr, layout=cute.make_layout((1, ))) + w2 = cute.make_tensor( + w2_ptr, + layout=cute.make_ordered_layout( + (fc2_n, fc2_k, l), order=(1, 0, 2))) + w2_sf = cute.make_tensor( + w2_sf_ptr, + layout=cute.make_ordered_layout( + (32, 4, fc2_n // 128, 4, fc2_scale_k // 4, + l), + order=(2, 1, 4, 0, 3, 5))) + out = cute.make_tensor( + out_ptr, + layout=cute.make_layout( + (num_tokens, fc2_n, 1), + stride=(fc2_n, 1, num_tokens * fc2_n))) + permuted = cute.make_tensor(permuted_ptr, + layout=cute.make_layout( + (m, ))) + routing = cute.make_tensor( + routing_ptr, + layout=cute.make_ordered_layout( + (num_tokens, top_k_cx), order=(1, 0))) + gemm(a, + b, + c, + sfa, + sfb, + sfc, + norm_const, + tile_grp, + tile_mn, + num_non_exiting, + fc1_alpha, + fc2_alpha, + ready, + fc1_sched, + fc2_sched, + w2, + out, + w2_sf, + permuted, + routing, + max_active_clusters=max_active_clusters, + stream=stream) + + compiled_gemm = cute.compile( + _fc12_wrapper, + fc1_a_ptr, + fc1_b_ptr, + fc1_c_ptr, + fc1_sfa_ptr, + fc1_sfb_ptr, + fc1_sfc_ptr, + fc1_norm_const_ptr, + tile_idx_to_group_idx_ptr, + tile_idx_to_mn_limit_ptr, + num_non_exiting_tiles_ptr, + fc1_alpha_ptr, + fc2_alpha_ptr, + fc1_ready_ptr, + fc1_scheduler_counter_ptr, + fc2_scheduler_counter_ptr, + fc2_b_ptr, + fc2_c_ptr, + fc2_sfb_ptr, + permuted_idx_to_expanded_idx_ptr, + fc2_routing_scales_ptr, + orig_m, + m, + fc1_n, + k, + l, + fc2_n, + fc2_k, + num_tokens, + max_active_clusters=max_active_clusters, + stream=stream, + ) + self.__class__.kernel_cache[cache_key] = compiled_gemm + else: + compiled_gemm = self.__class__.kernel_cache[cache_key] + + compiled_gemm( + fc1_a_ptr, + fc1_b_ptr, + fc1_c_ptr, + fc1_sfa_ptr, + fc1_sfb_ptr, + fc1_sfc_ptr, + fc1_norm_const_ptr, + tile_idx_to_group_idx_ptr, + tile_idx_to_mn_limit_ptr, + num_non_exiting_tiles_ptr, + fc1_alpha_ptr, + fc2_alpha_ptr, + fc1_ready_ptr, + fc1_scheduler_counter_ptr, + fc2_scheduler_counter_ptr, + fc2_b_ptr, + fc2_c_ptr, + fc2_sfb_ptr, + permuted_idx_to_expanded_idx_ptr, + fc2_routing_scales_ptr, + orig_m, + m, + fc1_n, + k, + l, + fc2_n, + fc2_k, + num_tokens, + stream=stream, + ) + return fc2_c + + def _run_nvfp4_fc12_fused_rubin( + input: torch.Tensor, + fc1_weight: torch.Tensor, + input_scale: torch.Tensor, + fc1_weight_scale: torch.Tensor, + fc1_alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + fc2_weight: torch.Tensor, + fc2_weight_scale: torch.Tensor, + fc2_alpha: torch.Tensor, + output: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + scaling_vector_size: int, + swiglu_limit: float, + precomputed_tactic: Optional[str], + expanded_idx_to_permuted_idx: torch.Tensor, + 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( + num_experts, + top_k, + num_local_experts, + local_expert_offset, + tile_size, + scaling_vector_size, + 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 + # feeds the in-op output memset). + inputs = [ + input, fc1_weight, input_scale, fc1_weight_scale, fc1_alpha, + tile_idx_to_group_idx, tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, num_non_exiting_tiles, + global_sf, fc2_weight, fc2_weight_scale, fc2_alpha, output, + token_final_scales, expanded_idx_to_permuted_idx + ] + if precomputed_tactic is None: + _, best_tactic = tuner.choose_one( + tuner_key, + [runner], + runner.get_tuning_config(), + inputs, + ) + else: + best_tactic = ast.literal_eval(precomputed_tactic) + return runner(inputs, tactic=best_tactic) + + @torch.library.custom_op( + "trtllm::cute_dsl_nvfp4_fc12_fused_rubin", + mutates_args=("output", ), + schema= + "(Tensor input, Tensor fc1_weight, Tensor input_scale, Tensor fc1_weight_scale, " + "Tensor fc1_alpha, Tensor tile_idx_to_group_idx, Tensor tile_idx_to_mn_limit, " + "Tensor permuted_idx_to_expanded_idx, Tensor num_non_exiting_tiles, Tensor global_sf, " + "Tensor fc2_weight, Tensor fc2_weight_scale, Tensor fc2_alpha, " + "Tensor(a13!) output, Tensor token_final_scales, " + "Tensor expanded_idx_to_permuted_idx, " + "SymInt num_experts, SymInt top_k, SymInt num_local_experts, " + "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, " + 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, - weight_0: torch.Tensor, - weight_1: torch.Tensor, + fc1_weight: torch.Tensor, input_scale: torch.Tensor, - weight_scale_0: torch.Tensor, - weight_scale_1: torch.Tensor, - alpha: torch.Tensor, + fc1_weight_scale: torch.Tensor, + fc1_alpha: torch.Tensor, tile_idx_to_group_idx: torch.Tensor, tile_idx_to_mn_limit: torch.Tensor, permuted_idx_to_expanded_idx: torch.Tensor, num_non_exiting_tiles: torch.Tensor, global_sf: torch.Tensor, + fc2_weight: torch.Tensor, + fc2_weight_scale: torch.Tensor, + fc2_alpha: torch.Tensor, + output: torch.Tensor, + token_final_scales: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, num_experts: int, top_k: int, num_local_experts: int, local_expert_offset: int, tile_size: int, - output_tensor: torch.Tensor, - output_sf_tensor: torch.Tensor, + swiglu_limit: float, + ep_size: int, + enable_alltoall: bool, scaling_vector_size: int = 16, + precomputed_tactic: Optional[str] = None, activation_type: int = int(ActivationType.Swiglu), - ) -> None: - return None + 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. + _run_nvfp4_fc12_fused_rubin( + input, fc1_weight, input_scale, fc1_weight_scale, fc1_alpha, + tile_idx_to_group_idx, tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, num_non_exiting_tiles, + global_sf, fc2_weight, fc2_weight_scale, fc2_alpha, output, + token_final_scales, num_experts, top_k, num_local_experts, + 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", + ActivationType(activation_type), situ_beta, + situ_linear_beta) + + @torch.library.register_fake( + "trtllm::cute_dsl_nvfp4_fc12_fused_rubin") + def _( + input: torch.Tensor, + fc1_weight: torch.Tensor, + input_scale: torch.Tensor, + fc1_weight_scale: torch.Tensor, + fc1_alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + fc2_weight: torch.Tensor, + fc2_weight_scale: torch.Tensor, + fc2_alpha: torch.Tensor, + output: torch.Tensor, + token_final_scales: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + swiglu_limit: float, + ep_size: int, + 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 # ---------------------------------------------------------------- # Rubin BF16/FP16 Gather + SwiGLU Fusion (FC1 layer) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py index b31ba370a78c..171cd4ceb615 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py @@ -901,13 +901,17 @@ def _chunk_kda_fwd( # Varlen with non-aligned seq lengths is a separate problem (Phase 2): # multi-seq varlen can't be host-padded without repacking memory. real_T = T - needs_eqlen_pad = (not is_varlen) and (T % BT != 0) + CPB_BT = 4 * BT # CHUNKS_PER_BLOCK * BT, the cgs_per_head divisibility unit + # The trigger is CPB_BT, not BT: the eqlen scheduler floors `cgs_per_head` + # and has no `chunk_idx < num_chunks` guard, so a BT-aligned T whose chunk + # count is not a multiple of CHUNKS_PER_BLOCK (T = 64, 128, 192, 320, ...) + # would silently drop its trailing chunks. + needs_eqlen_pad = (not is_varlen) and (T % CPB_BT != 0) if needs_eqlen_pad: if B != 1: raise NotImplementedError( - f"eqlen with B>1 and T % {BT} != 0 not supported (got B={B}, T={T})." + f"eqlen with B>1 and T % {CPB_BT} != 0 not supported (got B={B}, T={T})." ) - CPB_BT = 4 * BT # CHUNKS_PER_BLOCK * BT, the cgs_per_head divisibility unit T_padded = ((T + CPB_BT - 1) // CPB_BT) * CPB_BT # Pre-allocated padded scratch buffers (per (B,T_padded,H,K,dtype) cache # key). torch.cat would reallocate + copy the full 200MB q tensor every @@ -1001,16 +1005,11 @@ def _chunk_kda_fwd( if chunk_indices is None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = len(chunk_indices) - if NT < 4: - # The persistent K123 scheduler launches NT // 4 cooperative - # groups per head; fewer than 4 total chunks produces a - # zero-size grid (DSLCudaRuntimeError at launch). Callers must - # route such batches to the FLA path — see - # KDAKernelDispatch.prefill_chunk_kda. - raise ValueError( - f"kda_prefill requires >= 4 total varlen chunks (got {NT}); " - "route small varlen batches to the FLA fallback" - ) + # Any chunk count is launchable: the varlen scheduler rounds up + # (`total_cgs_per_head = ceil(num_chunks / CHUNKS_PER_BLOCK)`), launches + # the constant `NUM_SMS` grid, and guards every chunk loop with + # `chunk_idx < num_chunks`. Only eqlen needs divisibility, and it gets + # that from its own `CPB_BT` pad above. N_seqs = len(cu_seqlens) - 1 else: NT = T // BT diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 52d66018ebb6..cf7fc63ef8b3 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -2192,6 +2192,104 @@ def forward( return output +class Fp8PrequantizedSwapABGemmRunner(TunableRunner): + """Runs DeepGemm with pre-quantized FP8 activations and packed scales.""" + + tuning_config = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, deep_gemm_gen_tuning_buckets), ), + constraint_specs=(ConstraintSpec( + 1, 0, lambda input_shapes: input_shapes[0][0]), ), + exclude_from_cache=True, + ) + + def __init__(self, output_dtype: torch.dtype, + disable_ue8m0_cast: bool) -> None: + self.output_dtype = output_dtype + self.disable_ue8m0_cast = disable_ue8m0_cast + + def unique_id(self): + return ( + self.output_dtype, + self.disable_ue8m0_cast, + ) + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + ) -> List[int]: + return [0] + + def forward( + self, + inputs: List[torch.Tensor], + tactic: int = -1, + ) -> torch.Tensor: + del tactic + activation, activation_scale, weight, weight_scale = inputs + scale_m_aligned = fp4_utils.pad_up(activation_scale.size(0), 4) + if activation_scale.stride() != (1, scale_m_aligned): + # Dynamic autotuning recreates constrained integer tensors with a + # contiguous layout. Restore the MN-major packed-scale stride that + # the real quantizers return and DeepGemm requires. + normalized_scale = torch.empty_strided( + activation_scale.shape, (1, scale_m_aligned), + dtype=activation_scale.dtype, + device=activation_scale.device) + normalized_scale.copy_(activation_scale) + activation_scale = normalized_scale + output = torch.empty( + (activation.size(0), weight.size(0)), + device=activation.device, + dtype=self.output_dtype, + ) + deep_gemm.fp8_gemm_nt( + (activation, activation_scale), + (weight, weight_scale), + output, + disable_ue8m0_cast=self.disable_ue8m0_cast, + ) + return output + + +@torch.library.custom_op("trtllm::fp8_prequantized_swap_ab_gemm", + mutates_args=()) +def fp8_prequantized_swap_ab_gemm( + activation: torch.Tensor, + activation_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype = torch.bfloat16, + disable_ue8m0_cast: bool = False, +) -> torch.Tensor: + runner = Fp8PrequantizedSwapABGemmRunner(output_dtype, disable_ue8m0_cast) + _, best_tactic = AutoTuner.get().choose_one( + "trtllm::fp8_prequantized_swap_ab_gemm", + [runner], + Fp8PrequantizedSwapABGemmRunner.tuning_config, + [activation, activation_scale, weight, weight_scale], + ) + return runner( + inputs=[activation, activation_scale, weight, weight_scale], + tactic=best_tactic, + ) + + +@fp8_prequantized_swap_ab_gemm.register_fake +def _( + activation: torch.Tensor, + activation_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype = torch.bfloat16, + disable_ue8m0_cast: bool = False, +) -> torch.Tensor: + del activation_scale, weight_scale, disable_ue8m0_cast + return activation.new_empty((activation.size(0), weight.size(0)), + dtype=output_dtype) + + @torch.library.custom_op("trtllm::fp8_swap_ab_gemm", mutates_args=()) def fp8_swap_ab_gemm( input: torch.Tensor, 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 7ed708faaacd..d53dcd7c3ce3 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 @@ -2786,7 +2786,14 @@ def softmax( row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, row_max_new, 0) - elif cutlass.const_expr(arch >= Arch.sm_103 and arch <= Arch.sm_103f): + # SM107 (Rubin) shares this reduction path with SM103. The arm was + # dropped when this file was taken wholesale from main during the + # rebase (c130d75c42); ``rubin-advance`` gates the identical body on + # sm_101/sm_103/sm_107/sm_110. Only sm_107 is restored here -- the + # others were never in main's copy and are not Rubin. + elif cutlass.const_expr( + (arch >= Arch.sm_103 and arch <= Arch.sm_103f) + or (arch >= Arch.sm_107 and arch <= Arch.sm_107f)): tmem_load_red_atom = cute.make_copy_atom( tcgen05.copy.LdRed32x32bOp(tcgen05.copy.Repetition(64), redOp=tcgen05.TmemLoadRedOp.MAX), 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 adfd5cc78726..9e2477edf634 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 @@ -2759,7 +2759,14 @@ def softmax( # reduction for row_max row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, row_max_new, 0) - elif cutlass.const_expr(arch >= Arch.sm_103 and arch <= Arch.sm_103f): + # SM107 (Rubin) shares this reduction path with SM103. The arm was + # dropped when this file was taken wholesale from main during the + # rebase (c130d75c42); ``rubin-advance`` gates the identical body on + # sm_101/sm_103/sm_107/sm_110. Only sm_107 is restored here -- the + # others were never in main's copy and are not Rubin. + elif cutlass.const_expr( + (arch >= Arch.sm_103 and arch <= Arch.sm_103f) + or (arch >= Arch.sm_107 and arch <= Arch.sm_107f)): tmem_load_red_atom = cute.make_copy_atom( tcgen05.copy.LdRed32x32bOp(tcgen05.copy.Repetition(64), redOp=tcgen05.TmemLoadRedOp.MAX), diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py index 6dc6becd131a..d63eba9b6e9e 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py @@ -89,17 +89,27 @@ def __extract_mlir_values__(self): return values def __new_from_mlir_values__(self, values): - problem_shape_b = cutlass.new_from_mlir_values(self.problem_shape_b, - (values[0], )) - problem_shape_s = cutlass.new_from_mlir_values(self.problem_shape_s, - (values[1], )) - split_kv = cutlass.new_from_mlir_values(self.split_kv, (values[2], )) - problem_shape_b_fdd = cutlass.new_from_mlir_values( - self.problem_shape_b_fdd, (values[3], )) - problem_shape_s_fdd = cutlass.new_from_mlir_values( - self.problem_shape_s_fdd, (values[4], )) - split_kv_fdd = cutlass.new_from_mlir_values(self.split_kv_fdd, - (values[5], )) + # Slice per field by that field's own extraction width, in + # __extract_mlir_values__ order. FastDivmodDivisor occupies 2 SSA values + # on nvidia-cutlass-dsl-internal, so fixed indices mis-slice (cutlass #3243). + fields = ( + self.problem_shape_b, + self.problem_shape_s, + self.split_kv, + self.problem_shape_b_fdd, + self.problem_shape_s_fdd, + self.split_kv_fdd, + ) + rebuilt = [] + offset = 0 + for field in fields: + width = len(cutlass.extract_mlir_values(field)) + rebuilt.append( + cutlass.new_from_mlir_values( + field, tuple(values[offset:offset + width]))) + offset += width + (problem_shape_b, problem_shape_s, split_kv, problem_shape_b_fdd, + problem_shape_s_fdd, split_kv_fdd) = rebuilt return MLAStaticTileSchedulerParams( self.is_persistent, problem_shape_b, @@ -281,16 +291,21 @@ def __extract_mlir_values__(self): return values def __new_from_mlir_values__(self, values): - assert len(values) == 13 - new_params = cutlass.new_from_mlir_values(self.params, values[0:6]) - new_current_work_linear_idx = cutlass.new_from_mlir_values( - self.current_work_linear_idx, [values[6]]) - new_blk_coord = cutlass.new_from_mlir_values(self.blk_coord, - values[7:10]) - new_grid_shape = cutlass.new_from_mlir_values(self.grid_shape, - values[10:]) - return MLAStaticTileScheduler(new_params, new_current_work_linear_idx, - new_blk_coord, new_grid_shape) + # Width-driven slicing, as in MLAStaticTileSchedulerParams: params holds + # three FastDivmodDivisors, which are not 1 SSA value each on every + # cutlass-dsl build, so the old fixed 6/1/3/rest split mis-slices. + rebuilt = [] + offset = 0 + for field in (self.params, self.current_work_linear_idx, self.blk_coord, + self.grid_shape): + width = len(cutlass.extract_mlir_values(field)) + rebuilt.append( + cutlass.new_from_mlir_values( + field, list(values[offset:offset + width]))) + offset += width + assert offset == len(values), ( + f"MLAStaticTileScheduler consumed {offset} of {len(values)} values") + return MLAStaticTileScheduler(*rebuilt) def create_mla_static_tile_scheduler( diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py index 69984b7edc3c..2852baa5d745 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py @@ -47,10 +47,15 @@ griddepcontrol_launch_dependents, griddepcontrol_wait, is_power_of_2, + sigmoid_f32, silu_f32, ) -SUPPORTED_ACTIVATION_TYPES = (ActivationType.Swiglu, ActivationType.Relu2) +SUPPORTED_ACTIVATION_TYPES = ( + ActivationType.Swiglu, + ActivationType.Relu2, + ActivationType.SiTu, +) def validate_activation_type(activation_type) -> ActivationType: @@ -72,6 +77,8 @@ def validate_activation_type(activation_type) -> ActivationType: Supported fused activations (selected at construction via ``activation_type``): - ActivationType.Swiglu: C = up * silu(gate), where up/gate come from interleaved weight matrix B - ActivationType.Relu2: C = relu(alpha * x)^2 + - ActivationType.SiTu: C = (beta*tanh(gate/beta)*sigmoid(gate)) * (linear_beta*tanh(up/linear_beta)), + gated like Swiglu; requires situ_beta / situ_linear_beta Any other ``ActivationType`` value raises an assertion at construction time. @@ -278,13 +285,16 @@ def __init__( swiglu_limit: cutlass.Float32 = float("inf"), use_expert_counts: bool = False, num_local_experts: int = 0, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, ): """Initializes the configuration for a Blackwell blockscaled dense GEMM kernel with gather operation and fused activation. ``activation_type`` accepts a value from ``ActivationType``; only - ``ActivationType.Swiglu`` (gated path) and ``ActivationType.Relu2`` - (non-gated path) are currently supported. + ``ActivationType.Swiglu`` (gated path), ``ActivationType.Relu2`` + (non-gated path) and ``ActivationType.SiTu`` (gated path, Kimi K3) + are currently supported. This configuration includes several key aspects: @@ -318,8 +328,15 @@ def __init__( :param topk: Number of experts selected per token (used for token ID mapping). :type topk: cutlass.Int64 :param activation_type: Fused activation. Must be ``ActivationType.Swiglu`` - (gated, default) or ``ActivationType.Relu2`` (non-gated). + (gated, default), ``ActivationType.Relu2`` (non-gated) or + ``ActivationType.SiTu`` (gated). :type activation_type: ActivationType + :param situ_beta: Gate-side SiTU constant. Required for -- and only + valid with -- ``ActivationType.SiTu``. + :type situ_beta: Optional[float] + :param situ_linear_beta: Linear-side (up) SiTU constant. Required for -- + and only valid with -- ``ActivationType.SiTu``. + :type situ_linear_beta: Optional[float] """ self.sf_vec_size = sf_vec_size @@ -415,6 +432,35 @@ def __init__( assert self.topk == 1 assert self.num_local_experts > 0 + # SiTU constants. They are per-model scalars (not per-expert), so they + # are folded at trace time -- which also means they belong in the + # caller's compiled-kernel cache key. + 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( + "SiTU beta parameters must be positive, got " + f"{situ_beta} and {situ_linear_beta}." + ) + if self.has_swiglu_limit: + # Matches MegaMoE (both backends) and DeepGEMM, which reject + # activation_clamp together with SiTU. + raise ValueError( + "ActivationType.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) + def _setup_attributes(self): """Set up configurations that are dependent on GEMM inputs @@ -2482,6 +2528,9 @@ def kernel( if cutlass.const_expr(self.activation_type == ActivationType.Swiglu): acc_vec_gate = tTR_rAcc_gate.load() self._apply_swiglu_epilogue(acc_vec_up, acc_vec_gate, alpha_val, tCompute) + elif cutlass.const_expr(self.activation_type == ActivationType.SiTu): + acc_vec_gate = tTR_rAcc_gate.load() + self._apply_situ_epilogue(acc_vec_up, acc_vec_gate, alpha_val, tCompute) elif cutlass.const_expr(self.activation_type == ActivationType.Relu2): self._apply_relu2_epilogue(acc_vec_up, alpha_val, tCompute) @@ -2770,6 +2819,93 @@ def _apply_swiglu_epilogue( 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 ``kimi_k3_moe/_mlp.py::SituAndMul`` + (itself byte-identical to HF ``modeling_kimi.py``):: + + 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 + + ``up`` and ``gate`` come from the two interleaved accumulator subtiles + loaded by the caller, same as the SwiGLU epilogue. + + There is no packed tanh, so the vectorized path uses the identity + ``tanh(z) = 2 * sigmoid(2z) - 1`` (the same one ``utils.gelu_tanh_f32`` + uses) to stay on the packed f32x2 path -- calling a scalar tanh would + force the whole loop back to scalar. The reciprocals and ``2*beta`` + factors fold at trace time because both betas are ``const_expr``:: + + beta * tanh(x/beta) = beta * (2*sigmoid(2x/beta) - 1) + = 2*beta*sigmoid((2/beta)*x) - beta + """ + 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) + + # sigmoid(x) = rcp(1 + exp2(-x * log2e)), shared by both cores. + 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 + @cute.jit def _apply_relu2_epilogue( self, @@ -3452,8 +3588,8 @@ def wrapper( """Single-B wrapper. ``l`` is the number of experts in the (sole) B tensor. ``activation_type`` - must match the one passed to ``__init__``; only ``Swiglu`` and ``Relu2`` - are supported. + must match the one passed to ``__init__``; only ``Swiglu``, ``Relu2`` + and ``SiTu`` are supported. """ is_gated = is_gated_activation(activation_type) scale_k = k // scaling_vector_size diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py index 795934f776f1..d506a8314e69 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py @@ -44,6 +44,7 @@ # This file is copied and modified from cutlass https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/cute/core.py import ctypes +import math import os from typing import Union @@ -99,6 +100,12 @@ def __init__( assert int(self._pointer) % self._assumed_align == 0, ( f"pointer must be {self._assumed_align} bytes aligned") + def __add__(self, offset: int) -> Pointer: # type: ignore[override] + offset_bytes = offset * self._dtype.width // 8 + assumed_align = math.gcd(offset_bytes, self._assumed_align) + return _Pointer(self._pointer + offset_bytes, self._dtype, + self._addr_space, assumed_align) + def size_in_bytes(self) -> int: return ctypes.sizeof(ctypes.c_void_p(int(self._pointer))) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dsv4_qb_fusion/kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dsv4_qb_fusion/kernel.py index 1421e192eef1..9cbf8b8ff367 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dsv4_qb_fusion/kernel.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dsv4_qb_fusion/kernel.py @@ -90,7 +90,6 @@ from __future__ import annotations import os -from functools import partial from typing import NamedTuple, Optional, Tuple import cuda.bindings.driver as cuda_driver @@ -101,12 +100,16 @@ import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.utils.blockscaled_layout as blockscaled_utils import cutlass.utils.rubin_helpers as sm107_utils +from cutlass._mlir import ir as mlir_ir +from cutlass._mlir.dialects import llvm as mlir_llvm from cutlass._mlir.dialects import nvvm -from cutlass.cute.arch import nvvm_wrappers as _arch_nvvm_wrappers +from cutlass._mlir.dialects import vector as mlir_vector +from cutlass._mlir.dialects.nvvm import FPRoundingMode from cutlass.cute.experimental import iket from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 from cutlass.cute.nvgpu.common import CacheEvictionPriority from cutlass.cute.nvgpu.tcgen05.mma import CollectorOp +from cutlass.cutlass_dsl import dsl_user_op from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.utils.gemm import sm100 as gemm_sm100 @@ -124,22 +127,126 @@ HEAD_DIM = DEFAULT_QK_NOPE_HEAD_DIM + DEFAULT_QK_ROPE_HEAD_DIM -def _fma_packed_bf16x2_nvvm(res, src_a, src_b, src_c, *, rnd=None, ftz=None, loc=None, ip=None): - del ftz - src_a_bf16x2 = _arch_nvvm_wrappers.cvt_f32x2_bf16x2(src_a, loc=loc, ip=ip) - return _arch_nvvm_wrappers.nvvm.fma_packed_f32x2_bf16x2_f32x2_f32x2( - res, src_a_bf16x2, src_b, src_c, rnd=rnd, loc=loc, ip=ip +@dsl_user_op +def _prmt_trunc_bf16x2( + v0: cutlass.Float32, + v1: cutlass.Float32, + *, + loc=None, + ip=None, +): + """Truncate two fp32 to bf16 and pack into one bf16x2 register. + + A single PRMT selecting the high two bytes of each source + (selector 0x7632: result = [v0.b2, v0.b3, v1.b2, v1.b3]) -- the + same numerics as masking the low 16 mantissa bits (truncate + instead of rn), one instruction per pair, and no F2FP + conversion-pipe traffic. + + NOTE: builds the vector<2xbf16> via cutlass._mlir plumbing + (llvm.bitcast) -- there is no public wrapper accepting a + pre-packed bf16x2 operand; revisit when one appears.""" + packed_i32 = cute.arch.prmt( + v0.bitcast(cutlass.Int32, loc=loc, ip=ip), + v1.bitcast(cutlass.Int32, loc=loc, ip=ip), + cutlass.Int32(0x7632), + loc=loc, + ip=ip, + ) + return mlir_llvm.bitcast( + mlir_ir.VectorType.get([2], cutlass.BFloat16.mlir_type, loc=loc), + packed_i32, + loc=loc, + ip=ip, ) -_fma_packed_f32x2_bf16x2_f32x2_f32x2 = getattr( - cute.arch, - "fma_packed_f32x2_bf16x2_f32x2_f32x2", - partial( - _arch_nvvm_wrappers.calc_packed_f32x2_op, - calc_func=_fma_packed_bf16x2_nvvm, - ), -) +@dsl_user_op +def _prmt_trunc_bf16x2_neg_lo( + v0: cutlass.Float32, + v1: cutlass.Float32, + *, + loc=None, + ip=None, +): + """Like _prmt_trunc_bf16x2 but with the LOW bf16 half negated: + returns packed (-bf16(v0), bf16(v1)). + + Used for the rope rotation's mixed-sign multiply: (x1, x0) * + (-sin, +sin) == (-x1, x0) * (sin, sin), which replaces the f32 + sin negation (an FADD that LLVM hoists away from the FMA, + defeating ptxas's .NP-modifier folding) with one sign-bit LOP3 on + the packed operand and lets the b operand be a scalar broadcast.""" + packed_i32 = cute.arch.prmt( + v0.bitcast(cutlass.Int32, loc=loc, ip=ip), + v1.bitcast(cutlass.Int32, loc=loc, ip=ip), + cutlass.Int32(0x7632), + loc=loc, + ip=ip, + ) + negated = cute.arch.lop3( + cutlass.Int32(packed_i32), + cutlass.Int32(0x00008000), + cutlass.Int32(0), + 0x3C, # a XOR b + loc=loc, + ip=ip, + ) + return mlir_llvm.bitcast( + mlir_ir.VectorType.get([2], cutlass.BFloat16.mlir_type, loc=loc), + negated.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _fma_packed_f32x2_bf16x2_f32x2_f32x2( + vec_a, + src_b, + src_c, + *, + loc=None, + ip=None, +): + """fma_packed_f32x2_bf16x2_f32x2_f32x2 with a pre-packed bf16x2 + src_a (vector<2xbf16> ir.Value), skipping the wrapper's per-element + f32->bf16 conversion (the F2FP).""" + vec_res_type = mlir_ir.VectorType.get([2], cutlass.Float32.mlir_type, loc=loc) + vec_b = mlir_vector.from_elements( + vec_res_type, + tuple(v.ir_value(loc=loc, ip=ip) for v in src_b), + loc=loc, + ip=ip, + ) + vec_c = mlir_vector.from_elements( + vec_res_type, + tuple(v.ir_value(loc=loc, ip=ip) for v in src_c), + loc=loc, + ip=ip, + ) + rn = next(mode for mode in FPRoundingMode if str(mode) == "rn") + vec_res = nvvm.fma_packed_f32x2_bf16x2_f32x2_f32x2( + vec_res_type, + vec_a, + vec_b, + vec_c, + rnd=rn, + loc=loc, + ip=ip, + ) + return tuple( + cutlass.Float32( + mlir_vector.extract( + vec_res, + dynamic_position=[], + static_position=[i], + loc=loc, + ip=ip, + ) + ) + for i in range(2) + ) class S2TCopyBundle(NamedTuple): @@ -175,6 +282,7 @@ def __init__( rms_norm_eps: float = DEFAULT_RMS_NORM_EPS, max_batch: int = 128, tma_prefetch_dist: int = 0, + enable_pdl: bool = True, ): if int(max_batch) < 1: raise ValueError("max_batch must be >= 1") @@ -182,6 +290,7 @@ def __init__( if int(tma_prefetch_dist) < 0: raise ValueError("tma_prefetch_dist must be >= 0") self.tma_prefetch_dist = int(tma_prefetch_dist) + self.enable_pdl = bool(enable_pdl) if tuple(mma_inst_tile) not in ((128, 256), (256, 256)): raise ValueError("mma_inst_tile must be (128, 256) or (256, 256)") mma_m = mma_inst_tile[0] @@ -269,11 +378,12 @@ def __init__( self.f32_bytes = cutlass.Float32.width // 8 self.i32_bytes = cutlass.Int32.width // 8 self.rope_cache_row_floats = 2 * self.qk_rope_head_dim - self.rope_smem_row_words = self.qk_rope_head_dim - self.rope_cp_bytes = 16 - self.rope_cp_chunks = (self.rope_smem_row_words * self.f32_bytes) // self.rope_cp_bytes - self.rope_rows_per_cp = 32 // self.rope_cp_chunks - self.rope_cp_iters = 32 // self.rope_rows_per_cp + self.rope_smem_row_words = self.qk_rope_head_dim // 2 + self.rope_row_quads = self.qk_rope_head_dim // 4 + self.rope_half_quads = self.rope_row_quads // 2 + self.rope_ldg_row_threads = self.rope_half_quads + self.rope_ldg_rows_per_iter = 32 // self.rope_ldg_row_threads + self.rope_ldg_iters = 32 // self.rope_ldg_rows_per_iter def _check_cluster(self, shape_mn, name: str) -> None: cm, cn = shape_mn @@ -904,8 +1014,6 @@ def kernel_body( smem = utils.SmemAllocator() storage = smem.allocate(self.shared_storage) - # CUTLASS DSL 4.5 expresses multicast signaling through the consumer - # thread count; this is API adaptation only. num_mcast_ctas_a = cute.size(cluster_layout_vmnk.shape[2]) num_mcast_ctas_b = cute.size(cluster_layout_vmnk.shape[1]) num_tma_producers = num_mcast_ctas_a + num_mcast_ctas_b - 1 @@ -1071,372 +1179,191 @@ def kernel_body( work_tile = tile_sched.initial_work_tile_info() iket.range_pop() - if warp_idx == self.tma_warp_id: - cute.arch.setmaxregister_decrease(self.mainloop_reg_count) - ab_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.num_ab_stage + s_cu_seqlens = storage.sCuSeqlens.get_tensor(cute.make_layout(self.max_batch + 1)) + s_kv_lengths = storage.sKvLengths.get_tensor(cute.make_layout(self.max_batch)) + s_rope_t = storage.sRope.get_tensor( + self.rope_smem_layout.outer, + swizzle=self.rope_smem_layout.inner, + ) + if cutlass.const_expr(self.use_tma_store): + sOut = storage.sOut.get_tensor( + self.out_smem_layout_staged.outer, + swizzle=self.out_smem_layout_staged.inner, + ) + + if warp_idx < self.mma_warp_id: + cute.arch.setmaxregister_increase(self.epilog_reg_count) + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + for meta_i in cutlass.range(0, (batch_size + 128) // 128, 1, unroll=1): + meta_idx = tidx + meta_i * cutlass.Int32(128) + if meta_idx < batch_size + 1: + s_cu_seqlens[meta_idx] = cu_q_seqlens[meta_idx] + if meta_idx < batch_size: + s_kv_lengths[meta_idx] = kv_cache_lengths[meta_idx] + if cutlass.const_expr(quant_scale_qkv is None): + quant_scale_value = cutlass.Float32(1.0) + else: + quant_scale_value = quant_scale_qkv[0] + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + tCtAcc_epi = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + tCtAcc_x = gemm_sm100.transform_partitioned_tensor_layout(tCtAcc_epi) + tCgC_x = gemm_sm100.transform_partitioned_tensor_layout(tCgC) + epi_tile = (self.cta_tile_shape_mnk[0], self.epi_chunk) + tAcc_epi = cute.flat_divide(tCtAcc_x, epi_tile) + gC_epi = cute.flat_divide(tCgC_x, epi_tile) + copy_atom_t2r = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.epi_chunk)), + self.acc_dtype, + ) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + tTR_gC_part = thr_copy_t2r.partition_D(gC_epi) + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC_part[(None, None, None, 0, 0, 0, 0, 0)].shape, + self.acc_dtype, + ) + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tTR_rAcc_flat = cute.make_tensor(tTR_rAcc.iterator, cute.make_layout((self.epi_chunk,))) + tTR_rC_flat = cute.make_tensor(tTR_rC.iterator, cute.make_layout((self.epi_chunk,))) + simt_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), + self.c_dtype, + num_bits_per_copy=256, + l1c_evict_priority=CacheEvictionPriority.NO_ALLOCATE, + ) + tTR_rC_quads = cute.tiled_divide(tTR_rC_flat, (16,)) + sts_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.c_dtype, + num_bits_per_copy=128, + ) + pred_store = cute.make_rmem_tensor((1, *tTR_rC.shape[1:]), cutlass.Boolean) + if cutlass.const_expr(self.use_tma_store): + s_out_quads = cute.tiled_divide(sOut, (1, 16)) + gC_store = cute.local_tile(mC_tma, self.epi_store_tile, (None, None, None)) + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + cute.group_modes(sOut, 0, 2), + cute.group_modes(gC_store, 0, 2), + ) + + acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) + s_rope_quads = cute.tiled_divide(s_rope_t, (1, 4)) + cache_quads = cute.tiled_divide(cos_sin_cache, (1, 4)) + cs_vec_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + cutlass.Float32, + num_bits_per_copy=128, + ) + cs_ldg_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.Float32, + num_bits_per_copy=128, + load_cache_mode=cute.nvgpu.common.LoadCacheMode.GLOBAL, ) + lane = tidx & cutlass.Int32(31) + row_in_cta = warp_idx * cutlass.Int32(32) + lane + ldg_row_sel = lane >> cutlass.Int32(3) + ldg_chunk = lane & cutlass.Int32(self.rope_ldg_row_threads - 1) + warp_row0 = warp_idx * cutlass.Int32(32) + cache_qcols = cache_quads[((0, None), 0, None)] + cs_cache_addr = cache_qcols.iterator.toint().to(cutlass.Int64) + cs_stage_flat = cute.make_rmem_tensor((self.qk_rope_head_dim,), cutlass.Float32) + cs_stage = cute.tiled_divide(cs_stage_flat, (4,)) + cs_row_flat = cute.make_rmem_tensor((self.qk_rope_head_dim,), cutlass.Float32) + cs_row = cute.tiled_divide(cs_row_flat, (4,)) + while work_tile.is_valid_tile: - iket.range_push("tma_tile") - cur_tile_coord = work_tile.tile_idx - mma_tile_coord_m = cur_tile_coord[0] // cute.size(tiled_mma_akeep.thr_id.shape) - head_id = cur_tile_coord[1] + iket.range_push("epi_tile") - tAgA_slice = tAgA[(None, mma_tile_coord_m, None, 0)] - tBgB_slice = tBgB[(None, head_id, None, 0)] - tAgSFA_slice = tAgSFA[(None, mma_tile_coord_m, None, 0)] - tBgSFB_slice = tBgSFB[(None, head_id, None, 0)] + def cs_transpose_half(h): + for ldg_i in cutlass.range_constexpr(self.rope_ldg_iters): + cute.copy( + cs_vec_atom, + cs_stage[(None, h * self.rope_ldg_iters + ldg_i)], + s_rope_quads[ + ( + (0, None), + warp_row0 + + cutlass.Int32(self.rope_ldg_rows_per_iter * ldg_i) + + ldg_row_sel, + ldg_chunk, + ) + ], + ) + cute.arch.sync_warp() + for qc in cutlass.range_constexpr(self.rope_half_quads): + cute.copy( + cs_vec_atom, + s_rope_quads[((0, None), row_in_cta, qc)], + cs_row[(None, h * self.rope_half_quads + qc)], + ) + cute.arch.sync_warp() - if cutlass.const_expr(self.tma_prefetch_dist > 0): - for pf_k_tile in cutlass.range( - 0, - cutlass.min(self.tma_prefetch_dist, k_tile_cnt), - 1, - unroll=1, - ): - cute.prefetch(tma_atom_a, tAgA_slice[(None, pf_k_tile)]) - cute.prefetch(tma_atom_sfa, tAgSFA_slice[(None, pf_k_tile)]) - cute.prefetch(tma_atom_b, tBgB_slice[(None, pf_k_tile)]) - cute.prefetch(tma_atom_sfb, tBgSFB_slice[(None, pf_k_tile)]) + cur_tile_coord = work_tile.tile_idx + head_id = cur_tile_coord[1] + coord_m_cta = cur_tile_coord[0] * cutlass.Int32(self.cta_tile_shape_mnk[0]) + row = coord_m_cta + row_in_cta + row_in_bounds = row < m + if cutlass.const_expr(self.use_tma_store): + bSG_gC_tile = bSG_gC[(None, cur_tile_coord[0], None, 0)] + else: + pred_store[(0, 0, 0)] = row_in_bounds + mma_coord_m = cur_tile_coord[0] // cute.size(tiled_mma_akeep.thr_id.shape) + tTR_gC_tile = tTR_gC_part[(None, None, None, None, None, mma_coord_m, head_id, 0)] - ab_producer_state.reset_count() - peek_ab_empty_status = cutlass.Boolean(1) - if ab_producer_state.count < k_tile_cnt: - peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) - for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): - iket.range_push("tma_wait", k_tile) - ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) - iket.range_pop() - iket.range_push("tma_issue") - tma_bar = ab_pipeline.producer_get_barrier(ab_producer_state) - cute.copy( - tma_atom_a, - tAgA_slice[(None, ab_producer_state.count)], - tAsA[(None, ab_producer_state.index)], - tma_bar_ptr=tma_bar, - mcast_mask=a_full_mcast_mask, - ) - cute.copy( - tma_atom_sfa, - tAgSFA_slice[(None, ab_producer_state.count)], - tAsSFA[(None, ab_producer_state.index)], - tma_bar_ptr=tma_bar, - mcast_mask=sfa_full_mcast_mask, + iket.range_push("epi_position") + position = cutlass.Int32(-1) + if row_in_bounds: + candidate = self._position_for_row( + s_cu_seqlens, + s_kv_lengths, + helix_position_offsets, + row, + batch_size, ) + if candidate >= 0 and candidate < rope_positions: + position = candidate + iket.range_pop() - cute.copy( - tma_atom_b, - tBgB_slice[(None, ab_producer_state.count)], - tBsB[(None, ab_producer_state.index)], - tma_bar_ptr=tma_bar, - mcast_mask=b_full_mcast_mask, - ) - cute.copy( - tma_atom_sfb, - tBgSFB_slice[(None, ab_producer_state.count)], - tBsSFB[(None, ab_producer_state.index)], - tma_bar_ptr=tma_bar, - mcast_mask=sfb_full_mcast_mask, + iket.range_push("epi_cs_stage") + pos_qcol_base = position * cutlass.Int32(self.rope_cache_row_floats // 4) + src_qcols, src_addrs = [], [] + for ldg_i in cutlass.range_constexpr(self.rope_ldg_iters): + src_lane = cutlass.Int32(self.rope_ldg_rows_per_iter * ldg_i) + ldg_row_sel + src_qcol = cute.arch.shuffle_sync(pos_qcol_base, src_lane) + src_qcols.append(src_qcol) + src_addrs.append( + cs_cache_addr + (src_qcol + ldg_chunk).to(cutlass.Int64) * cutlass.Int64(16) ) - if cutlass.const_expr(self.tma_prefetch_dist > 0): - pf_k_tile = k_tile + self.tma_prefetch_dist - if pf_k_tile < k_tile_cnt: - cute.prefetch(tma_atom_a, tAgA_slice[(None, pf_k_tile)]) - cute.prefetch(tma_atom_sfa, tAgSFA_slice[(None, pf_k_tile)]) - cute.prefetch(tma_atom_b, tBgB_slice[(None, pf_k_tile)]) - cute.prefetch(tma_atom_sfb, tBgSFB_slice[(None, pf_k_tile)]) - iket.range_pop() - ab_producer_state.advance() - peek_ab_empty_status = cutlass.Boolean(1) - if ab_producer_state.count < k_tile_cnt: - peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) - iket.range_push("tma_wait_clc") - clc_pipeline.consumer_wait(clc_consumer_state) - iket.range_pop() - work_tile = tile_sched.get_current_work() - clc_pipeline.consumer_release(clc_consumer_state) - clc_consumer_state.advance() + for ldg_i in cutlass.range_constexpr(self.rope_ldg_iters): + if src_qcols[ldg_i] >= 0: + for h in cutlass.range_constexpr(2): + src_ptr = cute.make_ptr( + cutlass.Float32, + src_addrs[ldg_i] + cutlass.Int64(h * self.rope_half_quads * 16), + cute.AddressSpace.gmem, + assumed_align=16, + ) + cute.copy( + cs_ldg_atom, + cute.make_tensor(src_ptr, cute.make_layout((4,))), + cs_stage[(None, h * self.rope_ldg_iters + ldg_i)], + ) iket.range_pop() - ab_pipeline.producer_tail(ab_producer_state) - if warp_idx == self.sched_warp_id: - cute.arch.setmaxregister_decrease(self.mainloop_reg_count) - if is_first_cta_in_cluster: - clc_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.ProducerConsumer, - self.num_clc_stage, - ) - while work_tile.is_valid_tile: - iket.range_push("sched_tile") - iket.range_push("sched_wait_empty") - clc_pipeline.producer_acquire(clc_producer_state) - iket.range_pop() - iket.range_push("sched_query") - mbarrier_addr = clc_pipeline.producer_get_barrier(clc_producer_state) - tile_sched.advance_to_next_work(mbarrier_addr) - clc_producer_state.advance() + cs_transpose_half(1) - clc_pipeline.consumer_wait(clc_consumer_state) - work_tile = tile_sched.get_current_work() - clc_pipeline.consumer_release(clc_consumer_state) - clc_consumer_state.advance() - iket.range_pop() - iket.range_pop() - clc_pipeline.producer_tail(clc_producer_state) - - if warp_idx == self.mma_warp_id: - cute.arch.setmaxregister_decrease(self.mainloop_reg_count) - tmem.wait_for_alloc() - acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) - sfa_tmem_ptr = cute.recast_ptr( - acc_tmem_ptr + self.num_accumulator_tmem_cols, - dtype=self.sf_dtype, - ) - tCtSFA = cute.make_tensor(sfa_tmem_ptr, self.tCtSFA_layout) - sfb_tmem_ptr = cute.recast_ptr( - acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, - dtype=self.sf_dtype, - ) - tCtSFB = cute.make_tensor(sfb_tmem_ptr, self.tCtSFB_layout) - - sfa_s2t = self._mainloop_s2t_copy_and_partition(sSFA, tCtSFA) - sfb_s2t = self._mainloop_s2t_copy_and_partition(sSFB, tCtSFB) - - ab_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_ab_stage - ) - acc_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) - while work_tile.is_valid_tile: - iket.range_push("mma_tile") - ab_consumer_state.reset_count() - peek_ab_full_status = cutlass.Boolean(1) - if ab_consumer_state.count < k_tile_cnt and is_leader_cta: - peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) - iket.range_push("mma_wait_acc") - if is_leader_cta: - acc_pipeline.producer_acquire(acc_producer_state) - iket.range_pop() - - for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): - if is_leader_cta: - iket.range_push("mma_wait", k_tile) - ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) - iket.range_pop() - iket.range_push("mma_issue") - self._mainloop_s2t_copies(ab_consumer_state.index, sfa_s2t, sfb_s2t) - num_kblocks = cute.size(tCrA, mode=[2]) - for k_block in cutlass.range(num_kblocks, unroll_full=True): - a_kblk_crd = ( - None, - 0, - k_block, - ab_consumer_state.index, - ) - sfa_kblk_crd = (None, 0, k_block) - for n_set in cutlass.range_constexpr(2): - b_kblk_crd = ( - None, - n_set, - k_block, - ab_consumer_state.index, - ) - sfb_kblk_crd = (None, n_set, k_block) - tCtAcc_set = tCtAcc_base[(None, 0, n_set)] - if cutlass.const_expr(n_set == 0): - tiled_mma_akeep.set( - tcgen05.Field.ACCUMULATE, - k_tile != 0 or k_block != 0, - ) - cute.gemm( - tiled_mma_akeep, - tCtAcc_set, - [ - tCrA[a_kblk_crd], - tCtSFA[sfa_kblk_crd], - ], - [ - tCrB[b_kblk_crd], - tCtSFB[sfb_kblk_crd], - ], - tCtAcc_set, - ) - else: - tiled_mma_areuse.set( - tcgen05.Field.ACCUMULATE, - k_tile != 0 or k_block != 0, - ) - cute.gemm( - tiled_mma_areuse, - tCtAcc_set, - [ - tCrA[a_kblk_crd], - tCtSFA[sfa_kblk_crd], - ], - [ - tCrB[b_kblk_crd], - tCtSFB[sfb_kblk_crd], - ], - tCtAcc_set, - ) - ab_pipeline.consumer_release(ab_consumer_state) - iket.range_pop() - ab_consumer_state.advance() - peek_ab_full_status = cutlass.Boolean(1) - if ab_consumer_state.count < k_tile_cnt: - if is_leader_cta: - peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) - if is_leader_cta: - acc_pipeline.producer_commit(acc_producer_state) - acc_producer_state.advance() - iket.range_push("mma_wait_clc") - clc_pipeline.consumer_wait(clc_consumer_state) - iket.range_pop() - work_tile = tile_sched.get_current_work() - clc_pipeline.consumer_release(clc_consumer_state) - clc_consumer_state.advance() - iket.range_pop() - acc_pipeline.producer_tail(acc_producer_state) - - # CUTLASS DSL 4.5 cannot carry a cute.struct through a dynamic - # warp-role branch. Materialize SMEM tensors before that branch. - s_cu_seqlens = storage.sCuSeqlens.get_tensor(cute.make_layout(self.max_batch + 1)) - s_kv_lengths = storage.sKvLengths.get_tensor(cute.make_layout(self.max_batch)) - sOut = None - if cutlass.const_expr(self.use_tma_store): - sOut = storage.sOut.get_tensor( - self.out_smem_layout_staged.outer, - swizzle=self.out_smem_layout_staged.inner, - ) - s_rope_t = storage.sRope.get_tensor( - self.rope_smem_layout.outer, - swizzle=self.rope_smem_layout.inner, - ) - - if warp_idx == self.pad_warp_id: - cute.arch.setmaxregister_decrease(self.mainloop_reg_count) - - if warp_idx < self.mma_warp_id: - cute.arch.setmaxregister_increase(self.epilog_reg_count) - for meta_i in cutlass.range(0, (batch_size + 128) // 128, 1, unroll=1): - meta_idx = tidx + meta_i * cutlass.Int32(128) - if meta_idx < batch_size + 1: - s_cu_seqlens[meta_idx] = cu_q_seqlens[meta_idx] - if meta_idx < batch_size: - s_kv_lengths[meta_idx] = kv_cache_lengths[meta_idx] - - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - - if cutlass.const_expr(quant_scale_qkv is None): - quant_scale_value = cutlass.Float32(1.0) - else: - quant_scale_value = quant_scale_qkv[0] - - tCtAcc_epi = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) - tCtAcc_x = gemm_sm100.transform_partitioned_tensor_layout(tCtAcc_epi) - tCgC_x = gemm_sm100.transform_partitioned_tensor_layout(tCgC) - epi_tile = (self.cta_tile_shape_mnk[0], self.epi_chunk) - tAcc_epi = cute.flat_divide(tCtAcc_x, epi_tile) - gC_epi = cute.flat_divide(tCgC_x, epi_tile) - copy_atom_t2r = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.epi_chunk)), - self.acc_dtype, - ) - tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0)]) - thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) - tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) - tTR_gC_part = thr_copy_t2r.partition_D(gC_epi) - tTR_rAcc = cute.make_rmem_tensor( - tTR_gC_part[(None, None, None, 0, 0, 0, 0, 0)].shape, - self.acc_dtype, - ) - tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) - tTR_rAcc_flat = cute.make_tensor(tTR_rAcc.iterator, cute.make_layout((self.epi_chunk,))) - tTR_rC_flat = cute.make_tensor(tTR_rC.iterator, cute.make_layout((self.epi_chunk,))) - simt_atom = cute.make_copy_atom( - cute.nvgpu.CopyR2GOp(), - self.c_dtype, - num_bits_per_copy=256, - l1c_evict_priority=CacheEvictionPriority.NO_ALLOCATE, - ) - tTR_rC_quads = cute.tiled_divide(tTR_rC_flat, (16,)) - sts_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - self.c_dtype, - num_bits_per_copy=128, - ) - pred_store = cute.make_rmem_tensor((1, *tTR_rC.shape[1:]), cutlass.Boolean) - if cutlass.const_expr(self.use_tma_store): - s_out_quads = cute.tiled_divide(sOut, (1, 16)) - gC_store = cute.local_tile(mC_tma, self.epi_store_tile, (None, None, None)) - bSG_sC, bSG_gC = cpasync.tma_partition( - tma_atom_c, - 0, - cute.make_layout(1), - cute.group_modes(sOut, 0, 2), - cute.group_modes(gC_store, 0, 2), - ) - - acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) - s_rope_quads = cute.tiled_divide(s_rope_t, (1, 4)) - cache_quads = cute.tiled_divide(cos_sin_cache, (1, 4)) - cs_copy_atom = cute.make_copy_atom( - cpasync.CopyG2SOp(), cutlass.Float32, num_bits_per_copy=128 - ) - lane = tidx & cutlass.Int32(31) - row_in_cta = warp_idx * cutlass.Int32(32) + lane - cp_chunk = lane & cutlass.Int32(self.rope_cp_chunks - 1) - cp_row_sel = lane >> cutlass.Int32(4) - warp_row0 = warp_idx * cutlass.Int32(32) - cache_qcols = cache_quads[((0, None), 0, None)] - - while work_tile.is_valid_tile: - iket.range_push("epi_tile") - cur_tile_coord = work_tile.tile_idx - head_id = cur_tile_coord[1] - coord_m_cta = cur_tile_coord[0] * cutlass.Int32(self.cta_tile_shape_mnk[0]) - row = coord_m_cta + row_in_cta - row_in_bounds = row < m - if cutlass.const_expr(self.use_tma_store): - bSG_gC_tile = bSG_gC[(None, cur_tile_coord[0], None, 0)] - else: - pred_store[(0, 0, 0)] = row_in_bounds - mma_coord_m = cur_tile_coord[0] // cute.size(tiled_mma_akeep.thr_id.shape) - tTR_gC_tile = tTR_gC_part[(None, None, None, None, None, mma_coord_m, head_id, 0)] - - iket.range_push("epi_position") - position = cutlass.Int32(-1) - if row_in_bounds: - candidate = self._position_for_row( - s_cu_seqlens, - s_kv_lengths, - helix_position_offsets, - row, - batch_size, - ) - if candidate >= 0 and candidate < rope_positions: - position = candidate - iket.range_pop() - - iket.range_push("epi_cs_stage") - pos_qcol_base = position * cutlass.Int32(self.rope_cache_row_floats // 4) - for cp_i in cutlass.range_constexpr(self.rope_cp_iters): - src_lane = cutlass.Int32(2 * cp_i) + cp_row_sel - src_qcol = cute.arch.shuffle_sync(pos_qcol_base, src_lane) - if src_qcol >= 0: - cute.copy( - cs_copy_atom, - cache_qcols[(None, src_qcol + cp_chunk)], - s_rope_quads[((0, None), warp_row0 + src_lane, cp_chunk)], - ) - cute.arch.cp_async_commit_group() - iket.range_pop() - - iket.range_push("epi_wait_acc") - acc_pipeline.consumer_wait(acc_consumer_state) - iket.range_pop() + iket.range_push("epi_wait_acc") + acc_pipeline.consumer_wait(acc_consumer_state) + iket.range_pop() iket.range_push("epi_rmsnorm_reduce") f32_zero = self.acc_dtype(0.0) @@ -1463,60 +1390,53 @@ def kernel_body( norm_quant_scale = inv_rms * quant_scale_value iket.range_pop() - iket.range_push("epi_cs_drain") - cute.arch.cp_async_wait_group(0) if cutlass.const_expr(self.use_tma_store): if warp_idx == 0: cute.arch.cp_async_bulk_wait_group(0, read=True) - self.epilog_sync_barrier.arrive_and_wait() - iket.range_pop() + self.epilog_sync_barrier.arrive_and_wait() - for rope_rev in cutlass.range_constexpr(self.rope_chunks): - chunk = self.epi_chunks - 1 - rope_rev - if cutlass.const_expr(rope_rev != 0): - cute.copy( - tiled_copy_t2r, - tTR_tAcc[(None, None, None, 0, chunk)], - tTR_rAcc, - ) - iket.range_push("epi_rope") + def rope_chunk_body(chunk): col_base = chunk * self.epi_chunk for quad_i in cutlass.range_constexpr(self.epi_chunk // 4): elem0 = quad_i * 4 quad_idx = (col_base - self.qk_nope_head_dim) // 4 + quad_i - cs_quad = s_rope_quads[((0, None), row_in_cta, quad_idx)].load() + cs_quad = cs_row[(None, quad_idx)].load() for sub in cutlass.range_constexpr(2): e0 = elem0 + sub * 2 cos_value = cs_quad[sub * 2] sin_value = cs_quad[sub * 2 + 1] + z_pk = _prmt_trunc_bf16x2( + tTR_rAcc_flat[e0], + tTR_rAcc_flat[e0 + 1], + ) x0, x1 = _fma_packed_f32x2_bf16x2_f32x2_f32x2( - ( - tTR_rAcc_flat[e0], - tTR_rAcc_flat[e0 + 1], - ), + z_pk, (inv_rms, inv_rms), (f32_zero, f32_zero), ) + x_pk = _prmt_trunc_bf16x2(x0, x1) + x_pk_rev = _prmt_trunc_bf16x2_neg_lo(x1, x0) t_pair = _fma_packed_f32x2_bf16x2_f32x2_f32x2( - (x0, x1), + x_pk, (cos_value, cos_value), (f32_zero, f32_zero), ) y0, y1 = _fma_packed_f32x2_bf16x2_f32x2_f32x2( - (x1, x0), - (-sin_value, sin_value), + x_pk_rev, + (sin_value, sin_value), t_pair, ) + y_pk = _prmt_trunc_bf16x2(y0, y1) o0, o1 = _fma_packed_f32x2_bf16x2_f32x2_f32x2( - (y0, y1), + y_pk, (quant_scale_value, quant_scale_value), (f32_zero, f32_zero), ) tTR_rAcc_flat[e0] = o0 tTR_rAcc_flat[e0 + 1] = o1 - out_ssa = tTR_rAcc_flat.load() - iket.range_pop() - iket.range_push("epi_quant+stg") + return tTR_rAcc_flat.load() + + def store_chunk(chunk, out_ssa): tTR_rC_flat.store(out_ssa.to(self.c_dtype)) if cutlass.const_expr(self.use_tma_store): q0 = (chunk & 3) * 2 @@ -1534,7 +1454,12 @@ def kernel_body( tTR_gC_tile[(None, None, None, 0, chunk)], pred=pred_store, ) - iket.range_pop() + + rope_hi_chunk = self.epi_chunks - 1 + out_hi = rope_chunk_body(rope_hi_chunk) + store_chunk(rope_hi_chunk, out_hi) + + cs_transpose_half(0) for chunk_i in cutlass.range_constexpr(self.nope_chunks): chunk = self.nope_chunks - 1 - chunk_i cute.copy( @@ -1546,40 +1471,33 @@ def kernel_body( cute.arch.fence_view_async_tmem_load() for pair in cutlass.range_constexpr(self.epi_chunk // 2): e0 = pair * 2 + z_pk = _prmt_trunc_bf16x2( + tTR_rAcc_flat[e0], + tTR_rAcc_flat[e0 + 1], + ) o0, o1 = _fma_packed_f32x2_bf16x2_f32x2_f32x2( - ( - tTR_rAcc_flat[e0], - tTR_rAcc_flat[e0 + 1], - ), + z_pk, (norm_quant_scale, norm_quant_scale), (f32_zero, f32_zero), ) tTR_rAcc_flat[e0] = o0 tTR_rAcc_flat[e0 + 1] = o1 out_ssa = tTR_rAcc_flat.load() - if cutlass.const_expr(chunk_i == self.nope_chunks - 1): - nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC) - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() - iket.range_push("epi_quant+stg") - tTR_rC_flat.store(out_ssa.to(self.c_dtype)) - if cutlass.const_expr(self.use_tma_store): - q0 = (chunk & 3) * 2 - slot = (chunk >> 2) & (self.out_ring_slots - 1) - for h in cutlass.range_constexpr(2): - cute.copy( - sts_atom, - tTR_rC_quads[(None, h)], - s_out_quads[((0, None), row_in_cta, q0 + h, slot)], - ) - else: - cute.copy( - simt_atom, - tTR_rC, - tTR_gC_tile[(None, None, None, 0, chunk)], - pred=pred_store, - ) - iket.range_pop() + store_chunk(chunk, out_ssa) + + rope_lo_chunk = self.nope_chunks + cute.copy( + tiled_copy_t2r, + tTR_tAcc[(None, None, None, 0, rope_lo_chunk)], + tTR_rAcc, + ) + + cute.arch.fence_view_async_tmem_load() + nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC) + acc_pipeline.consumer_release(acc_consumer_state) + out_lo = rope_chunk_body(rope_lo_chunk) + acc_consumer_state.advance() + store_chunk(rope_lo_chunk, out_lo) if cutlass.const_expr(self.use_tma_store): iket.range_push("epi_fence+tmastg") @@ -1616,6 +1534,244 @@ def kernel_body( tmem.free(acc_tmem_ptr) iket.range_pop() + if warp_idx == self.tma_warp_id: + cute.arch.setmaxregister_decrease(self.mainloop_reg_count) + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + a_cache_policy = cute.CacheEvictionPriority.EVICT_LAST + a_cache_policy = cutlass.Int64(0x14F0000000000000) # L2 evict_last + ab_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + while work_tile.is_valid_tile: + iket.range_push("tma_tile") + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_m = cur_tile_coord[0] // cute.size(tiled_mma_akeep.thr_id.shape) + head_id = cur_tile_coord[1] + + tAgA_slice = tAgA[(None, mma_tile_coord_m, None, 0)] + tBgB_slice = tBgB[(None, head_id, None, 0)] + tAgSFA_slice = tAgSFA[(None, mma_tile_coord_m, None, 0)] + tBgSFB_slice = tBgSFB[(None, head_id, None, 0)] + + if cutlass.const_expr(self.tma_prefetch_dist > 0): + for pf_k_tile in cutlass.range( + 0, + cutlass.min(self.tma_prefetch_dist, k_tile_cnt), + 1, + unroll=1, + ): + cute.prefetch(tma_atom_a, tAgA_slice[(None, pf_k_tile)]) + cute.prefetch(tma_atom_sfa, tAgSFA_slice[(None, pf_k_tile)]) + cute.prefetch(tma_atom_b, tBgB_slice[(None, pf_k_tile)]) + cute.prefetch(tma_atom_sfb, tBgSFB_slice[(None, pf_k_tile)]) + + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + iket.range_push("tma_wait", k_tile) + ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) + iket.range_pop() + iket.range_push("tma_issue") + tma_bar = ab_pipeline.producer_get_barrier(ab_producer_state) + cute.copy( + tma_atom_a, + tAgA_slice[(None, ab_producer_state.count)], + tAsA[(None, ab_producer_state.index)], + tma_bar_ptr=tma_bar, + mcast_mask=a_full_mcast_mask, + cache_policy=a_cache_policy, + ) + cute.copy( + tma_atom_sfa, + tAgSFA_slice[(None, ab_producer_state.count)], + tAsSFA[(None, ab_producer_state.index)], + tma_bar_ptr=tma_bar, + mcast_mask=sfa_full_mcast_mask, + cache_policy=a_cache_policy, + ) + + cute.copy( + tma_atom_b, + tBgB_slice[(None, ab_producer_state.count)], + tBsB[(None, ab_producer_state.index)], + tma_bar_ptr=tma_bar, + mcast_mask=b_full_mcast_mask, + ) + cute.copy( + tma_atom_sfb, + tBgSFB_slice[(None, ab_producer_state.count)], + tBsSFB[(None, ab_producer_state.index)], + tma_bar_ptr=tma_bar, + mcast_mask=sfb_full_mcast_mask, + ) + if cutlass.const_expr(self.tma_prefetch_dist > 0): + pf_k_tile = k_tile + self.tma_prefetch_dist + if pf_k_tile < k_tile_cnt: + cute.prefetch(tma_atom_a, tAgA_slice[(None, pf_k_tile)]) + cute.prefetch(tma_atom_sfa, tAgSFA_slice[(None, pf_k_tile)]) + cute.prefetch(tma_atom_b, tBgB_slice[(None, pf_k_tile)]) + cute.prefetch(tma_atom_sfb, tBgSFB_slice[(None, pf_k_tile)]) + iket.range_pop() + ab_producer_state.advance() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + iket.range_push("tma_wait_clc") + clc_pipeline.consumer_wait(clc_consumer_state) + iket.range_pop() + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + iket.range_pop() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + ab_pipeline.producer_tail(ab_producer_state) + + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.mainloop_reg_count) + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, self.tCtSFA_layout) + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, self.tCtSFB_layout) + + sfa_s2t = self._mainloop_s2t_copy_and_partition(sSFA, tCtSFA) + sfb_s2t = self._mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + ab_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage + ) + acc_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) + while work_tile.is_valid_tile: + iket.range_push("mma_tile") + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + iket.range_push("mma_wait_acc") + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + iket.range_pop() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + if is_leader_cta: + iket.range_push("mma_wait", k_tile) + ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) + iket.range_pop() + iket.range_push("mma_issue") + self._mainloop_s2t_copies(ab_consumer_state.index, sfa_s2t, sfb_s2t) + num_kblocks = cute.size(tCrA, mode=[2]) + for k_block in cutlass.range(num_kblocks, unroll_full=True): + a_kblk_crd = ( + None, + 0, + k_block, + ab_consumer_state.index, + ) + sfa_kblk_crd = (None, 0, k_block) + for n_set in cutlass.range_constexpr(2): + b_kblk_crd = ( + None, + n_set, + k_block, + ab_consumer_state.index, + ) + sfb_kblk_crd = (None, n_set, k_block) + tCtAcc_set = tCtAcc_base[(None, 0, n_set)] + if cutlass.const_expr(n_set == 0): + tiled_mma_akeep.set( + tcgen05.Field.ACCUMULATE, + k_tile != 0 or k_block != 0, + ) + cute.gemm( + tiled_mma_akeep, + tCtAcc_set, + [ + tCrA[a_kblk_crd], + tCtSFA[sfa_kblk_crd], + ], + [ + tCrB[b_kblk_crd], + tCtSFB[sfb_kblk_crd], + ], + tCtAcc_set, + ) + else: + tiled_mma_areuse.set( + tcgen05.Field.ACCUMULATE, + k_tile != 0 or k_block != 0, + ) + cute.gemm( + tiled_mma_areuse, + tCtAcc_set, + [ + tCrA[a_kblk_crd], + tCtSFA[sfa_kblk_crd], + ], + [ + tCrB[b_kblk_crd], + tCtSFB[sfb_kblk_crd], + ], + tCtAcc_set, + ) + ab_pipeline.consumer_release(ab_consumer_state) + iket.range_pop() + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt: + if is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + iket.range_push("mma_wait_clc") + clc_pipeline.consumer_wait(clc_consumer_state) + iket.range_pop() + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + iket.range_pop() + acc_pipeline.producer_tail(acc_producer_state) + + if warp_idx == self.sched_warp_id: + cute.arch.setmaxregister_decrease(self.mainloop_reg_count) + if is_first_cta_in_cluster: + clc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.ProducerConsumer, + self.num_clc_stage, + ) + while work_tile.is_valid_tile: + iket.range_push("sched_tile") + iket.range_push("sched_wait_empty") + clc_pipeline.producer_acquire(clc_producer_state) + iket.range_pop() + iket.range_push("sched_query") + mbarrier_addr = clc_pipeline.producer_get_barrier(clc_producer_state) + tile_sched.advance_to_next_work(mbarrier_addr) + clc_producer_state.advance() + + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + iket.range_pop() + iket.range_pop() + clc_pipeline.producer_tail(clc_producer_state) + + if warp_idx == self.pad_warp_id: + cute.arch.setmaxregister_decrease(self.mainloop_reg_count) + @cute.jit def _mainloop_s2t_copy_and_partition(self, sSF: cute.Tensor, tSF: cute.Tensor) -> S2TCopyBundle: """Tiled S2T (UTCCP) copy for one SF tensor + its partitions.""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py index c26f544877b8..89351bf6c0d1 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py @@ -90,7 +90,11 @@ def silu_f32( return a * sigmoid_f32(a, fastmath=fastmath) -SUPPORTED_ACTIVATION_TYPES = (ActivationType.Swiglu, ActivationType.Relu2) +SUPPORTED_ACTIVATION_TYPES = ( + ActivationType.Swiglu, + ActivationType.SiTu, + ActivationType.Relu2, +) def validate_activation_type(activation_type) -> ActivationType: @@ -113,7 +117,7 @@ class S2TCopyBundle(NamedTuple): """ Rubin (SM107) persistent blockscaled contiguous grouped GEMM with token gather -and fused SwiGLU or Relu2 activation (FC1 of MoE). +and fused SwiGLU, SiTU, or Relu2 activation (FC1 of MoE). Compute: acc = alpha * (SFA * A[token_ids]) * (SFB * B) # GEMM @@ -141,7 +145,7 @@ class Sm107BlockScaledContiguousGatherGroupedGemmActFusionKernel: """Rubin (SM107) FC1: contiguous grouped blockscaled GEMM with token gather on A/SFA and activation fusion in the epilogue. - Supports both SwiGLU (gated) and Relu2 (non-gated) activations. + Supports SwiGLU and SiTU (gated), plus Relu2 (non-gated), activations. Builds on Sm107BlockScaledContiguousGroupedGemmKernel (persistent tile scheduling, warp specialization, B-reuse, tcgen05.mma block-scale, TMA @@ -159,7 +163,7 @@ class Sm107BlockScaledContiguousGatherGroupedGemmActFusionKernel: merged ab_pipeline (no relay warp needed). SFA is always loaded via CpAsync128.CG, then reorganized into SFA TMEM by transform warps via LDS + STTM (sfa_transform_pipeline). - - SwiGLU epilogue: C = up * silu(gate), where up/gate come from + - Gated epilogue (SwiGLU or SiTU): C = up * silu(gate), where up/gate come from interleaved accumulator at granularity=64 → output N is halved. - Optional NVFP4 quant: when c_dtype == Float4E2M1FN, the epilogue also generates SFC and quantizes the output. @@ -197,6 +201,8 @@ def __init__( use_pdl: bool = True, locality_domain_half_gemm: bool = False, activation_type: ActivationType = ActivationType.Swiglu, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, ): self.a_path = a_path # locality domain half-GEMM: two partitions write their N-half into a shared @@ -205,6 +211,23 @@ def __init__( self.sf_vec_size = sf_vec_size self.topk = topk self.activation_type = validate_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}." + ) + 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.is_gated = is_gated_activation(self.activation_type) if locality_domain_half_gemm and not self.is_gated: raise ValueError("Rubin locality domain half-GEMM currently supports SwiGLU only") @@ -3028,6 +3051,9 @@ def kernel( self._apply_swiglu_epilogue( acc_vec_up, acc_vec_gate, alpha_val, tCompute ) + elif cutlass.const_expr(self.activation_type == ActivationType.SiTu): + acc_vec_gate = tTR_rAcc_gate.load() + self._apply_situ_epilogue(acc_vec_up, acc_vec_gate, alpha_val, tCompute) elif cutlass.const_expr(self.activation_type == ActivationType.Relu2): self._apply_relu2_epilogue(acc_vec_up, alpha_val, tCompute) @@ -3283,6 +3309,93 @@ def _apply_swiglu_epilogue( acc_vec_gate_alpha = acc_vec_gate[i] * cutlass.Float32(alpha_val) 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 ``kimi_k3_moe/_mlp.py::SituAndMul`` + (itself byte-identical to HF ``modeling_kimi.py``):: + + 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 + + ``up`` and ``gate`` come from the two interleaved accumulator subtiles + loaded by the caller, same as the SwiGLU epilogue. + + There is no packed tanh, so the vectorized path uses the identity + ``tanh(z) = 2 * sigmoid(2z) - 1`` (the same one ``utils.gelu_tanh_f32`` + uses) to stay on the packed f32x2 path -- calling a scalar tanh would + force the whole loop back to scalar. The reciprocals and ``2*beta`` + factors fold at trace time because both betas are ``const_expr``:: + + beta * tanh(x/beta) = beta * (2*sigmoid(2x/beta) - 1) + = 2*beta*sigmoid((2/beta)*x) - beta + """ + 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) + + # sigmoid(x) = rcp(1 + exp2(-x * log2e)), shared by both cores. + 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 + @cute.jit def _apply_relu2_epilogue( self, diff --git a/tensorrt_llm/_torch/cute_dsl_utils.py b/tensorrt_llm/_torch/cute_dsl_utils.py index cb571f85b597..352e1de26e7a 100644 --- a/tensorrt_llm/_torch/cute_dsl_utils.py +++ b/tensorrt_llm/_torch/cute_dsl_utils.py @@ -1,7 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import platform +import sys +import types from ..logger import logger + +def _skip_legacy_cutlass_mlir_helpers() -> None: + """Keep Cutlass version discovery from importing its legacy helper tree.""" + legacy_name = "cutlass.base_dsl._mlir_helpers" + if legacy_name in sys.modules: + return + + # Cutlass uses pkgutil.walk_packages to hash its sources. The internal + # package also ships this unused legacy helper tree alongside the canonical + # cutlass._mlir_helpers package. Prevent pkgutil from descending into the + # legacy tree, which would register the same MLIR value casters twice. + legacy_module = types.ModuleType(legacy_name) + legacy_module.__path__ = [] + sys.modules[legacy_name] = legacy_module + + IS_CUTLASS_DSL_AVAILABLE = False # Whether the public CuTe DSL package provides the SM107/Rubin helper module. @@ -12,8 +45,8 @@ if platform.system() != "Windows": try: - import cutlass # noqa - import cutlass.cute as cute # noqa + from cutlass import cute # noqa + _skip_legacy_cutlass_mlir_helpers() logger.info(f"cutlass dsl is available") IS_CUTLASS_DSL_AVAILABLE = True diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index ba1a38fa6956..3a2e180664f1 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -713,6 +713,7 @@ def __init__( reduce_output: bool = True, # ROW parallel only skip_create_weights_in_init: bool = False, use_custom_cublas_mm: bool = False, + use_cute_dsl_bf16_gemm: bool = False, use_cute_dsl_blockscaling_mm: bool = False, lora: Optional[LoraLayer] = None, ): @@ -720,16 +721,17 @@ def __init__( in_features, out_features, bias, - dtype, - mapping, - tensor_parallel_mode, - gather_output, - quant_config, - weights_loading_config, - reduce_output, - skip_create_weights_in_init, - use_custom_cublas_mm, - lora, + dtype=dtype, + mapping=mapping, + tensor_parallel_mode=tensor_parallel_mode, + gather_output=gather_output, + quant_config=quant_config, + weights_loading_config=weights_loading_config, + reduce_output=reduce_output, + skip_create_weights_in_init=skip_create_weights_in_init, + use_custom_cublas_mm=use_custom_cublas_mm, + use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm, + lora=lora, use_cute_dsl_blockscaling_mm=use_cute_dsl_blockscaling_mm, ) @@ -739,7 +741,12 @@ def apply_linear(self, lora_params: Optional[dict] | None = None, layer_idx: Optional[int] | None = None): num_tokens = input.shape[0] - if (not self.has_any_quant and 1 <= num_tokens <= 16 + has_any_quant = self.has_any_quant + use_cute_dsl_bf16_gemm = (self.use_cute_dsl_bf16_gemm + and not has_any_quant and is_sm_100f() + and self.weight.dtype == torch.bfloat16) + if (not use_cute_dsl_bf16_gemm and not has_any_quant + and 1 <= num_tokens <= 16 and get_sm_version() not in [120, 121]): output = torch.ops.trtllm.dsv3_fused_a_gemm_op( input, self.weight.t(), bias, None) @@ -793,6 +800,7 @@ def __init__( skip_create_weights_in_init=model_config. skip_create_weights_in_init, use_custom_cublas_mm=True, + use_cute_dsl_bf16_gemm=model_config.use_cute_dsl_bf16_gemm, use_cute_dsl_blockscaling_mm=model_config. use_cute_dsl_blockscaling_mm, ) @@ -843,7 +851,8 @@ def __init__( quant_config=model_config.get_quant_config(), skip_create_weights_in_init=model_config. skip_create_weights_in_init, - use_custom_cublas_mm=True) + use_custom_cublas_mm=True, + use_cute_dsl_bf16_gemm=model_config.use_cute_dsl_bf16_gemm) class DeepseekV3Gate(nn.Module): diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index e854fc6f6769..0e1a7638a702 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -781,8 +781,8 @@ def load_o_a_proj(module_name: str, module) -> None: if o_a_proj_scale is not None: o_a_proj_scale = split_matrix_tp(o_a_proj_scale, tp_size, tp_rank, 0) - # Skip the BF16 dequant when the destination is FP8 (the cute_dsl - # FP8 BMM path on SM100 consumes the native FP8 weight directly). + # Skip BF16 dequant when the architecture-specific CuTe DSL BMM + # consumes the native FP8 weight directly. if o_a_proj_scale is not None and module.o_a_proj.dtype != torch.float8_e4m3fn: o_a_proj = weight_dequant( o_a_proj.reshape(-1, o_a_proj.shape[-1]).contiguous().cuda(), @@ -1598,6 +1598,7 @@ def __init__( config=model_config, overridden_tp_size=shared_tp_size, reduce_output=False, + use_cute_dsl_blockscaling_mm=model_config.use_cute_dsl_blockscaling_mm, swiglu_limit=swiglu_limit, ) @@ -2265,6 +2266,7 @@ def __init__( dtype=config.torch_dtype, quant_config=model_config.get_quant_config(), skip_create_weights_in_init=model_config.skip_create_weights_in_init, + use_cute_dsl_blockscaling_mm=model_config.use_cute_dsl_blockscaling_mm, ) self.h_proj = Linear( config.hidden_size, @@ -2273,6 +2275,7 @@ def __init__( dtype=config.torch_dtype, quant_config=model_config.get_quant_config(), skip_create_weights_in_init=model_config.skip_create_weights_in_init, + use_cute_dsl_blockscaling_mm=model_config.use_cute_dsl_blockscaling_mm, ) else: self.e_proj = Linear( @@ -2285,6 +2288,7 @@ def __init__( reduce_output=True, quant_config=model_config.get_quant_config(), skip_create_weights_in_init=model_config.skip_create_weights_in_init, + use_cute_dsl_blockscaling_mm=model_config.use_cute_dsl_blockscaling_mm, ) self.h_proj = Linear( config.hidden_size, @@ -2296,6 +2300,7 @@ def __init__( reduce_output=True, quant_config=model_config.get_quant_config(), skip_create_weights_in_init=model_config.skip_create_weights_in_init, + use_cute_dsl_blockscaling_mm=model_config.use_cute_dsl_blockscaling_mm, ) self.shared_head = DeepseekV4MTPHead(model_config) diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py index 24bf6d21b9a6..6bb5fcbc6289 100644 --- a/tensorrt_llm/_torch/modules/gated_mlp.py +++ b/tensorrt_llm/_torch/modules/gated_mlp.py @@ -5,9 +5,11 @@ import torch.nn.functional as F from torch import nn +from tensorrt_llm._utils import get_sm_version, is_sm_100f from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping +from ..cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE from ..distributed import AllReduceParams from ..model_config import ModelConfig from ..peft.lora.layer import LoraLayer, LoraModuleType @@ -119,6 +121,15 @@ def __init__( allreduce_strategy=config.allreduce_strategy, force_dynamic_quantization=config.force_dynamic_quantization, use_cute_dsl_blockscaling_mm=use_cute_dsl_blockscaling_mm, + # The fused CuteDSL NVFP4 SwiGLU epilogue applies no clamp, so a + # layer carrying a real ``swiglu_limit`` must stay on the Triton + # kernel. ``_is_plain_swiglu`` deliberately covers only alpha/beta; + # the limit is gated here (as on ``rubin-advance``). Without this + # the clamp is silently dropped -- wrong numerics, no error. + use_cute_dsl_nvfp4_swiglu_blackwell=( + use_cute_dsl_blockscaling_mm + and activation == F.silu and not bias + and (swiglu_limit is None or swiglu_limit == float("inf"))), use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm, disable_deep_gemm=disable_deep_gemm, fused_weight_shard_indices_mapping=gateup_shard_indices_mapping, @@ -214,16 +225,11 @@ def _is_plain_swiglu(self): def _can_fuse_gate_up_swiglu(self): """Check if fused GEMM + SwiGLU path is available. - Returns True when all conditions are met: - - CuteDSL blockscaling mode is enabled (implies Blackwell + CuteDSL) - - Activation is plain SwiGLU (F.silu), see _is_plain_swiglu - - gate_up_proj uses NVFP4 quantization - - gate_up_proj has no bias (bias not supported in fused kernel) + The projection owns the capability predicate because weight loading + must make exactly the same decision as forward dispatch. """ - return (self.use_cute_dsl_blockscaling_mm and self.activation == F.silu - and self._is_plain_swiglu() - and self.gate_up_proj.has_nvfp4_activation_quantization - and not self.gate_up_proj.has_bias) + return (self.activation == F.silu and self._is_plain_swiglu() + and self.gate_up_proj.can_use_cute_dsl_nvfp4_swiglu_blackwell()) def _can_fuse_gate_up_swiglu_fp4out(self): """Check if fused GEMM + SwiGLU with FP4 output path is available. @@ -237,6 +243,22 @@ def _can_fuse_gate_up_swiglu_fp4out(self): return False return is_static_nvfp4_input_eligible(self.down_proj) + def _can_fuse_swiglu_fp8_quant(self) -> bool: + """Check whether down projection can consume fused SwiGLU FP8 output.""" + # silu_and_mul_fp8_quantize_1x128_packed_ue8m0 takes the limit but has + # no alpha/beta, so a parameterized SwiGLU must stay unfused. MiniMax + # M3 SwiGLU-OAI reaches here with plain F.silu plus swiglu_alpha and + # swiglu_beta, so the activation check alone does not exclude it. + if not (self.activation == F.silu and self._is_plain_swiglu() + and self.down_proj.has_fp8_block_scales): + return False + if get_sm_version() == 107: + return (IS_CUTLASS_DSL_RUBIN_AVAILABLE + and (self.down_proj.use_cute_dsl_blockscaling_mm + or self.down_proj.disable_deep_gemm)) + return (is_sm_100f() and not self.down_proj.use_cute_dsl_blockscaling_mm + and not self.down_proj.disable_deep_gemm) + def _fused_gate_up_swiglu(self, x, fp4_out=False): """Fused FC1 GEMM + SwiGLU using CuteDSL dense kernel. @@ -315,6 +337,7 @@ def forward( return self.forward_lora(x, all_rank_num_tokens, final_all_reduce_params, lora_params) + fused_output_shape = None if self._can_fuse_gate_up_swiglu_fp4out(): # During torch.compile the token dim is a SymInt, so `m >= MIN_M` # would create a SymBool guard that breaks piecewise CUDA graph @@ -333,11 +356,21 @@ def forward( h2 = self._fused_gate_up_swiglu(x) else: h1 = self.gate_up_proj(x) - h2 = self._apply_activation(h1) + if self._can_fuse_swiglu_fp8_quant(): + if h1.dim() > 2: + fused_output_shape = h1.shape[:-1] + h1 = h1.reshape(-1, h1.shape[-1]) + use_r128c4_layout = get_sm_version() == 107 + h2 = torch.ops.trtllm.silu_and_mul_fp8_quantize_1x128_packed_ue8m0( + h1, self.swiglu_limit, use_r128c4_layout) + else: + h2 = self._apply_activation(h1) output = self.down_proj(h2, all_reduce_params=final_all_reduce_params, layer_idx=self.layer_idx) + if fused_output_shape is not None: + output = output.reshape(*fused_output_shape, output.shape[-1]) return output def forward_lora( diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index b3e29a427482..cdacad5064b6 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -35,7 +35,8 @@ from ..._utils import get_sm_version, is_sm_100f from ...models.modeling_utils import QuantConfig -from ..cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE +from ..cute_dsl_utils import (IS_CUTLASS_DSL_AVAILABLE, + IS_CUTLASS_DSL_RUBIN_AVAILABLE) from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, replace_parameter_and_save_metadata, unswizzle_sf) from .low_m_gemm import _should_apply_low_m_gemm, apply_low_m_gemm @@ -1235,6 +1236,38 @@ def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor]): # fp8_block_scaling_gemm does not support writing into an NCCL window # buffer; supports_nccl_symmetric_memory_window_output is False so the window path is bypassed. + if isinstance(input, tuple): + if len(input) != 2: + raise ValueError( + "Pre-quantized FP8 input must contain activation and scale") + activation, activation_scale = input + sm_version = get_sm_version() + uses_cute_dsl_rubin = (activation_scale.dtype == torch.uint8 + and sm_version == 107 + and IS_CUTLASS_DSL_RUBIN_AVAILABLE + and (module.use_cute_dsl_blockscaling_mm + or module.disable_deep_gemm)) + if uses_cute_dsl_rubin: + output = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin( + activation, module.weight, activation_scale, + module.weight_scale) + elif (activation_scale.dtype == torch.int32 and is_sm_100f() + and not module.disable_deep_gemm): + output = torch.ops.trtllm.fp8_prequantized_swap_ab_gemm( + activation, + activation_scale, + module.weight, + module.weight_scale, + disable_ue8m0_cast=True, + ) + else: + raise RuntimeError( + "Pre-quantized FP8 scale layout is incompatible with the " + "selected block-scale GEMM backend") + if bias is not None: + output = output + bias + return output + # Handle multi-dimensional inputs (e.g., 3D: batch, seq, hidden) # GEMM ops require 2D matrices original_shape = input.shape @@ -2024,7 +2057,7 @@ def process_weights_after_loading_fused_gate_up_linear( # interleaves in 64-row groups to match the kernel layout. # # Weight scales are similarly unswizzled, interleaved, and re-swizzled. - if not module.use_cute_dsl_blockscaling_mm: + if not module.can_use_cute_dsl_nvfp4_swiglu_blackwell(): return group_size = 64 @@ -3625,6 +3658,7 @@ def __init__( allreduce_strategy: AllReduceStrategy = AllReduceStrategy.AUTO, force_dynamic_quantization: bool = False, use_cute_dsl_blockscaling_mm: bool = False, + use_cute_dsl_nvfp4_swiglu_blackwell: bool = False, disable_deep_gemm: bool = False, fused_weight_shard_indices_mapping: Optional[dict] = None, nvfp4_allowed_backends: Optional[List[str]] = None, @@ -3635,6 +3669,9 @@ def __init__( ): """ Args: + use_cute_dsl_nvfp4_swiglu_blackwell: Allow this fused gate/up + projection to use the Blackwell-only NVFP4 GEMM + SwiGLU + kernel and its required interleaved weight layout. nvfp4_allowed_backends: List of backends to consider for NVFP4 GEMM auto-selection. Default (via config): ['cutlass', 'cublaslt', 'cuda_core'] - excludes cutedsl for faster build. Add 'cutedsl' for extreme performance at the cost of longer build time. @@ -3659,6 +3696,8 @@ def __init__( self.gather_output = gather_output self.force_dynamic_quantization = force_dynamic_quantization self.use_cute_dsl_blockscaling_mm = use_cute_dsl_blockscaling_mm + self.use_cute_dsl_nvfp4_swiglu_blackwell = \ + use_cute_dsl_nvfp4_swiglu_blackwell self.disable_deep_gemm = disable_deep_gemm self.fused_weight_shard_indices_mapping = fused_weight_shard_indices_mapping # Store NVFP4 GEMM allowed backends configuration @@ -4007,6 +4046,18 @@ def has_nvfp4(self): return self.quant_config is not None and self.quant_config.layer_quant_mode.has_nvfp4( ) + def can_use_cute_dsl_nvfp4_swiglu_blackwell(self) -> bool: + """Return whether this layer can use the Blackwell NVFP4 SwiGLU op. + + Keep this predicate shared by weight transformation and forward + dispatch so a fallback backend never consumes the fused layout. + """ + return (self.use_cute_dsl_nvfp4_swiglu_blackwell + and self.use_cute_dsl_blockscaling_mm + and IS_CUTLASS_DSL_AVAILABLE + and self.has_nvfp4_activation_quantization + and get_sm_version() in (100, 103) and not self.has_bias) + @property def has_nvfp4_activation_quantization(self): assert self._weights_created diff --git a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py index 28aaadfee290..e29231a972a9 100644 --- a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py +++ b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py @@ -875,7 +875,7 @@ def add(tactic): if m_tiles * ks <= max_grid_ctas: for bs in _fused_hc_mma_bigfuse_bs_options(M): add(("fused_half_mma", 0, ks, bs, 1)) - if M >= 64: + if M >= 64 and ks in _FUSED_HC_ALL_MMA_KS: add(("fused_all_mma", 0, ks, 0, 1)) if not mma_ok and M > 32: diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/attention.py b/tensorrt_llm/_torch/visual_gen/models/flux/attention.py index 31df009e0be8..0d3b3054fd0d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/attention.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/attention.py @@ -512,7 +512,9 @@ def _can_project_hidden_mlp_with_cute_dsl(self) -> bool: return False return ( - gate_up_proj.use_cute_dsl_blockscaling_mm + # Same predicate the weight transform uses, so dispatch can never + # consume a layout the loader did not produce. + gate_up_proj.can_use_cute_dsl_nvfp4_swiglu_blackwell() and gate_up_proj.has_nvfp4 and not gate_up_proj.has_bias and self._is_cute_dsl_swiglu_layout_compatible( diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/joint_proj.py b/tensorrt_llm/_torch/visual_gen/models/flux/joint_proj.py index 7bc2a4938f7b..38a9eb03b168 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/joint_proj.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/joint_proj.py @@ -257,6 +257,10 @@ def range_size(r): tensor_parallel_mode=TensorParallelMode.COLUMN, reduce_output=False, use_cute_dsl_blockscaling_mm=use_cute_dsl_blockscaling_mm, + # Flux runs a plain, unclamped SwiGLU, so it opts in to the + # fused Blackwell NVFP4 epilogue and the interleaved gate/up + # weight layout that epilogue reads. + use_cute_dsl_nvfp4_swiglu_blackwell=use_cute_dsl_blockscaling_mm, override_tp_sharding={ "gate": (local_mlp_hidden_start, local_mlp_hidden_end), "up": (local_mlp_hidden_start, local_mlp_hidden_end), diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py index 5ef9a9793c40..7a00afc2e7fb 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py @@ -30,7 +30,7 @@ from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.module import ( project_sparse_attn_output, ) -from tensorrt_llm._torch.attention.mla import MLA +from tensorrt_llm._torch.attention.mla import MLA, _is_cute_dsl_fp8_bmm_available from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_deepseekv3 import weight_dequant from tensorrt_llm._utils import get_sm_version @@ -66,6 +66,7 @@ def calculate_reference_deepseek_v4_o_proj( qk_rope_head_dim, device, is_fp8: bool = False, + quantize_o_a_input: bool = False, ): """ Reference implementation for DeepSeek-V4 output projection based on ref/model.py. @@ -80,6 +81,7 @@ def calculate_reference_deepseek_v4_o_proj( qk_rope_head_dim: Dimension of positional part device: Device to run on is_fp8: Whether test fp8 or bf16 + quantize_o_a_input: Whether o_a_proj consumes block-scaled FP8 input Returns: output: [num_tokens, hidden_size] projected output @@ -92,7 +94,7 @@ def calculate_reference_deepseek_v4_o_proj( # Reshape for grouped projection attn_out_grouped = attn_out_latent.view(num_tokens, n_local_groups, -1) - if is_fp8: + if quantize_o_a_input: attn_out_grouped = _per_token_fp8_quant_dequant( attn_out_grouped.transpose(0, 1).contiguous() ).transpose(0, 1) @@ -124,6 +126,7 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): device = torch.device("cuda") dtype = torch.bfloat16 + sm_version = get_sm_version() # Model configuration matching the reference model num_heads = 64 @@ -184,7 +187,11 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): pretrained_config=pretrained_config, sparse_attention_config=sparse_config, quant_config=quant_config, - use_cute_dsl_blockscaling_mm=dtype_str == "fp8", + # Mirror the production predicate exactly, so the test does not enable + # the cute-dsl path on a box where the DSL package is absent. + use_cute_dsl_blockscaling_mm=( + dtype_str == "fp8" and _is_cute_dsl_fp8_bmm_available(sm_version) + ), ) # Setup positional embedding params @@ -219,6 +226,12 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): assert not hasattr(mla, "v_b_proj") assert not hasattr(mla, "o_proj") + has_native_fp8_o_proj = _is_cute_dsl_fp8_bmm_available(sm_version) + if dtype_str == "fp8" and has_native_fp8_o_proj: + assert mla.o_a_proj.dtype == torch.float8_e4m3fn + assert mla.o_a_proj_dequant is None + assert not mla.use_cute_dsl_blockscaling_bmm + # Initialize weights nn_init_std = 0.02 with torch.no_grad(): @@ -245,12 +258,8 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): fp8_a_weight = fp8_a_weight.reshape(n_local_groups, o_lora_rank, dim) mla.o_a_proj.data = fp8_a_weight mla.o_a_proj_scale.data = fp8_a_scale - # mla.o_a_proj_dequant is None for DSv4 on SM100: PR #14254 - # decouples the FP8-native o_a_proj path from - # use_cute_dsl_blockscaling_bmm, so DSv4 unconditionally uses the - # fused inv-RoPE + FP8 quant + cute-dsl BMM chain and never needs - # the bf16-dequant fallback buffer. The reference path below uses - # o_a_proj_bf16 directly. + if not has_native_fp8_o_proj: + mla.o_a_proj_dequant.data = o_a_proj_bf16 # Initialize o_b_proj weights if dtype_str == "bf16": @@ -289,16 +298,19 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): o_a_proj_ref = mla.o_a_proj.data o_b_proj_weight_ref = mla.o_b_proj.weight.data else: - # Match the FP8-native o_a_proj path: the runtime BMM consumes - # quantized o_a_proj plus block scales, not the original BF16 weight. - o_a_proj_ref = ( - weight_dequant( - fp8_a_weight.reshape(-1, dim).contiguous(), - fp8_a_scale.contiguous(), + if has_native_fp8_o_proj: + # Match the native path, which consumes the quantized weight and + # its block scales rather than the original BF16 values. + o_a_proj_ref = ( + weight_dequant( + fp8_a_weight.reshape(-1, dim).contiguous(), + fp8_a_scale.contiguous(), + ) + .bfloat16() + .reshape(o_a_proj_bf16.shape) ) - .bfloat16() - .reshape(o_a_proj_bf16.shape) - ) + else: + o_a_proj_ref = o_a_proj_bf16 o_b_proj_weight_ref = fp8_b_weight_dequant freqs_cis = precompute_freqs_cis( @@ -321,6 +333,7 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): qk_rope_head_dim=qk_rope_head_dim, device=device, is_fp8=dtype_str == "fp8", + quantize_o_a_input=dtype_str == "fp8" and has_native_fp8_o_proj, ) # Validate output shapes diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py new file mode 100644 index 000000000000..f5d9adcbc977 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py @@ -0,0 +1,279 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exercise prior-state ownership through the actual indexer forward path.""" + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer +from tensorrt_llm._torch.modules.top_k import _MAX_RADIX_BLOCKS_PER_ROW, TopK, TopKImplementation + + +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "cuda", + marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA"), + ), + ], +) +def test_update_gvr_prior_from_prefill_uses_device_lengths(device: str) -> None: + """Seed the last prefill row using lengths on the selections' device.""" + top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR, gvr_self_sampling=False) + prefill_indices = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int32, device=device) + prior_indices = torch.zeros(3, 2, dtype=torch.int32, device=device) + + # Production passes the device seq_lens twin so the row gather stays async. + top_k.update_gvr_prior_from_prefill( + prefill_indices, + torch.tensor([2, 1], dtype=torch.int32, device=device), + prior_indices, + request_offset=1, + ) + + assert prior_indices.tolist() == [[0, 0], [2, 3], [4, 5]] + assert top_k.needs_gvr_prior + + +@pytest.mark.parametrize( + "implementation,self_sampling", + [ + (TopKImplementation.CUDA_RADIX, False), + (TopKImplementation.CUTE_DSL_RADIX, False), + (TopKImplementation.CUTE_DSL_GVR, True), + (TopKImplementation.CUTE_DSL_GVR, False), + ], +) +@pytest.mark.parametrize("phase", ["prefill", "decode", "mixed", "split_prefill", "split_decode"]) +@pytest.mark.parametrize("next_n", [1, 4]) +def test_indexer_forward_uses_prior_only_for_temporal_gvr( + implementation: TopKImplementation, self_sampling: bool, phase: str, next_n: int +) -> None: + """Preserve prior ownership and pass the device lengths through each phase.""" + topk = 2 + top_k = TopK(topk, decode_implementation=implementation, gvr_self_sampling=self_sampling) + num_contexts = 0 if phase == "decode" else 2 + num_generations = 0 if phase == "prefill" else 2 + num_ctx_tokens = 5 if num_contexts else 0 + num_gen_tokens = num_generations * next_n + total_tokens = num_ctx_tokens + num_gen_tokens + is_generation = {"split_prefill": False, "split_decode": True}.get(phase) + has_prefill = num_contexts > 0 and is_generation is not True + has_decode = num_generations > 0 and is_generation is not False + input_tokens = ( + total_tokens + if is_generation is None + else (num_gen_tokens if is_generation else num_ctx_tokens) + ) + cache_manager = SimpleNamespace( + quant_block_size=128, + layer_offsets={7: 1}, + get_indexer_k_cache_buffers=Mock(return_value=torch.empty(0)), + ) + metadata = SimpleNamespace( + kv_cache_manager=cache_manager, + num_contexts=num_contexts, + num_generations=num_generations, + num_ctx_tokens=num_ctx_tokens, + num_tokens=total_tokens, + seq_lens=torch.tensor(([2, 3] if num_contexts else []) + [next_n] * num_generations), + cuda_graph_buffers={}, + is_cuda_graph=False, + get_empty=lambda buffers, shape, **kwargs: torch.full(shape, -1, dtype=torch.int32), + skip_indexer_for_ctx_reqs=False, + skip_indexer_for_gen_reqs=False, + indexer_prefill_chunks=None, + num_ctx_kv_tokens=16, + cu_seqlen_ks=torch.zeros(num_ctx_tokens, dtype=torch.int32), + cu_seqlen_ke=torch.full((num_ctx_tokens,), 16, dtype=torch.int32), + use_expanded_buffers_for_mtp=False, + kv_lens_cuda_2d=torch.full((num_generations, next_n), 16, dtype=torch.int32), + indexer_k_cache_block_offsets=torch.empty((num_contexts + num_generations, 1)), + scheduler_metadata_buffer=None, + scheduler_metadata_buffer_full_next_n=None, + get_indexer_max_seq_len=lambda: 16, + kv_lens_cuda_runtime=torch.full((num_contexts + num_generations,), 16), + gen_indexer_kv_lens_cuda_runtime=torch.full((num_generations,), 16), + kv_lens_row_reorder=None, + # The native CUDA Radix decode path asserts the caller-owned aux + # workspaces are present before handing them to TopK. + radix_aux_indices=torch.empty( + (num_generations or 1, _MAX_RADIX_BLOCKS_PER_ROW, topk), dtype=torch.int32 + ), + radix_aux_logits=torch.empty( + (num_generations or 1, _MAX_RADIX_BLOCKS_PER_ROW, topk), dtype=torch.float32 + ), + ) + # Keep distinct storage so the caller test catches use of the host twin. + metadata.seq_lens_cuda = metadata.seq_lens.clone() + # Radix fallback and self-sampling deliberately have no prior attribute. + # The raw heuristic option remains True, as in the SM107 failure. + if top_k.needs_gvr_prior: + metadata.gvr_prior_indices = torch.full( + (2, num_contexts + num_generations, topk), -99, dtype=torch.int32 + ) + + def logits(q: torch.Tensor, *args: object, **kwargs: object) -> torch.Tensor: + return torch.ones((q.shape[0], 16)) + + def selections(rows: int) -> torch.Tensor: + return torch.arange(rows * topk, dtype=torch.int32).reshape(rows, topk) + + def decode( + scores: torch.Tensor, + sequence_lengths: torch.Tensor, + scan_lengths: torch.Tensor, + output: torch.Tensor, + n: int, + max_seq_len: int, + extra: dict[str, torch.Tensor | None] | None, + radix_aux_indices: torch.Tensor | None, + radix_aux_logits: torch.Tensor | None, + ) -> torch.Tensor: + # Without a prior the indexer hands TopK no GVR kwargs at all, so the + # dict itself is absent rather than carrying a ``None`` entry. + prior = extra["gvr_prior_indices"] if extra is not None else None + if top_k.needs_gvr_prior: + assert prior.shape == (num_generations, topk) + assert prior.data_ptr() == metadata.gvr_prior_indices[1].data_ptr() + assert torch.all(prior == -99) + else: + assert prior is None + assert n == next_n + output.copy_(selections(num_gen_tokens)) + return output + + def prefill( + scores: torch.Tensor, + starts: torch.Tensor, + ends: torch.Tensor, + output: torch.Tensor, + k: int, + ) -> None: + output.copy_(selections(num_ctx_tokens)) + + indexer = SimpleNamespace( + top_k=top_k, + index_topk=topk, + layer_idx=7, + _enable_heuristic_topk=True, + # The GVR emission closed loop has its own coverage; this test only + # pins prior ownership, so keep emission off. + use_gvr_emission=False, + mtp_index_share=False, + use_fp4=False, + use_cute_dsl_paged_mqa_logits=False, + _call_mqa_logits=logits, + _call_paged_mqa_logits=Mock(return_value=torch.ones((num_gen_tokens, 16))), + aux_stream=None, + ) + with ( + patch.object( + top_k, "update_gvr_prior_from_prefill", wraps=top_k.update_gvr_prior_from_prefill + ) as seed_call, + patch.object(top_k, "_forward_decode", side_effect=decode) as decode_call, + patch.object(torch.ops.trtllm, "indexer_topk_prefill", side_effect=prefill), + ): + result = Indexer.sparse_attn_indexer( + indexer, + metadata, + torch.empty((input_tokens, 2)), + torch.empty((input_tokens, 1, 2)), + torch.empty((16, 2)), + torch.empty(16), + torch.empty(input_tokens), + is_generation=is_generation, + ) + assert decode_call.call_count == int(has_decode) + # The indexer seeds unconditionally after prefill; TopK itself is the one + # that drops the update when the implementation keeps no prior, so the call + # count tracks prefill alone and the prior contents assert the ownership. + assert seed_call.call_count == int(has_prefill) + if seed_call.called: + assert seed_call.call_args.args[1].data_ptr() == metadata.seq_lens_cuda.data_ptr() + token_offset = num_ctx_tokens if is_generation is None else 0 + if has_prefill: + torch.testing.assert_close(result[:num_ctx_tokens], selections(num_ctx_tokens)) + if has_decode: + torch.testing.assert_close(result[token_offset:], selections(num_gen_tokens)) + if top_k.needs_gvr_prior: + expected = torch.full_like(metadata.gvr_prior_indices, -99) + if has_prefill: + expected[1, num_generations:] = selections(num_ctx_tokens)[[1, 4]] + if has_decode: + expected[1, :num_generations] = selections(num_gen_tokens)[next_n - 1 :: next_n] + torch.testing.assert_close(metadata.gvr_prior_indices, expected) + + +@pytest.mark.parametrize("batch_size", [1, 8]) +@pytest.mark.parametrize("next_n", [1, 4]) +@pytest.mark.parametrize("score_width", [8192, 262144]) +def test_self_sampling_gpu_exactness_and_graph_replay( + batch_size: int, next_n: int, score_width: int +) -> None: + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + from tensorrt_llm._utils import get_sm_version + + if ( + not torch.cuda.is_available() + or not IS_CUTLASS_DSL_AVAILABLE + or get_sm_version() not in (100, 103, 107) + ): + pytest.skip("Self-sampling GVR requires datacenter Blackwell or Rubin and CuTe DSL") + + topk = 512 + compress_ratio = 4 + rows = batch_size * next_n + generator = torch.Generator(device="cuda").manual_seed(42) + scores = torch.randn((rows, score_width), generator=generator, device="cuda") + lengths = score_width * compress_ratio - 68 * torch.arange( + batch_size, dtype=torch.int32, device="cuda" + ) + row_ids = torch.arange(rows, device="cuda") + valid_columns = (lengths[row_ids // next_n] - next_n + row_ids % next_n + 1) // compress_ratio + valid = torch.arange(score_width, device="cuda")[None, :] < valid_columns[:, None] + # Invalid tail values must never enter the result. + scores.masked_fill_(~valid, 1.0e6) + output = torch.empty((rows, topk), dtype=torch.int32, device="cuda") + top_k = TopK( + topk, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + compress_ratio=compress_ratio, + gvr_self_sampling=True, + ) + assert not top_k.needs_gvr_prior + + def run() -> None: + top_k( + scores, + output, + is_prefill=False, + sequence_lengths=lengths, + scan_lengths=lengths // compress_ratio, + next_n=next_n, + max_seq_len=score_width, + ) + + def check_result() -> None: + assert torch.all((output >= 0) & (output < valid_columns[:, None])) + sorted_indices = output.sort(dim=-1).values + assert torch.all(sorted_indices[:, 1:] != sorted_indices[:, :-1]) + actual = scores.gather(1, output.long()).sort(dim=-1, descending=True).values + expected = scores.masked_fill(~valid, float("-inf")).topk(topk, dim=-1).values + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + run() + torch.cuda.synchronize() + check_result() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + output.fill_(-1) + graph.replay() + torch.cuda.synchronize() + check_result() diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_metadata_topk_init.py b/tests/unittest/_torch/attention/sparse/dsa/test_metadata_topk_init.py new file mode 100644 index 000000000000..df1200650c45 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/dsa/test_metadata_topk_init.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Regression coverage for TopK dispatch during DSA metadata initialization.""" + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.dsa import metadata as dsa_metadata +from tensorrt_llm._torch.attention.backends.sparse.dsa.params import DSAMetadataParams + + +@pytest.mark.parametrize("enable_heuristic", [False, True]) +@pytest.mark.parametrize("use_self_sampling", [False, True]) +@pytest.mark.parametrize("dsl_available", [False, True]) +@pytest.mark.parametrize("sm_version", [90, 100, 103, 107]) +def test_topk_flags_initialized_before_buffer_allocation( + enable_heuristic: bool, + use_self_sampling: bool, + dsl_available: bool, + sm_version: int, +) -> None: + metadata = object.__new__(dsa_metadata.DSAtrtllmAttentionMetadata) + metadata.sparse_metadata_params = DSAMetadataParams( + indexer_max_chunk_size=8192, + max_sparse_topk=512, + index_head_dim=128, + enable_indexer_skip=False, + enable_heuristic_topk=enable_heuristic, + use_cute_dsl_topk=True, + use_cute_dsl_paged_mqa_logits=False, + q_split_threshold=8192, + use_self_sampling_topk=use_self_sampling, + ) + metadata.kv_cache_manager = SimpleNamespace( + tokens_per_block=128, + compressed_block_sizes={}, + get_cache_indices=Mock(), + ) + metadata.is_cuda_graph = False + enabled = enable_heuristic and sm_version >= 100 + self_sampling_supported = enabled and dsl_available and sm_version in (100, 103) + temporal_supported = enabled and dsl_available and sm_version in (100, 103) + + def check_flags(*, capture_graph: bool) -> None: + assert capture_graph is False + assert metadata.enable_gvr_topk is enabled + assert metadata.use_self_sampling_topk is (self_sampling_supported and use_self_sampling) + assert metadata.needs_gvr_prior is (temporal_supported and not use_self_sampling) + + metadata.create_buffers_for_mla_rope_append = Mock(side_effect=check_flags) + metadata.create_buffers_for_indexer = Mock(side_effect=check_flags) + + with ( + patch.object(dsa_metadata.TrtllmAttentionMetadata, "__post_init__"), + patch.object(dsa_metadata, "IS_CUTLASS_DSL_AVAILABLE", dsl_available), + patch.object(dsa_metadata, "get_sm_version", return_value=sm_version), + ): + metadata.__post_init__() + + metadata.create_buffers_for_mla_rope_append.assert_called_once_with(capture_graph=False) + metadata.create_buffers_for_indexer.assert_called_once_with(capture_graph=False) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("num_local_layers", [1, 3]) +def test_temporal_gvr_allocates_real_prior_buffers(num_local_layers: int) -> None: + """Allocate actual CUDA buffers and size the zeroed prior by local layers.""" + metadata = object.__new__(dsa_metadata.DSAtrtllmAttentionMetadata) + metadata.sparse_metadata_params = DSAMetadataParams( + indexer_max_chunk_size=32, + max_sparse_topk=512, + index_head_dim=128, + enable_indexer_skip=False, + enable_heuristic_topk=True, + use_cute_dsl_topk=True, + use_cute_dsl_paged_mqa_logits=False, + q_split_threshold=32, + use_self_sampling_topk=False, + ) + metadata.kv_cache_manager = SimpleNamespace( + tokens_per_block=128, + compressed_block_sizes={}, + get_cache_indices=Mock(), + max_blocks_per_seq=2, + num_local_layers=num_local_layers, + ) + metadata.draft_kv_cache_manager = None + metadata.is_cuda_graph = False + metadata.cuda_graph_buffers = None + metadata.max_num_sequences = 4 + metadata.max_num_tokens = 16 + metadata.max_draft_tokens = 3 + metadata.num_sms = 16 + metadata.enable_context_mla_with_cached_kv = False + # object.__new__ skips __init__, which is where this default is set. + metadata._radix_rows_per_sequence = 1 + with ( + patch.object(dsa_metadata.TrtllmAttentionMetadata, "__post_init__"), + patch.object(metadata, "create_buffers_for_mla_rope_append"), + patch.object(dsa_metadata, "IS_CUTLASS_DSL_AVAILABLE", True), + patch.object(dsa_metadata, "get_sm_version", return_value=100), + ): + metadata.__post_init__() + + assert metadata.enable_gvr_topk + assert metadata.needs_gvr_prior + assert metadata.gvr_prior_indices.shape == (num_local_layers, 4, 512) + assert metadata.gvr_prior_indices.dtype == torch.int32 + assert metadata.gvr_prior_indices.is_cuda + assert torch.count_nonzero(metadata.gvr_prior_indices).item() == 0 + assert metadata.kv_lens_row_reorder_buffer.shape == (4,) diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py index 642a41010074..9f52acd5bd04 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py @@ -2403,7 +2403,9 @@ def yarn_get_mscale(scale=1, mscale=1): @pytest.mark.skipif(get_sm_version() < 100, reason="DSv4 fused FP8 Q-quant requires SM100") @pytest.mark.parametrize("batch_name", _FUSED_Q_FP8_PREFILL_BATCHES) -def test_forward_sparse_mla_unified_fused_q_fp8(monkeypatch, batch_name): +@pytest.mark.parametrize("q_rope_applied", [False, True]) +def test_forward_sparse_mla_unified_fused_q_fp8(monkeypatch, batch_name, + q_rope_applied): """Regression test for the sparse-MLA context-branch fused FP8 Q-quant wiring. The wo-linear helper bypasses `_q_branch`, so we monkey-patch `forward_context_sparse_attn` to manually populate the fused buffers @@ -2430,16 +2432,41 @@ def patched_fwd(self, q, compressed_kv, k_pe, attn_metadata, output, quant_q_buffer = torch.empty((num_tokens, num_heads * head_dim), dtype=torch.float8_e4m3fn, device=q.device) - quant_q_buffer.view( - num_tokens, num_heads, - head_dim)[:, :, :nope_dim] = q_view[:, :, :nope_dim].float().to( + quant_q_view = quant_q_buffer.view(num_tokens, num_heads, head_dim) + quant_q_view[:, :, :nope_dim] = q_view[:, :, :nope_dim].float().to( + torch.float8_e4m3fn) + if q_rope_applied: + position_ids = kwargs["position_ids"].reshape(-1).long() + rope_dim = head_dim - nope_dim + cache = self.mqa.rotary_cos_sin.view(-1, rope_dim, 2) + coefficient = cache[position_ids, :rope_dim // 2].unsqueeze(1) + q_pe = q_view[:, :, nope_dim:].float() + q0, q1 = q_pe[..., 0::2], q_pe[..., 1::2] + cos, sin = coefficient[..., 0], coefficient[..., 1] + rotated = torch.stack( + [cos * q0 - sin * q1, cos * q1 + sin * q0], + dim=-1).flatten(-2) + quant_q_view[:, :, nope_dim:] = rotated.to(q.dtype).to( torch.float8_e4m3fn) + self._fused_q_pe = None + else: + self._fused_q_pe = q_view[:, :, nope_dim:].contiguous() self._fused_quant_q_buffer = quant_q_buffer - self._fused_q_pe = q_view[:, :, nope_dim:].contiguous() + self._fused_q_rope_applied = q_rope_applied self._quant_scale_qkv = torch.tensor([1.0], dtype=torch.float32, device=q.device) - q = torch.full_like(q, float('nan')) + if q_rope_applied: + # q_b GEMM fusion carries only compact q_lora once the complete + # rotated Q has been written to quant_q_buffer. + q = torch.full( + (num_tokens, self.q_lora_rank), + float("nan"), + dtype=q.dtype, + device=q.device, + ) + else: + q = torch.full_like(q, float("nan")) return original_fwd(self, q, compressed_kv, k_pe, attn_metadata, output, **kwargs) diff --git a/tests/unittest/_torch/executor/test_indexer_workspace_reserve.py b/tests/unittest/_torch/executor/test_indexer_workspace_reserve.py new file mode 100644 index 000000000000..06d238c53776 --- /dev/null +++ b/tests/unittest/_torch/executor/test_indexer_workspace_reserve.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from tensorrt_llm._torch.attention.backends.sparse.params import ( + get_indexer_mqa_logits_elem_budget, + get_indexer_mqa_logits_workspace_bytes, +) + + +def test_workspace_bytes_match_runtime_element_budget(monkeypatch): + monkeypatch.setenv("TLLM_INDEXER_MQA_LOGITS_ELEM_BUDGET", "1024") + assert get_indexer_mqa_logits_elem_budget() == 1024 + assert get_indexer_mqa_logits_workspace_bytes() == 4096 + + +def test_workspace_bytes_are_bounded_by_reachable_request_shape(monkeypatch): + monkeypatch.setenv("TLLM_INDEXER_MQA_LOGITS_ELEM_BUDGET", str(1 << 31)) + assert get_indexer_mqa_logits_workspace_bytes(4096, 4096) == 64 * 1024 * 1024 diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 8c253ba15026..f940d5fe7910 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -697,6 +697,7 @@ def fake_decoder_layer_init(self, model_config, *_args, **_kwargs): pretrained_config=config, mapping=Mapping(world_size=4, rank=2, tp_size=4), quant_config=quant_config, + use_cute_dsl_blockscaling_mm=True, ) mtp_layer = DeepseekV4MTP( @@ -715,6 +716,8 @@ def fake_decoder_layer_init(self, model_config, *_args, **_kwargs): assert mtp_layer.h_proj.out_features == config.hidden_size assert mtp_layer.e_proj.reduce_output is True assert mtp_layer.h_proj.reduce_output is True + assert mtp_layer.e_proj.use_cute_dsl_blockscaling_mm is True + assert mtp_layer.h_proj.use_cute_dsl_blockscaling_mm is True assert mtp_layer.e_proj.weight.dtype is torch.float8_e4m3fn assert mtp_layer.h_proj.weight.dtype is torch.float8_e4m3fn assert hasattr(mtp_layer.e_proj, "weight_scale") diff --git a/tests/unittest/_torch/modules/test_gated_mlp.py b/tests/unittest/_torch/modules/test_gated_mlp.py new file mode 100644 index 000000000000..da0b31ee3b1f --- /dev/null +++ b/tests/unittest/_torch/modules/test_gated_mlp.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from typing import Callable +from unittest.mock import Mock + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from tensorrt_llm._torch.modules import gated_mlp as gated_mlp_module +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP + + +def _make_gate_up_proj( + projected: torch.Tensor, + *, + partitioned: bool, +) -> nn.Module: + gate_up_proj = nn.Module() + gate_up_proj.has_nvfp4 = True + gate_up_proj.has_bias = False + gate_up_proj.partition_plan = SimpleNamespace(enabled=partitioned) + gate_up_proj.can_use_cute_dsl_nvfp4_swiglu_blackwell = Mock(return_value=not partitioned) + gate_up_proj.forward = Mock(return_value=projected) + return gate_up_proj + + +def _make_down_proj() -> nn.Module: + down_proj = nn.Module() + down_proj.has_fp8_qdq = False + down_proj.has_w4a8_nvfp4_fp8 = False + # The unfused branch calls ``_can_fuse_swiglu_fp8_quant``, whose first + # check reads this attribute, so a stand-in without it raises + # AttributeError before the assertions below are reached. + down_proj.has_fp8_block_scales = False + down_proj.forward = Mock(side_effect=lambda value, **kwargs: value + 1) + return down_proj + + +def test_gate_up_partition_falls_back_to_swiglu( + monkeypatch: pytest.MonkeyPatch, +) -> None: + mlp = GatedMLP(hidden_size=2, intermediate_size=2, bias=False) + mlp.use_cute_dsl_blockscaling_mm = True + projected = torch.tensor([[1.0, 2.0, 3.0, 4.0]]) + mlp.gate_up_proj = _make_gate_up_proj(projected, partitioned=True) + mlp.down_proj = _make_down_proj() + + gate, up = projected.chunk(2, dim=-1) + activated = F.silu(gate) * up + swiglu = Mock(side_effect=lambda value, **kwargs: F.silu(value[..., :2]) * value[..., 2:]) + monkeypatch.setattr(gated_mlp_module, "swiglu", swiglu) + fused_gate_up_swiglu = Mock(side_effect=AssertionError("Blackwell fused op must not run")) + monkeypatch.setattr(mlp, "_fused_gate_up_swiglu", fused_gate_up_swiglu) + + inputs = torch.tensor([[5.0, 6.0]]) + output = mlp(inputs) + + assert not mlp._can_fuse_gate_up_swiglu() + mlp.gate_up_proj.forward.assert_called_once_with(inputs) + # ``GatedMLP`` forwards all three SwiGLU shape parameters on every call; + # they are None here because this layer is plain SwiGLU. + swiglu.assert_called_once_with( + projected, swiglu_limit=None, swiglu_alpha=None, swiglu_beta=None + ) + mlp.down_proj.forward.assert_called_once() + down_args, down_kwargs = mlp.down_proj.forward.call_args + torch.testing.assert_close(down_args[0], activated) + assert down_kwargs == {"all_reduce_params": None, "layer_idx": None} + fused_gate_up_swiglu.assert_not_called() + torch.testing.assert_close(output, activated + 1) + + +@pytest.mark.parametrize( + "swiglu_limit, expected", + [(None, True), (float("inf"), True), (7.0, False)], +) +def test_swiglu_limit_controls_gate_up_fusion_capability( + swiglu_limit: float | None, + expected: bool, +) -> None: + mlp = GatedMLP( + hidden_size=2, + intermediate_size=2, + bias=False, + swiglu_limit=swiglu_limit, + use_cute_dsl_blockscaling_mm=True, + ) + + assert mlp.gate_up_proj.use_cute_dsl_nvfp4_swiglu_blackwell is expected + + +@pytest.mark.parametrize( + ("activation", "expected"), + [ + pytest.param(F.silu, True, id="plain-swiglu"), + pytest.param(lambda value: value, False, id="custom-swiglu-oai"), + ], +) +def test_activation_controls_fp8_quant_fusion_capability( + monkeypatch: pytest.MonkeyPatch, + activation: Callable[[torch.Tensor], torch.Tensor], + expected: bool, +) -> None: + mlp = GatedMLP( + hidden_size=2, + intermediate_size=2, + bias=False, + activation=activation, + ) + down_proj = nn.Module() + down_proj.has_fp8_block_scales = True + down_proj.use_cute_dsl_blockscaling_mm = True + down_proj.disable_deep_gemm = False + mlp.down_proj = down_proj + monkeypatch.setattr(gated_mlp_module, "get_sm_version", lambda: 107) + monkeypatch.setattr(gated_mlp_module, "IS_CUTLASS_DSL_RUBIN_AVAILABLE", True) + + assert mlp._can_fuse_swiglu_fp8_quant() is expected 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 69600d69f043..718e4b983613 100644 --- a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py +++ b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py @@ -23,6 +23,7 @@ from tensorrt_llm._torch.autotuner import AutoTuner, OptimizationProfile, TunableRunner from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + SITU_BETA_DISABLED, GroupedGemmInputsHelper, _get_sm107_nvfp4_default_mma_config, ) @@ -64,6 +65,31 @@ def swiglu_ref(x: torch.Tensor, swiglu_limit: float = float("inf")) -> torch.Ten return x * torch.nn.functional.silu(gate) +# Kimi K3 SiTU constants; deliberately non-unit and unequal, so that a swapped +# beta / linear_beta or a dropped constant cannot pass the test. +SITU_BETA = 2.5 +SITU_LINEAR_BETA = 0.7 + + +def situ_ref( + x: torch.Tensor, beta: float = SITU_BETA, linear_beta: float = SITU_LINEAR_BETA +) -> torch.Tensor: + """SiTU reference, in this kernel's ``[linear | gate]`` chunk order. + + ``tensorrt_llm._torch.modules.kimi_k3_moe._mlp.SituAndMul`` is the golden + definition, but it consumes the model's ``[gate | up]`` packing -- the two + halves in the opposite order from ``swiglu_ref`` and from the interleaved + FC1 weight this kernel reads. Keep the order explicit: silently swapping + the halves still produces plausible-looking numbers. + """ + up, gate = x.chunk(2, dim=-1) + gate = gate.to(torch.float32) + up = up.to(torch.float32) + situ_a = beta * torch.tanh(gate / beta) * torch.sigmoid(gate) + up = linear_beta * torch.tanh(up / linear_beta) + return (situ_a * up).to(x.dtype) + + def apply_activation_ref( x: torch.Tensor, activation_type: ActivationType, swiglu_limit: float = float("inf") ) -> torch.Tensor: @@ -71,6 +97,8 @@ def apply_activation_ref( return swiglu_ref(x, swiglu_limit) if activation_type == ActivationType.Relu2: return relu2(x) + if activation_type == ActivationType.SiTu: + return situ_ref(x) raise ValueError(f"Unsupported activation_type: {activation_type}") @@ -1136,8 +1164,8 @@ def test_nvfp4_grouped_gemm_swiglu_blackwell( @pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="cutlass-dsl is not available") @pytest.mark.parametrize( "activation_type", - [ActivationType.Swiglu, ActivationType.Relu2], - ids=["swiglu", "relu2"], + [ActivationType.Swiglu, ActivationType.Relu2, ActivationType.SiTu], + ids=["swiglu", "relu2", "situ"], ) @pytest.mark.parametrize("tile_size", [128, 256]) @pytest.mark.parametrize("ep_size", [1, 8, 32]) @@ -1155,11 +1183,14 @@ def test_nvfp4_gather_grouped_gemm_act_fusion_blackwell( This test validates the gather kernel which: 1. Uses LDGSTS for A/SFA loading with permuted_idx_to_expanded_idx 2. Performs GEMM with (interleaved for gated) weights - 3. Applies the fused activation (SwiGLU for gated, Relu2 for non-gated) + 3. Applies the fused activation (SwiGLU / SiTU for gated, Relu2 for + non-gated) 4. Quantizes output to FP4 with scale factor generation """ is_gated = is_gated_activation(activation_type) - swiglu_limit = 1.0 if is_gated else float("inf") + is_situ = activation_type == ActivationType.SiTu + # SiTU rejects the SwiGLU clamp (matching MegaMoE and DeepGEMM). + swiglu_limit = 1.0 if (is_gated and not is_situ) else float("inf") weight_n_multiplier = 2 if is_gated else 1 sf_vec_size = 16 hidden_size = 4096 @@ -1196,16 +1227,17 @@ def test_nvfp4_gather_grouped_gemm_act_fusion_blackwell( num_valid_permuted_tokens = total_num_padded_tokens.item() # Create input tensors (original size, not permuted) - a = torch.randint(-5, 5, (num_tokens, hidden_size), dtype=torch.int32, device="cuda").to( - torch.bfloat16 - ) + # Draw straight into bfloat16: the int32 staging buffer for `b` is twice + # the size of the tensor it produces, and at ep_size=1 the pair is ~100 GB + # live at once, which is what makes this grid OOM. + a = torch.randint(-5, 5, (num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda") b = torch.randint( -5, 5, (num_local_experts, interm_size * weight_n_multiplier, hidden_size), - dtype=torch.int32, + dtype=torch.bfloat16, device="cuda", - ).to(torch.bfloat16) + ) # Quantize inputs to FP4 a_global_sf = a.abs().max().float() / (448 * 6) @@ -1301,6 +1333,8 @@ def test_nvfp4_gather_grouped_gemm_act_fusion_blackwell( scaling_vector_size=sf_vec_size, activation_type=activation_type, swiglu_limit_scalar=swiglu_limit, + situ_beta=SITU_BETA if is_situ else -1.0, + situ_linear_beta=SITU_LINEAR_BETA if is_situ else -1.0, ) # Verify output (only compare valid tokens, skip padding tokens where permuted_idx_to_expanded_idx == -1) @@ -1538,6 +1572,203 @@ def test_nvfp4_gather_grouped_gemm_act_fusion_rubin( check_accuracy(c_sf_valid, c_sf_ref_valid, atol=1e-4, rtol=1e-4, percent=0.95) +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on SM 107 (Rubin) GPUs", +) +@pytest.mark.parametrize("tile_size", [128, 256]) +def test_nvfp4_gather_grouped_gemm_situ_rubin(tile_size: int): + """SiTU epilogue on the Rubin NVFP4 gather+grouped GEMM kernel. + + Reduced grid vs the SwiGLU test: enough to compile both tile sizes and + catch a dropped / swapped beta. The Blackwell test documents why the + swap check is restricted to (128 tokens, top_k=1, ep=1). + """ + num_tokens = 128 + top_k = 1 + ep_size = 1 + sf_vec_size = 16 + hidden_size = 4096 + interm_size = 8192 + num_experts = 256 + num_local_experts = num_experts // ep_size + + torch.manual_seed(20260727 + tile_size) + routing_logits = torch.randn(num_tokens, num_experts, device="cuda") + token_final_scales, token_selected_experts = routing_logits.topk(top_k, dim=-1) + token_selected_experts = token_selected_experts.to(torch.int32) + token_final_scales = token_final_scales.softmax(dim=-1).to(torch.float32) + token_selected_experts[0] = 0 + + ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles, + ) = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=0, + local_num_experts=num_local_experts, + tile_tokens_dim=tile_size, + ) + + max_num_permuted_tokens = permuted_idx_to_expanded_idx.size(0) + num_valid_permuted_tokens = total_num_padded_tokens.item() + + a = torch.randint(-5, 5, (num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda") + b = torch.randint( + -5, + 5, + (num_local_experts, interm_size * 2, hidden_size), + dtype=torch.bfloat16, + device="cuda", + ) + + a_global_sf = a.abs().max().float() / (448 * 6) + b_global_sf = b.abs().amax(dim=(1, 2)).float() / (448 * 6) + a, a_sf = torch.ops.trtllm.fp4_quantize(a, 1 / a_global_sf, sf_vec_size, False) + a = a.view(torch.float4_e2m1fn_x2) + a_sf_unswizzled = unswizzle_sf(a_sf, (num_tokens + 127) // 128 * 128, hidden_size)[:num_tokens] + b, b_sf = torch.ops.trtllm.fp4_quantize(b, 1 / b_global_sf, sf_vec_size, False) + b = b.view(torch.float4_e2m1fn_x2) + b_sf = b_sf.view(num_local_experts, interm_size * 2, hidden_size // sf_vec_size) + alpha = a_global_sf * b_global_sf + + b_interleaved = interleave_linear_and_gate(b.view(torch.uint8), group_size=64, dim=1).view( + torch.float4_e2m1fn_x2 + ) + b_sf_unswizzled = unswizzle_sf(b_sf, interm_size * 2, hidden_size).view( + num_local_experts, interm_size * 2, hidden_size // sf_vec_size + ) + b_sf_unswizzled_interleaved = interleave_linear_and_gate(b_sf_unswizzled, group_size=64, dim=1) + b_sf_interleaved = swizzle_sf(b_sf_unswizzled_interleaved, interm_size * 2, hidden_size).view( + num_local_experts, interm_size * 2, hidden_size // sf_vec_size + ) + + permuted_idx_to_expanded_idx_list = permuted_idx_to_expanded_idx.cpu().tolist() + tile_idx_to_mn_limit_list = tile_idx_to_mn_limit.cpu().tolist() + + a_gathered = torch.empty(max_num_permuted_tokens, hidden_size // 2, dtype=a.dtype) + a_sf_gathered = torch.empty( + max_num_permuted_tokens, hidden_size // sf_vec_size, dtype=a_sf.dtype + ) + for i in range(num_valid_permuted_tokens): + if i >= tile_idx_to_mn_limit_list[i // tile_size]: + continue + expanded_idx = permuted_idx_to_expanded_idx_list[i] + token_id = expanded_idx // top_k + a_gathered[i] = a[token_id] + a_sf_gathered[i] = a_sf_unswizzled[token_id] + a_gathered = a_gathered.to(a.device) + a_sf_gathered = a_sf_gathered.to(a.device) + + a_sf_gathered_swizzled = swizzle_sf( + a_sf_gathered.view(max_num_permuted_tokens, hidden_size // sf_vec_size), + max_num_permuted_tokens, + hidden_size, + ) + + c_ref = cute_dsl_nvfp4_grouped_gemm_ref( + a_gathered, + b, + a_sf_gathered_swizzled, + b_sf, + alpha, + tile_idx_to_group_idx, + num_non_exiting_tiles, + tile_size=tile_size, + output_dtype=torch.bfloat16, + scaling_vector_size=sf_vec_size, + ) + c_ref = situ_ref(c_ref) + global_sf = c_ref[:num_valid_permuted_tokens].abs().max().float() / (448 * 6) + c_ref, c_sf_ref = torch.ops.trtllm.fp4_quantize(c_ref, 1 / global_sf, sf_vec_size, False) + + op_kwargs = dict( + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=0, + tile_size=tile_size, + output_tensor=None, + output_sf_tensor=None, + scaling_vector_size=sf_vec_size, + activation_type=ActivationType.SiTu, + situ_beta=SITU_BETA, + situ_linear_beta=SITU_LINEAR_BETA, + ) + c, c_sf = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( + a, + b_interleaved, + a_sf_unswizzled, + b_sf_interleaved, + alpha, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + torch.tensor([1 / global_sf], dtype=torch.float32, device="cuda"), + **op_kwargs, + ) + + valid_token_mask = torch.zeros(num_valid_permuted_tokens, dtype=torch.bool, device="cuda") + for i in range(num_valid_permuted_tokens): + if i >= tile_idx_to_mn_limit_list[i // tile_size]: + continue + valid_token_mask[i] = True + + num_valid_tokens = valid_token_mask.sum().item() + if num_valid_tokens > 0: + c_valid = c[:num_valid_permuted_tokens].view(torch.uint8)[valid_token_mask] + c_ref_valid = c_ref[:num_valid_permuted_tokens][valid_token_mask] + check_accuracy(c_valid, c_ref_valid, atol=1e-4, rtol=1e-4, percent=0.95) + + c_sf_unswizzled = unswizzle_sf(c_sf, max_num_permuted_tokens, interm_size, sf_vec_size) + c_sf_ref_unswizzled = unswizzle_sf( + c_sf_ref, max_num_permuted_tokens, interm_size, sf_vec_size + ) + + c_sf_valid = [] + c_sf_ref_valid = [] + for i in range(num_valid_permuted_tokens): + if i >= tile_idx_to_mn_limit_list[i // tile_size]: + continue + c_sf_valid.append(c_sf_unswizzled[i]) + c_sf_ref_valid.append(c_sf_ref_unswizzled[i]) + + c_sf_valid = torch.cat(c_sf_valid) + c_sf_ref_valid = torch.cat(c_sf_ref_valid) + check_accuracy(c_sf_valid, c_sf_ref_valid, atol=1e-4, rtol=1e-4, percent=0.95) + + c_swapped, _ = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( + a, + b_interleaved, + a_sf_unswizzled, + b_sf_interleaved, + alpha, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + torch.tensor([1 / global_sf], dtype=torch.float32, device="cuda"), + **{ + **op_kwargs, + "situ_beta": SITU_LINEAR_BETA, + "situ_linear_beta": SITU_BETA, + }, + ) + swapped_valid = c_swapped[:num_valid_permuted_tokens].view(torch.uint8)[valid_token_mask] + assert not torch.equal(c_valid, swapped_valid), ( + "Swapping situ_beta and situ_linear_beta produced a bit-identical " + "output; the kernel is ignoring one of the SiTU constants." + ) + + @pytest.mark.skipif( get_sm_version() != 107, reason="This test is only supported on SM 107 (Rubin) GPUs", @@ -2645,12 +2876,16 @@ def _assert_rubin_moe_op_schema( "scaling_vector_size", "partition_id", "activation_type", + "situ_beta", + "situ_linear_beta", "precomputed_tactic", ), { "scaling_vector_size": 16, "partition_id": -1, "activation_type": int(ActivationType.Swiglu), + "situ_beta": SITU_BETA_DISABLED, + "situ_linear_beta": SITU_BETA_DISABLED, "precomputed_tactic": None, }, {"output_tensor", "output_sf_tensor"}, @@ -2917,10 +3152,14 @@ def test_rubin_bf16_moe_precomputed_tactic_fake_signatures(): "output_sf_tensor", "scaling_vector_size", "activation_type", + "situ_beta", + "situ_linear_beta", ), { "scaling_vector_size": 16, "activation_type": int(ActivationType.Swiglu), + "situ_beta": SITU_BETA_DISABLED, + "situ_linear_beta": SITU_BETA_DISABLED, }, {"output_tensor", "output_sf_tensor"}, id="nvfp4_fc1", @@ -3374,7 +3613,12 @@ def fake_moe_output_memset(*args, **kwargs): runner_args, runner_kwargs = runner_instances[0].init_call if quantized and is_fc1: assert runner_args == (1, 1, 1, 0, 128, 16) - assert runner_kwargs == {"activation_type": ActivationType.Swiglu} + # The disabled sentinel canonicalizes to None before the runner sees it. + assert runner_kwargs == { + "activation_type": ActivationType.Swiglu, + "situ_beta": None, + "situ_linear_beta": None, + } elif quantized: assert runner_args == (1, 1, 1, 0, 128, torch.bfloat16, 16) assert not runner_kwargs diff --git a/tests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py b/tests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py index 41b494e355ae..667fad95e691 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py +++ b/tests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py @@ -15,6 +15,7 @@ import os import subprocess import sys +from unittest import mock import pytest import torch @@ -24,6 +25,7 @@ import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils from tensorrt_llm._torch.autotuner import AutoTuner, autotune +from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops from tensorrt_llm._torch.cute_dsl_utils import (IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_RUBIN_AVAILABLE) @@ -470,15 +472,57 @@ def test_cute_dsl_mxfp8_gemm_rubin_k128_replicated_scales(): output = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin( a_fp8, b_fp8, a_sf, b_sf) - output = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin(a_fp8, b_fp8, a_sf, - b_sf) expected = a @ b.t() + alpha = cute_dsl_custom_ops._get_mxfp8_gemm_alpha(a.device) + seen_alphas = [] + real_get_alpha = cute_dsl_custom_ops._get_mxfp8_gemm_alpha + + def _get_alpha_spy(device): + cached = real_get_alpha(device) + seen_alphas.append(cached) + return cached + + # After the one-time warmup above, the steady-state call must reuse the + # cached scalar instead of filling a fresh one. + with mock.patch.object( + cute_dsl_custom_ops, + "_get_mxfp8_gemm_alpha", + side_effect=_get_alpha_spy, + ), mock.patch.object( + torch, + "ones", + side_effect=AssertionError("unexpected per-call fill"), + ): + output = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin( + a_fp8, b_fp8, a_sf, b_sf) + + assert len(seen_alphas) == 1 + assert all(seen is alpha for seen in seen_alphas) + assert {seen.data_ptr() for seen in seen_alphas} == {alpha.data_ptr()} diff = calc_diff(output, expected) assert diff < 1e-3 torch.testing.assert_close(output, expected, atol=1e-3, rtol=1e-3) +def test_mxfp8_alpha_cache_rejects_first_init_during_capture(): + """A capture cannot perform the per-device cache's one-time allocation.""" + device = torch.device("cuda", 31415) + assert device not in cute_dsl_custom_ops._MXFP8_GEMM_ALPHA_CACHE + + with mock.patch.object(torch.cuda, + "device") as device_guard, mock.patch.object( + torch.cuda, + "is_current_stream_capturing", + return_value=True): + with pytest.raises(RuntimeError, + match="run one eager GEMM warmup first"): + cute_dsl_custom_ops._get_mxfp8_gemm_alpha(device) + + device_guard.assert_called_once_with(device) + assert device not in cute_dsl_custom_ops._MXFP8_GEMM_ALPHA_CACHE + + @pytest.mark.skipif( getSMVersion() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, reason="The test requires SM107 and SM107 CuTe DSL support.", 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 new file mode 100644 index 000000000000..838b7d20123f --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_kimi_k3_fp8_weight_read_linear.py @@ -0,0 +1,247 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Coverage for Kimi K3's CuTe-only FP8 weight-read linear. + +This module is hand-built rather than a ``Linear`` + ``FP8BlockScalesLinearMethod``, +so the shipping FP8-block-scale tests never reach it. These tests pin the +single-scale loader contract and the specialized quant + CuTe path. +""" + +import pytest +import torch +from _torch.helpers import calc_diff, per_block_cast_to_fp8 +from utils.util import getSMVersion + +from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + _get_kimi_k3_mxfp8_tuning_buckets, + _kimi_k3_mxfp8_tuning_bucket, +) +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE +from tensorrt_llm._torch.models.modeling_kimi_linear import ( + _Fp8BlockScaleWeightReadLinear as K3Fp8Linear, +) + +# K3-representative (out, in) projections, all 128-aligned. +SHAPES = [(512, 1024), (2048, 1024)] +MS = [1, 32, 128] + +# The Rubin MXFP8 weight-read path needs both SM107 and the internal CuTe DSL +# build; keep the predicate local so the gate does not depend on a helper that +# lives outside this change. +RUBIN = pytest.mark.skipif( + getSMVersion() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, + reason="needs SM107 with Rubin CuTe DSL support", +) + + +def _ref(x, w): + return (x.float() @ w.float().t()).to(x.dtype) + + +def _check(out, expected, tag): + assert out.dtype == expected.dtype, tag + assert out.shape == expected.shape, tag + assert torch.isfinite(out).all(), f"{tag}: non-finite output" + diff = calc_diff(out, expected) + assert diff < 5e-3, f"{tag}: calc_diff={diff}" + + +def _make(out_features, in_features, seed=0): + torch.random.manual_seed(seed) + w = ( + torch.randn((out_features, in_features), device="cuda", dtype=torch.bfloat16) + / in_features**0.5 + ) + weight, weight_scale = K3Fp8Linear.quantize_weight(w) + return w, K3Fp8Linear(weight, weight_scale, out_features) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@RUBIN +@pytest.mark.parametrize("out_features, in_features", SHAPES) +@pytest.mark.parametrize("m", MS) +def test_forward_matches_bf16(m, out_features, in_features): + """The CuTe-only path must preserve the FP8 block-scale numerics.""" + w, mod = _make(out_features, in_features) + x = torch.randn((m, in_features), device="cuda", dtype=torch.bfloat16) + out = mod(x) + _check(out, _ref(x, w), f"m={m} n={out_features} k={in_features}") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@RUBIN +def test_forward_prequantized_matches_bf16(): + """D-Spark's fused KDA output must use the same CuTe scale ABI.""" + w, mod = _make(512, 1024, seed=9) + x = torch.randn((16, 1024), device="cuda", dtype=torch.bfloat16) + activation, activation_scale = torch.ops.trtllm.fp8_quantize_1x128_cutedsl_ue8m0(x) + + out = mod.forward_prequantized(activation, activation_scale) + + _check(out, _ref(x, w), "prequantized") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@RUBIN +def test_forward_uses_specialized_quant(): + """Every call must use the specialized quantizer and fine-M runner.""" + _, mod = _make(512, 1024, seed=8) + x = torch.randn((64, 1024), device="cuda", dtype=torch.bfloat16) + mod(x) + + quant_calls = [] + fine_grained_m = [] + real_quant = torch.ops.trtllm.fp8_quantize_1x128_cutedsl_ue8m0 + real_gemm = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin + + class _QuantSpy: + def __call__(self, *args, **kwargs): + quant_calls.append(args[0].shape[0]) + return real_quant(*args, **kwargs) + + class _GemmSpy: + def __call__(self, *args, **kwargs): + fine_grained_m.append(kwargs["fine_grained_m"]) + return real_gemm(*args, **kwargs) + + torch.ops.trtllm.fp8_quantize_1x128_cutedsl_ue8m0 = _QuantSpy() + torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin = _GemmSpy() + try: + mod(x) + mod(x) + finally: + torch.ops.trtllm.fp8_quantize_1x128_cutedsl_ue8m0 = real_quant + torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin = real_gemm + + assert quant_calls == [64, 64] + assert fine_grained_m == [True, True] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@RUBIN +def test_single_cute_scale_is_built_on_rubin(): + """The CuTe layout must be prepared at load time, not during capture.""" + _, mod = _make(2048, 1024, seed=5) + assert mod.weight_scale.numel() > 0 + assert mod.weight_scale.dtype is torch.uint8 + assert not hasattr(mod, "weight_scale_mx") + assert not hasattr(mod, "gemm_alpha") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@RUBIN +def test_placeholder_load_carries_cute_scale(): + """A loader-filled placeholder must receive the CuTe scale layout.""" + in_features, parts = 1024, [512, 512] + torch.random.manual_seed(6) + pairs = [] + 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) + pairs.append((q, s.float())) + + mod = K3Fp8Linear.empty_placeholder(sum(parts), in_features) + mod.load_checkpoint_pair(pairs) + assert not mod.is_placeholder + 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") +def test_weight_preparation_returns_only_cute_pair(): + """Both construction routes return exactly FP8 weight + CuTe scale.""" + torch.random.manual_seed(7) + w = torch.randn((512, 1024), device="cuda", dtype=torch.bfloat16) / 32 + assert len(K3Fp8Linear.quantize_weight(w)) == 2 + + q, s = per_block_cast_to_fp8(w) + assert len(K3Fp8Linear.prepare_checkpoint_scale(q, s.float())) == 2 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_unfilled_placeholder_raises(): + mod = K3Fp8Linear.empty_placeholder(256, 256) + x = torch.randn((2, 256), device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="never filled"): + mod(x) + + +def test_kimi_k3_fine_m_tuning_uses_hybrid_buckets(): + low_m = (1, 2, 4, 8, *range(16, 193, 16)) + assert _get_kimi_k3_mxfp8_tuning_buckets(17) == low_m[:5] + assert _get_kimi_k3_mxfp8_tuning_buckets(65) == (*low_m[:8], 80) + assert _get_kimi_k3_mxfp8_tuning_buckets(192) == low_m + assert _get_kimi_k3_mxfp8_tuning_buckets(193) == (*low_m, 256) + assert _get_kimi_k3_mxfp8_tuning_buckets(4096) == ( + *low_m, + 256, + 512, + 1024, + 2048, + 4096, + ) + assert tuple( + _kimi_k3_mxfp8_tuning_bucket(m) + for m in ( + 1, + 3, + 5, + 8, + 9, + 15, + 16, + 17, + 31, + 32, + 64, + 65, + 79, + 80, + 128, + 129, + 191, + 192, + 193, + 255, + 256, + 511, + 512, + ) + ) == ( + 1, + 2, + 4, + 8, + 8, + 8, + 16, + 16, + 16, + 32, + 64, + 80, + 80, + 80, + 128, + 144, + 192, + 192, + 256, + 256, + 256, + 512, + 512, + ) diff --git a/tests/unittest/_torch/visual_gen/test_flux_transformer.py b/tests/unittest/_torch/visual_gen/test_flux_transformer.py index f77662e76cb5..c839bd42fb63 100644 --- a/tests/unittest/_torch/visual_gen/test_flux_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_flux_transformer.py @@ -90,6 +90,13 @@ def _make_fake_flux2_parallel_attn(): pre_quant_scale=None, force_dynamic_quantization=False, ) + # Flux now dispatches on Linear.can_use_cute_dsl_nvfp4_swiglu_blackwell(), the + # same predicate that drives the gate/up interleave, so the stand-in has to + # expose it. Mirror the predicate's dependence on use_cute_dsl_blockscaling_mm + # to keep the guard tests below meaningful. + gate_up_proj.can_use_cute_dsl_nvfp4_swiglu_blackwell = ( + lambda: gate_up_proj.use_cute_dsl_blockscaling_mm + ) attn.to_qkv_mlp_proj = SimpleNamespace( tp_size=2, qkv_proj=object(),