diff --git a/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_instantiation.cu b/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_instantiation.cu index 6d27a11029a..e8d78329b76 100644 --- a/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_instantiation.cu +++ b/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_instantiation.cu @@ -44,22 +44,17 @@ template class CutlassMoeFCRunner<__nv_fp8_e4m3, cutlass::uint4b_t, __nv_bfloat1 #endif #endif #ifdef ENABLE_FP4 -template class CutlassMoeFCRunner; -template class CutlassMoeFCRunner; -template class CutlassMoeFCRunner<__nv_fp8_e4m3, Fp4Type, half>; -template class CutlassMoeFCRunner<__nv_fp8_e4m3, Fp4Type, half, half, half, false, - Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0>; -// PHASE3_POST_MMA_PLACEHOLDER: future post-MMA MXFP4 x FP8/MXFP8 paths should -// use FP16/BF16 external activation input plus online FP8/MXFP8 quantization, -// not the transitional direct-FP8 input runner. -template class CutlassMoeFCRunner; +template class CutlassMoeFCRunner<__nv_fp4_e2m1, __nv_fp4_e2m1, half>; +template class CutlassMoeFCRunner<__nv_fp4_e2m1, __nv_fp4_e2m1, half, half>; +template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, half>; +template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, half, half>; +template class CutlassMoeFCRunner; #ifdef ENABLE_BF16 -template class CutlassMoeFCRunner; -template class CutlassMoeFCRunner; -template class CutlassMoeFCRunner<__nv_fp8_e4m3, Fp4Type, __nv_bfloat16>; -template class CutlassMoeFCRunner<__nv_fp8_e4m3, Fp4Type, __nv_bfloat16, __nv_bfloat16, - __nv_bfloat16, false, Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0>; -template class CutlassMoeFCRunner<__nv_bfloat16, Fp4Type>; +template class CutlassMoeFCRunner<__nv_fp4_e2m1, __nv_fp4_e2m1, __nv_bfloat16>; +template class CutlassMoeFCRunner<__nv_fp4_e2m1, __nv_fp4_e2m1, __nv_bfloat16, __nv_bfloat16>; +template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, __nv_bfloat16>; +template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, __nv_bfloat16, __nv_bfloat16>; +template class CutlassMoeFCRunner<__nv_bfloat16, __nv_fp4_e2m1>; #endif #endif diff --git a/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh b/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh index 7decd3befc8..4c5a6aedbbb 100644 --- a/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh +++ b/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh @@ -15,7 +15,6 @@ */ #include -#include #include #include #include @@ -55,7 +54,6 @@ #include "tensorrt_llm/common/dataType.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" -#include "tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_prebuild.h" #include "tensorrt_llm/kernels/preQuantScaleKernel.h" #include "tensorrt_llm/kernels/quantization.cuh" @@ -77,51 +75,6 @@ namespace tensorrt_llm::kernels::cutlass_kernels { constexpr int CVT_ELTS_PER_THREAD = 8; -struct FloatMaxOp { - __device__ float operator()(float a, float b) const { return fmaxf(a, b); } -}; - -template -struct AbsMaxOp { - using Accum = float; - - __device__ static Accum zero() { return 0.0f; } - - __device__ static Accum abs(T value) { return fabsf(static_cast(value)); } - - __device__ static Accum update(Accum current, T value) { return fmaxf(current, abs(value)); } - - __device__ static float to_float(Accum value) { return value; } -}; - -template <> -struct AbsMaxOp { - using Accum = half; - - __device__ static Accum zero() { return __float2half(0.0f); } - - __device__ static Accum abs(half value) { return __habs(value); } - - __device__ static Accum update(Accum current, half value) { return __hmax(current, abs(value)); } - - __device__ static float to_float(Accum value) { return __half2float(value); } -}; - -template <> -struct AbsMaxOp<__nv_bfloat16> { - using Accum = __nv_bfloat16; - - __device__ static Accum zero() { return __float2bfloat16(0.0f); } - - __device__ static Accum abs(__nv_bfloat16 value) { return __habs(value); } - - __device__ static Accum update(Accum current, __nv_bfloat16 value) { - return __hmax(current, abs(value)); - } - - __device__ static float to_float(Accum value) { return __bfloat162float(value); } -}; - template auto dispatchNVFP44Over6Config(Fn&& fn) { bool const use4Over6 = tensorrt_llm::common::getEnvNVFP4Use4Over6(); @@ -1197,35 +1150,6 @@ float const** computeFP8DequantScale(float const** alpha_scale_ptr_array, return alpha_scale_ptr_array; } -__global__ void prepareProfilerFP8TokenScalePtrArrayKernel( - float const** alpha_scale_ptr_array, float const* fp8_token_scale, - int64_t const* expert_first_token_offset, int const num_experts_per_node, - int const expert_first_token_offset_stride) { - int const expert = blockIdx.x * blockDim.x + threadIdx.x; - int const sample = blockIdx.y; - if (expert >= num_experts_per_node) { - return; - } - - auto const* sample_offsets = - expert_first_token_offset + sample * expert_first_token_offset_stride; - alpha_scale_ptr_array[sample * num_experts_per_node + expert] = - fp8_token_scale + sample_offsets[expert]; -} - -float const** prepareProfilerFP8TokenScalePtrArray(float const** alpha_scale_ptr_array, - float const* fp8_token_scale, - int64_t const* expert_first_token_offset, - int const num_experts_per_node, - int const num_samples, cudaStream_t stream) { - int const threads = std::min(128, num_experts_per_node); - int const blocks = (num_experts_per_node + threads - 1) / threads; - prepareProfilerFP8TokenScalePtrArrayKernel<<>>( - alpha_scale_ptr_array, fp8_token_scale, expert_first_token_offset, num_experts_per_node, - num_experts_per_node + 1); - return alpha_scale_ptr_array; -} - template __device__ void setupFP4BlockScalingFactors( TmaWarpSpecializedGroupedGemmInput& layout_info, int expert, int gemm_m, int gemm_n, int gemm_k, @@ -1416,35 +1340,10 @@ __global__ void computeStridesTmaWarpSpecializedKernel( layout_info2.swap_ab ? gemm_m : gemm2_n, gemm2_k); } - bool const needs_zero_token_weight_desc = - layout_info1.int4_groupwise_params.enabled || layout_info2.int4_groupwise_params.enabled; - - auto compute_tma_strides = [&]() { - assert(gemm_m <= INT32_MAX); - assert(gemm1_n > 0 && gemm1_n <= INT32_MAX); - assert(gemm1_k > 0 && gemm1_k <= INT32_MAX); - assert(gemm2_n > 0 && gemm2_n <= INT32_MAX); - assert(gemm2_k > 0 && gemm2_k <= INT32_MAX); - computeTmaWarpSpecializedInputStrides(layout_info1, gemm_m, gemm1_n, gemm1_k, expert); - computeTmaWarpSpecializedInputStrides(layout_info2, gemm_m, gemm2_n, gemm2_k, expert); - }; - - if (gemm_m != 0 || needs_zero_token_weight_desc) { - compute_tma_strides(); - } - - if (gemm_m == 0 && needs_zero_token_weight_desc) { - layout_info1.ptr_weight[expert] = safe_inc_ptr(weights1, expert * (gemm1_n * gemm1_k)); - layout_info2.ptr_weight[expert] = safe_inc_ptr(weights2, expert * (gemm2_n * gemm2_k)); - } - // Skip expensive stride/pointer/SF setup for experts with no assigned tokens. // All problem shapes (including int4_groupwise) are initialized above so CUTLASS - // can correctly traverse the problem list. For groupwise W4 paths, weight - // strides/pointers are initialized above because the prebuilt global weight - // TMA descriptor uses group 0 as its base even when that expert has no tokens. - // The remaining work (alpha scales, block scaling factors, activation/output - // pointers) is only needed for active experts. + // can correctly traverse the problem list. The remaining work (alpha scales, + // block scaling factors, strides, pointers) is only needed for active experts. // For decode (1 token, top_k=8, 128 experts), this skips ~120 of 128 experts. if (gemm_m == 0) { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) @@ -1472,34 +1371,31 @@ __global__ void computeStridesTmaWarpSpecializedKernel( }; setupIfSelected(TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaledConfig{}, quant_params.fp4); - // FP8 x MXFP4 has plain FP8 activation and must not fabricate an activation - // block-scale tensor. The Humming-style pre-MMA path consumes folded weight - // scale through int4_groupwise metadata; future post-MMA FP8 x MXFP4 support - // should use a weight-scale-only path rather than MXFPX activation SF setup. + setupIfSelected(TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaledConfig{}, + quant_params.fp8_mxfp4); setupIfSelected(TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaledConfig{}, quant_params.mxfp8_mxfp4); setupIfSelected(TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaledConfig{}, quant_params.mxfp8_mxfp8); - auto const* fc1_weight_scale = - quant_params.fp8_mxfp4.fc1.weight_block_scale - ? reinterpret_cast( - quant_params.fp8_mxfp4.fc1.weight_block_scale) - : reinterpret_cast( - quant_params.groupwise.fc1.weight_scales); - auto const* fc2_weight_scale = - quant_params.fp8_mxfp4.fc2.weight_block_scale - ? reinterpret_cast( - quant_params.fp8_mxfp4.fc2.weight_block_scale) - : reinterpret_cast( - quant_params.groupwise.fc2.weight_scales); + assert(gemm_m <= INT32_MAX); + assert(gemm1_n > 0 && gemm1_n <= INT32_MAX); + assert(gemm1_k > 0 && gemm1_k <= INT32_MAX); + assert(gemm2_n > 0 && gemm2_n <= INT32_MAX); + assert(gemm2_k > 0 && gemm2_k <= INT32_MAX); + computeTmaWarpSpecializedInputStrides(layout_info1, gemm_m, gemm1_n, gemm1_k, expert); + computeTmaWarpSpecializedInputStrides(layout_info2, gemm_m, gemm2_n, gemm2_k, expert); computeTmaWarpSpecializedInputPointers( layout_info1, gemm_m, gemm1_n, gemm1_k, num_tokens_before_expert, expert, gemm1_in, weights1, - fc1_weight_scale, bias1, gemm1_output, nullptr, nullptr, expert); + reinterpret_cast( + quant_params.groupwise.fc1.weight_scales), + bias1, gemm1_output, nullptr, nullptr, expert); computeTmaWarpSpecializedInputPointers( layout_info2, gemm_m, gemm2_n, gemm2_k, num_tokens_before_expert, expert, gemm2_in, weights2, - fc2_weight_scale, bias2, gemm2_output, router_scales, permuted_row_to_unpermuted_row, expert); + reinterpret_cast( + quant_params.groupwise.fc2.weight_scales), + bias2, gemm2_output, router_scales, permuted_row_to_unpermuted_row, expert); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); #endif @@ -1539,9 +1435,7 @@ __global__ void expandInputRowsKernel( int64_t const* expert_first_token_offset, TmaWarpSpecializedGroupedGemmInput::ElementSF* fc1_act_sf_flat, TmaWarpSpecializedGroupedGemmInput::ElementSF const* input_sf, bool const swizzled_input_sf, - int64_t const num_experts_per_node, InputActivationsType const* prequant_scales = nullptr, - float* fp8_token_dequant_scale = nullptr, float const* fp8_token_residual_scale = nullptr, - float const** fp8_token_scale_ptr_array = nullptr) { + int64_t const num_experts_per_node, InputActivationsType const* prequant_scales = nullptr) { static_assert(BlockScalingType == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE || !PRE_QUANT_AWQ, "AWQ and Block Scaling are mutually exclusive"); @@ -1551,17 +1445,13 @@ __global__ void expandInputRowsKernel( !PRE_QUANT_AWQ; constexpr bool is_mxfp8_input = is_mxfp8 && std::is_same_v; constexpr bool need_mxfp8_quant = is_mxfp8 && !is_mxfp8_input; - constexpr bool need_per_token_fp8_quant = - std::is_same_v && - BlockScalingType == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE && - !PRE_QUANT_AWQ && !std::is_same_v; #ifdef ENABLE_FP4 constexpr bool is_nvfp4 = - std::is_same_v && + std::is_same_v && BlockScalingType == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4 && !PRE_QUANT_AWQ; - constexpr bool is_nvfp4_input = is_nvfp4 && std::is_same_v; + constexpr bool is_nvfp4_input = is_nvfp4 && std::is_same_v; constexpr bool need_nvfp4_quant = is_nvfp4 && !is_nvfp4_input; #else constexpr bool is_nvfp4 = false; @@ -1569,7 +1459,7 @@ __global__ void expandInputRowsKernel( constexpr bool need_nvfp4_quant = false; #endif - static_assert(need_nvfp4_quant || need_mxfp8_quant || need_per_token_fp8_quant || PRE_QUANT_AWQ || + static_assert(need_nvfp4_quant || need_mxfp8_quant || PRE_QUANT_AWQ || std::is_same_v, "Only NVFP4, MXFP8 and WINT4_AFP8 supports outputting a different format as part " "of the expansion"); @@ -1598,14 +1488,6 @@ __global__ void expandInputRowsKernel( TmaWarpSpecializedGroupedGemmInput::alignToSfDim(hidden_size, min_k_dim_alignment); int64_t const num_valid_tokens = expert_first_token_offset[num_experts_per_node]; - - if (fp8_token_scale_ptr_array && fp8_token_dequant_scale && blockIdx.x == 0) { - for (int expert = threadIdx.x; expert < num_experts_per_node; expert += blockDim.x) { - fp8_token_scale_ptr_array[expert] = - fp8_token_dequant_scale + expert_first_token_offset[expert]; - } - } - for (int64_t permuted_row = blockIdx.x; permuted_row < num_valid_tokens; permuted_row += gridDim.x) { int64_t const unpermuted_row = permuted_row_to_unpermuted_row[permuted_row]; @@ -1697,41 +1579,6 @@ __global__ void expandInputRowsKernel( dest_row_ptr[elem_index] = arrayConvert(frag_elems); } - } else if constexpr (need_per_token_fp8_quant) { - using AmaxOp = AbsMaxOp; - typename AmaxOp::Accum thread_amax = AmaxOp::zero(); - for (int elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) { - auto input_value = source_row_ptr[elem_index]; - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < ELEM_PER_THREAD; ++i) { - thread_amax = AmaxOp::update(thread_amax, input_value[i]); - } - } - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage reduce_storage; - __shared__ float shared_token_quant_scale; - float const row_amax = - BlockReduce(reduce_storage).Reduce(AmaxOp::to_float(thread_amax), FloatMaxOp{}); - if (threadIdx.x == 0) { - float const quant = row_amax > 0.0f ? (448.0f / row_amax) : 1.0f; - float const residual = - fp8_token_residual_scale ? fp8_token_residual_scale[permuted_row] : 1.0f; - shared_token_quant_scale = quant; - fp8_token_dequant_scale[permuted_row] = (1.0f / quant) * residual; - } - __syncthreads(); - - for (int elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) { - auto input_value = arrayConvert>( - source_row_ptr[elem_index]); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < ELEM_PER_THREAD; ++i) { - input_value[i] *= shared_token_quant_scale; - } - dest_row_ptr[elem_index] = - arrayConvert, OutputElem>(input_value); - } } else { for (int elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) { dest_row_ptr[elem_index] = source_row_ptr[elem_index]; @@ -1765,13 +1612,12 @@ void expandInputRowsKernelLauncher( bool use_per_expert_act_scale, int64_t* expert_first_token_offset, TmaWarpSpecializedGroupedGemmInput::ElementSF* fc1_act_sf_flat, TmaWarpSpecializedGroupedGemmInput::ElementSF const* input_sf, bool const swizzled_input_sf, - void const* prequant_scales, float* fp8_token_dequant_scale, - float const* fp8_token_residual_scale, float const** fp8_token_scale_ptr_array, bool enable_pdl, - cudaStream_t stream) { + void const* prequant_scales, bool enable_pdl, cudaStream_t stream) { #ifdef ENABLE_FP4 - TLLM_CHECK_WITH_INFO((std::is_same_v && fc1_act_sf_flat) || - !use_per_expert_act_scale, - "Per-expert act scale for FC1 is only supported for NVFP4 activations"); + TLLM_CHECK_WITH_INFO( + (std::is_same_v && fc1_act_sf_flat) || + !use_per_expert_act_scale, + "Per-expert act scale for FC1 is only supported for NVFP4 activations"); #endif static int64_t const smCount = tensorrt_llm::common::getMultiProcessorCount(); @@ -1786,20 +1632,13 @@ void expandInputRowsKernelLauncher( // Always MXFP8 if constexpr (std::is_same_v && !std::is_same_v) { - bool const use_per_token_fp8_quant = fp8_token_dequant_scale != nullptr; TLLM_CHECK_WITH_INFO(quant_params.mxfp8_mxfp4.fc1.weight_block_scale || - quant_params.mxfp8_mxfp8.fc1.weight_block_scale || prequant_scales || - use_per_token_fp8_quant, - "MXFP8 block scaling, prequant_scales or FP8 token scale output " - "parameters not provided"); + quant_params.mxfp8_mxfp8.fc1.weight_block_scale || prequant_scales, + "MXFP8 block scaling or prequant_scales parameters not provided"); return prequant_scales ? &expandInputRowsKernel< InputActivationsType, ExpandedActivationsType, TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE, true> - : use_per_token_fp8_quant - ? &expandInputRowsKernel< - InputActivationsType, ExpandedActivationsType, - TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE, false> : &expandInputRowsKernel< InputActivationsType, ExpandedActivationsType, TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX, false>; @@ -1819,7 +1658,7 @@ void expandInputRowsKernelLauncher( } else #endif #ifdef ENABLE_FP4 - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { TLLM_CHECK_WITH_INFO(quant_params.fp4.fc1.weight_block_scale, "NVFP4 block scaling is expected for FP4xFP4"); TLLM_CHECK_WITH_INFO(!prequant_scales, "NVFP4 is not supported for AWQ"); @@ -1856,8 +1695,7 @@ void expandInputRowsKernelLauncher( quant_params.fp4.fc1.act_global_scale, use_per_expert_act_scale, expert_first_token_offset, fc1_act_sf_flat, input_sf, swizzled_input_sf, num_experts_per_node, - reinterpret_cast(prequant_scales), - fp8_token_dequant_scale, fp8_token_residual_scale, fp8_token_scale_ptr_array); + reinterpret_cast(prequant_scales)); } #define INSTANTIATE_EXPAND_INPUT_ROWS(InputActivationsType, ExpandedActivationsType) \ @@ -1870,9 +1708,7 @@ void expandInputRowsKernelLauncher( int64_t* expert_first_token_offset, \ TmaWarpSpecializedGroupedGemmInput::ElementSF* fc1_act_sf_flat, \ TmaWarpSpecializedGroupedGemmInput::ElementSF const* input_sf, bool const swizzled_input_sf, \ - void const* prequant_scales, float* fp8_token_dequant_scale, \ - float const* fp8_token_residual_scale, float const** fp8_token_scale_ptr_array, \ - bool enable_pdl, cudaStream_t stream) + void const* prequant_scales, bool enable_pdl, cudaStream_t stream) // Instantiate the data types that are used by the external pytorch op // INSTANTIATE_EXPAND_INPUT_ROWS(float, float); @@ -2309,11 +2145,10 @@ __global__ __launch_bounds__(ACTIVATION_THREADS_PER_BLOCK) void doActivationKern ScaleBiasType const* bias_ptr, bool bias_is_broadcast, int64_t const* expert_first_token_offset, int num_experts_per_node, int64_t inter_size, float const* fc2_act_global_scale, bool use_per_expert_act_scale, TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_act_sf_flat, - float* fp8_token_dequant_scale, float const* fp8_token_residual_scale, ActivationParams activation_params) { #ifdef ENABLE_FP4 constexpr bool IsNVFP4 = - std::is_same_v && + std::is_same_v && BlockScalingType == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4; constexpr bool IsMXFP8 = std::is_same_v && @@ -2411,7 +2246,7 @@ __global__ __launch_bounds__(ACTIVATION_THREADS_PER_BLOCK) void doActivationKern if (activation_params.swiglu_limit) { fn.limit = gate_limit; } - auto compute_activation = [&](int64_t elem_index) { + for (int64_t elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) { auto fc1_value = arrayConvert(gemm_result_vec[elem_index + gated_off_vec]); if (bias_ptr) { @@ -2432,41 +2267,8 @@ __global__ __launch_bounds__(ACTIVATION_THREADS_PER_BLOCK) void doActivationKern return fn(fc1_value); } }(); - return gate_act; - }; - - bool const use_per_token_fp8_quant = - std::is_same_v && - BlockScalingType == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE && - fp8_token_dequant_scale != nullptr; - float token_quant_scale = quant_scale; - if (use_per_token_fp8_quant) { - float thread_amax = 0.0f; - for (int64_t elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) { - auto gate_act = compute_activation(elem_index); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < ComputeElem::kElements; ++i) { - thread_amax = fmaxf(thread_amax, fabsf(gate_act[i])); - } - } - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage reduce_storage; - __shared__ float shared_token_quant_scale; - float const row_amax = BlockReduce(reduce_storage).Reduce(thread_amax, FloatMaxOp{}); - if (tid == 0) { - float const quant = row_amax > 0.0f ? (448.0f / row_amax) : 1.0f; - float const residual = fp8_token_residual_scale ? fp8_token_residual_scale[token] : 1.0f; - shared_token_quant_scale = quant; - fp8_token_dequant_scale[token] = (1.0f / quant) * residual; - } - __syncthreads(); - token_quant_scale = shared_token_quant_scale; - } - - for (int64_t elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) { - auto post_act_val = compute_activation(elem_index) * token_quant_scale; + auto post_act_val = gate_act * quant_scale; if constexpr (IsNVFP4 || IsMXFP8) { // We use GemmOutputType as the intermediate compute type as that should always be @@ -2519,9 +2321,8 @@ void doActivation(T* output, GemmOutputType const* gemm_result, float const* fp8 int64_t const* expert_first_token_offset, int num_experts_per_node, int64_t inter_size, int64_t expanded_num_tokens, ActivationParams activation_type, QuantParams const& quant_params, bool use_per_expert_act_scale, - TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_act_sf_flat, - float* fp8_token_dequant_scale, float const* fp8_token_residual_scale, - bool enable_pdl, cudaStream_t stream) { + TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_act_sf_flat, bool enable_pdl, + cudaStream_t stream) { static int64_t const smCount = tensorrt_llm::common::getMultiProcessorCount(); // Note: Launching 8 blocks per SM can fully leverage the memory bandwidth (tested on B200). // N-dim SF padding has been removed (CUTLASS grouped GEMM never reads beyond @@ -2589,7 +2390,7 @@ void doActivation(T* output, GemmOutputType const* gemm_result, float const* fp8 TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType, TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE>{}; #ifdef ENABLE_FP4 - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { TLLM_CHECK_WITH_INFO(quant_params.fp4.fc2.weight_block_scale, "NVFP4 block scaling is expected for FP4xFP4"); return dispatchNVFP44Over6Config( @@ -2621,8 +2422,7 @@ void doActivation(T* output, GemmOutputType const* gemm_result, float const* fp8 cudaLaunchKernelEx(&config, fn, output, gemm_result, fp8_quant, bias, bias_is_broadcast, expert_first_token_offset, num_experts_per_node, inter_size, quant_params.fp4.fc2.act_global_scale, use_per_expert_act_scale, - fc2_act_sf_flat, fp8_token_dequant_scale, fp8_token_residual_scale, - activation_type); + fc2_act_sf_flat, activation_type); } // ============================== Lora Add Bias ================================= @@ -2797,32 +2597,18 @@ void dequantFP8(OutputType* output, InputType const* input, int64_t const* num_v output, input, num_valid_tokens_ptr, inter_size, scale, scale_is_dequant); } -template -std::unique_ptr -makeDeepSeekBlockScaleGemmRunnerIfSupported() { - // This runner is only for the DeepSeek BF16 activation x FP8 weight - // block-scale path. Mixed W4/MXFP4 paths still instantiate CutlassMoeFCRunner - // but should not construct a DeepSeek block-scale GEMM runner. - if constexpr (std::is_same_v && std::is_same_v && - std::is_same_v) { - return std::make_unique>(); - } else { - return nullptr; - } -} - template -CutlassMoeFCRunner +CutlassMoeFCRunner::CutlassMoeFCRunner() - : blockscale_gemm_runner_( - makeDeepSeekBlockScaleGemmRunnerIfSupported()) {} + : blockscale_gemm_runner_{ + std::make_unique>()} {} template + bool IsMXFPX, class Enable> std::map> -CutlassMoeFCRunner::getWorkspaceDeviceBufferSizes(int64_t const num_rows, int64_t const hidden_size, int64_t const inter_size, @@ -2904,12 +2690,6 @@ CutlassMoeFCRunner -size_t -CutlassMoeFCRunner::getWorkspaceSize(int64_t const num_rows, int64_t const hidden_size, - int64_t const inter_size, int const num_experts, - int const experts_per_token, - ActivationType activation_type, - MOEParallelismConfig parallelism_config, bool use_lora, - bool use_deepseek_fp8_block_scale, - bool use_mxfp8_act_scaling, bool min_latency_mode, - bool use_awq) { + bool IsMXFPX, class Enable> +size_t CutlassMoeFCRunner::getWorkspaceSize(int64_t const num_rows, + int64_t const hidden_size, + int64_t const inter_size, int const num_experts, + int const experts_per_token, + ActivationType activation_type, + MOEParallelismConfig parallelism_config, + bool use_lora, + bool use_deepseek_fp8_block_scale, + bool use_mxfp8_act_scaling, + bool min_latency_mode, bool use_awq) { int const ep_size = parallelism_config.ep_size; TLLM_CHECK_WITH_INFO(num_experts % ep_size == 0, "Number of experts must be a multiple of ep size"); @@ -3038,15 +2811,18 @@ CutlassMoeFCRunner -void CutlassMoeFCRunner< - T, WeightType, OutputType, InputType, BackBoneType, IsMXFPX, Sm90Wfp4Afp8Mode, - Enable>::configureWsPtrs(char* ws_ptr, int64_t const num_rows, int64_t const hidden_size, - int64_t const inter_size, int const num_experts_per_node, - int const experts_per_token, ActivationType activation_type, - MOEParallelismConfig parallelism_config, bool use_lora, - bool use_deepseek_fp8_block_scale, bool use_mxfp8_act_scaling, - bool min_latency_mode, bool use_awq) { + bool IsMXFPX, class Enable> +void CutlassMoeFCRunner::configureWsPtrs(char* ws_ptr, int64_t const num_rows, + int64_t const hidden_size, + int64_t const inter_size, + int const num_experts_per_node, + int const experts_per_token, + ActivationType activation_type, + MOEParallelismConfig parallelism_config, + bool use_lora, bool use_deepseek_fp8_block_scale, + bool use_mxfp8_act_scaling, bool min_latency_mode, + bool use_awq) { auto workspaces = getWorkspaceDeviceBufferSizes( num_rows, hidden_size, inter_size, num_experts_per_node, experts_per_token, activation_type, use_lora, use_deepseek_fp8_block_scale, use_mxfp8_act_scaling, min_latency_mode, use_awq); @@ -3105,11 +2881,7 @@ void CutlassMoeFCRunner< } alpha_scale_ptr_array_fc1_ = getWsPtr((float const*)(nullptr), "alpha_scale_ptr_array_fc1"); - if constexpr (use_wfp4afp8 && Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0) { - alpha_scale_ptr_array_fc2_ = alpha_scale_ptr_array_fc1_; - } else { - alpha_scale_ptr_array_fc2_ = getWsPtr((float const*)(nullptr), "alpha_scale_ptr_array_fc2"); - } + alpha_scale_ptr_array_fc2_ = getWsPtr((float const*)(nullptr), "alpha_scale_ptr_array_fc2"); // NOTE: We alias these, but if we fuse the quantization for GEMM2 into GEMM1 they will need // separated @@ -3121,10 +2893,6 @@ void CutlassMoeFCRunner< TLLM_CHECK(fc1_fp4_act_scale_ != nullptr); TLLM_CHECK(fc2_fp4_act_scale_ != nullptr); } - act_fp8_token_scale_ = getWsPtr(float{}, "act_fp8_token_scale"); - if constexpr (use_wfp4afp8 && Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0) { - TLLM_CHECK(act_fp8_token_scale_ != nullptr); - } tma_ws_grouped_gemm1_input_ = {}; tma_ws_grouped_gemm2_input_ = {}; @@ -3132,13 +2900,11 @@ void CutlassMoeFCRunner< tma_ws_grouped_gemm1_input_.configureWorkspace( getWsPtr(int8_t{}, "tma_ws_gemm1_workspace"), num_experts_per_node, getWsPtr(int8_t{}, "gemm_workspace"), workspaces.at("gemm_workspace").first, - getWsPtr(int8_t{}, "precomputed_scheduler_workspace"), - workspaces.at("precomputed_scheduler_workspace").first, getScalingType()); + getScalingType()); tma_ws_grouped_gemm2_input_.configureWorkspace( getWsPtr(int8_t{}, "tma_ws_gemm2_workspace"), num_experts_per_node, getWsPtr(int8_t{}, "gemm_workspace"), workspaces.at("gemm_workspace").first, - getWsPtr(int8_t{}, "precomputed_scheduler_workspace"), - workspaces.at("precomputed_scheduler_workspace").first, getScalingType()); + getScalingType()); } lora_fc1_result_ = {}; @@ -3169,9 +2935,9 @@ void CutlassMoeFCRunner< } template + bool IsMXFPX, class Enable> kernels::fp8_blockscale_gemm::CutlassFp8BlockScaleGemmRunnerInterface* -CutlassMoeFCRunner::getDeepSeekBlockScaleGemmRunner() const { TLLM_CHECK_WITH_INFO( (std::is_same_v && std::is_same_v), @@ -3182,19 +2948,16 @@ CutlassMoeFCRunner -void CutlassMoeFCRunner< - T, WeightType, OutputType, InputType, ScaleBiasType, IsMXFPX, Sm90Wfp4Afp8Mode, - Enable>::BlockScaleFC1(DeepSeekBlockScaleGemmRunner& gemm_runner, T const* const input, - T* const output, void* const gemm_output, - int64_t const* const expert_first_token_offset, - WeightType const* const fc1_expert_weights, - ScaleBiasType const* const fc1_expert_biases, - float const* const fc2_fp8_quant, int64_t const num_rows, - int64_t const expanded_num_rows, int64_t const hidden_size, - int64_t const inter_size, int const num_experts_per_node, - ActivationParams fc1_activation_type, QuantParams& quant_params, - bool enable_pdl, cudaStream_t stream) { + bool IsMXFPX, class Enable> +void CutlassMoeFCRunner:: + BlockScaleFC1(DeepSeekBlockScaleGemmRunner& gemm_runner, T const* const input, T* const output, + void* const gemm_output, int64_t const* const expert_first_token_offset, + WeightType const* const fc1_expert_weights, + ScaleBiasType const* const fc1_expert_biases, float const* const fc2_fp8_quant, + int64_t const num_rows, int64_t const expanded_num_rows, + int64_t const hidden_size, int64_t const inter_size, + int const num_experts_per_node, ActivationParams fc1_activation_type, + QuantParams& quant_params, bool enable_pdl, cudaStream_t stream) { bool const is_gated_activation = isGatedActivation(fc1_activation_type); int shape_n = is_gated_activation ? inter_size * 2 : inter_size; @@ -3212,30 +2975,26 @@ void CutlassMoeFCRunner< output, static_cast(gemm_output), fc2_fp8_quant, fc1_expert_biases, bias_is_broadcast, expert_first_token_offset, num_experts_per_node, inter_size, expanded_num_rows, fc1_activation_type, quant_params, use_per_expert_act_scale, - nullptr, nullptr, nullptr, enable_pdl, stream); + nullptr, enable_pdl, stream); sync_check_cuda_error(stream); } template -void CutlassMoeFCRunner< - T, WeightType, OutputType, InputType, ScaleBiasType, IsMXFPX, Sm90Wfp4Afp8Mode, - Enable>::BlockScaleFC2(DeepSeekBlockScaleGemmRunner& gemm_runner, T const* const input, - void* const gemm_output, OutputType* const final_output, - int64_t const* const expert_first_token_offset, - WeightType const* const fc2_expert_weights, - ScaleBiasType const* const fc2_expert_biases, - float const* const unpermuted_final_scales, - int const* const unpermuted_row_to_permuted_row, - int const* const permuted_row_to_unpermuted_row, - int const* const token_selected_experts, - int64_t const* const num_valid_tokens_ptr, int64_t const num_rows, - int64_t const expanded_num_rows, int64_t const hidden_size, - int64_t const unpadded_hidden_size, int64_t const inter_size, - int64_t const num_experts_per_node, int64_t const k, - MOEParallelismConfig parallelism_config, bool const enable_alltoall, - QuantParams& quant_params, bool enable_pdl, cudaStream_t stream) { + bool IsMXFPX, class Enable> +void CutlassMoeFCRunner:: + BlockScaleFC2( + DeepSeekBlockScaleGemmRunner& gemm_runner, T const* const input, void* const gemm_output, + OutputType* const final_output, int64_t const* const expert_first_token_offset, + WeightType const* const fc2_expert_weights, ScaleBiasType const* const fc2_expert_biases, + float const* const unpermuted_final_scales, int const* const unpermuted_row_to_permuted_row, + int const* const permuted_row_to_unpermuted_row, int const* const token_selected_experts, + int64_t const* const num_valid_tokens_ptr, int64_t const num_rows, + int64_t const expanded_num_rows, int64_t const hidden_size, + int64_t const unpadded_hidden_size, int64_t const inter_size, + int64_t const num_experts_per_node, int64_t const k, + MOEParallelismConfig parallelism_config, bool const enable_alltoall, + QuantParams& quant_params, bool enable_pdl, cudaStream_t stream) { int shape_n = hidden_size; int shape_k = inter_size; @@ -3255,15 +3014,15 @@ void CutlassMoeFCRunner< } template -T const* -CutlassMoeFCRunner::applyPrequantScale(void* smoothed_act, void const* permuted_data, - void const* prequant_scales, - int64_t const* num_valid_tokens_ptr, - int64_t const expanded_num_rows, - int64_t const seq_len, bool const use_awq, - cudaStream_t stream) { + bool IsMXFPX, class Enable> +T const* CutlassMoeFCRunner::applyPrequantScale(void* smoothed_act, + void const* permuted_data, + void const* prequant_scales, + int64_t const* num_valid_tokens_ptr, + int64_t const expanded_num_rows, + int64_t const seq_len, bool const use_awq, + cudaStream_t stream) { T const* gemm_input; bool use_prequant_scale_kernel = use_awq && !std::is_same_v; if (use_prequant_scale_kernel) { @@ -3285,29 +3044,22 @@ CutlassMoeFCRunner -void CutlassMoeFCRunner< - T, WeightType, OutputType, InputType, BackBoneType, IsMXFPX, Sm90Wfp4Afp8Mode, - Enable>::gemm1(MoeGemmRunner& gemm_runner, - DeepSeekBlockScaleGemmRunner* fp8_blockscale_gemm_runner, T const* const input, - T* const output, void* const intermediate_result, - int64_t const* const expert_first_token_offset, - TmaWarpSpecializedGroupedGemmInput const tma_ws_input_template, - WeightType const* const fc1_expert_weights, - ScaleBiasType const* const fc1_expert_biases, - int64_t const* const num_valid_tokens_ptr, - ScaleBiasType const* const fc1_int_scales, float const* const fc1_fp8_dequant, - float const* const fc2_fp8_quant, float* const act_fp8_token_scale, - TmaWarpSpecializedGroupedGemmInput::ElementSF const* fc1_fp4_act_flat, - TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_fp4_act_flat, - QuantParams quant_params, int64_t const num_rows, - int64_t const expanded_num_rows, int64_t const hidden_size, - int64_t const inter_size, int const num_experts_per_node, - ActivationParams fc1_activation_type, float const** alpha_scale_ptr_array, - bool bias_is_broadcast, cudaStream_t stream, - cutlass_extensions::CutlassGemmConfig config, bool min_latency_mode, - int* num_active_experts_per, int* active_expert_global_ids, bool enable_pdl) { + bool IsMXFPX, class Enable> +void CutlassMoeFCRunner::gemm1( + MoeGemmRunner& gemm_runner, + DeepSeekBlockScaleGemmRunner* fp8_blockscale_gemm_runner, T const* const input, T* const output, + void* const intermediate_result, int64_t const* const expert_first_token_offset, + TmaWarpSpecializedGroupedGemmInput const tma_ws_input_template, + WeightType const* const fc1_expert_weights, ScaleBiasType const* const fc1_expert_biases, + int64_t const* const num_valid_tokens_ptr, ScaleBiasType const* const fc1_int_scales, + float const* const fc1_fp8_dequant, float const* const fc2_fp8_quant, + TmaWarpSpecializedGroupedGemmInput::ElementSF const* fc1_fp4_act_flat, + TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_fp4_act_flat, QuantParams quant_params, + int64_t const num_rows, int64_t const expanded_num_rows, int64_t const hidden_size, + int64_t const inter_size, int const num_experts_per_node, ActivationParams fc1_activation_type, + float const** alpha_scale_ptr_array, bool bias_is_broadcast, cudaStream_t stream, + cutlass_extensions::CutlassGemmConfig config, bool min_latency_mode, + int* num_active_experts_per, int* active_expert_global_ids, bool enable_pdl) { if (fp8_blockscale_gemm_runner) { TLLM_CHECK(!min_latency_mode); Self::BlockScaleFC1(*fp8_blockscale_gemm_runner, input, output, intermediate_result, @@ -3354,30 +3106,27 @@ void CutlassMoeFCRunner< alpha_scale_ptr_array = computeFP8DequantScale(alpha_scale_ptr_array, num_experts_per_node, quant_params.groupwise.fc1.alpha, stream); } - if constexpr (use_wfp4afp8 && Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0) { - TLLM_CHECK(alpha_scale_ptr_array != nullptr); - } - auto universal_input = GroupedGemmInput{ - input, - total_tokens_including_expert, - /*weights*/ nullptr, - /*scales*/ nullptr, - /*zeros*/ nullptr, - /*biases*/ nullptr, - /*C*/ static_cast(gemm_output), - alpha_scale_ptr_array, - /*occupancy*/ nullptr, - fc1_activation_type, - num_rows, - /*N*/ int64_t(fc1_out_size), - /*K*/ hidden_size, - num_experts_per_node, - quant_params.groupwise.group_size, - /*bias_is_broadcast*/ true, - /*use_fused_moe*/ false, - stream, - config}; + auto universal_input = + GroupedGemmInput{input, + total_tokens_including_expert, + /*weights*/ nullptr, + /*scales*/ nullptr, + /*zeros*/ nullptr, + /*biases*/ nullptr, + /*C*/ nullptr, + alpha_scale_ptr_array, + /*occupancy*/ nullptr, + fc1_activation_type, + num_rows, + /*N*/ int64_t(fc1_out_size), + /*K*/ hidden_size, + num_experts_per_node, + quant_params.groupwise.group_size, + /*bias_is_broadcast*/ true, + /*use_fused_moe*/ false, + stream, + config}; gemm_runner.moeGemm(universal_input, tma_ws_input); sync_check_cuda_error(stream); @@ -3389,20 +3138,13 @@ void CutlassMoeFCRunner< ? quant_params.fp8_mxfp4.fc2.use_per_expert_act_scale : use_fp8 ? quant_params.fp8.fc2_use_per_expert_act_scale : false; - float* fp8_token_dequant_scale = nullptr; - float const* fp8_token_residual_scale = nullptr; - if constexpr (use_wfp4afp8 && Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0) { - TLLM_CHECK(quant_params.fp8_mxfp4.fc2.global_scale != nullptr); - fp8_token_dequant_scale = act_fp8_token_scale; - fp8_token_residual_scale = quant_params.fp8_mxfp4.fc2.global_scale; - } doActivation( reinterpret_cast(output), static_cast(gemm_output), fc2_fp8_quant, fc1_expert_biases, bias_is_broadcast, expert_first_token_offset, num_experts_per_node, inter_size, expanded_num_rows, fc1_activation_type, quant_params, use_per_expert_act_scale, - fc2_fp4_act_flat, fp8_token_dequant_scale, fp8_token_residual_scale, enable_pdl, stream); + fc2_fp4_act_flat, enable_pdl, stream); sync_check_cuda_error(stream); } else if (use_fp8) { @@ -3440,7 +3182,7 @@ void CutlassMoeFCRunner< output, static_cast(intermediate_result), fc2_fp8_quant, fc1_expert_biases, bias_is_broadcast, expert_first_token_offset, num_experts_per_node, inter_size, expanded_num_rows, fc1_activation_type, quant_params, use_per_expert_act_scale, - nullptr, nullptr, nullptr, enable_pdl, stream); + nullptr, enable_pdl, stream); sync_check_cuda_error(stream); } else if (!is_gated_activation) { @@ -3532,29 +3274,25 @@ void CutlassMoeFCRunner< } template -void CutlassMoeFCRunner:: - gemm2(MoeGemmRunner& - gemm_runner, - DeepSeekBlockScaleGemmRunner* fp8_blockscale_gemm_runner, T const* const input, - void* const gemm_output, OutputType* const final_output, - int64_t const* const expert_first_token_offset, - TmaWarpSpecializedGroupedGemmInput const tma_ws_input_template, - WeightType const* const fc2_expert_weights, ScaleBiasType const* const fc2_expert_biases, - ScaleBiasType const* const fc2_int_scales, float const* const fc2_fp8_dequant, - TmaWarpSpecializedGroupedGemmInput::ElementSF const* fc2_fp4_act_flat, - QuantParams quant_params, float const* const unpermuted_final_scales, - float const* const permuted_final_scales, int const* const unpermuted_row_to_permuted_row, - int const* permuted_row_to_unpermuted_row, int const* const token_selected_experts, - int64_t const* const num_valid_tokens_ptr, int64_t const num_rows, - int64_t const expanded_num_rows, int64_t const hidden_size, - int64_t const unpadded_hidden_size, int64_t const inter_size, - int const num_experts_per_node, int64_t const k, float const** alpha_scale_ptr_array, - bool use_lora, void* fc2_lora, cudaStream_t stream, - MOEParallelismConfig parallelism_config, bool const enable_alltoall, - cutlass_extensions::CutlassGemmConfig config, bool min_latency_mode, - int* num_active_experts_per, int* active_expert_global_ids, bool enable_pdl) { + bool IsMXFPX, class Enable> +void CutlassMoeFCRunner::gemm2( + MoeGemmRunner& gemm_runner, + DeepSeekBlockScaleGemmRunner* fp8_blockscale_gemm_runner, T const* const input, + void* const gemm_output, OutputType* const final_output, + int64_t const* const expert_first_token_offset, + TmaWarpSpecializedGroupedGemmInput const tma_ws_input_template, + WeightType const* const fc2_expert_weights, ScaleBiasType const* const fc2_expert_biases, + ScaleBiasType const* const fc2_int_scales, float const* const fc2_fp8_dequant, + TmaWarpSpecializedGroupedGemmInput::ElementSF const* fc2_fp4_act_flat, QuantParams quant_params, + float const* const unpermuted_final_scales, float const* const permuted_final_scales, + int const* const unpermuted_row_to_permuted_row, int const* permuted_row_to_unpermuted_row, + int const* const token_selected_experts, int64_t const* const num_valid_tokens_ptr, + int64_t const num_rows, int64_t const expanded_num_rows, int64_t const hidden_size, + int64_t const unpadded_hidden_size, int64_t const inter_size, int const num_experts_per_node, + int64_t const k, float const** alpha_scale_ptr_array, bool use_lora, void* fc2_lora, + cudaStream_t stream, MOEParallelismConfig parallelism_config, bool const enable_alltoall, + cutlass_extensions::CutlassGemmConfig config, bool min_latency_mode, + int* num_active_experts_per, int* active_expert_global_ids, bool enable_pdl) { int64_t const* total_tokens_including_expert = expert_first_token_offset + 1; bool const using_tma_ws_gemm2 = gemm_runner.isTmaWarpSpecialized(config); @@ -3597,9 +3335,6 @@ void CutlassMoeFCRunner(gemm_output), nullptr, static_cast(fc2_lora), false, expert_first_token_offset, num_experts_per_node, hidden_size, expanded_num_rows, - ActivationParams(ActivationType::Identity), {}, false, nullptr, nullptr, - nullptr, enable_pdl, stream); + ActivationParams(ActivationType::Identity), {}, false, nullptr, enable_pdl, + stream); sync_check_cuda_error(stream); } @@ -3668,18 +3403,13 @@ void CutlassMoeFCRunner + bool IsMXFPX, class Enable> bool CutlassMoeFCRunner::setupLoraWorkspace(int64_t expanded_num_rows, - int64_t num_rows, - int64_t inter_size, - int64_t hidden_size, - int start_expert, - bool is_gated_activation, - int num_experts_per_node, - bool needs_num_valid, - LoraParams& lora_params, - cudaStream_t stream) { + Enable>::setupLoraWorkspace(int64_t expanded_num_rows, int64_t num_rows, + int64_t inter_size, int64_t hidden_size, + int start_expert, bool is_gated_activation, + int num_experts_per_node, bool needs_num_valid, + LoraParams& lora_params, cudaStream_t stream) { std::vector& host_permuted_rows = host_lora_workspace_.host_permuted_rows; std::vector& host_permuted_fc1_weight_ptrs = host_lora_workspace_.host_permuted_fc1_weight_ptrs; @@ -3767,19 +3497,15 @@ bool CutlassMoeFCRunner + bool IsMXFPX, class Enable> auto CutlassMoeFCRunner::loraFC1(int64_t expanded_num_rows, - int64_t inter_size, int64_t hidden_size, - int num_experts_per_node, - int start_expert, - int64_t const* num_valid_tokens_ptr, - bool is_gated_activation, - ScaleBiasType const* fc1_expert_biases, - LoraParams& lora_params, - float const* input_fp8_dequant, - cudaStream_t stream) - -> ScaleBiasType const* { + Enable>::loraFC1(int64_t expanded_num_rows, int64_t inter_size, + int64_t hidden_size, int num_experts_per_node, + int start_expert, int64_t const* num_valid_tokens_ptr, + bool is_gated_activation, + ScaleBiasType const* fc1_expert_biases, + LoraParams& lora_params, float const* input_fp8_dequant, + cudaStream_t stream) -> ScaleBiasType const* { TLLM_CHECK_WITH_INFO(!act_fp4, "LoRA does not support FP4 activations"); std::vector& host_permuted_fc1_weight_ptrs = host_lora_workspace_.host_permuted_fc1_weight_ptrs; @@ -3853,16 +3579,13 @@ auto CutlassMoeFCRunner + bool IsMXFPX, class Enable> void CutlassMoeFCRunner::loraFC2(int64_t inter_size, int64_t hidden_size, - int num_experts_per_node, - int start_expert, - int64_t const* num_valid_tokens_ptr, - int64_t num_tokens, - LoraParams& lora_params, - float const* fc2_fp8_quant, - cudaStream_t stream) { + Enable>::loraFC2(int64_t inter_size, int64_t hidden_size, + int num_experts_per_node, int start_expert, + int64_t const* num_valid_tokens_ptr, int64_t num_tokens, + LoraParams& lora_params, float const* fc2_fp8_quant, + cudaStream_t stream) { std::vector& host_permuted_fc2_weight_ptrs = host_lora_workspace_.host_permuted_fc2_weight_ptrs; std::vector& host_permuted_fc2_lora_ranks = @@ -3899,22 +3622,20 @@ void CutlassMoeFCRunner -void CutlassMoeFCRunner< - T, WeightType, OutputType, InputType, BackBoneType, IsMXFPX, Sm90Wfp4Afp8Mode, - Enable>::runMoe(void const* input_activations_void, void const* input_sf_void, - bool const swizzled_input_sf, int const* token_selected_experts, - float const* token_final_scales, void const* fc1_expert_weights_void, - void const* fc1_expert_biases_void, ActivationParams fc1_activation_type, - void const* fc2_expert_weights_void, void const* fc2_expert_biases_void, - QuantParams quant_params, int64_t const num_rows, int64_t const hidden_size, - int64_t const unpadded_hidden_size, int64_t const inter_size, - int const full_num_experts, int const experts_per_token, char* workspace_ptr, - void* final_output_void, int* unpermuted_row_to_permuted_row, - MOEParallelismConfig parallelism_config, bool const enable_alltoall, - bool use_lora, LoraParams& lora_params, bool use_deepseek_fp8_block_scale, - bool use_mxfp8_act_scaling, bool min_latency_mode, - MoeMinLatencyParams& min_latency_params, bool enable_pdl, cudaStream_t stream) { + bool IsMXFPX, class Enable> +void CutlassMoeFCRunner:: + runMoe(void const* input_activations_void, void const* input_sf_void, + bool const swizzled_input_sf, int const* token_selected_experts, + float const* token_final_scales, void const* fc1_expert_weights_void, + void const* fc1_expert_biases_void, ActivationParams fc1_activation_type, + void const* fc2_expert_weights_void, void const* fc2_expert_biases_void, + QuantParams quant_params, int64_t const num_rows, int64_t const hidden_size, + int64_t const unpadded_hidden_size, int64_t const inter_size, int const full_num_experts, + int const experts_per_token, char* workspace_ptr, void* final_output_void, + int* unpermuted_row_to_permuted_row, MOEParallelismConfig parallelism_config, + bool const enable_alltoall, bool use_lora, LoraParams& lora_params, + bool use_deepseek_fp8_block_scale, bool use_mxfp8_act_scaling, bool min_latency_mode, + MoeMinLatencyParams& min_latency_params, bool enable_pdl, cudaStream_t stream) { static constexpr bool int_scales_required = std::is_same::value || std::is_same::value || use_wfp4a16; @@ -3947,10 +3668,7 @@ void CutlassMoeFCRunner< } auto const* input_fp8_dequant = quant_params.fp8.dequant_input; - auto const* fc2_wfp4afp8_quant_scale = - (use_wfp4afp8 && Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0) - ? nullptr - : quant_params.fp8_mxfp4.fc2.act_global_scale; + auto const* fc2_wfp4afp8_quant_scale = quant_params.fp8_mxfp4.fc2.act_global_scale; auto const* fc2_expert_biases = reinterpret_cast(fc2_expert_biases_void); auto* final_output = static_cast(final_output_void); @@ -4112,7 +3830,7 @@ void CutlassMoeFCRunner< reinterpret_cast(input_activations_void), fc1_result_, glu_inter_result_, expert_first_token_offset_, gemm1_tma_ws_input, fc1_expert_weights, fc1_expert_biases, num_valid_tokens_ptr, fc1_int_scales, fc1_fp8_dequant, - use_wfp4afp8 ? fc2_wfp4afp8_quant_scale : fc2_fp8_quant, act_fp8_token_scale_, + use_wfp4afp8 ? fc2_wfp4afp8_quant_scale : fc2_fp8_quant, input_sf /*input fp4 scale or expanded fp4 scale*/, fc2_fp4_act_scale_, quant_params, num_rows, expanded_num_rows, hidden_size, inter_size, num_experts_per_node, fc1_activation_type, alpha_scale_ptr_array_fc1_, !use_lora, @@ -4137,7 +3855,7 @@ void CutlassMoeFCRunner< sync_check_cuda_error(stream); } else { bool fused_prologue_result = false; - if (!use_sm90_mixed_input_gemm) { + if (!use_w4_groupwise) { // WAR: fusedBuildExpertMapsSortFirstToken kernel will lead to illegal memory access for // W4AFP8 fused_prologue_result = fusedBuildExpertMapsSortFirstToken( @@ -4178,22 +3896,12 @@ void CutlassMoeFCRunner< bool use_per_expert_act_scale = use_fp4 ? quant_params.fp4.fc1.use_per_expert_act_scale : false; T* gemm1_input_expand = use_w4afp8 ? reinterpret_cast(smoothed_act_) : reinterpret_cast(permuted_data_); - float* fp8_token_dequant_scale = nullptr; - float const* fp8_token_residual_scale = nullptr; - float const** act_fp8_token_scale_ptr_array = nullptr; - if constexpr (use_wfp4afp8 && Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0) { - TLLM_CHECK(quant_params.fp8_mxfp4.fc1.global_scale != nullptr); - fp8_token_dequant_scale = act_fp8_token_scale_; - fp8_token_residual_scale = quant_params.fp8_mxfp4.fc1.global_scale; - act_fp8_token_scale_ptr_array = alpha_scale_ptr_array_fc1_; - } expandInputRowsKernelLauncher( input_activations, gemm1_input_expand, token_topk_unpermuted_scales, permuted_token_final_scales_, permuted_row_to_unpermuted_row_, num_rows, hidden_size, experts_per_token, num_experts_per_node, quant_params, use_per_expert_act_scale, expert_first_token_offset_, fc1_fp4_act_scale_, input_sf, swizzled_input_sf, (use_w4afp8 && !use_fp8_input) ? quant_params.groupwise.fc1.act_scales : nullptr, - fp8_token_dequant_scale, fp8_token_residual_scale, act_fp8_token_scale_ptr_array, enable_pdl, stream); auto const* gemm1_input = gemm1_input_expand; @@ -4232,10 +3940,10 @@ void CutlassMoeFCRunner< glu_inter_result_, expert_first_token_offset_, gemm1_tma_ws_input, fc1_expert_weights, fc1_expert_biases, num_valid_tokens_ptr, fc1_int_scales, fc1_fp8_dequant, use_wfp4afp8 ? fc2_wfp4afp8_quant_scale : fc2_fp8_quant, - act_fp8_token_scale_, fc1_fp4_act_scale_, fc2_fp4_act_scale_, quant_params, - num_rows, expanded_num_rows, hidden_size, inter_size, num_experts_per_node, - fc1_activation_type, alpha_scale_ptr_array_fc1_, !use_lora, stream, *gemm1_config_, - false, nullptr, nullptr, enable_pdl); + fc1_fp4_act_scale_, fc2_fp4_act_scale_, quant_params, num_rows, expanded_num_rows, + hidden_size, inter_size, num_experts_per_node, fc1_activation_type, + alpha_scale_ptr_array_fc1_, !use_lora, stream, *gemm1_config_, false, nullptr, + nullptr, enable_pdl); sync_check_cuda_error(stream); if (use_lora) { @@ -4262,10 +3970,9 @@ void CutlassMoeFCRunner< } template + bool IsMXFPX, class Enable> std::pair -CutlassMoeFCRunner:: +CutlassMoeFCRunner:: computeStridesTmaWarpSpecialized( int64_t const* expert_first_token_offset, TmaWarpSpecializedGroupedGemmInput layout_info1, TmaWarpSpecializedGroupedGemmInput layout_info2, int64_t num_tokens, @@ -4290,28 +3997,23 @@ CutlassMoeFCRunner + bool IsMXFPX, class Enable> std::pair -CutlassMoeFCRunner:: +CutlassMoeFCRunner:: computeStridesTmaWarpSpecializedLowLatency( TmaWarpSpecializedGroupedGemmInput layout_info1, TmaWarpSpecializedGroupedGemmInput layout_info2, int64_t num_tokens, int64_t gemm1_n, @@ -4366,10 +4067,9 @@ CutlassMoeFCRunner + bool IsMXFPX, class Enable> std::pair -CutlassMoeFCRunner:: +CutlassMoeFCRunner:: setupTmaWarpSpecializedInputs(int64_t num_rows, int64_t expanded_num_rows, ActivationParams fc1_activation_type, int64_t hidden_size, int64_t unpadded_hidden_size, int64_t inter_size, @@ -4432,19 +4132,17 @@ CutlassMoeFCRunnerswap_ab; gemm2_tma_ws_input.swap_ab = gemm2_config_->swap_ab; - gemm1_tma_ws_input.precomputed_scheduler_total_routed_tokens = expanded_num_rows; - gemm2_tma_ws_input.precomputed_scheduler_total_routed_tokens = expanded_num_rows; TLLM_CHECK_WITH_INFO( - (gemm1_tma_ws_input.swap_ab && gemm2_tma_ws_input.swap_ab) || !use_sm90_mixed_input_gemm, - "Hopper mixed-input grouped GEMM requires swap_ab"); + (gemm1_tma_ws_input.swap_ab && gemm2_tma_ws_input.swap_ab) || !use_w4_groupwise, + "Hopper w4 mixed input groupwise requires swap_ab"); bool apply_bias = parallelism_config.tp_rank == 0; auto* fc2_bias = apply_bias ? fc2_expert_biases : nullptr; bool gemm2_using_finalize_fusion = gemm2_config_->epilogue_fusion_type == cutlass_extensions::CutlassGemmConfig::EpilogueFusionType::FINALIZE; - bool using_fused_finalize = use_fused_finalize_ && gemm2_using_finalize_fusion && - !use_sm90_mixed_input_gemm && !use_lora; + bool using_fused_finalize = + use_fused_finalize_ && gemm2_using_finalize_fusion && !use_w4_groupwise && !use_lora; TLLM_CHECK_WITH_INFO( using_fused_finalize == gemm2_using_finalize_fusion, "GEMM2 tactic requests finalize fusion, but the runner is not configured to use it"); @@ -4456,6 +4154,34 @@ CutlassMoeFCRunner= 12080 + __nv_fp8_e8m0 tmp; + tmp.__x = __nv_cvt_float_to_e8m0(1.0f, __NV_SATFINITE, cudaRoundPosInf); + std::memcpy(&weight_block_scale_value_int, &tmp, sizeof(tmp)); +#endif + + auto act_sf_rows = std::min(expanded_num_rows, num_rows * num_experts_per_node); + auto fc1_sf_offset = + getOffsetActivationSF(num_experts_per_node, act_sf_rows, hidden_size, + TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX); + auto fc2_sf_offset = + getOffsetActivationSF(num_experts_per_node, act_sf_rows, inter_size, + TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX); + auto max_size = std::max(fc1_sf_offset, fc2_sf_offset) * + sizeof(TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF); + check_cuda_error( + cudaMemsetAsync(fc1_fp4_act_scale_, weight_block_scale_value_int, max_size, stream)); + } + TLLM_CHECK_WITH_INFO(gemm1_input != gemm1_output, "Input and output buffers are overlapping"); return Self::computeStridesTmaWarpSpecialized( expert_first_token_offset_, gemm1_tma_ws_input, gemm2_tma_ws_input, num_rows, @@ -4679,13 +4405,12 @@ std::map> GemmProfilerBackend::getProfile mGroupSize > 0; bool is_fp8_act_quant = mDType == nvinfer1::DataType::kFP8; bool is_fp8_w_quant = mWType == nvinfer1::DataType::kFP8; - bool const is_native_wfp4afp8_family = isNativeWfp4Afp8Family(); - // This predicate identifies the SM90 FP8 activation x packed-MXFP4 storage - // family. Sm90Wfp4Afp8ScaleMode selects Humming/pre-MMA vs future post-MMA - // semantics; do not infer the semantic path from dtype/layout alone. - bool const is_sm90_wfp4afp8_family = isSm90Wfp4Afp8Family(); - checkSm90Wfp4Afp8ScaleMode(); - bool is_w4afp8_quant = is_int_groupwise_w_quant && is_fp8_act_quant && !is_sm90_wfp4afp8_family; + // nvllm still uses int64 because torch doesn't have fp4 yet. + // bool is_fp4_act_quant = mDType == nvinfer1::DataType::kFP4 || mDType == + // nvinfer1::DataType::kINT64; + bool is_fp4_w_quant = mWType == nvinfer1::DataType::kFP4 || mWType == nvinfer1::DataType::kINT64; + bool is_w4afp8_quant = is_int_groupwise_w_quant && is_fp8_act_quant; + // bool is_wfp4afp8_quant = is_fp4_w_quant && is_fp8_act_quant; bool is_wfp4a16_quant = (mDType == nvinfer1::DataType::kHALF || mDType == nvinfer1::DataType::kBF16) && mWType == nvinfer1::DataType::kUINT8; @@ -4696,7 +4421,7 @@ std::map> GemmProfilerBackend::getProfile if (is_int_w_quant) { quant_1_size = fc1_out_size * num_experts_per_node * dtype_bytes; quant_2_size = hidden_size * num_experts_per_node * dtype_bytes; - } else if ((is_int_groupwise_w_quant && !is_sm90_wfp4afp8_family) || is_wfp4a16_quant) { + } else if (is_int_groupwise_w_quant || is_wfp4a16_quant) { quant_1_size = fc1_out_size * num_experts_per_node * dtype_bytes * hidden_size / mGroupSize; quant_2_size = hidden_size * num_experts_per_node * dtype_bytes * inter_size / mGroupSize; } @@ -4712,47 +4437,18 @@ std::map> GemmProfilerBackend::getProfile } // FP4 sizes - bool const use_humming_pre_mma = isHummingPreMmaScaleMode(); - size_t const fp8_mxfp4_token_scale_size = num_expanded_tokens * sizeof(float); - bool const is_nvfp4_quant = - mSM >= 100 && (mDType == nvinfer1::DataType::kFP4 || mDType == nvinfer1::DataType::kINT64) && - (mWType == nvinfer1::DataType::kFP4 || mWType == nvinfer1::DataType::kINT64); - size_t quant_5_size = 0; - size_t quant_6_size = 0; - if (is_nvfp4_quant) { - quant_1_size = sizeof(float); - quant_2_size = getOffsetWeightSF(num_experts_per_node, inter_size, hidden_size, mScalingType) * - sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF); - quant_3_size = num_experts_per_node * sizeof(float); - quant_4_size = sizeof(float); - quant_5_size = getOffsetWeightSF(num_experts_per_node, hidden_size, inter_size, mScalingType) * - sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF); - quant_6_size = num_experts_per_node * sizeof(float); - } else if (is_native_wfp4afp8_family) { - quant_1_size = sizeof(float); - quant_2_size = getOffsetWeightSF(num_experts_per_node, inter_size, hidden_size, mScalingType) * - sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF); - quant_3_size = num_experts_per_node * sizeof(float); - quant_4_size = sizeof(float); - quant_5_size = getOffsetWeightSF(num_experts_per_node, hidden_size, inter_size, mScalingType) * - sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF); - quant_6_size = num_experts_per_node * sizeof(float); - } else if (is_sm90_wfp4afp8_family) { - quant_1_size = sizeof(float); - quant_2_size = - getOffsetWeightSF(num_experts_per_node, fc1_out_size, hidden_size, - TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX) * - sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF); - quant_3_size = - use_humming_pre_mma ? fp8_mxfp4_token_scale_size : num_experts_per_node * sizeof(float); - quant_4_size = sizeof(float); - quant_5_size = - getOffsetWeightSF(num_experts_per_node, hidden_size, inter_size, - TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX) * - sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF); - quant_6_size = - use_humming_pre_mma ? fp8_mxfp4_token_scale_size : num_experts_per_node * sizeof(float); - } + quant_1_size = is_fp4_w_quant ? sizeof(float) : quant_1_size; + quant_2_size = is_fp4_w_quant ? getOffsetWeightSF(num_experts_per_node, inter_size, hidden_size, + mScalingType) * + sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF) + : quant_2_size; + quant_3_size = is_fp4_w_quant ? num_experts_per_node * sizeof(float) : quant_3_size; + quant_4_size = is_fp4_w_quant ? sizeof(float) : quant_4_size; + size_t quant_5_size = is_fp4_w_quant ? getOffsetWeightSF(num_experts_per_node, hidden_size, + inter_size, mScalingType) * + sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF) + : 0; + size_t quant_6_size = is_fp4_w_quant ? num_experts_per_node * sizeof(float) : 0; // MXFP8xMXFP8 sizes: the FP8 branch above only reserves per-tensor float // scalars, but block-scaled MXFP8 needs per-expert weight block SFs and @@ -4801,15 +4497,7 @@ std::map> GemmProfilerBackend::getProfile size_t w4a8_alpha_size = (is_w4afp8_quant || is_wfp4a16_quant) ? num_experts_per_node * sizeof(float) : 0; size_t alpha_scale_ptr_array_size = num_experts_per_node * sizeof(float**); - if (use_humming_pre_mma) { - alpha_scale_ptr_array_size *= NUM_ROUTING_SAMPLES; - } size_t gemm_workspace_size = mInterface->getGemmWorkspaceSize(num_experts_per_node); - size_t precomputed_scheduler_workspace_size = - (is_tma_ws_input && isSm90MixedInputFamily()) - ? cutlass_kernels_oss::detail::precomputed_scheduler_workspace_size( - num_experts_per_node, num_expanded_tokens, std::max(hidden_size, fc1_out_size)) - : 0; // Routing info size_t expert_first_token_offset_size = @@ -4820,6 +4508,7 @@ std::map> GemmProfilerBackend::getProfile size_t permuted_size = mMinLatencyMode ? 0 : num_expanded_tokens * sizeof(int); size_t token_topk_unpermuted_scales_size = mMinLatencyMode ? 0 : num_expanded_tokens * sizeof(float); + int64_t const num_tokens_per_block = computeNumTokensPerBlock(maxM, num_experts_per_node); int64_t const num_blocks_per_seq = tensorrt_llm::common::ceilDiv(maxM, num_tokens_per_block); size_t const blocked_expert_counts_size = @@ -4861,8 +4550,6 @@ std::map> GemmProfilerBackend::getProfile ADD(blocked_expert_counts_cumsum); ADD(blocked_row_to_unpermuted_row); ADD(token_topk_unpermuted_scales); - ADD_NAME(act_fp8_token_scale, - isHummingPreMmaScaleMode() ? num_expanded_tokens * sizeof(float) : 0); ADD(num_active_experts_per_node); ADD(active_expert_global_ids); ADD(input); @@ -4880,7 +4567,6 @@ std::map> GemmProfilerBackend::getProfile ADD(w4a8_alpha); ADD(alpha_scale_ptr_array); ADD(fp4_act_scale_flat); - ADD(precomputed_scheduler_workspace); ADD(gemm_workspace); ADD(swiglu_alpha); ADD(swiglu_beta); @@ -4971,19 +4657,9 @@ void GemmProfilerBackend::prepareQuantParams(int num_tokens, char* workspace_ptr GET_WS_PTR(float const*, w4a8_alpha); #undef GET_WS_PTR - checkSm90Wfp4Afp8ScaleMode(); - - if (isSm90Wfp4Afp8Family()) { - TLLM_CHECK(quant_2 && quant_3 && quant_4 && quant_5 && quant_6); - mQuantParams = QuantParams::FP8MXFP4( - nullptr, static_cast(quant_2), - static_cast(quant_3), - isHummingPreMmaScaleMode() ? nullptr : static_cast(quant_4), - static_cast(quant_5), - static_cast(quant_6), false, false); - } else if ((mWType == nvinfer1::DataType::kINT8 || mWType == nvinfer1::DataType::kINT4 || - mWType == nvinfer1::DataType::kUINT8) && - mGroupSize < 0) { + if ((mWType == nvinfer1::DataType::kINT8 || mWType == nvinfer1::DataType::kINT4 || + mWType == nvinfer1::DataType::kUINT8) && + mGroupSize < 0) { TLLM_CHECK(quant_1 && quant_2); mQuantParams = QuantParams::Int(quant_1, quant_2); } else if (mWType == nvinfer1::DataType::kINT4 || mWType == nvinfer1::DataType::kUINT8) { @@ -5061,17 +4737,21 @@ void GemmProfilerBackend::prepareTmaWsInputs( return; } - bool const use_sm90_mixed_input_gemm = isSm90MixedInputFamily(); + bool use_w4afp8 = (mDType == nvinfer1::DataType::kFP8 && mWType == nvinfer1::DataType::kINT4); + bool use_wfp4a16 = + ((mDType == nvinfer1::DataType::kHALF || mDType == nvinfer1::DataType::kBF16) && + mWType == nvinfer1::DataType::kUINT8); + bool use_w4_groupwise = use_w4afp8 || use_wfp4a16; bool const use_finalize_fusion = fusion == TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE; bool const finalize_fusion_not_supported = !mInterface->use_fused_finalize_ || mMinLatencyMode || - use_sm90_mixed_input_gemm || + use_w4_groupwise || mGemmToProfile != GemmToProfile::GEMM_2; if (use_finalize_fusion && finalize_fusion_not_supported) { return; } - if (use_sm90_mixed_input_gemm && !swap_ab) { + if (use_w4_groupwise && !swap_ab) { return; } @@ -5096,7 +4776,6 @@ void GemmProfilerBackend::prepareTmaWsInputs( GET_WS_PTR(float*, token_topk_unpermuted_scales); GET_WS_PTR(int8_t*, tma_ws_input_workspace); GET_WS_PTR(void*, gemm_workspace); - GET_WS_PTR(void*, precomputed_scheduler_workspace); GET_WS_PTR(float*, alpha_scale_ptr_array); GET_WS_PTR(TmaWarpSpecializedGroupedGemmInput::ElementSF*, fp4_act_scale_flat); GET_WS_PTR(int*, num_active_experts_per_node); @@ -5118,10 +4797,8 @@ void GemmProfilerBackend::prepareTmaWsInputs( TmaWarpSpecializedGroupedGemmInput::workspaceSize(mNumExpertsPerNode, mScalingType); TmaWarpSpecializedGroupedGemmInput dummy_tma_ws_input; - dummy_tma_ws_input.configureWorkspace( - tma_ws_input_workspace, mNumExpertsPerNode, gemm_workspace, - workspaces.at("gemm_workspace").first, precomputed_scheduler_workspace, - workspaces.at("precomputed_scheduler_workspace").first, mScalingType); + dummy_tma_ws_input.configureWorkspace(tma_ws_input_workspace, mNumExpertsPerNode, gemm_workspace, + workspaces.at("gemm_workspace").first, mScalingType); dummy_tma_ws_input.enable_pdl = enable_pdl; // Set enable_pdl for dummy input tma_ws_input_workspace += tma_ws_size; @@ -5131,17 +4808,13 @@ void GemmProfilerBackend::prepareTmaWsInputs( tma_ws_input_workspace += workspace_index * tma_ws_size; size_t num_expanded_tokens = num_tokens * mK; - dummy_tma_ws_input.precomputed_scheduler_total_routed_tokens = num_expanded_tokens; for (int64_t i = 0; i < NUM_ROUTING_SAMPLES; i++) { // Note: Even though we have separate TMA WS inputs for finalize fusion on/off we reuse the same // pointers to save space. auto& cache_element = mTmaInputCache[use_finalize_fusion][swap_ab][i]; - cache_element.configureWorkspace( - tma_ws_input_workspace, mNumExpertsPerNode, gemm_workspace, - workspaces.at("gemm_workspace").first, precomputed_scheduler_workspace, - workspaces.at("precomputed_scheduler_workspace").first, mScalingType); + cache_element.configureWorkspace(tma_ws_input_workspace, mNumExpertsPerNode, gemm_workspace, + workspaces.at("gemm_workspace").first, mScalingType); cache_element.enable_pdl = enable_pdl; // Set enable_pdl for cache element - cache_element.precomputed_scheduler_total_routed_tokens = num_expanded_tokens; tma_ws_input_workspace += tma_ws_size; int64_t* expert_first_token_offset = @@ -5205,19 +4878,6 @@ void GemmProfilerBackend::prepare(int num_tokens, char* workspace_ptr_char, prepareRouting(num_tokens, workspace_ptr_char, enable_pdl, stream); prepareQuantParams(num_tokens, workspace_ptr_char, stream); - if (isHummingPreMmaScaleMode()) { - auto workspaces = getProfilerWorkspaces(num_tokens, mSM >= 90); - auto* expert_first_token_offset = reinterpret_cast( - workspace_ptr_char + workspaces.at("expert_first_token_offset").second); - auto* act_fp8_token_scale = - reinterpret_cast(workspace_ptr_char + workspaces.at("act_fp8_token_scale").second); - auto* alpha_scale_ptr_array = reinterpret_cast( - workspace_ptr_char + workspaces.at("alpha_scale_ptr_array").second); - prepareProfilerFP8TokenScalePtrArray(alpha_scale_ptr_array, act_fp8_token_scale, - expert_first_token_offset, mNumExpertsPerNode, - NUM_ROUTING_SAMPLES, stream); - sync_check_cuda_error(stream); - } for (auto fusion : {TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::NONE, TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE}) { for (auto swap_ab : {false, true}) { @@ -5267,7 +4927,6 @@ void GemmProfilerBackend::runProfiler(int original_num_tokens, Config const& tac GET_WS_PTR(float const*, token_topk_unpermuted_scales); auto const* token_topk_permuted_scales = token_topk_unpermuted_scales; - GET_WS_PTR(float*, act_fp8_token_scale); GET_WS_PTR_OFFSET(int*, num_active_experts_per_node, mSampleIndex); GET_WS_PTR_OFFSET(int*, active_expert_global_ids, (mSampleIndex * mNumExpertsPerNode)); @@ -5279,8 +4938,7 @@ void GemmProfilerBackend::runProfiler(int original_num_tokens, Config const& tac void const* weights_sel = mNeedWeights ? weights : expert_weights; GET_WS_PTR(void const*, bias); - GET_WS_PTR_OFFSET(float const**, alpha_scale_ptr_array, - (isHummingPreMmaScaleMode() ? mSampleIndex * mNumExpertsPerNode : 0)); + GET_WS_PTR(float const**, alpha_scale_ptr_array); GET_WS_PTR(TmaWarpSpecializedGroupedGemmInput::ElementSF*, fp4_act_scale_flat); GET_WS_PTR(void*, gemm_workspace); @@ -5291,14 +4949,17 @@ void GemmProfilerBackend::runProfiler(int original_num_tokens, Config const& tac #undef GET_WS_PTR_OFFSET #undef GET_WS_PTR - checkSm90Wfp4Afp8ScaleMode(); - TmaWarpSpecializedGroupedGemmInput tma_ws_input_template; if (tactic.is_tma_warp_specialized) { // Use non-finalize cache when finalize fusion is not supported for the current GEMM + bool use_w4afp8 = (mDType == nvinfer1::DataType::kFP8 && mWType == nvinfer1::DataType::kINT4); + bool use_wfp4a16 = + ((mDType == nvinfer1::DataType::kHALF || mDType == nvinfer1::DataType::kBF16) && + mWType == nvinfer1::DataType::kUINT8); + bool use_w4_groupwise = use_w4afp8 || use_wfp4a16; bool finalize_supported_this_gemm = (mGemmToProfile == GemmToProfile::GEMM_2) && mInterface->use_fused_finalize_ && !mMinLatencyMode && - !isSm90MixedInputFamily(); + !use_w4_groupwise; bool request_finalize = tactic.epilogue_fusion_type == cutlass_extensions::CutlassGemmConfig::EpilogueFusionType::FINALIZE; bool use_finalize_index = request_finalize && finalize_supported_this_gemm; @@ -5309,11 +4970,6 @@ void GemmProfilerBackend::runProfiler(int original_num_tokens, Config const& tac } mInterface->is_profiler = true; - auto const* profiler_fc2_activation_quant_scale = - isHummingPreMmaScaleMode() ? nullptr - : (mQuantParams.fp8_mxfp4.fc2.act_global_scale - ? mQuantParams.fp8_mxfp4.fc2.act_global_scale - : mQuantParams.fp8.quant_fc2); if (mGemmToProfile == GemmToProfile::GEMM_1) { mInterface->gemm1(input, // output, // @@ -5325,16 +4981,17 @@ void GemmProfilerBackend::runProfiler(int original_num_tokens, Config const& tac expert_first_token_offset + num_experts_per_node, // mQuantParams.wo.fc1_weight_scales, // mQuantParams.fp8.dequant_fc1, // - profiler_fc2_activation_quant_scale, // - act_fp8_token_scale, // - fp4_act_scale_flat, // - fp4_act_scale_flat, // - mQuantParams, // - original_num_tokens, // - expanded_num_tokens, // - mExpertHiddenSize, // - mExpertInterSize, // - num_experts_per_node, // + mQuantParams.fp8_mxfp4.fc2.act_global_scale + ? mQuantParams.fp8_mxfp4.fc2.act_global_scale + : mQuantParams.fp8.quant_fc2, // + fp4_act_scale_flat, // + fp4_act_scale_flat, // + mQuantParams, // + original_num_tokens, // + expanded_num_tokens, // + mExpertHiddenSize, // + mExpertInterSize, // + num_experts_per_node, // ActivationParams(mActivationType, swiglu_alpha, swiglu_beta, swiglu_limit), alpha_scale_ptr_array, // !mUseLora, // diff --git a/csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu b/csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu index 0a0847775cc..8ba445af69b 100644 --- a/csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu +++ b/csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu @@ -83,9 +83,7 @@ class DtypeUtils { class FusedMoeRunner : public tvm::ffi::ModuleObj { public: - template < - typename TypeAct, typename TypeWeight, bool NeedQuant = false, bool IsMXFPX = false, - kernels::Sm90Wfp4Afp8ScaleMode Sm90Wfp4Afp8Mode = kernels::Sm90Wfp4Afp8ScaleMode::kDisabled> + template std::unique_ptr switch_output_type(DLDataType output_type) { switch (encode_dlpack_dtype(output_type)) { case int64_code: // INT64 == FP4 @@ -97,22 +95,20 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { // return std::make_unique>(); case float16_code: if constexpr (NeedQuant) { - return std::make_unique>(); + return std::make_unique< + kernels::CutlassMoeFCRunner>(); } else { - return std::make_unique>(); + return std::make_unique< + kernels::CutlassMoeFCRunner>(); } #ifdef ENABLE_BF16 case bfloat16_code: if constexpr (NeedQuant) { - return std::make_unique< - kernels::CutlassMoeFCRunner>(); + return std::make_unique>(); } else { - return std::make_unique< - kernels::CutlassMoeFCRunner>(); + return std::make_unique>(); } #endif default: @@ -125,8 +121,7 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { FusedMoeRunner(DLDataType activation_dtype, DLDataType weight_dtype, DLDataType output_dtype, bool use_deepseek_fp8_block_scale, bool use_w4_group_scaling, - bool use_mxfp8_act_scaling, bool use_packed_weights, bool use_fused_finalize, - bool use_wfp4afp8_humming) { + bool use_mxfp8_act_scaling, bool use_packed_weights, bool use_fused_finalize) { mActivationDtype = activation_dtype; mWeightDtype = weight_dtype; mUsePackedWeights = use_packed_weights; @@ -135,24 +130,8 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { mUseW4GroupScaling = use_w4_group_scaling; mUseMxfp8ActScaling = use_mxfp8_act_scaling; mUseFusedFinalize = use_fused_finalize; - mUseWfp4Afp8Humming = use_wfp4afp8_humming; - mSm90Wfp4Afp8Mode = kernels::Sm90Wfp4Afp8ScaleMode::kDisabled; mInnerDimMultiplier = 1; - auto make_humming_runner = [&] { - mInnerDimMultiplier = 2; - mSm90Wfp4Afp8Mode = kernels::Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0; - TVM_FFI_ICHECK(mActivationDtype == dl_float16 || mActivationDtype == dl_bfloat16) - << "Humming-style MXFP4 x FP8 requires FP16/BF16 inputs and online FP8 activation " - "quantization."; - TVM_FFI_ICHECK(mActivationDtype == mOutputDtype) - << "Humming-style MXFP4 x FP8 online activation quantization currently requires " - "activation dtype and output dtype to match."; - mKernelRunner = - switch_output_type<__nv_fp8_e4m3, kernels::Fp4Type, true, false, - kernels::Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0>(mOutputDtype); - }; - // keep consistent with cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp if (mActivationDtype == dl_float16 && mWeightDtype == dl_float16) { mKernelRunner = std::make_shared>(); @@ -176,37 +155,9 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { } #endif #ifdef ENABLE_FP4 - int const sm = common::getSMVersion(); - if (sm >= 100 && (isWMxfp4AMxfp8Quant() || isWMxfp4AFp8Quant())) { + if (isWMxfp4AMxfp8Quant() || isWMxfp4AFp8Quant()) { mInnerDimMultiplier = 16; // 16 FP4 -> 1 LONG - mKernelRunner = switch_output_type<__nv_fp8_e4m3, kernels::Fp4Type>(mOutputDtype); - } - -#if 0 - // PHASE3_POST_MMA_PLACEHOLDER: future SM90 post-MMA MXFP4 paths should use - // uint8-packed MXFP4 storage and select their Sm90Wfp4Afp8ScaleMode explicitly. - // Enabling these paths also requires updating the shared predicates for - // the SM90 uint8-packed input contract. - if (sm == 90 && isWMxfp4AFp8Quant()) { - mInnerDimMultiplier = 2; - mSm90Wfp4Afp8Mode = kernels::Sm90Wfp4Afp8ScaleMode::kPostMmaFp8Act; - mKernelRunner = switch_output_type< - __nv_fp8_e4m3, kernels::Fp4Type, false, false, - kernels::Sm90Wfp4Afp8ScaleMode::kPostMmaFp8Act>(mOutputDtype); - } - - if (sm == 90 && isWMxfp4AMxfp8Quant()) { - mInnerDimMultiplier = 2; - mSm90Wfp4Afp8Mode = kernels::Sm90Wfp4Afp8ScaleMode::kPostMmaMxfp8Act; - mKernelRunner = switch_output_type< - __nv_fp8_e4m3, kernels::Fp4Type, false, false, - kernels::Sm90Wfp4Afp8ScaleMode::kPostMmaMxfp8Act>(mOutputDtype); - } -#endif - - if (isWMxfp4AFp8HummingQuant()) { - TVM_FFI_ICHECK_EQ(sm, 90) << "Humming-style MXFP4 x FP8 is only supported on SM90."; - make_humming_runner(); + mKernelRunner = switch_output_type<__nv_fp8_e4m3, __nv_fp4_e2m1>(mOutputDtype); } if (isNvfp4Quant()) { @@ -216,23 +167,23 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { #ifdef ENABLE_BF16 case bfloat16_code: #endif - mKernelRunner = - switch_output_type(mOutputDtype); + mKernelRunner = switch_output_type<__nv_fp4_e2m1, __nv_fp4_e2m1, true>(mOutputDtype); break; default: - mKernelRunner = - switch_output_type(mOutputDtype); + mKernelRunner = switch_output_type<__nv_fp4_e2m1, __nv_fp4_e2m1, false>(mOutputDtype); } } if (isWFP4A16Quant()) { - TVM_FFI_ICHECK_EQ(mActivationDtype, dl_bfloat16) - << "SM90 MXFP4 W4A16 supports BF16 activations only; FP16 is incompatible with the " - "interleaved weight layout."; mInnerDimMultiplier = 2; + if (mActivationDtype == dl_float16) { + mKernelRunner = std::make_shared>(); + } #ifdef ENABLE_BF16 - mKernelRunner = - std::make_shared>(); + else if (mActivationDtype == dl_bfloat16) { + mKernelRunner = + std::make_shared>(); + } #endif } @@ -376,7 +327,7 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { int64_t hidden_size = fc2_expert_weights.size(1); int64_t inter_size = fc2_expert_weights.size(2) * mInnerDimMultiplier; - if (isWMxfp4AMxfp8Quant() || isWMxfp4AFp8Quant() || isWMxfp4AFp8HummingQuant()) { + if (isWMxfp4AMxfp8Quant() || isWMxfp4AFp8Quant()) { // MXFP4 weights are required to bealigned to 128 bytes TVM_FFI_ICHECK_EQ(hidden_size % 128, 0) << "hidden_size must be divisible by 128 for MXFP4 weights"; @@ -436,9 +387,8 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { static_cast(experts_per_token), base_activation_type, parallelism_config, min_latency_mode, input.device(), workspace_buffer); - int64_t const routed_tokens = input.size(0) * token_selected_experts.size(1); auto const quant_params = getQuantParams(num_experts_on_rank, hidden_size, inter_size, - routed_tokens, quant_scales, base_activation_type); + quant_scales, base_activation_type); kernels::MoeMinLatencyParams min_latency_params{}; // TODO: support lora in the future @@ -628,9 +578,8 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { static_cast(experts_per_token), base_activation_type, parallelism_config, min_latency_mode, input.device(), workspace_buffer); - int64_t const routed_tokens = input.size(0) * token_selected_experts.size(1); auto const quant_params = getQuantParams(num_experts_on_rank, hidden_size, inter_size, - routed_tokens, quant_scales, base_activation_type); + quant_scales, base_activation_type); // TODO: support lora in the future ::tensorrt_llm::kernels::LoraParams lora_params{}; @@ -699,7 +648,7 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { isInt4Quant() ? TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::int4_group_size : -1; int64_t group_size = - (isWFP4A16Quant() || isWMxfp4AFp8HummingQuant()) + isWFP4A16Quant() ? TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::wfp4a16_group_size : group_size_; int const num_experts = static_cast(fc2_expert_weights.size(0) * ep_size); @@ -733,22 +682,19 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { activation_dtype = isNvfp4Quant() ? dl_int64 : activation_dtype; int64_t const unpadded_hidden_size_profiler = hidden_size; // HACK no padding by default #ifdef USING_OSS_CUTLASS_MOE_GEMM - mProfiler->init(*mKernelRunner.get(), mProfiler->mGemmToProfile, - DtypeUtils::dataType(activation_dtype), DtypeUtils::dataType(mWeightDtype), - DtypeUtils::dataType(mOutputDtype), num_experts, static_cast(top_k), - hidden_size, unpadded_hidden_size_profiler, inter_size, group_size, - activation_type, USE_BIAS, USE_LORA, min_latency_mode, - /*need_weights*/ false, parallelism_config, enable_alltoall, - mUseMxfp8ActScaling, mSm90Wfp4Afp8Mode); + mProfiler->init( + *mKernelRunner.get(), mProfiler->mGemmToProfile, DtypeUtils::dataType(activation_dtype), + DtypeUtils::dataType(mWeightDtype), DtypeUtils::dataType(mOutputDtype), num_experts, + static_cast(top_k), hidden_size, unpadded_hidden_size_profiler, inter_size, + group_size, activation_type, USE_BIAS, USE_LORA, min_latency_mode, + /*need_weights*/ false, parallelism_config, enable_alltoall, mUseMxfp8ActScaling); #else mProfiler->init(*mKernelRunner.get(), mProfiler->mGemmToProfile, DtypeUtils::dataType(activation_dtype), DtypeUtils::dataType(mWeightDtype), DtypeUtils::dataType(mOutputDtype), num_experts, static_cast(top_k), hidden_size, unpadded_hidden_size_profiler, inter_size, group_size, activation_type, USE_BIAS, USE_LORA, min_latency_mode, - /*need_weights*/ false, parallelism_config, - /*enable_alltoall*/ false, - /*use_mxfp8_act_scaling*/ false, mSm90Wfp4Afp8Mode); + /*need_weights*/ false, parallelism_config); #endif size_t profile_workspace_size = mProfiler->getWorkspaceSize(num_rows); @@ -799,12 +745,6 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { return static_cast( mKernelRunner->queryOccupancyForConfig(mAllProfiles[tactic_id])); }); - } else if (name == "get_valid_tactics_for_shape") { - return Function::FromTyped( - [this](int64_t stage, int64_t gemm_n, int64_t gemm_k) -> Array { - std::lock_guard lock(mMutex); - return getValidTacticsForShape(stage, gemm_n, gemm_k); - }); } else if (name == "run_moe") { return Function::FromTyped( [this](TensorView output, TensorView input, TensorView token_selected_experts, @@ -888,115 +828,14 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { bool mUseDeepSeekFP8BlockScaling = false; bool mUseW4GroupScaling = false; bool mUseMxfp8ActScaling = false; - bool mUseWfp4Afp8Humming = false; bool mUsePackedWeights = false; bool mUseFusedFinalize = true; - kernels::Sm90Wfp4Afp8ScaleMode mSm90Wfp4Afp8Mode = kernels::Sm90Wfp4Afp8ScaleMode::kDisabled; using Profile = tensorrt_llm::cutlass_extensions::CutlassGemmConfig; std::vector mAllProfiles; int64_t mGemm1TacticCount{0}; int64_t mGemm2TacticCount{0}; - bool isProfileShapeSupported(Profile const& profile, int64_t gemm_n, int64_t gemm_k) const { - int64_t tile_m = 0; - int64_t tile_n = 0; - int64_t tile_k = 0; - if (profile.sm_version == 90) { - auto const [m, n, k] = - tensorrt_llm::cutlass_extensions::enum_to_shape_tuple(profile.tile_config_sm90); - tile_m = m; - tile_n = n; - tile_k = k; - } else if (profile.sm_version == 100) { - auto const [m, n, k] = - tensorrt_llm::cutlass_extensions::enum_to_shape_tuple(profile.tile_config_sm100); - tile_m = m; - tile_n = n; - tile_k = k; - } else if (profile.sm_version == 120) { - auto const [m, n, k] = - tensorrt_llm::cutlass_extensions::enum_to_shape_tuple(profile.tile_config_sm120); - tile_m = m; - tile_n = n; - tile_k = k; - } - - if (tile_m <= 0 || tile_n <= 0 || tile_k <= 0 || gemm_n <= 0 || gemm_k <= 0) { - return false; - } - if (gemm_k < tile_k || gemm_k % tile_k != 0) { - return false; - } - if (gemm_n < tile_n) { - return false; - } - if (mUseW4GroupScaling && gemm_n % tile_m != 0) { - return false; - } - bool const is_single_warpgroup = - profile.mainloop_schedule == - tensorrt_llm::cutlass_extensions::MainloopScheduleType::SINGLE_WARPGROUP_PREFILL || - profile.mainloop_schedule == - tensorrt_llm::cutlass_extensions::MainloopScheduleType::SINGLE_WARPGROUP_ROLLING; - if (is_single_warpgroup) { - if (!mUseWfp4Afp8Humming || profile.sm_version != 90 || tile_m != 128 || tile_k != 128 || - (tile_n != 8 && tile_n != 16 && tile_n != 32 && tile_n != 40) || - profile.cluster_shape != - tensorrt_llm::cutlass_extensions::ClusterShape::ClusterShape_1x1x1 || - gemm_n % 128 != 0) { - return false; - } - if (profile.mainloop_schedule == - tensorrt_llm::cutlass_extensions::MainloopScheduleType::SINGLE_WARPGROUP_PREFILL) { - return gemm_k <= 384; - } - return gemm_k > 384; - } - if (isWFP4A16Quant()) { - if (tile_k == 256 && ((tile_m == 128 && tile_n == 256) || (tile_m == 256 && tile_n == 128))) { - return false; - } - if (tile_k == 512 && tile_n >= 128) { - return false; - } - } - return true; - } - - Array getValidTacticsForShape(int64_t stage, int64_t gemm_n, int64_t gemm_k) const { - int64_t begin = 0; - int64_t end = static_cast(mAllProfiles.size()); - if (stage == 1) { - end = mGemm1TacticCount; - } else if (stage == 2) { - begin = mGemm1TacticCount; - end = mGemm1TacticCount + mGemm2TacticCount; - } - - int64_t const total = static_cast(mAllProfiles.size()); - if (begin < 0) { - begin = 0; - } - if (begin > total) { - begin = total; - } - if (end < begin) { - end = begin; - } - if (end > total) { - end = total; - } - - Array tactics; - for (int64_t tactic_id = begin; tactic_id < end; ++tactic_id) { - if (!mUseW4GroupScaling || isProfileShapeSupported(mAllProfiles[tactic_id], gemm_n, gemm_k)) { - tactics.push_back(tactic_id); - } - } - return tactics; - } - void setRunnerProfiles(Optional> profile_ids) { if (mUseDeepSeekFP8BlockScaling) { auto config = tensorrt_llm::cutlass_extensions::CutlassGemmConfig( @@ -1023,7 +862,7 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { best_gemm1_profile = mAllProfiles.at(id1); } - // GEMM2 profiles use absolute indices in the combined profile array. + // GEMM2 index: accept absolute index and raise error if out of GEMM2 range auto id2 = profile_ids.value()[1]; if (id2 != -1) { TVM_FFI_ICHECK(id2 >= mGemm1TacticCount && id2 < mGemm1TacticCount + mGemm2TacticCount) @@ -1100,7 +939,7 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { } kernels::QuantParams getQuantParams( - int64_t num_experts_on_rank, int64_t hidden_size, int64_t inter_size, int64_t routed_tokens, + int64_t num_experts_on_rank, int64_t hidden_size, int64_t inter_size, Optional> quant_scales, ActivationType base_activation_type = ActivationType::Swiglu) const { if (isWMxfp8AMxfp8Quant()) { @@ -1207,9 +1046,9 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { fc2_quant.ndim() == 1); } else if (isWMxfp4AFp8Quant()) { TVM_FFI_ICHECK(quant_scales.has_value()) - << "Expecting quant scales for post-MMA MXFP4 x FP8 quantization"; + << "Expecting quant scales for W4A8_MXFP4_MXF8 quantization"; TVM_FFI_ICHECK_EQ(quant_scales.value().size(), 5) - << "Expecting 5 quant scales for post-MMA MXFP4 x FP8 quantization"; + << "Expecting 5 quant scales for W4A8_MXFP4_FP8 quantization"; auto const fc1_weight_block = quant_scales.value()[0]; auto const fc1_global = quant_scales.value()[1]; @@ -1271,56 +1110,6 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { static_cast(fc2_act_global.data_ptr()), static_cast(fc2_weight_block.data_ptr()), static_cast(fc2_global.data_ptr()), false, fc2_act_global.ndim() == 1); - } else if (isWMxfp4AFp8HummingQuant()) { - TVM_FFI_ICHECK(quant_scales.has_value()) - << "Expecting quant scales for Humming-style MXFP4 x FP8 quantization"; - TVM_FFI_ICHECK_EQ(quant_scales.value().size(), 5) - << "Expecting 5 quant scales for Humming-style MXFP4 x FP8 quantization"; - - auto const fc1_weight_block = quant_scales.value()[0]; - auto const fc1_token_scale = quant_scales.value()[1]; - auto const fc2_act_global = quant_scales.value()[2]; - auto const fc2_weight_block = quant_scales.value()[3]; - auto const fc2_token_scale = quant_scales.value()[4]; - - CHECK_INPUT_TYPE(fc1_weight_block, dl_int32); - CHECK_INPUT_TYPE(fc1_token_scale, dl_float32); - CHECK_INPUT_TYPE(fc2_act_global, dl_float32); - CHECK_INPUT_TYPE(fc2_weight_block, dl_int32); - CHECK_INPUT_TYPE(fc2_token_scale, dl_float32); - CHECK_DIM(5, fc1_weight_block); - CHECK_DIM(1, fc1_token_scale); - TVM_FFI_ICHECK_LE(fc2_act_global.ndim(), 1) - << "fc2 act global must be a scalar or 1-D tensor"; - CHECK_DIM(5, fc2_weight_block); - CHECK_DIM(1, fc2_token_scale); - int const fc1_n_mult = isGatedActivation(base_activation_type) ? 2 : 1; - TVM_FFI_ICHECK(fc1_weight_block.size(0) == num_experts_on_rank && - fc1_weight_block.size(1) * 64 == inter_size * fc1_n_mult && - fc1_weight_block.size(2) * 128 == hidden_size && - fc1_weight_block.size(3) == 16 && fc1_weight_block.size(4) == 4) - << "fc1 Humming-style folded weight scale must be " - "(num_experts_on_rank, inter_size" - << (fc1_n_mult == 2 ? " * 2" : "") << " / 64, hidden_size / 128, 16, 4)"; - TVM_FFI_ICHECK_EQ(fc1_token_scale.size(0), routed_tokens) - << "fc1 token scale must have one element per routed token"; - TVM_FFI_ICHECK(fc2_act_global.ndim() == 0 || fc2_act_global.size(0) == num_experts_on_rank) - << "fc2 act global must be scalar or (num_experts_on_rank,)"; - TVM_FFI_ICHECK(fc2_weight_block.size(0) == num_experts_on_rank && - fc2_weight_block.size(1) * 64 == hidden_size && - fc2_weight_block.size(2) * 128 == inter_size && - fc2_weight_block.size(3) == 16 && fc2_weight_block.size(4) == 4) - << "fc2 Humming-style folded weight scale must be " - "(num_experts_on_rank, hidden_size / 64, inter_size / 128, 16, 4)"; - TVM_FFI_ICHECK_EQ(fc2_token_scale.size(0), routed_tokens) - << "fc2 token scale must have one element per routed token"; - - return kernels::QuantParams::FP8MXFP4( - nullptr, - static_cast(fc1_weight_block.data_ptr()), - static_cast(fc1_token_scale.data_ptr()), nullptr, - static_cast(fc2_weight_block.data_ptr()), - static_cast(fc2_token_scale.data_ptr()), false, false); } else if (isWMxfp4AMxfp8Quant()) { #ifdef USING_OSS_CUTLASS_MOE_GEMM TVM_FFI_ICHECK(quant_scales.has_value()) @@ -1500,14 +1289,6 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { auto const& fc2_weight_zeros = quant_scales.value()[5]; auto const& fc1_alpha = quant_scales.value()[6]; auto const& fc2_alpha = quant_scales.value()[7]; - if (fc1_act_scales.numel() > 0) { - TVM_FFI_ICHECK_EQ(fc1_act_scales.numel(), hidden_size) - << "INT4xFP8 FC1 prequant scale must be shared across experts with shape [hidden_size]"; - } - if (fc2_act_scales.numel() > 0) { - TVM_FFI_ICHECK_EQ(fc2_act_scales.numel(), inter_size) - << "INT4xFP8 FC2 prequant scale must be shared across experts with shape [inter_size]"; - } int group_size = TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::int4_group_size; return kernels::QuantParams::GroupWise( group_size, static_cast(fc1_weight_scales.data_ptr()), @@ -1545,22 +1326,13 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { } bool isWFP4A16Quant() const { - return mUseW4GroupScaling && - (mActivationDtype == dl_float16 || mActivationDtype == dl_bfloat16) && - mWeightDtype == dl_uint8 && !mUsePackedWeights && !mUseWfp4Afp8Humming; + return mUseW4GroupScaling && mWeightDtype == dl_uint8 && !mUsePackedWeights; } bool isInt4Quant() const { return mWeightDtype == dl_uint8 && mUsePackedWeights; } bool isW4AFp8Quant() const { return mActivationDtype == dl_float8_e4m3fn && isInt4Quant(); } - bool isWMxfp4AFp8HummingQuant() const { - bool const supported_activation = - mActivationDtype == dl_float16 || mActivationDtype == dl_bfloat16; - return mUseWfp4Afp8Humming && mUseW4GroupScaling && supported_activation && - mWeightDtype == dl_uint8 && !mUsePackedWeights && !mUseMxfp8ActScaling; - } - bool isWMxfp4AFp8Quant() const { return mActivationDtype == dl_float8_e4m3fn && mWeightDtype == dl_int64 && !mUseMxfp8ActScaling; } @@ -1572,20 +1344,18 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { tvm::ffi::Module init(DLDataType activation_dtype, DLDataType weight_dtype, DLDataType output_dtype, bool use_deepseek_fp8_block_scale, bool use_w4_group_scaling, - bool use_mxfp8_act_scaling, bool use_packed_weights, bool use_fused_finalize, - bool use_wfp4afp8_humming) { + bool use_mxfp8_act_scaling, bool use_packed_weights, + bool use_fused_finalize) { auto ptr = tvm::ffi::make_object( activation_dtype, weight_dtype, output_dtype, use_deepseek_fp8_block_scale, - use_w4_group_scaling, use_mxfp8_act_scaling, use_packed_weights, use_fused_finalize, - use_wfp4afp8_humming); + use_w4_group_scaling, use_mxfp8_act_scaling, use_packed_weights, use_fused_finalize); return tvm::ffi::Module(ptr); } // Interleave a 4-bit packed weight tensor into the layout required by the // SM90 mixed-input MoE GEMM. Expected input shape (num_experts, n, // k / 2) uint8 on CUDA. Writes into an output tensor of the same shape. -// quant_type: 0 for INT4 (W4A8), 1 for FP4 (W4A16 / MXFP4 BF16), -// 2 for FP4 consumed by FP8/Humming-style pre-MMA scaling. +// quant_type: 0 for INT4 (W4A8), 1 for FP4 (W4A16 / MXFP4). void interleave_moe_weights_for_sm90_mixed_gemm(TensorView weight, TensorView weight_interleaved, int64_t quant_type) { CHECK_INPUT_TYPE(weight, dl_uint8); @@ -1600,16 +1370,12 @@ void interleave_moe_weights_for_sm90_mixed_gemm(TensorView weight, TensorView we << "weight and weight_interleaved must share n dim"; TVM_FFI_ICHECK_EQ(weight.size(2), weight_interleaved.size(2)) << "weight and weight_interleaved must share packed-k dim"; - TVM_FFI_ICHECK(quant_type == 0 || quant_type == 1 || quant_type == 2) - << "quant_type must be 0 (INT4), 1 (FP4), or 2 (FP4 for FP8), got " << quant_type; + TVM_FFI_ICHECK(quant_type == 0 || quant_type == 1) + << "quant_type must be 0 (INT4) or 1 (FP4), got " << quant_type; int64_t const num_experts = weight.size(0); int64_t const n = weight.size(1); int64_t const k = weight.size(2) * 2; - TVM_FFI_ICHECK_EQ(n % 16, 0) - << "weight n dimension must be divisible by 16 for SM90 mixed-gemm interleave"; - TVM_FFI_ICHECK_EQ(k % 64, 0) - << "logical K dimension must be divisible by 64 for SM90 mixed-gemm interleave"; int64_t const per_expert_bytes = n * (k / 2); auto stream = get_stream(weight.device()); @@ -1621,9 +1387,6 @@ void interleave_moe_weights_for_sm90_mixed_gemm(TensorView weight, TensorView we if (quant_type == 1) { tensorrt_llm::kernels::cutlass_kernels::interleave_fp4_weights_for_sm90_mixed_gemm( src_e, dst_e, static_cast(n), static_cast(k), stream); - } else if (quant_type == 2) { - tensorrt_llm::kernels::cutlass_kernels::interleave_fp4_fp8_weights_for_sm90_mixed_gemm( - src_e, dst_e, static_cast(n), static_cast(k), stream); } else { tensorrt_llm::kernels::cutlass_kernels::interleave_int4_weights_for_sm90_mixed_gemm( src_e, dst_e, static_cast(n), static_cast(k), stream); diff --git a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/detail/collective/mixed_input_utils.hpp b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/detail/collective/mixed_input_utils.hpp index 9650b8e6b7a..74a07035948 100644 --- a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/detail/collective/mixed_input_utils.hpp +++ b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/detail/collective/mixed_input_utils.hpp @@ -19,28 +19,13 @@ #include "cute/numeric/arithmetic_tuple.hpp" #include "cute/util/type_traits.hpp" #include "cutlass/cutlass.h" -#include "cutlass/detail/collective/mixed_input_utils.hpp" #include "cutlass/numeric_conversion.h" ///////////////////////////////////////////////////////////////////////////////////////////////// namespace cutlass::gemm::collective::detail { -constexpr int int4_group_size = 128; -constexpr int mxfp4_group_size = 32; - -template -struct DefaultWeightScaleGroupSize; - -template <> -struct DefaultWeightScaleGroupSize { - static constexpr int value = 32; -}; - -template <> -struct DefaultWeightScaleGroupSize { - static constexpr int value = 128; -}; +using namespace cute; typedef uint32_t __nv_fp4x8_storage_t; typedef uint32_t __nv_bf16x2_storage_t; @@ -48,57 +33,8 @@ typedef uint32_t __nv_int4x8_storage_t; typedef uint64_t __nv_fp8x8_storage_t; typedef cutlass::uint128_t __nv_bf16x8_storage_t; -// ----------------------------------------------------------------------- -// Interleaved version of the bits of four consecutive fp4 values (i.e. 16-bits): -// s000000eem000000 (1st fp4) -// s000000eem000000 (2nd fp4) -// s000000eem000000 (3rd fp4) -// 0sm0ee0000000000 (4th fp4) -// ----------------------------------------------------------------------- - -__device__ __inline__ __nv_bf16x8_storage_t psx_cvt_triton_fp4x8_to_bf16x8_interleaved( - const __nv_fp4x8_storage_t fp4x8) { - __nv_bf16x8_storage_t bf16x8_raw; - __nv_bfloat162* bf16x2_raw = reinterpret_cast<__nv_bfloat162*>(&bf16x8_raw); - - // 0x7e807e80 -> BF16 [126, 126] - uint32_t bias_raw = 0x7e807e80U; - __nv_bfloat162 bias = reinterpret_cast<__nv_bfloat162&>(bias_raw); - - __nv_fp4x8_storage_t first_fp4 = fp4x8 & 0x81C081C0U; - bf16x2_raw[0] = __hmul2(reinterpret_cast<__nv_bfloat162&>(first_fp4), bias); - - __nv_fp4x8_storage_t second_fp4 = (fp4x8 << 3) & 0x81C081C0U; - bf16x2_raw[1] = __hmul2(reinterpret_cast<__nv_bfloat162&>(second_fp4), bias); - - __nv_fp4x8_storage_t third_fp4 = (fp4x8 << 6) & 0x81C081C0U; - bf16x2_raw[2] = __hmul2(reinterpret_cast<__nv_bfloat162&>(third_fp4), bias); - - __nv_fp4x8_storage_t fourth_fp4; - __nv_fp4x8_storage_t fourth_fp4_s = (fp4x8 << 1) & 0x80008000U; - __nv_fp4x8_storage_t fourth_fp4_e = fp4x8 >> 3; - - static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; - asm volatile( - "{\n" - " lop3.b32 %0, %0, %1, %2, %3;\n" - "}\n" - : "+r"(fourth_fp4_e) - : "n"(0x01800180U), "r"(fourth_fp4_s), "n"(immLut)); - - __nv_fp4x8_storage_t fourth_fp4_m = fp4x8 >> 7; - - asm volatile( - "{\n" - " lop3.b32 %0, %1, %2, %3, %4;\n" - "}\n" - : "=r"(fourth_fp4) - : "r"(fourth_fp4_m), "n"(0x00400040U), "r"(fourth_fp4_e), "n"(immLut)); - - bf16x2_raw[3] = __hmul2(reinterpret_cast<__nv_bfloat162&>(fourth_fp4), bias); - - return bf16x8_raw; -} +constexpr int int4_group_size = 128; +constexpr int mxfp4_group_size = 32; inline __device__ unsigned prmt(unsigned hi, unsigned lo, unsigned select_code) { unsigned res = 0; @@ -116,7 +52,7 @@ inline __device__ unsigned prmt(unsigned hi, unsigned lo, unsigned select_code) __constant__ static __nv_fp8x4_storage_t HIGH_E4M3s_LUT_[2] = {0x03020100U, 0x03020100U}; __constant__ static __nv_fp8x4_storage_t LOW_E4M3s_LUT_[2] = {0xFFFEFC00U, 0xFFFEFC00U}; -__device__ __inline__ __nv_fp8x4_storage_t cvt_lut_fp4_to_bf16(const unsigned index) { +__device__ __inline__ __nv_fp8x4_storage_t cvt_lut_fp4_to_bf16(unsigned const index) { auto lane_id = threadIdx.x & 0x1; __nv_fp8x4_storage_t h4b_lut = HIGH_E4M3s_LUT_[lane_id]; __nv_fp8x4_storage_t l4b_lut = LOW_E4M3s_LUT_[lane_id]; @@ -128,10 +64,6 @@ __device__ __inline__ __nv_fp8x4_storage_t cvt_lut_fp4_to_bf16(const unsigned in __device__ __inline__ __nv_bf16x8_storage_t psx_cvt_lut_prmt_fp4x8_to_bf16x8_interleaved( const __nv_fp4x8_storage_t fp4x8) { - // interleaved version - // input fp4x8: 7564 3120 - // output bf16x8: 7654 3210 - __nv_bf16x8_storage_t bf16x8_raw; __nv_bf16x2_storage_t* bf16x2_raw = reinterpret_cast<__nv_bf16x2_storage_t*>(&bf16x8_raw); @@ -153,64 +85,6 @@ __device__ __inline__ __nv_bf16x8_storage_t psx_cvt_lut_prmt_fp4x8_to_bf16x8_int return bf16x8_raw; } -// FP4 E2M1 [0, 0.5, 1, 1.5] encoded as FP8 E4M3. -__constant__ static uint32_t FP4_POS_E4M3s_REG1_[2] = {0x3C383000, 0x3C383000}; -// FP4 E2M1 [2, 3, 4, 6] encoded as FP8 E4M3. -__constant__ static uint32_t FP4_POS_E4M3s_REG2_[2] = {0x4C484440, 0x4C484440}; - -__device__ __inline__ __nv_fp8x8_storage_t psx_cvt_lut_prmt_fp4x8_to_fp8x8( - const __nv_fp4x8_storage_t fp4x8) { - __nv_fp8x8_storage_t fp8x8_raw; - __nv_fp8x4_storage_t* fp8x4_raw = reinterpret_cast<__nv_fp8x4_storage_t*>(&fp8x8_raw); - - __nv_fp8x4_storage_t hb_sign_fp8x4 = (fp4x8 & 0x80808080U); - __nv_fp8x4_storage_t lb_sign_fp8x4 = (fp4x8 & 0x08080808U) << 4U; - - __nv_fp8x4_storage_t h4b_sign_fp8x4 = prmt(hb_sign_fp8x4, lb_sign_fp8x4, 0x7362U); - __nv_fp8x4_storage_t l4b_sign_fp8x4 = prmt(hb_sign_fp8x4, lb_sign_fp8x4, 0x5140U); - - // PRMT consumes only the low 16 bits of its selector in generic mode, so - // the low half does not need its high selector half cleared. - unsigned l4b_em_fp4x4 = fp4x8 & 0x77777777U; - unsigned h4b_em_fp4x4 = l4b_em_fp4x4 >> 16U; - - auto lane_id = threadIdx.x & 0x1; - uint32_t h4b_lut = FP4_POS_E4M3s_REG2_[lane_id]; - uint32_t l4b_lut = FP4_POS_E4M3s_REG1_[lane_id]; - __nv_fp8x4_storage_t h4b_em_fp8x4 = prmt(h4b_lut, l4b_lut, h4b_em_fp4x4); - __nv_fp8x4_storage_t l4b_em_fp8x4 = prmt(h4b_lut, l4b_lut, l4b_em_fp4x4); - - fp8x4_raw[0] = l4b_sign_fp8x4 | l4b_em_fp8x4; - fp8x4_raw[1] = h4b_sign_fp8x4 | h4b_em_fp8x4; - - return fp8x8_raw; -} - -__device__ __inline__ __nv_fp8x8_storage_t psx_cvt_lut_prmt_fp4x8_to_fp8x8_preprocessed_signs( - const __nv_fp4x8_storage_t fp4x8) { - __nv_fp8x8_storage_t fp8x8_raw; - __nv_fp8x4_storage_t* fp8x4_raw = reinterpret_cast<__nv_fp8x4_storage_t*>(&fp8x8_raw); - - // Offline preprocessing keeps each nibble's low 3 EM bits in place, but - // repacks signs so outputs 0..3 are already in byte bit7 and outputs 4..7 - // are in bit3 of each byte. That removes the runtime sign-gather PRMTs. - // PRMT consumes only the low 16 bits of its selector in generic mode, so - // the low half does not need its high selector half cleared. - unsigned l4b_em_fp4x4 = fp4x8 & 0x77777777U; - unsigned h4b_em_fp4x4 = l4b_em_fp4x4 >> 16U; - - auto lane_id = threadIdx.x & 0x1; - uint32_t h4b_lut = FP4_POS_E4M3s_REG2_[lane_id]; - uint32_t l4b_lut = FP4_POS_E4M3s_REG1_[lane_id]; - __nv_fp8x4_storage_t h4b_em_fp8x4 = prmt(h4b_lut, l4b_lut, h4b_em_fp4x4); - __nv_fp8x4_storage_t l4b_em_fp8x4 = prmt(h4b_lut, l4b_lut, l4b_em_fp4x4); - - fp8x4_raw[0] = (fp4x8 & 0x80808080U) | l4b_em_fp8x4; - fp8x4_raw[1] = ((fp4x8 << 4U) & 0x80808080U) | h4b_em_fp8x4; - - return fp8x8_raw; -} - // [ 0, 1, 2, 3] encoded as FP8 __constant__ static uint32_t POS_E4M3s_REG1_[2] = {0x44403800, 0x44403800}; // [ 4, 5, 6, 7] encoded as FP8 @@ -226,7 +100,7 @@ __device__ __inline__ __nv_fp8x8_storage_t psx_cvt_lut_prmt_int4x8_to_fp8x8( __nv_fp8x4_storage_t* fp8x4_raw = reinterpret_cast<__nv_fp8x4_storage_t*>(&fp8x8_raw); // View the input as reg - uint32_t reg = reinterpret_cast(int4x8); + uint32_t reg = reinterpret_cast(int4x8); // Determines if to get from the signed or unsigned candidates uint32_t sign = (reg & 0x88888888) >> 1; @@ -265,32 +139,6 @@ __device__ __inline__ __nv_fp8x8_storage_t psx_cvt_lut_prmt_int4x8_to_fp8x8( return fp8x8_raw; } -template -using MixedInputVoid = void; - -template -struct MixedInputFusedE8M0PreMmaScale { - static constexpr bool value = false; -}; - -template -struct MixedInputFusedE8M0PreMmaScale> { - static constexpr bool value = Collective::FusedE8M0PreMmaScale; -}; - -template -struct MixedInputFoldedWeightScaleStorage { - static constexpr bool value = false; -}; - -template -struct MixedInputFoldedWeightScaleStorage< - Collective, MixedInputVoid> { - static constexpr bool value = true; -}; - template struct MixedGroupedGemmInputUtils { private: @@ -299,25 +147,18 @@ struct MixedGroupedGemmInputUtils { using SmemLayoutA = typename Collective::SmemLayoutA; using SmemLayoutB = typename Collective::SmemLayoutB; using SmemLayoutScale = typename Collective::SmemLayoutScale; - using SmemLayoutActivationScale = typename Collective::SmemLayoutActivationScale; using SwappedElementA = typename Collective::SwappedElementA; using SwappedElementB = typename Collective::SwappedElementB; using RealSwappedElementA = typename Collective::RealSwappedElementA; using RealSwappedElementB = typename Collective::RealSwappedElementB; using ElementScale = typename Collective::ElementScale; using ElementZero = typename Collective::ElementZero; - using NonVoidElementActivationScale = typename Collective::NonVoidElementActivationScale; using SmemCopyAtomScale = typename Collective::SmemCopyAtomScale; static constexpr auto KernelConversionMode = Collective::KernelConversionMode; static constexpr auto ModeHasScales = Collective::ModeHasScales; static constexpr auto UseScaleLookupTable = Collective::UseScaleLookupTable; static constexpr auto UseFP4ToBF16LookupTable = Collective::UseFP4ToBF16LookupTable; - static constexpr auto UseFP4ToFP8LookupTable = Collective::UseFP4ToFP8LookupTable; static constexpr auto UseInt4ToFP8LookupTable = Collective::UseInt4ToFP8LookupTable; - static constexpr auto HasActivationScale = Collective::HasActivationScale; - static constexpr bool FusedE8M0PreMmaScale = MixedInputFusedE8M0PreMmaScale::value; - static constexpr bool HasFoldedWeightScaleStorage = - MixedInputFoldedWeightScaleStorage::value; public: static constexpr auto elements_per_smem_scale() { @@ -356,46 +197,23 @@ struct MixedGroupedGemmInputUtils { } static constexpr uint32_t compute_tma_transaction_bytes_extra() { - constexpr uint32_t bulk_copy_alignment_bytes = 16; if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { return 0; } else if constexpr (ModeHasScales) { constexpr uint32_t scale_tx_bytes = cutlass::bits_to_bytes(size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) * static_cast(cute::sizeof_bits_v)); + static_assert(scale_tx_bytes % 128 == 0, + "Each scale stage must be 128B aligned."); // required by TMA if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { - if constexpr (FusedE8M0PreMmaScale || HasFoldedWeightScaleStorage) { - static_assert(Collective::WeightScaleBulkCopyBytes % bulk_copy_alignment_bytes == 0, - "Each folded weight-scale bulk copy must be 16B aligned."); - } else { - static_assert(scale_tx_bytes % bulk_copy_alignment_bytes == 0, - "Each scale bulk copy must be 16B aligned."); - } - if constexpr (HasActivationScale) { - constexpr uint32_t activation_scale_tx_bytes = cutlass::bits_to_bytes( - size<0>(SmemLayoutActivationScale{}) * size<1>(SmemLayoutActivationScale{}) * - static_cast(cute::sizeof_bits_v)); - if constexpr (FusedE8M0PreMmaScale || HasFoldedWeightScaleStorage) { - return Collective::WeightScaleTransactionBytes + activation_scale_tx_bytes; - } else { - return scale_tx_bytes + activation_scale_tx_bytes; - } - } else { - if constexpr (FusedE8M0PreMmaScale || HasFoldedWeightScaleStorage) { - return Collective::WeightScaleTransactionBytes; - } else { - return scale_tx_bytes; - } - } + return scale_tx_bytes; } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { // Scale and zero share smem layout - static_assert(scale_tx_bytes % bulk_copy_alignment_bytes == 0, - "Each scale bulk copy must be 16B aligned."); constexpr uint32_t zero_tx_bytes = cutlass::bits_to_bytes(size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) * static_cast(cute::sizeof_bits_v)); - static_assert(zero_tx_bytes % bulk_copy_alignment_bytes == 0, - "Each zero bulk copy must be 16B aligned."); + static_assert(zero_tx_bytes % 128 == 0, + "Each zero stage must be 128B aligned."); // required by TMA return scale_tx_bytes + zero_tx_bytes; } else { static_assert(cutlass::detail::dependent_false, @@ -430,6 +248,7 @@ struct MixedGroupedGemmInputUtils { auto smem_tiled_copy_S = cute::get<0>(tiled_copy_and_views); auto tCrS_copy_view = cute::get<1>(tiled_copy_and_views); auto tCsS = cute::get<0>(partitioned_mma_extra_info); + copy(smem_tiled_copy_S, tCsS(_, _, k_block, read_stage), tCrS_copy_view(_, _, k_block)); if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { // Nothing extra to do @@ -495,6 +314,7 @@ struct MixedGroupedGemmInputUtils { Tensor const& scales_pos) { lookup_table_convert(src, dst, scales_neg, scales_pos); } + template CUTLASS_DEVICE static void lookup_table_convert( @@ -559,8 +379,7 @@ struct MixedGroupedGemmInputUtils { auto&& src_ = cute::recast<__nv_fp4x8_storage_t>(src)(0); auto&& dst_ = cute::recast<__nv_bf16x8_storage_t>(dst)(0); - // dst_ = psx_cvt_lut_prmt_fp4x8_to_bf16x8_interleaved(src_); - dst_ = psx_cvt_triton_fp4x8_to_bf16x8_interleaved(src_); + dst_ = psx_cvt_lut_prmt_fp4x8_to_bf16x8_interleaved(src_); } template - CUTLASS_DEVICE static void fp4tofp8_lookup_table_convert( // Accept mutable temporaries - Tensor const& src, Tensor&& dst) { - fp4tofp8_lookup_table_convert(src, dst); - } - - template - CUTLASS_DEVICE static void fp4tofp8_lookup_table_convert(Tensor const& src, - Tensor& dst) { - auto&& src_ = cute::recast<__nv_fp4x8_storage_t>(src)(0); - auto&& dst_ = cute::recast<__nv_fp8x8_storage_t>(dst)(0); - -#if defined(CUTLASS_MIXED_GEMM_FP4_FP8_PREPROCESSED_SIGNS) - dst_ = psx_cvt_lut_prmt_fp4x8_to_fp8x8_preprocessed_signs(src_); -#else - dst_ = psx_cvt_lut_prmt_fp4x8_to_fp8x8(src_); -#endif - } - - __device__ __inline__ static void fp4tofp8_fused_e8m0_pre_mma_convert_pair( - __nv_fp4x8_storage_t fp4x8_0, __nv_fp4x8_storage_t fp4x8_1, __nv_fp8x8_storage_t& fp8x8_raw_0, - __nv_fp8x8_storage_t& fp8x8_raw_1, uint32_t lo_exp_offset, uint32_t hi_exp_offset) { - // One WGMMA A operand lane contributes two fp4x8 registers whose low - // fp8x4 chunks share one row scale, and high chunks share the other. - __nv_fp8x4_storage_t* fp8x4_raw_0 = reinterpret_cast<__nv_fp8x4_storage_t*>(&fp8x8_raw_0); - __nv_fp8x4_storage_t* fp8x4_raw_1 = reinterpret_cast<__nv_fp8x4_storage_t*>(&fp8x8_raw_1); - - uint32_t const fp4_raw_0 = reinterpret_cast(fp4x8_0); - uint32_t const fp4_raw_1 = reinterpret_cast(fp4x8_1); - uint32_t const em_selector_0 = fp4_raw_0 & 0x77777777U; - uint32_t const em_selector_1 = fp4_raw_1 & 0x77777777U; - constexpr uint32_t fp4_codes_0_to_3_em_bias = 0x0c080000U; - constexpr uint32_t fp4_codes_4_to_7_em_bias = 0x1c181410U; - uint32_t const lo_l4b_exp_offseted_lut = - (lo_exp_offset * 0x08080800U) + fp4_codes_0_to_3_em_bias; - uint32_t const lo_h4b_exp_offseted_lut = - (lo_exp_offset * 0x08080808U) + fp4_codes_4_to_7_em_bias; - uint32_t const hi_l4b_exp_offseted_lut = - (hi_exp_offset * 0x08080800U) + fp4_codes_0_to_3_em_bias; - uint32_t const hi_h4b_exp_offseted_lut = - (hi_exp_offset * 0x08080808U) + fp4_codes_4_to_7_em_bias; - - uint32_t const lo_em_fp8x4_0 = - prmt(lo_h4b_exp_offseted_lut, lo_l4b_exp_offseted_lut, em_selector_0); - uint32_t const lo_em_fp8x4_1 = - prmt(lo_h4b_exp_offseted_lut, lo_l4b_exp_offseted_lut, em_selector_1); - -#if defined(CUTLASS_MIXED_GEMM_FP4_FP8_PREPROCESSED_SIGNS) - fp8x4_raw_0[0] = (fp4_raw_0 & 0x80808080U) | lo_em_fp8x4_0; - fp8x4_raw_1[0] = (fp4_raw_1 & 0x80808080U) | lo_em_fp8x4_1; -#else - uint32_t const hb_sign_fp8x4_0 = fp4_raw_0 & 0x80808080U; - uint32_t const hb_sign_fp8x4_1 = fp4_raw_1 & 0x80808080U; - uint32_t const lb_sign_fp8x4_0 = (fp4_raw_0 & 0x08080808U) << 4U; - uint32_t const lb_sign_fp8x4_1 = (fp4_raw_1 & 0x08080808U) << 4U; - uint32_t const l4b_sign_fp8x4_0 = prmt(hb_sign_fp8x4_0, lb_sign_fp8x4_0, 0x5140U); - uint32_t const l4b_sign_fp8x4_1 = prmt(hb_sign_fp8x4_1, lb_sign_fp8x4_1, 0x5140U); - uint32_t const h4b_sign_fp8x4_0 = prmt(hb_sign_fp8x4_0, lb_sign_fp8x4_0, 0x7362U); - uint32_t const h4b_sign_fp8x4_1 = prmt(hb_sign_fp8x4_1, lb_sign_fp8x4_1, 0x7362U); - - fp8x4_raw_0[0] = l4b_sign_fp8x4_0 | lo_em_fp8x4_0; - fp8x4_raw_1[0] = l4b_sign_fp8x4_1 | lo_em_fp8x4_1; -#endif - - uint32_t const hi_em_fp8x4_0 = - prmt(hi_h4b_exp_offseted_lut, hi_l4b_exp_offseted_lut, em_selector_0 >> 16U); - uint32_t const hi_em_fp8x4_1 = - prmt(hi_h4b_exp_offseted_lut, hi_l4b_exp_offseted_lut, em_selector_1 >> 16U); - -#if defined(CUTLASS_MIXED_GEMM_FP4_FP8_PREPROCESSED_SIGNS) - fp8x4_raw_0[1] = ((fp4_raw_0 << 4U) & 0x80808080U) | hi_em_fp8x4_0; - fp8x4_raw_1[1] = ((fp4_raw_1 << 4U) & 0x80808080U) | hi_em_fp8x4_1; -#else - fp8x4_raw_0[1] = h4b_sign_fp8x4_0 | hi_em_fp8x4_0; - fp8x4_raw_1[1] = h4b_sign_fp8x4_1 | hi_em_fp8x4_1; -#endif - } - - template - CUTLASS_DEVICE static void fp4tofp8_fused_e8m0_pre_mma_convert_pair( - Tensor const& src0, Tensor const& src1, - Tensor& dst0, Tensor& dst1, - uint32_t lo_exp_offset, uint32_t hi_exp_offset) { - auto&& src0_ = cute::recast<__nv_fp4x8_storage_t>(src0)(0); - auto&& src1_ = cute::recast<__nv_fp4x8_storage_t>(src1)(0); - auto&& dst0_ = cute::recast<__nv_fp8x8_storage_t>(dst0)(0); - auto&& dst1_ = cute::recast<__nv_fp8x8_storage_t>(dst1)(0); - - fp4tofp8_fused_e8m0_pre_mma_convert_pair(src0_, src1_, dst0_, dst1_, lo_exp_offset, - hi_exp_offset); - } - /// Utilities to dequantize A. template CUTLASS_DEVICE static void static_check_scale(Layout const& tensor) { static_assert(shape<0>(Layout{}) >= 4 && stride<0>(Layout{}) == 0, "At least 4 adjacent weights in a thread must share the same scale."); } + template CUTLASS_DEVICE static void static_check_scale(Tensor const& tensor) { static_check_scale(flatten(Layout{})); @@ -826,18 +553,24 @@ struct MixedGroupedGemmInputUtils { } } - template - CUTLASS_DEVICE static void convert_A_slot(Tensor const& src, - Tensor& dst) { + template + CUTLASS_DEVICE static void convert_A_kblock(Tensor const& tCrA_load, + Tensor& tCrA_mma, + int const k_block) { static_assert(is_rmem::value, "Input tensor for A conversion must come from registers"); static_assert(is_rmem::value, "Output tensor for A conversion must come from registers"); + static_assert(cosize_v == cosize_v); + static_assert(size_v == cosize_v); + static_assert(size_v == cosize_v); using SrcType = typename EngineIn::value_type; + Tensor src = tCrA_load(_, _, k_block); + Tensor dst = tCrA_mma(_, _, k_block); + CUTE_STATIC_ASSERT_V(size(src(_, 0)) == cosize(src(_, 0).layout()), "The first mode of tensor src must be contiguous in memory"); - CUTE_STATIC_ASSERT_V(size(src) == size(dst)); // try to make the size of the first mode equal to 32bit int constexpr NumValPerSrcReg = cute::min(decltype(size(src(_, 0)))::value, ceil_div(32, sizeof_bits_v)); @@ -849,8 +582,6 @@ struct MixedGroupedGemmInputUtils { for (int i = 0; i < size<1>(dst_vm); ++i) { if constexpr (UseFP4ToBF16LookupTable) { fp4tobf16_lookup_table_convert(src_vm(_, i), dst_vm(_, i)); - } else if constexpr (UseFP4ToFP8LookupTable) { - fp4tofp8_lookup_table_convert(src_vm(_, i), dst_vm(_, i)); } else if constexpr (UseInt4ToFP8LookupTable) { int4tofp8_lookup_table_convert(src_vm(_, i), dst_vm(_, i)); } else { @@ -859,165 +590,6 @@ struct MixedGroupedGemmInputUtils { } } - template - CUTLASS_DEVICE static void convert_A_kblock(Tensor const& tCrA_load, - Tensor& tCrA_mma, - int const k_block) { - Tensor src = tCrA_load(_, _, k_block); - Tensor dst = tCrA_mma(_, _, k_block); - convert_A_slot(src, dst); - } - - template - CUTLASS_DEVICE static void convert_A_kblock(Tensor const& tCrA_load, - Tensor& tCrA_mma, - cute::Int k_block) { - Tensor src = tCrA_load(_, _, k_block); - Tensor dst = tCrA_mma(_, _, k_block); - convert_A_slot(src, dst); - } - - template - CUTLASS_DEVICE static void convert_A_kblock_fused_e8m0_pre_mma_raw_scale_to_slot( - Tensor const& tCrA_load, Tensor& tCrA_mma_slot, - Tensor& scale_values, cute::Int) { - static_assert(FusedE8M0PreMmaScale, "This helper is only for fused e8m0 pre-MMA scale."); - static_assert(UseFP4ToFP8LookupTable, - "Fused e8m0 pre-MMA scale currently supports MXFP4 x FP8 only."); - static_assert(is_rmem::value, - "Input tensor for A conversion must come from registers"); - static_assert(is_rmem::value, - "Output tensor for A conversion must come from registers"); - static_assert(is_rmem::value, - "Scale tensor for A conversion must come from registers"); - using SrcType = typename EngineIn::value_type; - using ScaleScalar = ElementScale; - static_assert(cute::is_same_v, - "Raw fused e8m0 scale tensor must use scalar e8m0 elements."); - - Tensor src = tCrA_load(_, _, cute::Int{}); - Tensor dst = tCrA_mma_slot; - Tensor scales = scale_values(_, _, cute::Int{}); - - CUTE_STATIC_ASSERT_V(size(src(_, 0)) == cosize(src(_, 0).layout()), - "The first mode of tensor src must be contiguous in memory"); - CUTE_STATIC_ASSERT_V(size(src) == size(dst)); - CUTE_STATIC_ASSERT_V(size(src) == size(scales)); - - int constexpr NumValPerSrcReg = - cute::min(decltype(size(src(_, 0)))::value, ceil_div(32, sizeof_bits_v)); - Tensor src_vm = cute::group_modes<1, -1>(cute::zipped_divide(src, Int{})); - Tensor dst_vm = cute::group_modes<1, -1>(cute::zipped_divide(dst, Int{})); - Tensor scales_vm = - cute::group_modes<1, -1>(cute::zipped_divide(scales, Int{})); - - auto scale_values_0 = cute::filter(scales_vm(_, Int<0>{})); - constexpr int ScaleValueCount = decltype(size(scale_values_0))::value; - constexpr int DstVecCount = decltype(size<1>(dst_vm))::value; - static_assert(ScaleValueCount == 2 || ScaleValueCount == NumValPerSrcReg, - "Fused e8m0 pre-MMA raw scale expects either two compact row scales or one scale " - "per fp4 lane."); - static_assert((DstVecCount % 2) == 0, - "Fused e8m0 pre-MMA pair conversion expects an even number of fp4x8 operands."); - - constexpr int HiScaleIndex = (ScaleValueCount == NumValPerSrcReg) ? (NumValPerSrcReg / 2) : 1; - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < DstVecCount; i += 2) { - auto row_scales = cute::filter(scales_vm(_, i)); - ScaleScalar const lo_scale = row_scales(0); - ScaleScalar const hi_scale = row_scales(HiScaleIndex); - uint32_t const lo_exp_offset = static_cast(lo_scale.storage); - uint32_t const hi_exp_offset = static_cast(hi_scale.storage); - auto src_vec0 = src_vm(_, i); - auto src_vec1 = src_vm(_, i + 1); - auto dst_vec0 = dst_vm(_, i); - auto dst_vec1 = dst_vm(_, i + 1); - fp4tofp8_fused_e8m0_pre_mma_convert_pair(src_vec0, src_vec1, dst_vec0, dst_vec1, - lo_exp_offset, hi_exp_offset); - } - } - - template - CUTLASS_DEVICE static void cache_A_kblock_fused_e8m0_pre_mma_exp_offsets( - Tensor const& scales, cute::Int, cute::Int, - LoOffsetArray& lo_exp_offsets, HiOffsetArray& hi_exp_offsets) { - static_assert(FusedE8M0PreMmaScale, "This helper is only for fused e8m0 pre-MMA scale."); - using ScaleScalar = typename EngineScale::value_type; - constexpr int NumValPerSrcReg = 8; - Tensor scales_vm = - cute::group_modes<1, -1>(cute::zipped_divide(scales, Int{})); - static_assert(decltype(size<1>(scales_vm))::value == ScalePairCount * 2, - "Fused e8m0 pre-MMA scale tensor must match A operand pair layout."); - - cute::for_each(cute::make_seq{}, [&](auto pair_c) { - constexpr int pair = decltype(pair_c)::value; - constexpr int scale_vec = pair * 2; - Tensor row_scales = scales_vm(_, Int{}); - constexpr int ScaleValueCount = decltype(size(row_scales))::value; - static_assert(ScaleValueCount == 2 || ScaleValueCount == NumValPerSrcReg, - "Fused e8m0 pre-MMA raw scale expects either two compact row scales or one " - "scale per fp4 lane."); - constexpr int HiScaleIndex = (ScaleValueCount == NumValPerSrcReg) ? (NumValPerSrcReg / 2) : 1; - ScaleScalar const lo_scale = row_scales(0); - ScaleScalar const hi_scale = row_scales(HiScaleIndex); - constexpr int cache_index = KBlock * ScalePairCount + pair; - lo_exp_offsets[cache_index] = static_cast(lo_scale.storage); - hi_exp_offsets[cache_index] = static_cast(hi_scale.storage); - }); - } - - template - CUTLASS_DEVICE static void convert_A_kblock_fused_e8m0_pre_mma_exp_offsets_to_slot( - Tensor const& tCrA_load, Tensor& tCrA_mma_slot, - cute::Int, cute::Int, LoOffsetArray const& lo_exp_offsets, - HiOffsetArray const& hi_exp_offsets) { - static_assert(FusedE8M0PreMmaScale, "This helper is only for fused e8m0 pre-MMA scale."); - static_assert(UseFP4ToFP8LookupTable, - "Fused e8m0 pre-MMA scale currently supports MXFP4 x FP8 only."); - static_assert(is_rmem::value, - "Input tensor for A conversion must come from registers"); - static_assert(is_rmem::value, - "Output tensor for A conversion must come from registers"); - using SrcType = typename EngineIn::value_type; - - Tensor src = tCrA_load(_, _, cute::Int{}); - Tensor dst = tCrA_mma_slot; - - CUTE_STATIC_ASSERT_V(size(src(_, 0)) == cosize(src(_, 0).layout()), - "The first mode of tensor src must be contiguous in memory"); - CUTE_STATIC_ASSERT_V(size(src) == size(dst)); - - int constexpr NumValPerSrcReg = - cute::min(decltype(size(src(_, 0)))::value, ceil_div(32, sizeof_bits_v)); - Tensor src_vm = cute::group_modes<1, -1>(cute::zipped_divide(src, Int{})); - Tensor dst_vm = cute::group_modes<1, -1>(cute::zipped_divide(dst, Int{})); - - constexpr int DstVecCount = decltype(size<1>(dst_vm))::value; - static_assert((DstVecCount % 2) == 0, - "Fused e8m0 pre-MMA pair conversion expects an even number of fp4x8 operands."); - static_assert( - ScalePairCount * 2 == DstVecCount, - "Fused e8m0 pre-MMA scale cache must provide one scale pair per fp4x8 operand pair."); - - cute::for_each(cute::make_seq{}, [&](auto pair_c) { - constexpr int pair = decltype(pair_c)::value; - constexpr int i = pair * 2; - auto src_vec0 = src_vm(_, i); - auto src_vec1 = src_vm(_, i + 1); - auto dst_vec0 = dst_vm(_, i); - auto dst_vec1 = dst_vm(_, i + 1); - constexpr int cache_index = KBlock * ScalePairCount + pair; - uint32_t const lo_exp_offset = lo_exp_offsets[cache_index]; - uint32_t const hi_exp_offset = hi_exp_offsets[cache_index]; - fp4tofp8_fused_e8m0_pre_mma_convert_pair(src_vec0, src_vec1, dst_vec0, dst_vec1, - lo_exp_offset, hi_exp_offset); - }); - } - /// Utilities for any additional inputs inside of the TMA load template CUTLASS_DEVICE static auto partition_extra_tma_inputs(Params const& mainloop_params, diff --git a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/collective/default_epilogue_array_per_token_scale.hpp b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/collective/default_epilogue_array_per_token_scale.hpp deleted file mode 100644 index 4b6ffac935f..00000000000 --- a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/collective/default_epilogue_array_per_token_scale.hpp +++ /dev/null @@ -1,329 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - **************************************************************************************************/ - -#pragma once - -#include - -#include "cute/tensor.hpp" -#include "cutlass/cuda_host_adapter.hpp" -#include "cutlass/cutlass.h" -#include "cutlass/epilogue/collective/detail.hpp" -#include "cutlass/epilogue/dispatch_policy.hpp" -#include "cutlass/numeric_conversion.h" - -namespace cutlass::epilogue::collective { - -// Shared-memory exchange variant matching the small-K Humming writeback shape: -// accumulator owners scatter BF16 values into a compact token-major tile, then -// all 128 threads issue coalesced 16B global stores. -template -class SmemEpilogueArrayPerTokenScale { - public: - using CtaTileShapeMNK = CtaTileShapeMNK_; - using EpilogueSchedule = PtrArrayNoSmemWarpSpecialized; - using DispatchPolicy = EpilogueSchedule; - using ElementOutput = ElementD_; - using ElementAccumulator = ElementAccumulator_; - using ElementCompute = ElementAccumulator; - using ElementScalar = ElementScalar_; - using ElementC = ElementC_; - using StrideC = StrideC_; - using InternalStrideC = cute::remove_pointer_t; - using ElementD = ElementD_; - using StrideD = StrideD_; - using InternalStrideD = cute::remove_pointer_t; - using GmemTiledCopyC = void; - using GmemTiledCopyD = void; - - struct ThreadEpilogueOp { - using ElementOutput = ElementD_; - using ElementD = ElementD_; - using ElementAccumulator = ElementAccumulator_; - using ElementCompute = ElementAccumulator_; - }; - - static constexpr int TileM = cute::size<0>(CtaTileShapeMNK{}); - static constexpr int TileN = cute::size<1>(CtaTileShapeMNK{}); - static constexpr int TileElements = TileM * TileN; - static constexpr int OutputAlignmentBits = 128; - static constexpr int RequiredChannelMultiple = 128; - static constexpr int ElementsPerVector = 128 / cute::sizeof_bits_v; - static constexpr int VectorCount = TileElements / ElementsPerVector; - static constexpr int VectorsPerThread = VectorCount / NumThreadsPerWarpGroup; - static constexpr int kOutputAlignment = ElementsPerVector; - - static_assert(OutputAlignmentBits % cute::sizeof_bits_v == 0); - static_assert(TileElements % ElementsPerVector == 0); - static_assert(TileN % 8 == 0); - static_assert(TileM % ElementsPerVector == 0, - "Each vector store must remain inside one output row."); - static_assert( - VectorCount % NumThreadsPerWarpGroup == 0, - "The compact SMEM epilogue expects an integer number of output vectors per thread."); - static_assert(cute::is_same_v(InternalStrideD{})), cute::Int<1>>, - "The compact SMEM epilogue requires a unit-stride channel dimension."); - static_assert(cute::rank(InternalStrideC{}) == 3, "StrideC must be rank-3."); - static_assert(cute::rank(InternalStrideD{}) == 3, "StrideD must be rank-3."); - - struct SharedStorage { - alignas(16) ElementD output[TileElements]; - alignas(16) ElementCompute token_scale[TileN > 8 ? TileN : 1]; - }; - using TensorMapStorage = SharedStorage; - - struct ThreadArguments { - ElementScalar token_scale_default = ElementScalar(1); - ElementScalar const* const* token_scale_ptr_array = nullptr; - }; - - struct Arguments { - ThreadArguments thread{}; - ElementC const** ptr_C = nullptr; - StrideC dC{}; - ElementD** ptr_D = nullptr; - StrideD dD{}; - // ptr_D entries must address rows within this contiguous output allocation. - ElementD* output_base = nullptr; - int64_t output_channel_extent = 0; - int64_t output_row_stride = 0; - ElementCompute beta = ElementCompute(0); - }; - - struct Params { - ThreadArguments thread{}; - ElementD** ptr_D = nullptr; - StrideD dD{}; - int64_t output_channel_extent = 0; - int64_t output_row_stride = 0; - }; - - template - static constexpr Params to_underlying_arguments(ProblemShape const&, Arguments const& args, - void*) { - return {args.thread, args.ptr_D, args.dD, args.output_channel_extent, args.output_row_stride}; - } - - template - static size_t get_workspace_size(ProblemShape const&, Arguments const&, int) { - return 0; - } - - template - static Status initialize_workspace(ProblemShape const&, Arguments const&, void*, cudaStream_t, - CudaHostAdapter* = nullptr) { - return Status::kSuccess; - } - - template - static bool can_implement(ProblemShape problem_shapes, Arguments const& args) { - bool const no_source_or_beta = args.ptr_C == nullptr && args.beta == ElementCompute(0); - bool const valid_output_storage = - args.ptr_D != nullptr && args.dD != nullptr && args.output_base != nullptr && - (reinterpret_cast(args.output_base) % (OutputAlignmentBits / 8)) == 0; - bool const valid_output_shape = args.output_channel_extent > 0 && - (args.output_channel_extent % RequiredChannelMultiple) == 0 && - (args.output_channel_extent % TileM) == 0; - bool const valid_output_stride = args.output_row_stride >= args.output_channel_extent && - (args.output_row_stride % ElementsPerVector) == 0; - - bool host_shapes_match = true; - if (problem_shapes.is_host_problem_shape_available()) { - for (int group = 0; group < problem_shapes.groups(); ++group) { - auto const problem = problem_shapes.get_host_problem_shape(group); - int64_t const channel_extent = static_cast(cute::get<0>(problem)); - host_shapes_match = host_shapes_match && channel_extent == args.output_channel_extent && - (channel_extent % RequiredChannelMultiple) == 0; - } - } - - if (!no_source_or_beta) { - CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Compact epilogue does not support source C or beta.\n"); - } - if (!valid_output_storage) { - CUTLASS_TRACE_HOST( - " CAN IMPLEMENT: Output storage must provide a 16B-aligned base and pointer/stride " - "arrays.\n"); - } - if (!valid_output_shape) { - CUTLASS_TRACE_HOST( - " CAN IMPLEMENT: Output channel extent must be a positive multiple of 128 and TileM.\n"); - } - if (!valid_output_stride) { - CUTLASS_TRACE_HOST( - " CAN IMPLEMENT: Output row stride must cover the channel extent and be 16B aligned.\n"); - } - if (!host_shapes_match) { - CUTLASS_TRACE_HOST( - " CAN IMPLEMENT: Host problem shapes do not match the compact epilogue channel " - "contract.\n"); - } - - return no_source_or_beta && valid_output_storage && valid_output_shape && valid_output_stride && - host_shapes_match; - } - - CUTLASS_HOST_DEVICE - explicit SmemEpilogueArrayPerTokenScale(Params const& params) : params_(params) {} - - CUTLASS_DEVICE - bool is_source_needed() const { return false; } - - template - CUTLASS_DEVICE void operator()(ProblemShapeMNKL problem_shape_mnkl, BlockShapeMNK block_shape_mnk, - BlockCoordMNKL block_coord_mnkl, - cute::Tensor const& accumulators, - TiledMma tiled_mma, ResidueMNK, int thread_idx, - char* shared_storage_ptr) { - using namespace cute; - static_assert(is_same_v); - - auto M = get<0>(problem_shape_mnkl); - auto N = get<1>(problem_shape_mnkl); - auto [m_coord, n_coord, k_coord, group_coord] = block_coord_mnkl; - int const tile_m_origin = int(m_coord) * TileM; - int const tile_n_origin = int(n_coord) * TileN; - - auto stride_d = [&, group = group_coord]() { - if constexpr (!is_same_v) { - return detail::get_epilogue_stride(params_.dD[group]); - } else { - return detail::get_epilogue_stride(params_.dD); - } - }(); - - auto thread_mma = tiled_mma.get_thread_slice(thread_idx); - Tensor output_coordinates = make_identity_tensor(make_shape(M, N)); - Tensor tile_coordinates = - local_tile(output_coordinates, take<0, 2>(block_shape_mnk), make_coord(m_coord, n_coord)); - Tensor thread_coordinates = thread_mma.partition_C(tile_coordinates); - - SharedStorage& shared = *reinterpret_cast(shared_storage_ptr); - - if constexpr (TileN > 8) { - if (thread_idx < TileN) { - ElementScalar const* token_scales = params_.thread.token_scale_ptr_array - ? params_.thread.token_scale_ptr_array[group_coord] - : nullptr; - int const global_n = tile_n_origin + thread_idx; - shared.token_scale[thread_idx] = - global_n < N && token_scales - ? static_cast(token_scales[global_n]) - : static_cast(params_.thread.token_scale_default); - } - __syncthreads(); - } - - if constexpr (TileN == 8) { - ElementScalar const* token_scales = params_.thread.token_scale_ptr_array - ? params_.thread.token_scale_ptr_array[group_coord] - : nullptr; - NumericConverter convert; - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(accumulators); ++i) { - auto coordinate = thread_coordinates(i); - if (get<1>(coordinate) < N) { - int const local_m = int(get<0>(coordinate)) - tile_m_origin; - int const local_n = int(get<1>(coordinate)) - tile_n_origin; - ElementCompute const token_scale = - token_scales ? static_cast(token_scales[int(get<1>(coordinate))]) - : static_cast(params_.thread.token_scale_default); - shared.output[local_n * TileM + local_m] = - convert(static_cast(accumulators(i)) * token_scale); - } - } - } else { - static_assert(decltype(size(accumulators))::value % 4 == 0); - NumericArrayConverter convert; - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(accumulators); i += 4) { - auto coordinate_0 = thread_coordinates(i); - auto coordinate_1 = thread_coordinates(i + 1); - int const local_m_0 = int(get<0>(coordinate_0)) - tile_m_origin; - int const local_n_0 = int(get<1>(coordinate_0)) - tile_n_origin; - int const local_m_1 = int(get<0>(coordinate_1)) - tile_m_origin; - int const local_n_1 = int(get<1>(coordinate_1)) - tile_n_origin; - ElementCompute const token_scale_0 = shared.token_scale[local_n_0]; - ElementCompute const token_scale_1 = shared.token_scale[local_n_1]; - cutlass::Array scaled_accumulators_01{ - static_cast(accumulators(i)) * token_scale_0, - static_cast(accumulators(i + 1)) * token_scale_1}; - auto converted_01 = convert(scaled_accumulators_01); - - if (get<1>(coordinate_0) < N) { - shared.output[local_n_0 * TileM + local_m_0] = converted_01[0]; - } - if (get<1>(coordinate_1) < N) { - shared.output[local_n_1 * TileM + local_m_1] = converted_01[1]; - } - - auto coordinate_2 = thread_coordinates(i + 2); - auto coordinate_3 = thread_coordinates(i + 3); - int const local_m_2 = int(get<0>(coordinate_2)) - tile_m_origin; - int const local_m_3 = int(get<0>(coordinate_3)) - tile_m_origin; -#if !defined(NDEBUG) - // The SM90 GMMA C fragment repeats each token pair across two adjacent - // channel octets. Reuse those two scales for all four accumulators. - CUTLASS_ASSERT(int(get<1>(coordinate_2)) - tile_n_origin == local_n_0); - CUTLASS_ASSERT(int(get<1>(coordinate_3)) - tile_n_origin == local_n_1); -#endif - cutlass::Array scaled_accumulators_23{ - static_cast(accumulators(i + 2)) * token_scale_0, - static_cast(accumulators(i + 3)) * token_scale_1}; - auto converted_23 = convert(scaled_accumulators_23); - if (get<1>(coordinate_2) < N) { - shared.output[local_n_0 * TileM + local_m_2] = converted_23[0]; - } - if (get<1>(coordinate_3) < N) { - shared.output[local_n_1 * TileM + local_m_3] = converted_23[1]; - } - } - } - - __syncthreads(); - - using OutputVector = cutlass::Array; - auto const* shared_vectors = reinterpret_cast(shared.output); - ElementD* output = params_.ptr_D[group_coord]; - int64_t const stride_n = int64_t(get<1>(stride_d)); - -#if !defined(NDEBUG) - CUTLASS_ASSERT(int64_t(M) == params_.output_channel_extent); - CUTLASS_ASSERT((int64_t(M) % RequiredChannelMultiple) == 0); - CUTLASS_ASSERT(output != nullptr); - CUTLASS_ASSERT((reinterpret_cast(output) % (OutputAlignmentBits / 8)) == 0); - CUTLASS_ASSERT(stride_n == params_.output_row_stride); - CUTLASS_ASSERT((stride_n % ElementsPerVector) == 0); -#endif - - CUTLASS_PRAGMA_UNROLL - for (int vector_group = 0; vector_group < VectorsPerThread; ++vector_group) { - int const vector_idx = thread_idx + vector_group * NumThreadsPerWarpGroup; - int const element_idx = vector_idx * ElementsPerVector; - int const local_n = element_idx / TileM; - int const local_m = element_idx % TileM; - int const global_n = tile_n_origin + local_n; - - if (global_n < N) { - auto* output_vector = reinterpret_cast( - output + int64_t(tile_m_origin + local_m) + int64_t(global_n) * stride_n); - *output_vector = shared_vectors[vector_idx]; - } - } - - // Larger token tiles synchronize before the next tile scatters output, so - // that token-scale barrier also protects this output scratch from reuse. - if constexpr (TileN == 8) { - __syncthreads(); - } - } - - private: - Params params_; -}; - -} // namespace cutlass::epilogue::collective diff --git a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/collective/sm90_epilogue_array_tma_warpspecialized_mixed_input.hpp b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/collective/sm90_epilogue_array_tma_warpspecialized_mixed_input.hpp deleted file mode 100644 index 7b3d27371ee..00000000000 --- a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/collective/sm90_epilogue_array_tma_warpspecialized_mixed_input.hpp +++ /dev/null @@ -1,1191 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ -/*! \file - \brief Functor performing elementwise operations used by epilogues. -*/ - -#pragma once - -#include "cute/atom/copy_traits_sm90_tma.hpp" -#include "cute/tensor.hpp" -#include "cutlass/arch/barrier.h" -#include "cutlass/cuda_host_adapter.hpp" -#include "cutlass/cutlass.h" -#include "cutlass/detail/collective.hpp" -#include "cutlass/detail/layout.hpp" -#include "cutlass/epilogue/collective/collective_builder.hpp" -#include "cutlass/epilogue/collective/detail.hpp" -#include "cutlass/epilogue/dispatch_policy.hpp" -#include "cutlass/epilogue/fusion/callbacks.hpp" -#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp" -#include "cutlass/epilogue/thread/scale_type.h" -#include "cutlass/trace.h" - -///////////////////////////////////////////////////////////////////////////////////////////////// - -namespace tensorrt_llm::cutlass_extensions::epilogue::collective { - -using namespace cute; -using cutlass::Array; -using cutlass::canonical_warp_idx_sync; -using cutlass::CudaHostAdapter; -using cutlass::NumThreadsPerWarp; -using cutlass::NumWarpsPerWarpGroup; -using cutlass::epilogue::Sm90PtrArrayTmaWarpSpecialized; -namespace detail = cutlass::epilogue::collective::detail; -namespace fusion = cutlass::epilogue::fusion; - -///////////////////////////////////////////////////////////////////////////////////////////////// - -template -class Sm90MixedInputPtrArrayTmaWarpSpecializedEpilogue { - public: - // - // Type Aliases - // - using DispatchPolicy = - Sm90PtrArrayTmaWarpSpecialized; - using CtaTileMNK = CtaTileMNK_; - using EpilogueTile = EpilogueTile_; - using FusionCallbacks = FusionCallbacks_; - using ElementC = ElementC_; - using StrideC = StrideC_; - using InternalStrideC = cute::remove_pointer_t; - using ElementD = ElementD_; - using StrideD = StrideD_; - using InternalStrideD = cute::remove_pointer_t; - using CopyOpG2S = CopyOpG2S_; - using SmemLayoutAtomC = SmemLayoutAtomC_; - using CopyOpS2R = CopyOpS2R_; - using CopyOpS2G = CopyOpS2G_; - using SmemLayoutAtomD = SmemLayoutAtomD_; - using CopyOpR2S = CopyOpR2S_; - using CopyAtomC = CopyAtomC_; - using CopyOpR2R = CopyOpR2R_; - - using ThreadEpilogueOp = - typename cutlass::epilogue::fusion::FusionCallbacksTraits::Operation; - using GmemTiledCopyC = CopyOpG2S; - using GmemTiledCopyD = CopyOpS2G; - - static_assert(!is_layout::value && is_tuple::value, - "EpilogueTile must be a cute::Tile or cute::Shape"); - static_assert(cute::rank(CtaTileMNK{}) == 3, "CtaTileMNK must be rank-3: [CTA_M, CTA_N, CTA_K]"); - static_assert(cute::rank(EpilogueTile{}) == 2, - "EpilogueTile must be rank-2: [EPI_TILE_M, EPI_TILE_N]"); - static_assert(size<0>(CtaTileMNK{}) % size<0>(shape(EpilogueTile{})) == 0, - "EPI_TILE_M must divide CTA_M"); - static_assert(size<1>(CtaTileMNK{}) % size<1>(shape(EpilogueTile{})) == 0, - "EPI_TILE_N must divide CTA_N"); - static_assert(cute::rank(InternalStrideC{}) == 3, "StrideC must be rank-3: [M, N, L]"); - static_assert(cute::rank(InternalStrideD{}) == 3, "StrideD must be rank-3: [M, N, L]"); - - private: - constexpr static bool is_source_supported = not cute::is_void_v; - constexpr static bool is_destination_supported = not cute::is_void_v; - using NonVoidElementD = cute::conditional_t, ElementD>; - static_assert(not cute::is_void_v, "SmemElementD is void"); - using NonVoidElementC = cute::conditional_t; // prevents void ref breakages - - using SmemElementC = typename cutlass::detail::get_unpacked_element_type::type; - using SmemElementD = typename cutlass::detail::get_unpacked_element_type::type; - - constexpr static int StagesC = StagesC_; - constexpr static int StagesD = StagesD_; - constexpr static bool ReuseSmemC = ReuseSmemC_ and is_destination_supported; - constexpr static bool DelayTmaStore = DelayTmaStore_; - - constexpr static bool is_m_major_C = detail::is_m_major(); - constexpr static bool is_m_major_D = detail::is_m_major(); - - constexpr static bool is_im2col_C = cute::is_same_v; - constexpr static bool is_im2col_D = cute::is_same_v; - - // Check if register transformation is needed before copying register to shared memory. - constexpr static bool IsUseR2R = !cute::is_void_v; - - using SmemLayoutC = decltype(tile_to_shape( - SmemLayoutAtomC{}, - make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}), Int{}), - cute::conditional_t, Step<_1, _2, _3>>{})); - using SmemLayoutD = decltype(tile_to_shape( - SmemLayoutAtomD{}, - make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}), - Int{}), - cute::conditional_t, Step<_1, _2, _3>>{})); - - constexpr static bool support_smem_reuse = - is_source_supported && is_destination_supported && StagesD <= StagesC && - cosize(take<0, 2>(SmemLayoutC{})) == cosize(take<0, 2>(SmemLayoutD{})); - static_assert(not(ReuseSmemC && not support_smem_reuse), "Smem reuse requirements not met"); - - constexpr static size_t SmemAlignmentD = cutlass::detail::alignment_for_swizzle(SmemLayoutD{}); - constexpr static size_t SmemAlignmentC = cutlass::detail::alignment_for_swizzle(SmemLayoutC{}); - constexpr static size_t MaxSmemAlignment = cute::max(SmemAlignmentC, SmemAlignmentD); - - using SmemArrayTypeC = cute::ArrayEngine>; - using SmemArrayTypeD = cute::ArrayEngine>; - - using EmptyType = cute::tuple<>; - using SmemCStorage = - cute::conditional_t; - using SmemDStorage = cute::conditional_t; - - struct CollectiveStorageWithC { - alignas(SmemAlignmentC) ArrayEngine> smem_C; - alignas(SmemAlignmentD) ArrayEngine> smem_D; - }; - - union CollectiveStorageWithoutC { - cute::array smem_C; - alignas(SmemAlignmentD) ArrayEngine> smem_D; - }; - - union CollectiveStorageReuseC { - alignas(MaxSmemAlignment) ArrayEngine> smem_C; - alignas(MaxSmemAlignment) ArrayEngine> smem_D; - }; - - public: - // TMA pipeline for loading C - using LoadPipeline = cutlass::PipelineTransactionAsync; - using LoadPipelineState = cutlass::PipelineState; - constexpr static uint32_t TmaTransactionBytes = - (size(take<0, 2>(SmemLayoutC{})) * static_cast(sizeof_bits::value)) / - 8; - constexpr static bool RequiresTransactionBytes = true; - - constexpr static int NumEpilogueWarpGroups = NumEpilogueWarpGroups_; - - // TMA pipeline for storing D - using StorePipeline = - cute::conditional_t, - cutlass::PipelineTmaStore>; - using StorePipelineState = cutlass::PipelineState; - - struct SharedStorage { - struct TensorStorage { - using CollectiveStorage = cute::conditional_t< - not is_source_supported, CollectiveStorageWithoutC, - cute::conditional_t>; - CollectiveStorage collective; - - using FusionStorage = typename FusionCallbacks::SharedStorage; - FusionStorage thread; - } tensors; - - struct TensorMapStorage : cute::aligned_struct<128, _0> { - cute::TmaDescriptor smem_tensormap_C; - cute::array smem_tensormap_D; - } tensormaps; - - using PipelineStorage = typename LoadPipeline::SharedStorage; - PipelineStorage pipeline; - }; - using TensorStorage = typename SharedStorage::TensorStorage; - using TensorMapStorage = typename SharedStorage::TensorMapStorage; - using PipelineStorage = typename SharedStorage::PipelineStorage; - - static constexpr bool IsGroupedGemmKernel = !cute::is_same_v; - - // Host side epilogue arguments - struct Arguments { - typename FusionCallbacks::Arguments thread{}; - ElementC const** ptr_C = nullptr; - StrideC dC; - ElementD** ptr_D = nullptr; - StrideD dD; - }; - - // Device side epilogue params - struct Params { - using TMA_C = decltype(make_tma_copy( - CopyOpG2S{}, - make_tensor(make_gmem_ptr(static_cast(nullptr)), - repeat_like(InternalStrideC{}, int32_t(0)), InternalStrideC{}), - take<0, 2>(SmemLayoutC{}), EpilogueTile{}, _1{})); - - using TMA_D = decltype(make_tma_copy( - CopyOpS2G{}, - make_tensor(make_gmem_ptr(static_cast(nullptr)), - repeat_like(InternalStrideD{}, int32_t(0)), InternalStrideD{}), - take<0, 2>(SmemLayoutD{}), EpilogueTile{}, _1{})); - - typename FusionCallbacks::Params thread{}; - TMA_C tma_load_c; - TMA_D tma_store_d; - cute::TmaDescriptor* tensormaps; - ElementC const** ptr_C; - StrideC dC; - ElementD** ptr_D; - StrideD dD; - uint32_t tma_transaction_bytes = TmaTransactionBytes; - }; - - // - // Methods - // - - template - static constexpr Params to_underlying_arguments(ProblemShape const& problem_shape, - Arguments const& args, - [[maybe_unused]] void* workspace) { - // These tensor shapes (only applicable for grouped gemm) and pointers are only used to create - // tensormap/tma desc. These will be replaced with correct values before the initial tma load. - auto init_shape = - repeat_like(append<4>(typename ProblemShape::UnderlyingProblemShape{}, 1), int32_t(1)); - auto init_M = get<0>(init_shape); - auto init_N = get<1>(init_shape); - auto init_L = get<3>(init_shape); - - static_assert(!is_im2col_C and !is_im2col_D, "Im2Col not supported on C or D"); - - InternalStrideC stride_c; - InternalStrideD stride_d; - if constexpr (IsGroupedGemmKernel) { - // Strides for Grouped Gemm will be replaced prior to the first access regardless. - stride_c = InternalStrideC{}; - stride_d = InternalStrideD{}; - } else { - // Tensor shapes for Ptr-Array are initialized correctly only here. - auto problem_shape_MNKL = append<4>(problem_shape.get_host_problem_shape(0), 1); - init_M = get<0>(problem_shape_MNKL); - init_N = get<1>(problem_shape_MNKL); - init_L = get<3>(problem_shape_MNKL); - - stride_c = args.dC; - stride_d = args.dD; - } - - uint32_t transaction_bytes = TmaTransactionBytes; - typename Params::TMA_C tma_load_c{}; - if constexpr (is_source_supported) { - ElementC const* ptr_C_first_batch = reinterpret_cast(args.ptr_C); - Tensor tensor_c = - make_tensor(ptr_C_first_batch, - make_layout(make_shape(init_M, init_N, init_L), append<3>(stride_c, _0{}))); - tma_load_c = - make_tma_copy(CopyOpG2S{}, tensor_c, take<0, 2>(SmemLayoutC{}), EpilogueTile{}, _1{}); - } - - typename Params::TMA_D tma_store_d{}; - if constexpr (is_destination_supported) { - ElementD const* ptr_D_first_batch = reinterpret_cast(args.ptr_D); - Tensor tensor_d = - make_tensor(ptr_D_first_batch, - make_layout(make_shape(init_M, init_N, init_L), append<3>(stride_d, _0{}))); - tma_store_d = - make_tma_copy(CopyOpS2G{}, tensor_d, take<0, 2>(SmemLayoutD{}), EpilogueTile{}, _1{}); - } - - auto fusion_workspace = static_cast(workspace); - auto fusion_workspace_size = FusionCallbacks::get_workspace_size(problem_shape, args.thread); - auto tma_descriptor_workspace = reinterpret_cast( - static_cast(workspace) + fusion_workspace_size); - - return { - FusionCallbacks::to_underlying_arguments(problem_shape, args.thread, fusion_workspace), - tma_load_c, - tma_store_d, - tma_descriptor_workspace, - args.ptr_C, - args.dC, - args.ptr_D, - args.dD, - transaction_bytes, - }; - } - - template - static size_t get_workspace_size(ProblemShape const& problem_shape, Arguments const& args, - int sm_count) { - constexpr uint32_t NumInputTensors = - NumEpilogueWarpGroups + (cute::is_void_v ? 0 : 1); - auto descriptors_shape = cute::make_shape(sm_count, Int{}); - constexpr size_t SizeOfCuTensorMap = sizeof(cute::TmaDescriptor); - - // Allocate gmem space for input tensormaps per each SM, A tensormap copies followed by B - // tensormap copies - return (size(descriptors_shape) * SizeOfCuTensorMap) + - FusionCallbacks::get_workspace_size(problem_shape, args.thread); - } - - template - static cutlass::Status initialize_workspace(ProblemShape const& problem_shape, - Arguments const& args, void* workspace, - cudaStream_t stream, - CudaHostAdapter* cuda_adapter = nullptr) { - return FusionCallbacks::initialize_workspace(problem_shape, args.thread, workspace, stream, - cuda_adapter); - } - - template - static bool can_implement(ProblemShape problem_shape, [[maybe_unused]] Arguments const& args) { - bool implementable = true; - bool fusion_implementable = true; - - if (problem_shape.is_host_problem_shape_available()) { - for (int i = 0; i < problem_shape.groups(); ++i) { - auto problem_shape_MNKL = append<4>(problem_shape.get_host_problem_shape(i), 1); - auto [M, N, K, L] = problem_shape_MNKL; - - if constexpr (is_destination_supported) { - constexpr int tma_alignment_bits_D = - cutlass::detail::get_output_alignment_bits(); - constexpr int min_tma_aligned_elements_D = - tma_alignment_bits_D / cutlass::sizeof_bits::value; - implementable = - implementable && cutlass::detail::check_alignment( - cute::make_shape(M, N, L), InternalStrideD{}); - } - - if constexpr (is_source_supported) { - constexpr int tma_alignment_bits_C = - cutlass::detail::get_input_alignment_bits(); - constexpr int min_tma_aligned_elements_C = - tma_alignment_bits_C / cutlass::sizeof_bits::value; - implementable = - implementable && cutlass::detail::check_alignment( - cute::make_shape(M, N, L), InternalStrideC{}); - } - - fusion_implementable = - fusion_implementable && FusionCallbacks::can_implement(problem_shape_MNKL, args.thread); - } - } else { - CUTLASS_TRACE_HOST( - " CAN IMPLEMENT: Ignoring check to can implement because host problem shape is not " - "available.\n"); - } - - if (!implementable) { - CUTLASS_TRACE_HOST( - " CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for " - "TMA.\n"); - } - - if (!fusion_implementable) { - CUTLASS_TRACE_HOST( - " CAN IMPLEMENT: Problem Size doesn't meet the minimum requirements for " - "FusionCallbacks.\n"); - } - - bool beta_implementable = true; - - if (cute::is_void_v || args.ptr_C == nullptr) { - if constexpr (detail::has_beta::value) { - beta_implementable = args.thread.beta == 0.0; - } - if constexpr (detail::has_beta_ptr::value) { - beta_implementable = beta_implementable && args.thread.beta_ptr == nullptr; - } - if constexpr (detail::has_beta_ptr_array::value) { - beta_implementable = beta_implementable && args.thread.beta_ptr_array == nullptr; - } - } - - if (!beta_implementable) { - CUTLASS_TRACE_HOST( - " CAN IMPLEMENT: Beta/beta pointer was set, but epilogue is sourceless (void-C).\n"); - } - - return implementable && fusion_implementable && beta_implementable; - } - - template - CUTLASS_HOST_DEVICE static constexpr int get_load_pipe_increment(TileShapeMNK tile_shape_MNK) { - // Compute number of epilogue subtiles - return size<1>(zipped_divide(make_layout(take<0, 2>(tile_shape_MNK)), EpilogueTile{})); - } - - template - CUTLASS_HOST_DEVICE static constexpr int get_store_pipe_increment(TileShapeMNK tile_shape_MNK) { - return get_load_pipe_increment(tile_shape_MNK); - } - - CUTLASS_HOST_DEVICE - Sm90MixedInputPtrArrayTmaWarpSpecializedEpilogue(Params const& params_, - TensorStorage& shared_tensors) - : params(params_), fusion_callbacks(params_.thread, shared_tensors.thread) {} - - CUTLASS_DEVICE - bool is_producer_load_needed() const { return fusion_callbacks.is_producer_load_needed(); } - - CUTLASS_DEVICE auto load_init(Params const& params, TensorMapStorage& shared_tensormaps, - int32_t sm_count, int32_t sm_idx) { - // Initialize tma for loading - constexpr bool IsLoad = true; - auto load_tensormaps = tensormaps_init(params, shared_tensormaps, sm_count, sm_idx, 0); - return load_tensormaps; - } - - template )> - CUTLASS_DEVICE auto load(LoadPipeline load_pipeline, LoadPipelineState load_pipe_producer_state, - ProblemShapeMNKL problem_shape_mnkl, TileShapeMNK tile_shape_MNK, - TileCoordMNKL tile_coord_mnkl, TiledMma tiled_mma, int thread_idx, - TensorStorage& shared_tensors, TensorMapC const& load_tensormap, - int subtile_idx = -1, bool wait_until_load_finishes = false) { - using namespace cute; - - // Indexing variables - auto [M, N, K, L] = problem_shape_mnkl; - auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl; - - static_assert(!is_im2col_D, "Do not support im2col"); - - auto coord_shape = append<3>(make_shape(m_coord, n_coord), Int<0>{}); - - // Represent the full source tensor, slice to get the tile this CTA is currently responsible for - Tensor mC_mn = - params.tma_load_c.get_tma_tensor(append<3>(make_shape(M, N), Int<1>{})); // (M,N,L) - Tensor mC = coalesce(mC_mn, take<0, 2>(CtaTileMNK{})); - Tensor gC = local_tile(mC, take<0, 2>(CtaTileMNK{}), coord_shape); // (CTA_M,CTA_N) - - // Apply epilogue subtile, get matching smem tensor - auto ptr_sC = shared_tensors.collective.smem_C.begin(); - Tensor gC_epi = flat_divide(gC, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) - Tensor sC_epi = - make_tensor(make_smem_ptr(ptr_sC), SmemLayoutC{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_C) - - // Prepare the thread(b)lock's (G)mem to (S)mem TMA tiled copy (bGS_) - ThrCopy thrblk_g2s = params.tma_load_c.get_slice(Int<0>{}); - Tensor bGS_gC = thrblk_g2s.partition_S(gC_epi); // (G2S,G2S_M,G2S_N,EPI_M,EPI_N) - Tensor bGS_sC = thrblk_g2s.partition_D(sC_epi); // (G2S,G2S_M,G2S_N,PIPE_C) - - // Get the fusion callbacks for the producer load warp - auto pld_args = cutlass::epilogue::fusion::detail::ProducerLoadArgs{ - problem_shape_mnkl, CtaTileMNK{}, tile_coord_mnkl, tiled_mma, EpilogueTile{}, thread_idx}; - auto pld_callbacks = fusion_callbacks.get_producer_load_callbacks(pld_args); - bool is_C_load_needed = is_source_supported && fusion_callbacks.is_C_load_needed(); - - LoadPipelineState last_load_producer_state = load_pipe_producer_state; - - // Predication for TMA load (one thread issues TMA load) - bool issue_tma_load = cute::elect_one_sync(); - - // Pre-loop fusion callback entry point - pld_callbacks.begin(); - - LoadPipelineState prior_state = load_pipe_producer_state; - - bool did_load = false; - - CUTLASS_PRAGMA_UNROLL - for (int epi_n = 0; epi_n < size<3>(gC_epi); ++epi_n) { - CUTLASS_PRAGMA_UNROLL - for (int epi_m = 0; epi_m < size<2>(gC_epi); ++epi_m) { - if (subtile_idx != -1 && - (epi_n * static_cast(size<2>(gC_epi)) + epi_m) != subtile_idx) { - continue; - } - - // Acquire the lock for this stage - constexpr uint16_t mcast_mask = 0; - uint64_t* tma_barrier = load_pipeline.producer_get_barrier(load_pipe_producer_state); - - load_pipeline.producer_acquire(load_pipe_producer_state); - - // Loop fusion callback entry point - pld_callbacks.step(tma_barrier, epi_m, epi_n, load_pipe_producer_state.count(), - issue_tma_load); - - // Execute the TMA load for C if needed - if (is_C_load_needed) { - if (issue_tma_load) { - copy(params.tma_load_c.with(load_tensormap, *tma_barrier, mcast_mask), - bGS_gC(_, _, _, epi_m, epi_n), bGS_sC(_, _, _, load_pipe_producer_state.index())); - load_pipeline.producer_expect_transaction(load_pipe_producer_state); - } - last_load_producer_state = load_pipe_producer_state; - did_load = true; - } - - // Commit TMA loads for this stage and release the lock - load_pipeline.producer_commit(load_pipe_producer_state); - ++load_pipe_producer_state; - } - } - - // Post-loop fusion callback entry point - pld_callbacks.end(); - - if (wait_until_load_finishes && did_load) { - LoadPipelineState epi_load_pipe_tma_consumer_state = {last_load_producer_state.index(), - !last_load_producer_state.phase(), - last_load_producer_state.count()}; - load_pipeline.consumer_wait(epi_load_pipe_tma_consumer_state); - } - - return load_pipe_producer_state; - } - - CUTLASS_DEVICE auto load_tail(LoadPipeline load_pipeline, - LoadPipelineState load_pipe_producer_state) { - if (!fusion_callbacks.is_producer_load_needed()) { - return load_pipe_producer_state; - } - - bool issue_tma_load = cute::elect_one_sync(); - if (issue_tma_load) { - load_pipeline.producer_tail(load_pipe_producer_state); - } - - return load_pipe_producer_state; - } - - template - CUTLASS_DEVICE auto store(LoadPipeline load_pipeline, LoadPipelineState load_pipe_consumer_state, - StorePipeline store_pipeline, - StorePipelineState store_pipe_producer_state, - ProblemShapeMNKL problem_shape_mnkl, TileShapeMNK tile_shape_MNK, - TileCoordMNKL tile_coord_mnkl, - cute::Tensor accumulators, TiledMma tiled_mma, - int thread_idx, TensorStorage& shared_tensors, - TensorMapD const& store_tensormap, int subtile_idx = -1) { - using namespace cute; - using ElementAccumulator = typename AccEngine::value_type; - using ElementCompute_ = - typename cutlass::epilogue::fusion::FusionCallbacksTraits::ElementCompute; - using ElementCompute = - cute::conditional_t, ElementAccumulator, ElementCompute_>; - - static_assert(is_rmem::value, "Accumulator must be RF resident."); - static_assert(rank(AccLayout{}) == 3, "Accumulator must be MMA-partitioned: (MMA,MMA_M,MMA_N)"); - static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4"); - static_assert(is_static::value, "TileShapeMNK must be static"); - static_assert(rank(TileShapeMNK{}) == 3, "TileShapeMNK must be rank 3"); - static_assert(rank(TileCoordMNKL{}) == 4, "TileCoordMNKL must be rank 4"); - - // Indexing variables - auto [M, N, K, L] = problem_shape_mnkl; - auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl; - - static_assert(!is_im2col_D, "Do not support im2col"); - - auto coord_shape = append<3>(make_shape(m_coord, n_coord), Int<0>{}); - - // Represent the full output tensor, slice to get the tile this CTA is responsible for - Tensor mD_mn = - params.tma_store_d.get_tma_tensor(append<3>(make_shape(M, N), Int<1>{})); // (M,N,L) - - Tensor mD = coalesce(mD_mn, take<0, 2>(CtaTileMNK{})); - Tensor gD = local_tile(mD, take<0, 2>(CtaTileMNK{}), coord_shape); // (CTA_M,CTA_N) - - // Apply epilogue subtiling - Tensor gD_epi = flat_divide(gD, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) - - // Construct the corresponding pipelined smem tensors - auto ptr_sC = shared_tensors.collective.smem_C.begin(); - auto ptr_sD = shared_tensors.collective.smem_D.begin(); - Tensor sC_epi = cute::as_position_independent_swizzle_tensor( - make_tensor(make_smem_ptr(ptr_sC), SmemLayoutC{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_C) - Tensor sD_epi = cute::as_position_independent_swizzle_tensor( - make_tensor(make_smem_ptr(ptr_sD), SmemLayoutD{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D) - - TiledCopy tiled_copy_C_atom = make_tiled_copy_C_atom(CopyAtomC{}, tiled_mma); - - // (t)hread-partition for (r)egister to (r)egister copy (tRR_) - TiledCopy tiled_r2r = [&]() { - if constexpr (IsUseR2R) { - return make_tiled_copy_S(Copy_Atom{}, tiled_copy_C_atom); - } else { - return make_tiled_copy_S( - Copy_Atom, ElementCompute>{}, - tiled_copy_C_atom); - } - }(); - ThrCopy thread_r2r = tiled_r2r.get_slice(thread_idx); - - // (t)hread-partition for (r)egister to (s)mem copy (tRS_) - TiledCopy tiled_r2s = [&]() { - if constexpr (IsUseR2R) { - return make_tiled_copy_D(Copy_Atom{}, tiled_r2r); - } else { - return make_tiled_copy_S(Copy_Atom{}, tiled_copy_C_atom); - } - }(); - ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx); - Tensor tRS_rAcc = thread_r2s.retile_S(accumulators); // ((R2S,R2S_V),MMA_M,MMA_N) - Tensor tRS_sD = thread_r2s.partition_D(sD_epi); // (R2S,R2S_M,R2S_N,PIPE_D) - - auto mma_tile_m = size<0>(TileShapeMNK{}) / size<1>(tRS_rAcc); - auto mma_tile_n = size<1>(TileShapeMNK{}) / size<2>(tRS_rAcc); - auto epi_tile_m = size<0>(EpilogueTile{}); - auto epi_tile_n = size<1>(EpilogueTile{}); - - // Allocate D registers - Layout tRS_rD_layout = make_layout(take<0, 3>(shape(thread_r2s.partition_S(sD_epi)))); - Tensor tRS_rD = make_tensor(tRS_rD_layout); // (R2S,R2S_M,R2S_N) - - // Vectorized fragment view - constexpr int FragmentSize = DispatchPolicy::FragmentSize; - Tensor tRS_rAcc_frg = recast>(tRS_rAcc); - Tensor tRS_rD_frg = recast>(tRS_rD); - CUTE_STATIC_ASSERT(size<0>(tRS_rAcc) % FragmentSize == 0, - "Fragment size does not vectorize properly"); - - // (t)hread-partition for (s)mem to (r)egister copy (tSR_) - TiledCopy tiled_s2r = - make_tiled_copy_S(Copy_Atom{}, tiled_copy_C_atom); - ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx); - Tensor tSR_sC = thread_s2r.partition_S(sC_epi); // (S2R,S2R_M,S2R_N,PIPE_C) - Layout tSR_rC_layout = thread_s2r.retile_D(tRS_rD).layout(); // (S2R,S2R_M,S2R_N) - - // Allocate C registers - // If C smem load is a non-vectorized dst(i) = src(i) then we can allocate C registers directly - // in the compute type to eliminate some redundant pack+unpack instruction sequences for - // sub-word types - constexpr bool IsDirectS2R = - cute::is_same_v> && - decltype(max_common_vector(tSR_rC_layout, tSR_sC.layout()))::value <= 1; - using RegisterElementC = cute::conditional_t; - Tensor tRS_rC = make_tensor(tRS_rD_layout); // (R2S,R2S_M,R2S_N) - Tensor tSR_rC = thread_s2r.retile_D(tRS_rC); // (S2R,S2R_M,S2R_N) - - // thread(b)lock-partition for (s)mem to (g)mem copy (bSG_) - ThrCopy thrblk_s2g = params.tma_store_d.get_slice(Int<0>{}); - Tensor bSG_sD = thrblk_s2g.partition_S(sD_epi); // (S2G,S2G_M,S2G_N,PIPE_D) - Tensor bSG_gD = thrblk_s2g.partition_D(gD_epi); // (S2G,S2G_M,S2G_N,EPI_M,EPI_N) - - // OOB predication for tile quantization "residue" - // Absolute coordinate tensors (dynamic) - Tensor mD_crd = make_identity_tensor(make_shape(M, N)); // (M,N) - Tensor cD_mn = local_tile(mD_crd, take<0, 2>(CtaTileMNK{}), - make_coord(m_coord, n_coord)); // (CTA_M,CTA_N) - Tensor tRS_cD_mn = thread_r2s.partition_S( - flat_divide(cD_mn, EpilogueTile{})); // (R2S,R2S_M,R2S_N,EPI_M,EPI_N) - // Relative coordinate tensors (static) - Tensor cD = make_coord_tensor(cD_mn.layout()); // (CTA_M,CTA_N) - Tensor tRS_cD = make_coord_tensor(tRS_cD_mn.layout()); // (R2S,R2S_M,R2S_N,EPI_M,EPI_N) - // Subtract the global "bottom right" corner from the local "top left" corner to get the max - // relative coordinate - auto residue_cD = make_coord(M, N) - cD_mn(_0{}); // (m,n) - auto residue_tRS_cD = make_coord(M, N) - tRS_cD_mn(_0{}); // (m,n) - - CUTE_STATIC_ASSERT(epi_tile_m % mma_tile_m == 0, "MMA_TILE_M must divide EPI_TILE_M"); - - CUTE_STATIC_ASSERT(mma_tile_n % epi_tile_n == 0, "EPI_TILE_N must divide MMA_TILE_N"); - // Get TiledCopy for partition reference when consumer store. - TiledCopy tiled_copy_partition_ref = - make_tiled_copy_S(Copy_Atom{}, tiled_copy_C_atom); - // Get the fusion callbacks for the consumer store warps - constexpr bool RefSrc = true; // Register tensors reference R2S copy src layout - auto cst_args = cutlass::epilogue::fusion::detail::ConsumerStoreArgs{problem_shape_mnkl, - CtaTileMNK{}, - tile_coord_mnkl, - tiled_mma, - EpilogueTile{}, - tiled_copy_partition_ref, - cD, - residue_cD, - tRS_cD, - residue_tRS_cD, - tRS_rC, - thread_idx}; - auto cst_callbacks = fusion_callbacks.template get_consumer_store_callbacks(cst_args); - bool is_producer_load_needed = fusion_callbacks.is_producer_load_needed(); - bool is_C_load_needed = is_source_supported && fusion_callbacks.is_C_load_needed(); - - // Thread synchronizer for previously issued waits or fences - // to ensure visibility of smem reads/writes to threads or TMA unit - auto synchronize = [&]() { - cutlass::arch::NamedBarrier::sync(size(TiledMma{}), - cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); - }; - - // Predication for TMA store (one warp issues TMA store) - bool issue_tma_store = (thread_idx / NumThreadsPerWarp) == 0; - - // In the reuse smem configuration we have StagesC smem buffers and at most StagesD committed - // TMA stores in flight. The TMA store pipeline producer acquire returns when at most StagesD-1 - // committed stores are in-flight, so we can only guarantee store completion after StagesD - // iterations, then we can begin issuing releases on the smem buffer locks. - // store_pipe_producer_state tracks the acquire and load_pipe_consumer_state tracks the release, - // in circular buffer fashion. - LoadPipelineState load_wait_state = load_pipe_consumer_state; - if constexpr (ReuseSmemC) { - load_wait_state = store_pipe_producer_state; - load_wait_state.phase_ ^= 1; - } - - // We can delay issue of TMA store by one iteration to achieve better interleaving of non-TMA - // instructions Sync requirements of smem reuse may preclude this optimization Delayed stores - // cause delayed stage releases which causes deadlock when StagesC == StagesD - int epi_m_prev = 0, epi_n_prev = 0; - static_assert(not(DelayTmaStore and ReuseSmemC and StagesC <= StagesD), - "This TMA epilogue configuration will deadlock"); - - // The TMA store sequence for one subtile iteration - auto tma_store_fn = [&](int epi_m, int epi_n) { - // Write the tile from smem to gmem with TMA - cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA - synchronize(); // ensure all threads have issued their async fence - if constexpr (is_destination_supported) { - if (issue_tma_store) { - copy(params.tma_store_d.with(store_tensormap), - bSG_sD(_, _, _, store_pipe_producer_state.index()), bSG_gD(_, _, _, epi_m, epi_n)); - } - } - - // Post async fence, pre TMA commit callback entry point - cst_callbacks.tma_store(epi_m, epi_n, store_pipe_producer_state.count(), issue_tma_store); - - // Commit the TMA stores for this stage - if (issue_tma_store) { - store_pipeline.producer_commit(store_pipe_producer_state); - } - ++store_pipe_producer_state; - ++issued_stores; - - // Wait for the next smem buffer to be available - if (issue_tma_store) { - store_pipeline.producer_acquire(store_pipe_producer_state); - } - synchronize(); - - if constexpr (ReuseSmemC) { - // producer_acquire returns when at most StagesD-1 committed stores are pending - bool store_finished = issued_stores > StorePipeline::UnacquiredStages; - // Let dma warp know earliest smem buffer is consumed and empty after StagesD producer - // commits - if (store_finished) { - if (is_producer_load_needed) { - load_pipeline.consumer_release(load_pipe_consumer_state); - } - ++load_pipe_consumer_state; - } - } - }; - - // - // BEGIN EPILOGUE - // - - // Pre-loop fusion callback entry point - cst_callbacks.begin(); - if (cst_callbacks.begin_sync_needed()) { - synchronize(); - } - - // For each output tile - CUTLASS_PRAGMA_UNROLL - for (int epi_n = 0; epi_n < size<3>(gD_epi); ++epi_n) { - CUTLASS_PRAGMA_UNROLL - for (int epi_m = 0; epi_m < size<2>(gD_epi); ++epi_m) { - bool is_first_iteration = epi_m == 0 && epi_n == 0; - bool is_last_iteration = epi_m == size<2>(gD_epi) - 1 && epi_n == size<3>(gD_epi) - 1; - - if (subtile_idx != -1 && - (epi_n * static_cast(size<2>(gD_epi)) + epi_m) != subtile_idx) { - continue; - } - - cst_callbacks.begin_loop(epi_m, epi_n); - - if (is_producer_load_needed) { - // Wait for the producer load to fill smem - load_pipeline.consumer_wait(load_wait_state); - - if (is_C_load_needed) { - // Copy source tile from smem to register - copy(tiled_s2r, tSR_sC(_, _, _, load_wait_state.index()), tSR_rC); - } - } - - // First loop fusion callback entry point - cst_callbacks.previsit(epi_m, epi_n, load_wait_state.count(), is_producer_load_needed); - - if (is_producer_load_needed) { - if constexpr (not ReuseSmemC) { - // Let producer load warp know smem buffers are consumed and empty - cutlass::arch::fence_view_async_shared(); - load_pipeline.consumer_release(load_pipe_consumer_state); - ++load_pipe_consumer_state; - } - ++load_wait_state; - } - - int mma_m = epi_m; - int mma_n = (epi_n * size<1>(EpilogueTile{})) / mma_tile_n; - Tensor tRS_rAcc_frg_mn = tRS_rAcc_frg(_, mma_m, mma_n); - - // Vectorized fragment loop with visitor callback entry point - int epi_n_in_mma = epi_n % (mma_tile_n / epi_tile_n); - int r2s_v = epi_n_in_mma * size(tRS_rD_frg); - CUTLASS_PRAGMA_UNROLL - for (int epi_v = 0; epi_v < size(tRS_rD_frg); ++epi_v) { - tRS_rD_frg(epi_v) = - cst_callbacks.visit(tRS_rAcc_frg_mn(r2s_v + epi_v), epi_v, epi_m, epi_n); - } - // The latest we can delay the TMA store is right before the smem store of the next - // iteration since the current TMA store needs to be committed before we can acquire the - // next smem buffer - if constexpr (DelayTmaStore) { - // Issue TMA stores for the previous subtile - if (not is_first_iteration and subtile_idx == -1) { - tma_store_fn(epi_m_prev, epi_n_prev); - } - epi_m_prev = epi_m; - epi_n_prev = epi_n; - } - - // Smem reduction callback entry point using current store buffer for workspace - cst_callbacks.reduce(sD_epi(_, _, store_pipe_producer_state.index()), synchronize, epi_m, - epi_n, is_last_iteration, tRS_rD_frg); - - // Copy tile from register to regiser if needed - if constexpr (IsUseR2R) { - // retile source and destination for tiled_r2r - Tensor tRR_rD_src = thread_r2r.retile_S(tRS_rD); // (R2R,R2R_M,R2R_N,EPI_M,EPI_N) - Tensor tRR_rD_dst = thread_r2r.retile_D(tRS_rD); // (R2R,R2R_M,R2R_N,EPI_M,EPI_N) - - // Output needs register shuffling before copying to shared memory. - copy(tiled_r2r, tRR_rD_src, tRR_rD_dst); - } - - // Copy tile from register to smem - if constexpr (is_destination_supported) { - copy(tiled_r2s, tRS_rD, tRS_sD(_, _, _, store_pipe_producer_state.index())); - } - - // Post reduction, pre TMA store callback entry point - constexpr bool issue_smem_store = true; // No smem store predication - cst_callbacks.postreduce(epi_m, epi_n, store_pipe_producer_state.count(), issue_smem_store); - - if constexpr (not DelayTmaStore) { - // Issue TMA stores for this subtile - tma_store_fn(epi_m, epi_n); - } - - cst_callbacks.end_loop(epi_m, epi_n); - - } // for epi_m - } // for epi_n - - if constexpr (DelayTmaStore) { - // Issue TMA stores for the last subtile - tma_store_fn(epi_m_prev, epi_n_prev); - } - - // Post-loop fusion callback entry point - cst_callbacks.end(); - - return cute::make_tuple(load_pipe_consumer_state, store_pipe_producer_state); - } - - CUTLASS_DEVICE auto store_tail(LoadPipeline load_pipeline, - LoadPipelineState load_pipe_consumer_state, - StorePipeline store_pipeline, - StorePipelineState store_pipe_producer_state) { - // wait for all TMA stores to complete - store_pipeline.producer_tail(store_pipe_producer_state); - // reset store counter - issued_stores = 0; - - if constexpr (ReuseSmemC) { - if (fusion_callbacks.is_producer_load_needed()) { - // Issue releases on up to StagesD-1 previously issued TMA stores - constexpr int release_stages = - cute::min(StorePipeline::UnacquiredStages, get_load_pipe_increment(CtaTileMNK{})); - CUTLASS_PRAGMA_UNROLL - for (int stage = 0; stage < release_stages; ++stage) { - load_pipeline.consumer_release(load_pipe_consumer_state); - ++load_pipe_consumer_state; - } - } - } - - return cute::make_tuple(load_pipe_consumer_state, store_pipe_producer_state); - } - - CUTLASS_DEVICE auto store_init(Params const& params, TensorMapStorage& shared_tensormaps, - int32_t sm_count, int32_t sm_idx, int32_t warp_group_idx) { - int warp_idx_in_warp_group = canonical_warp_idx_sync() % NumWarpsPerWarpGroup; - // Since only one warp issues TMA store, we only need that one warp to initialize tensormaps - if (warp_idx_in_warp_group == 0) { - // Initialize tma - constexpr bool IsLoad = false; - auto store_tensormaps = - tensormaps_init(params, shared_tensormaps, sm_count, sm_idx, warp_group_idx); - return store_tensormaps; - } - TmaDescriptor* null_tma_desc = nullptr; - return cute::make_tuple(null_tma_desc); - } - - // - // Methods to perform different parts of TMA/Tensormap modifications - // - - template - CUTLASS_DEVICE auto tensormaps_init(Params const& params, TensorMapStorage& shared_tensormaps, - int32_t sm_count, int32_t sm_idx, int32_t warp_group_idx) { - constexpr uint32_t NumInputTensors = - NumEpilogueWarpGroups + (cute::is_void_v ? 0 : 1); - Layout desc_layout = make_layout(make_shape(sm_count, Int{})); - - Tensor gmem_tensormap = make_tensor(params.tensormaps, desc_layout); // (SMs, NumInputTensors) - - if constexpr (IsLoad) { - if (is_source_supported) { - constexpr int C_tensormap_index = NumEpilogueWarpGroups; - Tensor pC_tensormap = - make_tensor(params.tma_load_c.get_tma_descriptor(), Int<1>{}, Int<1>{}); - Tensor sC_tensormap = - make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_C), Int<1>{}, Int<1>{}); - - if (cute::elect_one_sync()) { - // Bringing tensormaps from params to smem for modification later - copy(recast(pC_tensormap), recast(sC_tensormap)); - } - __syncwarp(); - return cute::make_tuple(&gmem_tensormap(sm_idx, C_tensormap_index)); - } - TmaDescriptor* null_tma_desc = nullptr; - return cute::make_tuple(null_tma_desc); - } else { - Tensor pD_tensormap = - make_tensor(params.tma_store_d.get_tma_descriptor(), Int<1>{}, Int<1>{}); - Tensor sD_tensormap = make_tensor( - make_smem_ptr(&shared_tensormaps.smem_tensormap_D[warp_group_idx]), Int<1>{}, Int<1>{}); - - if (cute::elect_one_sync()) { - // Bringing tensormaps from params to smem for modification later - copy(recast(pD_tensormap), recast(sD_tensormap)); - } - __syncwarp(); - return cute::make_tuple(&gmem_tensormap(sm_idx, warp_group_idx)); - } - } - - // Replace address for the global tensor (to be done by single thread) - template - CUTLASS_DEVICE void tensormaps_replace_global_address(TensorMapStorage& shared_tensormaps, - Params const& params, int32_t next_batch, - int32_t warp_group_idx) { - // Replacing global_address for the next batch - if constexpr (IsLoad) { - if constexpr (is_source_supported) { - if (params.ptr_C != nullptr) { - cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormaps.smem_tensormap_C, - params.ptr_C[next_batch]); - } - } - } else if constexpr (is_destination_supported) { - cute::tma_descriptor_replace_addr_in_shared_mem( - shared_tensormaps.smem_tensormap_D[warp_group_idx], params.ptr_D[next_batch]); - } - } - - // Replace dim and strides for the global tensor - used only for Grouped GEMM (to be done by - // single thread) - template - CUTLASS_DEVICE void tensormaps_replace_global_tensor_properties( - TensorMapStorage& shared_tensormaps, Params const& params, int32_t next_group, - ProblemShape_MNKL problem_shape_mnkl, int32_t warp_group_idx) { - const uint32_t M = get<0>(problem_shape_mnkl); - const uint32_t N = get<1>(problem_shape_mnkl); - // Replace all dims for consistency - constexpr int MaxTensorRank = 5; - cute::array prob_shape = {1, 1, 1, 1, 1}; - cute::array prob_stride = {0, 0, 0, 0, 0}; - - if constexpr (IsLoad) { - if constexpr (is_source_supported) { - if (params.dC != nullptr) { - ElementC const* ptr_C = nullptr; - Tensor tensor_c = - make_tensor(ptr_C, make_layout(make_shape(M, N, Int<1>{}), params.dC[next_group])); - - cute::detail::fill_tma_gmem_shape_stride(params.tma_load_c, tensor_c, prob_shape, - prob_stride); - // Convert strides to byte strides - for (uint64_t& stride : prob_stride) { - stride = (stride * sizeof_bits_v) / 8; - } - cute::tma_descriptor_replace_dims_strides_in_shared_mem( - shared_tensormaps.smem_tensormap_C, prob_shape, prob_stride); - } - } - } else if constexpr (is_destination_supported) { - ElementD const* ptr_D = nullptr; - Tensor tensor_d = - make_tensor(ptr_D, make_layout(make_shape(M, N, Int<1>{}), params.dD[next_group])); - - cute::detail::fill_tma_gmem_shape_stride(params.tma_store_d, tensor_d, prob_shape, - prob_stride); - // Convert strides to byte strides - for (uint64_t& stride : prob_stride) { - stride = (stride * sizeof_bits_v) / 8; - } - - cute::tma_descriptor_replace_dims_strides_in_shared_mem( - shared_tensormaps.smem_tensormap_D[warp_group_idx], prob_shape, prob_stride); - } - } - - template - CUTLASS_DEVICE void tensormaps_perform_update(TensorMapStorage& shared_tensormaps, - Params const& params, - cute::TmaDescriptor const* tensormap, - ProblemShape_MNKL problem_shape_mnkl, - int32_t next_batch, int32_t warp_group_idx) { - if (cute::elect_one_sync()) { - // Replacing global_address for the next batch - tensormaps_replace_global_address(shared_tensormaps, params, next_batch, - warp_group_idx); - - if constexpr (IsGroupedGemmKernel) { - // Replacing global dims and strides for the next batch - tensormaps_replace_global_tensor_properties(shared_tensormaps, params, next_batch, - problem_shape_mnkl, warp_group_idx); - } - } - } - - template - CUTLASS_DEVICE void tensormaps_cp_fence_release(TensorMapStorage& shared_tensormaps, - cute::TmaDescriptor const* tensormap, - const int32_t warp_group_idx = 0) { - // Entire warp must do this (ie its aligned) - if constexpr (IsLoad) { - if constexpr (is_source_supported) { - tma_descriptor_cp_fence_release(tensormap, shared_tensormaps.smem_tensormap_C); - } - } else if constexpr (is_destination_supported) { - tma_descriptor_cp_fence_release(tensormap, - shared_tensormaps.smem_tensormap_D[warp_group_idx]); - } - } - - template - CUTLASS_DEVICE void tensormaps_fence_acquire(cute::TmaDescriptor const* tensormap) { - if constexpr (IsLoad) { - if constexpr (is_source_supported) { - cute::tma_descriptor_fence_acquire(tensormap); - } - } else { - cute::tma_descriptor_fence_acquire(tensormap); - } - } - - private: - Params const& params; - FusionCallbacks fusion_callbacks; - int issued_stores = 0; -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - -template > -struct MixedInputSm90TmaEpilogueBuilder { - static_assert(cute::is_same_v, - "Mixed-input TMA epilogue builder is SM90-only."); - static_assert(cutlass::epilogue::collective::detail::sm90_is_ptr_array_tma_v, - "Mixed-input TMA epilogue builder expects a ptr-array TMA schedule."); - - private: - static_assert(detail::is_aligned(), - "C/D should meet TMA alignment requirement."); - - using ElementD = cute::conditional_t, - fusion::get_element_aux_t, ElementD_>; - using ElementC = cute::conditional_t, ElementD, ElementC_>; - using GmemLayoutTagC = - cute::conditional_t, GmemLayoutTagD, GmemLayoutTagC_>; - - using EpilogueTile_MN = - decltype(detail::sm90_compute_tile_shape_or_override()); - using DispatchPolicy = - decltype(detail::sm90_get_tma_dispatch_policy()); - - using GmemStrideTypeC = cutlass::detail::TagToStrideC_t; - using GmemStrideTypeD = cutlass::detail::TagToStrideC_t; - using UnderlyingGmemStrideTypeC = cute::remove_pointer_t; - using UnderlyingGmemStrideTypeD = cute::remove_pointer_t; - - using CopyOpS2G = cute::conditional_t, - SM90_TMA_STORE_IM2COL, SM90_TMA_STORE>; - using CopyOpG2S = cute::conditional_t, - SM90_TMA_LOAD_IM2COL, SM90_TMA_LOAD>; - - using CopyAtomC = - cute::conditional_t(EpilogueTile_MN{}) % 16 == 0, - Copy_Atom, - cute::conditional_t(EpilogueTile_MN{}) % 8 == 0, - Copy_Atom, void>>; - static_assert(!cute::is_same_v, "CopyAtomC cannot be void."); - - using FusionCallbacks = typename cutlass::epilogue::collective::detail::CallbacksBuilder< - DispatchPolicy, FusionOpOrCallbacks, TileShape_MNK, EpilogueTile_MN, - ElementAccumulator>::Callbacks; - - public: - using CollectiveOp = Sm90MixedInputPtrArrayTmaWarpSpecializedEpilogue< - DispatchPolicy::StagesC, DispatchPolicy::StagesD, DispatchPolicy::FragmentSize, - DispatchPolicy::ReuseSmemC, DispatchPolicy::DelayTmaStore, - DispatchPolicy::NumEpilogueWarpGroups, TileShape_MNK, EpilogueTile_MN, ElementC_, - GmemStrideTypeC, ElementD_, GmemStrideTypeD, FusionCallbacks, CopyOpG2S, - decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom()), - decltype(detail::sm90_get_smem_load_op_for_source()), - CopyOpS2G, - decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom()), - decltype(detail::sm90_get_smem_store_op_for_accumulator()), - CopyAtomC, void>; -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace tensorrt_llm::cutlass_extensions::epilogue::collective - -///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/fusion/sm90_ptr_array_per_token_scale_callbacks_tma_warpspecialized.hpp b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/fusion/sm90_ptr_array_per_token_scale_callbacks_tma_warpspecialized.hpp deleted file mode 100644 index 8c9d9ddd49c..00000000000 --- a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/fusion/sm90_ptr_array_per_token_scale_callbacks_tma_warpspecialized.hpp +++ /dev/null @@ -1,100 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -/*! \file - \brief Pointer-array row-scale fusion callbacks for the sm90 TMA warp-specialized epilogue. -*/ - -#pragma once - -#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp" - -///////////////////////////////////////////////////////////////////////////////////////////////// - -namespace cutlass::epilogue::fusion { - -///////////////////////////////////////////////////////////////////////////////////////////////// - -template , - FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest> -struct PtrArrayPerTokenScaledAcc - : ScaledAcc { - static constexpr int AlignmentScalar = AlignmentScalar_; -}; - -template , - FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest> -using Sm90PtrArrayPerTokenScaledAcc = - Sm90EVT, - Sm90RowBroadcast<0, CtaTileShapeMNK, ElementScalar*, ElementCompute, Stride<_0, _1, _0>, - AlignmentScalar>, - Sm90AccFetch>; - -template -struct FusionCallbacks< - epilogue::Sm90PtrArrayTmaWarpSpecialized, - fusion::PtrArrayPerTokenScaledAcc, - CtaTileShapeMNK, EpilogueTile> - : Sm90PtrArrayPerTokenScaledAcc< - CtaTileShapeMNK, typename cutlass::detail::get_unpacked_element_type::type, - ElementCompute, ElementScalar, AlignmentScalar, RoundStyle> { - using Impl = Sm90PtrArrayPerTokenScaledAcc< - CtaTileShapeMNK, typename cutlass::detail::get_unpacked_element_type::type, - ElementCompute, ElementScalar, AlignmentScalar, RoundStyle>; - - struct Arguments { - ElementScalar token_scale_default = ElementScalar(1); - ElementScalar const* const* token_scale_ptr_array = nullptr; - - using StrideTokenScale = Stride<_0, _1, _0>; - StrideTokenScale dTokenScale = {_0{}, _1{}, _0{}}; - - operator typename Impl::Arguments() const { - return {{token_scale_ptr_array, token_scale_default, dTokenScale}, {}, {}}; - } - }; - - using Impl::Impl; -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace cutlass::epilogue::fusion - -///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/collective/builders/sm90_gmma_builder_mixed_input.inl b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/collective/builders/sm90_gmma_builder_mixed_input.inl index f15aff928e8..750df369092 100644 --- a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/collective/builders/sm90_gmma_builder_mixed_input.inl +++ b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/collective/builders/sm90_gmma_builder_mixed_input.inl @@ -34,69 +34,14 @@ namespace cutlass::gemm::collective { ///////////////////////////////////////////////////////////////////////////////////////////////// -namespace detail { - -template -constexpr int compute_stage_count_or_override_folded_weight_scale(cute::Int stage_count) { - return stages; -} - -template -constexpr int compute_stage_count_or_override_folded_weight_scale(StageCount stage_count) { - return stages; -} - -template -constexpr int compute_stage_count_or_override_folded_weight_scale( - StageCountAutoCarveout stage_count) { - constexpr auto mainloop_pipeline_bytes = - sizeof(typename cutlass::PipelineTmaAsync<1>::SharedStorage); - constexpr auto a_bits = cute::sizeof_bits_v; - constexpr auto b_bits = cute::sizeof_bits_v; - constexpr auto s_bits = cute::sizeof_bits_v; - constexpr auto z_bits = get_bits_for_possibly_void_element(); - constexpr int scale_group_size = DefaultWeightScaleGroupSize::value; - - static_assert((size<2>(TileShapeMNK{}) % scale_group_size) == 0, - "Folded weight-scale storage requires TileK to cover complete scale groups."); - - constexpr auto scale_bytes = cutlass::bits_to_bytes(s_bits * size<0>(TileShapeMNK{}) * - size<2>(TileShapeMNK{}) / scale_group_size); - constexpr auto zero_bytes = cutlass::bits_to_bytes(z_bits * size<0>(TileShapeMNK{})); - static_assert(scale_bytes % 16 == 0, - "Folded weight-scale bulk copy must be at least 16B aligned."); - static_assert(zero_bytes % 128 == 0, "Zero bytes must be a multiple of 128"); - - constexpr int stage_bytes_ = - cutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + - cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + - scale_bytes + zero_bytes; - - constexpr int stage_bytes = - cutlass::round_up(stage_bytes_, alignment) + static_cast(mainloop_pipeline_bytes); - constexpr int carveout_bytes = cutlass::round_up(carveout_bytes_, alignment); - constexpr int capacity_bytes = capacity_bytes_ / alignment * alignment; - - constexpr int computed_stage_count = (capacity_bytes - carveout_bytes) / stage_bytes; - return computed_stage_count < 2 ? 2 : computed_stage_count; -} - -} // namespace detail - -///////////////////////////////////////////////////////////////////////////////////////////////// - // GMMA_TMA_WS_RS template + class ClusterShape_MNK, class StageCountType, class KernelScheduleType> struct CollectiveBuilderMixedInput< arch::Sm90, arch::OpClassTensorOp, ElementA_, GmemLayoutATag_, AlignmentA, ElementB_, GmemLayoutBTag_, AlignmentB, ElementAccumulator, TileShape_MNK, ClusterShape_MNK, - StageCountType, KernelScheduleType, ScaleMode, + StageCountType, KernelScheduleType, cute::enable_if_t< (cute::is_same_v || cute::is_same_v || @@ -104,7 +49,7 @@ struct CollectiveBuilderMixedInput< cute::is_same_v || cute::is_same_v) && (detail::is_use_rmem_A() || - // ConvertAndScale + // ConvertAndScale and ConvertAndScaleWithZero cute::is_tuple::value || cute::is_tuple::value || // DirectConvert sizeof_bits::value != sizeof_bits::value)>> { @@ -161,6 +106,7 @@ struct CollectiveBuilderMixedInput< static constexpr bool IsATransformed = cute::is_tuple::value; using ElementScale = cute::conditional_t; using ElementZero = cute::conditional_t; + static_assert(is_static::value); static_assert(is_static::value); static_assert( @@ -231,8 +177,8 @@ struct CollectiveBuilderMixedInput< cutlass::detail::alignment_for_swizzle(SmemLayoutAtomB{}); static constexpr int SmemAlignment = static_cast(cute::max(SmemAlignmentA, SmemAlignmentB)); - // Array mixed-input GEMM keeps only A/B TMA descriptors; folded scales are loaded by bulk copy. - static constexpr size_t TensorMapStorage = sizeof(cute::TmaDescriptor) * size_t(IsMixedInput) * 2; + // Handle mixed dtype array GEMM's size of tensor map storage. + static constexpr size_t TensorMapStorage = sizeof(cute::TmaDescriptor) * size_t(IsMixedInput) * 4; static constexpr int KernelSmemCarveout = static_cast(TensorMapStorage); static constexpr int Sm90ReducedSmemCapacityBytes = detail::sm90_smem_capacity_bytes - KernelSmemCarveout; @@ -240,33 +186,21 @@ struct CollectiveBuilderMixedInput< static constexpr int PipelineStages = IsMixedInput ? (IsArrayOfPointersGemm - ? detail::compute_stage_count_or_override_folded_weight_scale< + ? detail::compute_stage_count_or_override_single_affine_transformed_input< Sm90ReducedSmemCapacityBytes, RealElementA, RealElementB, ElementScale, ElementZero, TileShape_MNK, SmemAlignment>(StageCountType{}) - : detail::compute_stage_count_or_override_folded_weight_scale< + : detail::compute_stage_count_or_override_single_affine_transformed_input< detail::sm90_smem_capacity_bytes, RealElementA, RealElementB, ElementScale, ElementZero, TileShape_MNK, SmemAlignment>(StageCountType{})) : detail::compute_stage_count_or_override( StageCountType{}); - static constexpr bool UseFusedE8M0PreMmaScale = ScaleMode == MixedInputScaleMode::kPreMmaE8M0; - static_assert( - !UseFusedE8M0PreMmaScale || (IsArrayOfPointersGemm && IsATransformed && - cute::is_same_v && - cute::is_same_v), - "Pre-MMA E8M0 scale mode is only implemented for grouped MXFP4 weight x FP8 activation."); - - using ArrayMixedInputDispatchPolicy = - cute::conditional_t, - MainloopSm90ArrayTmaGmmaWarpSpecializedMixedInput< - PipelineStages, ClusterShape_MNK, KernelScheduleType>>; - using DispatchPolicy = cute::conditional_t< IsMixedInput, - cute::conditional_t, MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput< PipelineStages, ClusterShape_MNK, KernelScheduleType>>, MainloopSm90TmaGmmaRmemAWarpSpecialized + class KernelScheduleType, class Enable = void> struct CollectiveBuilderMixedInput { static_assert(sizeof(ElementA) == 0, "Could not build a collective for given parameters."); }; diff --git a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/collective/collective_mma_array_mixed_input.hpp b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/collective/collective_mma_array_mixed_input.hpp index a7c1023cd02..7ca25def0dd 100644 --- a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/collective/collective_mma_array_mixed_input.hpp +++ b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/collective/collective_mma_array_mixed_input.hpp @@ -16,31 +16,6 @@ #pragma once #include "cutlass/detail/dependent_false.hpp" -#include "cutlass/gemm/dispatch_policy.hpp" - -///////////////////////////////////////////////////////////////////////////////////////////////// - -namespace cutlass::gemm { - -///////////////////////////////////////////////////////////////////////////////////////////////// - -// Pre-MMA scale variant -template , - class KernelSchedule = KernelPtrArrayTmaWarpSpecializedCooperative> -struct MainloopSm90ArrayTmaGmmaWarpSpecializedMixedInputPreScale { - constexpr static int Stages = Stages_; - using ClusterShape = ClusterShape_; - using ArchTag = arch::Sm90; - using Schedule = KernelSchedule; - static_assert( - cute::is_same_v || - cute::is_same_v, - "KernelSchedule must be one of the Ptr-Array or Grouped GEMM TMA warp specialized policies"); -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace cutlass::gemm ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -48,13 +23,6 @@ namespace cutlass::gemm::collective { ///////////////////////////////////////////////////////////////////////////////////////////////// -enum class MixedInputScaleMode { - kPostMma = 0, - kPreMmaE8M0, -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - template , float, ElementScale>; using NonVoidElementZero = cute::conditional_t, float, ElementZero>; - static constexpr bool HasActivationScale = false; - using ElementActivationScale = void; - using NonVoidElementActivationScale = cutlass::float_ue8m0_t; using StrideA = StrideA_; using InternalStrideA = cute::remove_pointer_t; @@ -125,20 +122,22 @@ struct CollectiveMmaArrayMixedInput< "Scale must be MN major [Col Major if A is scaled, Row Major if B is scaled]."); static constexpr bool IsMXFP4 = cute::is_same_v; - static constexpr bool IsInt4Weight = cute::is_same_v; - static_assert(IsMXFP4 || IsInt4Weight, - "Folded weight scale is only defined for MXFP4 and INT4 weights."); - static constexpr int ScalingGroupSize = detail::DefaultWeightScaleGroupSize::value; + // Group size 128 for int4 weights + // Group size 32 for mxfp4 weights + static constexpr int ScalingGroupSize = + IsMXFP4 ? detail::mxfp4_group_size : detail::int4_group_size; using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{})); using TiledMma = TiledMma_; using ElementAccumulator = typename TiledMma::ValTypeC; using GmemTiledCopyA = GmemTiledCopyA_; using GmemTiledCopyB = GmemTiledCopyB_; + using GmemTiledCopyScale = cute::SM90_TMA_LOAD; using SmemLayoutAtomA = SmemLayoutAtomA_; using SmemLayoutAtomB = SmemLayoutAtomB_; using SmemCopyAtomA = SmemCopyAtomA_; using SmemCopyAtomB = SmemCopyAtomB_; + using SmemCopyAtomScale = Copy_Atom; // We must ensure the type to be scaled goes to RF static constexpr bool SwapAB = !IsATransformed; @@ -171,6 +170,9 @@ struct CollectiveMmaArrayMixedInput< static constexpr int IsSubbyteA = cute::sizeof_bits_v < 8; using TmaElementA = cute::conditional_t; + using TmaElementScale = + uint_bit_t>; // in case we have array. translating to uint + // to satisfy tma descriptor's specialization using MainloopPipeline = cutlass::PipelineTmaAsync; using PipelineState = cutlass::PipelineState; @@ -178,6 +180,11 @@ struct CollectiveMmaArrayMixedInput< static constexpr int NumProducerThreadEvents = 1; + using SmemLayoutAtomScale = + Layout(SwappedSmemLayoutAtomA{})), cute::Int<1>>>; + using ScaleTileShape = + decltype(make_shape(shape<0>(TileShape{}), shape<1>(SmemLayoutAtomScale{}))); + static_assert(cute::rank(SwappedSmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); static_assert((size<0>(TileShape{}) % size<0>(SwappedSmemLayoutAtomA{})) == 0, @@ -192,121 +199,24 @@ struct CollectiveMmaArrayMixedInput< static_assert((size<2>(TileShape{}) % size<1>(SwappedSmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert(rank(SmemLayoutAtomScale{}) == 2, "SmemLayoutAtomScale must be rank 2"); + static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomScale{})) == 0, + "SmemLayoutAtomScale must equal the tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomScale{})) == 0, + "SmemLayoutAtomScale must evenly divide tile k shape."); + /// Tile along modes in a way that maximizes the TMA box size. using SmemLayoutA = decltype(detail::get_smem_layout( SwappedSmemLayoutAtomA{}, select<0, 2>(TileShape{}), InternalSwappedStrideA{})); using SmemLayoutB = decltype(detail::get_smem_layout( SwappedSmemLayoutAtomB{}, select<1, 2>(TileShape{}), InternalSwappedStrideB{})); - using SmemLayoutAtomScale = - Layout(SwappedSmemLayoutAtomA{})), cute::Int<1>>>; - using ScaleTileShape = - decltype(make_shape(shape<0>(TileShape{}), shape<1>(SmemLayoutAtomScale{}))); + // It is assumed that the scales and zero-points share the same smem layout using SmemLayoutScale = decltype(tile_to_shape( SmemLayoutAtomScale{}, make_shape(shape<0>(ScaleTileShape{}), shape<1>(ScaleTileShape{}), Int{}), cute::conditional_t<::cutlass::gemm::detail::is_major<0, NonVoidStrideScale>(), Step<_2, _1, _3>, Step<_1, _2, _3>>{})); - using SmemLayoutActivationScale = void; - using SmemCopyAtomScale = Copy_Atom; - - using WeightScaleRawElement = NonVoidElementScale; - static_assert( - !cutlass::detail::is_Array_v, - "Folded weight scale uses scalar scale storage; packed scale arrays are not supported."); - static_assert(cute::is_void_v, - "Folded weight scale storage does not support zero-point."); - static constexpr int WeightScaleLogicalMPerFoldBlock = 64; - static constexpr int WeightScaleLogicalKPerFoldBlock = 128; - static constexpr int WeightScalePhysicalColsPerFoldBlock = - 128 / cutlass::sizeof_bits::value; - static constexpr int WeightScaleScaleGroupsPerFoldBlock = - WeightScaleLogicalKPerFoldBlock / ScalingGroupSize; - static constexpr int WeightScaleMSlicesPerFoldBlock = - WeightScalePhysicalColsPerFoldBlock / WeightScaleScaleGroupsPerFoldBlock; - static constexpr int WeightScaleFoldedMPerFoldBlock = - WeightScaleLogicalMPerFoldBlock / WeightScaleMSlicesPerFoldBlock; - static constexpr int WeightScaleMBlocksPerTile = - size<0>(TileShape{}) / WeightScaleLogicalMPerFoldBlock; - static constexpr int WeightScaleKBlocksPerTile = - size<2>(TileShape{}) / WeightScaleLogicalKPerFoldBlock; - static_assert(size<0>(TileShape{}) % WeightScaleLogicalMPerFoldBlock == 0, - "Folded weight scale requires TileShapeM to be a multiple of 64."); - static_assert(size<2>(TileShape{}) % WeightScaleLogicalKPerFoldBlock == 0, - "Folded weight scale requires TileShapeK to be a multiple of 128."); - static_assert(WeightScalePhysicalColsPerFoldBlock % WeightScaleScaleGroupsPerFoldBlock == 0, - "Folded weight scale requires each M slice to contain an integer number of scale " - "groups."); - static_assert(WeightScaleLogicalMPerFoldBlock % WeightScaleMSlicesPerFoldBlock == 0, - "Folded weight scale M slices must evenly divide the logical M block."); - static_assert(WeightScalePhysicalColsPerFoldBlock * - cutlass::sizeof_bits::value == - 128, - "Folded weight scale must expose 16B per folded-M coordinate."); - static constexpr int WeightScaleRawElementsPerFoldBlock = - WeightScaleLogicalMPerFoldBlock * WeightScaleLogicalKPerFoldBlock / ScalingGroupSize; - static constexpr int WeightScaleRawElementsPerStage = - WeightScaleRawElementsPerFoldBlock * WeightScaleMBlocksPerTile * WeightScaleKBlocksPerTile; - static constexpr uint32_t WeightScaleFoldBlockBytes = cutlass::bits_to_bytes( - WeightScaleRawElementsPerFoldBlock * cutlass::sizeof_bits::value); - static constexpr uint32_t WeightScaleBulkCopyBytes = - WeightScaleFoldBlockBytes * WeightScaleKBlocksPerTile; - static constexpr uint32_t WeightScaleTransactionBytes = cutlass::bits_to_bytes( - WeightScaleRawElementsPerStage * cutlass::sizeof_bits::value); - static_assert(WeightScaleBulkCopyBytes % 16 == 0, - "Folded weight-scale bulk copy size must be 16B aligned."); - - static constexpr bool IsInt4Fp8Path = cute::is_same_v && - cute::is_same_v; - static constexpr bool IsMxfp4Bf16Path = IsMXFP4 && cute::is_same_v; - static constexpr bool IsMxfp4Fp8Path = - IsMXFP4 && cute::is_same_v; - - static constexpr bool UseDirectSmemWeightScale = - ((IsMxfp4Fp8Path && ((size<0>(TileShape{}) == 64 && size<1>(TileShape{}) == 64 && - size<0>(ClusterShape{}) == 1 && size<1>(ClusterShape{}) == 1) || - (size<0>(TileShape{}) == 64 && size<1>(TileShape{}) == 128 && - size<2>(TileShape{}) == 256 && size<0>(ClusterShape{}) == 1 && - size<1>(ClusterShape{}) == 1) || - (size<0>(TileShape{}) == 128 && size<1>(TileShape{}) == 256 && - size<2>(TileShape{}) == 256) || - (size<0>(TileShape{}) == 256 && size<1>(TileShape{}) == 128 && - size<2>(TileShape{}) == 256) || - (size<0>(TileShape{}) == 128 && size<1>(TileShape{}) == 64 && - size<2>(TileShape{}) == 512 && - (size<0>(ClusterShape{}) != 1 || size<1>(ClusterShape{}) != 1)))) || - (IsMxfp4Bf16Path && ((size<0>(TileShape{}) == 64 && size<1>(TileShape{}) == 64 && - size<0>(ClusterShape{}) == 1 && size<1>(ClusterShape{}) == 1) || - (size<0>(TileShape{}) == 64 && size<1>(TileShape{}) == 128 && - size<2>(TileShape{}) == 256 && size<0>(ClusterShape{}) == 1 && - size<1>(ClusterShape{}) == 1))) || - (IsInt4Fp8Path && - ((size<0>(TileShape{}) == 64 && size<1>(TileShape{}) == 64 && size<2>(TileShape{}) == 256 && - size<0>(ClusterShape{}) == 1 && size<1>(ClusterShape{}) == 1) || - (size<0>(TileShape{}) == 128 && size<1>(TileShape{}) == 64 && - size<2>(TileShape{}) == 512 && size<0>(ClusterShape{}) == 1 && - size<1>(ClusterShape{}) == 2)))); - using SmemLayoutWeightScaleRaw = - Layout, Int, - Int, Int, Int>, - Stride<_1, Int, - Int, - Int, - Int>>; - using SmemLayoutWeightScaleExpanded = Layout< - Shape, Int, - Int>, - Shape, - Shape, Int>>, - Int>, - Stride< - Stride, Int, - Int>, - Stride<_0, Stride<_1, Int>>, - Int>>; static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 2 or more."); @@ -321,21 +231,28 @@ struct CollectiveMmaArrayMixedInput< cute::is_same_v, "GmemTiledCopy - invalid SM90 TMA copy atom specified."); + // To relax them, we need to handle loading more than 1 row of scales for every main loop + // iteration. We must also handle updating the pipeline transaction bytes on the fly. + static_assert(size<1>(SmemLayoutAtomScale{}) == 1, "size<1>(SmemLayoutAtomScale) must be 1."); + private: static constexpr ConversionMode get_conversion_mode() { if constexpr (cute::is_void_v) { return ConversionMode::DirectConvert; - } else { + } else if constexpr (cute::is_void_v) { return ConversionMode::ConvertAndScale; + } else { + return ConversionMode::ConvertAndScaleWithZero; } } - int current_group_idx_ = 0; - cute::TmaDescriptor const* current_tma_desc_b_ = nullptr; + bool TensormapUpdateShapesStridesForAandScale = true; public: static constexpr ConversionMode KernelConversionMode = get_conversion_mode(); - static constexpr bool ModeHasScales = KernelConversionMode == ConversionMode::ConvertAndScale; + static constexpr bool ModeHasScales = + KernelConversionMode == ConversionMode::ConvertAndScale || + KernelConversionMode == ConversionMode::ConvertAndScaleWithZero; static constexpr bool UseScaleLookupTable = KernelConversionMode == ConversionMode::ConvertAndScale && cutlass::detail::is_Array_v; @@ -343,13 +260,9 @@ struct CollectiveMmaArrayMixedInput< KernelConversionMode == ConversionMode::ConvertAndScale && cute::is_same_v && cute::is_same_v; - static constexpr bool UseFP4ToFP8LookupTable = - KernelConversionMode == ConversionMode::ConvertAndScale && - cute::is_same_v && - cute::is_same_v; static constexpr bool UseInt4ToFP8LookupTable = KernelConversionMode == ConversionMode::ConvertAndScale && - cute::is_same_v && + cute::is_same_v && cute::is_same_v; static constexpr size_t SmemAlignmentA = cutlass::detail::alignment_for_swizzle(SmemLayoutA{}); static constexpr size_t SmemAlignmentB = cutlass::detail::alignment_for_swizzle(SmemLayoutB{}); @@ -358,17 +271,24 @@ struct CollectiveMmaArrayMixedInput< static_assert(SmemAlignmentA >= 128 and SmemAlignmentB >= 128, "Require at least 128B alignment"); struct SharedStorage { - static constexpr int scale_elements = cute::cosize_v; + static constexpr int scale_elements = Utils::elements_per_smem_scale(); + static constexpr int zero_elements = Utils::elements_per_smem_zero(); struct TensorStorage { CUTE_ALIGNAS(SmemAlignmentA) cute::ArrayEngine> smem_A; CUTE_ALIGNAS(SmemAlignmentB) cute::ArrayEngine> smem_B; - cute::ArrayEngine smem_scale; + cute::ArrayEngine smem_scale; + cute::ArrayEngine smem_zero; } tensors; - struct TensorMapStorage {}; + struct TensorMapStorage { + cute::TmaDescriptor smem_tensormap_A; + cute::TmaDescriptor smem_tensormap_B; + cute::TmaDescriptor smem_tensormap_scale; + cute::TmaDescriptor smem_tensormap_zero; + }; using PipelineStorage = typename MainloopPipeline::SharedStorage; PipelineStorage pipeline; @@ -379,8 +299,6 @@ struct CollectiveMmaArrayMixedInput< using PipelineStorage = typename SharedStorage::PipelineStorage; static constexpr bool IsGroupedGemmKernel = !cute::is_same_v; - static constexpr bool RequiresTensormapUpdateOnBatchChange = false; - static constexpr bool RequiresPrebuiltTensormapAcquireOnBatchChange = IsGroupedGemmKernel; // kernel Arguments // Host side kernel arguments @@ -392,24 +310,14 @@ struct CollectiveMmaArrayMixedInput< ElementScale const** ptr_S = nullptr; NonVoidStrideScale const* dS{}; int chunk_size = 0; - cute::TmaDescriptor const* ptr_A_prebuilt_tma_desc = nullptr; - cute::TmaDescriptor const* ptr_B_prebuilt_tma_descs = nullptr; + ElementZero const** ptr_Z = nullptr; }; // Device side kernel params struct Params { - // For grouped GEMM with non-layout stride: replace static-zero L stride with - // a static non-zero value so the A/weight TMA descriptor is created as 3D - // and can select the expert through the L coordinate. Int<32> is the - // minimum static value that becomes 16 bytes after FP4 subbyte upcast. - using TmaStrideA = - cute::conditional_t::value, - decltype(cute::make_stride(cute::get<0>(InternalSwappedStrideA{}), - cute::get<1>(InternalSwappedStrideA{}), - cute::Int<32>{})), - InternalSwappedStrideA>; - using LayoutA = - decltype(detail::get_gmem_layout(repeat_like(TmaStrideA{}, int32_t(0)), TmaStrideA{})); + // Assumption: StrideA is congruent with Problem_MK + using LayoutA = decltype(detail::get_gmem_layout( + repeat_like(InternalSwappedStrideA{}, int32_t(0)), InternalSwappedStrideA{})); using LayoutB = decltype(detail::get_gmem_layout( repeat_like(InternalSwappedStrideB{}, int32_t(0)), InternalSwappedStrideB{})); @@ -429,21 +337,40 @@ struct CollectiveMmaArrayMixedInput< make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})), size<0>(ClusterShape{}))); // mcast along M mode for this N load, if any + using TMA_Scale = decltype(make_tma_copy( + GmemTiledCopyScale{}, + make_tensor(detail::get_logical_ptr(static_cast(nullptr)), + repeat_like(NonVoidStrideScale{}, int32_t(0)), NonVoidStrideScale{}), + SmemLayoutScale{}(_, _, cute::Int<0>{}), ScaleTileShape{}, + _1{})); // mcast along N mode for this M load, if any. Scale is ALWAYS loaded with A for RF + // kernel + + using TMA_Zero = decltype(make_tma_copy( + GmemTiledCopyScale{}, + make_tensor(detail::get_logical_ptr(static_cast(nullptr)), + repeat_like(NonVoidStrideScale{}, int32_t(0)), NonVoidStrideScale{}), + SmemLayoutScale{}(_, _, cute::Int<0>{}), ScaleTileShape{}, + _1{})); // mcast along N mode for this M load, if any. Scale is ALWAYS loaded with A for RF + // kernel + TMA_A tma_load_a; TMA_B tma_load_b; uint32_t tma_transaction_bytes = TmaTransactionBytes; - cute::TmaDescriptor const* ptr_A_prebuilt_tma_desc; - cute::TmaDescriptor const* ptr_B_prebuilt_tma_descs; + TMA_Scale tma_load_scale; + TMA_Zero tma_load_zero; + void* tensormaps; SwappedElementA const** ptr_A; SwappedStrideA ptr_dA; SwappedElementB const** ptr_B; SwappedStrideB ptr_dB; NonVoidElementScale const** ptr_S; NonVoidStrideScale const* dS; + NonVoidElementZero const** ptr_Z; + int64_t scale_k; int chunk_size; + int reload_factor = (chunk_size + size<2>(TileShape{}) - 1) / size<2>(TileShape{}); InternalSwappedStrideA dA; InternalSwappedStrideB dB; - int num_groups; }; // @@ -520,14 +447,8 @@ struct CollectiveMmaArrayMixedInput< ptr_dA = SwappedStrideA{}; ptr_dB = SwappedStrideB{}; } - // Grouped A/weight uses TmaStrideA to force a 3D descriptor. The descriptor - // is later rebuilt once with the real expert count and reused for all groups. - typename Params::TmaStrideA tma_dA; - if constexpr (!IsGroupedGemmKernel || cute::is_layout::value) { - tma_dA = dA; - } - Tensor tensor_a = make_tensor( - ptr_A_first_batch, detail::get_gmem_layout(make_shape(init_M, init_K, mock_L), tma_dA)); + Tensor tensor_a = make_tensor(ptr_A_first_batch, + detail::get_gmem_layout(make_shape(init_M, init_K, mock_L), dA)); Tensor tensor_b = make_tensor(ptr_B_first_batch, detail::get_gmem_layout(make_shape(init_N, init_K, mock_L), dB)); @@ -539,33 +460,68 @@ struct CollectiveMmaArrayMixedInput< make_tma_copy(GmemTiledCopyB{}, tensor_b, SmemLayoutB{}(_, _, cute::Int<0>{}), make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})), size<0>(ClusterShape{})); // mcast along M mode for this N load, if any - int num_groups_val = 1; - if constexpr (IsGroupedGemmKernel) { - num_groups_val = problem_shapes.groups(); - } - auto args_setup = [&](auto ptr_A, auto ptr_B, int chunk_size = 0) -> Params { + typename Params::TMA_Scale tma_load_scale{}; + typename Params::TMA_Zero tma_load_zero{}; + + void* tensormaps = workspace; + auto args_setup = [&](auto ptr_A, auto ptr_B, int64_t scale_k = 0, int chunk_size = 0, + int reload_factor = 1) -> Params { return {tma_load_a, tma_load_b, TmaTransactionBytes, - args.ptr_A_prebuilt_tma_desc, - args.ptr_B_prebuilt_tma_descs, + tma_load_scale, + tma_load_zero, + tensormaps, reinterpret_cast(ptr_A), ptr_dA, reinterpret_cast(ptr_B), ptr_dB, reinterpret_cast(args.ptr_S), args.dS, + reinterpret_cast(args.ptr_Z), + scale_k, chunk_size, + reload_factor, dA, - dB, - num_groups_val}; + dB}; }; if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { return SwapAB ? args_setup(args.ptr_B, args.ptr_A) : args_setup(args.ptr_A, args.ptr_B); } else if constexpr (ModeHasScales) { - return SwapAB ? args_setup(args.ptr_B, args.ptr_A, args.chunk_size) - : args_setup(args.ptr_A, args.ptr_B, args.chunk_size); + auto fake_scale_k = 1; + ElementScale const* ptr_S = reinterpret_cast(args.ptr_S); + StrideScale dS{}; + Tensor tensor_scale = make_tensor(detail::get_logical_ptr(ptr_S), + make_layout(make_shape(init_M, fake_scale_k, mock_L), dS)); + tma_load_scale = make_tma_copy( + GmemTiledCopyScale{}, tensor_scale, SmemLayoutScale{}(_, _, cute::Int<0>{}), + ScaleTileShape{}, _1{}); // mcast along N mode for this M load, if any + + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + return SwapAB + ? args_setup(args.ptr_B, args.ptr_A, fake_scale_k, args.chunk_size, + (args.chunk_size + size<2>(TileShape{}) - 1) / size<2>(TileShape{})) + : args_setup( + args.ptr_A, args.ptr_B, fake_scale_k, args.chunk_size, + (args.chunk_size + size<2>(TileShape{}) - 1) / size<2>(TileShape{})); + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + ElementZero const* ptr_Z = reinterpret_cast(args.ptr_Z); + Tensor tensor_zero = make_tensor(detail::get_logical_ptr(ptr_Z), + make_layout(make_shape(init_M, fake_scale_k, mock_L), dS)); + tma_load_zero = make_tma_copy(GmemTiledCopyScale{}, tensor_zero, + SmemLayoutScale{}(_, _, cute::Int<0>{}), ScaleTileShape{}, + _1{}); // mcast along N mode for this M load, if any + return SwapAB + ? args_setup(args.ptr_B, args.ptr_A, fake_scale_k, args.chunk_size, + (args.chunk_size + size<2>(TileShape{}) - 1) / size<2>(TileShape{})) + : args_setup( + args.ptr_A, args.ptr_B, fake_scale_k, args.chunk_size, + (args.chunk_size + size<2>(TileShape{}) - 1) / size<2>(TileShape{})); + } else { + static_assert(cutlass::detail::dependent_false, + "Conversion mode not handled in to_underlying_arguments."); + } } else { static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in to_underlying_arguments."); @@ -575,10 +531,29 @@ struct CollectiveMmaArrayMixedInput< template static size_t get_workspace_size(ProblemShape const& problem_shape, Arguments const& args, int sm_count) { - (void)problem_shape; - (void)args; - (void)sm_count; - return 0; + constexpr size_t SizeOfCuTensorMap = sizeof(cute::TmaDescriptor); + + // Calculating workspace size + auto calculate_workspace_size = [SizeOfCuTensorMap, sm_count](uint32_t num_input_tensors) { + return num_input_tensors * SizeOfCuTensorMap * sm_count; + }; + + if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { + // Allocate gmem space for input tensormaps per each SM, A tensormap copies followed by B + // tensormap copies + return calculate_workspace_size(2); + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + // Allocate gmem space for input tensormaps per each SM, A tensormap copies followed by B + // tensormap copies, followed by scale tensormap copies + return calculate_workspace_size(3); + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + // Allocate gmem space for input tensormaps per each SM, A tensormap copies followed by B + // tensormap copies, followed by scale and zeros tensormap copies + return calculate_workspace_size(4); + } else { + static_assert(cutlass::detail::dependent_false, + "Conversion mode not handled in get_workspace_size."); + } } template @@ -621,14 +596,32 @@ struct CollectiveMmaArrayMixedInput< detail::get_gmem_layout(cute::make_shape(N, K, L), dB)); if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { implementable = implementable && (args.ptr_S == nullptr); + implementable = implementable && (args.ptr_Z == nullptr); } else if constexpr (ModeHasScales) { int const scale_mn = SwapAB ? N : M; + int const scale_k = (K + args.chunk_size - 1) / args.chunk_size; + constexpr int min_tma_aligned_elements_scale = + tma_alignment_bits / cutlass::sizeof_bits::value; + implementable = + implementable && cutlass::detail::check_alignment( + cute::make_shape(scale_mn, scale_k, L), StrideScale{}); + implementable = implementable && + (args.chunk_size == K || ((args.chunk_size % size<2>(TileShape{})) == 0)); implementable = implementable && args.chunk_size != 0; implementable = implementable && (args.ptr_S != nullptr); - implementable = implementable && (args.chunk_size == ScalingGroupSize); - implementable = implementable && ((scale_mn % size<0>(TileShape{})) == 0); - implementable = implementable && ((K % size<2>(TileShape{})) == 0); - implementable = implementable && ((K % WeightScaleLogicalKPerFoldBlock) == 0); + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + implementable = implementable && (args.ptr_Z == nullptr); + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + constexpr int min_tma_aligned_elements_zero = + tma_alignment_bits / cutlass::sizeof_bits::value; + implementable = + implementable && cutlass::detail::check_alignment( + cute::make_shape(scale_mn, scale_k, L), StrideScale{}); + implementable = implementable && (args.ptr_Z != nullptr); + } else { + static_assert(cutlass::detail::dependent_false, + "Conversion mode not handled in can_implement."); + } } else { static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in can_implement."); @@ -641,10 +634,6 @@ struct CollectiveMmaArrayMixedInput< " CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for " "TMA.\n"); } - if constexpr (IsGroupedGemmKernel) { - implementable = implementable && args.ptr_A_prebuilt_tma_desc != nullptr; - implementable = implementable && args.ptr_B_prebuilt_tma_descs != nullptr; - } return implementable; } @@ -672,25 +661,40 @@ struct CollectiveMmaArrayMixedInput< // TMA requires special handling of strides to deal with coord codomain mapping // Represent the full tensors -- get these from TMA - auto A_L = mainloop_params.num_groups; - auto B_L = mock_L; Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor( - shape(detail::get_gmem_layout(make_shape(M, K, A_L), mainloop_params.dA))); // (m,k,l) + shape(detail::get_gmem_layout(make_shape(M, K, mock_L), mainloop_params.dA))); // (m,k,l) Tensor mB_nkl = mainloop_params.tma_load_b.get_tma_tensor( - shape(detail::get_gmem_layout(make_shape(N, K, B_L), mainloop_params.dB))); // (n,k,l) + shape(detail::get_gmem_layout(make_shape(N, K, mock_L), mainloop_params.dB))); // (n,k,l) // Make tiled views, defer the slice Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_, _, _), Step<_1, X, _1>{}); // (BLK_M,BLK_K,m,k,l) Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,n,k,l) - int const scale_total_k128_blocks = int(K) / WeightScaleLogicalKPerFoldBlock; if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { return cute::make_tuple(gA_mkl, gB_nkl); } else if constexpr (ModeHasScales) { - return cute::make_tuple(gA_mkl, gB_nkl, static_cast(nullptr), - int64_t(0), scale_total_k128_blocks); + // The real scale_k that actually works + // auto scale_k = K / mainloop_params.chunk_size; + auto scale_k = K / ScalingGroupSize; + + Tensor mS_mkl = mainloop_params.tma_load_scale.get_tma_tensor( + make_shape(M, scale_k, L)); // (m,scale_k,l) + Tensor gS_mkl = local_tile(mS_mkl, ScaleTileShape{}, + make_coord(_, _)); // (BLK_M,BLK_Scale_K,m,scale_k,l) + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + return cute::make_tuple(gA_mkl, gB_nkl, gS_mkl); + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + Tensor mZ_mkl = mainloop_params.tma_load_zero.get_tma_tensor( + make_shape(M, scale_k, L)); // (m,scale_k,l) + Tensor gZ_mkl = local_tile(mZ_mkl, ScaleTileShape{}, + make_coord(_, _)); // (BLK_M,BLK_Scale_K,m,scale_k,l) + return cute::make_tuple(gA_mkl, gB_nkl, gS_mkl, gZ_mkl); + } else { + static_assert(cutlass::detail::dependent_false, + "Conversion mode not handled in load_init."); + } } else { static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in load_init."); @@ -710,8 +714,11 @@ struct CollectiveMmaArrayMixedInput< static_assert(sizeof...(Ts) == 2, "Direct convert needs two inputs"); static_assert(sizeof...(TMs) == 2, "Direct convert needs two tensormaps"); } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { - static_assert(sizeof...(Ts) == 5, "Folded scaled convert needs five inputs"); - static_assert(sizeof...(TMs) == 2, "Folded scaled convert needs two tensormaps"); + static_assert(sizeof...(Ts) == 3, "Scaled convert needs three inputs"); + static_assert(sizeof...(TMs) == 3, "Scaled convert needs three tensormaps"); + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + static_assert(sizeof...(Ts) == 4, "Scaled and zero convert needs four inputs"); + static_assert(sizeof...(TMs) == 4, "Scaled and zero convert needs four tensormaps"); } else { static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in TMA load."); @@ -740,10 +747,8 @@ struct CollectiveMmaArrayMixedInput< // Partition the inputs based on the current block coordinates. auto [m_coord, n_coord, k_coord, l_coord] = blk_coord; - auto a_l_coord = current_group_idx_; - auto b_l_coord = cute::Int<0>{}; - Tensor gA = gA_mkl(_, _, m_coord, _, a_l_coord); // (BLK_M,BLK_K,k) - Tensor gB = gB_nkl(_, _, n_coord, _, b_l_coord); // (BLK_N,BLK_K,k) + Tensor gA = gA_mkl(_, _, m_coord, _, l_coord); // (BLK_M,BLK_K,k) + Tensor gB = gB_nkl(_, _, n_coord, _, l_coord); // (BLK_N,BLK_K,k) // Applies the mapping from block_tma_a Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k) @@ -754,6 +759,7 @@ struct CollectiveMmaArrayMixedInput< uint16_t mcast_mask_a = 0; uint16_t mcast_mask_b = 0; + uint16_t mcast_mask_s = 0; // Issue TmaLoads // Maps the tile -> block, value @@ -771,6 +777,9 @@ struct CollectiveMmaArrayMixedInput< } } + auto extra_input_partitions = Utils::partition_extra_tma_inputs( + mainloop_params, load_inputs, shared_tensors, cluster_local_block_id, m_coord, l_coord); + // Mainloop CUTLASS_PRAGMA_NO_UNROLL for (; k_tile_count > 0; --k_tile_count) { @@ -786,41 +795,44 @@ struct CollectiveMmaArrayMixedInput< int write_stage = smem_pipe_write.index(); if (cute::elect_one_sync()) { - copy(mainloop_params.tma_load_a.with(mainloop_params.ptr_A_prebuilt_tma_desc, *tma_barrier, - mcast_mask_a), + copy(mainloop_params.tma_load_a.with(get<0>(input_tensormaps), *tma_barrier, mcast_mask_a), tAgA(_, _, _, *k_tile_iter), tAsA(_, _, _, write_stage)); - copy(mainloop_params.tma_load_b.with(current_tma_desc_b_, *tma_barrier, mcast_mask_b), + copy(mainloop_params.tma_load_b.with(get<1>(input_tensormaps), *tma_barrier, mcast_mask_b), tBgB(_, _, _, *k_tile_iter), tBsB(_, _, _, write_stage)); } if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { // Nothing extra to do. } else if constexpr (ModeHasScales) { + // scale copy + auto tSgS = get<0>(extra_input_partitions); + auto tSsS = get<1>(extra_input_partitions); + + // Temporary factor which will determine which k tile to reload from gmem. Needed so we + // don't modify tma transaction bytes on the fly. We must do a ceiling divide here to + // correctly handle with chunk_size == K. In that case, we don't require that K is a + // multiple of the threadblock tile K + int const scale_load_k = *k_tile_iter / 1; + // const int scale_load_k = *k_tile_iter / mainloop_params.reload_factor; // This will + // always be 0 when chunk_size == K. if (cute::elect_one_sync()) { - auto scale_ptr = get<2>(load_inputs); - int const scale_k_tile = *k_tile_iter; - int const scale_total_k128_blocks = get<4>(load_inputs); - int const scale_k128_offset = scale_k_tile * WeightScaleKBlocksPerTile; - int const scale_m64_offset = - int(m_coord) * int(size<0>(TileShape{})) / WeightScaleLogicalMPerFoldBlock; - auto* scale_base = reinterpret_cast(scale_ptr); - Tensor sSRaw = make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), - SmemLayoutWeightScaleRaw{}); - - CUTLASS_PRAGMA_UNROLL - for (int local_m64_block = 0; local_m64_block < WeightScaleMBlocksPerTile; - ++local_m64_block) { - int const m64_block = scale_m64_offset + local_m64_block; - int64_t const scale_gmem_fold_block = - int64_t(m64_block) * int64_t(scale_total_k128_blocks) + int64_t(scale_k128_offset); - int64_t const scale_gmem_offset = - scale_gmem_fold_block * int64_t(WeightScaleRawElementsPerFoldBlock); - auto* scale_gmem_addr = reinterpret_cast(scale_base + scale_gmem_offset); - auto* scale_smem_addr = - static_cast(&sSRaw(0, 0, local_m64_block, 0, write_stage)); - cute::SM90_BULK_COPY_G2S::copy(scale_gmem_addr, - reinterpret_cast(tma_barrier), - scale_smem_addr, WeightScaleBulkCopyBytes); + copy(mainloop_params.tma_load_scale.with(get<2>(input_tensormaps), *tma_barrier, + mcast_mask_s), + tSgS(_, _, _, scale_load_k), tSsS(_, _, _, write_stage)); + } + + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + // Nothing extra to do + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + auto tZgZ = get<2>(extra_input_partitions); + auto tZsZ = get<3>(extra_input_partitions); + if (cute::elect_one_sync()) { + copy(mainloop_params.tma_load_zero.with(get<3>(input_tensormaps), *tma_barrier, + mcast_mask_s), + tZgZ(_, _, _, scale_load_k), tZsZ(_, _, _, write_stage)); } + } else { + static_assert(cutlass::detail::dependent_false, + "Conversion mode not handled for TMA copy op."); } } else { static_assert(cutlass::detail::dependent_false, @@ -855,79 +867,14 @@ struct CollectiveMmaArrayMixedInput< if constexpr (cute::is_same_v) { cutlass::float_ue8m0_t scale_ue8m0 = scale; - uint32_t temp = static_cast(scale_ue8m0.storage) << 23; - return cutlass::detail::copy_bits(temp); + uint32_t temp = 0; + temp = (temp | *reinterpret_cast(&scale_ue8m0)) << 23; + return *reinterpret_cast(&temp); } else { return static_cast(scale); } } - template - CUTLASS_DEVICE float load_weight_scale(ScaleTensor const& tCrS, int scale_idx, int mma_m, int m) { - return scale_convertor(tCrS(make_coord(make_tuple(0, m, 0), mma_m, scale_idx))); - } - - template - CUTLASS_DEVICE float load_weight_scale_smem(ScaleTensor const& tCsS, int scale_idx, int mma_m, - int m, int read_stage) { - return scale_convertor(tCsS(make_coord(make_tuple(0, m, 0), mma_m, scale_idx, read_stage))); - } - - template - CUTLASS_DEVICE void apply_groupwise_scale(AccumTensor& accum, IntermTensor const& intermediate, - ScaleTensor const& tCrS, int scale_idx, - bool is_first_accum) { - multiply_add fma_op; - - CUTLASS_PRAGMA_UNROLL - for (int mma_m = 0; mma_m < size<1>(accum); mma_m++) { - CUTLASS_PRAGMA_UNROLL - for (int m = 0; m < size<0, 1>(accum); m++) { - float scale_val = load_weight_scale(tCrS, scale_idx, mma_m, m); - CUTLASS_PRAGMA_UNROLL - for (int n = 0; n < size<0, 2>(accum); n++) { - CUTLASS_PRAGMA_UNROLL - for (int e = 0; e < size<0, 0>(accum); e++) { - auto coord = make_coord(make_tuple(e, m, n), mma_m, 0); - if (is_first_accum) { - accum(coord) = intermediate(coord) * scale_val; - } else { - accum(coord) = fma_op(intermediate(coord), scale_val, accum(coord)); - } - } - } - } - } - } - - template - CUTLASS_DEVICE void apply_groupwise_scale_smem(AccumTensor& accum, - IntermTensor const& intermediate, - ScaleTensor const& tCsS, int scale_idx, - int read_stage, bool is_first_accum) { - multiply_add fma_op; - - CUTLASS_PRAGMA_UNROLL - for (int mma_m = 0; mma_m < size<1>(accum); mma_m++) { - CUTLASS_PRAGMA_UNROLL - for (int m = 0; m < size<0, 1>(accum); m++) { - float scale_val = load_weight_scale_smem(tCsS, scale_idx, mma_m, m, read_stage); - CUTLASS_PRAGMA_UNROLL - for (int n = 0; n < size<0, 2>(accum); n++) { - CUTLASS_PRAGMA_UNROLL - for (int e = 0; e < size<0, 0>(accum); e++) { - auto coord = make_coord(make_tuple(e, m, n), mma_m, 0); - if (is_first_accum) { - accum(coord) = intermediate(coord) * scale_val; - } else { - accum(coord) = fma_op(intermediate(coord), scale_val, accum(coord)); - } - } - } - } - } - } - /// Perform a collective-scoped matrix multiply-accumulate /// Consumer Perspective template @@ -1005,21 +952,10 @@ struct CollectiveMmaArrayMixedInput< Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA_load); // (CPY,CPY_M,CPY_K) - constexpr int MmaKPerKBlock = cute::get<0, 1>(tCsB.shape())(); - constexpr int NumMMAsPerChunk = ScalingGroupSize / MmaKPerKBlock; - constexpr int KBlockMaxForScale = size<2>(TileShape{}) / MmaKPerKBlock; - constexpr int NumChunksPerTileK = cute::size<1>(sA.shape())() / ScalingGroupSize; - - Tensor sS = make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), - SmemLayoutWeightScaleExpanded{}); - Tensor tCsS = mma_thread_slice.partition_A(sS); - Tensor tCrS = make_tensor( - mma_thread_slice.partition_fragment_A(sS(_, _, Int<0>{})).layout()); - - using SmemCopyAtomScaleRaw = Copy_Atom; - auto smem_tiled_copy_S = make_tiled_copy_A(SmemCopyAtomScaleRaw{}, tiled_mma); - auto smem_thr_copy_S = smem_tiled_copy_S.get_thread_slice(warp_group_thread_idx); - Tensor tCrS_copy_view = smem_thr_copy_S.retile_D(tCrS); + // Partition of thread -> shared and thread -> RF + auto partitioned_extra_info = Utils::partition_extra_mma_info(mma_thread_slice, shared_tensors); + auto copy_partitions_extra_info = + Utils::retile_extra_mma_info(tiled_mma, partitioned_extra_info, warp_group_thread_idx); CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // CPY_M CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCrA_copy_view)); // CPY_K @@ -1049,13 +985,9 @@ struct CollectiveMmaArrayMixedInput< auto ptr = recast_ptr(tCrA_load_LDSM.data()); auto old_shape = tCrA_load_LDSM.shape(); - auto tCrA_load_4b_layout = make_layout( - make_shape(size<0>(old_shape), get<1>(old_shape), - make_shape(ABBitWidthRatio{}, size<2>(old_shape))), - make_stride(Int<1>{}, size<0>(old_shape) * ABBitWidthRatio{}, - make_stride(size<0>(old_shape), - size<0>(old_shape) * ABBitWidthRatio{} * size<1>(old_shape)))); - Tensor tCrA_load_4b_packed = make_tensor(ptr, tCrA_load_4b_layout); + auto new_shape = + make_shape(size<0>(old_shape), get<1>(old_shape), size<2>(old_shape) * ABBitWidthRatio{}); + Tensor tCrA_load_4b_packed = make_tensor(ptr, make_layout(new_shape)); ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// @@ -1066,46 +998,16 @@ struct CollectiveMmaArrayMixedInput< // We release buffers to producer warps(dma load) with some mmas in flight PipelineState smem_pipe_release = smem_pipe_read; + multiply_add fma; + + constexpr int NumMMAsPerChunk = ScalingGroupSize / cute::get<0, 1>(tCsB.shape())(); + constexpr int NumChunksPerTileK = cute::size<1>(sA.shape())() / ScalingGroupSize; cute::array intermediate_array; constexpr int K_BLOCK_MAX = size<2>(tCrA_load); - static_assert(K_BLOCK_MAX == KBlockMaxForScale, - "Scale and A operand K-block counts must match."); constexpr int K_WAIT_MAX = cute::min(K_BLOCK_MAX - 1, 7); static_assert(K_BLOCK_MAX >= 4, "Consider increasing TileShapeK"); - auto copy_scale_for_tile = [&](int read_stage) { - if constexpr (UseDirectSmemWeightScale) { - (void)read_stage; - } else { - if constexpr (NumMMAsPerChunk == 1) { - cute::for_each(cute::make_seq{}, [&](auto k_block_c) { - constexpr int k_block = decltype(k_block_c)::value; - cute::copy(smem_tiled_copy_S, tCsS(_, _, Int{}, read_stage), - tCrS_copy_view(_, _, Int{})); - }); - } else { - cute::for_each(cute::make_seq{}, [&](auto scale_idx_c) { - constexpr int scale_idx = decltype(scale_idx_c)::value; - constexpr int k_block = scale_idx * NumMMAsPerChunk; - cute::copy(smem_tiled_copy_S, tCsS(_, _, Int{}, read_stage), - tCrS_copy_view(_, _, Int{})); - }); - } - } - }; - - auto scale_intermediate = [&](auto const& intermediate, int scale_idx, int read_stage, - bool is_first_accum) { - int const weight_scale_idx = scale_idx * NumMMAsPerChunk; - if constexpr (UseDirectSmemWeightScale) { - apply_groupwise_scale_smem(accum, intermediate, tCsS, weight_scale_idx, read_stage, - is_first_accum); - } else { - apply_groupwise_scale(accum, intermediate, tCrS, weight_scale_idx, is_first_accum); - } - }; - ConsumerToken barrier_token = {BarrierStatus::WaitAgain}; // First k tile { @@ -1124,56 +1026,95 @@ struct CollectiveMmaArrayMixedInput< } // src: tCrA_load, dst: tCrA_mma - Utils::convert_A_kblock(tCrA_load_4b_packed, tCrA_mma, cute::Int<0>{}); + Utils::convert_A_kblock(tCrA_load_4b_packed, tCrA_mma, 0); // Unroll the K mode manually to set scale D to 1 - cute::for_each(cute::make_seq{}, [&](auto chunk_id_c) { - constexpr int chunk_id = decltype(chunk_id_c)::value; + CUTLASS_PRAGMA_UNROLL + for (int chunk_id = 0; chunk_id < NumChunksPerTileK; ++chunk_id) { tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; - cute::for_each(cute::make_seq{}, [&](auto mma_id_c) { - constexpr int mma_id = decltype(mma_id_c)::value; - constexpr int k_block = chunk_id * NumMMAsPerChunk + mma_id; + CUTLASS_PRAGMA_UNROLL + for (int mma_id = 0; mma_id < NumMMAsPerChunk; ++mma_id) { + int k_block = chunk_id * NumMMAsPerChunk + mma_id; warpgroup_arrive(); // (V,M) x (V,N) => (V,M,N) - cute::gemm(tiled_mma, tCrA_mma(_, _, cute::Int{}), - tCrB(_, _, cute::Int{}, read_stage), intermediate_array[chunk_id]); + cute::gemm(tiled_mma, tCrA_mma(_, _, k_block), tCrB(_, _, k_block, read_stage), + intermediate_array[chunk_id]); tiled_mma.accumulate_ = GMMA::ScaleOut::One; - if constexpr (k_block == 0) { - copy_scale_for_tile(read_stage); + if (k_block == 0) { + Utils::copy_tensors_SFA(partitioned_extra_info, copy_partitions_extra_info, 0, + read_stage); } - if constexpr (k_block < K_BLOCK_MAX - 2) { + if (k_block < K_BLOCK_MAX - 2) { Utils::copy_tensors_A(smem_tiled_copy_A_LDSM, tCsA_LDSM, tCrA_copy_view_LDSM, k_block + 2, read_stage); } - if constexpr (k_block < K_BLOCK_MAX - 1) { - Utils::convert_A_kblock(tCrA_load_4b_packed, tCrA_mma, cute::Int{}); + if (k_block < K_BLOCK_MAX - 1) { + Utils::convert_A_kblock(tCrA_load_4b_packed, tCrA_mma, k_block + 1); } - }); + } warpgroup_commit_batch(); - if constexpr (chunk_id > 0) { + if (chunk_id > 0) { warpgroup_wait<1>(); - constexpr int chunk_id_ = chunk_id - 1; + int chunk_id_ = chunk_id - 1; warpgroup_fence_operand(intermediate_array[chunk_id_]); - scale_intermediate(intermediate_array[chunk_id_], chunk_id_, read_stage, chunk_id_ == 0); + // Apply the group-wise scaling + // tCrS ((4, _2, _2), MMA_M, _1) + // accum ((2, _2, _2), MMA_M, _1) + auto tCrS = cute::get<1>(partitioned_extra_info); + for (int mma_m = 0; mma_m < size<1>(accum); mma_m++) { + for (int m = 0; m < size<0, 1>(accum); m++) { + auto scale_coord = make_coord(make_tuple(0, m, 0), mma_m, 0); + for (int n = 0; n < size<0, 2>(accum); n++) { + for (int e = 0; e < size<0, 0>(accum); e++) { + auto accum_coord = make_coord(make_tuple(e, m, n), mma_m, 0); + + if (chunk_id_ == 0) { + accum(accum_coord) = intermediate_array[chunk_id_](accum_coord) * + scale_convertor(tCrS(scale_coord)[0]); + } else { + accum(accum_coord) = + fma(intermediate_array[chunk_id_](accum_coord), + scale_convertor(tCrS(scale_coord)[chunk_id_]), accum(accum_coord)); + } + } + } + } + } } - }); + } warpgroup_wait<0>(); - constexpr int chunk_id_ = NumChunksPerTileK - 1; + int chunk_id_ = NumChunksPerTileK - 1; warpgroup_fence_operand(intermediate_array[chunk_id_]); - scale_intermediate(intermediate_array[chunk_id_], chunk_id_, read_stage, - NumChunksPerTileK == 1); + // Apply the group-wise scaling + // tCrS ((4, _2, _2), MMA_M, _1) + // accum ((2, _2, _2), MMA_M, _1) + auto tCrS = cute::get<1>(partitioned_extra_info); + for (int mma_m = 0; mma_m < size<1>(accum); mma_m++) { + for (int m = 0; m < size<0, 1>(accum); m++) { + auto scale_coord = make_coord(make_tuple(0, m, 0), mma_m, 0); + for (int n = 0; n < size<0, 2>(accum); n++) { + for (int e = 0; e < size<0, 0>(accum); e++) { + auto accum_coord = make_coord(make_tuple(e, m, n), mma_m, 0); + + accum(accum_coord) = + fma(intermediate_array[chunk_id_](accum_coord), + scale_convertor(tCrS(scale_coord)[chunk_id_]), accum(accum_coord)); + } + } + } + } --k_tile_count; if (k_tile_count > 0) { @@ -1186,7 +1127,7 @@ struct CollectiveMmaArrayMixedInput< Utils::copy_tensors_A(smem_tiled_copy_A_LDSM, tCsA_LDSM, tCrA_copy_view_LDSM, 1, smem_pipe_read.index()); - Utils::convert_A_kblock(tCrA_load_4b_packed, tCrA_mma, cute::Int<0>{}); + Utils::convert_A_kblock(tCrA_load_4b_packed, tCrA_mma, 0); } } @@ -1205,32 +1146,33 @@ struct CollectiveMmaArrayMixedInput< ++smem_pipe_read; // Unroll the K mode manually to set scale D to 1 - cute::for_each(cute::make_seq{}, [&](auto chunk_id_c) { - constexpr int chunk_id = decltype(chunk_id_c)::value; + CUTLASS_PRAGMA_UNROLL + for (int chunk_id = 0; chunk_id < NumChunksPerTileK; ++chunk_id) { tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; - cute::for_each(cute::make_seq{}, [&](auto mma_id_c) { - constexpr int mma_id = decltype(mma_id_c)::value; - constexpr int k_block = chunk_id * NumMMAsPerChunk + mma_id; + CUTLASS_PRAGMA_UNROLL + for (int mma_id = 0; mma_id < NumMMAsPerChunk; ++mma_id) { + int k_block = chunk_id * NumMMAsPerChunk + mma_id; warpgroup_arrive(); // (V,M) x (V,N) => (V,M,N) - cute::gemm(tiled_mma, tCrA_mma(_, _, cute::Int{}), - tCrB(_, _, cute::Int{}, read_stage), intermediate_array[chunk_id]); + cute::gemm(tiled_mma, tCrA_mma(_, _, k_block), tCrB(_, _, k_block, read_stage), + intermediate_array[chunk_id]); tiled_mma.accumulate_ = GMMA::ScaleOut::One; - if constexpr (k_block == K_BLOCK_MAX - 1) { + if (k_block == K_BLOCK_MAX - 1) { pipeline.consumer_release( smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it ++smem_pipe_release; } - if constexpr (k_block == 0) { + if (k_block == 0) { barrier_token = pipeline.consumer_try_wait(smem_pipe_read); - copy_scale_for_tile(read_stage); + Utils::copy_tensors_SFA(partitioned_extra_info, copy_partitions_extra_info, 0, + read_stage); } - if constexpr (k_block == K_BLOCK_MAX - 1) { + if (k_block == K_BLOCK_MAX - 1) { // The last k_block pipeline.consumer_wait(smem_pipe_read, barrier_token); @@ -1244,29 +1186,59 @@ struct CollectiveMmaArrayMixedInput< warpgroup_fence_operand(intermediate_array[chunk_id]); - scale_intermediate(intermediate_array[chunk_id], chunk_id, read_stage, false); + // Apply the group-wise scaling + auto tCrS = cute::get<1>(partitioned_extra_info); + for (int mma_m = 0; mma_m < size<1>(accum); mma_m++) { + for (int m = 0; m < size<0, 1>(accum); m++) { + auto scale_coord = make_coord(make_tuple(0, m, 0), mma_m, 0); + for (int n = 0; n < size<0, 2>(accum); n++) { + for (int e = 0; e < size<0, 0>(accum); e++) { + auto accum_coord = make_coord(make_tuple(e, m, n), mma_m, 0); + + accum(accum_coord) = + fma(intermediate_array[chunk_id](accum_coord), + scale_convertor(tCrS(scale_coord)[chunk_id]), accum(accum_coord)); + } + } + } + } - Utils::convert_A_kblock(tCrA_load_4b_packed, tCrA_mma, cute::Int<0>{}); + Utils::convert_A_kblock(tCrA_load_4b_packed, tCrA_mma, 0); } else { - if constexpr (k_block < K_BLOCK_MAX - 2) { + if (k_block < K_BLOCK_MAX - 2) { Utils::copy_tensors_A(smem_tiled_copy_A_LDSM, tCsA_LDSM, tCrA_copy_view_LDSM, k_block + 2, read_stage); } - Utils::convert_A_kblock(tCrA_load_4b_packed, tCrA_mma, cute::Int{}); + Utils::convert_A_kblock(tCrA_load_4b_packed, tCrA_mma, k_block + 1); } - }); + } warpgroup_commit_batch(); - if constexpr (chunk_id > 0) { + if (chunk_id > 0) { warpgroup_wait<1>(); - constexpr int chunk_id_ = chunk_id - 1; + int chunk_id_ = chunk_id - 1; warpgroup_fence_operand(intermediate_array[chunk_id_]); - scale_intermediate(intermediate_array[chunk_id_], chunk_id_, read_stage, false); + // Apply the group-wise scaling + auto tCrS = cute::get<1>(partitioned_extra_info); + for (int mma_m = 0; mma_m < size<1>(accum); mma_m++) { + for (int m = 0; m < size<0, 1>(accum); m++) { + auto scale_coord = make_coord(make_tuple(0, m, 0), mma_m, 0); + for (int n = 0; n < size<0, 2>(accum); n++) { + for (int e = 0; e < size<0, 0>(accum); e++) { + auto accum_coord = make_coord(make_tuple(e, m, n), mma_m, 0); + + accum(accum_coord) = + fma(intermediate_array[chunk_id_](accum_coord), + scale_convertor(tCrS(scale_coord)[chunk_id_]), accum(accum_coord)); + } + } + } + } } - }); + } } { @@ -1280,45 +1252,60 @@ struct CollectiveMmaArrayMixedInput< tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; // Unroll the K mode manually to set scale D to 1 - cute::for_each(cute::make_seq{}, [&](auto k_block_c) { - constexpr int k_block = decltype(k_block_c)::value; - + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) { warpgroup_arrive(); // (V,M) x (V,N) => (V,M,N) - cute::gemm(tiled_mma, tCrA_mma(_, _, cute::Int{}), - tCrB(_, _, cute::Int{}, read_stage), intermediate); + cute::gemm(tiled_mma, tCrA_mma(_, _, k_block), tCrB(_, _, k_block, read_stage), + intermediate); tiled_mma.accumulate_ = GMMA::ScaleOut::One; - if constexpr (k_block == 0) { - copy_scale_for_tile(read_stage); + if (k_block == 0) { + Utils::copy_tensors_SFA(partitioned_extra_info, copy_partitions_extra_info, 0, + read_stage); } - if constexpr (k_block == K_BLOCK_MAX - 1) { + if (k_block == K_BLOCK_MAX - 1) { // release prior barrier pipeline.consumer_release( smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it ++smem_pipe_release; } - if constexpr (k_block < K_BLOCK_MAX - 2) { + if (k_block < K_BLOCK_MAX - 2) { Utils::copy_tensors_A(smem_tiled_copy_A_LDSM, tCsA_LDSM, tCrA_copy_view_LDSM, k_block + 2, read_stage); } - if constexpr (k_block < K_BLOCK_MAX - 1) { - Utils::convert_A_kblock(tCrA_load_4b_packed, tCrA_mma, cute::Int{}); + if (k_block < K_BLOCK_MAX - 1) { + Utils::convert_A_kblock(tCrA_load_4b_packed, tCrA_mma, k_block + 1); } - if constexpr ((k_block + 1) % NumMMAsPerChunk == 0) { + if ((k_block + 1) % NumMMAsPerChunk == 0) { tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; warpgroup_commit_batch(); warpgroup_wait<0>(); warpgroup_fence_operand(intermediate); - constexpr int scale_idx = k_block / NumMMAsPerChunk; - scale_intermediate(intermediate, scale_idx, read_stage, false); + // Apply the group-wise scaling + auto tCrS = cute::get<1>(partitioned_extra_info); + for (int mma_m = 0; mma_m < size<1>(accum); mma_m++) { + for (int m = 0; m < size<0, 1>(accum); m++) { + auto scale_coord = make_coord(make_tuple(0, m, 0), mma_m, 0); + for (int n = 0; n < size<0, 2>(accum); n++) { + for (int e = 0; e < size<0, 0>(accum); e++) { + auto accum_coord = make_coord(make_tuple(e, m, n), mma_m, 0); + int scale_idx = k_block / NumMMAsPerChunk; + + accum(accum_coord) = + fma(intermediate(accum_coord), scale_convertor(tCrS(scale_coord)[scale_idx]), + accum(accum_coord)); + } + } + } + } } - }); + } } } @@ -1344,45 +1331,195 @@ struct CollectiveMmaArrayMixedInput< // Methods to perform different parts of TMA/Tensormap modifications // CUTLASS_DEVICE auto tensormaps_init(Params const& mainloop_params, - [[maybe_unused]] TensorMapStorage& shared_tensormaps, - [[maybe_unused]] int32_t sm_count, int32_t sm_idx) { - (void)sm_idx; + TensorMapStorage& shared_tensormaps, int32_t sm_count, + int32_t sm_idx) { + cute::TmaDescriptor* gmem_tensormap = + reinterpret_cast(mainloop_params.tensormaps); + + cute::TmaDescriptor* tma_desc_a = &gmem_tensormap[sm_idx]; + cute::TmaDescriptor* tma_desc_b = &gmem_tensormap[sm_idx + sm_count]; + cute::TmaDescriptor* tma_desc_scale = &gmem_tensormap[sm_idx + 2 * sm_count]; + cute::TmaDescriptor* tma_desc_zero = &gmem_tensormap[sm_idx + 3 * sm_count]; + + // Bringing tensormaps from params to smem for modification later + Tensor pA_tensormap = + make_tensor(mainloop_params.tma_load_a.get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor sA_tensormap = + make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_A), Int<1>{}, Int<1>{}); + Tensor pB_tensormap = + make_tensor(mainloop_params.tma_load_b.get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor sB_tensormap = + make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_B), Int<1>{}, Int<1>{}); + + if (cute::elect_one_sync()) { + copy(recast(pA_tensormap), recast(sA_tensormap)); + copy(recast(pB_tensormap), recast(sB_tensormap)); + } + + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + Tensor pS_tensormap = + make_tensor(mainloop_params.tma_load_scale.get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor sS_tensormap = + make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_scale), Int<1>{}, Int<1>{}); + if (cute::elect_one_sync()) { + copy(recast(pS_tensormap), recast(sS_tensormap)); + } + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + Tensor pZ_tensormap = + make_tensor(mainloop_params.tma_load_zero.get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor sZ_tensormap = + make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_zero), Int<1>{}, Int<1>{}); + if (cute::elect_one_sync()) { + copy(recast(pZ_tensormap), recast(sZ_tensormap)); + } + } else if constexpr (KernelConversionMode != ConversionMode::DirectConvert) { + static_assert(cutlass::detail::dependent_false, + "Conversion mode not handled in tensormaps_init."); + } + + __syncwarp(); + if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { - return cute::make_tuple(mainloop_params.ptr_A_prebuilt_tma_desc, - mainloop_params.ptr_B_prebuilt_tma_descs); + return cute::make_tuple(tma_desc_a, tma_desc_b); } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { - return cute::make_tuple(mainloop_params.ptr_A_prebuilt_tma_desc, - mainloop_params.ptr_B_prebuilt_tma_descs); + return cute::make_tuple(tma_desc_a, tma_desc_b, tma_desc_scale); + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + return cute::make_tuple(tma_desc_a, tma_desc_b, tma_desc_scale, tma_desc_zero); } else { static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in tensormaps_init."); } } + // Replace address for the global tensor (to be done by single thread) template - CUTLASS_DEVICE void tensormaps_cp_fence_release( - [[maybe_unused]] TensorMapStorage& shared_tensormaps, - [[maybe_unused]] cute::tuple const& input_tensormaps) {} - - // The entire warp must call this function collectively. - template - CUTLASS_DEVICE void tensormaps_fence_acquire(cute::tuple const& input_tensormaps) { - cute::tma_descriptor_fence_acquire(get<0>(input_tensormaps)); - cute::tma_descriptor_fence_acquire(current_tma_desc_b_); + CUTLASS_DEVICE void tensormaps_replace_global_address(TensorMapStorage& shared_tensormaps, + Params const& mainloop_params, + cute::tuple const& input_tensormaps, + int32_t next_batch) { + // Replacing global_address for the next batch + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormaps.smem_tensormap_B, + mainloop_params.ptr_B[next_batch]); + + if (TensormapUpdateShapesStridesForAandScale) { + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormaps.smem_tensormap_A, + mainloop_params.ptr_A[next_batch]); + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormaps.smem_tensormap_scale, + mainloop_params.ptr_S[next_batch]); + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormaps.smem_tensormap_zero, + mainloop_params.ptr_Z[next_batch]); + } else if constexpr (KernelConversionMode != ConversionMode::DirectConvert) { + static_assert(cutlass::detail::dependent_false, + "Conversion mode not handled in tensormaps_replace_global_address."); + } + } else { + cute::tma_descriptor_replace_addr_in_global_mem(get<0>(input_tensormaps), + mainloop_params.ptr_A[next_batch]); + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + cute::tma_descriptor_replace_addr_in_global_mem(get<2>(input_tensormaps), + mainloop_params.ptr_S[next_batch]); + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + cute::tma_descriptor_replace_addr_in_global_mem(get<3>(input_tensormaps), + mainloop_params.ptr_Z[next_batch]); + } else if constexpr (KernelConversionMode != ConversionMode::DirectConvert) { + static_assert(cutlass::detail::dependent_false, + "Conversion mode not handled in tensormaps_replace_global_address."); + } + } } - template - CUTLASS_DEVICE InputTensors tensors_perform_update( - InputTensors const& input_tensors, Params const& mainloop_params, - [[maybe_unused]] ProblemShape_MNKL problem_shape_mnkl, int32_t next_batch) { - current_group_idx_ = next_batch; - current_tma_desc_b_ = mainloop_params.ptr_B_prebuilt_tma_descs + next_batch; - if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { - return cute::make_tuple(get<0>(input_tensors), get<1>(input_tensors), - mainloop_params.ptr_S[next_batch], get<3>(input_tensors), - get<4>(input_tensors)); + // Replace dim and strides for the global tensor - used only for Grouped GEMM (to be done by + // single thread) + template + CUTLASS_DEVICE void tensormaps_replace_global_tensor_properties( + TensorMapStorage& shared_tensormaps, Params const& mainloop_params, int32_t next_group, + ProblemShape_MNKL problem_shape_mnkl) { + const uint32_t M = get<0>(problem_shape_mnkl); + const uint32_t N = get<1>(problem_shape_mnkl); + const uint32_t K = get<2>(problem_shape_mnkl); + + // Replace all dims for consistency + constexpr int MaxTensorRank = 5; + cute::array prob_shape_A = {1, 1, 1, 1, 1}; + cute::array prob_stride_A = {0, 0, 0, 0, 0}; + cute::array prob_shape_B = {1, 1, 1, 1, 1}; + cute::array prob_stride_B = {0, 0, 0, 0, 0}; + cute::array prob_shape_scale = {1, 1, 1, 1, 1}; + cute::array prob_stride_scale = {0, 0, 0, 0, 0}; + cute::array prob_shape_zero = {1, 1, 1, 1, 1}; + cute::array prob_stride_zero = {0, 0, 0, 0, 0}; + + SwappedElementB const* ptr_B = nullptr; + Tensor tensor_b = make_tensor( + ptr_B, + detail::get_gmem_layout(make_shape(N, K, Int<1>{}), mainloop_params.ptr_dB[next_group])); + cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_b, tensor_b, prob_shape_B, + prob_stride_B); + + for (uint64_t& stride : prob_stride_B) { + stride = (stride * sizeof_bits_v) / 8; + } + + cute::tma_descriptor_replace_dims_strides_in_shared_mem(shared_tensormaps.smem_tensormap_B, + prob_shape_B, prob_stride_B); + + if (TensormapUpdateShapesStridesForAandScale) { + SwappedElementA const* ptr_A = nullptr; + Tensor tensor_a = make_tensor( + ptr_A, + detail::get_gmem_layout(make_shape(M, K, Int<1>{}), mainloop_params.ptr_dA[next_group])); + cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_a, tensor_a, prob_shape_A, + prob_stride_A); + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + NonVoidElementScale const* ptr_S = nullptr; + // auto scale_k = K / mainloop_params.chunk_size; + auto scale_k = K / ScalingGroupSize; + Tensor tensor_scale = + make_tensor(detail::get_logical_ptr(ptr_S), make_shape(M, scale_k, Int<1>{}), + mainloop_params.dS[next_group]); + cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_scale, tensor_scale, + prob_shape_scale, prob_stride_scale); + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + ElementZero const* ptr_Z = nullptr; + // auto scale_k = K / mainloop_params.chunk_size; + auto scale_k = K / ScalingGroupSize; + Tensor tensor_zero = + make_tensor(detail::get_logical_ptr(ptr_Z), make_shape(M, scale_k, Int<1>{}), + mainloop_params.dS[next_group]); + cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_zero, tensor_zero, + prob_shape_zero, prob_stride_zero); + } else if constexpr (KernelConversionMode != ConversionMode::DirectConvert) { + static_assert( + cutlass::detail::dependent_false, + "Conversion mode not handled in tensormaps_replace_global_tensor_properties."); + } + + // Convert strides to byte strides + for (uint64_t& stride : prob_stride_A) { + stride = (stride * sizeof_bits_v) / 8; + } + cute::tma_descriptor_replace_dims_strides_in_shared_mem(shared_tensormaps.smem_tensormap_A, + prob_shape_A, prob_stride_A); + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + for (uint64_t& stride : prob_stride_scale) { + stride = (stride * sizeof_bits_v) / 8; + } + cute::tma_descriptor_replace_dims_strides_in_shared_mem( + shared_tensormaps.smem_tensormap_scale, prob_shape_scale, prob_stride_scale); + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + for (uint64_t& stride : prob_stride_zero) { + stride = (stride * sizeof_bits_v) / 8; + } + cute::tma_descriptor_replace_dims_strides_in_shared_mem( + shared_tensormaps.smem_tensormap_zero, prob_shape_zero, prob_stride_zero); + } else if constexpr (KernelConversionMode != ConversionMode::DirectConvert) { + static_assert( + cutlass::detail::dependent_false, + "Conversion mode not handled in tensormaps_replace_global_tensor_properties."); + } } - return input_tensors; } template @@ -1391,11 +1528,71 @@ struct CollectiveMmaArrayMixedInput< cute::tuple const& input_tensormaps, ProblemShape_MNKL problem_shape_mnkl, int32_t next_batch) { - (void)shared_tensormaps; - (void)mainloop_params; - (void)input_tensormaps; - (void)problem_shape_mnkl; - (void)next_batch; + if (cute::elect_one_sync()) { + // Replacing global_address for the next batch + tensormaps_replace_global_address(shared_tensormaps, mainloop_params, input_tensormaps, + next_batch); + + if constexpr (IsGroupedGemmKernel) { + // Replacing global dims and strides for the next batch + tensormaps_replace_global_tensor_properties(shared_tensormaps, mainloop_params, next_batch, + problem_shape_mnkl); + } + } + } + + template + CUTLASS_DEVICE void tensormaps_cp_fence_release(TensorMapStorage& shared_tensormaps, + cute::tuple const& input_tensormaps) { + // [None][fix] Fix W4A8 MoE kernel issue + // https://github.com/NVIDIA/TensorRT-LLM/pull/7072 + if (cute::elect_one_sync()) { + cute::tma_desc_commit_group(); + cute::tma_desc_wait_group(); + } + + // Entire warp must do this (i.e. it's aligned) + tma_descriptor_cp_fence_release(get<1>(input_tensormaps), shared_tensormaps.smem_tensormap_B); + + if (TensormapUpdateShapesStridesForAandScale) { + TensormapUpdateShapesStridesForAandScale = false; + + tma_descriptor_cp_fence_release(get<0>(input_tensormaps), shared_tensormaps.smem_tensormap_A); + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + tma_descriptor_cp_fence_release(get<2>(input_tensormaps), + shared_tensormaps.smem_tensormap_scale); + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + tma_descriptor_cp_fence_release(get<3>(input_tensormaps), + shared_tensormaps.smem_tensormap_zero); + } else if constexpr (KernelConversionMode != ConversionMode::DirectConvert) { + static_assert(cutlass::detail::dependent_false, + "Conversion mode not handled in tensormaps_cp_fence_release."); + } + } else { + tma_descriptor_fence_release(); + } + } + + // The entire warp must call this function collectively (that is, the instructions are aligned) + template + CUTLASS_DEVICE void tensormaps_fence_acquire(cute::tuple const& input_tensormaps) { + cute::tma_descriptor_fence_acquire(get<0>(input_tensormaps)); + cute::tma_descriptor_fence_acquire(get<1>(input_tensormaps)); + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + cute::tma_descriptor_fence_acquire(get<2>(input_tensormaps)); + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + cute::tma_descriptor_fence_acquire(get<3>(input_tensormaps)); + } else if constexpr (KernelConversionMode != ConversionMode::DirectConvert) { + static_assert(cutlass::detail::dependent_false, + "Conversion mode not handled in tensormaps_fence_acquire."); + } + } + + template + CUTLASS_DEVICE InputTensors tensors_perform_update( + InputTensors const& input_tensors, [[maybe_unused]] Params const& mainloop_params, + [[maybe_unused]] ProblemShape_MNKL problem_shape_mnkl, [[maybe_unused]] int32_t next_batch) { + return input_tensors; } }; diff --git a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/collective/sm90_mma_array_tma_gmma_rs_warpspecialized_mixed_input_prescale.hpp b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/collective/sm90_mma_array_tma_gmma_rs_warpspecialized_mixed_input_prescale.hpp deleted file mode 100644 index 17457d94a5d..00000000000 --- a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/collective/sm90_mma_array_tma_gmma_rs_warpspecialized_mixed_input_prescale.hpp +++ /dev/null @@ -1,1310 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include "cute/algorithm/functional.hpp" -#include "cute/algorithm/gemm.hpp" -#include "cute/arch/cluster_sm90.hpp" -#include "cute/arch/copy_sm90.hpp" -#include "cute/atom/mma_atom.hpp" -#include "cute/numeric/arithmetic_tuple.hpp" -#include "cutlass/cuda_host_adapter.hpp" -#include "cutlass/cutlass.h" -#include "cutlass/gemm/dispatch_policy.hpp" -#include "cutlass/numeric_types.h" -#include "cutlass/pipeline/pipeline.hpp" -#include "cutlass/trace.h" -#include "cutlass_extensions/detail/collective/mixed_input_utils.hpp" - -///////////////////////////////////////////////////////////////////////////////////////////////// - -namespace cutlass::gemm::collective { -using namespace cute; - -///////////////////////////////////////////////////////////////////////////////////////////////// - -// WarpSpecialized Mainloop -template -struct CollectiveMmaArrayMixedInput, - TileShape_, ElementAOptionalTuple, StrideA_, - ElementBOptionalTuple, StrideB_, TiledMma_, GmemTiledCopyA_, - SmemLayoutAtomA_, SmemCopyAtomA_, TransformA_, GmemTiledCopyB_, - SmemLayoutAtomB_, SmemCopyAtomB_, TransformB_> { - public: - enum class ConversionMode { DirectConvert, ConvertAndScale, ConvertAndScaleWithZero }; - - // - // Type Aliases - // - using DispatchPolicy = - MainloopSm90ArrayTmaGmmaWarpSpecializedMixedInputPreScale; - using TileShape = TileShape_; - using KernelSchedule = KernelSchedule_; - - private: - template - friend struct detail::MixedGroupedGemmInputUtils; - using CollectiveType = - CollectiveMmaArrayMixedInput; - using Utils = detail::MixedGroupedGemmInputUtils; - - // - // Type Aliases - // - using ScaleA = detail::deduce_mixed_width_dtype_t<1, ElementAOptionalTuple>; - using ScaleB = detail::deduce_mixed_width_dtype_t<1, ElementBOptionalTuple>; - using ZeroA = detail::deduce_mixed_width_dtype_t<2, ElementAOptionalTuple>; - using ZeroB = detail::deduce_mixed_width_dtype_t<2, ElementBOptionalTuple>; - - public: - using ElementA = detail::deduce_mixed_width_dtype_t<0, ElementAOptionalTuple>; - using ElementB = detail::deduce_mixed_width_dtype_t<0, ElementBOptionalTuple>; - static constexpr bool IsANarrow = sizeof_bits::value < sizeof_bits::value; - static constexpr bool HasWeightScale = !cute::is_void_v; - static constexpr bool HasZeroB = !cute::is_void_v; - static_assert( - IsANarrow, - "SM90 mixed-input mainloop expects the first operand to be the narrow transformed weight."); - static_assert(HasWeightScale, "The transformed weight operand must carry mixed-input scale."); - static_assert(!HasZeroB, "Activation operand must not carry zero-point."); - static constexpr bool IsATransformed = true; - using ElementScale = ScaleA; - using ElementZero = ZeroA; - - using StrideA = StrideA_; - using InternalStrideA = cute::remove_pointer_t; - using StrideB = StrideB_; - using InternalStrideB = cute::remove_pointer_t; - - static constexpr bool IsMXFP4 = cute::is_same_v; - static constexpr bool HasActivationScale = !cute::is_void_v && - cute::is_same_v && - cute::is_same_v; - using ElementActivationScale = cute::conditional_t; - // For cases where we can't have a void type, we can use this to allow the code to compile when - // the scale / zero is void. - using NonVoidElementScale = - cute::conditional_t, float, ElementScale>; - using NonVoidElementZero = cute::conditional_t, float, ElementZero>; - using NonVoidElementActivationScale = - cute::conditional_t, cutlass::float_ue8m0_t, - ElementActivationScale>; - // The GEMM kernel consumes weight scales in Ktile-major, MN-contiguous form. - // MXFP8 activation scales stay in their raw M-major, K-contiguous form and - // are loaded with a separate TMA descriptor. - using StrideScale = cute::Stride, int64_t, int64_t>; - using NonVoidStrideScale = cute::conditional_t, - cute::Stride<_1, int64_t, int64_t>, StrideScale>; - using StrideActivationScale = cute::Stride, int64_t>; - - static_assert( - (IsATransformed && (cutlass::gemm::detail::is_k_major() || - is_layout::value || is_layout::value)) || - (!IsATransformed && (cutlass::gemm::detail::is_k_major() || - is_layout::value || is_layout::value)), - "The transformed type must be K-major."); - - static_assert((IsATransformed && (sizeof(ElementB) == 2)) || - (!IsATransformed && (sizeof(ElementA) == 2)) || - ((cutlass::gemm::detail::is_k_major() || is_layout::value || - is_layout::value) && - (cutlass::gemm::detail::is_k_major() || is_layout::value || - is_layout::value)), - "The unscaled element must be 2 bytes OR both inputs must be K-major"); - - static_assert(cutlass::gemm::detail::is_mn_major(), - "Scale tensor consumed by the GEMM kernel must be MN major."); - - static constexpr int ScalingGroupSize = detail::DefaultWeightScaleGroupSize::value; - - using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{})); - using TiledMma = TiledMma_; - using ElementAccumulator = typename TiledMma::ValTypeC; - using GmemTiledCopyA = GmemTiledCopyA_; - using GmemTiledCopyB = GmemTiledCopyB_; - using SmemLayoutAtomA = SmemLayoutAtomA_; - using SmemLayoutAtomB = SmemLayoutAtomB_; - using SmemCopyAtomA = SmemCopyAtomA_; - using SmemCopyAtomB = SmemCopyAtomB_; - using WeightScaleRawElement = NonVoidElementScale; - using SmemCopyAtomScale = Copy_Atom; - - // We must ensure the type to be scaled goes to RF - static constexpr bool SwapAB = !IsATransformed; - using SwappedStrideA = cute::conditional_t; - using SwappedStrideB = cute::conditional_t; - using InternalSwappedStrideA = cute::conditional_t; - using InternalSwappedStrideB = cute::conditional_t; - using SwappedSmemLayoutAtomA = cute::conditional_t; - using SwappedSmemLayoutAtomB = cute::conditional_t; - using SwappedSmemCopyAtomA = cute::conditional_t; - using SwappedSmemCopyAtomB = cute::conditional_t; - // TMA converts f32 input to tf32 when copying from GMEM to SMEM - // For all other types, cast to size equivalent uint type to avoid any rounding by TMA. - static constexpr bool ConvertF32toTF32A = cute::is_same_v; - static constexpr bool ConvertF32toTF32B = cute::is_same_v; - using ConvertedElementA = - cute::conditional_t>>; - using ConvertedElementB = - cute::conditional_t>>; - using RealSwappedElementA = cute::conditional_t; - using RealSwappedElementB = cute::conditional_t; - using SwappedElementA = cute::conditional_t; - using SwappedElementB = cute::conditional_t; - - using TransformA = TransformA_; - using TransformB = TransformB_; - using SwappedTransformA = cute::conditional_t; - using SwappedTransformB = cute::conditional_t; - using ArchTag = typename DispatchPolicy::ArchTag; - - static constexpr int IsSubbyteA = cute::sizeof_bits_v < 8; - using TmaElementA = cute::conditional_t; - using MainloopPipeline = cutlass::PipelineTmaAsync; - using PipelineState = cutlass::PipelineState; - using PipelineParams = typename MainloopPipeline::Params; - - static constexpr int NumProducerThreadEvents = 1; - - static constexpr int RawActScaleChunksPerTileK = size<2>(TileShape{}) / ScalingGroupSize; - // Keep the derived value nonzero so the modulo static_assert below remains - // well-formed; RawActScaleChunksPerTileK carries the actual validity check. - static constexpr int ActScaleChunksPerTileK = - HasActivationScale ? ((RawActScaleChunksPerTileK > 0) ? RawActScaleChunksPerTileK : 1) : 1; - static constexpr int ActScaleTmaAlignmentChunks = - HasActivationScale ? (128 / cutlass::sizeof_bits::value) : 1; - static constexpr bool ActScaleTmaUsesMinWindow = - ActScaleChunksPerTileK <= ActScaleTmaAlignmentChunks; - static constexpr int ActScaleTmaChunks = - ActScaleTmaUsesMinWindow ? ActScaleTmaAlignmentChunks : ActScaleChunksPerTileK; - static_assert(!HasActivationScale || RawActScaleChunksPerTileK > 0, - "Activation scale TMA requires TileShapeK to cover at least one scale group."); - static_assert( - !HasActivationScale || - (ActScaleTmaUsesMinWindow ? (ActScaleTmaAlignmentChunks % ActScaleChunksPerTileK == 0) - : (ActScaleChunksPerTileK % ActScaleTmaAlignmentChunks == 0)), - "Activation scale TileShapeK must divide or be a multiple of the 16B TMA scale window."); - static constexpr int WeightScaleLogicalMPerFoldBlock = 64; - static constexpr int WeightScaleLogicalKPerFoldBlock = 128; - static constexpr int WeightScaleFoldedMPerFoldBlock = 16; - static constexpr int WeightScaleMSlicesPerFoldBlock = - WeightScaleLogicalMPerFoldBlock / WeightScaleFoldedMPerFoldBlock; - static constexpr int WeightScaleScaleGroupsPerFoldBlock = - WeightScaleLogicalKPerFoldBlock / ScalingGroupSize; - static constexpr int WeightScalePhysicalColsPerFoldBlock = - WeightScaleMSlicesPerFoldBlock * WeightScaleScaleGroupsPerFoldBlock; - static constexpr int WeightScaleMBlocksPerTile = - size<0>(TileShape{}) / WeightScaleLogicalMPerFoldBlock; - static constexpr int WeightScaleKBlocksPerTile = - size<2>(TileShape{}) / WeightScaleLogicalKPerFoldBlock; - static_assert(size<0>(TileShape{}) % WeightScaleLogicalMPerFoldBlock == 0, - "Folded weight scale requires TileShapeM to be a multiple of 64."); - static_assert(size<2>(TileShape{}) % WeightScaleLogicalKPerFoldBlock == 0, - "Folded weight scale requires TileShapeK to be a multiple of 128."); - static_assert(WeightScaleLogicalMPerFoldBlock % WeightScaleFoldedMPerFoldBlock == 0, - "Folded weight scale M dimension must evenly divide the logical M block."); - static_assert(WeightScalePhysicalColsPerFoldBlock * - cutlass::sizeof_bits::value == - 128, - "Folded weight scale must expose 16B per folded-M coordinate."); - static constexpr int ScaleNRawElementsPerStage = size<1>(TileShape{}) * ActScaleTmaChunks; - static constexpr int ScaleNElementsPerStage = size<1>(TileShape{}); - static constexpr int WeightScaleRawElementsPerFoldBlock = - WeightScaleLogicalMPerFoldBlock * WeightScaleLogicalKPerFoldBlock / ScalingGroupSize; - static constexpr int WeightScaleRawElementsPerStage = - WeightScaleRawElementsPerFoldBlock * WeightScaleMBlocksPerTile * WeightScaleKBlocksPerTile; - static constexpr uint32_t WeightScaleFoldBlockBytes = cutlass::bits_to_bytes( - WeightScaleRawElementsPerFoldBlock * cutlass::sizeof_bits::value); - static constexpr uint32_t WeightScaleBulkCopyBytes = - WeightScaleFoldBlockBytes * WeightScaleKBlocksPerTile; - static constexpr uint32_t WeightScaleTransactionBytes = cutlass::bits_to_bytes( - WeightScaleRawElementsPerStage * cutlass::sizeof_bits::value); - static_assert(WeightScaleBulkCopyBytes % 16 == 0, - "Folded weight-scale bulk copy size must be 16B aligned."); - - using SmemLayoutAtomScale = - Layout(SwappedSmemLayoutAtomA{})), cute::Int<1>>>; - using ScaleTileShape = - decltype(make_shape(shape<0>(TileShape{}), shape<1>(SmemLayoutAtomScale{}))); - - static_assert(cute::rank(SwappedSmemLayoutAtomA{}) == 2, - "SmemLayoutAtom must be rank 2 (M/N, K)"); - static_assert((size<0>(TileShape{}) % size<0>(SwappedSmemLayoutAtomA{})) == 0, - "SmemLayoutAtom must evenly divide tile shape."); - static_assert((size<2>(TileShape{}) % size<1>(SwappedSmemLayoutAtomA{})) == 0, - "SmemLayoutAtom must evenly divide tile shape."); - - static_assert(cute::rank(SwappedSmemLayoutAtomB{}) == 2, - "SmemLayoutAtom must be rank 2 (M/N, K)"); - static_assert((size<1>(TileShape{}) % size<0>(SwappedSmemLayoutAtomB{})) == 0, - "SmemLayoutAtom must evenly divide tile shape."); - static_assert((size<2>(TileShape{}) % size<1>(SwappedSmemLayoutAtomB{})) == 0, - "SmemLayoutAtom must evenly divide tile shape."); - - static_assert(rank(SmemLayoutAtomScale{}) == 2, "SmemLayoutAtomScale must be rank 2"); - static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomScale{})) == 0, - "SmemLayoutAtomScale must equal the tile shape."); - static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomScale{})) == 0, - "SmemLayoutAtomScale must evenly divide tile k shape."); - - /// Tile along modes in a way that maximizes the TMA box size. - using SmemLayoutA = decltype(detail::get_smem_layout( - SwappedSmemLayoutAtomA{}, select<0, 2>(TileShape{}), InternalSwappedStrideA{})); - using SmemLayoutB = decltype(detail::get_smem_layout( - SwappedSmemLayoutAtomB{}, select<1, 2>(TileShape{}), InternalSwappedStrideB{})); - - // It is assumed that weight scales and zero-points share the same smem layout. - using SmemLayoutScale = decltype(tile_to_shape( - SmemLayoutAtomScale{}, - make_shape(shape<0>(ScaleTileShape{}), shape<1>(ScaleTileShape{}), Int{}), - cute::conditional_t<::cutlass::gemm::detail::is_major<0, NonVoidStrideScale>(), - Step<_2, _1, _3>, Step<_1, _2, _3>>{})); - using SmemLayoutWeightScaleRaw = - Layout, Int, - Int, Int, Int>, - Stride<_1, Int, - Int, - Int, - Int>>; - using SmemLayoutWeightScaleExpanded = Layout< - Shape, Int, - Int>, - Shape, - Shape, Int>>, - Int>, - Stride< - Stride, Int, - Int>, - Stride<_0, Stride<_1, Int>>, - Int>>; - // MXFP8 activation scales are independent from MXFP4 weight scales. They are - // stored in raw M-major, K-contiguous form and TMA-loaded into this raw scale - // layout: (BLK_N, ActScaleTmaChunks, PIPE). The TMA window is 16B-aligned: - // smaller compute Ktiles reuse a subrange, while larger Ktiles must already - // span a 16B-aligned activation-scale row. - using SmemLayoutActivationScale = - Layout(TileShape{})), Int, Int>, - Stride, _1, Int>>; - - static_assert(DispatchPolicy::Stages >= 2, - "Specialization requires Stages set to value 2 or more."); - static_assert( - not cute::is_base_of::value && - cute::is_base_of::value, - "MMA atom must source A from rmem and B operand from smem_desc for this mainloop."); - static_assert(cute::is_same_v || - cute::is_same_v, - "GmemTiledCopy - invalid SM90 TMA copy atom specified."); - static_assert(cute::is_same_v || - cute::is_same_v, - "GmemTiledCopy - invalid SM90 TMA copy atom specified."); - - // To relax them, we need to handle loading more than 1 row of scales for every main loop - // iteration. We must also handle updating the pipeline transaction bytes on the fly. - static_assert(size<1>(SmemLayoutAtomScale{}) == 1, "size<1>(SmemLayoutAtomScale) must be 1."); - - private: - static constexpr ConversionMode get_conversion_mode() { - if constexpr (cute::is_void_v) { - return ConversionMode::DirectConvert; - } else if constexpr (cute::is_void_v) { - return ConversionMode::ConvertAndScale; - } else { - return ConversionMode::ConvertAndScaleWithZero; - } - } - - int current_group_idx_ = 0; - cute::TmaDescriptor const* current_tma_desc_b_ = nullptr; - - public: - static constexpr ConversionMode KernelConversionMode = get_conversion_mode(); - // MixedInputUtils consumes these traits for shared-memory sizing and layout - // selection. Prescale only supports the FP4->FP8 scale-table case below. - static constexpr bool ModeHasScales = - KernelConversionMode == ConversionMode::ConvertAndScale || - KernelConversionMode == ConversionMode::ConvertAndScaleWithZero; - static constexpr bool FusedE8M0PreMmaScale = true; - static_assert(!HasActivationScale, - "The prescale collective expects activation scale, if any, to be handled outside " - "the mainloop."); - static constexpr bool UseScaleLookupTable = false; - static constexpr bool UseFP4ToBF16LookupTable = - KernelConversionMode == ConversionMode::ConvertAndScale && - cute::is_same_v && - cute::is_same_v; - static constexpr bool UseFP4ToFP8LookupTable = - KernelConversionMode == ConversionMode::ConvertAndScale && - cute::is_same_v && - cute::is_same_v; - static constexpr bool UseInt4ToFP8LookupTable = - KernelConversionMode == ConversionMode::ConvertAndScale && - cute::is_same_v && - cute::is_same_v; - static_assert(UseFP4ToFP8LookupTable && cute::is_same_v, - "Fused e8m0 pre-MMA scale is only implemented for MXFP4 x FP8 with folded scalar " - "e8m0 scales."); - static constexpr size_t SmemAlignmentA = cutlass::detail::alignment_for_swizzle(SmemLayoutA{}); - static constexpr size_t SmemAlignmentB = cutlass::detail::alignment_for_swizzle(SmemLayoutB{}); - static constexpr size_t SmemAlignmentScale = cute::max(SmemAlignmentA, SmemAlignmentB); - - static_assert(SmemAlignmentA >= 128 and SmemAlignmentB >= 128, "Require at least 128B alignment"); - - struct SharedStorage { - static constexpr int scale_elements = cute::cosize_v; - static constexpr int zero_elements = 0; - static constexpr int activation_scale_elements = 0; - struct TensorStorage { - CUTE_ALIGNAS(SmemAlignmentA) - cute::ArrayEngine> smem_A; - CUTE_ALIGNAS(SmemAlignmentB) - cute::ArrayEngine> smem_B; - // Keep the member layout aligned with mixed_input.hpp for online collective - // switching. Prescale only stages weight e8m0 scale; zero and activation - // scale storage are intentionally empty. - cute::ArrayEngine smem_scale; - cute::ArrayEngine - smem_activation_scale; - cute::ArrayEngine smem_zero; - } tensors; - - struct TensorMapStorage {}; - - using PipelineStorage = typename MainloopPipeline::SharedStorage; - PipelineStorage pipeline; - }; - using TensorStorage = typename SharedStorage::TensorStorage; - using TensorMapStorage = typename SharedStorage::TensorMapStorage; - using PipelineStorage = typename SharedStorage::PipelineStorage; - - static constexpr bool IsGroupedGemmKernel = !cute::is_same_v; - static constexpr bool RequiresTensormapUpdateOnBatchChange = false; - static constexpr bool RequiresPrebuiltTensormapAcquireOnBatchChange = IsGroupedGemmKernel; - - // Host side kernel arguments. Keep this parameter surface aligned with - // mixed_input.hpp so callers can switch collectives without rebuilding the - // argument-preparation path; unsupported prescale paths are rejected in - // can_implement instead of by deleting fields here. - struct Arguments { - ElementA const** ptr_A; - StrideA dA; - ElementB const** ptr_B; - StrideB dB; - ElementScale const** ptr_S = nullptr; - NonVoidStrideScale const* dS{}; - int chunk_size = 0; - ElementZero const** ptr_Z = nullptr; - NonVoidElementActivationScale const** ptr_ActivationScale = nullptr; - StrideActivationScale const* dActivationScale{}; - cute::TmaDescriptor const* ptr_A_prebuilt_tma_desc = nullptr; - cute::TmaDescriptor const* ptr_B_prebuilt_tma_descs = nullptr; - cute::TmaDescriptor const* ptr_ActivationScale_prebuilt_tma_descs = nullptr; - }; - - // Device side kernel params - struct Params { - // For grouped GEMM with non-layout stride: replace static-zero L stride (_0) with - // a static non-zero value so the TMA descriptor includes the L dimension at creation. - // Int<32> is the minimum static value that after subbyte upcast<2> (FP4→uint8_t) - // produces Int<16> = 16 bytes, satisfying cuTensorMapEncodeTiled's 16-byte alignment. - // Being fully static, all CuTe coordinate computations remain compile-time optimizable. - using TmaStrideA = - cute::conditional_t::value, - decltype(cute::make_stride(cute::get<0>(InternalSwappedStrideA{}), - cute::get<1>(InternalSwappedStrideA{}), - cute::Int<32>{})), - InternalSwappedStrideA>; - using LayoutA = - decltype(detail::get_gmem_layout(repeat_like(TmaStrideA{}, int32_t(0)), TmaStrideA{})); - using LayoutB = decltype(detail::get_gmem_layout( - repeat_like(InternalSwappedStrideB{}, int32_t(0)), InternalSwappedStrideB{})); - - using TMA_A = decltype(make_tma_copy( - GmemTiledCopyA{}, - make_tensor(detail::get_logical_ptr(static_cast(nullptr)), - LayoutA{}), - SmemLayoutA{}(_, _, cute::Int<0>{}), - make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})), - size<1>(ClusterShape{}))); // mcast along N mode for this M load, if any - // Assumption: StrideB is congruent with Problem_NK - using TMA_B = decltype(make_tma_copy( - GmemTiledCopyB{}, - make_tensor(detail::get_logical_ptr(static_cast(nullptr)), - LayoutB{}), - SmemLayoutB{}(_, _, cute::Int<0>{}), - make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})), - size<0>(ClusterShape{}))); // mcast along M mode for this N load, if any - using LayoutActivationScale = decltype(detail::get_gmem_layout( - repeat_like(StrideActivationScale{}, int32_t(0)), StrideActivationScale{})); - using TMA_ActivationScale_ = decltype(make_tma_copy( - SM90_TMA_LOAD{}, - make_tensor( - detail::get_logical_ptr(static_cast(nullptr)), - LayoutActivationScale{}), - SmemLayoutActivationScale{}(_, _, cute::Int<0>{}), - make_shape(shape<1>(TileShape{}), Int{}), Int<1>{})); - using TMA_ActivationScale = - cute::conditional_t>; - - TMA_A tma_load_a; - TMA_B tma_load_b; - TMA_ActivationScale tma_load_activation_scale; - uint32_t tma_transaction_bytes = TmaTransactionBytes; - SwappedElementA const** ptr_A; - SwappedStrideA ptr_dA; - SwappedElementB const** ptr_B; - SwappedStrideB ptr_dB; - cute::TmaDescriptor const* ptr_A_prebuilt_tma_desc; - cute::TmaDescriptor const* ptr_B_prebuilt_tma_descs; - NonVoidElementScale const** ptr_S; - NonVoidStrideScale const* dS; - NonVoidElementActivationScale const** ptr_ActivationScale; - StrideActivationScale const* dActivationScale; - cute::TmaDescriptor const* ptr_ActivationScale_prebuilt_tma_descs; - NonVoidElementZero const** ptr_Z; - int64_t scale_k; - int chunk_size; - int reload_factor = (chunk_size + size<2>(TileShape{}) - 1) / size<2>(TileShape{}); - InternalSwappedStrideA dA; - InternalSwappedStrideB dB; - int num_groups; - }; - - // - // Methods - // - - template - static constexpr Params to_underlying_arguments(ProblemShape problem_shapes, - Arguments const& args, - [[maybe_unused]] void* workspace) { - // These tensor shapes (only applicable for grouped gemm) and pointers are only used to create - // tensormap/tma desc. These will be replaced with correct values before the initial tma load. - auto init_shape = repeat_like(typename ProblemShape::UnderlyingProblemShape{}, int32_t(1)); - auto init_M = get<0>(init_shape); - auto init_N = get<1>(init_shape); - auto init_K = get<2>(init_shape); - - if constexpr (SwapAB) { - init_M = get<1>(init_shape); - init_N = get<0>(init_shape); - } - // Batches/Groups are managed by using appropriate pointers to input matrices - const uint32_t mock_L = 1; - SwappedElementA const* ptr_A_first_batch; - SwappedElementB const* ptr_B_first_batch; - SwappedStrideA ptr_dA; - SwappedStrideB ptr_dB; - InternalSwappedStrideA dA; - InternalSwappedStrideB dB; - - if constexpr (not SwapAB) { - ptr_A_first_batch = reinterpret_cast(args.ptr_A); - ptr_B_first_batch = reinterpret_cast(args.ptr_B); - } else { - ptr_A_first_batch = reinterpret_cast(args.ptr_B); - ptr_B_first_batch = reinterpret_cast(args.ptr_A); - } - - if constexpr (IsGroupedGemmKernel) { - // Strides for Grouped Gemm will be replaced prior to the first access regardless. - if constexpr (not SwapAB) { - ptr_dA = args.dA; - ptr_dB = args.dB; - } else { - ptr_dA = args.dB; - ptr_dB = args.dA; - } - dA = InternalSwappedStrideA{}; - if constexpr (is_layout::value) { - dA = make_layout(transform_leaf(dA.shape(), - [](auto x) { - if constexpr (not is_static_v) { - return static_cast(1); - } else { - return x; - } - }), - dA.stride()); - } - dB = InternalSwappedStrideB{}; - } else { - // Tensor shapes for Ptr-Array are initialized correctly only here. - auto problem_shape_MNK = problem_shapes.get_host_problem_shape(0); - init_M = get<0>(problem_shape_MNK); - init_N = get<1>(problem_shape_MNK); - init_K = get<2>(problem_shape_MNK); - - if constexpr (not SwapAB) { - dA = args.dA; - dB = args.dB; - } else { - dA = args.dB; - dB = args.dA; - } - ptr_dA = SwappedStrideA{}; - ptr_dB = SwappedStrideB{}; - } - // For grouped GEMM: use TmaStrideA (with static _1 L stride) so the TMA descriptor - // is created as 3D, enabling coordinate-based group selection. - typename Params::TmaStrideA tma_dA; - if constexpr (!IsGroupedGemmKernel || cute::is_layout::value) { - tma_dA = dA; - } - Tensor tensor_a = make_tensor( - ptr_A_first_batch, detail::get_gmem_layout(make_shape(init_M, init_K, mock_L), tma_dA)); - Tensor tensor_b = make_tensor(ptr_B_first_batch, - detail::get_gmem_layout(make_shape(init_N, init_K, mock_L), dB)); - - typename Params::TMA_A tma_load_a = make_tma_copy( - GmemTiledCopyA{}, tensor_a, SmemLayoutA{}(_, _, cute::Int<0>{}), - make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})), - size<1>(ClusterShape{})); // mcast along N mode for this M load, if any - typename Params::TMA_B tma_load_b = - make_tma_copy(GmemTiledCopyB{}, tensor_b, SmemLayoutB{}(_, _, cute::Int<0>{}), - make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})), - size<0>(ClusterShape{})); // mcast along M mode for this N load, if any - typename Params::TMA_ActivationScale tma_load_activation_scale{}; - - int num_groups_val = 1; - if constexpr (IsGroupedGemmKernel) { - num_groups_val = problem_shapes.groups(); - } - auto args_setup = [&](auto ptr_A, auto ptr_B, int64_t scale_k = 0, int chunk_size = 0, - int reload_factor = 1) -> Params { - return {tma_load_a, - tma_load_b, - tma_load_activation_scale, - TmaTransactionBytes, - reinterpret_cast(ptr_A), - ptr_dA, - reinterpret_cast(ptr_B), - ptr_dB, - args.ptr_A_prebuilt_tma_desc, - args.ptr_B_prebuilt_tma_descs, - reinterpret_cast(args.ptr_S), - args.dS, - args.ptr_ActivationScale, - args.dActivationScale, - args.ptr_ActivationScale_prebuilt_tma_descs, - reinterpret_cast(args.ptr_Z), - scale_k, - chunk_size, - reload_factor, - dA, - dB, - num_groups_val}; - }; - - // Prescale keeps the historical scale_k field in Params so the argument - // surface stays aligned with mixed_input.hpp. Runtime scale addressing is - // computed directly from ptr_S, chunk_size, and the current K tile. - int64_t scale_k_placeholder = 1; - return SwapAB ? args_setup(args.ptr_B, args.ptr_A, scale_k_placeholder, args.chunk_size, - (args.chunk_size + size<2>(TileShape{}) - 1) / size<2>(TileShape{})) - : args_setup(args.ptr_A, args.ptr_B, scale_k_placeholder, args.chunk_size, - (args.chunk_size + size<2>(TileShape{}) - 1) / size<2>(TileShape{})); - } - - template - static size_t get_workspace_size([[maybe_unused]] ProblemShape const& problem_shape, - [[maybe_unused]] Arguments const& args, - [[maybe_unused]] int sm_count) { - return 0; - } - - template - static cutlass::Status initialize_workspace( - [[maybe_unused]] ProblemShape const& problem_shape, [[maybe_unused]] Arguments const& args, - [[maybe_unused]] void* workspace, [[maybe_unused]] cudaStream_t stream, - [[maybe_unused]] CudaHostAdapter* cuda_adapter = nullptr) { - return cutlass::Status::kSuccess; - } - - template - CUTLASS_HOST_DEVICE static bool can_implement(ProblemShape problem_shapes, - Arguments const& args) { - constexpr int tma_alignment_bits = 128; - constexpr int min_tma_aligned_elements_A = - tma_alignment_bits / cutlass::sizeof_bits::value; - constexpr int min_tma_aligned_elements_B = - tma_alignment_bits / cutlass::sizeof_bits::value; - - bool implementable = true; - if constexpr (IsGroupedGemmKernel) { - implementable = implementable && (args.ptr_A_prebuilt_tma_desc != nullptr); - implementable = implementable && (args.ptr_B_prebuilt_tma_descs != nullptr); - } - if (problem_shapes.is_host_problem_shape_available()) { - // Check alignment for all problem sizes - for (int i = 0; i < problem_shapes.groups(); i++) { - auto problem_shape_MNKL = append<4>(problem_shapes.get_host_problem_shape(i), 1); - auto [M, N, K, L] = problem_shape_MNKL; - if constexpr (!cute::is_pointer_v>) { - implementable = - implementable && cutlass::detail::check_alignment( - detail::get_gmem_layout(cute::make_shape(M, K, L), args.dA)); - } - if constexpr (!cute::is_pointer_v>) { - implementable = - implementable && cutlass::detail::check_alignment( - detail::get_gmem_layout(cute::make_shape(N, K, L), args.dB)); - } - const int scale_mn = SwapAB ? N : M; - if (args.chunk_size == 0) { - implementable = false; - } else { - implementable = implementable && (args.chunk_size == ScalingGroupSize); - implementable = implementable && ((scale_mn % size<0>(TileShape{})) == 0); - implementable = implementable && ((K % size<2>(TileShape{})) == 0); - } - implementable = implementable && (args.ptr_S != nullptr); - implementable = implementable && (args.ptr_Z == nullptr); - implementable = implementable && (args.ptr_ActivationScale == nullptr); - } - } - - if (!implementable) { - CUTLASS_TRACE_HOST( - " CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for " - "TMA.\n"); - } - return implementable; - } - - static constexpr int K_PIPE_MAX = DispatchPolicy::Stages; - static constexpr int K_PIPE_MMAS = 1; - static constexpr uint32_t TmaTransactionBytesMK = Utils::compute_tma_transaction_bytes_mk(); - static constexpr uint32_t TmaTransactionBytesNK = Utils::compute_tma_transaction_bytes_nk(); - static constexpr uint32_t TmaTransactionBytesExtra = Utils::compute_tma_transaction_bytes_extra(); - static constexpr uint32_t TmaTransactionBytes = - TmaTransactionBytesMK + TmaTransactionBytesNK + TmaTransactionBytesExtra; - - // Set up the data needed by this collective for load and mma. - // Returns a tuple of tensors. The collective and the kernel layer have the contract that the - // returned tuple must contain at least two elements, with the first two elements being: - // gA_mkl - The tma tensor, A after a local tile so it has shape (BLK_M,BLK_K,m,k,l) - // gB_nkl - The tma tensor, B after a local tile so it has shape (BLK_N,BLK_K,n,k,l) - // The rest of the tensors can be specified as needed by this collective. - template - CUTLASS_DEVICE auto load_init(ProblemShape_MNKL const& problem_shape_MNKL, - Params const& mainloop_params) const { - using X = Underscore; - // Separate out problem shape for convenience - auto [M, N, K, L] = problem_shape_MNKL; - const int32_t mock_L = 1; - - // TMA requires special handling of strides to deal with coord codomain mapping - // Represent the full tensors -- get these from TMA - // In this mixed-input layout, the transformed A operand is the offline weight tensor and the B - // operand is the activation tensor. A keeps a grouped L dimension for the - // fixed weight layout; B and activation scales are retargeted per expert. - auto A_L = mainloop_params.num_groups; - auto B_L = mock_L; - Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor( - shape(detail::get_gmem_layout(make_shape(M, K, A_L), mainloop_params.dA))); // (m,k,l) - Tensor mB_nkl = mainloop_params.tma_load_b.get_tma_tensor( - shape(detail::get_gmem_layout(make_shape(N, K, B_L), mainloop_params.dB))); // (n,k,l) - int const scale_total_k128_blocks = int(K) / WeightScaleLogicalKPerFoldBlock; - - // Make tiled views, defer the slice - Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_, _, _), - Step<_1, X, _1>{}); // (BLK_M,BLK_K,m,k,l) - Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_, _, _), - Step{}); // (BLK_N,BLK_K,n,k,l) - return cute::make_tuple(gA_mkl, gB_nkl, scale_total_k128_blocks); - } - - ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// - // Perform a collective-scoped matrix multiply-accumulate - // Producer Perspective - template - CUTLASS_DEVICE void load(Params const& mainloop_params, MainloopPipeline pipeline, - PipelineState smem_pipe_write, cute::tuple const& load_inputs, - [[maybe_unused]] cute::tuple const& input_tensormaps, - BlockCoord const& blk_coord, KTileIterator k_tile_iter, int k_tile_count, - int thread_idx, uint32_t block_rank_in_cluster, - TensorStorage& shared_tensors) { - static_assert(sizeof...(Ts) == 3, - "Fused pre-MMA scale needs three inputs (gA, gB, total_k128_blocks)"); - static_assert(sizeof...(TMs) == 2, "Only A and B tensormaps needed"); - - Tensor sA_ = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), - SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) - Tensor sB_ = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), - SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) - Tensor sA = as_position_independent_swizzle_tensor(sA_); // (BLK_M,BLK_K,PIPE) - Tensor sB = as_position_independent_swizzle_tensor(sB_); // (BLK_N,BLK_K,PIPE) - - // - // Prepare the TMA loads for A and B - // - - constexpr uint32_t cluster_shape_x = get<0>(typename DispatchPolicy::ClusterShape()); - uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, - block_rank_in_cluster / cluster_shape_x}; - - Tensor gA_mkl = get<0>(load_inputs); - Tensor gB_nkl = get<1>(load_inputs); - int const scale_total_k128_blocks = get<2>(load_inputs); - - auto block_tma_a = mainloop_params.tma_load_a.get_slice(cluster_local_block_id.y); - auto block_tma_b = mainloop_params.tma_load_b.get_slice(cluster_local_block_id.x); - - // Partition the inputs based on the current block coordinates. - auto [m_coord, n_coord, k_coord, l_coord] = blk_coord; - auto a_l_coord = current_group_idx_; - auto b_l_coord = cute::Int<0>{}; - Tensor gA = gA_mkl(_, _, m_coord, _, a_l_coord); // (BLK_M,BLK_K,k) - Tensor gB = gB_nkl(_, _, n_coord, _, b_l_coord); // (BLK_N,BLK_K,k) - - // Applies the mapping from block_tma_a - Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k) - Tensor tAsA = block_tma_a.partition_D(sA); // (TMA,TMA_M,TMA_K,PIPE) - - Tensor tBgB = block_tma_b.partition_S(gB); // (TMA,TMA_N,TMA_K,k) - Tensor tBsB = block_tma_b.partition_D(sB); // (TMA,TMA_N,TMA_K,PIPE) - - Tensor sSRaw = make_tensor( - make_smem_ptr(reinterpret_cast(shared_tensors.smem_scale.begin())), - SmemLayoutWeightScaleRaw{}); - - uint16_t mcast_mask_a = 0; - uint16_t mcast_mask_b = 0; - - // Issue TmaLoads - // Maps the tile -> block, value - if constexpr (cute::is_same_v) { - auto block_layout = Layout{}; // (m,n) -> block_id - for (int n = 0; n < size<1>(block_layout); ++n) { - mcast_mask_a |= (uint16_t(1) << block_layout(cluster_local_block_id.x, n, Int<0>{})); - } - } - - if constexpr (cute::is_same_v) { - auto block_layout = Layout{}; // (m,n) -> block_id - for (int m = 0; m < size<0>(block_layout); ++m) { - mcast_mask_b |= (uint16_t(1) << block_layout(m, cluster_local_block_id.y, Int<0>{})); - } - } - - // Mainloop - CUTLASS_PRAGMA_NO_UNROLL - for (; k_tile_count > 0; --k_tile_count) { - // LOCK smem_pipe_write for _writing_ - pipeline.producer_acquire(smem_pipe_write); - - // - // Copy gmem to smem for *k_tile_iter - // - - using BarrierType = typename MainloopPipeline::ProducerBarrierType; - BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write); - - int write_stage = smem_pipe_write.index(); - if (cute::elect_one_sync()) { - // TMA for A and B - copy(mainloop_params.tma_load_a.with(mainloop_params.ptr_A_prebuilt_tma_desc, *tma_barrier, - mcast_mask_a), - tAgA(_, _, _, *k_tile_iter), tAsA(_, _, _, write_stage)); - copy(mainloop_params.tma_load_b.with(current_tma_desc_b_, *tma_barrier, mcast_mask_b), - tBgB(_, _, _, *k_tile_iter), tBsB(_, _, _, write_stage)); - } - - int const scale_k128_offset = int(*k_tile_iter) * WeightScaleKBlocksPerTile; - int const scale_m64_offset = - int(m_coord) * int(size<0>(TileShape{})) / WeightScaleLogicalMPerFoldBlock; - auto* scale_base = - reinterpret_cast(mainloop_params.ptr_S[current_group_idx_]); - - auto scale_gmem_fold_block = [&](int m64_block, int k128_block) { - return int64_t(m64_block) * int64_t(scale_total_k128_blocks) + int64_t(k128_block); - }; - auto issue_scale_bulk_copy = [&](int local_m64_block) { - int const m64_block = scale_m64_offset + local_m64_block; - int64_t const scale_gmem_offset = scale_gmem_fold_block(m64_block, scale_k128_offset) * - int64_t(WeightScaleRawElementsPerFoldBlock); - auto* scale_gmem_addr = reinterpret_cast(scale_base + scale_gmem_offset); - auto* scale_smem_addr = static_cast(&sSRaw(0, 0, local_m64_block, 0, write_stage)); - cute::SM90_BULK_COPY_G2S::copy(scale_gmem_addr, reinterpret_cast(tma_barrier), - scale_smem_addr, WeightScaleBulkCopyBytes); - }; - - if (cute::elect_one_sync()) { - CUTLASS_PRAGMA_UNROLL - for (int local_m64_block = 0; local_m64_block < WeightScaleMBlocksPerTile; - ++local_m64_block) { - issue_scale_bulk_copy(local_m64_block); - } - } - ++k_tile_iter; - - // Advance smem_pipe_write - ++smem_pipe_write; - } - } - ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// - // Perform a Producer Epilogue to prevent early exit of blocks in a Cluster - CUTLASS_DEVICE void load_tail(MainloopPipeline pipeline, PipelineState smem_pipe_write) { - int lane_predicate = cute::elect_one_sync(); - - // Issue the epilogue waits - if (lane_predicate) { - // This helps avoid early exit of blocks in Cluster. - // Waits for all stages to either be released (all - // Consumer UNLOCKs), or if the stage was never used - // then it would just be acquired since the phase was - // still inverted from make_producer_start_state. - pipeline.producer_tail(smem_pipe_write); - } - } - struct NoopReleasedStageProducer { - CUTLASS_DEVICE void operator()() const {} - }; - - /// Perform a collective-scoped matrix multiply-accumulate - /// Consumer Perspective - template - CUTLASS_DEVICE void mma(MainloopPipeline pipeline, PipelineState smem_pipe_read, - FrgTensorC& accum, int k_tile_count, int thread_idx, - TensorStorage& shared_tensors, Params const& mainloop_params) { - NoopReleasedStageProducer released_stage_producer; - mma_with_released_stage_producer(pipeline, smem_pipe_read, accum, k_tile_count, thread_idx, - shared_tensors, mainloop_params, released_stage_producer); - } - - // The compact single-warpgroup kernel refills a stage immediately after the - // current tile safely releases it. Regular kernels compile the no-op callback away. - template - CUTLASS_DEVICE void mma_with_released_stage_producer( - MainloopPipeline pipeline, PipelineState smem_pipe_read, FrgTensorC& accum, int k_tile_count, - int thread_idx, TensorStorage& shared_tensors, Params const& mainloop_params, - ReleasedStageProducer& released_stage_producer) { - static_assert(is_rmem::value, "C tensor must be rmem resident."); - static_assert(cute::rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3."); - static_assert(cute::rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3."); - static_assert(cute::rank(SwappedSmemLayoutAtomA{}) == 2, - "SwappedSmemLayoutAtomA must be rank 2."); - static_assert(cute::rank(SwappedSmemLayoutAtomB{}) == 2, - "SwappedSmemLayoutAtomB must be rank 2."); - static_assert( - !cute::is_void_v, - "SM90 GMMA mainloops must specify a non-void copy atom for smem sourced instructions."); - static_assert( - cute::is_void_v, - "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions."); - - // Obtain warp index - int warp_idx = canonical_warp_idx_sync(); - [[maybe_unused]] int warp_group_thread_idx = thread_idx % 128; - - Tensor sA_ = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), - SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) - Tensor sA = as_position_independent_swizzle_tensor(sA_); // (BLK_M,BLK_K,PIPE) - - Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), - SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) - - // - // Define C accumulators and A/B partitioning - // - - // Layout of warp group to thread mapping - - static_assert(stride<0>(typename TiledMma::BLayout{}) == 0 and - size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup, - "Stride of the first mode must be 0 and the size of the mode must be " - "NumThreadsPerWarpGroup"); - - constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup; - Layout warp_group_thread_layout = - make_layout(Int{}, Int{}); - - int warp_group_idx = thread_idx / NumThreadsPerWarpGroup; - - TiledMma tiled_mma; - auto mma_thread_slice = tiled_mma.get_thread_slice(thread_idx); - Tensor tCsA = mma_thread_slice.partition_A(sA); - auto mma_warpgroup_slice = tiled_mma.get_slice(warp_group_thread_layout(warp_group_idx)); - - // Allocate fragments and descriptors - Tensor tCrA_mma = - mma_thread_slice.partition_fragment_A(sA(_, _, Int<0>{})); // (MMA,MMA_M,MMA_K,PIPE) - Tensor tCrA_load = [&] { - if constexpr (not is_layout::value) { - // Make register tensor with MMA layout - return make_fragment_like(tCrA_mma); - } else { - // Make register tensor matching smem layout, converter will take care of de-swizzling - return make_tensor_like(tCsA(_, _, _, Int<0>{})); - } - }(); - Tensor tCsB = mma_warpgroup_slice.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE) - // tCrB is just a view of the tensor tCsB - Tensor tCrB = mma_warpgroup_slice.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE) - - // - // Copy Atom A retiling - // - auto smem_tiled_copy_A = make_tiled_copy_A(SwappedSmemCopyAtomA{}, tiled_mma); - auto smem_thr_copy_A = smem_tiled_copy_A.get_thread_slice(warp_group_thread_idx); - - Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA_load); // (CPY,CPY_M,CPY_K) - - CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // CPY_M - CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCrA_copy_view)); // CPY_K - CUTE_STATIC_ASSERT_V(size<1>(tCrA_mma) == size<1>(accum)); // MMA_M - CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum)); // N - CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // K - CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB)); // PIPE - CUTE_STATIC_ASSERT_V(Int{} == size<2>(sA)); // PIPE - CUTE_STATIC_ASSERT_V(Int{} == size<2>(sB)); // PIPE - - ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// - - using SmemCopyAtomA_LDSM = Copy_Atom; - - auto smem_tiled_copy_A_LDSM = make_tiled_copy_A(SmemCopyAtomA_LDSM{}, tiled_mma); - auto smem_thr_copy_A_LDSM = smem_tiled_copy_A_LDSM.get_thread_slice(thread_idx); - - Tensor sA_LDSM = recast(sA); - auto tCsA_LDSM = smem_thr_copy_A_LDSM.partition_S(sA_LDSM); - - using ABBitWidthRatio = Int / sizeof_bits_v>; - auto tCrA_load_LDSM_shape = - replace<2>(tCrA_mma.shape(), size(get<2>(tCrA_mma.shape())) / ABBitWidthRatio{}); - Tensor tCrA_load_LDSM = make_fragment_like(tCrA_load_LDSM_shape); - Tensor tCrA_copy_view_LDSM = - smem_thr_copy_A_LDSM.retile_D(tCrA_load_LDSM); // (CPY,CPY_M,CPY_K) - - auto ptr = recast_ptr(tCrA_load_LDSM.data()); - auto old_shape = tCrA_load_LDSM.shape(); - // LDSM packs two 4-bit K sub-blocks before advancing to the next MMA_M - // slice. Preserve that nested K order so MMA_M > 1 does not alias K. - auto tCrA_load_4b_layout = make_layout( - make_shape(size<0>(old_shape), get<1>(old_shape), - make_shape(ABBitWidthRatio{}, size<2>(old_shape))), - make_stride(Int<1>{}, size<0>(old_shape) * ABBitWidthRatio{}, - make_stride(size<0>(old_shape), - size<0>(old_shape) * ABBitWidthRatio{} * size<1>(old_shape)))); - Tensor tCrA_load_4b_packed = make_tensor(ptr, tCrA_load_4b_layout); - - ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// - - Tensor sSRaw = make_tensor( - make_smem_ptr(reinterpret_cast(shared_tensors.smem_scale.begin())), - SmemLayoutWeightScaleExpanded{}); - Tensor tCsSRaw = mma_thread_slice.partition_A(sSRaw); - - PipelineState smem_pipe_release = smem_pipe_read; - - constexpr int K_BLOCK_MAX = size<2>(tCrA_load); - constexpr int K_COMMIT_GROUP_SIZE = 4; - constexpr int K_COMMIT_GROUPS = (K_BLOCK_MAX + K_COMMIT_GROUP_SIZE - 1) / K_COMMIT_GROUP_SIZE; - constexpr int K_WAIT_MAX = (K_COMMIT_GROUPS - 1 < 7) ? K_COMMIT_GROUPS - 1 : 7; - // Large-N tiles expose scale smem->RF latency; small-N best configs keep - // the rolling copy to avoid extending scale register lifetime. - constexpr bool PreloadAllScaleKblocks = size<1>(TileShape{}) >= 128; - static_assert(K_BLOCK_MAX >= 4, "Consider increasing TileShapeK"); - static_assert(ScalingGroupSize % cute::get<0, 1>(tCsB.shape())() == 0, - "Fused e8m0 pre-MMA scale requires scale groups to align to MMA K blocks."); - Tensor tCrA_scale_probe = tCrA_load_4b_packed(_, _, Int<0>{}); - Tensor tCrA_scale_probe_vm = - cute::group_modes<1, -1>(cute::zipped_divide(tCrA_scale_probe, Int<8>{})); - constexpr int ScalePairCount = decltype(size<1>(tCrA_scale_probe_vm))::value / 2; - static_assert(ScalePairCount > 0, - "Fused e8m0 pre-MMA scale cache expects at least one fp4x8 operand pair."); - // TileM256 has enough scale pairs per K block that the compact offset - // cache creates a longer dependency chain than keeping the expanded scale - // tensor in RF. Smaller M tiles still prefer the compact offset cache. - constexpr bool UseExpandedScaleRFForLargeM = size<0>(TileShape{}) >= 256; - Tensor tCrA_scale = make_fragment_like(tCrA_load_4b_packed); - cute::array lo_exp_offsets; - cute::array hi_exp_offsets; - - ConsumerToken barrier_token = {BarrierStatus::WaitAgain}; - auto copy_scale_kblock = [&](auto k_block_c, int read_stage) { - constexpr int k_block = decltype(k_block_c)::value; - if constexpr (k_block < size<2>(tCsSRaw.shape())) { - Tensor scales = tCsSRaw(_, _, k_block_c, read_stage); - if constexpr (UseExpandedScaleRFForLargeM) { - copy(scales, tCrA_scale(_, _, k_block_c)); - } else { - Utils::cache_A_kblock_fused_e8m0_pre_mma_exp_offsets( - scales, k_block_c, Int{}, lo_exp_offsets, hi_exp_offsets); - } - } - }; - auto copy_scale_for_mma = [&](auto k_block_c, int read_stage) { - if constexpr (PreloadAllScaleKblocks) { - if constexpr (decltype(k_block_c)::value == 0) { - cute::for_each(cute::make_seq{}, [&](auto preload_k_block_c) { - copy_scale_kblock(preload_k_block_c, read_stage); - }); - } - } else { - copy_scale_kblock(k_block_c, read_stage); - } - }; - auto convert_A_kblock_static = [&](auto k_block_c, int read_stage) { - auto tCrA_mma_slot = tCrA_mma(_, _, k_block_c); - if constexpr (UseExpandedScaleRFForLargeM) { - Utils::convert_A_kblock_fused_e8m0_pre_mma_raw_scale_to_slot( - tCrA_load_4b_packed, tCrA_mma_slot, tCrA_scale, k_block_c); - } else { - Utils::convert_A_kblock_fused_e8m0_pre_mma_exp_offsets_to_slot( - tCrA_load_4b_packed, tCrA_mma_slot, k_block_c, Int{}, lo_exp_offsets, - hi_exp_offsets); - } - }; - auto commit_mma_group = [&] { - warpgroup_commit_batch(); - // A operand slots are reused by the next K tile. Commit four adjacent K - // blocks as one group and keep only the tail groups outstanding. The - // wait is FIFO: after the last group of tile T, the first group of T is - // retired before tile T+1 overwrites slots 0..3. Subsequent commits in - // tile T+1 keep retiring older tail groups before their slots are reused. - warpgroup_wait(); - }; - auto maybe_commit_mma_group = [&](auto k_block_c) { - constexpr int k_block = decltype(k_block_c)::value; - if constexpr (((k_block + 1) % K_COMMIT_GROUP_SIZE == 0) || (k_block == K_BLOCK_MAX - 1)) { - commit_mma_group(); - } - }; - - // First K tile. - { - barrier_token = pipeline.consumer_try_wait(smem_pipe_read); - pipeline.consumer_wait(smem_pipe_read, barrier_token); - - int read_stage = smem_pipe_read.index(); - - ++smem_pipe_read; - barrier_token = pipeline.consumer_try_wait(smem_pipe_read); - - Utils::copy_tensors_A(smem_tiled_copy_A_LDSM, tCsA_LDSM, tCrA_copy_view_LDSM, 0, read_stage); - copy_scale_for_mma(cute::Int<0>{}, read_stage); - if (K_BLOCK_MAX > 1) { - Utils::copy_tensors_A(smem_tiled_copy_A_LDSM, tCsA_LDSM, tCrA_copy_view_LDSM, 1, - read_stage); - copy_scale_for_mma(cute::Int<1>{}, read_stage); - } - - convert_A_kblock_static(cute::Int<0>{}, read_stage); - - tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; - warpgroup_arrive(); - cute::gemm(tiled_mma, tCrA_mma(_, _, cute::Int<0>{}), tCrB(_, _, cute::Int<0>{}, read_stage), - accum); - maybe_commit_mma_group(cute::Int<0>{}); - tiled_mma.accumulate_ = GMMA::ScaleOut::One; - - Utils::copy_tensors_A(smem_tiled_copy_A_LDSM, tCsA_LDSM, tCrA_copy_view_LDSM, 2, read_stage); - copy_scale_for_mma(cute::Int<2>{}, read_stage); - convert_A_kblock_static(cute::Int<1>{}, read_stage); - - cute::for_each(cute::make_seq{}, [&](auto i) { - constexpr int k_block = decltype(i)::value + 1; - warpgroup_arrive(); - cute::gemm(tiled_mma, tCrA_mma(_, _, cute::Int{}), - tCrB(_, _, cute::Int{}, read_stage), accum); - maybe_commit_mma_group(cute::Int{}); - - if constexpr (k_block < K_BLOCK_MAX - 2) { - Utils::copy_tensors_A(smem_tiled_copy_A_LDSM, tCsA_LDSM, tCrA_copy_view_LDSM, k_block + 2, - read_stage); - copy_scale_for_mma(cute::Int{}, read_stage); - } - if constexpr (k_block < K_BLOCK_MAX - 1) { - convert_A_kblock_static(cute::Int{}, read_stage); - } - }); - - --k_tile_count; - if (k_tile_count > 0) { - pipeline.consumer_wait(smem_pipe_read, barrier_token); - - int const next_read_stage = smem_pipe_read.index(); - Utils::copy_tensors_A(smem_tiled_copy_A_LDSM, tCsA_LDSM, tCrA_copy_view_LDSM, 0, - next_read_stage); - copy_scale_for_mma(cute::Int<0>{}, next_read_stage); - Utils::copy_tensors_A(smem_tiled_copy_A_LDSM, tCsA_LDSM, tCrA_copy_view_LDSM, 1, - next_read_stage); - copy_scale_for_mma(cute::Int<1>{}, next_read_stage); - - // The rolling wait after the last commit has retired the oldest group - // from the previous tile, which is the group that reads A slots 0..3. - convert_A_kblock_static(cute::Int<0>{}, next_read_stage); - } else { - warpgroup_wait<0>(); - } - } - - if (k_tile_count == 0) { - return; - } - - CUTLASS_PRAGMA_NO_UNROLL - for (; k_tile_count > 1; --k_tile_count) { - int read_stage = smem_pipe_read.index(); - ++smem_pipe_read; - - cute::for_each(cute::make_seq{}, [&](auto i) { - constexpr int k_block = decltype(i)::value; - warpgroup_arrive(); - cute::gemm(tiled_mma, tCrA_mma(_, _, cute::Int{}), - tCrB(_, _, cute::Int{}, read_stage), accum); - maybe_commit_mma_group(cute::Int{}); - - if constexpr (k_block == K_BLOCK_MAX - 1) { - pipeline.consumer_release(smem_pipe_release); - ++smem_pipe_release; - released_stage_producer(); - } - - if constexpr (k_block == 0) { - barrier_token = pipeline.consumer_try_wait(smem_pipe_read); - } - - if constexpr (k_block == K_BLOCK_MAX - 1) { - pipeline.consumer_wait(smem_pipe_read, barrier_token); - int const next_read_stage = smem_pipe_read.index(); - Utils::copy_tensors_A(smem_tiled_copy_A_LDSM, tCsA_LDSM, tCrA_copy_view_LDSM, 0, - next_read_stage); - copy_scale_for_mma(cute::Int<0>{}, next_read_stage); - Utils::copy_tensors_A(smem_tiled_copy_A_LDSM, tCsA_LDSM, tCrA_copy_view_LDSM, 1, - next_read_stage); - copy_scale_for_mma(cute::Int<1>{}, next_read_stage); - - // The rolling wait after the last commit has retired the previous - // tile's first A-slot group. Later A-slot groups are retired by - // subsequent grouped commits before their convert points. - convert_A_kblock_static(cute::Int<0>{}, next_read_stage); - } else { - if constexpr (k_block < K_BLOCK_MAX - 2) { - Utils::copy_tensors_A(smem_tiled_copy_A_LDSM, tCsA_LDSM, tCrA_copy_view_LDSM, - k_block + 2, read_stage); - copy_scale_for_mma(cute::Int{}, read_stage); - } - convert_A_kblock_static(cute::Int{}, read_stage); - } - }); - } - - { - int read_stage = smem_pipe_read.index(); - - cute::for_each(cute::make_seq{}, [&](auto i) { - constexpr int k_block = decltype(i)::value; - warpgroup_arrive(); - cute::gemm(tiled_mma, tCrA_mma(_, _, cute::Int{}), - tCrB(_, _, cute::Int{}, read_stage), accum); - maybe_commit_mma_group(cute::Int{}); - - if constexpr (k_block == K_BLOCK_MAX - 1) { - pipeline.consumer_release(smem_pipe_release); - ++smem_pipe_release; - released_stage_producer(); - } - - if constexpr (k_block < K_BLOCK_MAX - 2) { - Utils::copy_tensors_A(smem_tiled_copy_A_LDSM, tCsA_LDSM, tCrA_copy_view_LDSM, k_block + 2, - read_stage); - copy_scale_for_mma(cute::Int{}, read_stage); - } - if constexpr (k_block < K_BLOCK_MAX - 1) { - convert_A_kblock_static(cute::Int{}, read_stage); - } - }); - - warpgroup_wait<0>(); - } - } - ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// - /// Perform a Consumer Epilogue to release all buffers - CUTLASS_DEVICE void mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, - int k_tile_count) { - // Prologue GMMAs - int prologue_mma_count = 1; - k_tile_count -= prologue_mma_count; - - smem_pipe_release.advance(k_tile_count); - - for (int count = 0; count < prologue_mma_count; ++count) { - pipeline.consumer_release( - smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it - ++smem_pipe_release; - } - } - ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// - // - // Methods to perform different parts of TMA/Tensormap modifications - // - CUTLASS_DEVICE auto tensormaps_init(Params const& mainloop_params, - [[maybe_unused]] TensorMapStorage& shared_tensormaps, - [[maybe_unused]] int32_t sm_count, - [[maybe_unused]] int32_t sm_idx) { - return cute::make_tuple(mainloop_params.ptr_A_prebuilt_tma_desc, - mainloop_params.ptr_B_prebuilt_tma_descs); - } - - template - CUTLASS_DEVICE void tensormaps_perform_update( - [[maybe_unused]] TensorMapStorage& shared_tensormaps, - [[maybe_unused]] Params const& mainloop_params, - [[maybe_unused]] cute::tuple const& input_tensormaps, - [[maybe_unused]] ProblemShape_MNKL problem_shape_mnkl, [[maybe_unused]] int32_t next_batch) {} - - template - CUTLASS_DEVICE void tensormaps_cp_fence_release( - [[maybe_unused]] TensorMapStorage& shared_tensormaps, - [[maybe_unused]] cute::tuple const& input_tensormaps) {} - - // The entire warp must call this function collectively (that is, the instructions are aligned) - template - CUTLASS_DEVICE void tensormaps_fence_acquire(cute::tuple const& input_tensormaps) { - cute::tma_descriptor_fence_acquire(get<0>(input_tensormaps)); - cute::tma_descriptor_fence_acquire(current_tma_desc_b_); - } - - template - CUTLASS_DEVICE InputTensors tensors_perform_update( - InputTensors const& input_tensors, [[maybe_unused]] Params const& mainloop_params, - [[maybe_unused]] ProblemShape_MNKL problem_shape_mnkl, [[maybe_unused]] int32_t next_batch) { - current_group_idx_ = next_batch; - current_tma_desc_b_ = mainloop_params.ptr_B_prebuilt_tma_descs + next_batch; - return input_tensors; - } -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace cutlass::gemm::collective - -///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_single_warpgroup_persistent.hpp b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_single_warpgroup_persistent.hpp deleted file mode 100644 index 35f1b35b0cc..00000000000 --- a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_single_warpgroup_persistent.hpp +++ /dev/null @@ -1,421 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ -#pragma once - -#include "cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong_precomputed.hpp" - -namespace cutlass::gemm::kernel { - -enum class SingleWarpgroupPipelineMode { PrefillAll, RollingRefill }; - -template -using SingleWarpgroupPersistentBase = GemmUniversalPrecomputedScheduler< - ProblemShape, CollectiveMainloop, CollectiveEpilogue, - detail::PersistentTileSchedulerSm90GroupPrecomputed>; - -// Small-K persistent kernel that keeps the existing mixed-input collectives -// intact while one warpgroup overlaps the next output tile with current MMA. -template -class SingleWarpgroupPersistentGemm - : public SingleWarpgroupPersistentBase { - private: - using Base = - SingleWarpgroupPersistentBase; - - public: - using ProblemShape = typename Base::ProblemShape; - using CollectiveMainloop = typename Base::CollectiveMainloop; - using CollectiveEpilogue = typename Base::CollectiveEpilogue; - using TileShape = typename Base::TileShape; - using TiledMma = typename Base::TiledMma; - using ArchTag = typename Base::ArchTag; - using InternalStrideA = typename Base::InternalStrideA; - using InternalStrideB = typename Base::InternalStrideB; - using InternalStrideC = typename Base::InternalStrideC; - using InternalStrideD = typename Base::InternalStrideD; - using ClusterShape = typename Base::ClusterShape; - using TileScheduler = typename Base::TileScheduler; - using Arguments = typename Base::Arguments; - using Params = typename Base::Params; - using SharedStorage = typename Base::SharedStorage; - - static constexpr uint32_t MaxThreadsPerBlock = NumThreadsPerWarpGroup; - static constexpr uint32_t MinBlocksPerMultiprocessor = MinCtasPerMultiprocessor_; - static constexpr int PrefetchNextTileStages = PrefetchNextTileStages_; - static constexpr SingleWarpgroupPipelineMode PipelineMode = PipelineMode_; - static constexpr int SharedStorageSize = sizeof(SharedStorage); - - static_assert(MinCtasPerMultiprocessor_ > 0, - "Single-warpgroup persistent GEMM requires a positive CTA/SM target."); - - static dim3 get_block_shape() { return dim3(MaxThreadsPerBlock, 1, 1); } - - static bool can_implement(Arguments const& args) { - bool implementable = Base::can_implement(args); - if constexpr (PipelineMode == SingleWarpgroupPipelineMode::PrefillAll) { - auto problem_shape = args.problem_shape; - if (problem_shape.is_host_problem_shape_available()) { - constexpr int MaxPrefillK = - CollectiveMainloop::DispatchPolicy::Stages * cute::size<2>(TileShape{}); - for (int group = 0; group < problem_shape.groups(); ++group) { - implementable &= cute::get<2>(problem_shape.get_host_problem_shape(group)) <= MaxPrefillK; - } - } - } - return implementable; - } - - CUTLASS_DEVICE - void operator()(Params const& params, char* smem_buf) { - using namespace cute; - using X = Underscore; - -#if !defined(__CUDA_ARCH_FEAT_SM90_ALL) - printf( - "ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. " - "Aborting.\n"); -#else - static_assert(size(TiledMma{}) == NumThreadsPerWarpGroup, - "Single-warpgroup persistent GEMM requires a 128-thread TiledMma."); - static_assert(size(ClusterShape{}) == 1, - "The single-warpgroup kernel supports only a 1x1x1 cluster."); - static_assert(Base::IsGroupedGemmKernel, - "The single-warpgroup kernel supports grouped GEMM only."); - static_assert( - PrefetchNextTileStages > 0 && - PrefetchNextTileStages <= CollectiveMainloop::DispatchPolicy::Stages, - "Cross-tile prefetch depth must be positive and fit inside the mainloop stage ring."); - static_assert(rank(InternalStrideA{}) == 3 && rank(InternalStrideB{}) == 3, - "Mainloop strides must be rank-3."); - static_assert(rank(InternalStrideC{}) == 3 && rank(InternalStrideD{}) == 3, - "Epilogue strides must be rank-3."); - - SharedStorage& shared_storage = *reinterpret_cast(smem_buf); - int const thread_idx = int(threadIdx.x); - int const lane_idx = canonical_lane_idx(); - int const warp_idx = canonical_warp_idx_sync(); - int const mma_thread_idx = thread_idx; - uint32_t const block_rank_in_cluster = cute::block_rank_in_cluster(); - - using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline; - typename MainloopPipeline::Params mainloop_pipeline_params; - mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::ProducerConsumer; - mainloop_pipeline_params.is_leader = thread_idx == 0; - mainloop_pipeline_params.num_consumers = NumThreadsPerWarpGroup; - mainloop_pipeline_params.num_producers = CollectiveMainloop::NumProducerThreadEvents; - mainloop_pipeline_params.transaction_bytes = params.mainloop.tma_transaction_bytes; - MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params, - ClusterShape{}); - - using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline; - typename EpiLoadPipeline::Params epi_load_pipeline_params; - epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer; - epi_load_pipeline_params.dst_blockid = block_rank_in_cluster; - epi_load_pipeline_params.producer_arv_count = NumThreadsPerWarp; - epi_load_pipeline_params.consumer_arv_count = NumThreadsPerWarpGroup; - if constexpr (CollectiveEpilogue::RequiresTransactionBytes) { - epi_load_pipeline_params.transaction_bytes = params.epilogue.tma_transaction_bytes; - } - EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params); - - using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline; - typename EpiStorePipeline::Params epi_store_pipeline_params; - epi_store_pipeline_params.always_wait = true; - EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params); - - typename CollectiveMainloop::PipelineState mainloop_pipe_consumer_state; - typename CollectiveEpilogue::LoadPipelineState epi_load_pipe_consumer_state; - PipelineState mainloop_pipe_producer_state = - cutlass::make_producer_start_state(); - PipelineState epi_store_pipe_producer_state = - cutlass::make_producer_start_state(); - - __syncthreads(); - - TiledMma tiled_mma; - auto const blk_shape = TileShape{}; - TileScheduler scheduler{params.scheduler}; - CollectiveMainloop collective_mainloop; - CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue); - - // This wrapper intentionally has no epilogue producer warp. The Humming-style - // token-scale callback reads its small row scale directly during store. - if (collective_epilogue.is_producer_load_needed()) { - return; - } - - auto work_tile_info = scheduler.initial_work_tile_info(ClusterShape{}); - if (!work_tile_info.is_valid()) { - return; - } - - auto problem_shape_MNKL = - append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), 1); - auto load_inputs = collective_mainloop.load_init(problem_shape_MNKL, params.mainloop); - static_assert(tuple_size_v >= 2, - "load_init must return at least A and B tensors."); - Tensor gA_mkl = get<0>(load_inputs); - Tensor gB_nkl = get<1>(load_inputs); - - int32_t const logical_sm_idx = int32_t(blockIdx.x + blockIdx.y * gridDim.x); - int32_t const logical_sm_count = params.hw_info.sm_count; - auto input_tensormaps = collective_mainloop.tensormaps_init( - params.mainloop, shared_storage.tensormaps.mainloop, logical_sm_count, logical_sm_idx); - - constexpr int EpilogueDescriptorSlot = 0; - auto epi_store_tensormap = get<0>( - collective_epilogue.store_init(params.epilogue, shared_storage.tensormaps.epilogue, - logical_sm_count, logical_sm_idx, EpilogueDescriptorSlot)); - - int32_t current_group = -1; - constexpr bool IsEpiLoad = false; - int current_prefetched_stages = 0; - - while (work_tile_info.is_valid()) { - int32_t const next_group = work_tile_info.L_idx; - bool const did_group_change = next_group != current_group; - if (did_group_change) { - problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(next_group), 1); - } - if (did_group_change && warp_idx == 0) { - load_inputs = collective_mainloop.tensors_perform_update(load_inputs, params.mainloop, - problem_shape_MNKL, next_group); - collective_mainloop.tensormaps_fence_acquire(input_tensormaps); - - collective_epilogue.template tensormaps_perform_update( - shared_storage.tensormaps.epilogue, params.epilogue, epi_store_tensormap, - problem_shape_MNKL, next_group, EpilogueDescriptorSlot); - __syncwarp(); - collective_epilogue.template tensormaps_cp_fence_release( - shared_storage.tensormaps.epilogue, epi_store_tensormap, EpilogueDescriptorSlot); - } - current_group = next_group; - - auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl)); - auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl)); - auto producer_blk_coord = make_coord(m_coord, n_coord, _, Int<0>{}); - auto epilogue_blk_coord = - make_coord(m_coord, n_coord, _, idx2crd(next_group, shape<4>(gB_nkl))); - - int const work_k_tile_count = - TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, blk_shape); - auto work_k_tile_start = TileScheduler::get_work_k_tile_start(work_tile_info); - auto k_tile_iter = - make_coord_iterator(idx2crd(work_k_tile_start, shape<3>(gA_mkl)), shape<3>(gA_mkl)); - - auto accumulators = partition_fragment_C(tiled_mma, take<0, 2>(blk_shape)); - auto next_work_tile_info = work_tile_info; - auto next_load_inputs = load_inputs; - auto next_problem_shape_MNKL = problem_shape_MNKL; - int next_prefetched_stages = 0; - - CUTLASS_PRAGMA_UNROLL - for (int stage = 0; stage < current_prefetched_stages; ++stage) { - ++k_tile_iter; - } - - int current_prefill_stage_count = work_k_tile_count; - int current_k_tiles_to_refill = 0; - if constexpr (PipelineMode == SingleWarpgroupPipelineMode::RollingRefill) { - current_prefill_stage_count = work_k_tile_count < CollectiveMainloop::DispatchPolicy::Stages - ? work_k_tile_count - : CollectiveMainloop::DispatchPolicy::Stages; - current_k_tiles_to_refill = work_k_tile_count - current_prefill_stage_count; - } - int const current_k_tiles_to_produce = - current_prefill_stage_count - current_prefetched_stages; - CUTLASS_ASSERT(current_k_tiles_to_produce >= 0); - if (current_k_tiles_to_produce > 0 && warp_idx == 0) { - collective_mainloop.load(params.mainloop, mainloop_pipeline, mainloop_pipe_producer_state, - load_inputs, input_tensormaps, producer_blk_coord, k_tile_iter, - current_k_tiles_to_produce, lane_idx, block_rank_in_cluster, - shared_storage.tensors.mainloop); - mainloop_pipe_producer_state.advance(current_k_tiles_to_produce); - } - auto current_refill_k_tile_iter = k_tile_iter; - CUTLASS_PRAGMA_UNROLL - for (int stage = 0; stage < current_k_tiles_to_produce; ++stage) { - ++current_refill_k_tile_iter; - } - - auto next_work = scheduler.fetch_next_work(work_tile_info); - next_work_tile_info = get<0>(next_work); - - auto next_producer_blk_coord = producer_blk_coord; - auto next_work_k_tile_start = work_k_tile_start; - int next_work_k_tile_count = 0; - bool next_group_change = false; - bool next_group_mainloop_state_ready = true; - - if (next_work_tile_info.is_valid()) { - next_group_change = next_work_tile_info.L_idx != current_group; - next_group_mainloop_state_ready = !next_group_change; - if (next_group_change) { - next_problem_shape_MNKL = - append<4>(params.problem_shape.get_problem_shape(next_work_tile_info.L_idx), 1); - } - auto next_m_coord = idx2crd(next_work_tile_info.M_idx, shape<2>(gA_mkl)); - auto next_n_coord = idx2crd(next_work_tile_info.N_idx, shape<2>(gB_nkl)); - next_producer_blk_coord = make_coord(next_m_coord, next_n_coord, _, Int<0>{}); - if (next_group_change) { - next_work_k_tile_count = TileScheduler::get_work_k_tile_count( - next_work_tile_info, next_problem_shape_MNKL, blk_shape); - next_work_k_tile_start = TileScheduler::get_work_k_tile_start(next_work_tile_info); - } else { - next_work_k_tile_count = work_k_tile_count; - next_work_k_tile_start = work_k_tile_start; - } - - if constexpr (PipelineMode == SingleWarpgroupPipelineMode::PrefillAll) { - if (next_group_change && warp_idx == 0) { - next_load_inputs = collective_mainloop.tensors_perform_update( - next_load_inputs, params.mainloop, next_problem_shape_MNKL, - next_work_tile_info.L_idx); - collective_mainloop.tensormaps_fence_acquire(input_tensormaps); - } - next_group_mainloop_state_ready = true; - } - } - - auto next_k_tile_iter = - make_coord_iterator(idx2crd(next_work_k_tile_start, shape<3>(gA_mkl)), shape<3>(gA_mkl)); - - int const available_prefetch_stages = - next_work_k_tile_count < work_k_tile_count ? next_work_k_tile_count : work_k_tile_count; - int const next_prefetch_stage_count = - next_work_tile_info.is_valid() - ? (available_prefetch_stages < PrefetchNextTileStages_ ? available_prefetch_stages - : PrefetchNextTileStages_) - : 0; - int next_k_tiles_to_produce = next_prefetch_stage_count; - auto produce_released_stage = [&] { - if constexpr (PipelineMode == SingleWarpgroupPipelineMode::RollingRefill) { - if (current_k_tiles_to_refill > 0) { - if (warp_idx == 0) { - collective_mainloop.load(params.mainloop, mainloop_pipeline, - mainloop_pipe_producer_state, load_inputs, input_tensormaps, - producer_blk_coord, current_refill_k_tile_iter, 1, lane_idx, - block_rank_in_cluster, shared_storage.tensors.mainloop); - ++current_refill_k_tile_iter; - ++mainloop_pipe_producer_state; - } - --current_k_tiles_to_refill; - return; - } - } - - if (next_k_tiles_to_produce > 0) { - if constexpr (PipelineMode == SingleWarpgroupPipelineMode::RollingRefill) { - if (!next_group_mainloop_state_ready) { - if (warp_idx == 0) { - next_load_inputs = collective_mainloop.tensors_perform_update( - next_load_inputs, params.mainloop, next_problem_shape_MNKL, - next_work_tile_info.L_idx); - collective_mainloop.tensormaps_fence_acquire(input_tensormaps); - } - next_group_mainloop_state_ready = true; - } - } - if (warp_idx == 0) { - collective_mainloop.load( - params.mainloop, mainloop_pipeline, mainloop_pipe_producer_state, next_load_inputs, - input_tensormaps, next_producer_blk_coord, next_k_tile_iter, 1, lane_idx, - block_rank_in_cluster, shared_storage.tensors.mainloop); - ++next_k_tile_iter; - ++mainloop_pipe_producer_state; - } - --next_k_tiles_to_produce; - } - }; - - collective_mainloop.mma_with_released_stage_producer( - mainloop_pipeline, mainloop_pipe_consumer_state, accumulators, work_k_tile_count, - mma_thread_idx, shared_storage.tensors.mainloop, params.mainloop, produce_released_stage); - collective_mainloop.mma_tail(mainloop_pipeline, mainloop_pipe_consumer_state, - work_k_tile_count); - CUTLASS_ASSERT(current_k_tiles_to_refill == 0); - produce_released_stage(); - - if constexpr (PipelineMode == SingleWarpgroupPipelineMode::RollingRefill) { - if (next_work_tile_info.is_valid() && !next_group_mainloop_state_ready) { - if (warp_idx == 0) { - next_load_inputs = collective_mainloop.tensors_perform_update( - next_load_inputs, params.mainloop, next_problem_shape_MNKL, - next_work_tile_info.L_idx); - collective_mainloop.tensormaps_fence_acquire(input_tensormaps); - } - next_group_mainloop_state_ready = true; - } - } - next_prefetched_stages = next_prefetch_stage_count - next_k_tiles_to_produce; - mainloop_pipe_consumer_state.advance(work_k_tile_count); - - TileScheduler::fixup(params.scheduler, work_tile_info, accumulators, 1, 0); - - if (did_group_change && warp_idx == 0) { - collective_epilogue.template tensormaps_fence_acquire(epi_store_tensormap); - } - - auto [epi_load_pipe_consumer_state_next, epi_store_pipe_producer_state_next] = - collective_epilogue.store(epi_load_pipeline, epi_load_pipe_consumer_state, - epi_store_pipeline, epi_store_pipe_producer_state, - problem_shape_MNKL, blk_shape, epilogue_blk_coord, accumulators, - tiled_mma, mma_thread_idx, shared_storage.tensors.epilogue, - epi_store_tensormap, work_tile_info.reduction_subtile_idx()); - epi_load_pipe_consumer_state = epi_load_pipe_consumer_state_next; - epi_store_pipe_producer_state = epi_store_pipe_producer_state_next; - - auto store_tail_states = - collective_epilogue.store_tail(epi_load_pipeline, epi_load_pipe_consumer_state, - epi_store_pipeline, epi_store_pipe_producer_state); - epi_load_pipe_consumer_state = get<0>(store_tail_states); - epi_store_pipe_producer_state = get<1>(store_tail_states); - - if (next_work_tile_info.is_valid()) { - load_inputs = next_load_inputs; - problem_shape_MNKL = next_problem_shape_MNKL; - current_group = next_work_tile_info.L_idx; - } - work_tile_info = next_work_tile_info; - current_prefetched_stages = next_prefetched_stages; - } - - if (warp_idx == 0) { - collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state); - } -#endif - } -}; - -} // namespace cutlass::gemm::kernel diff --git a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative_precomputed.hpp b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative_precomputed.hpp deleted file mode 100644 index 45a420cc2f1..00000000000 --- a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative_precomputed.hpp +++ /dev/null @@ -1,928 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ -#pragma once - -#include - -#include "cute/arch/cluster_sm90.hpp" -#include "cute/tensor.hpp" -#include "cutlass/arch/mma_sm90.h" -#include "cutlass/arch/reg_reconfig.h" -#include "cutlass/cutlass.h" -#include "cutlass/epilogue/collective/detail.hpp" -#include "cutlass/fast_math.h" -#include "cutlass/gemm/dispatch_policy.hpp" -#include "cutlass/gemm/gemm.h" -#include "cutlass/gemm/group_array_problem_shape.hpp" -#include "cutlass/gemm/kernel/gemm_universal_decl.h" -#include "cutlass/gemm/kernel/sm90_tile_scheduler.hpp" -#include "cutlass/gemm/kernel/tile_scheduler.hpp" -#include "cutlass/kernel_hardware_info.hpp" -#include "cutlass/pipeline/pipeline.hpp" -#include "cutlass/trace.h" -#include "cutlass/workspace.h" -#include "cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_warpspecialized_precomputed_decl.hpp" -#include "cutlass_extensions/gemm/kernel/sm90_tile_scheduler_group_precomputed.hpp" - -/////////////////////////////////////////////////////////////////////////////// - -namespace cutlass::gemm::kernel { - -/////////////////////////////////////////////////////////////////////////////// - -#ifndef CUTLASS_EXTENSIONS_PRECOMPUTED_TENSORMAP_TRAITS_HPP_ -#define CUTLASS_EXTENSIONS_PRECOMPUTED_TENSORMAP_TRAITS_HPP_ - -template -struct RequiresBatchTensormapUpdate { - static constexpr bool value = true; -}; - -template -struct RequiresBatchTensormapUpdate< - CollectiveMainloop, - std::void_t> { - static constexpr bool value = CollectiveMainloop::RequiresTensormapUpdateOnBatchChange; -}; - -template -struct RequiresBatchTensormapAcquire { - static constexpr bool value = false; -}; - -template -struct RequiresBatchTensormapAcquire< - CollectiveMainloop, - std::void_t> { - static constexpr bool value = CollectiveMainloop::RequiresPrebuiltTensormapAcquireOnBatchChange; -}; - -#endif - -/////////////////////////////////////////////////////////////////////////////// - -template -class GemmUniversalPrecomputedScheduler< - ProblemShape_, CollectiveMainloop_, CollectiveEpilogue_, TileScheduler_, - cute::enable_if_t>> { - public: - // - // Type Aliases - // - using ProblemShape = ProblemShape_; - static_assert(rank(typename ProblemShape::UnderlyingProblemShape{}) == 3 or - rank(typename ProblemShape::UnderlyingProblemShape{}) == 4, - "ProblemShape{} should be or "); - - static_assert(cute::is_base_of_v); - - static constexpr bool IsGdcEnabled = false; - - // Mainloop derived types - using CollectiveMainloop = CollectiveMainloop_; - using TileShape = typename CollectiveMainloop::TileShape; - using TiledMma = typename CollectiveMainloop::TiledMma; - using ArchTag = typename CollectiveMainloop::ArchTag; - using ElementA = typename CollectiveMainloop::ElementA; - using StrideA = typename CollectiveMainloop::StrideA; - using InternalStrideA = typename CollectiveMainloop::InternalStrideA; - using ElementB = typename CollectiveMainloop::ElementB; - using InternalStrideB = typename CollectiveMainloop::InternalStrideB; - using StrideB = typename CollectiveMainloop::StrideB; - using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy; - using Schedule = typename DispatchPolicy::Schedule; - using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator; - using ClusterShape = typename DispatchPolicy::ClusterShape; - using MainloopArguments = typename CollectiveMainloop::Arguments; - using MainloopParams = typename CollectiveMainloop::Params; - - // Epilogue derived types - using CollectiveEpilogue = CollectiveEpilogue_; - using ElementC = typename CollectiveEpilogue::ElementC; - using StrideC = typename CollectiveEpilogue::StrideC; - using InternalStrideC = typename CollectiveEpilogue::InternalStrideC; - using ElementD = typename CollectiveEpilogue::ElementD; - using StrideD = typename CollectiveEpilogue::StrideD; - using InternalStrideD = typename CollectiveEpilogue::InternalStrideD; - using EpilogueArguments = typename CollectiveEpilogue::Arguments; - using EpilogueParams = typename CollectiveEpilogue::Params; - - static_assert(ArchTag::kMinComputeCapability >= 90); - static_assert(cute::is_void_v, - "Ptr-Array Cooperative and Grouped Gemm Cooperative kernel only supports the " - "default scheduler."); - - static constexpr bool IsGroupedGemmKernel = !cute::is_same_v; - - static_assert(IsGroupedGemmKernel, - "Precomputed grouped scheduler kernel is only for grouped ptr-array GEMM."); - - using SchedulerTag = GroupScheduler; - using TileScheduler = detail::PersistentTileSchedulerSm90GroupPrecomputed; - using TileSchedulerArguments = typename TileScheduler::Arguments; - using TileSchedulerParams = typename TileScheduler::Params; - - static constexpr uint32_t NumLoadWarpGroups = 1; - static constexpr uint32_t NumMmaThreads = size(TiledMma{}); - static constexpr uint32_t NumMmaWarpGroups = NumMmaThreads / NumThreadsPerWarpGroup; - static constexpr uint32_t MaxThreadsPerBlock = - NumMmaThreads + (NumLoadWarpGroups * NumThreadsPerWarpGroup); - static constexpr uint32_t MinBlocksPerMultiprocessor = 1; - static constexpr uint32_t NumProducerThreads = CollectiveMainloop::NumProducerThreadEvents; - - /// Register requirement for Load and Math WGs - static constexpr uint32_t LoadRegisterRequirement = 40; - static constexpr uint32_t MmaRegisterRequirement = 232; - - // 1 stage ordered sequence between mainloop and epilogue producer load threads - using LoadWarpOrderBarrier = cutlass::OrderedSequenceBarrier<1, 2>; - - // Kernel level shared memory storage - struct SharedStorage { - struct TensorStorage : cute::aligned_struct<128, _1> { - using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage; - using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage; - - MainloopTensorStorage mainloop; - EpilogueTensorStorage epilogue; - } tensors; - - struct PipelineStorage : cute::aligned_struct<16, _1> { - using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage; - using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage; - - alignas(16) MainloopPipelineStorage mainloop; - alignas(16) EpiLoadPipelineStorage epi_load; - alignas(16) typename LoadWarpOrderBarrier::SharedStorage load_order; - } pipelines; - - struct TensorMapStorage : cute::aligned_struct<128, _1> { - using MainloopTensorMapStorage = typename CollectiveMainloop::TensorMapStorage; - using EpilogueTensorMapStorage = typename CollectiveEpilogue::TensorMapStorage; - - alignas(128) MainloopTensorMapStorage mainloop; - alignas(128) EpilogueTensorMapStorage epilogue; - } tensormaps; - }; - - static constexpr int SharedStorageSize = sizeof(SharedStorage); - - // Device side arguments - struct Arguments { - GemmUniversalMode mode{}; - ProblemShape problem_shape{}; - MainloopArguments mainloop{}; - EpilogueArguments epilogue{}; - KernelHardwareInfo hw_info{}; - TileSchedulerArguments scheduler{}; - }; - - // Kernel entry point API - struct Params { - GemmUniversalMode mode{}; - ProblemShape problem_shape{}; - MainloopParams mainloop{}; - EpilogueParams epilogue{}; - KernelHardwareInfo hw_info{}; - TileSchedulerParams scheduler{}; - void* workspace{nullptr}; - }; - - // - // Methods - // - - // Convert to underlying arguments. In this case, a simple copy for the aliased type. - static Params to_underlying_arguments(Arguments const& args, void* workspace) { - CUTLASS_TRACE_HOST("to_underlying_arguments():"); - - ProblemShape problem_shapes = args.problem_shape; - - // Get SM count if needed, otherwise use user supplied SM count - int sm_count = args.hw_info.sm_count; - if (sm_count <= 0) { - CUTLASS_TRACE_HOST( - " WARNING: Arguments do not include a valid SM count.\n" - " For optimal performance, populate the arguments KernelHardwareInfo struct with the SM " - "count."); - sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id); - } - CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " - << sm_count); - - // Get maximum number of clusters that could co-exist on the target device - int max_active_clusters = args.hw_info.max_active_clusters; - if (max_active_clusters <= 0) { - max_active_clusters = 0; - CUTLASS_TRACE_HOST( - " WARNING: Arguments do not include a valid max cluster count.\n" - " For optimal performance, populate the arguments KernelHardwareInfo struct with the " - "max_active_clusters."); - } else { - CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid cluster count to " - << max_active_clusters); - } - - KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count, max_active_clusters}; - - // Calculate workspace pointers - uint8_t* workspace_ptr = reinterpret_cast(workspace); - size_t workspace_offset = 0; - - void* epilogue_workspace = workspace_ptr + workspace_offset; - workspace_offset += - CollectiveEpilogue::get_workspace_size(problem_shapes, args.epilogue, sm_count); - workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); - - void* mainloop_workspace = workspace_ptr + workspace_offset; - workspace_offset += - CollectiveMainloop::get_workspace_size(problem_shapes, args.mainloop, sm_count); - workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); - - void* scheduler_workspace = workspace_ptr + workspace_offset; - workspace_offset += - TileScheduler::template get_workspace_size( - args.scheduler, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, - NumMmaWarpGroups); - workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); - - TileSchedulerParams scheduler; - if constexpr (IsGroupedGemmKernel) { - scheduler = - TileScheduler::to_underlying_arguments(problem_shapes, TileShape{}, ClusterShape{}, - hw_info, args.scheduler, scheduler_workspace); - } else { - scheduler = TileScheduler::to_underlying_arguments(problem_shapes.get_host_problem_shape(), - TileShape{}, ClusterShape{}, hw_info, - args.scheduler, scheduler_workspace); - } - - return {args.mode, - problem_shapes, - CollectiveMainloop::to_underlying_arguments(problem_shapes, args.mainloop, - mainloop_workspace), - CollectiveEpilogue::to_underlying_arguments(problem_shapes, args.epilogue, - epilogue_workspace), - hw_info, - scheduler, - workspace}; - } - - static bool can_implement(Arguments const& args) { - bool implementable = true; - if constexpr (IsGroupedGemmKernel) { - // Group GEMM currently only supports rank-3 problem shapes - implementable &= (args.mode == GemmUniversalMode::kGrouped && - rank(typename ProblemShape::UnderlyingProblemShape{}) == 3); - } else { - implementable &= (args.mode == GemmUniversalMode::kArray && - rank(typename ProblemShape::UnderlyingProblemShape{}) == 4); - } - if (!implementable) { - CUTLASS_TRACE_HOST( - " CAN IMPLEMENT: Arguments or Problem Shape don't meet the requirements for Ptr Array " - "Gemm or Grouped Gemm.\n"); - return implementable; - } - implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop); - implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue); - implementable &= TileScheduler::can_implement(args.scheduler); - return implementable; - } - - static size_t get_workspace_size(Arguments const& args) { - size_t workspace_size = 0; - constexpr uint32_t NumEpilogueSubTiles = - CollectiveEpilogue::get_store_pipe_increment(TileShape{}); - - // Get SM count if needed, otherwise use user supplied SM count - int sm_count = args.hw_info.sm_count; - if (sm_count <= 0) { - CUTLASS_TRACE_HOST( - " WARNING: Arguments do not include a valid SM count.\n" - " For optimal performance, populate the arguments KernelHardwareInfo struct with the SM " - "count."); - sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id); - } - - workspace_size += - CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue, sm_count); - workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment); - - workspace_size += - CollectiveMainloop::get_workspace_size(args.problem_shape, args.mainloop, sm_count); - workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment); - - workspace_size += - TileScheduler::template get_workspace_size( - args.scheduler, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, - NumMmaWarpGroups, NumEpilogueSubTiles); - workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment); - - return workspace_size; - } - - static cutlass::Status initialize_workspace(Arguments const& args, void* workspace = nullptr, - cudaStream_t stream = nullptr, - CudaHostAdapter* cuda_adapter = nullptr) { - Status status = Status::kSuccess; - uint8_t* workspace_ptr = reinterpret_cast(workspace); - size_t workspace_offset = 0; - constexpr uint32_t NumEpilogueSubTiles = - CollectiveEpilogue::get_store_pipe_increment(TileShape{}); - static constexpr uint32_t NumAccumulatorMtxs = 1; - - status = CollectiveEpilogue::initialize_workspace( - args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream, cuda_adapter); - workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue, - args.hw_info.sm_count); - workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); - if (status != Status::kSuccess) { - return status; - } - - status = CollectiveMainloop::initialize_workspace( - args.problem_shape, args.mainloop, workspace_ptr + workspace_offset, stream, cuda_adapter); - workspace_offset += CollectiveMainloop::get_workspace_size(args.problem_shape, args.mainloop, - args.hw_info.sm_count); - workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); - if (status != Status::kSuccess) { - return status; - } - - status = - TileScheduler::template initialize_workspace( - args.scheduler, workspace_ptr + workspace_offset, stream, - typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, - NumEpilogueSubTiles, NumAccumulatorMtxs, cuda_adapter); - workspace_offset += - TileScheduler::template get_workspace_size( - args.scheduler, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, - NumMmaWarpGroups, NumEpilogueSubTiles); - workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); - if (status != Status::kSuccess) { - return status; - } - - return status; - } - - // Computes the kernel launch grid shape based on runtime parameters - static dim3 get_grid_shape(Params const& params) { - // Given device SM count, set grid size s.t. we do not launch more thread blocks than we can run - // concurrently - TileSchedulerArguments args{}; - args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN - ? TileScheduler::RasterOrderOptions::AlongN - : TileScheduler::RasterOrderOptions::AlongM; - dim3 grid_shape; - if constexpr (IsGroupedGemmKernel) { - grid_shape = TileScheduler::get_grid_shape(params.scheduler, params.problem_shape, - TileShape{}, ClusterShape{}, params.hw_info, args); - } else { - grid_shape = TileScheduler::get_grid_shape(params.scheduler, - params.problem_shape.get_host_problem_shape(), - TileShape{}, ClusterShape{}, params.hw_info, args); - } - return grid_shape; - } - - static dim3 get_block_shape() { return dim3(MaxThreadsPerBlock, 1, 1); } - - CUTLASS_DEVICE - void operator()(Params const& params, char* smem_buf) { - using namespace cute; - using X = Underscore; - -// Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a. -#if !defined(__CUDA_ARCH_FEAT_SM90_ALL) - printf( - "ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. " - "Aborting.\n"); -#else - - // Preconditions - static_assert(size(TiledMma{}) == 256, - "Cooperative kernel must have TiledMMA operating using 256 threads."); - static_assert(size<0>(TileShape{}) >= 128, - "Cooperative kernel requires Tile Size to be greater than or equal to 128 along " - "the M-dimension."); - static_assert(NumMmaWarpGroups == 2, - "Cooperative kernels currently only support NumMmaWarpGroups == 2"); - - if constexpr (cutlass::epilogue::collective::detail::sm90_is_ptr_array_tma_dispatch_policy_v< - typename CollectiveEpilogue::DispatchPolicy>) { - static_assert(NumMmaWarpGroups == CollectiveEpilogue::NumEpilogueWarpGroups, - "Tiled MmA does not match expected warp groups performing the epilogue"); - } - - static_assert( - cute::rank(InternalStrideA{}) == 3, - "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>."); - static_assert( - cute::rank(InternalStrideB{}) == 3, - "StrideB must be rank-3: [N, K, L]. If batch mode is not needed, set L stride to Int<0>."); - static_assert( - cute::rank(InternalStrideC{}) == 3, - "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>."); - static_assert( - cute::rank(InternalStrideD{}) == 3, - "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>."); - - /* In the Cooperative kernel, Consumer0 and Consumer1 collaborate on the same tile */ - enum class WarpGroupRole { Producer = 0, Consumer0 = 1, Consumer1 = 2 }; - enum class ProducerWarpRole { Mainloop = 0, Warp1 = 1, Epilogue = 2, Warp3 = 3 }; - - // Kernel level shared memory storage - SharedStorage& shared_storage = *reinterpret_cast(smem_buf); - - int thread_idx = int(threadIdx.x); - int lane_idx = canonical_lane_idx(); - int warp_idx = canonical_warp_idx_sync(); - int warp_idx_in_warp_group = warp_idx % NumWarpsPerWarpGroup; - int warp_group_thread_idx = thread_idx % NumThreadsPerWarpGroup; - int mma_thread_idx = thread_idx % size(TiledMma{}); - auto warp_group_idx = canonical_warp_group_idx(); - auto warp_group_role = WarpGroupRole(warp_group_idx); - auto producer_warp_role = ProducerWarpRole(warp_idx_in_warp_group); - int lane_predicate = cute::elect_one_sync(); - uint32_t block_rank_in_cluster = cute::block_rank_in_cluster(); - - // Note: Tma Descriptor Prefetch (from either const or param) is not applicable here - - // Mainloop Load pipeline - using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline; - typename MainloopPipeline::Params mainloop_pipeline_params; - if (warp_group_role == WarpGroupRole::Producer && - producer_warp_role == ProducerWarpRole::Mainloop) { - mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer; - } - if (warp_group_role == WarpGroupRole::Consumer0 || - warp_group_role == WarpGroupRole::Consumer1) { - mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer; - } - mainloop_pipeline_params.is_leader = warp_group_thread_idx == 0; - mainloop_pipeline_params.num_consumers = NumMmaThreads; - mainloop_pipeline_params.num_producers = NumProducerThreads; - mainloop_pipeline_params.transaction_bytes = params.mainloop.tma_transaction_bytes; - MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params, - ClusterShape{}); - - // Epilogue Load pipeline - using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline; - typename EpiLoadPipeline::Params epi_load_pipeline_params; - if (warp_group_role == WarpGroupRole::Producer && - producer_warp_role == ProducerWarpRole::Epilogue) { - epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer; - } - if (warp_group_role == WarpGroupRole::Consumer0 || - warp_group_role == WarpGroupRole::Consumer1) { - epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer; - } - epi_load_pipeline_params.dst_blockid = cute::block_rank_in_cluster(); - epi_load_pipeline_params.producer_arv_count = NumThreadsPerWarp; - epi_load_pipeline_params.consumer_arv_count = size(TiledMma{}); - if constexpr (CollectiveEpilogue::RequiresTransactionBytes) { - epi_load_pipeline_params.transaction_bytes = params.epilogue.tma_transaction_bytes; - } - EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params); - - // Epilogue Store pipeline - using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline; - typename EpiStorePipeline::Params epi_store_pipeline_params; - epi_store_pipeline_params.always_wait = true; - EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params); - - typename LoadWarpOrderBarrier::Params params_load_order_barrier; - params_load_order_barrier.group_id = producer_warp_role == ProducerWarpRole::Mainloop ? 0 : 1; - params_load_order_barrier.group_size = NumThreadsPerWarp; - LoadWarpOrderBarrier load_order_barrier(shared_storage.pipelines.load_order, - params_load_order_barrier); - - // Initialize starting pipeline states for the collectives - // Epilogue store pipe is producer-only (consumer is TMA unit, waits via scoreboarding) - typename CollectiveMainloop::PipelineState mainloop_pipe_consumer_state; - typename CollectiveEpilogue::LoadPipelineState epi_load_pipe_consumer_state; - - // For the DMA Load (producer) we start with an opposite phase - // i.e., we skip all waits since we know that the buffer is indeed empty - PipelineState mainloop_pipe_producer_state = - cutlass::make_producer_start_state(); - PipelineState epi_load_pipe_producer_state = - cutlass::make_producer_start_state(); - PipelineState epi_store_pipe_producer_state = - cutlass::make_producer_start_state(); - - auto cluster_wait_fn = []() { - // We need this to guarantee that the Pipeline init is visible - // To all producers and consumer thread blocks in the Cluster - if constexpr (size(ClusterShape{}) > 1) { - cute::cluster_arrive_relaxed(); - return []() { cute::cluster_wait(); }; - } else { - __syncthreads(); - return []() {}; // do nothing - } - }(); - - // Get the appropriate blocks for this thread block -- potential for thread block locality - TiledMma tiled_mma; - const auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K) - const auto c_tile_count = CollectiveEpilogue::get_load_pipe_increment(blk_shape); - const auto d_tile_count = CollectiveEpilogue::get_store_pipe_increment(blk_shape); - - TileScheduler scheduler{params.scheduler}; - - // In a warp specialized kernel, collectives expose data movement and compute operations - // separately - CollectiveMainloop collective_mainloop; - CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue); - - // Wait for all thread blocks in the Cluster - cluster_wait_fn(); - - auto work_tile_info = scheduler.initial_work_tile_info(ClusterShape{}); - if (not work_tile_info.is_valid()) { - // When problem shapes are only on device, the grid launched may be larger than the total - // number of blocks across groups - return; - } - - // Optionally append 1s until problem shape is rank-4 in case it is only rank-3 (MNK) - auto problem_shape_MNKL = - append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), 1); - - // Prepare and partition the input tensors. Expects a tuple of tensors where: - // get<0>(load_inputs) is the tma tensor A after local tiling so that it has shape - // (BLK_M,BLK_K,m,k,l) get<1>(load_inputs) is the tma tensor B after local tiling so that it has - // shape (BLK_N,BLK_K,n,k,l) - auto load_inputs = collective_mainloop.load_init(problem_shape_MNKL, params.mainloop); - static_assert(cute::tuple_size_v >= 2, - "Output of load_init must have at least two elements (A, B)"); - - // Extract out partitioned A and B. - Tensor gA_mkl = get<0>(load_inputs); - Tensor gB_nkl = get<1>(load_inputs); - - // Get pipeline stage increments from tensor shapes - auto k_tile_count = size<3>(gA_mkl); - - if (warp_group_role == WarpGroupRole::Producer) { - cutlass::arch::warpgroup_reg_dealloc(); - - // Mainloop Producer Warp - if (producer_warp_role == ProducerWarpRole::Mainloop) { - int32_t curr_batch = idx2crd( - work_tile_info.L_idx, shape<4>(gB_nkl)); // Usually just returns work_tile_info.L_idx; - int32_t const mock_l_coord = 0; - int32_t const sm_idx = blockIdx.x + (blockIdx.y * gridDim.x); - int32_t const sm_count = params.hw_info.sm_count; - - // Fetch a copy of tensormaps for the CTA - auto input_tensormaps = collective_mainloop.tensormaps_init( - params.mainloop, shared_storage.tensormaps.mainloop, sm_count, sm_idx); - // Update tensormap for the initial batch for the CTA - if (work_tile_info.is_valid()) { - collective_mainloop.tensormaps_perform_update(shared_storage.tensormaps.mainloop, - params.mainloop, input_tensormaps, - problem_shape_MNKL, curr_batch); - // Ensure warp is converged before issuing tensormap fence release - __syncwarp(); - // Entire warp must do this (i.e. it's aligned) - collective_mainloop.tensormaps_cp_fence_release(shared_storage.tensormaps.mainloop, - input_tensormaps); - } - - bool do_load_order_arrive = true; - bool did_batch_change = true; - bool needs_tensormap_acquire = work_tile_info.is_valid(); - while (work_tile_info.is_valid()) { - if (!TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) { - auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info); - work_tile_info = next_work_tile_info; - continue; - } - - // Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape - auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl)); - auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl)); - auto blk_coord = make_coord(m_coord, n_coord, _, mock_l_coord); - - // Get the number of K tiles to compute for this work as well as the starting K tile - // offset of the work. - auto work_k_tile_count = - TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, blk_shape); - auto work_k_tile_start = TileScheduler::get_work_k_tile_start(work_tile_info); - auto k_tile_iter = cute::make_coord_iterator(idx2crd(work_k_tile_start, shape<3>(gA_mkl)), - shape<3>(gA_mkl)); - - if (did_batch_change) { - load_inputs = collective_mainloop.tensors_perform_update( - load_inputs, params.mainloop, problem_shape_MNKL, curr_batch); - if (needs_tensormap_acquire || - RequiresBatchTensormapAcquire::value) { - collective_mainloop.tensormaps_fence_acquire(input_tensormaps); - needs_tensormap_acquire = false; - } - } - - collective_mainloop.load(params.mainloop, mainloop_pipeline, mainloop_pipe_producer_state, - load_inputs, input_tensormaps, blk_coord, k_tile_iter, - work_k_tile_count, lane_idx, block_rank_in_cluster, - shared_storage.tensors.mainloop); - // Update starting pipeline state for the next tile - // Wait for the last TMA stage to complete loading, before issuing tensormap updates - mainloop_pipe_producer_state.advance(work_k_tile_count - 1); - - // Signal for the epilogue load warp to begin - if (do_load_order_arrive) { - load_order_barrier.arrive(); - do_load_order_arrive = false; - } - - // Get next work tile - auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info); - work_tile_info = next_work_tile_info; - auto next_batch = idx2crd(work_tile_info.L_idx, - shape<4>(gB_nkl)); // Usually just returns work_tile_info.L_idx - did_batch_change = next_batch != curr_batch; - if (work_tile_info.is_valid() && did_batch_change) { - curr_batch = next_batch; - if constexpr (IsGroupedGemmKernel) { - problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(curr_batch), 1); - } - if constexpr (RequiresBatchTensormapUpdate::value) { - // Purpose of this pipeline state is to make sure TMA loads have finished before doing - // descriptor updates Since this state is waiting for loads to finish, it must start - // in the inverted phase. - typename CollectiveMainloop::PipelineState mainloop_pipe_tma_consumer_state = { - mainloop_pipe_producer_state.index(), !mainloop_pipe_producer_state.phase(), - mainloop_pipe_producer_state.count()}; - mainloop_pipeline.consumer_wait(mainloop_pipe_tma_consumer_state); - collective_mainloop.tensormaps_perform_update(shared_storage.tensormaps.mainloop, - params.mainloop, input_tensormaps, - problem_shape_MNKL, curr_batch); - // Ensure warp is converged before issuing tensor replace - __syncwarp(); - // Entire warp must do this (i.e. it's aligned) - collective_mainloop.tensormaps_cp_fence_release(shared_storage.tensormaps.mainloop, - input_tensormaps); - needs_tensormap_acquire = true; - } - } - // Advance the producer state for the last remaining stage that was being waited for above - mainloop_pipe_producer_state.advance(1); - } // Scheduler work fetch loop - - // Make sure all Consumer Warp Groups have been waited upon - collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state); - } // Mainloop Producer Warp End - - // Epilogue Producer Warp - else if (producer_warp_role == ProducerWarpRole::Epilogue && - collective_epilogue.is_producer_load_needed()) { - int32_t const sm_idx = blockIdx.x + (blockIdx.y * gridDim.x); - int32_t const sm_count = params.hw_info.sm_count; - - auto epi_load_tensormap = get<0>(collective_epilogue.load_init( - params.epilogue, shared_storage.tensormaps.epilogue, sm_count, sm_idx)); - - bool did_batch_change = true; - constexpr bool IsEpiLoad = true; - - if (work_tile_info.is_valid()) { - collective_epilogue.template tensormaps_perform_update( - shared_storage.tensormaps.epilogue, params.epilogue, epi_load_tensormap, - problem_shape_MNKL, work_tile_info.L_idx, 0); - - // Converge before issuing tensormap fence release since fence is aligned - __syncwarp(); - collective_epilogue.template tensormaps_cp_fence_release( - shared_storage.tensormaps.epilogue, epi_load_tensormap, 0); - } - - load_order_barrier.wait(); - - while (work_tile_info.is_valid()) { - int32_t curr_batch = work_tile_info.L_idx; - - // Get next work tile - auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info); - - if (TileScheduler::compute_epilogue(work_tile_info, params.scheduler)) { - if constexpr (IsGroupedGemmKernel) { - problem_shape_MNKL = - append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), 1); - } - - // Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape - auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl)); - auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl)); - auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl)); - auto blk_coord = make_coord(m_coord, n_coord, _, l_coord); - - if (did_batch_change) { - collective_epilogue.template tensormaps_fence_acquire(epi_load_tensormap); - } - - bool wait = work_tile_info.is_valid() && curr_batch != next_work_tile_info.L_idx; - - epi_load_pipe_producer_state = collective_epilogue.load( - epi_load_pipeline, epi_load_pipe_producer_state, problem_shape_MNKL, blk_shape, - blk_coord, tiled_mma, lane_idx, shared_storage.tensors.epilogue, epi_load_tensormap, - work_tile_info.reduction_subtile_idx(), wait); - } - - work_tile_info = next_work_tile_info; - did_batch_change = curr_batch != work_tile_info.L_idx; - - if (work_tile_info.is_valid() && did_batch_change) { - if constexpr (IsGroupedGemmKernel) { - problem_shape_MNKL = - append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), 1); - } - - // tensormap update - { - collective_epilogue.template tensormaps_perform_update( - shared_storage.tensormaps.epilogue, params.epilogue, epi_load_tensormap, - problem_shape_MNKL, work_tile_info.L_idx, 0); - - // Converge before issuing tensormap fence release since fence is aligned - __syncwarp(); - collective_epilogue.template tensormaps_cp_fence_release( - shared_storage.tensormaps.epilogue, epi_load_tensormap, 0); - } - } - - } // Scheduler work fetch loop - - // Make sure all Consumer Warp Groups have been waited upon - collective_epilogue.load_tail(epi_load_pipeline, epi_load_pipe_producer_state); - } // Epilogue Producer Warp End - } // Producer Warp Group End - - else if (warp_group_role == WarpGroupRole::Consumer0 || - warp_group_role == WarpGroupRole::Consumer1) { - cutlass::arch::warpgroup_reg_alloc(); - - // Index of warp group within consumer warp groups - int consumer_warp_group_idx = warp_group_role == WarpGroupRole::Consumer0 ? 0 : 1; - - int32_t const sm_idx = blockIdx.x + (blockIdx.y * gridDim.x); - int32_t const sm_count = params.hw_info.sm_count; - // Do we potentially issue tail arrives for TMA stores, if epilogue load is waiting for it - bool do_store_tail = false; - // Get a copy of tensormaps - auto epi_store_tensormap = - get<0>(collective_epilogue.store_init(params.epilogue, shared_storage.tensormaps.epilogue, - sm_count, sm_idx, consumer_warp_group_idx)); - - bool did_batch_change = true; - constexpr bool IsEpiLoad = false; - - if (work_tile_info.is_valid()) { - if (warp_idx_in_warp_group == 0) { - collective_epilogue.template tensormaps_perform_update( - shared_storage.tensormaps.epilogue, params.epilogue, epi_store_tensormap, - problem_shape_MNKL, work_tile_info.L_idx, consumer_warp_group_idx); - - // Converge before issuing tensormap fence release since fence is aligned - __syncwarp(); - collective_epilogue.template tensormaps_cp_fence_release( - shared_storage.tensormaps.epilogue, epi_store_tensormap, consumer_warp_group_idx); - } - } - - while (work_tile_info.is_valid()) { - if constexpr (IsGroupedGemmKernel) { - problem_shape_MNKL = - append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), 1); - } - - int32_t curr_batch = work_tile_info.L_idx; - - // Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape - auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl)); - auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl)); - auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl)); - auto blk_coord = make_coord(m_coord, n_coord, _, l_coord); - auto work_k_tile_count = - TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, blk_shape); - - // Allocate the accumulators for the (M,N) blk_shape - // - // MSVC CTAD breaks if we say "Tensor" here, so we use "auto" instead. - auto accumulators = - partition_fragment_C(tiled_mma, take<0, 2>(blk_shape)); // (MMA,MMA_M,MMA_N) - - static_assert( - cute::is_same_v>); - if (TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) { - collective_mainloop.mma(mainloop_pipeline, mainloop_pipe_consumer_state, accumulators, - work_k_tile_count, mma_thread_idx, - shared_storage.tensors.mainloop, params.mainloop); - - // Make sure the math instructions are done and free buffers before entering the epilogue - collective_mainloop.mma_tail(mainloop_pipeline, mainloop_pipe_consumer_state, - work_k_tile_count); - - // Update starting mainloop pipeline state for the next tile - mainloop_pipe_consumer_state.advance(work_k_tile_count); - } - - // Perform reduction across splits, if needed - TileScheduler::fixup(params.scheduler, work_tile_info, accumulators, NumMmaWarpGroups, - consumer_warp_group_idx); - - if (did_batch_change) { - collective_epilogue.template tensormaps_fence_acquire(epi_store_tensormap); - } - - if (TileScheduler::compute_epilogue(work_tile_info, params.scheduler)) { - // Epilogue and write to gD - auto [epi_load_pipe_consumer_state_next, epi_store_pipe_producer_state_next] = - collective_epilogue.store( - epi_load_pipeline, epi_load_pipe_consumer_state, epi_store_pipeline, - epi_store_pipe_producer_state, problem_shape_MNKL, blk_shape, blk_coord, - accumulators, tiled_mma, mma_thread_idx, shared_storage.tensors.epilogue, - epi_store_tensormap, work_tile_info.reduction_subtile_idx()); - - epi_load_pipe_consumer_state = epi_load_pipe_consumer_state_next; - epi_store_pipe_producer_state = epi_store_pipe_producer_state_next; - do_store_tail = true; - } - - // Get next work tile - auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info); - work_tile_info = next_work_tile_info; - - did_batch_change = curr_batch != work_tile_info.L_idx; - if (work_tile_info.is_valid() && did_batch_change) { - if constexpr (IsGroupedGemmKernel) { - problem_shape_MNKL = - append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), 1); - } - if (warp_idx_in_warp_group == 0) { - collective_epilogue.template tensormaps_perform_update( - shared_storage.tensormaps.epilogue, params.epilogue, epi_store_tensormap, - problem_shape_MNKL, work_tile_info.L_idx, consumer_warp_group_idx); - - // Converge before issuing tensormap fence release since fence is aligned - __syncwarp(); - collective_epilogue.template tensormaps_cp_fence_release( - shared_storage.tensormaps.epilogue, epi_store_tensormap, consumer_warp_group_idx); - } - } - - } // Scheduler work fetch loop - - // Cooperative only needs TMA to complete at the very end of the kernel - if (do_store_tail) { - collective_epilogue.store_tail(epi_load_pipeline, epi_load_pipe_consumer_state, - epi_store_pipeline, epi_store_pipe_producer_state); - } - } // Consumer Warp Groups End -#endif - } -}; - -/////////////////////////////////////////////////////////////////////////////// - -} // namespace cutlass::gemm::kernel diff --git a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong_precomputed.hpp b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong_precomputed.hpp deleted file mode 100644 index 09138b78883..00000000000 --- a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong_precomputed.hpp +++ /dev/null @@ -1,998 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ -#pragma once - -#include - -#include "cute/arch/cluster_sm90.hpp" -#include "cute/tensor.hpp" -#include "cutlass/arch/mma_sm90.h" -#include "cutlass/arch/reg_reconfig.h" -#include "cutlass/cutlass.h" -#include "cutlass/epilogue/collective/detail.hpp" -#include "cutlass/fast_math.h" -#include "cutlass/gemm/dispatch_policy.hpp" -#include "cutlass/gemm/gemm.h" -#include "cutlass/gemm/group_array_problem_shape.hpp" -#include "cutlass/gemm/kernel/gemm_universal_decl.h" -#include "cutlass/gemm/kernel/sm90_tile_scheduler.hpp" -#include "cutlass/gemm/kernel/tile_scheduler.hpp" -#include "cutlass/kernel_hardware_info.hpp" -#include "cutlass/pipeline/pipeline.hpp" -#include "cutlass/trace.h" -#include "cutlass/workspace.h" -#include "cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_warpspecialized_precomputed_decl.hpp" -#include "cutlass_extensions/gemm/kernel/sm90_tile_scheduler_group_precomputed.hpp" - -/////////////////////////////////////////////////////////////////////////////// - -namespace cutlass::gemm::kernel { - -/////////////////////////////////////////////////////////////////////////////// - -#ifndef CUTLASS_EXTENSIONS_PRECOMPUTED_TENSORMAP_TRAITS_HPP_ -#define CUTLASS_EXTENSIONS_PRECOMPUTED_TENSORMAP_TRAITS_HPP_ - -template -struct RequiresBatchTensormapUpdate { - static constexpr bool value = true; -}; - -template -struct RequiresBatchTensormapUpdate< - CollectiveMainloop, - std::void_t> { - static constexpr bool value = CollectiveMainloop::RequiresTensormapUpdateOnBatchChange; -}; - -template -struct RequiresBatchTensormapAcquire { - static constexpr bool value = false; -}; - -template -struct RequiresBatchTensormapAcquire< - CollectiveMainloop, - std::void_t> { - static constexpr bool value = CollectiveMainloop::RequiresPrebuiltTensormapAcquireOnBatchChange; -}; - -#endif - -/////////////////////////////////////////////////////////////////////////////// - -template -class GemmUniversalPrecomputedScheduler< - ProblemShape_, CollectiveMainloop_, CollectiveEpilogue_, TileScheduler_, - cute::enable_if_t>> { - public: - // - // Type Aliases - // - using ProblemShape = ProblemShape_; - static_assert(rank(typename ProblemShape::UnderlyingProblemShape{}) == 3 or - rank(typename ProblemShape::UnderlyingProblemShape{}) == 4, - "ProblemShape{} should be or "); - - static_assert(cute::is_base_of_v); - - static constexpr bool IsGdcEnabled = false; - - // Mainloop derived types - using CollectiveMainloop = CollectiveMainloop_; - using TileShape = typename CollectiveMainloop::TileShape; - using TiledMma = typename CollectiveMainloop::TiledMma; - using ArchTag = typename CollectiveMainloop::ArchTag; - using ElementA = typename CollectiveMainloop::ElementA; - using StrideA = typename CollectiveMainloop::StrideA; - using InternalStrideA = typename CollectiveMainloop::InternalStrideA; - using ElementB = typename CollectiveMainloop::ElementB; - using InternalStrideB = typename CollectiveMainloop::InternalStrideB; - using StrideB = typename CollectiveMainloop::StrideB; - using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy; - using Schedule = typename DispatchPolicy::Schedule; - using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator; - using ClusterShape = typename DispatchPolicy::ClusterShape; - using MainloopArguments = typename CollectiveMainloop::Arguments; - using MainloopParams = typename CollectiveMainloop::Params; - - // Epilogue derived types - using CollectiveEpilogue = CollectiveEpilogue_; - using ElementC = typename CollectiveEpilogue::ElementC; - using StrideC = typename CollectiveEpilogue::StrideC; - using InternalStrideC = typename CollectiveEpilogue::InternalStrideC; - using ElementD = typename CollectiveEpilogue::ElementD; - using StrideD = typename CollectiveEpilogue::StrideD; - using InternalStrideD = typename CollectiveEpilogue::InternalStrideD; - using EpilogueArguments = typename CollectiveEpilogue::Arguments; - using EpilogueParams = typename CollectiveEpilogue::Params; - - static_assert(ArchTag::kMinComputeCapability >= 90); - - static constexpr bool IsGroupedGemmKernel = !cute::is_same_v; - - static_assert(IsGroupedGemmKernel, - "Precomputed grouped scheduler kernel is only for grouped ptr-array GEMM."); - - using SchedulerTag = GroupScheduler; - using TileScheduler = - cute::conditional_t, - detail::PersistentTileSchedulerSm90GroupPrecomputed, - TileScheduler_>; - - using TileSchedulerArguments = typename TileScheduler::Arguments; - using TileSchedulerParams = typename TileScheduler::Params; - - static constexpr uint32_t NumLoadWarpGroups = 1; - static constexpr uint32_t NumMmaWarpGroups = 2; - static constexpr uint32_t MaxThreadsPerBlock = - CUTE_STATIC_V(size(TiledMma{})) + (NumMmaWarpGroups * NumThreadsPerWarpGroup); - static constexpr uint32_t MinBlocksPerMultiprocessor = 1; - static constexpr uint32_t NumProducerThreads = CollectiveMainloop::NumProducerThreadEvents; - - /// Register requirement for Load and Math WGs - static constexpr uint32_t LoadRegisterRequirement = 40; - static constexpr uint32_t MmaRegisterRequirement = 232; - - // 1 stage ordered sequence between mainloop and epilogue producer load threads - using LoadWarpOrderBarrier = cutlass::OrderedSequenceBarrier<1, 2>; - - // Order Sequence barrier with two stages: one for Mainloop and one for Epilogue - static constexpr uint32_t StagesPerMathWarpGroup = 2; - using MathWarpGroupOrderBarrier = - cutlass::OrderedSequenceBarrier; - using MathWarpGroupOrderBarrierSharedStorage = - cutlass::PipelineDetail::OrderedSequenceBarrierSharedStorage< - MathWarpGroupOrderBarrier::SequenceDepth, MathWarpGroupOrderBarrier::SequenceLength>; - - // Kernel level shared memory storage - struct SharedStorage { - struct TensorStorage : cute::aligned_struct<128, _1> { - using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage; - using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage; - - MainloopTensorStorage mainloop; - EpilogueTensorStorage epilogue; - } tensors; - - struct PipelineStorage : cute::aligned_struct<16, _1> { - using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage; - using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage; - using MathWarpGroupOrderBarrierStorage = MathWarpGroupOrderBarrierSharedStorage; - - alignas(16) MainloopPipelineStorage mainloop; - alignas(16) EpiLoadPipelineStorage epi_load; - alignas(16) typename LoadWarpOrderBarrier::SharedStorage load_order; - alignas(16) MathWarpGroupOrderBarrierStorage math_wg_order; - } pipelines; - - struct TensorMapStorage : cute::aligned_struct<128, _1> { - using MainloopTensorMapStorage = typename CollectiveMainloop::TensorMapStorage; - using EpilogueTensorMapStorage = typename CollectiveEpilogue::TensorMapStorage; - - alignas(128) MainloopTensorMapStorage mainloop; - alignas(128) EpilogueTensorMapStorage epilogue; - } tensormaps; - }; - - static constexpr int SharedStorageSize = sizeof(SharedStorage); - - // Device side arguments - struct Arguments { - GemmUniversalMode mode{}; - ProblemShape problem_shape{}; - MainloopArguments mainloop{}; - EpilogueArguments epilogue{}; - KernelHardwareInfo hw_info{}; - TileSchedulerArguments scheduler{}; - }; - - // Kernel entry point API - struct Params { - GemmUniversalMode mode{}; - ProblemShape problem_shape{}; - MainloopParams mainloop{}; - EpilogueParams epilogue{}; - KernelHardwareInfo hw_info{}; - TileSchedulerParams scheduler{}; - void* workspace{nullptr}; - }; - - // - // Methods - // - - // Convert to underlying arguments. In this case, a simple copy for the aliased type. - static Params to_underlying_arguments(Arguments const& args, void* workspace) { - CUTLASS_TRACE_HOST("to_underlying_arguments():"); - - ProblemShape problem_shapes = args.problem_shape; - - // Get SM count if needed, otherwise use user supplied SM count - int sm_count = args.hw_info.sm_count; - if (sm_count <= 0) { - CUTLASS_TRACE_HOST( - " WARNING: Arguments do not include a valid SM count.\n" - " For optimal performance, populate the arguments KernelHardwareInfo struct with the SM " - "count."); - sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id); - } - CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " - << sm_count); - - // Get maximum number of clusters that could co-exist on the target device - int max_active_clusters = args.hw_info.max_active_clusters; - if (max_active_clusters <= 0) { - max_active_clusters = 0; - CUTLASS_TRACE_HOST( - " WARNING: Arguments do not include a valid max cluster count.\n" - " For optimal performance, populate the arguments KernelHardwareInfo struct with the " - "max_active_clusters."); - } else { - CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid cluster count to " - << max_active_clusters); - } - - KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count, max_active_clusters}; - - // Calculate workspace pointers - uint8_t* workspace_ptr = reinterpret_cast(workspace); - size_t workspace_offset = 0; - - void* epilogue_workspace = workspace_ptr + workspace_offset; - workspace_offset += - CollectiveEpilogue::get_workspace_size(problem_shapes, args.epilogue, sm_count); - workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); - - void* mainloop_workspace = workspace_ptr + workspace_offset; - workspace_offset += - CollectiveMainloop::get_workspace_size(problem_shapes, args.mainloop, sm_count); - workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); - - void* scheduler_workspace = workspace_ptr + workspace_offset; - workspace_offset += - TileScheduler::template get_workspace_size( - args.scheduler, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, - NumMmaWarpGroups); - workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); - - // Precompute the sub tiles numbers in epilogue, pass into tile scheduler. Therefore it will be - // used in separate reduction scheme for streamk case, NumEpilogueSubTiles default value is 1, - // which means subtile will not be used, therefore separate reduction will not be enabled. - constexpr uint32_t NumEpilogueSubTiles = - CollectiveEpilogue::get_store_pipe_increment(TileShape{}); - TileSchedulerParams scheduler; - if constexpr (IsGroupedGemmKernel) { - scheduler = TileScheduler::to_underlying_arguments(problem_shapes, TileShape{}, - ClusterShape{}, hw_info, args.scheduler, - scheduler_workspace, NumEpilogueSubTiles); - } else { - scheduler = TileScheduler::to_underlying_arguments( - problem_shapes.get_host_problem_shape(), TileShape{}, ClusterShape{}, hw_info, - args.scheduler, scheduler_workspace, NumEpilogueSubTiles); - } - - return {args.mode, - problem_shapes, - CollectiveMainloop::to_underlying_arguments(problem_shapes, args.mainloop, - mainloop_workspace), - CollectiveEpilogue::to_underlying_arguments(problem_shapes, args.epilogue, - epilogue_workspace), - hw_info, - scheduler, - workspace}; - } - - static bool can_implement(Arguments const& args) { - bool implementable = true; - if constexpr (IsGroupedGemmKernel) { - // Group GEMM currently only supports rank-3 problem shapes - implementable &= (args.mode == GemmUniversalMode::kGrouped && - rank(typename ProblemShape::UnderlyingProblemShape{}) == 3); - } else { - implementable &= (args.mode == GemmUniversalMode::kArray && - rank(typename ProblemShape::UnderlyingProblemShape{}) == 4); - } - if (!implementable) { - CUTLASS_TRACE_HOST( - " CAN IMPLEMENT: Arguments or Problem Shape don't meet the requirements for Ptr Array " - "Gemm or Grouped Gemm.\n"); - return implementable; - } - implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop); - implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue); - implementable &= TileScheduler::can_implement(args.scheduler); - return implementable; - } - - static size_t get_workspace_size(Arguments const& args) { - size_t workspace_size = 0; - constexpr uint32_t NumEpilogueSubTiles = - CollectiveEpilogue::get_store_pipe_increment(TileShape{}); - - // Get SM count if needed, otherwise use user supplied SM count - int sm_count = args.hw_info.sm_count; - if (sm_count <= 0) { - CUTLASS_TRACE_HOST( - " WARNING: Arguments do not include a valid SM count.\n" - " For optimal performance, populate the arguments KernelHardwareInfo struct with the SM " - "count."); - sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id); - } - - workspace_size += - CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue, sm_count); - workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment); - - workspace_size += - CollectiveMainloop::get_workspace_size(args.problem_shape, args.mainloop, sm_count); - workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment); - - workspace_size += - TileScheduler::template get_workspace_size( - args.scheduler, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, - NumMmaWarpGroups, NumEpilogueSubTiles); - workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment); - - return workspace_size; - } - - static cutlass::Status initialize_workspace(Arguments const& args, void* workspace = nullptr, - cudaStream_t stream = nullptr, - CudaHostAdapter* cuda_adapter = nullptr) { - Status status = Status::kSuccess; - uint8_t* workspace_ptr = reinterpret_cast(workspace); - size_t workspace_offset = 0; - constexpr uint32_t NumEpilogueSubTiles = - CollectiveEpilogue::get_store_pipe_increment(TileShape{}); - static constexpr uint32_t NumAccumulatorMtxs = 1; - - status = CollectiveEpilogue::initialize_workspace( - args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream, cuda_adapter); - workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue, - args.hw_info.sm_count); - workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); - if (status != Status::kSuccess) { - return status; - } - - status = CollectiveMainloop::initialize_workspace( - args.problem_shape, args.mainloop, workspace_ptr + workspace_offset, stream, cuda_adapter); - workspace_offset += CollectiveMainloop::get_workspace_size(args.problem_shape, args.mainloop, - args.hw_info.sm_count); - workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); - if (status != Status::kSuccess) { - return status; - } - - status = - TileScheduler::template initialize_workspace( - args.scheduler, workspace_ptr + workspace_offset, stream, - typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, - NumEpilogueSubTiles, NumAccumulatorMtxs, cuda_adapter); - workspace_offset += - TileScheduler::template get_workspace_size( - args.scheduler, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, - NumMmaWarpGroups, NumEpilogueSubTiles); - workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); - if (status != Status::kSuccess) { - return status; - } - return status; - } - - // Computes the kernel launch grid shape based on runtime parameters - static dim3 get_grid_shape(Params const& params) { - // Given device SM count, set grid size s.t. we do not launch more thread blocks than we can run - // concurrently - TileSchedulerArguments args{}; - args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN - ? TileScheduler::RasterOrderOptions::AlongN - : TileScheduler::RasterOrderOptions::AlongM; - dim3 grid_shape; - if constexpr (IsGroupedGemmKernel) { - grid_shape = TileScheduler::get_grid_shape(params.scheduler, params.problem_shape, - TileShape{}, ClusterShape{}, params.hw_info, args); - } else { - grid_shape = TileScheduler::get_grid_shape(params.scheduler, - params.problem_shape.get_host_problem_shape(), - TileShape{}, ClusterShape{}, params.hw_info, args); - } - return grid_shape; - } - - static dim3 get_block_shape() { return dim3(MaxThreadsPerBlock, 1, 1); } - - CUTLASS_DEVICE - void operator()(Params const& params, char* smem_buf) { - using namespace cute; - using X = Underscore; - -// Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a. -#if !defined(__CUDA_ARCH_FEAT_SM90_ALL) - printf( - "ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. " - "Aborting.\n"); -#else - - // Preconditions - static_assert(size(TiledMma{}) == 128, - "Pingpong kernel must have TiledMMA operating using 128 threads."); - static_assert(NumMmaWarpGroups == 2, - "Pingpong kernels currently only support NumMmaWarpGroups == 2"); - - if constexpr (cutlass::epilogue::collective::detail::sm90_is_ptr_array_tma_dispatch_policy_v< - typename CollectiveEpilogue::DispatchPolicy>) { - static_assert(NumMmaWarpGroups == CollectiveEpilogue::NumEpilogueWarpGroups, - "Tiled MmA does not match expected warp groups performing the epilogue"); - } - - static_assert( - cute::rank(InternalStrideA{}) == 3, - "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>."); - static_assert( - cute::rank(InternalStrideB{}) == 3, - "StrideB must be rank-3: [N, K, L]. If batch mode is not needed, set L stride to Int<0>."); - static_assert( - cute::rank(InternalStrideC{}) == 3, - "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>."); - static_assert( - cute::rank(InternalStrideD{}) == 3, - "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>."); - - enum class WarpGroupRole { Producer = 0, Consumer0 = 1, Consumer1 = 2 }; - enum class ProducerWarpRole { Mainloop = 0, Warp1 = 1, Epilogue = 2, Warp3 = 3 }; - - // Kernel level shared memory storage - SharedStorage& shared_storage = *reinterpret_cast(smem_buf); - - int thread_idx = int(threadIdx.x); - int lane_idx = canonical_lane_idx(); - int warp_idx = canonical_warp_idx_sync(); - int warp_idx_in_warp_group = warp_idx % NumWarpsPerWarpGroup; - int warp_group_thread_idx = thread_idx % NumThreadsPerWarpGroup; - int mma_thread_idx = thread_idx % size(TiledMma{}); - auto warp_group_idx = canonical_warp_group_idx(); - auto warp_group_role = WarpGroupRole(warp_group_idx); - auto producer_warp_role = ProducerWarpRole(warp_idx_in_warp_group); - int lane_predicate = cute::elect_one_sync(); - uint32_t block_rank_in_cluster = cute::block_rank_in_cluster(); - - // Note: Tma Descriptor Prefetch (from either const or param) is not applicable here - - // Mainloop Load pipeline - using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline; - typename MainloopPipeline::Params mainloop_pipeline_params; - if (warp_group_role == WarpGroupRole::Producer && - producer_warp_role == ProducerWarpRole::Mainloop) { - mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer; - } - if (warp_group_role == WarpGroupRole::Consumer0 || - warp_group_role == WarpGroupRole::Consumer1) { - mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer; - } - mainloop_pipeline_params.is_leader = warp_group_thread_idx == 0; - mainloop_pipeline_params.num_consumers = NumThreadsPerWarpGroup; - mainloop_pipeline_params.num_producers = NumProducerThreads; - mainloop_pipeline_params.transaction_bytes = params.mainloop.tma_transaction_bytes; - MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params, - ClusterShape{}); - - // Epilogue Load pipeline - using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline; - typename EpiLoadPipeline::Params epi_load_pipeline_params; - if (warp_group_role == WarpGroupRole::Producer && - producer_warp_role == ProducerWarpRole::Epilogue) { - epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer; - } - if (warp_group_role == WarpGroupRole::Consumer0 || - warp_group_role == WarpGroupRole::Consumer1) { - epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer; - } - epi_load_pipeline_params.dst_blockid = cute::block_rank_in_cluster(); - epi_load_pipeline_params.producer_arv_count = NumThreadsPerWarp; - epi_load_pipeline_params.consumer_arv_count = NumThreadsPerWarpGroup; - if constexpr (CollectiveEpilogue::RequiresTransactionBytes) { - epi_load_pipeline_params.transaction_bytes = params.epilogue.tma_transaction_bytes; - } - EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params); - - // Epilogue Store pipeline - using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline; - typename EpiStorePipeline::Params epi_store_pipeline_params; - epi_store_pipeline_params.always_wait = true; - EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params); - - typename LoadWarpOrderBarrier::Params params_load_order_barrier; - params_load_order_barrier.group_id = producer_warp_role == ProducerWarpRole::Mainloop ? 0 : 1; - params_load_order_barrier.group_size = NumThreadsPerWarp; - LoadWarpOrderBarrier load_order_barrier(shared_storage.pipelines.load_order, - params_load_order_barrier); - - typename MathWarpGroupOrderBarrier::Params params_math_wg_order_barrier; - // DMA Load WG will not participate in these Ordered Barrier syncs - params_math_wg_order_barrier.group_id = - warp_group_idx - static_cast(WarpGroupRole::Consumer0); - params_math_wg_order_barrier.group_size = - NumThreadsPerWarpGroup; // Number of threads / participants in a group - MathWarpGroupOrderBarrier math_wg_order_barrier(shared_storage.pipelines.math_wg_order, - params_math_wg_order_barrier); - - // Initialize starting pipeline states for the collectives - // Epilogue store pipe is producer-only (consumer is TMA unit, waits via scoreboarding) - typename CollectiveMainloop::PipelineState mainloop_pipe_consumer_state; - typename CollectiveEpilogue::LoadPipelineState epi_load_pipe_consumer_state; - - // For the DMA Load (producer) we start with an opposite phase - // i.e., we skip all waits since we know that the buffer is indeed empty - PipelineState mainloop_pipe_producer_state = - cutlass::make_producer_start_state(); - PipelineState epi_load_pipe_producer_state = - cutlass::make_producer_start_state(); - PipelineState epi_store_pipe_producer_state = - cutlass::make_producer_start_state(); - - auto cluster_wait_fn = []() { - // We need this to guarantee that the Pipeline init is visible - // To all producers and consumer thread blocks in the Cluster - if constexpr (size(ClusterShape{}) > 1) { - cute::cluster_arrive_relaxed(); - return []() { cute::cluster_wait(); }; - } else { - __syncthreads(); - return []() {}; // do nothing - } - }(); - - // Get the appropriate blocks for this thread block -- potential for thread block locality - TiledMma tiled_mma; - const auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K) - const auto c_tile_count = CollectiveEpilogue::get_load_pipe_increment(blk_shape); - const auto d_tile_count = CollectiveEpilogue::get_store_pipe_increment(blk_shape); - - TileScheduler scheduler{params.scheduler}; - - // In a warp specialized kernel, collectives expose data movement and compute operations - // separately - CollectiveMainloop collective_mainloop; - CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue); - - // Wait for all thread blocks in the Cluster - cluster_wait_fn(); - auto work_tile_info = scheduler.initial_work_tile_info(ClusterShape{}); - - if (not work_tile_info.is_valid()) { - // When problem shapes are only on device, the grid launched may be larger than the total - // number of blocks across groups - return; - } - - // Optionally append 1s until problem shape is rank-4 in case it is only rank-3 (MNK) - auto problem_shape_MNKL = - append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), 1); - - // Consumer1 is not on the critical path at prologue. - if (warp_group_role == WarpGroupRole::Consumer1) [[unlikely]] { - // Advance 2nd Math WG to the next work tile for the startup - const auto k_tile_count = - TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, blk_shape); - - auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info); - work_tile_info = next_work_tile_info; - if (!work_tile_info.is_valid()) { - return; - } - - // Advance 2nd Math WG pipeline states to the end of 1st Math WG - mainloop_pipe_consumer_state.advance(k_tile_count); - epi_load_pipe_consumer_state.advance(c_tile_count); - epi_store_pipe_producer_state.advance(d_tile_count); - - problem_shape_MNKL = - append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), 1); - } - - // Prepare and partition the input tensors. Expects a tuple of tensors where: - // get<0>(load_inputs) is the tma tensor A after local tiling so that it has shape - // (BLK_M,BLK_K,m,k,l) get<1>(load_inputs) is the tma tensor B after local tiling so that it has - // shape (BLK_N,BLK_K,n,k,l) - auto load_inputs = collective_mainloop.load_init(problem_shape_MNKL, params.mainloop); - static_assert(cute::tuple_size_v >= 2, - "Output of load_init must have at least two elements (A, B)"); - - // Extract out partitioned A and B. - Tensor gA_mkl = get<0>(load_inputs); - Tensor gB_nkl = get<1>(load_inputs); - - // Get pipeline stage increments from tensor shapes - auto k_tile_count = size<3>(gA_mkl); - - if (warp_group_role == WarpGroupRole::Producer) { - cutlass::arch::warpgroup_reg_dealloc(); - - // Mainloop Producer Warp - if (producer_warp_role == ProducerWarpRole::Mainloop) { - int32_t curr_batch = idx2crd( - work_tile_info.L_idx, shape<4>(gB_nkl)); // Usually just returns work_tile_info.L_idx; - int32_t const mock_l_coord = 0; - int32_t const sm_idx = blockIdx.x + (blockIdx.y * gridDim.x); - int32_t const sm_count = params.hw_info.sm_count; - - // Fetch a copy of tensormaps for the CTA - auto input_tensormaps = collective_mainloop.tensormaps_init( - params.mainloop, shared_storage.tensormaps.mainloop, sm_count, sm_idx); - - // Update tensormap for the initial batch for the CTA - collective_mainloop.tensormaps_perform_update(shared_storage.tensormaps.mainloop, - params.mainloop, input_tensormaps, - problem_shape_MNKL, curr_batch); - // Ensure warp is converged before issuing tensormap fence release - __syncwarp(); - // Entire warp must do this (i.e. it's aligned) - collective_mainloop.tensormaps_cp_fence_release(shared_storage.tensormaps.mainloop, - input_tensormaps); - - bool do_load_order_arrive = true; - bool did_batch_change = true; - bool needs_tensormap_acquire = work_tile_info.is_valid(); - while (work_tile_info.is_valid()) { - if (!TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) { - auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info); - work_tile_info = next_work_tile_info; - continue; - } - - // Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape - auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl)); - auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl)); - auto blk_coord = make_coord(m_coord, n_coord, _, mock_l_coord); - - // Get the number of K tiles to compute for this work as well as the starting K tile - // offset of the work. - auto work_k_tile_count = - TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, blk_shape); - auto work_k_tile_start = TileScheduler::get_work_k_tile_start(work_tile_info); - auto k_tile_iter = cute::make_coord_iterator(idx2crd(work_k_tile_start, shape<3>(gA_mkl)), - shape<3>(gA_mkl)); - - if (did_batch_change) { - load_inputs = collective_mainloop.tensors_perform_update( - load_inputs, params.mainloop, problem_shape_MNKL, curr_batch); - if (needs_tensormap_acquire || - RequiresBatchTensormapAcquire::value) { - collective_mainloop.tensormaps_fence_acquire(input_tensormaps); - needs_tensormap_acquire = false; - } - } - - collective_mainloop.load(params.mainloop, mainloop_pipeline, mainloop_pipe_producer_state, - load_inputs, input_tensormaps, blk_coord, k_tile_iter, - work_k_tile_count, lane_idx, block_rank_in_cluster, - shared_storage.tensors.mainloop); - // Update starting pipeline state for the next tile - // Wait for the last TMA stage to complete loading, before issuing tensormap updates - mainloop_pipe_producer_state.advance(work_k_tile_count - 1); - - // Signal for the epilogue load warp to begin - if (do_load_order_arrive) { - load_order_barrier.arrive(); - do_load_order_arrive = false; - } - - // Get next work tile - auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info); - work_tile_info = next_work_tile_info; - auto next_batch = idx2crd(work_tile_info.L_idx, - shape<4>(gB_nkl)); // Usually just returns work_tile_info.L_idx - did_batch_change = next_batch != curr_batch; - if (work_tile_info.is_valid() && did_batch_change) { - curr_batch = next_batch; - if constexpr (IsGroupedGemmKernel) { - problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(curr_batch), 1); - } - if constexpr (RequiresBatchTensormapUpdate::value) { - // Purpose of this pipeline state is to make sure TMA loads have finished before doing - // descriptor updates Since this state is waiting for loads to finish, it must start - // in the inverted phase. - typename CollectiveMainloop::PipelineState mainloop_pipe_tma_consumer_state = { - mainloop_pipe_producer_state.index(), !mainloop_pipe_producer_state.phase(), - mainloop_pipe_producer_state.count()}; - mainloop_pipeline.consumer_wait(mainloop_pipe_tma_consumer_state); - collective_mainloop.tensormaps_perform_update(shared_storage.tensormaps.mainloop, - params.mainloop, input_tensormaps, - problem_shape_MNKL, curr_batch); - // Ensure warp is converged before issuing tensor replace - __syncwarp(); - // Entire warp must do this (i.e. it's aligned) - collective_mainloop.tensormaps_cp_fence_release(shared_storage.tensormaps.mainloop, - input_tensormaps); - needs_tensormap_acquire = true; - } - } - // Advance the producer state for the last remaining stage that was being waited for above - mainloop_pipe_producer_state.advance(1); - } // Scheduler work fetch loop - - // Make sure all Consumer Warp Groups have been waited upon - collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state); - } // Mainloop Producer Warp End - // Epilogue Producer Warp - else if (producer_warp_role == ProducerWarpRole::Epilogue && - collective_epilogue.is_producer_load_needed()) { - int32_t const sm_idx = blockIdx.x + (blockIdx.y * gridDim.x); - int32_t const sm_count = params.hw_info.sm_count; - - auto epi_load_tensormap = get<0>(collective_epilogue.load_init( - params.epilogue, shared_storage.tensormaps.epilogue, sm_count, sm_idx)); - - bool did_batch_change = true; - constexpr bool IsEpiLoad = true; - - if (work_tile_info.is_valid()) { - collective_epilogue.template tensormaps_perform_update( - shared_storage.tensormaps.epilogue, params.epilogue, epi_load_tensormap, - problem_shape_MNKL, work_tile_info.L_idx, 0); - - // Converge before issuing tensormap fence release since fence is aligned - __syncwarp(); - collective_epilogue.template tensormaps_cp_fence_release( - shared_storage.tensormaps.epilogue, epi_load_tensormap, 0); - } - - load_order_barrier.wait(); - - while (work_tile_info.is_valid()) { - int32_t curr_batch = work_tile_info.L_idx; - - // Get next work tile - auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info); - - if (TileScheduler::compute_epilogue(work_tile_info, params.scheduler)) { - if constexpr (IsGroupedGemmKernel) { - problem_shape_MNKL = - append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), 1); - } - - // Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape - auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl)); - auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl)); - auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl)); - auto blk_coord = make_coord(m_coord, n_coord, _, l_coord); - - if (did_batch_change) { - collective_epilogue.template tensormaps_fence_acquire(epi_load_tensormap); - } - - bool wait = work_tile_info.is_valid() && curr_batch != next_work_tile_info.L_idx; - - epi_load_pipe_producer_state = collective_epilogue.load( - epi_load_pipeline, epi_load_pipe_producer_state, problem_shape_MNKL, blk_shape, - blk_coord, tiled_mma, lane_idx, shared_storage.tensors.epilogue, epi_load_tensormap, - work_tile_info.reduction_subtile_idx(), wait); - } - - work_tile_info = next_work_tile_info; - did_batch_change = curr_batch != work_tile_info.L_idx; - - if (work_tile_info.is_valid() && did_batch_change) { - if constexpr (IsGroupedGemmKernel) { - problem_shape_MNKL = - append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), 1); - } - - // tensormap update - { - collective_epilogue.template tensormaps_perform_update( - shared_storage.tensormaps.epilogue, params.epilogue, epi_load_tensormap, - problem_shape_MNKL, work_tile_info.L_idx, 0); - - // Converge before issuing tensormap fence release since fence is aligned - __syncwarp(); - collective_epilogue.template tensormaps_cp_fence_release( - shared_storage.tensormaps.epilogue, epi_load_tensormap, 0); - } - } - - } // Scheduler work fetch loop - - // Make sure all Consumer Warp Groups have been waited upon - collective_epilogue.load_tail(epi_load_pipeline, epi_load_pipe_producer_state); - } // Epilogue Producer Warp End - } // Producer Warp Group End - - else if (warp_group_role == WarpGroupRole::Consumer0 || - warp_group_role == WarpGroupRole::Consumer1) { - cutlass::arch::warpgroup_reg_alloc(); - - // Index of warp group within consumer warp groups - int consumer_warp_group_idx = warp_group_role == WarpGroupRole::Consumer0 ? 0 : 1; - - int32_t const sm_idx = blockIdx.x + (blockIdx.y * gridDim.x); - int32_t const sm_count = params.hw_info.sm_count; - // Do we potentially issue tail arrives for TMA stores, if epilogue load is waiting for it - bool do_store_tail = false; - // Get a copy of tensormaps - auto epi_store_tensormap = - get<0>(collective_epilogue.store_init(params.epilogue, shared_storage.tensormaps.epilogue, - sm_count, sm_idx, consumer_warp_group_idx)); - - bool did_batch_change = true; - constexpr bool IsEpiLoad = false; - - if (work_tile_info.is_valid()) { - if (warp_idx_in_warp_group == 0) { - collective_epilogue.template tensormaps_perform_update( - shared_storage.tensormaps.epilogue, params.epilogue, epi_store_tensormap, - problem_shape_MNKL, work_tile_info.L_idx, consumer_warp_group_idx); - - // Converge before issuing tensormap fence release since fence is aligned - __syncwarp(); - collective_epilogue.template tensormaps_cp_fence_release( - shared_storage.tensormaps.epilogue, epi_store_tensormap, consumer_warp_group_idx); - } - } - - while (work_tile_info.is_valid()) { - if constexpr (IsGroupedGemmKernel) { - problem_shape_MNKL = - append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), 1); - } - - int32_t curr_batch = work_tile_info.L_idx; - - // Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape - auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl)); - auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl)); - auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl)); - auto blk_coord = make_coord(m_coord, n_coord, _, l_coord); - auto work_k_tile_count = - TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, blk_shape); - - // Allocate the accumulators for the (M,N) blk_shape - // - // MSVC CTAD breaks if we say "Tensor" here, so we use "auto" instead. - auto accumulators = - partition_fragment_C(tiled_mma, take<0, 2>(blk_shape)); // (MMA,MMA_M,MMA_N) - - if (TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) { - math_wg_order_barrier.wait(); - - collective_mainloop.mma(mainloop_pipeline, mainloop_pipe_consumer_state, accumulators, - work_k_tile_count, mma_thread_idx, - shared_storage.tensors.mainloop, params.mainloop); - - math_wg_order_barrier.arrive(); - - // Make sure the math instructions are done and free buffers before entering the epilogue - collective_mainloop.mma_tail(mainloop_pipeline, mainloop_pipe_consumer_state, - work_k_tile_count); - - math_wg_order_barrier.wait(); - - // Update starting mainloop pipeline state for the next tile - mainloop_pipe_consumer_state.advance(work_k_tile_count); - } - - // Perform reduction across splits, if needed - TileScheduler::fixup(params.scheduler, work_tile_info, accumulators, NumMmaWarpGroups, - consumer_warp_group_idx); - - if (did_batch_change && warp_idx_in_warp_group == 0) { - collective_epilogue.template tensormaps_fence_acquire(epi_store_tensormap); - } - - if (TileScheduler::compute_epilogue(work_tile_info, params.scheduler)) { - // Epilogue and write to gD - auto [epi_load_pipe_consumer_state_next, epi_store_pipe_producer_state_next] = - collective_epilogue.store( - epi_load_pipeline, epi_load_pipe_consumer_state, epi_store_pipeline, - epi_store_pipe_producer_state, problem_shape_MNKL, blk_shape, blk_coord, - accumulators, tiled_mma, mma_thread_idx, shared_storage.tensors.epilogue, - epi_store_tensormap, work_tile_info.reduction_subtile_idx()); - - epi_load_pipe_consumer_state = epi_load_pipe_consumer_state_next; - epi_store_pipe_producer_state = epi_store_pipe_producer_state_next; - do_store_tail = true; - } - - // Get next work tile - auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info); - - work_tile_info = next_work_tile_info; - - // Skip a tile for pingpong - if (work_tile_info.is_valid()) { - if constexpr (IsGroupedGemmKernel) { - problem_shape_MNKL = - append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), 1); - } - work_k_tile_count = - TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, blk_shape); - mainloop_pipe_consumer_state.advance(work_k_tile_count); - - // Go to next tile - auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info); - work_tile_info = next_work_tile_info; - } - - did_batch_change = curr_batch != work_tile_info.L_idx; - if (work_tile_info.is_valid() && did_batch_change) { - if constexpr (IsGroupedGemmKernel) { - problem_shape_MNKL = - append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), 1); - } - if (warp_idx_in_warp_group == 0) { - collective_epilogue.template tensormaps_perform_update( - shared_storage.tensormaps.epilogue, params.epilogue, epi_store_tensormap, - problem_shape_MNKL, work_tile_info.L_idx, consumer_warp_group_idx); - - // Converge before issuing tensormap fence release since fence is aligned - __syncwarp(); - collective_epilogue.template tensormaps_cp_fence_release( - shared_storage.tensormaps.epilogue, epi_store_tensormap, consumer_warp_group_idx); - } - } - - // TMA store pipeline wait is only visible to TMA-issuing warp, so for multiple-consumer - // kernels we need to wait for all TMA stores to complete before issuing consumer order - // barrier arrives to ensure next math consumer doesn't overwrite smem of in-flight TMA - // stores of current consumer. - auto [epi_load_pipe_consumer_state_next_, epi_store_pipe_producer_state_next_] = - collective_epilogue.store_tail(epi_load_pipeline, epi_load_pipe_consumer_state, - epi_store_pipeline, epi_store_pipe_producer_state); - - // Update starting load/store pipeline states for the next tile - // state has already been incremented by 1 tile in collective calls, advance once again for - // ping pong - epi_load_pipe_consumer_state = epi_load_pipe_consumer_state_next_; - epi_store_pipe_producer_state = epi_store_pipe_producer_state_next_; - epi_load_pipe_consumer_state.advance(c_tile_count); - epi_store_pipe_producer_state.advance(d_tile_count); - - // Cue for next Math WG's Epilogue to start - math_wg_order_barrier.arrive(); - - } // Scheduler work fetch loop - } // Consumer Warp Groups End -#endif - } -}; - -/////////////////////////////////////////////////////////////////////////////// - -} // namespace cutlass::gemm::kernel diff --git a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_warpspecialized_precomputed_decl.hpp b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_warpspecialized_precomputed_decl.hpp deleted file mode 100644 index aa692bd4dc8..00000000000 --- a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_warpspecialized_precomputed_decl.hpp +++ /dev/null @@ -1,39 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ -#pragma once - -namespace cutlass::gemm::kernel { - -template -class GemmUniversalPrecomputedScheduler; - -} // namespace cutlass::gemm::kernel diff --git a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/sm90_tile_scheduler_group_precomputed.hpp b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/sm90_tile_scheduler_group_precomputed.hpp deleted file mode 100644 index a633ef0a8c3..00000000000 --- a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/sm90_tile_scheduler_group_precomputed.hpp +++ /dev/null @@ -1,514 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ -#pragma once - -#include "cute/arch/cluster_sm90.hpp" -#include "cute/layout.hpp" -#include "cute/tensor.hpp" -#include "cutlass/arch/barrier.h" -#include "cutlass/fast_math.h" -#include "cutlass/gemm/kernel/tile_scheduler_params.h" -#include "cutlass/gemm_coord.hpp" -#include "cutlass/kernel_hardware_info.hpp" -#include "cutlass/pipeline/pipeline.hpp" - -namespace cutlass::gemm::kernel::detail { - -/////////////////////////////////////////////////////////////////////////////// - -struct PrecomputedGroupWorkTile { - static constexpr uint64_t ChannelBits = 20; - static constexpr uint64_t TokenBits = 24; - static constexpr uint64_t ExpertBits = 19; - static constexpr uint64_t ChannelMask = (uint64_t(1) << ChannelBits) - 1; - static constexpr uint64_t TokenMask = (uint64_t(1) << TokenBits) - 1; - static constexpr uint64_t ExpertMask = (uint64_t(1) << ExpertBits) - 1; - static constexpr uint64_t TokenShift = ChannelBits; - static constexpr uint64_t ExpertShift = ChannelBits + TokenBits; - static constexpr uint64_t Invalid = ~uint64_t(0); - - CUTLASS_HOST_DEVICE - static bool fits(uint64_t channel_idx, uint64_t token_idx, uint64_t expert_idx) { - return channel_idx <= ChannelMask && token_idx <= TokenMask && expert_idx <= ExpertMask; - } - - CUTLASS_DEVICE - static uint64_t pack(uint64_t channel_idx, uint64_t token_idx, uint64_t expert_idx) { - if (!fits(channel_idx, token_idx, expert_idx)) { - asm volatile("trap;"); - } - - return uint64_t(channel_idx) | (uint64_t(token_idx) << TokenShift) | - (uint64_t(expert_idx) << ExpertShift); - } - - CUTLASS_DEVICE - static bool is_invalid(uint64_t packed) { return packed == Invalid; } - - CUTLASS_DEVICE - static int32_t channel_idx(uint64_t packed) { return static_cast(packed & ChannelMask); } - - CUTLASS_DEVICE - static int32_t token_idx(uint64_t packed) { - return static_cast((packed >> TokenShift) & TokenMask); - } - - CUTLASS_DEVICE - static int32_t expert_idx(uint64_t packed) { - return static_cast((packed >> ExpertShift) & ExpertMask); - } -}; - -template -struct PrecomputedWorkMapStride {}; - -template <> -struct PrecomputedWorkMapStride { - uint32_t precomputed_work_tiles_per_worker = 0; -}; - -// Persistent Thread Block (TB) scheduler -template -class PersistentTileSchedulerSm90GroupPrecomputed { - // - // Data members - // - - private: - using WorkLinearIdx = uint32_t; - WorkLinearIdx current_work_linear_idx_ = 0; - WorkLinearIdx total_grid_size_ = 0; - - public: - struct WorkTileInfo { - int32_t M_idx = 0; - int32_t N_idx = 0; - int32_t L_idx = 0; - int32_t is_valid_tile = 0; - - CUTLASS_HOST_DEVICE - bool is_valid() const { return is_valid_tile != 0; } - - CUTLASS_HOST_DEVICE - static WorkTileInfo invalid_work_tile() { return {-1, -1, -1, 0}; } - - CUTLASS_HOST_DEVICE - bool is_final_split(uint32_t k_tiles_per_output_tile) const { return true; } - - CUTLASS_HOST_DEVICE - int32_t reduction_subtile_idx() const { return -1; } - }; - - using ProblemShape = typename GroupProblemShape::UnderlyingProblemShape; - using ParamsBase = PersistentTileSchedulerSm90GroupParams; - struct Params : ParamsBase, PrecomputedWorkMapStride { - uint64_t const* precomputed_work_tiles_ = nullptr; - - void initialize_precomputed(dim3 problem_blocks, GemmCoord cluster_shape, - KernelHardwareInfo const& hw_info, int max_swizzle_size, - typename ParamsBase::RasterOrderOptions raster_order_option) { - CUTLASS_UNUSED(hw_info); - - auto problem_blocks_m = round_up(problem_blocks.x, cluster_shape.m()); - auto problem_blocks_n = round_up(problem_blocks.y, cluster_shape.n()); - - this->blocks_across_problem_ = problem_blocks.x * problem_blocks.y * problem_blocks.z; - this->pre_processed_problem_shapes = true; - this->max_swizzle_size_ = max_swizzle_size; - this->raster_order_ = ParamsBase::get_rasterization_order(problem_blocks_m, problem_blocks_n, - raster_order_option); - - this->cluster_shape_ = cluster_shape; - } - }; - using RasterOrder = typename Params::RasterOrder; - using RasterOrderOptions = typename Params::RasterOrderOptions; - static constexpr bool IsDynamicPersistent = false; - - // We need to hard code the number of stages here since the scheduling is static - // and it can benefit from a larger number of stages without worrying about imbalances. - - using Pipeline = PipelineAsync; - - // Call out the types here to work around a bug in MSVC. - - // using PipelineStorage = typename Pipeline::SharedStorage; - // using PipelineState = typename Pipeline::PipelineState; - using PipelineStorage = - cutlass::PipelineDetail::PipelineAsyncSharedStorage; - using PipelineState = - cutlass::PipelineDetail::PipelineAsyncPipelineState; - - using ThrottlePipeline = PipelineEmpty; - using ThrottlePipelineStorage = typename PipelineEmpty::SharedStorage; - using SchedulerResponse = WorkTileInfo; - - class SharedStorage { - public: - CUTLASS_DEVICE PipelineStorage pipeline() { return pipeline_; } - // Pipeline throttle is not needed here as the scheduling is not dynamic. - CUTLASS_DEVICE ThrottlePipelineStorage throttle_pipeline() { return ThrottlePipelineStorage{}; } - CUTLASS_DEVICE SchedulerResponse* data() { return data_; } - - private: - alignas(16) PipelineStorage pipeline_; - alignas(16) SchedulerResponse data_[SchedulerPipelineStageCount]; - }; - - struct Arguments : PrecomputedWorkMapStride { - int max_swizzle_size = 1; - // Not applying Heuristics for Grouped problems, since largest dimension can change per group - RasterOrderOptions raster_order = RasterOrderOptions::AlongM; - uint64_t const* precomputed_work_tiles = nullptr; - }; - - // Sink scheduler params as a member - Params scheduler_params; - void* response_ptr_ = nullptr; - - // - // Methods - // - - template - static Params to_underlying_arguments( - GroupProblemShape problem_shapes, TileShape tile_shape, ClusterShape cluster_shape, - KernelHardwareInfo const& hw_info, Arguments const& arguments, - [[maybe_unused]] void* workspace = nullptr, - [[maybe_unused]] const uint32_t epilogue_subtile = 1, - [[maybe_unused]] uint32_t ktile_start_alignment_count = 1u) { - // We only need the tile and cluster shape during scheduler setup, so let FTAD do the magic - static_assert(cute::is_static::value); - static_assert(cute::is_static::value); - - dim3 problem_blocks = - get_tiled_cta_shape_mnl(problem_shapes, hw_info, tile_shape, cluster_shape); - - CUTLASS_ASSERT(arguments.precomputed_work_tiles != nullptr); - Params params; - params.initialize_precomputed(problem_blocks, to_gemm_coord(cluster_shape), hw_info, - arguments.max_swizzle_size, RasterOrderOptions::AlongM); - params.precomputed_work_tiles_ = arguments.precomputed_work_tiles; - if constexpr (ChunkMajorWorkMap) { - params.precomputed_work_tiles_per_worker = arguments.precomputed_work_tiles_per_worker; - } - - return params; - } - - // Given the inputs, computes the physical grid we should launch. - template - CUTLASS_HOST_DEVICE static dim3 get_grid_shape([[maybe_unused]] Params const& params, - GroupProblemShape const& problem_shapes, - TileShape tile_shape, ClusterShape cluster_shape, - KernelHardwareInfo hw_info, Arguments arguments, - bool truncate_by_problem_size = true) { - dim3 problem_blocks = - get_tiled_cta_shape_mnl(problem_shapes, hw_info, tile_shape, cluster_shape); - - return Params::get_grid_shape(problem_blocks, to_gemm_coord(cluster_shape), hw_info, - arguments.max_swizzle_size, RasterOrderOptions::AlongM, - /* truncate_by_problem_size = */ true); - } - - // Given the inputs, computes the total number of output blocks this problem will compute over - // Note that this is only the logical size of our grid, not the physical grid we will actually - // launch. - template - CUTLASS_HOST_DEVICE static dim3 get_tiled_cta_shape_mnl(GroupProblemShape const& problem_shapes, - KernelHardwareInfo hw_info, - BlockShape cta_shape, - ClusterShape cluster_shape) { - int groups = problem_shapes.groups(); - uint32_t total_ctas = 0; - uint32_t cta_in_N_dim = 1; // We linearize the blocks across all the problems here - - // If host problem shapes are not provided. - if (!problem_shapes.is_host_problem_shape_available()) { - total_ctas = hw_info.sm_count; - } - // If host problem shapes are provided, make a better decision about possibility to launch - // smaller grid. - else { - for (int group = 0; group < groups; group++) { - auto ctas_along_m = - cute::size(cute::ceil_div(cute::shape<0>(problem_shapes.get_host_problem_shape(group)), - cute::shape<0>(cta_shape))); - auto ctas_along_n = - cute::size(cute::ceil_div(cute::shape<1>(problem_shapes.get_host_problem_shape(group)), - cute::shape<1>(cta_shape))); - if (ctas_along_m <= 0) ctas_along_m = 1; - if (ctas_along_n <= 0) ctas_along_n = 1; - auto problem_blocks_m = round_up(ctas_along_m, cute::get<0>(cluster_shape)); - auto problem_blocks_n = round_up(ctas_along_n, cute::get<1>(cluster_shape)); - total_ctas += problem_blocks_m * problem_blocks_n; - } - } - - return Params::get_tiled_cta_shape_mnl(to_gemm_coord(cluster_shape), total_ctas, cta_in_N_dim); - } - - static bool can_implement(Arguments const& args) { - bool implementable = args.precomputed_work_tiles != nullptr; - if constexpr (ChunkMajorWorkMap) { - implementable &= args.precomputed_work_tiles_per_worker > 0; - } - return implementable; - } - - PersistentTileSchedulerSm90GroupPrecomputed() = default; - - // Note: constructing this tile scheduler can touch global memory that was - // written to by the prior kernel. - CUTLASS_DEVICE explicit PersistentTileSchedulerSm90GroupPrecomputed(Params const& params_) - : scheduler_params(params_) { - // MSVC requires protecting use of CUDA-specific nonstandard syntax, - // like blockIdx and gridDim, with __CUDA_ARCH__. -#if defined(__CUDA_ARCH__) - CUTLASS_ASSERT(scheduler_params.precomputed_work_tiles_ != nullptr); - WorkLinearIdx const worker_idx = - WorkLinearIdx(blockIdx.x) * WorkLinearIdx(gridDim.y) + WorkLinearIdx(blockIdx.y) + - WorkLinearIdx(blockIdx.z) * WorkLinearIdx(gridDim.x) * WorkLinearIdx(gridDim.y); - if constexpr (ChunkMajorWorkMap) { - current_work_linear_idx_ = worker_idx * scheduler_params.precomputed_work_tiles_per_worker; - total_grid_size_ = 1; - } else { - current_work_linear_idx_ = worker_idx; - total_grid_size_ = - WorkLinearIdx(gridDim.x) * WorkLinearIdx(gridDim.y) * WorkLinearIdx(gridDim.z); - } - -#else - CUTLASS_ASSERT(false && "This line should never be reached"); -#endif - } - - CUTLASS_DEVICE explicit PersistentTileSchedulerSm90GroupPrecomputed( - Params const& params_, SchedulerResponse* response_ptr) - : scheduler_params(params_), response_ptr_(response_ptr) { - // MSVC requires protecting use of CUDA-specific nonstandard syntax, - // like blockIdx and gridDim, with __CUDA_ARCH__. -#if defined(__CUDA_ARCH__) - CUTLASS_ASSERT(scheduler_params.precomputed_work_tiles_ != nullptr); - WorkLinearIdx const worker_idx = - WorkLinearIdx(blockIdx.x) * WorkLinearIdx(gridDim.y) + WorkLinearIdx(blockIdx.y) + - WorkLinearIdx(blockIdx.z) * WorkLinearIdx(gridDim.x) * WorkLinearIdx(gridDim.y); - if constexpr (ChunkMajorWorkMap) { - current_work_linear_idx_ = worker_idx * scheduler_params.precomputed_work_tiles_per_worker; - total_grid_size_ = 1; - } else { - current_work_linear_idx_ = worker_idx; - total_grid_size_ = - WorkLinearIdx(gridDim.x) * WorkLinearIdx(gridDim.y) * WorkLinearIdx(gridDim.z); - } - -#else - CUTLASS_ASSERT(false && "This line should never be reached"); -#endif - } - - CUTLASS_DEVICE - WorkTileInfo get_current_work() { - return get_current_work_for_linear_idx(current_work_linear_idx_); - } - - CUTLASS_DEVICE - WorkTileInfo get_current_work_for_linear_idx(WorkLinearIdx linear_idx) { - return get_precomputed_work_tile(linear_idx, scheduler_params.precomputed_work_tiles_); - } - - CUTLASS_DEVICE - static WorkTileInfo get_precomputed_work_tile(uint64_t linear_idx, - uint64_t const* precomputed_work_tiles) { - uint64_t const packed = __ldg(precomputed_work_tiles + linear_idx); - if (PrecomputedGroupWorkTile::is_invalid(packed)) { - return WorkTileInfo::invalid_work_tile(); - } - - return {PrecomputedGroupWorkTile::channel_idx(packed), - PrecomputedGroupWorkTile::token_idx(packed), - PrecomputedGroupWorkTile::expert_idx(packed), 1}; - } - - template - CUTLASS_DEVICE auto advance_to_next_work( - TileSchedulerPipeline& scheduler_pipeline, - TileSchedulerPipelineState scheduler_pipe_producer_state, uint32_t advance_count = 1, - CallbackBeforeCommit callback_before_commit = [](WorkTileInfo info) { return info; }) { - current_work_linear_idx_ += total_grid_size_ * WorkLinearIdx(advance_count); - auto work_tile = get_current_work_for_linear_idx(current_work_linear_idx_); - using WorkTileWithCallbackInfo = decltype(callback_before_commit(work_tile)); - WorkTileWithCallbackInfo work_tile_with_callback_info = work_tile; - scheduler_pipeline.producer_acquire(scheduler_pipe_producer_state); - if (work_tile_with_callback_info.is_valid()) { - work_tile_with_callback_info = callback_before_commit(work_tile); - } - - if (cute::elect_one_sync()) { - reinterpret_cast( - response_ptr_)[scheduler_pipe_producer_state.index()] = work_tile_with_callback_info; - cutlass::arch::fence_view_async_shared(); - scheduler_pipeline.producer_commit(scheduler_pipe_producer_state); - } - return cute::make_tuple(work_tile_with_callback_info, true); - } - - CUTLASS_DEVICE - void advance_to_next_work() { current_work_linear_idx_ += total_grid_size_; } - - CUTLASS_DEVICE - void advance_to_next_work(uint32_t advance_count) { - current_work_linear_idx_ += total_grid_size_ * WorkLinearIdx(advance_count); - } - - // Returns whether the block assigned this work should compute the epilogue for the corresponding - // output tile. For the basic tile scheduler, this is always true. - CUTLASS_HOST_DEVICE - static bool compute_epilogue(WorkTileInfo const&, Params const&) { return true; } - - // Performs the reduction across splits for a given output tile. Since this scheduler does - // not split output tiles, no reduction is needed. - template - CUTLASS_DEVICE static void fixup(Params const&, WorkTileInfo const&, FrgTensorC&, uint32_t, - uint32_t) {} - - // Returns whether the current WorkTileInfo passed in should continue to be used. Since - // this scheduler only schedules work in units of single, full output tiles, the WorkTileInfo - // passed in should not be used after having been processed. - CUTLASS_DEVICE - static bool continue_current_work(WorkTileInfo&) { return false; } - - // The basic tile scheduler does not require any additional workspace - template - static size_t get_workspace_size(Arguments const&, ProblemShape, KernelHardwareInfo const&, - uint32_t, const uint32_t = 1, uint32_t = 1) { - return 0; - } - - template - static cutlass::Status initialize_workspace(Arguments const&, void*, cudaStream_t, ProblemShape, - KernelHardwareInfo const&, uint32_t, - const uint32_t = 1, uint32_t = 1, - CudaHostAdapter* cuda_adapter = nullptr) { - return Status::kSuccess; - } - - template - CUTLASS_HOST_DEVICE static int get_work_k_tile_count(WorkTileInfo const& work_tile_info, - ProblemShape_MNKL problem_shape, - TileShape tile_shape) { - // All work units returned by this scheduler cover the entire K iteration - // space of the output tile assigned to the work unit. - return cute::size(cute::ceil_div(cute::get<2>(problem_shape), cute::get<2>(tile_shape))); - } - - CUTLASS_HOST_DEVICE - static uint32_t get_work_k_tile_start(WorkTileInfo const&) { - // All work units returned by this scheduler start from K tile 0 - return 0u; - } - - CUTLASS_DEVICE - static bool need_separate_reduction(Params const& params) { return false; } - - CUTLASS_DEVICE - bool is_work_tile_for_reduction(WorkTileInfo const& work_tile_info, Params const& params) { - return false; - } - - CUTLASS_DEVICE - uint32_t epilgoue_subtile_idx(WorkTileInfo const& work_tile_info, Params const& params) const { - return 0; - } - - template - CUTLASS_DEVICE void separate_reduction(Params const& params, WorkTileInfo const& work_tile_info, - FrgTensorC& accumulators, uint32_t num_barriers, - uint32_t barrier_idx) {} - - // Shares the accumulator set with peers in the global workspace - template - CUTLASS_DEVICE static void share(Params const& params, WorkTileInfo const& work_tile_info, - FrgTensorC& accumulators, uint32_t num_barriers, - uint32_t barrier_idx) {} - - CUTLASS_DEVICE - static bool valid_warpgroup_in_work_tile(WorkTileInfo const& work_tile_info) { return true; } - - CUTLASS_DEVICE - static bool requires_separate_reduction(Params const& params) { return false; } - - // Kernel helper function to get next work tile - template - CUTLASS_DEVICE auto fetch_next_work(WorkTileWithCallbackInfo work_tile_with_callback_info, - TileSchedulerPipeline& scheduler_pipeline, - TileSchedulerPipelineState scheduler_pipe_consumer_state) { - if (continue_current_work(work_tile_with_callback_info)) { - return cute::make_tuple(work_tile_with_callback_info, true); - } - scheduler_pipeline.consumer_wait(scheduler_pipe_consumer_state); - work_tile_with_callback_info = reinterpret_cast( - response_ptr_)[scheduler_pipe_consumer_state.index()]; - cutlass::arch::fence_view_async_shared(); - scheduler_pipeline.consumer_release(scheduler_pipe_consumer_state); - - return cute::make_tuple(work_tile_with_callback_info, true); - } - - CUTLASS_DEVICE - auto fetch_next_work(WorkTileInfo work_tile_info) { - if (continue_current_work(work_tile_info)) { - return cute::make_tuple(work_tile_info, true); - } - - advance_to_next_work(); - return cute::make_tuple(get_current_work(), true); - } - - // Returns the initial work tile info that will be computed over - template - CUTLASS_DEVICE auto initial_work_tile_info( - ClusterShape, CallbackBeforeCommit callback_before_commit = [](WorkTileInfo response) { - return response; - }) { - auto work_tile = get_current_work_for_linear_idx(current_work_linear_idx_); - using WorkTileWithCallbackInfo = decltype(callback_before_commit(work_tile)); - WorkTileWithCallbackInfo work_tile_with_callback_info = work_tile; - if (work_tile_with_callback_info.is_valid()) { - work_tile_with_callback_info = callback_before_commit(work_tile); - } - return work_tile_with_callback_info; - } -}; - -} // namespace cutlass::gemm::kernel::detail diff --git a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm_configs.h b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm_configs.h index 579e88dc408..ee39030b783 100644 --- a/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm_configs.h +++ b/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm_configs.h @@ -106,40 +106,20 @@ enum class CutlassTileConfigSM90 : int { // CTA configs for M=64 CtaShape64x16x128B = shape_tuple_to_enum(64, 16, 128), - CtaShape64x16x256B = shape_tuple_to_enum(64, 16, 256), - CtaShape64x16x512B = shape_tuple_to_enum(64, 16, 512), CtaShape64x32x128B = shape_tuple_to_enum(64, 32, 128), - CtaShape64x32x256B = shape_tuple_to_enum(64, 32, 256), - CtaShape64x32x512B = shape_tuple_to_enum(64, 32, 512), CtaShape64x64x128B = shape_tuple_to_enum(64, 64, 128), - CtaShape64x64x256B = shape_tuple_to_enum(64, 64, 256), - CtaShape64x64x512B = shape_tuple_to_enum(64, 64, 512), CtaShape64x128x128B = shape_tuple_to_enum(64, 128, 128), - CtaShape64x128x256B = shape_tuple_to_enum(64, 128, 256), - CtaShape64x128x512B = shape_tuple_to_enum(64, 128, 512), CtaShape64x256x128B = shape_tuple_to_enum(64, 256, 128), // CTA configs for M=128 - CtaShape128x8x128B = shape_tuple_to_enum(128, 8, 128), CtaShape128x16x128B = shape_tuple_to_enum(128, 16, 128), - CtaShape128x16x256B = shape_tuple_to_enum(128, 16, 256), - CtaShape128x16x512B = shape_tuple_to_enum(128, 16, 512), CtaShape128x32x128B = shape_tuple_to_enum(128, 32, 128), - CtaShape128x32x256B = shape_tuple_to_enum(128, 32, 256), - CtaShape128x32x512B = shape_tuple_to_enum(128, 32, 512), - CtaShape128x40x128B = shape_tuple_to_enum(128, 40, 128), CtaShape128x64x128B = shape_tuple_to_enum(128, 64, 128), - CtaShape128x64x256B = shape_tuple_to_enum(128, 64, 256), - CtaShape128x64x512B = shape_tuple_to_enum(128, 64, 512), CtaShape128x128x128B = shape_tuple_to_enum(128, 128, 128), - CtaShape128x128x256B = shape_tuple_to_enum(128, 128, 256), - CtaShape128x128x512B = shape_tuple_to_enum(128, 128, 512), CtaShape128x256x128B = shape_tuple_to_enum(128, 256, 128), - CtaShape128x256x256B = shape_tuple_to_enum(128, 256, 256), // CTA configs for M=256 CtaShape256x128x128B = shape_tuple_to_enum(256, 128, 128), - CtaShape256x128x256B = shape_tuple_to_enum(256, 128, 256), CtaShape256x256x128B = shape_tuple_to_enum(256, 256, 128), }; @@ -203,9 +183,7 @@ enum class MainloopScheduleType { // architectures, this defaults to the "legacy" main loop schedule. PINGPONG, COOPERATIVE, - WARPSPECIALIZED, - SINGLE_WARPGROUP_PREFILL, - SINGLE_WARPGROUP_ROLLING + WARPSPECIALIZED }; static auto get_mainloop_schedule_name(MainloopScheduleType schedule) { @@ -217,10 +195,6 @@ static auto get_mainloop_schedule_name(MainloopScheduleType schedule) { return "cooperative"; } else if (schedule == MainloopScheduleType::WARPSPECIALIZED) { return "warpspecialized"; - } else if (schedule == MainloopScheduleType::SINGLE_WARPGROUP_PREFILL) { - return "single_warpgroup_prefill"; - } else if (schedule == MainloopScheduleType::SINGLE_WARPGROUP_ROLLING) { - return "single_warpgroup_rolling"; } return "unknown schedule"; } diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp index 4eae529bb68..4b0e4f56c19 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp @@ -203,21 +203,10 @@ std::vector get_candidate_tiles_sm90( if (config & CutlassGemmConfig::GROUPED_GEMM) { if (config & CutlassGemmConfig::WEIGHT_ONLY) { return { - CutlassTileConfigSM90::CtaShape64x16x128B, CutlassTileConfigSM90::CtaShape64x16x256B, - CutlassTileConfigSM90::CtaShape64x16x512B, CutlassTileConfigSM90::CtaShape64x32x128B, - CutlassTileConfigSM90::CtaShape64x32x256B, CutlassTileConfigSM90::CtaShape64x32x512B, - CutlassTileConfigSM90::CtaShape64x64x128B, CutlassTileConfigSM90::CtaShape64x64x256B, - CutlassTileConfigSM90::CtaShape64x64x512B, CutlassTileConfigSM90::CtaShape64x128x128B, - CutlassTileConfigSM90::CtaShape64x128x256B, CutlassTileConfigSM90::CtaShape64x128x512B, - CutlassTileConfigSM90::CtaShape128x16x128B, CutlassTileConfigSM90::CtaShape128x16x256B, - CutlassTileConfigSM90::CtaShape128x16x512B, CutlassTileConfigSM90::CtaShape128x32x128B, - CutlassTileConfigSM90::CtaShape128x32x256B, CutlassTileConfigSM90::CtaShape128x32x512B, - CutlassTileConfigSM90::CtaShape128x64x128B, CutlassTileConfigSM90::CtaShape128x64x256B, - CutlassTileConfigSM90::CtaShape128x64x512B, CutlassTileConfigSM90::CtaShape128x128x128B, - CutlassTileConfigSM90::CtaShape128x128x256B, CutlassTileConfigSM90::CtaShape128x128x512B, - CutlassTileConfigSM90::CtaShape128x256x128B, CutlassTileConfigSM90::CtaShape128x256x256B, - CutlassTileConfigSM90::CtaShape256x128x128B, CutlassTileConfigSM90::CtaShape256x128x256B, - CutlassTileConfigSM90::CtaShape256x256x128B}; + CutlassTileConfigSM90::CtaShape64x16x128B, CutlassTileConfigSM90::CtaShape64x32x128B, + CutlassTileConfigSM90::CtaShape64x64x128B, CutlassTileConfigSM90::CtaShape64x128x128B, + CutlassTileConfigSM90::CtaShape128x16x128B, CutlassTileConfigSM90::CtaShape128x32x128B, + CutlassTileConfigSM90::CtaShape128x64x128B, CutlassTileConfigSM90::CtaShape128x128x128B}; } else { return { CutlassTileConfigSM90::CtaShape128x16x128B, CutlassTileConfigSM90::CtaShape128x32x128B, @@ -239,20 +228,12 @@ bool sm90_supports_coop(CutlassTileConfigSM90 const tile) { #ifdef FAST_BUILD return false; #else - auto const [tile_m, tile_n, tile_k] = enum_to_shape_tuple(tile); - if (tile_m == 128 && (tile_n == 16 || tile_n == 32 || tile_n == 64)) { - return tile_k == 128 || tile_k == 256 || tile_k == 512; - } - if (tile_m == 128 && (tile_n == 128 || tile_n == 256)) { - return tile_k == 128 || tile_k == 256; - } - if (tile_m == 256 && tile_n == 128) { - return tile_k == 128 || tile_k == 256; - } - if (tile_m == 256 && tile_n == 256) { - return tile_k == 128; - } - return false; + std::set valid_tiles{ + CutlassTileConfigSM90::CtaShape128x16x128B, CutlassTileConfigSM90::CtaShape128x32x128B, + CutlassTileConfigSM90::CtaShape128x64x128B, CutlassTileConfigSM90::CtaShape128x128x128B, + CutlassTileConfigSM90::CtaShape128x256x128B, CutlassTileConfigSM90::CtaShape256x128x128B, + CutlassTileConfigSM90::CtaShape256x256x128B}; + return valid_tiles.count(tile) == 1; #endif } @@ -262,12 +243,11 @@ bool sm90_supports_mcast_along_m(CutlassTileConfigSM90 const tile) { #ifdef FAST_BUILD return false; #else - auto const [tile_m, tile_n, tile_k] = enum_to_shape_tuple(tile); - bool const supported_k = tile_k == 128 || tile_k == 256 || tile_k == 512; - bool const supported_mn = (tile_m == 128 && (tile_n == 16 || tile_n == 32 || tile_n == 64 || - tile_n == 128 || tile_n == 256)) || - (tile_m == 256 && tile_n == 128); - return supported_mn && supported_k; + std::set valid_tiles{ + CutlassTileConfigSM90::CtaShape128x16x128B, CutlassTileConfigSM90::CtaShape128x32x128B, + CutlassTileConfigSM90::CtaShape128x64x128B, CutlassTileConfigSM90::CtaShape128x128x128B, + CutlassTileConfigSM90::CtaShape128x256x128B, CutlassTileConfigSM90::CtaShape256x128x128B}; + return valid_tiles.count(tile) == 1; #endif } @@ -277,11 +257,11 @@ bool sm90_supports_mcast_along_n(CutlassTileConfigSM90 const tile) { #ifdef FAST_BUILD return false; #else - auto const [tile_m, tile_n, tile_k] = enum_to_shape_tuple(tile); - bool const supported_k = tile_k == 128 || tile_k == 256 || tile_k == 512; - bool const supported_mn = ((tile_m == 64 || tile_m == 128) && (tile_n == 128 || tile_n == 256)) || - (tile_m == 256 && tile_n == 128); - return supported_mn && supported_k; + std::set valid_tiles{ + CutlassTileConfigSM90::CtaShape64x128x128B, CutlassTileConfigSM90::CtaShape64x256x128B, + CutlassTileConfigSM90::CtaShape128x128x128B, CutlassTileConfigSM90::CtaShape128x256x128B, + CutlassTileConfigSM90::CtaShape256x128x128B}; + return valid_tiles.count(tile) == 1; #endif } @@ -411,6 +391,11 @@ std::vector get_candidate_configs_sm90( bool const has_coop_supported = sm90_supports_coop(tile_config); std::set mainloop_schedules; if (has_coop_supported) { + // Due to the limitation on the number of registers on SM, + // cooperative scheduler does not support CtaShape128x128x128B + // for mixed-dtype (W4A16) grouped GEMM. Skip the tile entirely + // to avoid register overflow. + if (tile_config == CutlassTileConfigSM90::CtaShape128x128x128B) continue; mainloop_schedules.insert(MainloopScheduleType::COOPERATIVE); } else { mainloop_schedules.insert(MainloopScheduleType::PINGPONG); diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h index 6037ba4b94a..bdf29ce7a99 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h @@ -19,6 +19,9 @@ #include #include #include +#if CUDA_VERSION >= 12080 +#include +#endif #include #include "cutlass/bfloat16.h" @@ -26,9 +29,6 @@ #include "cutlass/float_subbyte.h" #include "cutlass/half.h" #include "tensorrt_llm/common/NvInferRuntime.h" -#if defined(ENABLE_FP4) -#include "tensorrt_llm/kernels/cutlass_kernels/fp4_compat.h" -#endif namespace tensorrt_llm { namespace kernels { @@ -93,13 +93,14 @@ struct TllmToCutlassTypeAdapter<__nv_fp8_e5m2> { }; #endif -#if defined(ENABLE_FP4) && !defined(COMPILE_HOPPER_TMA_GEMMS) && \ - !defined(COMPILE_HOPPER_TMA_GROUPED_GEMMS) && !defined(CUTLASS_ENABLE_GDC_FOR_SM90) +#if defined(ENABLE_FP4) +#if CUDA_VERSION >= 12080 template <> -struct TllmToCutlassTypeAdapter { +struct TllmToCutlassTypeAdapter<__nv_fp4_e2m1> { using type = cutlass::float_e2m1_t; }; #endif +#endif /////////////////////////////////////////////////////////////////////////////////////////////////// // Cutlass to Tllm @@ -133,13 +134,14 @@ struct CutlassToTllmTypeAdapter { }; #endif -#if defined(ENABLE_FP4) && !defined(COMPILE_HOPPER_TMA_GEMMS) && \ - !defined(COMPILE_HOPPER_TMA_GROUPED_GEMMS) && !defined(CUTLASS_ENABLE_GDC_FOR_SM90) +#if defined(ENABLE_FP4) +#if CUDA_VERSION >= 12080 template <> struct CutlassToTllmTypeAdapter { - using type = Fp4Type; + using type = __nv_fp4_e2m1; }; #endif +#endif /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/fp4_compat.h b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/fp4_compat.h deleted file mode 100644 index 12eca02d805..00000000000 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/fp4_compat.h +++ /dev/null @@ -1,39 +0,0 @@ -/* - * Copyright (c) 2026 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include - -#if defined(COMPILE_HOPPER_TMA_GEMMS) || defined(COMPILE_HOPPER_TMA_GROUPED_GEMMS) || \ - defined(CUTLASS_ENABLE_GDC_FOR_SM90) -#include "cutlass/float_subbyte.h" - -namespace tensorrt_llm::kernels::cutlass_kernels { -using Fp4Type = cutlass::float_e2m1_t; -} // namespace tensorrt_llm::kernels::cutlass_kernels - -#else -#if CUDA_VERSION < 12080 -#error "Native FP4 paths require CUDA 12.8 or newer outside the Hopper CUTLASS FP4 path." -#endif -#include - -namespace tensorrt_llm::kernels::cutlass_kernels { -using Fp4Type = __nv_fp4_e2m1; -} // namespace tensorrt_llm::kernels::cutlass_kernels - -#endif diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h index 3761e32cf12..14261cf1fea 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h @@ -18,7 +18,6 @@ #include #include -#include #include #include @@ -33,7 +32,7 @@ #include "tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm_configs.h" #ifdef ENABLE_FP4 -#include "tensorrt_llm/kernels/cutlass_kernels/fp4_compat.h" +#include #endif namespace tensorrt_llm::kernels::cutlass_kernels { @@ -216,9 +215,6 @@ struct TmaWarpSpecializedGroupedGemmInput { uint8_t* gemm_workspace = nullptr; size_t gemm_workspace_size = 0; - uint8_t* precomputed_scheduler_workspace = nullptr; - size_t precomputed_scheduler_workspace_size = 0; - int64_t precomputed_scheduler_total_routed_tokens = 0; // Whether to enable PDL (Programmatic Dependent Launch). bool enable_pdl{}; @@ -228,9 +224,7 @@ struct TmaWarpSpecializedGroupedGemmInput { static size_t workspaceSize(int num_experts, FpXBlockScalingType scaling_type); void configureWorkspace(int8_t* start_ptr, int num_experts, void* gemm_workspace, - size_t gemm_workspace_size, void* precomputed_scheduler_workspace, - size_t precomputed_scheduler_workspace_size, - FpXBlockScalingType scaling_type); + size_t gemm_workspace_size, FpXBlockScalingType scaling_type); bool isValid() const { return stride_act != nullptr && ptr_act != nullptr; } @@ -247,31 +241,22 @@ constexpr bool isGatedActivation(ActivationType activation_type) { activation_type == ActivationType::GegluTanh; } -enum class Sm90Wfp4Afp8ScaleMode : uint8_t { - // Native SM100+ FP8 x FP4 paths leave this SM90-only mode disabled. - kDisabled = 0, - kHummingPreMmaE8M0, - kPostMmaFp8Act, - kPostMmaMxfp8Act, -}; - template + bool IsMXFPX = false> class MoeGemmRunner { public: MoeGemmRunner(); #if defined(ENABLE_FP4) #if defined(ENABLE_BF16) - static constexpr bool use_wfp4a16 = std::is_same_v && + static constexpr bool use_wfp4a16 = std::is_same_v && (std::is_same_v || std::is_same_v); #else static constexpr bool use_wfp4a16 = - std::is_same_v && std::is_same_v; + std::is_same_v && std::is_same_v; #endif #else static constexpr bool use_wfp4a16 = false; @@ -281,7 +266,7 @@ class MoeGemmRunner { (std::is_same_v || std::is_same_v) && !std::is_same_v #if defined(ENABLE_FP4) - && !std::is_same_v + && !std::is_same_v #endif ; static constexpr bool use_w4afp8 = @@ -290,25 +275,18 @@ class MoeGemmRunner { static constexpr bool use_fp8 = false; static constexpr bool use_w4afp8 = false; #endif + static constexpr bool use_mxfp8 = use_fp8 && IsMXFPX; + + static constexpr bool use_w4_groupwise = use_w4afp8 || use_wfp4a16; + #if defined(ENABLE_FP4) - static constexpr bool use_fp4 = std::is_same_v; + static constexpr bool use_fp4 = std::is_same_v; static constexpr bool use_wfp4afp8 = - std::is_same_v && std::is_same_v; + std::is_same_v && std::is_same_v; #else static constexpr bool use_fp4 = false; static constexpr bool use_wfp4afp8 = false; #endif - static constexpr bool use_mxfp8 = use_fp8 && IsMXFPX; - static constexpr bool use_sm90_wfp4afp8 = - use_wfp4afp8 && Sm90Wfp4Afp8Mode != Sm90Wfp4Afp8ScaleMode::kDisabled; - static constexpr bool use_sm90_humming_pre_mma = - use_sm90_wfp4afp8 && Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0; - static_assert(Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kDisabled || use_wfp4afp8, - "Sm90Wfp4Afp8ScaleMode is only valid for FP8 activation x FP4 weight."); - static_assert(!use_sm90_wfp4afp8 || !IsMXFPX, - "FP8 activation x FP4 weight uses Sm90Wfp4Afp8ScaleMode, not generic IsMXFPX."); - - static constexpr bool use_sm90_mixed_input_gemm = use_w4afp8 || use_wfp4a16 || use_sm90_wfp4afp8; void moeGemmBiasAct(GroupedGemmInput inputs, TmaWarpSpecializedGroupedGemmInput hopper_inputs); diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h index c19603110be..429626bd8c6 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h @@ -25,7 +25,7 @@ #include "tensorrt_llm/common/quantization.h" #include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" #ifdef ENABLE_FP4 -#include "tensorrt_llm/kernels/cutlass_kernels/fp4_compat.h" +#include #endif #include @@ -482,7 +482,6 @@ class CutlassMoeFCRunnerInterface { void const* const fc1_expert_weights, void const* const fc1_expert_biases, int64_t const* const num_valid_tokens_ptr, void const* const fc1_int_scales, float const* const fc1_fp8_dequant, float const* const fc2_fp8_quant, - float* const act_fp8_token_scale, TmaWarpSpecializedGroupedGemmInput::ElementSF const* fc1_fp4_act_flat, TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_fp4_act_flat, QuantParams quant_params, int64_t const num_rows, @@ -554,23 +553,20 @@ template + bool IsMXFPX = false, typename Enable = void> class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { using DeepSeekBlockScaleGemmRunner = tensorrt_llm::kernels::fp8_blockscale_gemm::CutlassFp8BlockScaleGemmRunnerInterface; using ScaleBiasType = BackBoneType; - using Self = CutlassMoeFCRunner; + using Self = CutlassMoeFCRunner; #if defined(ENABLE_FP4) #if defined(ENABLE_BF16) - static constexpr bool use_wfp4a16 = std::is_same_v && + static constexpr bool use_wfp4a16 = std::is_same_v && (std::is_same_v || std::is_same_v); #else static constexpr bool use_wfp4a16 = - std::is_same_v && std::is_same_v; + std::is_same_v && std::is_same_v; #endif #else static constexpr bool use_wfp4a16 = false; @@ -578,11 +574,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { #if defined(ENABLE_FP8) static constexpr bool use_fp8 = (std::is_same_v || std::is_same_v) && - !std::is_same_v -#if defined(ENABLE_FP4) - && !std::is_same_v -#endif - ; + !std::is_same_v; static constexpr bool use_w4afp8 = std::is_same_v && std::is_same_v; static constexpr bool use_fp8_input = std::is_same_v; @@ -594,14 +586,15 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { static constexpr bool use_fp8 = false; static constexpr bool use_w4afp8 = false; #endif + static constexpr bool use_w4_groupwise = use_w4afp8 || use_wfp4a16; #if defined(ENABLE_FP4) - static constexpr bool act_fp4 = std::is_same_v; - static constexpr bool weight_fp4 = std::is_same_v; + static constexpr bool act_fp4 = std::is_same_v; + static constexpr bool weight_fp4 = std::is_same_v; static constexpr bool use_wfp4afp8 = std::is_same_v && weight_fp4; static constexpr bool use_fp4 = act_fp4 && weight_fp4; - static_assert(!std::is_same_v, + static_assert(!std::is_same_v, "Current logic requires backbone type to be >=16-bits"); - static_assert(!std::is_same_v, + static_assert(!std::is_same_v, "Current logic requires output type to be >=16-bits"); #else static constexpr bool act_fp4 = false; @@ -611,20 +604,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { #endif static constexpr bool use_mxfp8 = use_fp8 && IsMXFPX; - static constexpr bool use_sm90_wfp4afp8 = - use_wfp4afp8 && Sm90Wfp4Afp8Mode != Sm90Wfp4Afp8ScaleMode::kDisabled; - static constexpr bool use_sm90_humming_pre_mma = - use_sm90_wfp4afp8 && Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0; - static constexpr bool use_native_wfp4afp8 = - use_wfp4afp8 && Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kDisabled; - static_assert(Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kDisabled || use_wfp4afp8, - "Sm90Wfp4Afp8ScaleMode is only valid for FP8 activation x FP4 weight."); - static_assert(!use_sm90_wfp4afp8 || !IsMXFPX, - "FP8 activation x FP4 weight uses Sm90Wfp4Afp8ScaleMode, not generic IsMXFPX."); - - static constexpr bool use_sm90_mixed_input_gemm = use_w4afp8 || use_wfp4a16 || use_sm90_wfp4afp8; - static constexpr bool use_block_scaling = - use_fp4 || use_mxfp8 || (use_wfp4afp8 && !use_sm90_humming_pre_mma); + static constexpr bool use_block_scaling = use_fp4 || use_wfp4afp8 || use_mxfp8; // This should leave the variable unchanged in any currently supported configuration using UnfusedGemmOutputType = BackBoneType; @@ -702,8 +682,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { cudaStream_t stream) override; // We make these GEMM1 & GEMM2 static because they need to be stateless for the profiler to work - static void gemm1(MoeGemmRunner& gemm_runner, + static void gemm1(MoeGemmRunner& gemm_runner, // This argument must not be null if fp8 block scaling is being used. // The gemm_runner will be ignored in that case. NOTE: it would // be great if we could consolidate gemm_runner and fp8_blockscale_gemm_runner. @@ -717,7 +696,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { ScaleBiasType const* const fc1_expert_biases, int64_t const* const num_valid_tokens_ptr, ScaleBiasType const* const fc1_int_scales, float const* const fc1_fp8_dequant, - float const* const fc2_fp8_quant, float* const act_fp8_token_scale, + float const* const fc2_fp8_quant, TmaWarpSpecializedGroupedGemmInput::ElementSF const* fc1_fp4_act_flat, TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_fp4_act_flat, QuantParams quant_params, int64_t const num_rows, @@ -729,8 +708,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { int* num_active_experts_per, int* active_expert_global_ids, bool enable_pdl); static void gemm2( - MoeGemmRunner& - gemm_runner, + MoeGemmRunner& gemm_runner, DeepSeekBlockScaleGemmRunner* fp8_blockscale_gemm_runner, T const* const input, void* const gemm_output, OutputType* const final_output, int64_t const* const expert_first_token_offset, @@ -758,7 +736,6 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { void const* const fc1_expert_weights, void const* const fc1_expert_biases, int64_t const* const num_valid_tokens_ptr, void const* const fc1_int_scales, float const* const fc1_fp8_dequant, float const* const fc2_fp8_quant, - float* const act_fp8_token_scale, TmaWarpSpecializedGroupedGemmInput::ElementSF const* fc1_fp4_act_flat, TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_fp4_act_flat, QuantParams quant_params, int64_t const num_rows, int64_t const expanded_num_rows, @@ -774,11 +751,11 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { tma_ws_input_template, static_cast(fc1_expert_weights), static_cast(fc1_expert_biases), num_valid_tokens_ptr, static_cast(fc1_int_scales), fc1_fp8_dequant, - fc2_fp8_quant, act_fp8_token_scale, fc1_fp4_act_flat, fc2_fp4_act_flat, - quant_params, num_rows, expanded_num_rows, hidden_size, inter_size, - num_experts_per_node, fc1_activation_type, alpha_scale_ptr_array, - bias_is_broadcast, stream, config, min_latency_mode, num_active_experts_per, - active_expert_global_ids, enable_pdl); + fc2_fp8_quant, fc1_fp4_act_flat, fc2_fp4_act_flat, quant_params, num_rows, + expanded_num_rows, hidden_size, inter_size, num_experts_per_node, + fc1_activation_type, alpha_scale_ptr_array, bias_is_broadcast, stream, + config, min_latency_mode, num_active_experts_per, active_expert_global_ids, + enable_pdl); } void gemm2(void const* const input, void* const gemm_output, void* const final_output, @@ -933,19 +910,17 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { bool mayHaveFinalizeFused() const { return moe_gemm_runner_.supportsTmaWarpSpecialized() && moe_gemm_runner_.getSM() >= 90 && - use_fused_finalize_ && !use_sm90_mixed_input_gemm; + use_fused_finalize_ && !use_w4_groupwise; } static bool mayHaveFinalizeFused(int sm) { using RunnerType = decltype(moe_gemm_runner_); - return RunnerType::supportsTmaWarpSpecialized(sm) && sm >= 90 && !use_sm90_mixed_input_gemm; + return RunnerType::supportsTmaWarpSpecialized(sm) && sm >= 90 && !use_w4_groupwise; } // TODO: This should eventually take the quant params to give more flexibility static auto getScalingType() { - constexpr bool use_mxfp8_act_scale = - use_wfp4afp8 && Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kPostMmaMxfp8Act; - return (use_native_wfp4afp8 || use_mxfp8_act_scale || use_mxfp8) + return (use_wfp4afp8 || use_mxfp8) ? TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX : use_fp4 ? TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4 : TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE; @@ -997,8 +972,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { int64_t const expanded_num_rows, int64_t const seq_len, bool const use_awq, cudaStream_t stream); - MoeGemmRunner - moe_gemm_runner_; + MoeGemmRunner moe_gemm_runner_; std::unique_ptr blockscale_gemm_runner_; std::optional gemm1_config_; @@ -1018,9 +992,9 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { void* glu_inter_result_{}; void* fc2_result_{}; T* fc1_result_{}; + // TODO If we fuse the quantization for GEMM2 into GEMM1 we will need two pointers TmaWarpSpecializedGroupedGemmInput::ElementSF* fc1_fp4_act_scale_; TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_fp4_act_scale_; - float* act_fp8_token_scale_{}; float const** alpha_scale_ptr_array_fc1_ = nullptr; float const** alpha_scale_ptr_array_fc2_ = nullptr; ScaleBiasType* lora_input_{}; @@ -1056,9 +1030,8 @@ struct GemmProfilerBackend { int num_experts, int k, int64_t hidden_size, int64_t unpadded_hidden_size, int64_t inter_size, int64_t group_size, ActivationType activation_type, bool bias, bool use_lora, bool min_latency_mode, bool need_weights, - MOEParallelismConfig parallelism_config, bool const enable_alltoall = false, - bool use_mxfp8_act_scaling = false, - Sm90Wfp4Afp8ScaleMode sm90_wfp4afp8_mode = Sm90Wfp4Afp8ScaleMode::kDisabled) { + MOEParallelismConfig parallelism_config, bool const enable_alltoall, + bool use_mxfp8_act_scaling = false) { mInterface = &runner; mGemmToProfile = gemm_to_profile; mDType = dtype; @@ -1078,13 +1051,11 @@ struct GemmProfilerBackend { mNeedWeights = need_weights; mParallelismConfig = parallelism_config; mEnableAlltoall = enable_alltoall; - mSm90Wfp4Afp8Mode = sm90_wfp4afp8_mode; mSM = common::getSMVersion(); mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE; - bool const use_sm90_mxfp8_act_scale = - mSm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kPostMmaMxfp8Act; - if (isNativeWfp4Afp8Family() || use_sm90_mxfp8_act_scale) { + if (dtype == nvinfer1::DataType::kFP8 && + (wtype == nvinfer1::DataType::kFP4 || wtype == nvinfer1::DataType::kINT64)) { mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX; } else if ((dtype == nvinfer1::DataType::kFP4 || dtype == nvinfer1::DataType::kINT64) && (wtype == nvinfer1::DataType::kFP4 || wtype == nvinfer1::DataType::kINT64)) { @@ -1144,45 +1115,10 @@ struct GemmProfilerBackend { bool mUseLora{}; bool mMinLatencyMode{}; bool mNeedWeights{}; - Sm90Wfp4Afp8ScaleMode mSm90Wfp4Afp8Mode = Sm90Wfp4Afp8ScaleMode::kDisabled; TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType mScalingType{}; private: - bool isNativeWfp4Afp8Family() const { - return mSM >= 100 && mDType == nvinfer1::DataType::kFP8 && - (mWType == nvinfer1::DataType::kFP4 || mWType == nvinfer1::DataType::kINT64); - } - - bool isSm90Wfp4Afp8Family() const { - return mSM == 90 && mDType == nvinfer1::DataType::kFP8 && - mWType == nvinfer1::DataType::kUINT8 && - mGroupSize == - TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::wfp4a16_group_size; - } - - bool isHummingPreMmaScaleMode() const { - return mSM == 90 && mSm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0; - } - - bool isSm90MixedInputFamily() const { - bool const is_fp8_groupwise_weight_family = - mDType == nvinfer1::DataType::kFP8 && - (mWType == nvinfer1::DataType::kINT8 || mWType == nvinfer1::DataType::kINT4 || - mWType == nvinfer1::DataType::kUINT8) && - mGroupSize > 0; - bool const is_wfp4a16_family = - (mDType == nvinfer1::DataType::kHALF || mDType == nvinfer1::DataType::kBF16) && - mWType == nvinfer1::DataType::kUINT8; - return mSM == 90 && (is_fp8_groupwise_weight_family || is_wfp4a16_family); - } - - void checkSm90Wfp4Afp8ScaleMode() const { - TLLM_CHECK_WITH_INFO( - isSm90Wfp4Afp8Family() == (mSm90Wfp4Afp8Mode != Sm90Wfp4Afp8ScaleMode::kDisabled), - "Sm90Wfp4Afp8ScaleMode must be set exactly for SM90 FP8 activation x packed MXFP4 weight."); - } - void prepareRouting(int num_tokens, char* workspace, bool enable_pdl, cudaStream_t stream); void prepareQuantParams(int num_tokens, char* workspace, cudaStream_t stream); void prepareTmaWsInputs(int num_tokens, char* workspace, void const* expert_weights, diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h index 228422793e0..01f107d095c 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h @@ -22,7 +22,7 @@ #include "tensorrt_llm/common/quantization.h" #include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" #ifdef ENABLE_FP4 -#include "tensorrt_llm/kernels/cutlass_kernels/fp4_compat.h" +#include #endif #include @@ -65,9 +65,7 @@ void expandInputRowsKernelLauncher( bool use_per_expert_act_scale, int64_t* expert_first_token_offset, TmaWarpSpecializedGroupedGemmInput::ElementSF* fc1_act_sf_flat, TmaWarpSpecializedGroupedGemmInput::ElementSF const* input_sf, bool const swizzled_input_sf, - void const* prequant_scales, float* fp8_token_dequant_scale, - float const* fp8_token_residual_scale, float const** fp8_token_scale_ptr_array, bool enable_pdl, - cudaStream_t stream); + void const* prequant_scales, bool enable_pdl, cudaStream_t stream); template void finalizeMoeRoutingKernelLauncher( diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl index 804a9d2cebb..46c1ee73c8a 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl @@ -42,7 +42,7 @@ #include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" // #include #ifdef ENABLE_FP4 -#include "tensorrt_llm/kernels/cutlass_kernels/fp4_compat.h" +#include #endif #include #include @@ -173,7 +173,7 @@ using SafeFP8 = __nv_fp8_e4m3; using SafeFP8 = void; #endif #ifdef ENABLE_FP4 -using SafeFP4 = Fp4Type; +using SafeFP4 = __nv_fp4_e2m1; #else struct SafeFP4 {}; #endif diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.h b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.h index 5d76ea0ff94..91d12ef0e7f 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.h +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.h @@ -14,12 +14,9 @@ * limitations under the License. */ -#pragma once - #include #include "../../include/moe_gemm_kernels.h" -#include "cutlass_extensions/gemm/collective/collective_mma_array_mixed_input.hpp" #include "cutlass_extensions/gemm_configs.h" #include "cutlass_extensions/weight_only_quant_op.h" @@ -30,26 +27,13 @@ using tensorrt_llm::kernels::cutlass_kernels::GroupedGemmInput; using tensorrt_llm::kernels::cutlass_kernels::TmaWarpSpecializedGroupedGemmInput; template + typename EpilogueScheduleType, cutlass::WeightOnlyQuantOp QuantOp> void sm90_generic_mixed_moe_gemm_kernelLauncher( tensorrt_llm::kernels::cutlass_kernels::GroupedGemmInput inputs, TmaWarpSpecializedGroupedGemmInput hopper_inputs, int sm_count_, size_t* workspace_size); -template -void sm90_generic_mixed_moe_small_k_kernelLauncher( - tensorrt_llm::kernels::cutlass_kernels::GroupedGemmInput - inputs, - TmaWarpSpecializedGroupedGemmInput hopper_inputs, int sm_count_, size_t* workspace_size); - } // namespace cutlass_kernels_oss } // namespace kernels } // namespace tensorrt_llm diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl index b722a317c62..78f91bbd639 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl @@ -19,8 +19,6 @@ #pragma GCC diagnostic ignored "-Wstrict-aliasing" #endif // __GNUC__ -#include - #include "cutlass/epilogue/collective/collective_builder.hpp" #include "cutlass/epilogue/collective/default_epilogue.hpp" #include "cutlass/epilogue/thread/linear_combination.h" @@ -43,14 +41,8 @@ #include "cutlass/util/reference/host/tensor_norm.h" #include "cutlass/util/tensor_view_io.h" #include "cutlass_extensions/compute_occupancy.h" -#include "cutlass_extensions/epilogue/collective/default_epilogue_array_per_token_scale.hpp" -#include "cutlass_extensions/epilogue/collective/sm90_epilogue_array_tma_warpspecialized_mixed_input.hpp" -#include "cutlass_extensions/epilogue/fusion/sm90_ptr_array_per_token_scale_callbacks_tma_warpspecialized.hpp" #include "cutlass_extensions/epilogue_helpers.h" #include "cutlass_extensions/gemm/collective/collective_builder_mixed_input.hpp" -#include "cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_single_warpgroup_persistent.hpp" -#include "cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative_precomputed.hpp" -#include "cutlass_extensions/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong_precomputed.hpp" #include "cutlass_extensions/gemm_configs.h" #ifdef __GNUC__ // Check if the compiler is GCC or Clang @@ -58,7 +50,6 @@ #endif // __GNUC__ #include "moe_gemm_tma_ws_mixed_input_launcher.h" -#include "moe_gemm_tma_ws_mixed_input_prebuild.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/logger.h" @@ -69,83 +60,15 @@ namespace tensorrt_llm { namespace kernels { namespace cutlass_kernels_oss { using namespace tensorrt_llm::kernels::cutlass_kernels; -#ifdef ENABLE_FP4 -using SafeFP4 = Fp4Type; -#else -struct SafeFP4 {}; -#endif namespace tk = tensorrt_llm::common; namespace tkc = tensorrt_llm::cutlass_extensions; using namespace cute; -namespace mixed_input_detail { - -template -struct EpilogueSelector; - -template -struct EpilogueSelector { - using Type = typename tensorrt_llm::cutlass_extensions::epilogue::collective:: - MixedInputSm90TmaEpilogueBuilder< - cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, TileShape, ClusterShape, - cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator, ElementAccumulator, - ElementC, typename cutlass::layout::LayoutTranspose::type*, AlignmentC, ElementD, - typename cutlass::layout::LayoutTranspose::type*, AlignmentD, EpilogueSchedule, - FusionOperation>::CollectiveOp; -}; - -template -struct EpilogueSelector { - using EpilogueLayoutC = typename cutlass::layout::LayoutTranspose::type; - using EpilogueLayoutD = typename cutlass::layout::LayoutTranspose::type; - using Epilogue = cutlass::epilogue::collective::SmemEpilogueArrayPerTokenScale< - TileShape, ElementC, cutlass::detail::TagToStrideC_t, ElementD, - cutlass::detail::TagToStrideC_t, ElementAccumulator, ElementAccumulator>; - using Type = cutlass::epilogue::collective::detail::Sm90TmaWarpSpecializedAdapter; -}; - -template -struct GemmKernelSelector; - -template -struct GemmKernelSelector { - using Type = - cutlass::gemm::kernel::GemmUniversalPrecomputedScheduler; -}; - -template -struct GemmKernelSelector { - using Type = cutlass::gemm::kernel::SingleWarpgroupPersistentGemm< - ProblemShape, CollectiveMainloop, CollectiveEpilogue, CtasPerSm, 3, - RollingRefill ? cutlass::gemm::kernel::SingleWarpgroupPipelineMode::RollingRefill - : cutlass::gemm::kernel::SingleWarpgroupPipelineMode::PrefillAll>; -}; - -} // namespace mixed_input_detail - template -void sm90_generic_mixed_moe_gemm_kernelLauncher_impl( + typename EpilogueScheduleType, cutlass::WeightOnlyQuantOp QuantOp> +void sm90_generic_mixed_moe_gemm_kernelLauncher( GroupedGemmInput inputs, TmaWarpSpecializedGroupedGemmInput hopper_inputs, int sm_count_, size_t* workspace_size) { TLLM_LOG_DEBUG(__PRETTY_FUNCTION__); @@ -179,12 +102,14 @@ void sm90_generic_mixed_moe_gemm_kernelLauncher_impl( using StrideB = cute::remove_pointer_t>; // Scale configuration - constexpr bool use_mxfp4_weight = std::is_same_v; - constexpr int group_size = use_mxfp4_weight ? cutlass::gemm::collective::detail::mxfp4_group_size - : cutlass::gemm::collective::detail::int4_group_size; + constexpr bool use_wfp4a16 = std::is_same_v; + constexpr int group_size = use_wfp4a16 ? cutlass::gemm::collective::detail::mxfp4_group_size + : cutlass::gemm::collective::detail::int4_group_size; + constexpr int PackedScalesNum = get<2>(CTAShape{}) / group_size; using ElementScale = - std::conditional_t; + using ElementScalePacked = cutlass::Array; using LayoutScale = cutlass::layout::RowMajor; // C/D matrix configuration @@ -205,6 +130,8 @@ void sm90_generic_mixed_moe_gemm_kernelLauncher_impl( cutlass::arch::Sm90; // Tag indicating the minimum SM that supports the intended feature using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag using TileShape = CTAShape; // Threadblock-level tile size + using StageCountType = + cutlass::gemm::collective::StageCountAuto; // Stage count maximized based on the tile size using KernelSchedule = std::conditional_t< std::is_same_v, cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong, @@ -213,55 +140,29 @@ void sm90_generic_mixed_moe_gemm_kernelLauncher_impl( std::is_same_v, cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong, cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative>; // Epilogue to launch - constexpr bool use_fused_e8m0_scale = - ScaleMode == cutlass::gemm::collective::MixedInputScaleMode::kPreMmaE8M0; - constexpr bool use_single_warpgroup = - KernelType == tkc::MainloopScheduleType::SINGLE_WARPGROUP_PREFILL || - KernelType == tkc::MainloopScheduleType::SINGLE_WARPGROUP_ROLLING; - constexpr bool use_rolling_refill = - KernelType == tkc::MainloopScheduleType::SINGLE_WARPGROUP_ROLLING; - constexpr int SmallKTileN = cute::size<1>(TileShape{}); - constexpr int SmallKCtasPerSm = SmallKTileN <= 16 ? 5 : (SmallKTileN == 32 ? 4 : 3); - - static_assert(!use_single_warpgroup || use_fused_e8m0_scale, - "The single-warpgroup kernel is only valid for pre-MMA E8M0 scaling."); - static_assert(!use_single_warpgroup || cute::size(ClusterShape{}) == 1, - "The single-warpgroup kernel requires a 1x1x1 cluster."); - static_assert( - !use_single_warpgroup || - (cute::size<0>(TileShape{}) == 128 && cute::size<2>(TileShape{}) == 128 && - (SmallKTileN == 8 || SmallKTileN == 16 || SmallKTileN == 32 || SmallKTileN == 40)), - "Unsupported single-warpgroup tile shape."); - - using FusionOperation = - std::conditional_t, - cutlass::epilogue::fusion::LinearCombination< - ElementD, ElementAccumulator, ElementC, ElementAccumulator>>; - - using CollectiveEpilogue = typename mixed_input_detail::EpilogueSelector< - use_single_warpgroup, TileShape, ClusterShape, ElementAccumulator, ElementC, LayoutC, - AlignmentC, ElementD, LayoutD, AlignmentD, EpilogueSchedule, FusionOperation>::Type; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, TileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator, ElementAccumulator, + ElementC, typename cutlass::layout::LayoutTranspose::type*, AlignmentC, ElementD, + typename cutlass::layout::LayoutTranspose::type*, AlignmentD, + EpilogueSchedule>::CollectiveOp; // =========================================================== MIXED INPUT WITH SCALES // =========================================================================== The Scale // information must get paired with the operand that will be scaled. In this example, B is scaled // so we make a tuple of B's information and the scale information. - using StageCountType = - std::conditional_t, - cutlass::gemm::collective::StageCountAutoCarveout( - sizeof(typename CollectiveEpilogue::SharedStorage))>>; using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilderMixedInput< - ArchTag, OperatorClass, cute::tuple, LayoutB_Transpose*, AlignmentB, - ElementA, LayoutA_Transpose*, AlignmentA, ElementAccumulator, TileShape, ClusterShape, - StageCountType, KernelSchedule, ScaleMode>::CollectiveOp; + ArchTag, OperatorClass, cute::tuple, LayoutB_Transpose*, + AlignmentB, ElementA, LayoutA_Transpose*, AlignmentA, ElementAccumulator, TileShape, + ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelSchedule>::CollectiveOp; - using ProblemShape = cutlass::gemm::GroupProblemShape>; using GemmKernel = - typename mixed_input_detail::GemmKernelSelector::Type; + cutlass::gemm::kernel::GemmUniversal>, + CollectiveMainloop, CollectiveEpilogue>; using GemmGrouped = cutlass::gemm::device::GemmUniversalAdapter; using StrideC = typename GemmKernel::InternalStrideC; @@ -273,62 +174,40 @@ void sm90_generic_mixed_moe_gemm_kernelLauncher_impl( Args arguments; decltype(arguments.epilogue.thread) fusion_args; - if constexpr (use_fused_e8m0_scale) { - fusion_args.token_scale_default = ElementAccumulator(1); - fusion_args.token_scale_ptr_array = inputs.alpha_scales; - } else { - fusion_args.alpha = use_mxfp4_weight ? 1 : 0; - fusion_args.beta = 0; - fusion_args.alpha_ptr = nullptr; - fusion_args.beta_ptr = nullptr; - fusion_args.alpha_ptr_array = use_mxfp4_weight ? nullptr : inputs.alpha_scales; - fusion_args.beta_ptr_array = nullptr; - // One alpha and beta per each group - fusion_args.dAlpha = {cute::_0{}, cute::_0{}, use_mxfp4_weight ? 0 : 1}; - fusion_args.dBeta = {cute::_0{}, cute::_0{}, use_mxfp4_weight ? 0 : 1}; - } + fusion_args.alpha = use_wfp4a16 ? 1 : 0; + fusion_args.beta = 0; + fusion_args.alpha_ptr = nullptr; + fusion_args.beta_ptr = nullptr; + fusion_args.alpha_ptr_array = use_wfp4a16 ? nullptr : inputs.alpha_scales; + fusion_args.beta_ptr_array = nullptr; + // One alpha and beta per each group + fusion_args.dAlpha = {cute::_0{}, cute::_0{}, use_wfp4a16 ? 0 : 1}; + fusion_args.dBeta = {cute::_0{}, cute::_0{}, use_wfp4a16 ? 0 : 1}; cutlass::KernelHardwareInfo hw_info; hw_info.device_id = 0; - hw_info.sm_count = use_single_warpgroup ? sm_count_ * SmallKCtasPerSm : sm_count_; - - if constexpr (use_single_warpgroup) { - arguments = Args{ - cutlass::gemm::GemmUniversalMode::kGrouped, - {inputs.num_experts, hopper_inputs.int4_groupwise_params.shape.problem_shapes, nullptr}, - {reinterpret_cast(hopper_inputs.ptr_weight), - reinterpret_cast(hopper_inputs.stride_weight), - reinterpret_cast(hopper_inputs.ptr_act), - reinterpret_cast(hopper_inputs.stride_act), - reinterpret_cast(hopper_inputs.int4_groupwise_params.ptr_s_a), - reinterpret_cast(hopper_inputs.int4_groupwise_params.stride_s_a), group_size}, - {fusion_args, nullptr, reinterpret_cast(hopper_inputs.stride_c), - reinterpret_cast(hopper_inputs.ptr_d), - reinterpret_cast(hopper_inputs.stride_d), reinterpret_cast(inputs.C), - inputs.n, inputs.n, ElementAccumulator(0)}, - hw_info}; - } else { - arguments = Args{ - cutlass::gemm::GemmUniversalMode::kGrouped, - {inputs.num_experts, hopper_inputs.int4_groupwise_params.shape.problem_shapes, nullptr}, - {reinterpret_cast(hopper_inputs.ptr_weight), - reinterpret_cast(hopper_inputs.stride_weight), - reinterpret_cast(hopper_inputs.ptr_act), - reinterpret_cast(hopper_inputs.stride_act), - reinterpret_cast(hopper_inputs.int4_groupwise_params.ptr_s_a), - reinterpret_cast(hopper_inputs.int4_groupwise_params.stride_s_a), group_size}, - {fusion_args, reinterpret_cast(hopper_inputs.ptr_c), - reinterpret_cast(hopper_inputs.stride_c), - reinterpret_cast(hopper_inputs.ptr_d), - reinterpret_cast(hopper_inputs.stride_d)}, - hw_info}; - } + hw_info.sm_count = sm_count_; + + arguments = Args{ + cutlass::gemm::GemmUniversalMode::kGrouped, + {inputs.num_experts, hopper_inputs.int4_groupwise_params.shape.problem_shapes, nullptr}, + {reinterpret_cast(hopper_inputs.ptr_weight), + reinterpret_cast(hopper_inputs.stride_weight), + reinterpret_cast(hopper_inputs.ptr_act), + reinterpret_cast(hopper_inputs.stride_act), + reinterpret_cast(hopper_inputs.int4_groupwise_params.ptr_s_a), + reinterpret_cast(hopper_inputs.int4_groupwise_params.stride_s_a), group_size}, + {fusion_args, reinterpret_cast(hopper_inputs.ptr_c), + reinterpret_cast(hopper_inputs.stride_c), + reinterpret_cast(hopper_inputs.ptr_d), + reinterpret_cast(hopper_inputs.stride_d)}, + hw_info}; // Optimize tile scheduling for better L2 locality using RasterOrderOptions = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90Params::RasterOrderOptions; - arguments.scheduler.max_swizzle_size = detail::kPrecomputedSchedulerMaxSwizzle; - arguments.scheduler.raster_order = RasterOrderOptions::AlongM; + arguments.scheduler.max_swizzle_size = 2; + arguments.scheduler.raster_order = RasterOrderOptions::Heuristic; assert(group_size == int(inputs.groupwise_quant_group_size)); if (workspace_size != nullptr) { @@ -336,45 +215,6 @@ void sm90_generic_mixed_moe_gemm_kernelLauncher_impl( return; } - if constexpr (use_single_warpgroup) { - TLLM_CHECK_WITH_INFO(inputs.k > 0 && inputs.k % 128 == 0, - "Single-warpgroup GEMM requires K to be a positive multiple of 128."); - if constexpr (use_rolling_refill) { - TLLM_CHECK_WITH_INFO(inputs.k > 384, - "Rolling-refill single-warpgroup GEMM requires K > 384."); - } else { - TLLM_CHECK_WITH_INFO(inputs.k <= 384, "Prefill single-warpgroup GEMM requires K <= 384."); - } - TLLM_CHECK_WITH_INFO(inputs.n > 0 && inputs.n % 128 == 0, - "Single-warpgroup GEMM requires output channels to be 128 aligned."); - TLLM_CHECK_WITH_INFO( - inputs.C != nullptr && reinterpret_cast(inputs.C) % 16 == 0, - "Single-warpgroup GEMM requires a 16B-aligned contiguous output base."); - } - - static constexpr int CurrentTileShapeM = cute::size<0>(TileShape{}); - static constexpr int CurrentTileShapeN = cute::size<1>(TileShape{}); - static constexpr int CurrentClusterShapeM = cute::size<0>(ClusterShape{}); - static constexpr int CurrentClusterShapeN = cute::size<1>(ClusterShape{}); - int64_t const total_routed_tokens = hopper_inputs.precomputed_scheduler_total_routed_tokens; - TLLM_CHECK_WITH_INFO(total_routed_tokens >= 0, - "Precomputed scheduler requires a nonnegative routed token count."); - if (total_routed_tokens == 0) { - return; - } - auto precomputed_workspace = - detail::partition_precomputed_scheduler_workspace( - hopper_inputs, inputs.num_experts, total_routed_tokens, inputs.n, hw_info.sm_count); - arguments.scheduler.precomputed_work_tiles = precomputed_workspace.work_tiles; - if constexpr (use_single_warpgroup) { - arguments.scheduler.precomputed_work_tiles_per_worker = - precomputed_workspace.work_tiles_per_worker; - } - arguments.mainloop.ptr_A_prebuilt_tma_desc = precomputed_workspace.prebuilt_tma_desc_A; - arguments.mainloop.ptr_B_prebuilt_tma_descs = precomputed_workspace.prebuilt_tma_desc_B; - if (gemm.get_workspace_size(arguments) > hopper_inputs.gemm_workspace_size) { TLLM_LOG_ERROR("[Mixed dtype WS grouped GEMM] given workspace size insufficient, %d < %d.", gemm.get_workspace_size(arguments), hopper_inputs.gemm_workspace_size); @@ -399,12 +239,6 @@ void sm90_generic_mixed_moe_gemm_kernelLauncher_impl( throw std::runtime_error("[Mixed dtype WS grouped GEMM] " + err_msg); } - detail::build_precomputed_work_tile_map( - precomputed_workspace, hopper_inputs.int4_groupwise_params.shape.problem_shapes, - inputs.num_experts, total_routed_tokens, inputs.n, gemm.params().mainloop, inputs.stream); - auto run_status = gemm.run(inputs.stream); if (run_status != cutlass::Status::kSuccess) { std::string err_msg = "Failed to run cutlass mixed dtype WS grouped gemm. Error: " + @@ -414,36 +248,6 @@ void sm90_generic_mixed_moe_gemm_kernelLauncher_impl( return; } -template -void sm90_generic_mixed_moe_gemm_kernelLauncher( - GroupedGemmInput inputs, - TmaWarpSpecializedGroupedGemmInput hopper_inputs, int sm_count_, size_t* workspace_size) { - sm90_generic_mixed_moe_gemm_kernelLauncher_impl< - T, WeightType, GemmOutputType, EpilogueTag, CTAShape, ClusterShape, MainloopScheduleType, - EpilogueScheduleType, QuantOp, ScaleMode, tkc::MainloopScheduleType::AUTO>( - inputs, hopper_inputs, sm_count_, workspace_size); -} - -template -void sm90_generic_mixed_moe_small_k_kernelLauncher( - GroupedGemmInput inputs, - TmaWarpSpecializedGroupedGemmInput hopper_inputs, int sm_count_, size_t* workspace_size) { - static_assert(KernelType == tkc::MainloopScheduleType::SINGLE_WARPGROUP_PREFILL || - KernelType == tkc::MainloopScheduleType::SINGLE_WARPGROUP_ROLLING, - "Small-K launcher requires a single-warpgroup schedule."); - sm90_generic_mixed_moe_gemm_kernelLauncher_impl< - T, WeightType, GemmOutputType, EpilogueTag, CTAShape, ClusterShape, MainloopScheduleType, - EpilogueScheduleType, QuantOp, ScaleMode, KernelType>(inputs, hopper_inputs, sm_count_, - workspace_size); -} - } // namespace cutlass_kernels_oss } // namespace kernels } // namespace tensorrt_llm diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_prebuild.h b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_prebuild.h deleted file mode 100644 index 893506b0b26..00000000000 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_prebuild.h +++ /dev/null @@ -1,513 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include -#include - -#include "../../include/moe_gemm_kernels.h" -#include "cute/tensor.hpp" -#include "cutlass/gemm/group_array_problem_shape.hpp" -#include "cutlass/kernel_hardware_info.hpp" -#include "cutlass_extensions/gemm/kernel/sm90_tile_scheduler_group_precomputed.hpp" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" - -namespace tensorrt_llm { -namespace kernels { -namespace cutlass_kernels_oss { -namespace detail { - -using namespace cute; - -using TmaWarpSpecializedGroupedGemmInput = - tensorrt_llm::kernels::cutlass_kernels::TmaWarpSpecializedGroupedGemmInput; -using PrecomputedWorkTileCodec = cutlass::gemm::kernel::detail::PrecomputedGroupWorkTile; - -static constexpr int kPrecomputedSchedulerThreads = 128; -static constexpr int kPrecomputedSchedulerMaxSwizzle = 2; -static constexpr uint64_t kPrecomputedSchedulerSentinelTiles = 4096; - -CUTLASS_HOST_DEVICE uint64_t div_round_up(uint64_t value, uint64_t divisor) { - return (value + divisor - 1) / divisor; -} - -CUTLASS_HOST_DEVICE uint64_t round_up_to_multiple(uint64_t value, uint64_t multiple) { - return div_round_up(value, multiple) * multiple; -} - -inline size_t align_bytes(size_t value, size_t alignment) { - return ((value + alignment - 1) / alignment) * alignment; -} - -inline int log_swizzle_size(uint64_t problem_blocks_m, uint64_t problem_blocks_n, - int max_swizzle_size) { - int const min_cta_dim = - static_cast(problem_blocks_m < problem_blocks_n ? problem_blocks_m : problem_blocks_n); - if (max_swizzle_size >= 8 && min_cta_dim >= 6) { - return 3; - } - if (max_swizzle_size >= 4 && min_cta_dim >= 3) { - return 2; - } - if (max_swizzle_size >= 2 && min_cta_dim >= 2) { - return 1; - } - return 0; -} - -template -inline uint64_t max_work_tiles_from_total_tokens(int num_experts, int64_t total_routed_tokens, - int64_t channels, int max_swizzle_size) { - if (num_experts <= 0 || total_routed_tokens <= 0 || channels <= 0) { - return 0; - } - - int const swizzle_log = - log_swizzle_size(static_cast(total_routed_tokens), 1, max_swizzle_size); - uint64_t const swizzle = uint64_t(1) << swizzle_log; - uint64_t const channel_multiple = swizzle * uint64_t(ClusterShapeM); - uint64_t const token_tile_group = swizzle * uint64_t(ClusterShapeN); - uint64_t const tokens_per_padded_group = uint64_t(TileShapeN) * token_tile_group; - uint64_t const channel_tiles = - div_round_up(static_cast(channels), uint64_t(TileShapeM)); - uint64_t const padded_channel_tiles = round_up_to_multiple(channel_tiles, channel_multiple); - uint64_t const total_tokens = static_cast(total_routed_tokens); - uint64_t const nonempty_experts = - total_tokens < uint64_t(num_experts) ? total_tokens : uint64_t(num_experts); - uint64_t const extra_tokens = total_tokens - nonempty_experts; - uint64_t const max_token_tile_rows = - token_tile_group * (nonempty_experts + extra_tokens / tokens_per_padded_group); - - return padded_channel_tiles * max_token_tile_rows; -} - -inline size_t precomputed_scheduler_workspace_size(int num_experts, int64_t total_routed_tokens, - int64_t max_channels) { - uint64_t const regular_max_work_tiles = max_work_tiles_from_total_tokens<64, 16, 2, 2>( - num_experts, total_routed_tokens, max_channels, kPrecomputedSchedulerMaxSwizzle); - uint64_t const single_warpgroup_max_work_tiles = max_work_tiles_from_total_tokens<128, 8, 1, 1>( - num_experts, total_routed_tokens, max_channels, kPrecomputedSchedulerMaxSwizzle); - uint64_t const max_work_tiles = regular_max_work_tiles > single_warpgroup_max_work_tiles - ? regular_max_work_tiles - : single_warpgroup_max_work_tiles; - // Chunk-major storage rounds every worker chunk up independently and appends - // one sentinel per worker. Its capacity is strictly below max_tiles + 2 * workers. - uint64_t const work_tile_capacity = max_work_tiles + 2 * kPrecomputedSchedulerSentinelTiles; - - size_t bytes = 0; - bytes += align_bytes(size_t(work_tile_capacity) * sizeof(uint64_t), 128); - bytes += align_bytes(sizeof(cute::TmaDescriptor), 128); - bytes += - align_bytes(size_t(num_experts > 0 ? num_experts : 1) * sizeof(cute::TmaDescriptor), 128); - return bytes; -} - -struct PrecomputedSchedulerWorkspace { - uint64_t* work_tiles = nullptr; - cute::TmaDescriptor* prebuilt_tma_desc_A = nullptr; - cute::TmaDescriptor* prebuilt_tma_desc_B = nullptr; - size_t required_bytes = 0; - dim3 gemm_grid_shape = dim3(1, 1, 1); - uint32_t work_tiles_per_worker = 0; -}; - -template -inline PrecomputedSchedulerWorkspace partition_precomputed_scheduler_workspace( - TmaWarpSpecializedGroupedGemmInput const& hopper_inputs, int num_experts, - int64_t total_routed_tokens, int64_t channels, int sm_count) { - using ProblemShape = TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::ProblemShapeInt; - using Scheduler = - cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90GroupPrecomputed; - using SchedulerParams = typename Scheduler::Params; - - cutlass::KernelHardwareInfo hw_info; - hw_info.device_id = 0; - hw_info.sm_count = sm_count; - cutlass::gemm::GemmCoord cluster_shape(ClusterShapeM, ClusterShapeN, 1); - dim3 const problem_blocks = - SchedulerParams::get_tiled_cta_shape_mnl(cluster_shape, static_cast(sm_count), 1); - dim3 const gemm_grid_shape = SchedulerParams::get_grid_shape( - problem_blocks, cluster_shape, hw_info, kPrecomputedSchedulerMaxSwizzle, - SchedulerParams::RasterOrderOptions::AlongM, true); - - uint64_t const max_work_tiles = - max_work_tiles_from_total_tokens( - num_experts, total_routed_tokens, channels, kPrecomputedSchedulerMaxSwizzle); - uint64_t const sentinel_count = - uint64_t(gemm_grid_shape.x) * uint64_t(gemm_grid_shape.y) * uint64_t(gemm_grid_shape.z); - TLLM_CHECK_WITH_INFO(sentinel_count > 0, - "Precomputed scheduler requires at least one logical worker."); - TLLM_CHECK_WITH_INFO(gemm_grid_shape.z == 1, - "Precomputed grouped scheduler work-map requires a 2D launch grid."); - uint32_t work_tiles_per_worker = 0; - uint64_t work_tile_capacity = max_work_tiles + sentinel_count; - if constexpr (ChunkMajorWorkMap) { - TLLM_CHECK_WITH_INFO( - sentinel_count <= kPrecomputedSchedulerSentinelTiles, - "Single-warpgroup logical worker count exceeds precomputed scheduler workspace bound."); - uint64_t const work_tiles_per_worker_u64 = - (max_work_tiles + sentinel_count - 1) / sentinel_count + 1; - TLLM_CHECK_WITH_INFO(work_tiles_per_worker_u64 <= uint64_t(0xffffffffu), - "Precomputed scheduler worker chunk exceeds uint32_t index range."); - work_tiles_per_worker = static_cast(work_tiles_per_worker_u64); - work_tile_capacity = sentinel_count * uint64_t(work_tiles_per_worker); - } - TLLM_CHECK_WITH_INFO(work_tile_capacity <= uint64_t(0xffffffffu), - "Precomputed scheduler work-map exceeds uint32_t index range."); - - size_t const work_tiles_bytes = align_bytes(size_t(work_tile_capacity) * sizeof(uint64_t), 128); - size_t const prebuilt_a_bytes = align_bytes(sizeof(cute::TmaDescriptor), 128); - size_t const prebuilt_b_bytes = - align_bytes(size_t(num_experts > 0 ? num_experts : 1) * sizeof(cute::TmaDescriptor), 128); - size_t const required_bytes = work_tiles_bytes + prebuilt_a_bytes + prebuilt_b_bytes; - - TLLM_CHECK_WITH_INFO( - hopper_inputs.precomputed_scheduler_workspace != nullptr, - "Precomputed scheduler workspace must be configured for mixed dtype TMA WS GEMM."); - TLLM_CHECK_WITH_INFO( - required_bytes <= hopper_inputs.precomputed_scheduler_workspace_size, - "Precomputed scheduler workspace is too small for selected mixed dtype TMA WS GEMM config."); - TLLM_CHECK_WITH_INFO(num_experts <= int(PrecomputedWorkTileCodec::ExpertMask + 1), - "Precomputed scheduler work-map expert index exceeds packed limit."); - - auto* base = hopper_inputs.precomputed_scheduler_workspace; - PrecomputedSchedulerWorkspace workspace; - workspace.work_tiles = reinterpret_cast(base); - workspace.prebuilt_tma_desc_A = reinterpret_cast(base + work_tiles_bytes); - workspace.prebuilt_tma_desc_B = - reinterpret_cast(base + work_tiles_bytes + prebuilt_a_bytes); - workspace.required_bytes = required_bytes; - workspace.gemm_grid_shape = gemm_grid_shape; - workspace.work_tiles_per_worker = work_tiles_per_worker; - return workspace; -} - -template -__device__ __forceinline__ uint64_t make_work_tile_static(uint64_t global_linear_idx, - uint64_t local_linear_idx, int group_idx, - uint64_t problem_blocks_m, - int swizzle_log, int gemm_grid_x, - int gemm_grid_y) { - uint64_t const cluster_shape_major = uint64_t(ClusterShapeM); - uint64_t const cluster_shape_minor = uint64_t(ClusterShapeN); - uint64_t const total_grid_size = uint64_t(gemm_grid_x) * uint64_t(gemm_grid_y); - uint64_t const worker_id = total_grid_size == 0 ? 0 : global_linear_idx % total_grid_size; - uint64_t const cluster_minor_offset = worker_id % uint64_t(gemm_grid_y); - - uint64_t const blk_per_grid_dim = local_linear_idx / cluster_shape_minor; - uint64_t const cluster_id = blk_per_grid_dim / cluster_shape_major; - uint64_t const cluster_major_offset = blk_per_grid_dim % cluster_shape_major; - - uint64_t const swizzle = uint64_t(1) << swizzle_log; - uint64_t const offset = cluster_id & (swizzle - 1); - uint64_t const extra = cluster_id >> swizzle_log; - uint64_t const curr_group_cluster_blk_major = problem_blocks_m / cluster_shape_major; - uint64_t const cluster_idx_minor_div_swizzle = extra / curr_group_cluster_blk_major; - uint64_t const cluster_idx_major = extra % curr_group_cluster_blk_major; - uint64_t const cluster_idx_minor = cluster_idx_minor_div_swizzle * swizzle + offset; - - uint64_t const minor_work_idx = cluster_idx_minor * cluster_shape_minor + cluster_minor_offset; - uint64_t const major_work_idx = cluster_idx_major * cluster_shape_major + cluster_major_offset; - return PrecomputedWorkTileCodec::pack(major_work_idx, minor_work_idx, uint64_t(group_idx)); -} - -struct PrecomputedGroupInfo { - uint64_t problem_blocks_m = 0; - uint64_t group_tiles = 0; -}; - -template -__device__ __forceinline__ PrecomputedGroupInfo get_group_info_static(Problem const& problem, - int swizzle_log) { - uint64_t const ctas_along_m = - (uint64_t(cute::get<0>(problem)) + uint64_t(TileShapeM) - 1) / uint64_t(TileShapeM); - uint64_t const ctas_along_n = - (uint64_t(cute::get<1>(problem)) + uint64_t(TileShapeN) - 1) / uint64_t(TileShapeN); - uint64_t const swizzle = uint64_t(1) << swizzle_log; - uint64_t const m_multiple = swizzle * uint64_t(ClusterShapeM); - uint64_t const n_multiple = swizzle * uint64_t(ClusterShapeN); - uint64_t const problem_blocks_m = round_up_to_multiple(ctas_along_m, m_multiple); - uint64_t const problem_blocks_n = round_up_to_multiple(ctas_along_n, n_multiple); - - if (problem_blocks_m > uint64_t(PrecomputedWorkTileCodec::ChannelMask + 1) || - problem_blocks_n > uint64_t(PrecomputedWorkTileCodec::TokenMask + 1)) { - asm volatile("trap;"); - } - - return {problem_blocks_m, problem_blocks_m * problem_blocks_n}; -} - -static constexpr int kPrebuiltTmaDescriptorScratchCount = 2; -static constexpr size_t kPrebuiltTmaDescriptorScratchBytes = - kPrebuiltTmaDescriptorScratchCount * sizeof(cute::TmaDescriptor); -static constexpr int kPrebuiltTmaDescriptorSlotA = 0; -static constexpr int kPrebuiltTmaDescriptorSlotB = 1; -static constexpr int kPrebuiltTmaDescriptorWarpA = 0; -static constexpr int kPrebuiltTmaDescriptorWarpB = 1; - -CUTE_DEVICE void publish_prebuilt_tma_descriptor(cute::TmaDescriptor const* gmem_desc_ptr, - cute::TmaDescriptor& smem_desc, - int publisher_warp) { - if ((threadIdx.x >> 5) == publisher_warp) { - __syncwarp(); - if (cute::elect_one_sync()) { - cute::tma_desc_commit_group(); - cute::tma_desc_wait_group(); - } - cute::tma_descriptor_cp_fence_release(gmem_desc_ptr, smem_desc); - __syncwarp(); - } -} - -template -__device__ __forceinline__ void build_prebuilt_tma_descriptors( - MainloopParams const& mainloop_params, Problem const& problem, int group, - cute::TmaDescriptor* smem_tma_desc, cute::TmaDescriptor* prebuilt_tma_desc_A, - cute::TmaDescriptor* prebuilt_tma_desc_B) { - if (group == 0) { - cute::TmaDescriptor& smem_desc = smem_tma_desc[kPrebuiltTmaDescriptorSlotA]; - if (threadIdx.x == kPrebuiltTmaDescriptorWarpA * 32) { - constexpr int MaxTensorRank = 5; - cute::array prob_shape_A = {1, 1, 1, 1, 1}; - cute::array prob_stride_A = {0, 0, 0, 0, 0}; - using PtrA = std::remove_reference_t; - PtrA ptr_A = nullptr; - uint32_t const M = static_cast(cute::get<0>(problem)); - uint32_t const K = static_cast(cute::get<2>(problem)); - auto dA_group = mainloop_params.ptr_dA[group]; - auto stride_m = cute::get<0>(dA_group); - auto stride_k = cute::get<1>(dA_group); - int64_t const term_m = static_cast(M) * static_cast(stride_m); - int64_t const term_k = static_cast(K) * static_cast(stride_k); - int64_t const stride_l = term_m > term_k ? term_m : term_k; - auto full_layout = - make_layout(make_shape(M, K, static_cast(mainloop_params.num_groups)), - cute::make_stride(stride_m, stride_k, stride_l)); - Tensor tensor_a = make_tensor(ptr_A, full_layout); - - smem_desc = *mainloop_params.tma_load_a.get_tma_descriptor(); - cute::tma_descriptor_replace_addr_in_shared_mem(smem_desc, mainloop_params.ptr_A[0]); - cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_a, tensor_a, prob_shape_A, - prob_stride_A); - - using ElementA = std::remove_cv_t>; - for (uint64_t& stride : prob_stride_A) { - stride = (stride * cutlass::sizeof_bits::value) / 8; - } - cute::tma_descriptor_replace_dims_strides_in_shared_mem(smem_desc, prob_shape_A, - prob_stride_A); - } - publish_prebuilt_tma_descriptor(&prebuilt_tma_desc_A[0], smem_desc, - kPrebuiltTmaDescriptorWarpA); - } - - if (cute::get<1>(problem) == 0) { - return; - } - - { - cute::TmaDescriptor& smem_desc = smem_tma_desc[kPrebuiltTmaDescriptorSlotB]; - if (threadIdx.x == kPrebuiltTmaDescriptorWarpB * 32) { - constexpr int MaxTensorRank = 5; - cute::array prob_shape_B = {1, 1, 1, 1, 1}; - cute::array prob_stride_B = {0, 0, 0, 0, 0}; - using PtrB = std::remove_reference_t; - PtrB ptr_B = nullptr; - uint32_t const N = static_cast(cute::get<1>(problem)); - uint32_t const K = static_cast(cute::get<2>(problem)); - auto dB_group = mainloop_params.ptr_dB[group]; - auto stride_n = cute::get<0>(dB_group); - auto stride_k = cute::get<1>(dB_group); - auto full_layout = make_layout(make_shape(N, K, uint32_t(1)), - cute::make_stride(stride_n, stride_k, int64_t(0))); - Tensor tensor_b = make_tensor(ptr_B, full_layout); - - smem_desc = *mainloop_params.tma_load_b.get_tma_descriptor(); - cute::tma_descriptor_replace_addr_in_shared_mem(smem_desc, mainloop_params.ptr_B[group]); - cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_b, tensor_b, prob_shape_B, - prob_stride_B); - - using ElementB = std::remove_cv_t>; - for (uint64_t& stride : prob_stride_B) { - stride = (stride * cutlass::sizeof_bits::value) / 8; - } - cute::tma_descriptor_replace_dims_strides_in_shared_mem(smem_desc, prob_shape_B, - prob_stride_B); - } - publish_prebuilt_tma_descriptor(&prebuilt_tma_desc_B[group], smem_desc, - kPrebuiltTmaDescriptorWarpB); - } -} - -template -__global__ void build_precomputed_work_tile_map_kernel( - Problem const* problem_shapes, int groups, int swizzle_log, int gemm_grid_x, int gemm_grid_y, - uint32_t work_tiles_per_worker, uint64_t* work_tiles, MainloopParams mainloop_params, - cute::TmaDescriptor* prebuilt_tma_desc_A, cute::TmaDescriptor* prebuilt_tma_desc_B) { - int const tid = threadIdx.x; - uint64_t const total_grid_size = uint64_t(gemm_grid_x) * uint64_t(gemm_grid_y); - - if (groups <= 0) { - if (blockIdx.x == 0) { - for (uint64_t i = uint64_t(tid); i < total_grid_size; i += uint64_t(blockDim.x)) { - uint64_t const storage_idx = ChunkMajorWorkMap ? i * uint64_t(work_tiles_per_worker) : i; - work_tiles[storage_idx] = PrecomputedWorkTileCodec::Invalid; - } - } - return; - } - - int const group = int(blockIdx.x); - if (group >= groups) { - return; - } - - extern __shared__ __align__(64) unsigned char shared_storage[]; - cute::TmaDescriptor* smem_tma_desc = reinterpret_cast(shared_storage); - unsigned long long* prefix_partials = - reinterpret_cast(shared_storage + kPrebuiltTmaDescriptorScratchBytes); - unsigned long long* total_partials = nullptr; - unsigned long long* group_info_storage = nullptr; - if constexpr (ChunkMajorWorkMap) { - total_partials = prefix_partials + blockDim.x; - group_info_storage = total_partials + blockDim.x; - } else { - group_info_storage = prefix_partials + blockDim.x; - } - - if (tid == 0) { - PrecomputedGroupInfo const info = - get_group_info_static( - problem_shapes[group], swizzle_log); - group_info_storage[0] = static_cast(info.problem_blocks_m); - group_info_storage[1] = static_cast(info.group_tiles); - } - - build_prebuilt_tma_descriptors(mainloop_params, problem_shapes[group], group, smem_tma_desc, - prebuilt_tma_desc_A, prebuilt_tma_desc_B); - - uint64_t prefix_sum = 0; - if constexpr (ChunkMajorWorkMap) { - uint64_t total_sum = 0; - for (int scan_group = tid; scan_group < groups; scan_group += blockDim.x) { - PrecomputedGroupInfo const info = - get_group_info_static( - problem_shapes[scan_group], swizzle_log); - total_sum += info.group_tiles; - if (scan_group < group) { - prefix_sum += info.group_tiles; - } - } - total_partials[tid] = static_cast(total_sum); - } else { - for (int prefix_group = tid; prefix_group < group; prefix_group += blockDim.x) { - PrecomputedGroupInfo const info = - get_group_info_static( - problem_shapes[prefix_group], swizzle_log); - prefix_sum += info.group_tiles; - } - } - - prefix_partials[tid] = static_cast(prefix_sum); - __syncthreads(); - - for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) { - if (tid < offset) { - prefix_partials[tid] += prefix_partials[tid + offset]; - if constexpr (ChunkMajorWorkMap) { - total_partials[tid] += total_partials[tid + offset]; - } - } - __syncthreads(); - } - - uint64_t const group_start = static_cast(prefix_partials[0]); - uint64_t const problem_blocks_m = static_cast(group_info_storage[0]); - uint64_t const group_tiles = static_cast(group_info_storage[1]); - uint64_t total_tiles = 0; - uint64_t tiles_per_worker = 0; - if constexpr (ChunkMajorWorkMap) { - total_tiles = static_cast(total_partials[0]); - tiles_per_worker = total_tiles == 0 ? 1 : (total_tiles + total_grid_size - 1) / total_grid_size; - } - - for (uint64_t local_tile = uint64_t(tid); local_tile < group_tiles; - local_tile += uint64_t(blockDim.x)) { - uint64_t const global_tile = group_start + local_tile; - uint64_t storage_idx = global_tile; - if constexpr (ChunkMajorWorkMap) { - uint64_t const worker_idx = global_tile / tiles_per_worker; - uint64_t const worker_tile_idx = global_tile % tiles_per_worker; - storage_idx = worker_idx * uint64_t(work_tiles_per_worker) + worker_tile_idx; - } - work_tiles[storage_idx] = make_work_tile_static( - global_tile, local_tile, group, problem_blocks_m, swizzle_log, gemm_grid_x, gemm_grid_y); - } - - if (group == groups - 1) { - [[maybe_unused]] uint64_t const sentinel_start = group_start + group_tiles; - for (uint64_t i = uint64_t(tid); i < total_grid_size; i += uint64_t(blockDim.x)) { - if constexpr (ChunkMajorWorkMap) { - uint64_t const worker_start = i * tiles_per_worker; - uint64_t const worker_tile_count = - worker_start < total_tiles - ? ((total_tiles - worker_start < tiles_per_worker) ? total_tiles - worker_start - : tiles_per_worker) - : 0; - work_tiles[i * uint64_t(work_tiles_per_worker) + worker_tile_count] = - PrecomputedWorkTileCodec::Invalid; - } else { - work_tiles[sentinel_start + i] = PrecomputedWorkTileCodec::Invalid; - } - } - } -} - -template -inline void build_precomputed_work_tile_map(PrecomputedSchedulerWorkspace const& workspace, - Problem const* problem_shapes, int groups, - int64_t total_routed_tokens, int64_t channels, - MainloopParams const& mainloop_params, - cudaStream_t stream) { - uint64_t const max_work_tiles = - max_work_tiles_from_total_tokens( - groups, total_routed_tokens, channels, kPrecomputedSchedulerMaxSwizzle); - int const swizzle_log = log_swizzle_size(max_work_tiles, 1, kPrecomputedSchedulerMaxSwizzle); - dim3 const scheduler_grid(groups > 0 ? groups : 1); - size_t const scheduler_smem = - kPrebuiltTmaDescriptorScratchBytes + - size_t((ChunkMajorWorkMap ? kPrecomputedSchedulerThreads * 2 : kPrecomputedSchedulerThreads) + - 2) * - sizeof(unsigned long long); - build_precomputed_work_tile_map_kernel - <<>>( - problem_shapes, groups, swizzle_log, workspace.gemm_grid_shape.x, - workspace.gemm_grid_shape.y, workspace.work_tiles_per_worker, workspace.work_tiles, - mainloop_params, workspace.prebuilt_tma_desc_A, workspace.prebuilt_tma_desc_B); - TLLM_CUDA_CHECK(cudaPeekAtLastError()); -} - -} // namespace detail -} // namespace cutlass_kernels_oss -} // namespace kernels -} // namespace tensorrt_llm diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp4.cu b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp4.cu index d5dd4762e56..c1d40e33ac8 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp4.cu +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp4.cu @@ -18,6 +18,6 @@ namespace tensorrt_llm::kernels::cutlass_kernels { #if defined(ENABLE_BF16) && defined(ENABLE_FP4) -template class MoeGemmRunner<__nv_bfloat16, Fp4Type, __nv_bfloat16>; +template class MoeGemmRunner<__nv_bfloat16, __nv_fp4_e2m1, __nv_bfloat16>; #endif } // namespace tensorrt_llm::kernels::cutlass_kernels diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp16_fp4.cu b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp16_fp4.cu index 08371ed36e5..ce4b57cc696 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp16_fp4.cu +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp16_fp4.cu @@ -18,6 +18,6 @@ namespace tensorrt_llm::kernels::cutlass_kernels { #if defined(ENABLE_FP4) -template class MoeGemmRunner; +template class MoeGemmRunner; #endif } // namespace tensorrt_llm::kernels::cutlass_kernels diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp4_fp4.cu b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp4_fp4.cu index 7ce6d044d75..a8c11e06928 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp4_fp4.cu +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp4_fp4.cu @@ -18,9 +18,9 @@ namespace tensorrt_llm::kernels::cutlass_kernels { #ifdef ENABLE_FP4 -template class MoeGemmRunner; +template class MoeGemmRunner<__nv_fp4_e2m1, __nv_fp4_e2m1, half>; #ifdef ENABLE_BF16 -template class MoeGemmRunner; +template class MoeGemmRunner<__nv_fp4_e2m1, __nv_fp4_e2m1, __nv_bfloat16>; #endif #endif } // namespace tensorrt_llm::kernels::cutlass_kernels diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp8_fp4.cu b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp8_fp4.cu index 04939f2fba5..6bc740c5fac 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp8_fp4.cu +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp8_fp4.cu @@ -18,21 +18,9 @@ namespace tensorrt_llm::kernels::cutlass_kernels { #ifdef ENABLE_FP4 -template class MoeGemmRunner<__nv_fp8_e4m3, Fp4Type, half>; -template class MoeGemmRunner<__nv_fp8_e4m3, Fp4Type, half, half, false, - Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0>; -template class MoeGemmRunner<__nv_fp8_e4m3, Fp4Type, half, half, false, - Sm90Wfp4Afp8ScaleMode::kPostMmaFp8Act>; -template class MoeGemmRunner<__nv_fp8_e4m3, Fp4Type, half, half, false, - Sm90Wfp4Afp8ScaleMode::kPostMmaMxfp8Act>; +template class MoeGemmRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, half>; #ifdef ENABLE_BF16 -template class MoeGemmRunner<__nv_fp8_e4m3, Fp4Type, __nv_bfloat16>; -template class MoeGemmRunner<__nv_fp8_e4m3, Fp4Type, __nv_bfloat16, __nv_bfloat16, false, - Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0>; -template class MoeGemmRunner<__nv_fp8_e4m3, Fp4Type, __nv_bfloat16, __nv_bfloat16, false, - Sm90Wfp4Afp8ScaleMode::kPostMmaFp8Act>; -template class MoeGemmRunner<__nv_fp8_e4m3, Fp4Type, __nv_bfloat16, __nv_bfloat16, false, - Sm90Wfp4Afp8ScaleMode::kPostMmaMxfp8Act>; +template class MoeGemmRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, __nv_bfloat16>; #endif #endif } // namespace tensorrt_llm::kernels::cutlass_kernels diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_mixed_utils.cu b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_mixed_utils.cu index b32ce1460ce..8b146366c78 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_mixed_utils.cu +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_mixed_utils.cu @@ -30,76 +30,30 @@ __global__ void interleave_fp4_weights_for_sm90_mixed_gemm_kernel(uint8_t* fp4_w int lane_id = threadIdx.x; int row_id = block_id / 8 * 16 + block_id % 8; - int mma_id = lane_id / 4; + int mma_id = lane_id / 8; int dst_row_id = row_id + (mma_id % 2) * 8; - int interleaved_lane_id = lane_id / 8 * 16 + (lane_id % 4) * 4; + int interleaved_lane_id = lane_id / 16 * 16 + (lane_id % 4) * 4 + (lane_id % 8) / 4 * 2; - int col_id = partition_id * 32 + mma_id * 8 + lane_id % 4; + int col_id = partition_id * 32 + lane_id; int dst_col_id = partition_id * 32 + interleaved_lane_id; - int first_fp4_id = row_id * cols / 2 + col_id; - int second_fp4_id = (row_id + 8) * cols / 2 + col_id; - int third_fp4_id = first_fp4_id + 4; - int fourth_fp4_id = second_fp4_id + 4; - - uint32_t fp4x8_raw = 0; - uint8_t* fp4x2 = reinterpret_cast(&fp4x8_raw); - fp4x2[0] = fp4_weight[first_fp4_id]; - fp4x2[1] = fp4_weight[second_fp4_id]; - fp4x2[2] = fp4_weight[third_fp4_id]; - fp4x2[3] = fp4_weight[fourth_fp4_id]; - - uint32_t fp4x8_interleaved = 0; - uint32_t mask; - - mask = 0b00000000000000000000000010000000; - fp4x8_interleaved |= (fp4x8_raw & mask) << 24; - mask = 0b00000000000000000000000001110000; - fp4x8_interleaved |= (fp4x8_raw & mask) << 18; - mask = 0b00000000000000000000000000001000; - fp4x8_interleaved |= (fp4x8_raw & mask) << 12; - mask = 0b00000000000000000000000000000111; - fp4x8_interleaved |= (fp4x8_raw & mask) << 6; - - mask = 0b00000000000000001000000000000000; - fp4x8_interleaved |= (fp4x8_raw & mask) << 13; - mask = 0b00000000000000000111000000000000; - fp4x8_interleaved |= (fp4x8_raw & mask) << 7; - mask = 0b00000000000000000000100000000000; - fp4x8_interleaved |= (fp4x8_raw & mask) << 1; - mask = 0b00000000000000000000011100000000; - fp4x8_interleaved |= (fp4x8_raw & mask) >> 5; - - mask = 0b00000000100000000000000000000000; - fp4x8_interleaved |= (fp4x8_raw & mask) << 2; - mask = 0b00000000011100000000000000000000; - fp4x8_interleaved |= (fp4x8_raw & mask) >> 4; - mask = 0b00000000000010000000000000000000; - fp4x8_interleaved |= (fp4x8_raw & mask) >> 10; - mask = 0b00000000000001110000000000000000; - fp4x8_interleaved |= (fp4x8_raw & mask) >> 16; - - mask = 0b10000000000000000000000000000000; - fp4x8_interleaved |= (fp4x8_raw & mask) >> 1; - mask = 0b00010000000000000000000000000000; - fp4x8_interleaved |= (fp4x8_raw & mask) << 1; - mask = 0b01100000000000000000000000000000; - fp4x8_interleaved |= (fp4x8_raw & mask) >> 3; - mask = 0b00001000000000000000000000000000; - fp4x8_interleaved |= (fp4x8_raw & mask) >> 13; - mask = 0b00000001000000000000000000000000; - fp4x8_interleaved |= (fp4x8_raw & mask) >> 11; - mask = 0b00000110000000000000000000000000; - fp4x8_interleaved |= (fp4x8_raw & mask) >> 15; + int index_a = row_id * cols / 2 + col_id; + int index_b = (row_id + 8) * cols / 2 + col_id; + + uint8_t fp4x2_a = fp4_weight[index_a]; + uint8_t fp4x2_b = fp4_weight[index_b]; + + uint8_t fp4_temp_a = (fp4x2_a & 0xF0U) >> 4; + uint8_t fp4_temp_b = (fp4x2_b & 0x0FU) << 4; + + fp4x2_a = (fp4x2_a & 0x0FU) | fp4_temp_b; + fp4x2_b = (fp4x2_b & 0xF0U) | fp4_temp_a; int dst_id = dst_row_id * cols / 2 + dst_col_id; - uint8_t* fp4x2_interleaved = reinterpret_cast(&fp4x8_interleaved); - fp4_weight_interleaved[dst_id] = fp4x2_interleaved[0]; - fp4_weight_interleaved[dst_id + 1] = fp4x2_interleaved[1]; - fp4_weight_interleaved[dst_id + 2] = fp4x2_interleaved[2]; - fp4_weight_interleaved[dst_id + 3] = fp4x2_interleaved[3]; + fp4_weight_interleaved[dst_id] = fp4x2_a; + fp4_weight_interleaved[dst_id + 1] = fp4x2_b; } } } @@ -137,69 +91,16 @@ __global__ void interleave_int4_weights_for_sm90_mixed_gemm_kernel(uint8_t* int4 } } -__device__ __forceinline__ uint32_t preprocess_fp4x8_signs_for_fp8(uint32_t fp4x8) { - uint32_t const em = fp4x8 & 0x77777777U; - uint32_t const signs = ((fp4x8 & 0x00000008U) << 4U) | ((fp4x8 & 0x00000080U) << 8U) | - ((fp4x8 & 0x00000800U) << 12U) | ((fp4x8 & 0x00008000U) << 16U) | - ((fp4x8 & 0x00080000U) >> 16U) | ((fp4x8 & 0x00800000U) >> 12U) | - ((fp4x8 & 0x08000000U) >> 8U) | ((fp4x8 & 0x80000000U) >> 4U); - return em | signs; -} - -__global__ void interleave_fp4_fp8_weights_for_sm90_mixed_gemm_kernel( - uint8_t* fp4_weight, uint8_t* fp4_weight_interleaved, int const rows, int const cols) { - uint16_t* uint16_ptr = reinterpret_cast(fp4_weight); - uint16_t* uint16_interleaved_ptr = reinterpret_cast(fp4_weight_interleaved); - - for (int block_id = blockIdx.x; block_id < rows / 2; block_id += gridDim.x) { - for (int partition_id = threadIdx.y; partition_id < cols / 64; partition_id += blockDim.y) { - int lane_id = threadIdx.x; - - int row_id = block_id / 8 * 16 + block_id % 8; - int dst_row_id = row_id + (lane_id % 8) / 4 * 8; - - int mma_id = lane_id / 8; - int interleaved_lane_id = mma_id * 8 + lane_id % 4 * 2; - - int col_id = partition_id * 16 + lane_id; - int dst_col_id = partition_id * 16 + interleaved_lane_id; - - int src_id_a = row_id * cols / 4 + col_id; - int src_id_b = (row_id + 8) * cols / 4 + col_id; - - uint16_t packed_4b_a = uint16_ptr[src_id_a]; - uint16_t packed_4b_b = uint16_ptr[src_id_b]; - - uint32_t fp4x8 = uint32_t(packed_4b_a) | (uint32_t(packed_4b_b) << 16U); - fp4x8 = preprocess_fp4x8_signs_for_fp8(fp4x8); - packed_4b_a = uint16_t(fp4x8); - packed_4b_b = uint16_t(fp4x8 >> 16U); - - int dst_id = dst_row_id * cols / 4 + dst_col_id; - uint16_interleaved_ptr[dst_id] = packed_4b_a; - uint16_interleaved_ptr[dst_id + 1] = packed_4b_b; - } - } -} - ///////////////////////////////////////////////////////////////////////////////////////////////////////// void interleave_fp4_weights_for_sm90_mixed_gemm(uint8_t* fp4_weight, uint8_t* fp4_weight_interleaved, int const rows, int const cols, cudaStream_t stream) { - dim3 block(16, 32); + dim3 block(32, 32); interleave_fp4_weights_for_sm90_mixed_gemm_kernel<<<1024, block, 0, stream>>>( fp4_weight, fp4_weight_interleaved, rows, cols); } -void interleave_fp4_fp8_weights_for_sm90_mixed_gemm(uint8_t* fp4_weight, - uint8_t* fp4_weight_interleaved, int const rows, - int const cols, cudaStream_t stream) { - dim3 block(16, 32); - interleave_fp4_fp8_weights_for_sm90_mixed_gemm_kernel<<<1024, block, 0, stream>>>( - fp4_weight, fp4_weight_interleaved, rows, cols); -} - void interleave_int4_weights_for_sm90_mixed_gemm(uint8_t* int4_weight, uint8_t* int4_weight_interleaved, int const rows, int const cols, cudaStream_t stream) { diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_mixed_utils.h b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_mixed_utils.h index 108155a3a5c..314f7bcd326 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_mixed_utils.h +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_mixed_utils.h @@ -27,9 +27,6 @@ namespace cutlass_kernels { void interleave_fp4_weights_for_sm90_mixed_gemm(uint8_t* weight, uint8_t* weight_interleaved, int rows, int cols, cudaStream_t stream = 0); -void interleave_fp4_fp8_weights_for_sm90_mixed_gemm(uint8_t* weight, uint8_t* weight_interleaved, - int rows, int cols, cudaStream_t stream = 0); - void interleave_int4_weights_for_sm90_mixed_gemm(uint8_t* weight, uint8_t* weight_interleaved, int rows, int cols, cudaStream_t stream = 0); diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h index beae58f6656..859cb1e24fa 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h @@ -68,9 +68,6 @@ #include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" namespace tensorrt_llm::kernels::cutlass_kernels_oss { -#if defined(ENABLE_FP4) -using tensorrt_llm::kernels::cutlass_kernels::Fp4Type; -#endif // ============================= Variable batched Gemm things =========================== template ::value || cutlass::platform::is_same::value || #if defined(ENABLE_FP4) - cutlass::platform::is_same::value || + cutlass::platform::is_same::value || #endif cutlass::platform::is_same::value); @@ -256,7 +253,7 @@ static void dispatch(GroupedGemmInput; + constexpr bool isFp4 = std::is_same_v; #else constexpr bool isFp4 = false; #endif @@ -535,17 +532,17 @@ void dispatchMoeGemmToCutlass( namespace tensorrt_llm::kernels::cutlass_kernels { template + bool IsMXFPX> std::vector -MoeGemmRunner::getConfigs( +MoeGemmRunner::getConfigs( bool supports_finalize_fusion) const { return getConfigs(sm_, supports_finalize_fusion); } template + bool IsMXFPX> std::vector -MoeGemmRunner::getConfigs( +MoeGemmRunner::getConfigs( int sm, bool supports_finalize_fusion) { std::vector candidate_configs = getTmaWarpSpecializedConfigs(sm, supports_finalize_fusion); @@ -555,9 +552,9 @@ MoeGemmRunner -std::vector MoeGemmRunner< - T, WeightType, OutputType, ScaleBiasType, IsMXFPX, Sm90Wfp4Afp8Mode>::getAmpereConfigs(int sm) { + bool IsMXFPX> +std::vector +MoeGemmRunner::getAmpereConfigs(int sm) { using tensorrt_llm::cutlass_extensions::CutlassGemmConfig; static constexpr auto weight_only_flag = std::is_same::value ? CutlassGemmConfig::NONE : CutlassGemmConfig::WEIGHT_ONLY; @@ -584,11 +581,10 @@ std::vector MoeGemmRunner< } template + bool IsMXFPX> std::vector -MoeGemmRunner::getTmaWarpSpecializedConfigs(int sm, - bool supports_finalize_fusion) { +MoeGemmRunner::getTmaWarpSpecializedConfigs( + int sm, bool supports_finalize_fusion) { using tensorrt_llm::cutlass_extensions::CutlassGemmConfig; static constexpr auto weight_only_flag = std::is_same::value ? CutlassGemmConfig::NONE : CutlassGemmConfig::WEIGHT_ONLY; @@ -600,9 +596,8 @@ MoeGemmRunner( @@ -673,45 +668,28 @@ MoeGemmRunner -bool MoeGemmRunner:: - isTmaWarpSpecialized(cutlass_extensions::CutlassGemmConfig gemm_config) const { + bool IsMXFPX> +bool MoeGemmRunner::isTmaWarpSpecialized( + cutlass_extensions::CutlassGemmConfig gemm_config) const { bool config_is_tma_warp_specialized = gemm_config.is_tma_warp_specialized; return supportsTmaWarpSpecialized() && config_is_tma_warp_specialized; } template -bool MoeGemmRunner::supportsTmaWarpSpecialized(int sm) { + bool IsMXFPX> +bool MoeGemmRunner::supportsTmaWarpSpecialized( + int sm) { return (sm == 90 && tensorrt_llm::kernels::cutlass_kernels::isValidHopperMOESpecialisation()) || @@ -723,18 +701,16 @@ bool MoeGemmRunner -int MoeGemmRunner::getSM() - const { + bool IsMXFPX> +int MoeGemmRunner::getSM() const { return this->sm_; } // currently support sm80 bf16/fp16 gate activation, only set predication tensor for m direction template -bool MoeGemmRunner::supportsFusedGatedActivation(ActivationType activation_type, - int gemm_n, int gemm_k) const { + bool IsMXFPX> +bool MoeGemmRunner::supportsFusedGatedActivation( + ActivationType activation_type, int gemm_n, int gemm_k) const { constexpr bool ENABLE_FUSED_GATED_ACTIVATION = true; return (activation_type == ActivationType::Swiglu || activation_type == ActivationType::Geglu) && std::is_same_v && !std::is_same_v && !use_fp8 && @@ -743,18 +719,17 @@ bool MoeGemmRunner -bool MoeGemmRunner:: - isFusedGatedActivation(cutlass_extensions::CutlassGemmConfig gemm_config, - ActivationType activation_type, int gemm_n, int gemm_k) const { + bool IsMXFPX> +bool MoeGemmRunner::isFusedGatedActivation( + cutlass_extensions::CutlassGemmConfig gemm_config, ActivationType activation_type, int gemm_n, + int gemm_k) const { return supportsFusedGatedActivation(activation_type, gemm_n, gemm_k) && !gemm_config.is_tma_warp_specialized; } template -MoeGemmRunner::MoeGemmRunner() { + bool IsMXFPX> +MoeGemmRunner::MoeGemmRunner() { int device{-1}; tensorrt_llm::common::check_cuda_error(cudaGetDevice(&device)); sm_ = tensorrt_llm::common::getSMVersion(); @@ -763,11 +738,11 @@ MoeGemmRunner + bool IsMXFPX> template -void MoeGemmRunner:: - dispatchToArch(GroupedGemmInput inputs, - TmaWarpSpecializedGroupedGemmInput hopper_inputs) { +void MoeGemmRunner::dispatchToArch( + GroupedGemmInput inputs, + TmaWarpSpecializedGroupedGemmInput hopper_inputs) { static_assert( std::is_same_v, "Separate Scale/Bias type is not supported. This is assumed to be the gemm output type"); @@ -783,7 +758,7 @@ void MoeGemmRunner= 75 && sm_ < 80) { #if defined(ENABLE_FP4) - constexpr bool is_fp4 = std::is_same_v; + constexpr bool is_fp4 = std::is_same_v; #else constexpr bool is_fp4 = false; #endif @@ -796,7 +771,7 @@ void MoeGemmRunner= 80 && sm_ < 90) { #if defined(ENABLE_FP4) - constexpr bool is_fp4 = std::is_same_v; + constexpr bool is_fp4 = std::is_same_v; #else constexpr bool is_fp4 = false; #endif @@ -834,7 +809,7 @@ void MoeGemmRunner() && - !use_sm90_mixed_input_gemm) { + !use_w4_groupwise) { // We allow both tma warp specialized and SM80 configurations to coexist because for some // cases with small numbers of tokens SM80 is faster. We check here to see which is selected if (inputs.gemm_config.sm_version >= 90) { @@ -918,13 +893,12 @@ void MoeGemmRunner( + T, WeightType, ScaleBiasType, cutlass_extensions::EpilogueOpDefault, 1>( inputs, hopper_inputs, multi_processor_count_, nullptr); return; } @@ -962,9 +936,9 @@ void MoeGemmRunner -size_t MoeGemmRunner::getMaxWorkspaceSize(int num_experts) const { + bool IsMXFPX> +size_t MoeGemmRunner::getMaxWorkspaceSize( + int num_experts) const { if (num_experts != num_experts_) { TLLM_LOG_TRACE("Calling getMaxWorkspaceSize() with a new expert count %d vs %d", num_experts, num_experts_); @@ -975,12 +949,13 @@ size_t MoeGemmRunner -size_t MoeGemmRunner::calcMaxWorkspaceSize(int num_experts) const { - if constexpr (use_sm90_mixed_input_gemm) { - return cutlass_kernels_oss::calcMaxWorkspaceSizeTmaWarpSpecializedMixedInput< - T, WeightType, OutputType, Sm90Wfp4Afp8Mode>(num_experts, multi_processor_count_); + bool IsMXFPX> +size_t MoeGemmRunner::calcMaxWorkspaceSize( + int num_experts) const { + if constexpr (use_w4_groupwise) { + return cutlass_kernels_oss::calcMaxWorkspaceSizeTmaWarpSpecializedMixedInput( + num_experts, multi_processor_count_); } if (!supportsTmaWarpSpecialized()) { return 0; @@ -1034,19 +1009,19 @@ size_t MoeGemmRunner + bool IsMXFPX> template -void MoeGemmRunner::runGemm( +void MoeGemmRunner::runGemm( GroupedGemmInput inputs, TmaWarpSpecializedGroupedGemmInput hopper_inputs) { dispatchToArch(inputs, hopper_inputs); } template -void MoeGemmRunner:: - moeGemmBiasAct(GroupedGemmInput inputs, - TmaWarpSpecializedGroupedGemmInput hopper_inputs) { + bool IsMXFPX> +void MoeGemmRunner::moeGemmBiasAct( + GroupedGemmInput inputs, + TmaWarpSpecializedGroupedGemmInput hopper_inputs) { switch (inputs.activation_type) { case ActivationType::Relu: runGemm(inputs, hopper_inputs); @@ -1079,17 +1054,17 @@ void MoeGemmRunner -void MoeGemmRunner::moeGemm( + bool IsMXFPX> +void MoeGemmRunner::moeGemm( GroupedGemmInput inputs, TmaWarpSpecializedGroupedGemmInput hopper_inputs) { runGemm(inputs, hopper_inputs); } template -int MoeGemmRunner:: - queryOccupancyForConfig(cutlass_extensions::CutlassGemmConfig const& config) { + bool IsMXFPX> +int MoeGemmRunner::queryOccupancyForConfig( + cutlass_extensions::CutlassGemmConfig const& config) { // TMA warp-specialized configs (Hopper/Blackwell native) do not use the Ampere GroupedGEMM // occupancy path, so we conservatively report them as supported (occupancy > 0). if (config.is_tma_warp_specialized) { diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws.h b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws.h index 54d2a2d8979..d2c5da271ca 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws.h +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws.h @@ -61,9 +61,6 @@ #include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" namespace tensorrt_llm::kernels::cutlass_kernels_oss { -#if defined(ENABLE_FP4) -using tensorrt_llm::kernels::cutlass_kernels::Fp4Type; -#endif using tensorrt_llm::kernels::cutlass_kernels::TmaWarpSpecializedGroupedGemmInput; using EpilogueFusion = TmaWarpSpecializedGroupedGemmInput::EpilogueFusion; @@ -75,7 +72,7 @@ auto getDispatchFunctionForSM100(cutlass_extensions::EpilogueScheduleType epilog auto select_dynamic_cga = [epilogue_schedule](auto dynamic_cga_t) { #if defined(ENABLE_FP4) constexpr bool is_block_scaled = - IsMXFPX || std::is_same_v || std::is_same_v; + IsMXFPX || std::is_same_v || std::is_same_v; #else constexpr bool is_block_scaled = IsMXFPX; #endif @@ -165,7 +162,7 @@ void dispatchMoeGemmFinalDispatchTmaWarpSpecialized( else { #if defined(ENABLE_FP4) constexpr static bool is_wfp4afp8 = - std::is_same_v && std::is_same_v; + std::is_same_v && std::is_same_v; #else constexpr static bool is_wfp4afp8 = false; #endif @@ -246,7 +243,7 @@ constexpr bool are_tile_shapes_supported_sm100() { if constexpr (Arch::kMinComputeCapability == 103) { #if defined(ENABLE_FP4) - return std::is_same_v && std::is_same_v && + return std::is_same_v && std::is_same_v && TileM == 128 && (TileN == 128 || TileN == 256); #else return false; @@ -258,7 +255,8 @@ constexpr bool are_tile_shapes_supported_sm100() { } #ifdef ENABLE_FP4 - if constexpr (std::is_same_v || std::is_same_v) { + if constexpr (std::is_same_v || + std::is_same_v) { // if (TileN % 64 != 0 || TileN < 128) // { // return false; @@ -457,8 +455,8 @@ void dispatchMoeGemmSelectTileShapeTmaWarpSpecialized( } #if defined(ENABLE_FP4) && defined(COMPILE_BLACKWELL_SM103_TMA_GROUPED_GEMMS) // Check this before SM100 because we fall back to SM100 if not NVFP4 - else if (gemm_config.sm_version == 103 && std::is_same_v && - std::is_same_v) { + else if (gemm_config.sm_version == 103 && std::is_same_v && + std::is_same_v) { if constexpr (kernels::cutlass_kernels::isValidBlackwellMOESpecialisation< T, WeightType, EpilogueTag, FUSION>()) { switch (gemm_config.tile_config_sm100) { diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws_mixed_dtype.h b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws_mixed_dtype.h index d0655852a68..eaaedf42580 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws_mixed_dtype.h +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws_mixed_dtype.h @@ -59,11 +59,7 @@ namespace tensorrt_llm::kernels::cutlass_kernels_oss { -#if defined(ENABLE_FP4) -using tensorrt_llm::kernels::cutlass_kernels::Fp4Type; -#endif using tensorrt_llm::kernels::cutlass_kernels::GroupedGemmInput; -using tensorrt_llm::kernels::cutlass_kernels::Sm90Wfp4Afp8ScaleMode; using tensorrt_llm::kernels::cutlass_kernels::TmaWarpSpecializedGroupedGemmInput; namespace tk = tensorrt_llm::common; namespace tkc = tensorrt_llm::cutlass_extensions; @@ -71,37 +67,32 @@ namespace tkc = tensorrt_llm::cutlass_extensions; using namespace cute; template + typename CTAShape, typename ClusterShape> void sm90_dispatch_mainloop_schedules( GroupedGemmInput inputs, TmaWarpSpecializedGroupedGemmInput hopper_inputs, int sm_count_, size_t* workspace_size) { TLLM_LOG_DEBUG(__PRETTY_FUNCTION__); #ifdef COMPILE_HOPPER_TMA_GROUPED_GEMMS - if constexpr (KernelType != tkc::MainloopScheduleType::AUTO) { - TLLM_CHECK_WITH_INFO(inputs.gemm_config.mainloop_schedule == KernelType, - "Single-warpgroup GEMM schedule does not match its compiled policy."); - sm90_generic_mixed_moe_small_k_kernelLauncher< - T, WeightType, GemmOutputType, EpilogueTag, CTAShape, ClusterShape, - cutlass::gemm::KernelTmaWarpSpecializedPingpong, - cutlass::epilogue::TmaWarpSpecializedCooperative, - cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY, ScaleMode, KernelType>( - inputs, hopper_inputs, sm_count_, workspace_size); - return; - } - switch (inputs.gemm_config.mainloop_schedule) { case tkc::MainloopScheduleType::COOPERATIVE: if constexpr (get<0>(CTAShape{}) < 128) { TLLM_THROW("COOPERATIVE is only enabled when tile M >= 128."); } else { - sm90_generic_mixed_moe_gemm_kernelLauncher< - T, WeightType, GemmOutputType, EpilogueTag, CTAShape, ClusterShape, - cutlass::gemm::KernelTmaWarpSpecializedCooperative, - cutlass::epilogue::TmaWarpSpecializedCooperative, - cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY, ScaleMode>( - inputs, hopper_inputs, sm_count_, workspace_size); + if constexpr ((get<0>(CTAShape{}) == 128) && get<1>(CTAShape{}) == 128) { + sm90_generic_mixed_moe_gemm_kernelLauncher< + T, WeightType, GemmOutputType, EpilogueTag, CTAShape, ClusterShape, + cutlass::gemm::KernelTmaWarpSpecializedPingpong, + cutlass::epilogue::TmaWarpSpecializedCooperative, + cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY>(inputs, hopper_inputs, sm_count_, + workspace_size); + } else { + sm90_generic_mixed_moe_gemm_kernelLauncher< + T, WeightType, GemmOutputType, EpilogueTag, CTAShape, ClusterShape, + cutlass::gemm::KernelTmaWarpSpecializedCooperative, + cutlass::epilogue::TmaWarpSpecializedCooperative, + cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY>(inputs, hopper_inputs, sm_count_, + workspace_size); + } } break; @@ -110,8 +101,8 @@ void sm90_dispatch_mainloop_schedules( T, WeightType, GemmOutputType, EpilogueTag, CTAShape, ClusterShape, cutlass::gemm::KernelTmaWarpSpecializedPingpong, cutlass::epilogue::TmaWarpSpecializedCooperative, - cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY, ScaleMode>(inputs, hopper_inputs, - sm_count_, workspace_size); + cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY>(inputs, hopper_inputs, sm_count_, + workspace_size); break; default: TLLM_THROW( @@ -128,41 +119,31 @@ void sm90_dispatch_mainloop_schedules( } template + typename CTAShape> void sm90_dispatch_moe_mixed_dtype_gemm_config( GroupedGemmInput inputs, TmaWarpSpecializedGroupedGemmInput hopper_inputs, int sm_count_, size_t* workspace_size) { TLLM_LOG_DEBUG(__PRETTY_FUNCTION__); - if constexpr (KernelType != tkc::MainloopScheduleType::AUTO) { - TLLM_CHECK_WITH_INFO(inputs.gemm_config.cluster_shape == tkc::ClusterShape::ClusterShape_1x1x1, - "Single-warpgroup GEMM requires a 1x1x1 cluster."); - sm90_dispatch_mainloop_schedules, ScaleMode, KernelType>( - inputs, hopper_inputs, sm_count_, workspace_size); - return; - } - switch (inputs.gemm_config.cluster_shape) { case tkc::ClusterShape::ClusterShape_1x1x1: sm90_dispatch_mainloop_schedules, ScaleMode>(inputs, hopper_inputs, - sm_count_, workspace_size); + Shape<_1, _1, _1>>(inputs, hopper_inputs, sm_count_, + workspace_size); break; case tkc::ClusterShape::ClusterShape_2x1x1: sm90_dispatch_mainloop_schedules, ScaleMode>(inputs, hopper_inputs, - sm_count_, workspace_size); + Shape<_2, _1, _1>>(inputs, hopper_inputs, sm_count_, + workspace_size); break; case tkc::ClusterShape::ClusterShape_1x2x1: sm90_dispatch_mainloop_schedules, ScaleMode>(inputs, hopper_inputs, - sm_count_, workspace_size); + Shape<_1, _2, _1>>(inputs, hopper_inputs, sm_count_, + workspace_size); break; case tkc::ClusterShape::ClusterShape_2x2x1: sm90_dispatch_mainloop_schedules, ScaleMode>(inputs, hopper_inputs, - sm_count_, workspace_size); + Shape<_2, _2, _1>>(inputs, hopper_inputs, sm_count_, + workspace_size); break; default: TLLM_THROW( @@ -172,62 +153,7 @@ void sm90_dispatch_moe_mixed_dtype_gemm_config( } template -void sm90_dispatch_moe_mixed_dtype_gemm_with_small_k( - GroupedGemmInput inputs, - TmaWarpSpecializedGroupedGemmInput hopper_inputs, int sm_count_, size_t* workspace_size) { - switch (inputs.gemm_config.mainloop_schedule) { - case tkc::MainloopScheduleType::AUTO: - case tkc::MainloopScheduleType::PINGPONG: - case tkc::MainloopScheduleType::COOPERATIVE: - case tkc::MainloopScheduleType::WARPSPECIALIZED: - sm90_dispatch_moe_mixed_dtype_gemm_config(inputs, hopper_inputs, - sm_count_, workspace_size); - break; - case tkc::MainloopScheduleType::SINGLE_WARPGROUP_PREFILL: - sm90_dispatch_moe_mixed_dtype_gemm_config< - T, WeightType, GemmOutputType, EpilogueTag, CTAShape, ScaleMode, - tkc::MainloopScheduleType::SINGLE_WARPGROUP_PREFILL>(inputs, hopper_inputs, sm_count_, - workspace_size); - break; - case tkc::MainloopScheduleType::SINGLE_WARPGROUP_ROLLING: - sm90_dispatch_moe_mixed_dtype_gemm_config< - T, WeightType, GemmOutputType, EpilogueTag, CTAShape, ScaleMode, - tkc::MainloopScheduleType::SINGLE_WARPGROUP_ROLLING>(inputs, hopper_inputs, sm_count_, - workspace_size); - break; - default: - TLLM_THROW("Unsupported mixed-input mainloop schedule."); - } -} - -template -void sm90_dispatch_moe_mixed_dtype_gemm_small_k_only( - GroupedGemmInput inputs, - TmaWarpSpecializedGroupedGemmInput hopper_inputs, int sm_count_, size_t* workspace_size) { - switch (inputs.gemm_config.mainloop_schedule) { - case tkc::MainloopScheduleType::SINGLE_WARPGROUP_PREFILL: - sm90_dispatch_moe_mixed_dtype_gemm_config< - T, WeightType, GemmOutputType, EpilogueTag, CTAShape, ScaleMode, - tkc::MainloopScheduleType::SINGLE_WARPGROUP_PREFILL>(inputs, hopper_inputs, sm_count_, - workspace_size); - break; - case tkc::MainloopScheduleType::SINGLE_WARPGROUP_ROLLING: - sm90_dispatch_moe_mixed_dtype_gemm_config< - T, WeightType, GemmOutputType, EpilogueTag, CTAShape, ScaleMode, - tkc::MainloopScheduleType::SINGLE_WARPGROUP_ROLLING>(inputs, hopper_inputs, sm_count_, - workspace_size); - break; - default: - TLLM_THROW("This tile shape is only available for a single-warpgroup small-K kernel."); - } -} - -template + int PackedScalesNum> void sm90_dispatch_moe_mixed_dtype_gemm_to_cutlass( GroupedGemmInput inputs, TmaWarpSpecializedGroupedGemmInput hopper_inputs, int sm_count_, size_t* workspace_size) { @@ -236,87 +162,72 @@ void sm90_dispatch_moe_mixed_dtype_gemm_to_cutlass( // perform the best for mixed type gemms. #if defined(ENABLE_FP4) - static constexpr size_t ExpectedActivationBytes = - Sm90Wfp4Afp8Mode != Sm90Wfp4Afp8ScaleMode::kDisabled - ? 1 - : (std::is_same_v ? 2 : 1); - TLLM_CHECK(sizeof(T) == ExpectedActivationBytes); + constexpr int Ntile = (std::is_same_v) ? 64 : 128; + constexpr int Ktile = + (std::is_same_v) ? 128 : 128 * PackedScalesNum / sizeof(T); + TLLM_CHECK(sizeof(T) == (std::is_same_v) ? 2 : 1); #else + constexpr int Ntile = 128; + constexpr int Ktile = 128 * PackedScalesNum / sizeof(T); TLLM_CHECK(sizeof(T) == 1); #endif - static_cast(UnusedScalePackFactor); - static constexpr auto ScaleMode = - Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0 - ? cutlass::gemm::collective::MixedInputScaleMode::kPreMmaE8M0 - : cutlass::gemm::collective::MixedInputScaleMode::kPostMma; -#define DISPATCH_MIXED_DTYPE_MOE_TILE(ENUM_NAME, TILE_M, TILE_N, TILE_K) \ - case tkc::CutlassTileConfigSM90::ENUM_NAME: \ - sm90_dispatch_moe_mixed_dtype_gemm_config, Int, Int>, \ - ScaleMode>(inputs, hopper_inputs, sm_count_, \ - workspace_size); \ - break - -#define DISPATCH_MIXED_DTYPE_MOE_TILE_WITH_SMALL_K(ENUM_NAME, TILE_M, TILE_N, TILE_K) \ - case tkc::CutlassTileConfigSM90::ENUM_NAME: \ - if constexpr (Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0) { \ - sm90_dispatch_moe_mixed_dtype_gemm_with_small_k< \ - T, WeightType, GemmOutputType, EpilogueTag, \ - Shape, Int, Int>, ScaleMode>(inputs, hopper_inputs, \ - sm_count_, workspace_size); \ - } else { \ - sm90_dispatch_moe_mixed_dtype_gemm_config, Int, Int>, \ - ScaleMode>(inputs, hopper_inputs, sm_count_, \ - workspace_size); \ - } \ - break - -#define DISPATCH_MIXED_DTYPE_MOE_SMALL_K_TILE(ENUM_NAME, TILE_M, TILE_N, TILE_K) \ - case tkc::CutlassTileConfigSM90::ENUM_NAME: \ - if constexpr (Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0) { \ - sm90_dispatch_moe_mixed_dtype_gemm_small_k_only< \ - T, WeightType, GemmOutputType, EpilogueTag, \ - Shape, Int, Int>, ScaleMode>(inputs, hopper_inputs, \ - sm_count_, workspace_size); \ - } else { \ - TLLM_THROW("Single-warpgroup small-K tile is only valid for Humming pre-MMA scale."); \ - } \ - break + using _Ntile = Int; + using _Ktile = Int; switch (inputs.gemm_config.tile_config_sm90) { - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape64x16x128B, 64, 16, 128); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape64x16x256B, 64, 16, 256); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape64x16x512B, 64, 16, 512); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape64x32x128B, 64, 32, 128); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape64x32x256B, 64, 32, 256); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape64x32x512B, 64, 32, 512); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape64x64x128B, 64, 64, 128); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape64x64x256B, 64, 64, 256); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape64x64x512B, 64, 64, 512); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape64x128x128B, 64, 128, 128); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape64x128x256B, 64, 128, 256); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape64x128x512B, 64, 128, 512); - DISPATCH_MIXED_DTYPE_MOE_SMALL_K_TILE(CtaShape128x8x128B, 128, 8, 128); - DISPATCH_MIXED_DTYPE_MOE_TILE_WITH_SMALL_K(CtaShape128x16x128B, 128, 16, 128); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape128x16x256B, 128, 16, 256); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape128x16x512B, 128, 16, 512); - DISPATCH_MIXED_DTYPE_MOE_TILE_WITH_SMALL_K(CtaShape128x32x128B, 128, 32, 128); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape128x32x256B, 128, 32, 256); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape128x32x512B, 128, 32, 512); - DISPATCH_MIXED_DTYPE_MOE_SMALL_K_TILE(CtaShape128x40x128B, 128, 40, 128); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape128x64x128B, 128, 64, 128); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape128x64x256B, 128, 64, 256); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape128x64x512B, 128, 64, 512); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape128x128x128B, 128, 128, 128); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape128x128x256B, 128, 128, 256); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape128x128x512B, 128, 128, 512); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape128x256x128B, 128, 256, 128); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape128x256x256B, 128, 256, 256); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape256x128x128B, 256, 128, 128); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape256x128x256B, 256, 128, 256); - DISPATCH_MIXED_DTYPE_MOE_TILE(CtaShape256x256x128B, 256, 256, 128); + case tkc::CutlassTileConfigSM90::CtaShape64x16x128B: + sm90_dispatch_moe_mixed_dtype_gemm_config>(inputs, hopper_inputs, + sm_count_, workspace_size); + break; + case tkc::CutlassTileConfigSM90::CtaShape64x32x128B: + sm90_dispatch_moe_mixed_dtype_gemm_config>(inputs, hopper_inputs, + sm_count_, workspace_size); + break; + case tkc::CutlassTileConfigSM90::CtaShape64x64x128B: + sm90_dispatch_moe_mixed_dtype_gemm_config>(inputs, hopper_inputs, + sm_count_, workspace_size); + break; + case tkc::CutlassTileConfigSM90::CtaShape64x128x128B: + sm90_dispatch_moe_mixed_dtype_gemm_config>( + inputs, hopper_inputs, sm_count_, workspace_size); + break; + // case tkc::CutlassTileConfigSM90::CtaShape64x256x128B: + // sm90_dispatch_moe_mixed_dtype_gemm_config>(inputs, hopper_inputs, sm_count_, workspace_size); break; + case tkc::CutlassTileConfigSM90::CtaShape128x16x128B: + sm90_dispatch_moe_mixed_dtype_gemm_config>( + inputs, hopper_inputs, sm_count_, workspace_size); + break; + case tkc::CutlassTileConfigSM90::CtaShape128x32x128B: + sm90_dispatch_moe_mixed_dtype_gemm_config>( + inputs, hopper_inputs, sm_count_, workspace_size); + break; + case tkc::CutlassTileConfigSM90::CtaShape128x64x128B: + sm90_dispatch_moe_mixed_dtype_gemm_config>( + inputs, hopper_inputs, sm_count_, workspace_size); + break; + case tkc::CutlassTileConfigSM90::CtaShape128x128x128B: + sm90_dispatch_moe_mixed_dtype_gemm_config>( + inputs, hopper_inputs, sm_count_, workspace_size); + break; + // case tkc::CutlassTileConfigSM90::CtaShape128x256x128B: + // sm90_dispatch_moe_mixed_dtype_gemm_config>(inputs, hopper_inputs, sm_count_, workspace_size); break; + // case tkc::CutlassTileConfigSM90::CtaShape256x128x128B: + // sm90_dispatch_moe_mixed_dtype_gemm_config>(inputs, hopper_inputs, sm_count_, workspace_size); break; + // case tkc::CutlassTileConfigSM90::CtaShape256x256x128B: + // sm90_dispatch_moe_mixed_dtype_gemm_config>(inputs, hopper_inputs, sm_count_, workspace_size); break; case tkc::CutlassTileConfigSM90::Undefined: TLLM_THROW( "[Mixed dtype MoE GEMM][sm90_dispatch_moe_mixed_dtype_gemm_to_cutlass] gemm config " @@ -336,36 +247,27 @@ void sm90_dispatch_moe_mixed_dtype_gemm_to_cutlass( "GEMM."); break; } - -#undef DISPATCH_MIXED_DTYPE_MOE_TILE -#undef DISPATCH_MIXED_DTYPE_MOE_TILE_WITH_SMALL_K -#undef DISPATCH_MIXED_DTYPE_MOE_SMALL_K_TILE } -template +template size_t calcMaxWorkspaceSizeTmaWarpSpecializedMixedInput(int num_experts, int sm_count_) { size_t count = 0; #if defined(ENABLE_FP4) - constexpr int Ktile = (std::is_same_v) ? 256 : 512; + constexpr int Ktile = (std::is_same_v) ? 256 : 512; #else constexpr int Ktile = 512; #endif using _Ktile = Int; - static constexpr auto ScaleMode = - Sm90Wfp4Afp8Mode == Sm90Wfp4Afp8ScaleMode::kHummingPreMmaE8M0 - ? cutlass::gemm::collective::MixedInputScaleMode::kPreMmaE8M0 - : cutlass::gemm::collective::MixedInputScaleMode::kPostMma; #ifdef COMPILE_HOPPER_TMA_GROUPED_GEMMS GroupedGemmInput inputs{}; inputs.num_experts = num_experts; - sm90_generic_mixed_moe_gemm_kernelLauncher< - T, WeightType, OutputType, tensorrt_llm::cutlass_extensions::EpilogueOpDefault, - Shape<_128, _64, _Ktile>, Shape<_1, _1, _1>, - cutlass::gemm::KernelTmaWarpSpecializedCooperative, - cutlass::epilogue::TmaWarpSpecializedCooperative, - cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY, ScaleMode>( + sm90_generic_mixed_moe_gemm_kernelLauncher, Shape<_1, _1, _1>, + cutlass::gemm::KernelTmaWarpSpecializedCooperative, + cutlass::epilogue::TmaWarpSpecializedCooperative, + cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY>( inputs, TmaWarpSpecializedGroupedGemmInput{}, sm_count_, &count); #endif return count; diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu index 83f07bc2924..52cd03887b5 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu @@ -78,10 +78,10 @@ size_t TmaWarpSpecializedGroupedGemmInput::workspaceSize(int num_experts, return tensorrt_llm::common::calculateTotalWorkspaceSize(buffers.data(), buffers.size()); } -void TmaWarpSpecializedGroupedGemmInput::configureWorkspace( - int8_t* start_ptr, int num_experts, void* gemm_workspace, size_t gemm_workspace_size, - void* precomputed_scheduler_workspace, size_t precomputed_scheduler_workspace_size, - FpXBlockScalingType scaling_type) { +void TmaWarpSpecializedGroupedGemmInput::configureWorkspace(int8_t* start_ptr, int num_experts, + void* gemm_workspace, + size_t gemm_workspace_size, + FpXBlockScalingType scaling_type) { auto buffers = workspaceBuffers(num_experts, scaling_type); std::array pointers{}; TLLM_CHECK_WITH_INFO(pointers.size() == buffers.size(), @@ -125,9 +125,6 @@ void TmaWarpSpecializedGroupedGemmInput::configureWorkspace( this->gemm_workspace = reinterpret_cast(gemm_workspace); this->gemm_workspace_size = gemm_workspace_size; - this->precomputed_scheduler_workspace = - reinterpret_cast(precomputed_scheduler_workspace); - this->precomputed_scheduler_workspace_size = precomputed_scheduler_workspace_size; } void TmaWarpSpecializedGroupedGemmInput::setFinalizeFusionParams(void* final_output, @@ -178,8 +175,6 @@ std::string TmaWarpSpecializedGroupedGemmInput::toString() const { << ", with Stride: " << (PrintType)fpX_block_scaling_factors_stride_weight << "\n"; ss << "Gemm Workspace: " << (PrintType)gemm_workspace << ", with Size: " << gemm_workspace_size << "\n"; - ss << "Precomputed Scheduler Workspace: " << (PrintType)precomputed_scheduler_workspace - << ", with Size: " << precomputed_scheduler_workspace_size << "\n"; } return ss.str(); diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_tma_warp_specialized_traits.h b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_tma_warp_specialized_traits.h index f3b9c457338..fb9ae80f2f7 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_tma_warp_specialized_traits.h +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_tma_warp_specialized_traits.h @@ -21,7 +21,7 @@ #include "cutlass_extensions/epilogue_helpers.h" #ifdef ENABLE_FP4 -#include "tensorrt_llm/kernels/cutlass_kernels/fp4_compat.h" +#include #endif namespace tensorrt_llm::kernels::cutlass_kernels { @@ -34,10 +34,10 @@ template ::value && + return ((cutlass::platform::is_same::value && cutlass::platform::is_same::value) || (cutlass::platform::is_same::value && - cutlass::platform::is_same::value)) && + cutlass::platform::is_same::value)) && cutlass::platform::is_same::value; #else return false; @@ -57,7 +57,7 @@ constexpr bool isValidBlackwellMOESpecialisation() { return (cutlass::platform::is_same::value || #if defined(ENABLE_FP4) (cutlass::platform::is_same::value && - cutlass::platform::is_same::value) + cutlass::platform::is_same::value) #else false #endif @@ -83,11 +83,12 @@ constexpr bool isValidHopperMOESpecialisation() { (cutlass::platform::is_same::value && cutlass::platform::is_same::value) #ifdef ENABLE_FP4 - || cutlass::platform::is_same::value + || (cutlass::platform::is_same<__nv_fp4_e2m1, WeightType>::value && + !cutlass::platform::is_same::value) #endif - ) + ) #ifdef ENABLE_FP4 - && !cutlass::platform::is_same::value + && !cutlass::platform::is_same::value #endif && cutlass::platform::is_same::value; #else @@ -114,7 +115,7 @@ template constexpr bool isValidAmpereMOESpecialisation() { #ifdef ENABLE_FP4 - return !std::is_same_v && !std::is_same_v; + return !std::is_same_v && !std::is_same_v; #else return true; // Default to true #endif diff --git a/docs/api/fused_moe.rst b/docs/api/fused_moe.rst index ef1654bfc7f..044dbab2ef1 100644 --- a/docs/api/fused_moe.rst +++ b/docs/api/fused_moe.rst @@ -38,14 +38,9 @@ Utility Functions reorder_rows_for_gated_act_gemm interleave_moe_weights_for_sm90_mixed_gemm interleave_moe_scales_for_sm90_mixed_gemm - preprocess_moe_weights_for_sm90_mixed_gemm_humming fused_topk_deepseek hash_topk -The E8M0 range-clamping, residual-scale factorization, and FP4 payload-rewrite -scheme used by ``preprocess_moe_weights_for_sm90_mixed_gemm_humming`` is adapted -from `Humming `_. - Multi-LoRA MoE (BGMV) --------------------- diff --git a/flashinfer/fused_moe/__init__.py b/flashinfer/fused_moe/__init__.py index 6e08cd04e69..3453afbabec 100644 --- a/flashinfer/fused_moe/__init__.py +++ b/flashinfer/fused_moe/__init__.py @@ -53,6 +53,8 @@ convert_to_block_layout, cutlass_fused_moe, cutlass_fused_moe_workspace_size, + interleave_moe_scales_for_sm90_mixed_gemm, + interleave_moe_weights_for_sm90_mixed_gemm, gen_cutlass_fused_moe_sm120_module, gen_cutlass_fused_moe_sm103_module, gen_cutlass_fused_moe_sm100_module, @@ -71,12 +73,6 @@ trtllm_mxint4_block_scale_routed_moe, ) -from .prepare import ( - interleave_moe_scales_for_sm90_mixed_gemm, - interleave_moe_weights_for_sm90_mixed_gemm, - preprocess_moe_weights_for_sm90_mixed_gemm_humming, -) - from ..tllm_enums import ( ActivationType, Fp8QuantizationType, @@ -164,7 +160,6 @@ "cutlass_fused_moe_workspace_size", "interleave_moe_scales_for_sm90_mixed_gemm", "interleave_moe_weights_for_sm90_mixed_gemm", - "preprocess_moe_weights_for_sm90_mixed_gemm_humming", "gen_cutlass_fused_moe_sm120_module", "gen_cutlass_fused_moe_sm103_module", "gen_cutlass_fused_moe_sm100_module", diff --git a/flashinfer/fused_moe/core.py b/flashinfer/fused_moe/core.py index f9f24e2ca43..db8cfe6be3a 100644 --- a/flashinfer/fused_moe/core.py +++ b/flashinfer/fused_moe/core.py @@ -81,12 +81,6 @@ register_custom_op, register_fake_op, ) - -# These helpers moved to prepare.py; keep aliases here for backward compatibility. -from .prepare import ( - interleave_moe_scales_for_sm90_mixed_gemm as interleave_moe_scales_for_sm90_mixed_gemm, - interleave_moe_weights_for_sm90_mixed_gemm as interleave_moe_weights_for_sm90_mixed_gemm, -) from .utils import ( get_hybrid_num_tokens_buckets, make_hybrid_bucket_mapper, @@ -327,17 +321,7 @@ def get_cutlass_fused_moe_module(backend: str = "100", use_fast_build: bool = Fa class MoERunner(TunableRunner): # avoid overhead of creating a new runner in forward pass runner_dict: Dict[ - Tuple[ - torch.dtype, - torch.dtype, - torch.dtype, - bool, - bool, - bool, - bool, - bool, - bool, - ], + Tuple[torch.dtype, torch.dtype, torch.dtype, bool, bool, bool, bool, bool], Any, ] = dict() tuning_config = TuningConfig( @@ -372,7 +356,6 @@ def __init__( activation_type: ActivationType, use_packed_weights: bool, use_fused_finalize: bool, - use_wfp4afp8_humming: bool, ): self.x_dtype = x_dtype self.weight_dtype = weight_dtype @@ -388,7 +371,6 @@ def __init__( self.use_deepseek_fp8_block_scale = use_deepseek_fp8_block_scale self.use_w4_group_scaling = use_w4_group_scaling self.use_mxfp8_act_scaling = use_mxfp8_act_scaling - self.use_wfp4afp8_humming = use_wfp4afp8_humming self.min_latency_mode = min_latency_mode self.enable_pdl = enable_pdl self.use_packed_weights = use_packed_weights @@ -402,7 +384,6 @@ def __init__( use_mxfp8_act_scaling, use_packed_weights, use_fused_finalize, - use_wfp4afp8_humming, ) self.activation_type = activation_type # Set by tuning flow to indicate which GEMM stage (1 or 2) to filter tactics for @@ -418,7 +399,6 @@ def __init__( use_mxfp8_act_scaling, use_packed_weights, use_fused_finalize, - use_wfp4afp8_humming, ) self.fused_moe_runner = MoERunner.runner_dict[instance_key] @@ -475,52 +455,7 @@ def get_valid_tactics( # a sentinel so the autotuner contract is never violated with an empty list. if not all_tactics: return [-1] - valid_tactics = valid_tactics if valid_tactics else all_tactics - - if not self.use_w4_group_scaling: - return valid_tactics - - if stage not in (1, 2): - return valid_tactics - - x, fc1_expert_weights, _, fc2_expert_weights, _ = inputs - if stage == 1: - gemm_n = int(fc1_expert_weights.shape[1]) - gemm_k = int(x.shape[1]) - else: - gemm_n = int(fc2_expert_weights.shape[1]) - if fc2_expert_weights.dtype == torch.uint8: - gemm_k = int(fc2_expert_weights.shape[2]) * 2 - elif fc2_expert_weights.dtype == torch.int64: - gemm_k = int(fc2_expert_weights.shape[2]) * 16 - else: - gemm_k = int(fc2_expert_weights.shape[2]) - - try: - get_valid_tactics_for_shape = ( - self.fused_moe_runner.get_valid_tactics_for_shape - ) - shape_valid_tactics = set( - int(t) - for t in get_valid_tactics_for_shape( - int(stage), int(gemm_n), int(gemm_k) - ) - ) - except AttributeError: - return valid_tactics - except Exception as e: - logger.warning( - "get_valid_tactics_for_shape failed for stage %s, N=%d, K=%d: %s; " - "including occupancy-valid tactics in autotuner", - stage, - gemm_n, - gemm_k, - e, - ) - return valid_tactics - - filtered_tactics = [t for t in valid_tactics if t in shape_valid_tactics] - return filtered_tactics if filtered_tactics else valid_tactics + return valid_tactics if valid_tactics else all_tactics def forward( self, @@ -608,12 +543,13 @@ def cutlass_fused_moe( activation_type: ActivationType = ActivationType.Swiglu, use_packed_weights: bool = False, use_fused_finalize: bool = True, - use_wfp4afp8_humming: bool = False, profile_ids: Optional[List[int]] = None, workspace_buffer: Optional[torch.Tensor] = None, ) -> List[torch.Tensor]: if enable_pdl is None: enable_pdl = device_support_pdl(input.device) + tuner = AutoTuner.get() + MoERunner.refine_tuning_config(tune_max_num_tokens) # allocate workspace for profiling moe_runner = MoERunner( @@ -636,50 +572,39 @@ def cutlass_fused_moe( activation_type=activation_type, use_packed_weights=use_packed_weights, use_fused_finalize=use_fused_finalize, - use_wfp4afp8_humming=use_wfp4afp8_humming, ) - if profile_ids is None: - tuner = AutoTuner.get() - MoERunner.refine_tuning_config(tune_max_num_tokens) - - # Limit tactics to GEMM1 during tuning - moe_runner.gemm_idx_for_tuning = 1 - _, gemm_tactic_1 = tuner.choose_one( - "trtllm::fused_moe::gemm1", - [moe_runner], - MoERunner.tuning_config, - [ - input, - fc1_expert_weights, - fc1_expert_biases, - fc2_expert_weights, - fc2_expert_biases, - ], - gemm_idx=1, - ) + # Limit tactics to GEMM1 during tuning + moe_runner.gemm_idx_for_tuning = 1 + _, gemm_tactic_1 = tuner.choose_one( + "trtllm::fused_moe::gemm1", + [moe_runner], + MoERunner.tuning_config, + [ + input, + fc1_expert_weights, + fc1_expert_biases, + fc2_expert_weights, + fc2_expert_biases, + ], + gemm_idx=1, + ) - # Limit tactics to GEMM2 during tuning - moe_runner.gemm_idx_for_tuning = 2 - _, gemm_tactic_2 = tuner.choose_one( - "trtllm::fused_moe::gemm2", - [moe_runner], - MoERunner.tuning_config, - [ - input, - fc1_expert_weights, - fc1_expert_biases, - fc2_expert_weights, - fc2_expert_biases, - ], - gemm_idx=2, - ) - else: - if len(profile_ids) != 2: - raise ValueError( - "profile_ids must contain [gemm1_profile, gemm2_profile]" - ) - gemm_tactic_1, gemm_tactic_2 = profile_ids + # Limit tactics to GEMM2 during tuning + moe_runner.gemm_idx_for_tuning = 2 + _, gemm_tactic_2 = tuner.choose_one( + "trtllm::fused_moe::gemm2", + [moe_runner], + MoERunner.tuning_config, + [ + input, + fc1_expert_weights, + fc1_expert_biases, + fc2_expert_weights, + fc2_expert_biases, + ], + gemm_idx=2, + ) run_moe = ( moe_runner.fused_moe_runner.run_moe_min_latency @@ -782,7 +707,6 @@ def _fake_cutlass_fused_moe( activation_type: ActivationType = ActivationType.Swiglu, use_packed_weights: bool = False, use_fused_finalize: bool = True, - use_wfp4afp8_humming: bool = False, profile_ids: Optional[List[int]] = None, workspace_buffer: Optional[torch.Tensor] = None, ) -> List[torch.Tensor]: @@ -824,7 +748,6 @@ def _cutlass_fused_moe_workspace_size( use_mxfp8_act_scaling: bool = False, use_fused_finalize: bool = True, use_packed_weights: bool = False, - use_wfp4afp8_humming: bool = False, ) -> int: enable_pdl = device_support_pdl(torch.device("cuda")) moe_runner = MoERunner( @@ -847,7 +770,6 @@ def _cutlass_fused_moe_workspace_size( activation_type=activation_type, use_packed_weights=use_packed_weights, use_fused_finalize=use_fused_finalize, - use_wfp4afp8_humming=use_wfp4afp8_humming, ) return int( moe_runner.fused_moe_runner.get_workspace_size( @@ -875,6 +797,116 @@ def _cutlass_fused_moe_workspace_size( ) +@flashinfer_api +def interleave_moe_scales_for_sm90_mixed_gemm( + scales: torch.Tensor, + group_size: int = 32, +) -> torch.Tensor: + """Interleave MXFP4 block scales for the SM90 mixed-input MoE GEMM. + + The kernel expects scales in layout + ``(num_experts, K // (group_size * 4), rows * 4)`` rather than the natural + ``(num_experts, rows, K // group_size)`` produced by the MXFP4 quantizer. + This helper performs the reshape + permute equivalent to TensorRT-LLM's + ``WFP4A16FusedMoEMethod.load_quant_scales`` (PR #12451), with the fixed + interleave factor of ``128 // group_size`` used for MXFP4. + + Parameters + ---------- + scales : torch.Tensor + ``[num_experts, rows, K // group_size]`` uint8 tensor of E8M0 block + scales. + group_size : int + MXFP4 quantization group size (default 32). + + Returns + ------- + torch.Tensor + Contiguous uint8 tensor with shape + ``[num_experts, K // (group_size * factor), rows * factor]`` + where ``factor = 128 // group_size``. + """ + if scales.dim() != 3: + raise ValueError( + f"scales must be 3D (num_experts, rows, K/group_size); got {tuple(scales.shape)}" + ) + if scales.dtype != torch.uint8: + raise ValueError(f"scales must be uint8 (E8M0); got {scales.dtype}") + + factor = 128 // group_size + if factor < 1 or 128 % group_size != 0: + raise ValueError( + f"group_size={group_size} must divide 128 (interleave factor = 128 // group_size)" + ) + e, rows, kgs = scales.shape + if kgs % factor != 0: + raise ValueError( + f"K/group_size={kgs} must be divisible by interleave factor {factor}" + ) + tmp = ( + scales.reshape(e, rows, kgs // factor, factor).permute(0, 2, 1, 3).contiguous() + ) + return tmp.reshape(e, kgs // factor, rows * factor) + + +@flashinfer_api +def interleave_moe_weights_for_sm90_mixed_gemm( + weight: torch.Tensor, + quant_type: str = "fp4", +) -> torch.Tensor: + """Interleave 4-bit packed MoE weights for the SM90 mixed-input GEMM. + + The SM90 mixed-dtype MoE GEMM (used by ``cutlass_fused_moe`` with + ``use_w4_group_scaling=True``) expects weights in a specific interleaved + layout; without preprocessing, the LUT-based FP4→BF16 conversion reads + bytes from the wrong positions and the output diverges from a dequantized + reference for any K > 128. TensorRT-LLM's W4A16 MoE runs the equivalent + preprocessing at weight-load time (see + ``interleave_4bit_weights_for_Hopper_mixed_gemm`` in TRT-LLM PR #12451). + + Parameters + ---------- + weight : torch.Tensor + ``[num_experts, n, k // 2]`` uint8 CUDA tensor (4-bit values packed + two-per-byte). + quant_type : str + ``"fp4"`` for MXFP4 (the W4A16 path) or ``"int4"`` for INT4 (the + W4A8 path). + + Returns + ------- + torch.Tensor + A new uint8 tensor with the same shape as ``weight`` holding the + interleaved layout. Feed this directly as ``fc1_expert_weights`` / + ``fc2_expert_weights`` to :func:`cutlass_fused_moe`. + """ + if weight.dim() != 3: + raise ValueError( + f"weight must be 3D (num_experts, n, k/2); got shape {tuple(weight.shape)}" + ) + if weight.dtype != torch.uint8: + raise ValueError(f"weight must be uint8 (packed 4-bit); got {weight.dtype}") + if not weight.is_cuda: + raise ValueError("weight must live on CUDA") + + qtype_map = {"fp4": 1, "int4": 0} + if quant_type not in qtype_map: + raise ValueError( + f"quant_type must be one of {list(qtype_map)}; got {quant_type!r}" + ) + + weight = weight.contiguous() + out = torch.empty_like(weight) + + major, minor = get_compute_capability(weight.device) + device_arch = f"{major * 10 + minor}" + module = get_cutlass_fused_moe_module(device_arch) + module.interleave_moe_weights_for_sm90_mixed_gemm( + weight, out, qtype_map[quant_type] + ) + return out + + # ref: https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py#L121 @flashinfer_api(trace=cutlass_fused_moe_trace) def cutlass_fused_moe( @@ -904,7 +936,6 @@ def cutlass_fused_moe( use_mxfp8_act_scaling: bool = False, min_latency_mode: bool = False, use_packed_weights: bool = False, - use_wfp4afp8_humming: bool = False, tune_max_num_tokens: int = 8192, enable_pdl: Optional[bool] = None, activation_type: ActivationType = ActivationType.Swiglu, @@ -1016,11 +1047,6 @@ def cutlass_fused_moe( use_packed_weights : bool = False Whether to use packed uint4x2 weights passed as packed uint8 values. Defaults to False. - use_wfp4afp8_humming : bool = False - Selects the Humming-style MXFP4-weight x FP8-activation Hopper path with pre-MMA E8M0 - scale fusion. This flag is separate from W4A16 because both paths use uint8 FP4 weight - storage and ``use_w4_group_scaling=True``. - tune_max_num_tokens : int = 8192 Maximum number of tokens for tuning. Defaults to 8192. @@ -1093,11 +1119,6 @@ def cutlass_fused_moe( major, minor = get_compute_capability(input.device) device_arch = f"{major * 10 + minor}" - if use_wfp4afp8_humming and device_arch != "90": - raise NotImplementedError( - "Humming-style MXFP4 x FP8 fused MoE is only implemented for SM90." - ) - if min_latency_mode: raise NotImplementedError("min latency mode not yet implemented for Blackwell.") @@ -1163,7 +1184,6 @@ def cutlass_fused_moe( enable_pdl=enable_pdl, activation_type=activation_type, use_fused_finalize=use_fused_finalize, - use_wfp4afp8_humming=use_wfp4afp8_humming, profile_ids=profile_ids, workspace_buffer=workspace_buffer, ) @@ -1190,7 +1210,6 @@ def cutlass_fused_moe_workspace_size( use_mxfp8_act_scaling: bool = False, use_fused_finalize: bool = True, use_packed_weights: bool = False, - use_wfp4afp8_humming: bool = False, device: Optional[torch.device] = None, ) -> int: """Return the workspace buffer size in bytes required by :func:`cutlass_fused_moe`. @@ -1288,7 +1307,6 @@ def cutlass_fused_moe_workspace_size( use_mxfp8_act_scaling=use_mxfp8_act_scaling, use_fused_finalize=use_fused_finalize, use_packed_weights=use_packed_weights, - use_wfp4afp8_humming=use_wfp4afp8_humming, ) diff --git a/flashinfer/fused_moe/prepare.py b/flashinfer/fused_moe/prepare.py index 5c4a2d0c576..c83859a7bc7 100644 --- a/flashinfer/fused_moe/prepare.py +++ b/flashinfer/fused_moe/prepare.py @@ -1,4 +1,4 @@ -"""First-class weight-preparation helpers for the unified MoE API. +"""First-class NVFP4 weight-preparation helpers for the unified MoE API. Copyright (c) 2026 by FlashInfer team. @@ -24,26 +24,14 @@ ``TrtllmFp4Config.prepare_weights(...)`` / ``CuteDslConfig.prepare_weights(...)`` / ``TrtllmBf16Config.prepare_weights(...)`` / ... static helpers (see ``api.py``). -The SM90 Humming-style MXFP4 x FP8 helper is currently exposed as a flat helper -for the CUTLASS fused-MoE path. """ from __future__ import annotations -import functools -import struct from typing import Dict, Optional, Tuple, Union import torch -from ..api_logging import flashinfer_api -from ..trace.templates.moe import ( - sm90_mixed_gemm_humming_weight_preprocess_trace_dispatch, - sm90_mixed_gemm_scale_interleave_trace, - sm90_mixed_gemm_weight_interleave_trace, -) -from ..utils import get_compute_capability - # Module-level permute-index cache. Permute indices depend only on weight # dims, so the cache is safe to reuse across shapes and calls. _TRTLLM_PERMUTE_CACHE: dict = {} @@ -51,356 +39,6 @@ _TRTLLM_FP8_PER_TENSOR_PERMUTE_CACHE: dict = {} -# The E8M0 range clamp and residual-scale factorization are adapted from -# Humming's HummingLayer.may_process_fused_e8m0_scale: -# https://github.com/inclusionAI/humming/blob/f6241bba8d507c19ca9ce4e5958a5d0641fc8eb4/humming/layer.py#L322-L362 -def _preprocess_humming_e8m0_weight_scale( - raw_scale: torch.Tensor, - max_range: int = 11, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Clamp Humming fused-E8M0 scales into offset, residual, and FP4 delta. - - The Humming layer computes this range clamp independently per expert. The - offset tensor is consumed by the pre-MMA FP4->E4M3 conversion; residual is - one FP32 scale per expert; delta rewrites clamped FP4 payload values. - """ - if raw_scale.dim() != 3: - raise ValueError( - "raw_scale must be 3D (num_experts, rows, K/32); " - f"got shape {tuple(raw_scale.shape)}" - ) - if raw_scale.dtype != torch.uint8: - raise ValueError(f"raw_scale must be uint8 E8M0 bytes; got {raw_scale.dtype}") - if not raw_scale.is_cuda: - raise ValueError("raw_scale must live on CUDA") - # The fused conversion adds max_range + 1 to FP4 exponent code 3; - # E4M3 exponent code 15 therefore limits max_range to 11. - if max_range < 0 or max_range > 11: - raise ValueError(f"max_range must be in [0, 11]; got {max_range}") - - num_experts = raw_scale.shape[0] - scale_view = raw_scale.contiguous().view(num_experts, -1) - scale_max = scale_view.max(dim=1, keepdim=True).values - scale_min = scale_view.min(dim=1, keepdim=True).values - scale_range = scale_max - scale_min - max_range_tensor = torch.tensor( - max_range, dtype=torch.uint8, device=raw_scale.device - ) - scale_range = torch.minimum(scale_range, max_range_tensor) - scale_min_new = scale_max - scale_range - - clamped_scale = scale_view.maximum(scale_min_new) - delta_scale_offsets = (clamped_scale - scale_view).to(torch.uint8) - offset = torch.bitwise_and(clamped_scale - scale_min_new + 1, 0x0F).to(torch.uint8) - residual = torch.exp2(scale_min_new.squeeze(1).to(torch.float32) - 127.0) * 0.5 - return ( - offset.view_as(raw_scale).contiguous(), - residual.contiguous(), - delta_scale_offsets.view_as(raw_scale).contiguous(), - ) - - -# The delta-scale FP4 payload rewrite semantics are adapted from Humming's -# process_mxfp4_w4a8 implementation: -# https://github.com/inclusionAI/humming/blob/f6241bba8d507c19ca9ce4e5958a5d0641fc8eb4/humming/include/humming/kernel/process_mxfp4.cuh#L6-L69 -@functools.cache -def _humming_mxfp4_w4a8_rewrite_lut_cpu() -> torch.Tensor: - def float_from_bits(bits: int) -> float: - return struct.unpack("f", struct.pack("I", bits & 0xFFFFFFFF))[0] - - def bits_from_float(value: float) -> int: - return struct.unpack("I", struct.pack("f", value))[0] - - def dequant_fp4_val(code: int) -> float: - sign = (code & 0x8) << 28 - other = (code & 0x7) << 22 - return float_from_bits(sign | other) - - def quant_to_fp4_val(value: float) -> int: - value_bits = bits_from_float(value) - mask = 0x81C00000 - rz_bits = value_bits & mask - ru_bits = (value_bits + 0x00200000) & mask - rz_value = float_from_bits(rz_bits) - ru_value = float_from_bits(ru_bits) - rounded_bits = ( - ru_bits if abs(value - rz_value) >= abs(value - ru_value) else rz_bits - ) - return ((rounded_bits & 0x80000000) >> 28) | ((rounded_bits & 0x01C00000) >> 22) - - lut = torch.empty((256, 16), dtype=torch.uint8) - for delta in range(256): - scale_factor = float_from_bits(0x3F800000 - (delta << 23)) - for code in range(16): - normalized_code = 0 if code == 8 else code - if delta: - normalized_code = quant_to_fp4_val( - dequant_fp4_val(normalized_code) * scale_factor - ) - lut[delta, code] = normalized_code - return lut - - -def _process_humming_mxfp4_w4a8_payload( - weight: torch.Tensor, - delta_scale_offsets: torch.Tensor, -) -> torch.Tensor: - if weight.dim() != 3: - raise ValueError( - "weight must be 3D (num_experts, rows, K/2); " - f"got shape {tuple(weight.shape)}" - ) - if weight.dtype != torch.uint8: - raise ValueError(f"weight must be packed uint8 FP4 payload; got {weight.dtype}") - if not weight.is_cuda: - raise ValueError("weight must live on CUDA") - if delta_scale_offsets.shape[0] != weight.shape[0]: - raise ValueError( - "delta_scale_offsets and weight must have the same num_experts; " - f"got {delta_scale_offsets.shape[0]} and {weight.shape[0]}" - ) - expected_delta_shape = ( - weight.shape[0], - weight.shape[1], - weight.shape[2] * 2 // 32, - ) - if tuple(delta_scale_offsets.shape) != expected_delta_shape: - raise ValueError( - "delta_scale_offsets must have shape " - f"{expected_delta_shape}; got {tuple(delta_scale_offsets.shape)}" - ) - if delta_scale_offsets.dtype != torch.uint8: - raise ValueError( - f"delta_scale_offsets must be uint8; got {delta_scale_offsets.dtype}" - ) - - lut = _humming_mxfp4_w4a8_rewrite_lut_cpu().to(weight.device) - lo = weight & 0x0F - hi = (weight >> 4) & 0x0F - fp4_codes = torch.stack([lo, hi], dim=-1).reshape(*weight.shape[:-1], -1) - delta = delta_scale_offsets.repeat_interleave(32, dim=-1).to(torch.long) - rewritten = lut[delta, fp4_codes.to(torch.long)] - processed = rewritten[..., 0::2] | (rewritten[..., 1::2] << 4) - return processed.contiguous() - - -@flashinfer_api(trace=sm90_mixed_gemm_humming_weight_preprocess_trace_dispatch) -def preprocess_moe_weights_for_sm90_mixed_gemm_humming( - weight: torch.Tensor, - raw_scale: torch.Tensor, - max_range: int = 11, - *, - interleave: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Prepare MXFP4 weights for the SM90 Humming-style FP8 activation path. - - Parameters - ---------- - weight : torch.Tensor - ``[num_experts, rows, K // 2]`` CUDA uint8 tensor containing packed - MXFP4 payload values. - raw_scale : torch.Tensor - ``[num_experts, rows, K // 32]`` CUDA uint8 tensor containing original - E8M0 MXFP4 weight scales. - max_range : int - Maximum per-expert E8M0 exponent range kept in the pre-MMA FP4->E4M3 - offset. Humming uses 11 for FP8 activation. - interleave : bool - If true, return tensors ready for ``cutlass_fused_moe``. If false, - return the logical processed weight and logical offset scale; this is - useful for validation against a dequantized or Humming reference. - - Returns - ------- - Tuple[torch.Tensor, torch.Tensor, torch.Tensor] - ``(weight_out, scale_out, residual)``. With ``interleave=True``, - ``weight_out`` is the SM90 mixed-input weight layout and ``scale_out`` - is the folded scale layout. With ``interleave=False``, they are the - logical processed packed weight and logical offset scale. ``residual`` - is one FP32 value per expert and should be folded into the routed-token - activation scale together with Humming's fixed ``2^6`` compensation. - - Notes - ----- - The E8M0 range clamp, residual-scale factorization, and FP4 payload-rewrite - scheme are adapted from `Humming `_. - """ - if weight.dim() != 3: - raise ValueError( - "weight must be 3D (num_experts, rows, K/2); " - f"got shape {tuple(weight.shape)}" - ) - k = weight.shape[2] * 2 - if k % 32 != 0: - raise ValueError(f"weight K dimension must be divisible by 32; got K={k}") - expected_scale_shape = ( - weight.shape[0], - weight.shape[1], - k // 32, - ) - if tuple(raw_scale.shape) != expected_scale_shape: - raise ValueError( - f"raw_scale must have shape {expected_scale_shape}; " - f"got {tuple(raw_scale.shape)}" - ) - if raw_scale.device != weight.device: - raise ValueError( - "raw_scale and weight must be on the same device; " - f"got {raw_scale.device} and {weight.device}" - ) - - offset, residual, delta_scale_offsets = _preprocess_humming_e8m0_weight_scale( - raw_scale, max_range - ) - processed_weight = _process_humming_mxfp4_w4a8_payload( - weight.contiguous(), delta_scale_offsets - ) - if not interleave: - return processed_weight, offset, residual - - return ( - interleave_moe_weights_for_sm90_mixed_gemm(processed_weight, "fp4_fp8"), - interleave_moe_scales_for_sm90_mixed_gemm(offset), - residual, - ) - - -@flashinfer_api(trace=sm90_mixed_gemm_scale_interleave_trace) -def interleave_moe_scales_for_sm90_mixed_gemm( - scales: torch.Tensor, - group_size: int = 32, -) -> torch.Tensor: - """Fold weight scales for the SM90 mixed-input MoE GEMM. - - Parameters - ---------- - scales : torch.Tensor - ``[num_experts, rows, K // group_size]`` tensor of scalar weight scales. - MXFP4 uses uint8 E8M0 scales with ``group_size=32``; W4A8 uses bf16 - bit-pattern scales with ``group_size=128``. - group_size : int - Weight quantization group size. - - Returns - ------- - torch.Tensor - Contiguous tensor with shape - ``[num_experts, rows // 64, K // 128, folded_m, physical_cols]``. - ``physical_cols`` is the number of scale elements in 16B and - ``folded_m`` is derived so each 64x128 logical scale block is stored as - a 16B-contiguous folded block. - """ - if scales.dim() != 3: - raise ValueError( - f"scales must be 3D (num_experts, rows, K/group_size); got {tuple(scales.shape)}" - ) - - if group_size <= 0 or 128 % group_size != 0: - raise ValueError(f"group_size={group_size} must be positive and divide 128") - scale_groups_per_k128 = 128 // group_size - element_bits = scales.element_size() * 8 - physical_cols = 128 // element_bits - if physical_cols < 1 or 128 % element_bits != 0: - raise ValueError( - f"scale dtype {scales.dtype} has unsupported element size {element_bits} bits" - ) - if physical_cols % scale_groups_per_k128 != 0: - raise ValueError( - f"scale dtype {scales.dtype} and group_size={group_size} do not form " - "an integer folded M slice" - ) - m_slices_per_m64 = physical_cols // scale_groups_per_k128 - if 64 % m_slices_per_m64 != 0: - raise ValueError( - f"folded M slices {m_slices_per_m64} must divide the logical M64 block" - ) - folded_m = 64 // m_slices_per_m64 - - e, rows, kgs = scales.shape - if rows % 64 != 0: - raise ValueError(f"scale rows={rows} must be divisible by 64") - if kgs % scale_groups_per_k128 != 0: - raise ValueError( - f"K/group_size={kgs} must be divisible by scale groups per K128 block " - f"{scale_groups_per_k128}" - ) - k128_blocks = kgs // scale_groups_per_k128 - return ( - scales.reshape( - e, - rows // 64, - m_slices_per_m64, - folded_m, - k128_blocks, - scale_groups_per_k128, - ) - .permute(0, 1, 4, 3, 2, 5) - .contiguous() - .reshape(e, rows // 64, k128_blocks, folded_m, physical_cols) - ) - - -@flashinfer_api(trace=sm90_mixed_gemm_weight_interleave_trace) -def interleave_moe_weights_for_sm90_mixed_gemm( - weight: torch.Tensor, - quant_type: str = "fp4", -) -> torch.Tensor: - """Interleave 4-bit packed MoE weights for the SM90 mixed-input GEMM. - - The SM90 mixed-dtype MoE GEMM (used by ``cutlass_fused_moe`` with - ``use_w4_group_scaling=True``) expects weights in a specific interleaved - layout; without preprocessing, the LUT-based FP4->BF16 conversion reads - bytes from the wrong positions and the output diverges from a dequantized - reference for any K > 128. TensorRT-LLM's W4A16 MoE runs the equivalent - preprocessing at weight-load time (see - ``interleave_4bit_weights_for_Hopper_mixed_gemm`` in TRT-LLM PR #12451). - - Parameters - ---------- - weight : torch.Tensor - ``[num_experts, n, k // 2]`` uint8 CUDA tensor (4-bit values packed - two-per-byte). - quant_type : str - ``"fp4"`` for MXFP4 (the W4A16 path), ``"fp4_fp8"`` for MXFP4 consumed - by the FP8/Humming-style pre-MMA-scale path, or ``"int4"`` for INT4 - (the W4A8 path). - - Returns - ------- - torch.Tensor - A new uint8 tensor with the same shape as ``weight`` holding the - interleaved layout. Feed this directly as ``fc1_expert_weights`` / - ``fc2_expert_weights`` to :func:`cutlass_fused_moe`. - """ - if weight.dim() != 3: - raise ValueError( - f"weight must be 3D (num_experts, n, k/2); got shape {tuple(weight.shape)}" - ) - if weight.dtype != torch.uint8: - raise ValueError(f"weight must be uint8 (packed 4-bit); got {weight.dtype}") - if not weight.is_cuda: - raise ValueError("weight must live on CUDA") - - qtype_map = {"fp4": 1, "fp4_fp8": 2, "int4": 0} - if quant_type not in qtype_map: - raise ValueError( - f"quant_type must be one of {list(qtype_map)}; got {quant_type!r}" - ) - - weight = weight.contiguous() - out = torch.empty_like(weight) - - from .core import get_cutlass_fused_moe_module - - major, minor = get_compute_capability(weight.device) - device_arch = f"{major * 10 + minor}" - module = get_cutlass_fused_moe_module(device_arch) - module.interleave_moe_weights_for_sm90_mixed_gemm( - weight, out, qtype_map[quant_type] - ) - return out - - def prepare_trtllm_fp4_weights( w1_bf16: torch.Tensor, w2_bf16: torch.Tensor, diff --git a/flashinfer/jit/fused_moe.py b/flashinfer/jit/fused_moe.py index 256a2aef190..b2c3972dba7 100644 --- a/flashinfer/jit/fused_moe.py +++ b/flashinfer/jit/fused_moe.py @@ -117,10 +117,9 @@ def gen_cutlass_fused_moe_sm90_module(use_fast_build: bool = False) -> JitSpec: "-DENABLE_BF16", "-DENABLE_FP8", "-DENABLE_FP8_BLOCK_SCALE" if is_cuda_version_at_least("12.8") else "", - "-DENABLE_FP4", + "-DENABLE_FP4" if is_cuda_version_at_least("12.8") else "", "-DUSING_OSS_CUTLASS_MOE_GEMM", "-DCUTLASS_ENABLE_GDC_FOR_SM90=1", - "-DCUTLASS_MIXED_GEMM_FP4_FP8_PREPROCESSED_SIGNS=1", ] return gen_cutlass_fused_moe_module(nvcc_flags, "90", use_fast_build) @@ -159,62 +158,63 @@ def gen_cutlass_fused_moe_module( except Exception as e: raise RuntimeError(f"Failed to generate Cutlass kernels: {e}") from e - sources = [ - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp8_uint4.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp8_fp8.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp8_fp4.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp4_fp4.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp32_fp32.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp16_uint8.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp16_uint4.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp16_fp16.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_uint8.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_uint4.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp8.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_bf16.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp4.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp16_fp4.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_mixed_utils.cu", - jit_env.FLASHINFER_CSRC_DIR - / "fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu", - jit_env.FLASHINFER_CSRC_DIR / "fused_moe/cutlass_backend/deepgemm_jit_setup.cu", - jit_env.FLASHINFER_CSRC_DIR - / "fused_moe/cutlass_backend/cutlass_fused_moe_instantiation.cu", - *(output_dir / kernel for kernel in output_dir.rglob("*.generated.cu")), - jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/envUtils.cpp", - jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/logger.cpp", - jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/stringUtils.cpp", - jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/tllmException.cpp", - jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/memoryUtils.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/preQuantScaleKernel.cu", - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp", - jit_env.FLASHINFER_CSRC_DIR / "nv_internal/tensorrt_llm/kernels/lora/lora.cpp", - ] - return gen_jit_spec( f"fused_moe_{device_arch}", - sources, + [ + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp8_uint4.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp8_fp8.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp8_fp4.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp4_fp4.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp32_fp32.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp16_uint8.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp16_uint4.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp16_fp16.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_uint8.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_uint4.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp8.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_bf16.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp4.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp16_fp4.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_mixed_utils.cu", + jit_env.FLASHINFER_CSRC_DIR + / "fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu", + jit_env.FLASHINFER_CSRC_DIR + / "fused_moe/cutlass_backend/deepgemm_jit_setup.cu", + jit_env.FLASHINFER_CSRC_DIR + / "fused_moe/cutlass_backend/cutlass_fused_moe_instantiation.cu", + # Add all generated kernels + *(output_dir / kernel for kernel in output_dir.rglob("*.generated.cu")), + jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/envUtils.cpp", + jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/logger.cpp", + jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/stringUtils.cpp", + jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/tllmException.cpp", + jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/memoryUtils.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/preQuantScaleKernel.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal/tensorrt_llm/kernels/lora/lora.cpp", + ], extra_cuda_cflags=nvcc_flags, extra_cflags=["-DFAST_BUILD"] if use_fast_build else [], extra_ldflags=["-lnvrtc"], diff --git a/flashinfer/jit/gemm/cutlass/generate_kernels.py b/flashinfer/jit/gemm/cutlass/generate_kernels.py index 1013fc5a57a..32484d15ff2 100644 --- a/flashinfer/jit/gemm/cutlass/generate_kernels.py +++ b/flashinfer/jit/gemm/cutlass/generate_kernels.py @@ -17,6 +17,7 @@ EpilogueScheduleTag, EpilogueScheduleSuffixes, ) +from ...cpp_ext import is_cuda_version_at_least ################################################################################ @@ -120,16 +121,6 @@ def GetDataTypeNames(type, is_mx_fpx=None): DataType.u4: "cutlass::uint4b_t", } -MixedInputScaleModeTag = { - "post_mma": "cutlass::gemm::collective::MixedInputScaleMode::kPostMma", - "pre_mma_e8m0": "cutlass::gemm::collective::MixedInputScaleMode::kPreMmaE8M0", -} - -Sm90MixedInputKernelScheduleTag = { - "single_warpgroup_prefill": "tensorrt_llm::cutlass_extensions::MainloopScheduleType::SINGLE_WARPGROUP_PREFILL", - "single_warpgroup_rolling": "tensorrt_llm::cutlass_extensions::MainloopScheduleType::SINGLE_WARPGROUP_ROLLING", -} - ################################################################################ # A data structure holding all info to instantiate gemm launchers in TRT LLM. @@ -153,8 +144,6 @@ def __init__( epi_schedule, epi_fusion=None, is_mx_fpx=False, - mixed_input_scale_mode="post_mma", - sm90_mixed_input_kernel_type="warp_specialized", dynamic_cga=False, swap_ab=False, ): @@ -176,8 +165,6 @@ def __init__( self.epi_schedule = epi_schedule self.epi_fusion = epi_fusion self.is_mx_fpx = is_mx_fpx - self.mixed_input_scale_mode = mixed_input_scale_mode - self.sm90_mixed_input_kernel_type = sm90_mixed_input_kernel_type self.swap_ab = swap_ab def __repr__(self): @@ -210,10 +197,6 @@ def __repr__(self): "_mxfpx_" if self.is_mx_fpx else "", "_swap_ab" if self.swap_ab else "", ) - if self.mixed_input_scale_mode != "post_mma": - hopper_suffix += f"_{self.mixed_input_scale_mode}" - if self.sm90_mixed_input_kernel_type != "warp_specialized": - hopper_suffix += f"_{self.sm90_mixed_input_kernel_type}" if self.arch >= 90: return kernel_prefix + hopper_suffix @@ -255,35 +238,14 @@ def instantiate_operation_tma_warp_specialized(operation): {out_tag}*, int, int, int, const int, tensorrt_llm::cutlass_extensions::CutlassGemmConfig, char*, size_t, cudaStream_t, int* );""" elif operation.gemm_kind == GemmKind.Grouped: - if operation.arch == 90 and operation.act_type != operation.weight_type: + if operation.act_type != operation.weight_type and ( + operation.act_type != DataType.e4m3 or operation.weight_type != e2m1 + ): # Mixed MoE GEMM weight_tag = CudaTypeName[operation.weight_type] - optional_template_args = [] - if operation.mixed_input_scale_mode != "post_mma": - optional_template_args.append( - MixedInputScaleModeTag[operation.mixed_input_scale_mode] - ) - if operation.sm90_mixed_input_kernel_type != "warp_specialized": - if operation.mixed_input_scale_mode == "post_mma": - optional_template_args.append( - MixedInputScaleModeTag[operation.mixed_input_scale_mode] - ) - optional_template_args.append( - Sm90MixedInputKernelScheduleTag[ - operation.sm90_mixed_input_kernel_type - ] - ) - optional_template_args = "".join( - f", {arg}" for arg in optional_template_args - ) - launcher_name = ( - "sm90_generic_mixed_moe_gemm_kernelLauncher" - if operation.sm90_mixed_input_kernel_type == "warp_specialized" - else "sm90_generic_mixed_moe_small_k_kernelLauncher" - ) instantiation = f""" -template void {launcher_name}<{act_tag}, {weight_tag}, {out_tag}, -{epi_tag}, {cute_cta_shape}, {cute_cga_shape}, {kernel_sched}, {epi_sched}, {quant_op}{optional_template_args}> ( +template void sm90_generic_mixed_moe_gemm_kernelLauncher<{act_tag}, {weight_tag}, {out_tag}, +{epi_tag}, {cute_cta_shape}, {cute_cga_shape}, {kernel_sched}, {epi_sched}, {quant_op}> ( GroupedGemmInput<{act_tag}, {weight_tag}, {out_tag}, {out_tag}>inputs, TmaWarpSpecializedGroupedGemmInput hopper_inputs, int sm_count_, size_t* workspace_size); """ else: @@ -696,43 +658,34 @@ def generate_sm90_mixed_type_grouped_gemm_operations(is_arch_enabled): (DataType.e4m3, DataType.u4, DataType.bf16, DataType.bf16, DataType.bf16), ] - supported_dtypes_fp4 = [ - ( - DataType.bf16, - e2m1, - DataType.ue8m0, - DataType.bf16, - DataType.bf16, - ), - ( - DataType.e4m3, - e2m1, - DataType.ue8m0, - DataType.f16, - DataType.f16, - ), - ( - DataType.e4m3, - e2m1, - DataType.ue8m0, - DataType.bf16, - DataType.bf16, - ), - ] + if is_cuda_version_at_least("12.8"): + supported_dtypes_fp4 = [ + (DataType.f16, DataType.e2m1, DataType.ue8m0, DataType.f16, DataType.f16), + ( + DataType.bf16, + DataType.e2m1, + DataType.ue8m0, + DataType.bf16, + DataType.bf16, + ), + ] + else: + supported_dtypes_fp4 = [] quant_ops = [TrtLlm_QuantOp.finegrained_scale_only] epi_tags = [TrtLlm_EpilogueTag.epilogue_op_default] - cta_shapes_mnk_mixed_input = list(product([64], [16, 32, 64, 128], [128, 256, 512])) - cta_shapes_mnk_mixed_input.extend( - product([128], [16, 32, 64, 128], [128, 256, 512]) - ) - cta_shapes_mnk_mixed_input.extend([(128, 256, k_tile) for k_tile in [128, 256]]) - cta_shapes_mnk_mixed_input.extend([(256, 128, k_tile) for k_tile in [128, 256]]) - cta_shapes_mnk_mixed_input.append((256, 256, 128)) - cta_shapes_mnk_int4 = list(cta_shapes_mnk_mixed_input) - cta_shapes_mnk_fp4 = list(cta_shapes_mnk_mixed_input) + M_TILES = [64, 128] # Currently M tile must be 128 for Grouped GEMM + N_TILES = [16, 32, 64, 128] + K_TILES = [128, 256, 512] + cta_shapes_mnk_int4 = list(product(M_TILES, N_TILES, K_TILES)) + + M_TILES = [64, 128] # Currently M tile must be 128 for Grouped GEMM + N_TILES = [16, 32, 64] + K_TILES = [128, 256] + cta_shapes_mnk_fp4 = list(product(M_TILES, N_TILES, K_TILES)) + cta_shapes_mnk_fp4.append((128, 128, 128)) warp_shape = [0, 0, 0] # ignored except for naming stages = 0 # auto @@ -749,9 +702,6 @@ def generate_sm90_mixed_type_grouped_gemm_operations(is_arch_enabled): operations = list() for dtype_combo, quant_op, epi_tag, cta_shape_mnk, cga_shape in partial_args: - is_fp8_mxfp4 = dtype_combo[0] == DataType.e4m3 and dtype_combo[1] == e2m1 - mixed_input_scale_mode = "pre_mma_e8m0" if is_fp8_mxfp4 else "post_mma" - use_coop = cta_shape_mnk[0] >= 128 mainloop_schedules = ( [ @@ -763,6 +713,13 @@ def generate_sm90_mixed_type_grouped_gemm_operations(is_arch_enabled): ) epi_schedule = EpilogueScheduleType.TmaWarpSpecializedCooperative for mainloop_schedule in mainloop_schedules: + if ( + cta_shape_mnk[0] == 128 + and cta_shape_mnk[1] == 128 + and mainloop_schedule + == KernelScheduleType.TmaWarpSpecializedCooperative + ): + continue moe_gemm_operation = TrtLlm_GemmLauncher( GemmKind.Grouped, arch, @@ -775,36 +732,8 @@ def generate_sm90_mixed_type_grouped_gemm_operations(is_arch_enabled): cga_shape, mainloop_schedule, epi_schedule, - mixed_input_scale_mode=mixed_input_scale_mode, ) operations.append(moe_gemm_operation) - - small_k_shapes = [(128, token_tile, 128) for token_tile in [8, 16, 32, 40]] - small_k_kernel_types = [ - "single_warpgroup_prefill", - "single_warpgroup_rolling", - ] - for dtype_combo in supported_dtypes_fp4: - if dtype_combo[0] != DataType.e4m3: - continue - for cta_shape_mnk, kernel_type in product(small_k_shapes, small_k_kernel_types): - operations.append( - TrtLlm_GemmLauncher( - GemmKind.Grouped, - arch, - *dtype_combo, - TrtLlm_QuantOp.finegrained_scale_only, - TrtLlm_EpilogueTag.epilogue_op_default, - cta_shape_mnk, - warp_shape, - 3, - (1, 1, 1), - KernelScheduleType.TmaWarpSpecializedPingpong, - EpilogueScheduleType.TmaWarpSpecializedCooperative, - mixed_input_scale_mode="pre_mma_e8m0", - sm90_mixed_input_kernel_type=kernel_type, - ) - ) return operations @@ -1132,14 +1061,15 @@ def has_arch(sm): def should_skip(op): return False # All kernels have a public implementation - # SM90 mixed-input grouped GEMMs have a dedicated launcher. + # The mixed dtype grouped gemm for w4afp8 has a different launcher def is_mixed_dtype_grouped(op): if isinstance(op, GemmSm80LauncherConfig): return False + # Only w4a8fp8 and not wfp4afp8 return ( - (op.arch == 90) - and (op.act_type != op.weight_type) + (op.act_type != op.weight_type) and (op.gemm_kind == GemmKind.Grouped) + and (op.act_type != DataType.e4m3 or op.weight_type != e2m1) ) # Fix OOM error in CI. If len(operations) is more than GROUP_SIZE, it will be split into multiple sub groups. diff --git a/flashinfer/trace/templates/moe.py b/flashinfer/trace/templates/moe.py index 310dc3c5707..4c7117807f5 100644 --- a/flashinfer/trace/templates/moe.py +++ b/flashinfer/trace/templates/moe.py @@ -30,138 +30,6 @@ from ._init_helpers import fp8_block_quant_1d, fp8_block_quant_2d from .quantize import _fp4_quantize_reference -# --------------------------------------------------------------------------- -# SM90 mixed-input weight preparation -# --------------------------------------------------------------------------- - -_SM90_MIXED_WEIGHT_PREP_AXES: dict[str, Var | Const] = { - "num_experts": Var(description="Number of local experts."), - "rows": Var(description="Logical weight rows per expert."), - "packed_k": Var(description="Packed 4-bit K bytes per row."), - "scale_groups": Var(description="Weight-scale groups per row."), -} - -_SM90_MIXED_WEIGHT_PREP_INPUTS: dict[str, Tensor | Scalar] = { - "weight": Tensor( - ["num_experts", "rows", "packed_k"], - description="Packed 4-bit expert weights.", - ), - "raw_scale": Tensor( - ["num_experts", "rows", "scale_groups"], - description="Logical E8M0 weight scales.", - ), - "max_range": Scalar("int32", description="Maximum fused E8M0 exponent range."), - "interleave": Scalar( - "bool", description="Whether to emit the SM90 physical layout." - ), -} - -sm90_mixed_gemm_humming_weight_preprocess_trace = TraceTemplate( - op_type="moe_preprocess", - name_prefix="sm90_mixed_gemm_humming_weight_preprocess", - description="Prepare Humming-style MXFP4 weights and scales for SM90 mixed-input MoE.", - axes={ - **_SM90_MIXED_WEIGHT_PREP_AXES, - "m64_blocks": Var(description="Folded 64-row scale blocks."), - "k128_blocks": Var(description="Folded K128 scale blocks."), - "folded_m": Var(description="Physical folded-M scale extent."), - "physical_cols": Var(description="Physical scale columns per folded block."), - }, - inputs=_SM90_MIXED_WEIGHT_PREP_INPUTS, - outputs={ - "processed_weight": Tensor( - ["num_experts", "rows", "packed_k"], dtype_from="weight" - ), - "processed_scale": Tensor( - ["num_experts", "m64_blocks", "k128_blocks", "folded_m", "physical_cols"], - dtype_from="raw_scale", - ), - "residual": Tensor(["num_experts"], dtype="float32"), - }, - constraints=["packed_k == scale_groups * 16"], - tags=["moe:sm90", "quantization:mxfp4"], -) - -sm90_mixed_gemm_humming_weight_preprocess_logical_trace = TraceTemplate( - op_type="moe_preprocess", - name_prefix="sm90_mixed_gemm_humming_weight_preprocess_logical", - description="Prepare logical Humming-style MXFP4 payloads and exponent offsets.", - axes=_SM90_MIXED_WEIGHT_PREP_AXES, - inputs=_SM90_MIXED_WEIGHT_PREP_INPUTS, - outputs={ - "processed_weight": Tensor( - ["num_experts", "rows", "packed_k"], dtype_from="weight" - ), - "processed_scale": Tensor( - ["num_experts", "rows", "scale_groups"], dtype_from="raw_scale" - ), - "residual": Tensor(["num_experts"], dtype="float32"), - }, - constraints=["packed_k == scale_groups * 16"], - tags=["moe:sm90", "quantization:mxfp4"], -) - - -def sm90_mixed_gemm_humming_weight_preprocess_trace_dispatch(**kwargs): - return ( - sm90_mixed_gemm_humming_weight_preprocess_trace - if kwargs.get("interleave", True) - else sm90_mixed_gemm_humming_weight_preprocess_logical_trace - ) - - -sm90_mixed_gemm_humming_weight_preprocess_trace_dispatch.templates = [ # type: ignore[attr-defined] - sm90_mixed_gemm_humming_weight_preprocess_trace, - sm90_mixed_gemm_humming_weight_preprocess_logical_trace, -] - -sm90_mixed_gemm_scale_interleave_trace = TraceTemplate( - op_type="moe_preprocess", - name_prefix="sm90_mixed_gemm_scale_interleave", - description="Fold logical weight scales into the SM90 mixed-input layout.", - axes={ - "num_experts": Var(description="Number of local experts."), - "rows": Var(description="Logical scale rows per expert."), - "scale_groups": Var(description="Logical scale groups per row."), - "m64_blocks": Var(description="Folded 64-row scale blocks."), - "k128_blocks": Var(description="Folded K128 scale blocks."), - "folded_m": Var(description="Physical folded-M scale extent."), - "physical_cols": Var(description="Physical scale columns per folded block."), - }, - inputs={ - "scales": Tensor(["num_experts", "rows", "scale_groups"]), - "group_size": Scalar("int32", description="Weight quantization group size."), - }, - outputs={ - "interleaved_scales": Tensor( - ["num_experts", "m64_blocks", "k128_blocks", "folded_m", "physical_cols"], - dtype_from="scales", - ) - }, - tags=["moe:sm90", "layout:folded_scale"], -) - -sm90_mixed_gemm_weight_interleave_trace = TraceTemplate( - op_type="moe_preprocess", - name_prefix="sm90_mixed_gemm_weight_interleave", - description="Interleave packed 4-bit weights for the SM90 mixed-input MoE GEMM.", - axes={ - "num_experts": Var(description="Number of local experts."), - "rows": Var(description="Logical weight rows per expert."), - "packed_k": Var(description="Packed 4-bit K bytes per row."), - }, - inputs={ - "weight": Tensor(["num_experts", "rows", "packed_k"]), - "quant_type": Scalar("str", description="4-bit payload/interleave variant."), - }, - outputs={ - "interleaved_weight": Tensor( - ["num_experts", "rows", "packed_k"], dtype_from="weight" - ) - }, - tags=["moe:sm90", "layout:mixed_input"], -) - # --------------------------------------------------------------------------- # Shared GEMM computation helper # --------------------------------------------------------------------------- diff --git a/tests/moe/test_trtllm_cutlass_fused_moe.py b/tests/moe/test_trtllm_cutlass_fused_moe.py index f88ed4569f0..461808e0530 100644 --- a/tests/moe/test_trtllm_cutlass_fused_moe.py +++ b/tests/moe/test_trtllm_cutlass_fused_moe.py @@ -15,8 +15,6 @@ """ from contextlib import nullcontext -import os -import struct import pytest from flashinfer.fused_moe.core import ActivationType @@ -273,7 +271,7 @@ def torch_moe_w4a8( scale1 = fc1_input_scale[expert_id] if fc1_pre_quant_scale is not None: - expert_inputs_scaled = expert_inputs * fc1_pre_quant_scale + expert_inputs_scaled = expert_inputs * fc1_pre_quant_scale[expert_id] else: expert_inputs_scaled = expert_inputs inp_q = ( @@ -293,7 +291,7 @@ def torch_moe_w4a8( if fc2_input_scale is not None: scale2 = fc2_input_scale[expert_id] if fc2_pre_quant_scale is not None: - inter_scaled = inter * fc2_pre_quant_scale + inter_scaled = inter * fc2_pre_quant_scale[expert_id] else: inter_scaled = inter inter_q = ( @@ -1851,8 +1849,23 @@ def test_moe_bf16_mxfp4( output=flash_output, ) - dq_mfxp4_w1 = _dequant_mxfp4_on_device(w1, w1_scale) - dq_mfxp4_w2 = _dequant_mxfp4_on_device(w2, w2_scale) + dq_mfxp4_w1 = ( + dequant_mxfp4_batches_host( + w1.cpu(), + w1_scale.cpu(), + ) + .cuda() + .to(torch.bfloat16) + ) + + dq_mfxp4_w2 = ( + dequant_mxfp4_batches_host( + w2.cpu(), + w2_scale.cpu(), + ) + .cuda() + .to(torch.bfloat16) + ) # Use original weights for reference computation ref_output = compute_with_experts( @@ -1917,12 +1930,12 @@ def test_moe_w4a8( torch.randn(e, n, k // group_size, dtype=dtype, device="cuda") * affine_coeff ) - # The current W4A8 contract shares prequant scales across experts. - w1_pre_quant_scale = torch.rand(k, dtype=dtype, device="cuda") * 0.1 + 0.95 - w2_pre_quant_scale = torch.rand(n, dtype=dtype, device="cuda") * 0.1 + 0.95 - w3_pre_quant_scale = torch.rand(k, dtype=dtype, device="cuda") * 0.1 + 0.95 + # per channel pre quant scales + w1_pre_quant_scale = torch.rand(e, k, dtype=dtype, device="cuda") * 0.1 + 0.95 + w2_pre_quant_scale = torch.rand(e, n, dtype=dtype, device="cuda") * 0.1 + 0.95 + w3_pre_quant_scale = torch.rand(e, k, dtype=dtype, device="cuda") * 0.1 + 0.95 - input_scale = torch.rand(1, dtype=torch.float32, device="cuda") * 0.2 + 0.1 + input_scale = torch.rand(e, 1, dtype=torch.float32, device="cuda") * 0.2 + 0.1 weight_scale_2 = torch.ones(e, 1, dtype=torch.float32, device="cuda") fc1_weights = torch.cat([w3_weight, w1_weight], dim=1) @@ -1937,13 +1950,27 @@ def test_moe_w4a8( fc2_weights.contiguous().view(torch.uint8), "int4" ) + def interleave_weights(w: torch.Tensor, dim: int) -> torch.Tensor: + # Factors are chosen based on TRTLLM's quantization.py + interleave_factor = 4 if dim % 512 == 0 else (2 if dim % 256 == 0 else 1) + s = w.shape + w_interleaved = ( + w.reshape(s[0], s[1], s[2] // interleave_factor, interleave_factor) + .permute(0, 2, 1, 3) + .reshape(s[0], s[2] // interleave_factor, s[1] * interleave_factor) + .contiguous() + ) + return w_interleaved + w3_w1_scales = torch.cat([w3_scale, w1_scale], dim=1) + w3_w1_scales_int = interleave_weights(w3_w1_scales, k) + w2_scales_int = interleave_weights(w2_scale, n) # act scales w3_w1_pre_quant_max = torch.max(w1_pre_quant_scale, w3_pre_quant_scale) w3_w1_input_scale_max = input_scale.max() fc31_act_scale = (w3_w1_pre_quant_max / w3_w1_input_scale_max).to(dtype) - fc2_act_scale = (w2_pre_quant_scale / input_scale).to(dtype) + fc2_act_scale = (w2_pre_quant_scale / input_scale).to(dtype).unsqueeze(-1) fc31_alpha = (weight_scale_2.squeeze(-1) * w3_w1_input_scale_max).float() fc2_alpha = (weight_scale_2.squeeze(-1) * input_scale.squeeze(-1)).float() @@ -1951,16 +1978,21 @@ def test_moe_w4a8( zero_1 = torch.empty(0, dtype=dtype, device="cuda") zero_2 = torch.empty(0, dtype=dtype, device="cuda") - # SM90 mixed-input kernels read INT4 weight scales as bf16 bit patterns in - # folded 64x128 scale blocks. - w3_w1_scales_out = fused_moe.interleave_moe_scales_for_sm90_mixed_gemm( - w3_w1_scales.to(torch.bfloat16).view(dtype), group_size + # SM90 requires bfloat16 bit patterns + sm = ( + torch.cuda.get_device_capability()[0] * 10 + + torch.cuda.get_device_capability()[1] ) - w2_scales_out = fused_moe.interleave_moe_scales_for_sm90_mixed_gemm( - w2_scale.to(torch.bfloat16).view(dtype), group_size - ) - fc31_act_out = fc31_act_scale.to(torch.bfloat16).view(dtype) - fc2_act_out = fc2_act_scale.to(torch.bfloat16).view(dtype) + if sm >= 90: + w3_w1_scales_out = w3_w1_scales_int.to(torch.bfloat16).view(dtype) + w2_scales_out = w2_scales_int.to(torch.bfloat16).view(dtype) + fc31_act_out = fc31_act_scale.to(torch.bfloat16).view(dtype) + fc2_act_out = fc2_act_scale.to(torch.bfloat16).view(dtype) + else: + w3_w1_scales_out = w3_w1_scales_int.to(dtype) + w2_scales_out = w2_scales_int.to(dtype) + fc31_act_out = fc31_act_scale + fc2_act_out = fc2_act_scale quant_scales = ( w3_w1_scales_out, @@ -2018,12 +2050,6 @@ def test_moe_w4a8( w31_weight_dequant = torch.stack(w31_weight_list, dim=0) # [e, 2N, K] w2_weight_dequant = torch.stack(w2_weight_list, dim=0) # [e, K, N] - input_scale_for_ref = torch.full( - (num_experts,), - input_scale.item(), - dtype=input_scale.dtype, - device=input_scale.device, - ) ref_output = torch_moe_w4a8( num_experts, x, @@ -2031,8 +2057,8 @@ def test_moe_w4a8( w2_weight_dequant, selected_experts, routing_weights, - fc1_input_scale=input_scale_for_ref, - fc2_input_scale=input_scale_for_ref, + fc1_input_scale=input_scale.squeeze(-1), + fc2_input_scale=input_scale.squeeze(-1), fc1_pre_quant_scale=torch.max(w1_pre_quant_scale, w3_pre_quant_scale), fc2_pre_quant_scale=w2_pre_quant_scale, fc1_weight_scale_2=weight_scale_2.squeeze(-1), @@ -2653,6 +2679,18 @@ def test_moe_mxfp8_mxfp4_ndim_padding_safety( ) +# ============================================================================ +# SM90 mixed-input MoE tests — PR #3084 +# +# Exercise the W4A16 (MXFP4 x BF16) and W4A8 (INT4 x FP8) paths with the +# preprocessing helpers exposed by this PR: weights go through +# ``interleave_moe_weights_for_sm90_mixed_gemm``, MXFP4 block scales go +# through ``interleave_moe_scales_for_sm90_mixed_gemm``, and W4A8 weight +# scales use a local group-wise reshape+permute (factor = 4 / 2 / 1 based on +# whether K is divisible by 512 / 256) to match the W4A8 kernel layout. +# ============================================================================ + + _MXFP4_LUT = ( 0.0, 0.5, @@ -2691,250 +2729,6 @@ def _dequant_mxfp4_on_device( return (values * scale).to(torch.bfloat16) -def _dequant_mxfp4_humming_prescale_on_device( - w_fp4: torch.Tensor, exp_offset: torch.Tensor -) -> torch.Tensor: - """Reference for the Humming-style MXFP4 fast path. - - ``exp_offset`` is not the original e8m0 scale. It is the preprocessed - offset byte consumed by the pre-MMA FP4->E4M3 conversion. - """ - lo = w_fp4 & 0x0F - hi = (w_fp4 >> 4) & 0x0F - fp4_code = torch.stack([lo, hi], dim=-1).reshape(*w_fp4.shape[:-1], -1) - offset = exp_offset.repeat_interleave(32, dim=-1).to(torch.int32) - em_code = (fp4_code & 0x07).to(torch.int32) - em = torch.zeros_like(em_code) - em = torch.where(em_code == 1, offset * 8, em) - em = torch.where(em_code == 2, offset * 8 + 0x08, em) - em = torch.where(em_code == 3, offset * 8 + 0x0C, em) - em = torch.where(em_code >= 4, offset * 8 + 0x10 + (em_code - 4) * 4, em) - sign = (fp4_code.to(torch.int32) & 0x08) << 4 - fp8_raw = (sign | em).to(torch.uint8).contiguous() - return fp8_raw.view(torch.float8_e4m3fn).to(torch.float32) - - -def _make_humming_e8m0_weight_scale( - shape: tuple[int, ...], - device: torch.device, - low: int = 114, - high: int = 128, -) -> torch.Tensor: - """Generate deterministic raw E8M0 scale bytes for Humming preprocessing.""" - numel = 1 - for dim in shape: - numel *= dim - values = torch.arange(numel, device=device, dtype=torch.int32) - return (low + values.remainder(high - low)).to(torch.uint8).reshape(shape) - - -def _reference_humming_e8m0_weight_scale( - raw_scale: torch.Tensor, - max_range: int = 11, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Self-contained reference for the Humming-style scale clamp.""" - num_experts = raw_scale.shape[0] - scale_view = raw_scale.contiguous().view(num_experts, -1) - scale_max = scale_view.max(dim=1, keepdim=True).values - scale_min = scale_view.min(dim=1, keepdim=True).values - max_range_tensor = torch.tensor( - max_range, dtype=torch.uint8, device=raw_scale.device - ) - scale_range = torch.minimum(scale_max - scale_min, max_range_tensor) - scale_min_new = scale_max - scale_range - - clamped_scale = scale_view.maximum(scale_min_new) - delta_scale_offsets = (clamped_scale - scale_view).to(torch.uint8) - offset = torch.bitwise_and(clamped_scale - scale_min_new + 1, 0x0F).to(torch.uint8) - residual = torch.exp2(scale_min_new.squeeze(1).to(torch.float32) - 127.0) * 0.5 - return ( - offset.view_as(raw_scale).contiguous(), - residual.contiguous(), - delta_scale_offsets.view_as(raw_scale).contiguous(), - ) - - -def _humming_payload_rewrite_lut(device: torch.device) -> torch.Tensor: - def float_from_bits(bits: int) -> float: - return struct.unpack("f", struct.pack("I", bits & 0xFFFFFFFF))[0] - - def bits_from_float(value: float) -> int: - return struct.unpack("I", struct.pack("f", value))[0] - - def dequant_fp4_val(code: int) -> float: - sign = (code & 0x8) << 28 - other = (code & 0x7) << 22 - return float_from_bits(sign | other) - - def quant_to_fp4_val(value: float) -> int: - value_bits = bits_from_float(value) - mask = 0x81C00000 - rz_bits = value_bits & mask - ru_bits = (value_bits + 0x00200000) & mask - rz_value = float_from_bits(rz_bits) - ru_value = float_from_bits(ru_bits) - rounded_bits = ( - ru_bits if abs(value - rz_value) >= abs(value - ru_value) else rz_bits - ) - return ((rounded_bits & 0x80000000) >> 28) | ((rounded_bits & 0x01C00000) >> 22) - - lut = torch.empty((256, 16), dtype=torch.uint8) - for delta in range(256): - scale_factor = float_from_bits(0x3F800000 - (delta << 23)) - for code in range(16): - normalized_code = 0 if code == 8 else code - if delta: - normalized_code = quant_to_fp4_val( - dequant_fp4_val(normalized_code) * scale_factor - ) - lut[delta, code] = normalized_code - return lut.to(device) - - -def _reference_humming_payload_rewrite( - weight: torch.Tensor, - delta_scale_offsets: torch.Tensor, -) -> torch.Tensor: - lut = _humming_payload_rewrite_lut(weight.device) - lo = weight & 0x0F - hi = (weight >> 4) & 0x0F - fp4_codes = torch.stack([lo, hi], dim=-1).reshape(*weight.shape[:-1], -1) - delta = delta_scale_offsets.repeat_interleave(32, dim=-1).to(torch.long) - rewritten = lut[delta, fp4_codes.to(torch.long)] - return (rewritten[..., 0::2] | (rewritten[..., 1::2] << 4)).contiguous() - - -def _assert_humming_payload_rewrite_bit_equal( - weight: torch.Tensor, - delta_scale_offsets: torch.Tensor, - processed_weight: torch.Tensor, -) -> None: - reference_weight = _reference_humming_payload_rewrite( - weight, delta_scale_offsets.contiguous() - ) - if not torch.equal(processed_weight, reference_weight): - raise AssertionError("Humming-style payload rewrite is not bit-exact") - - -PHASE3_HUMMING_E2E_CASES = { - "small": { - "seed": 7, - "e": 2, - "m": 4, - "n": 512, - "k": 512, - "top_k": 2, - "raw_scale": (118, 122), - "torch_ref_tolerance": (5e-2, 1e-3), - "torch_ref_max_bad": 0, - "description": "default small smoke, offset range 1..4", - }, - "wide_offset": { - "seed": 7, - "e": 2, - "m": 4, - "n": 512, - "k": 512, - "top_k": 2, - "raw_scale": (114, 128), - "torch_ref_tolerance": (2e-1, 5e-1), - # The PyTorch dequant reference is not the golden for this wider dynamic - # range. This small bad-count budget prevents sparse BF16/FP8 - # accumulation differences from hiding broad correctness failures. - "torch_ref_max_bad": 8, - "description": "small smoke with full Humming offset range 1..12", - }, - "small_k384": { - "seed": 17, - "e": 8, - "m": 16, - "n": 384, - "k": 384, - "top_k": 2, - "raw_scale": (118, 122), - "torch_ref_tolerance": (5e-2, 1e-3), - "torch_ref_max_bad": 0, - "description": "single-warpgroup prefill-all boundary at three K128 tiles", - }, - "small_k768": { - "seed": 19, - "e": 8, - "m": 16, - "n": 768, - "k": 768, - "top_k": 2, - "raw_scale": (118, 122), - "torch_ref_tolerance": (5e-2, 1e-3), - "torch_ref_max_bad": 0, - "description": "single-warpgroup rolling-refill coverage at six K128 tiles", - }, - "e256_config": { - "seed": 11, - "e": 256, - "m": 8, - "n": 256, - "k": 4096, - "top_k": 1, - "raw_scale": (118, 122), - "torch_ref_tolerance": (5e-2, 1e-3), - "torch_ref_max_bad": 0, - "description": "E256-style FC1 N512/K4096 full-tactic coverage", - }, - "e32_config": { - "seed": 13, - "e": 32, - "m": 8, - "n": 2048, - "k": 4096, - "top_k": 1, - "raw_scale": (118, 122), - "torch_ref_tolerance": (5e-2, 1e-3), - "torch_ref_max_bad": 0, - "description": "E32-style FC1 N4096/K4096 full-tactic coverage", - }, -} - - -def _assert_close_with_error_stats( - actual: torch.Tensor, - expected: torch.Tensor, - *, - label: str, - rtol: float, - atol: float, - max_bad: int = 0, - print_stats: bool = False, -) -> None: - actual_f = actual.to(torch.float32) - expected_f = expected.to(torch.float32) - abs_error = (actual_f - expected_f).abs() - allowed = atol + rtol * expected_f.abs() - bad = abs_error > allowed - bad_count = int(bad.sum().item()) - flat = abs_error.flatten() - stats = { - "max_abs": float(flat.max().item()) if flat.numel() else 0.0, - "mean_abs": float(flat.mean().item()) if flat.numel() else 0.0, - "p95_abs": float(torch.quantile(flat, 0.95).item()) if flat.numel() else 0.0, - "p99_abs": float(torch.quantile(flat, 0.99).item()) if flat.numel() else 0.0, - "bad_count": bad_count, - "total": flat.numel(), - } - if print_stats: - print( - f"{label}: max_abs={stats['max_abs']:.6g} " - f"mean_abs={stats['mean_abs']:.6g} p95_abs={stats['p95_abs']:.6g} " - f"p99_abs={stats['p99_abs']:.6g} bad={bad_count}/{flat.numel()}" - ) - if bad_count > max_bad: - raise AssertionError( - f"{label} exceeded tolerance: bad={bad_count}/{flat.numel()} " - f"(max_bad={max_bad}), max_abs={stats['max_abs']:.6g}, " - f"mean_abs={stats['mean_abs']:.6g}, p95_abs={stats['p95_abs']:.6g}, " - f"p99_abs={stats['p99_abs']:.6g}, rtol={rtol}, atol={atol}" - ) - - def _compute_with_active_experts( active_experts, x, @@ -3143,232 +2937,13 @@ def test_moe_bf16_mxfp4_hopper_activations( ) -@pytest.mark.skipif( - not is_sm90a_supported(torch.device("cuda")), - reason="FP8xMXFP4 pre-MMA scale MoE (Hopper mixed-input) requires SM90", -) -@pytest.mark.parametrize( - "case_name,case", - list(PHASE3_HUMMING_E2E_CASES.items()), - ids=list(PHASE3_HUMMING_E2E_CASES.keys()), -) -@pytest.mark.parametrize("use_autotune", [False, True]) -def test_moe_fp8_mxfp4_humming_prescale_hopper_correctness( - case_name, case, use_autotune -): - torch.manual_seed(case["seed"]) - device = torch.device("cuda") - e, m, n, k, top_k = ( - case["e"], - case["m"], - case["n"], - case["k"], - case["top_k"], - ) - output_dtype = torch.bfloat16 - - x_fp32 = torch.randn(m, k, dtype=torch.float32, device=device) * 0.05 - x = x_fp32.to(output_dtype) - w1 = torch.randint(0, 256, (e, 2 * n, k // 2), device=device, dtype=torch.uint8) - w2 = torch.randint(0, 256, (e, k, n // 2), device=device, dtype=torch.uint8) - - # Humming-style preprocessing constrains the original E8M0 scale range and - # stores only a small exponent offset for the pre-MMA FP4->E4M3 conversion. - # The residual is supplied through the GEMM epilogue routed-token scale. - raw_scale_low, raw_scale_high = case["raw_scale"] - w1_raw_scale = _make_humming_e8m0_weight_scale( - (e, 2 * n, k // 32), device, low=raw_scale_low, high=raw_scale_high - ) - w2_raw_scale = _make_humming_e8m0_weight_scale( - (e, k, n // 32), device, low=raw_scale_low, high=raw_scale_high - ) - w1_processed, w1_exp_offset, w1_residual = ( - fused_moe.preprocess_moe_weights_for_sm90_mixed_gemm_humming( - w1, w1_raw_scale, interleave=False - ) - ) - w2_processed, w2_exp_offset, w2_residual = ( - fused_moe.preprocess_moe_weights_for_sm90_mixed_gemm_humming( - w2, w2_raw_scale, interleave=False - ) - ) - if case_name == "small": - w1_api_il, w1_api_scale_il, w1_api_residual = ( - fused_moe.preprocess_moe_weights_for_sm90_mixed_gemm_humming( - w1, w1_raw_scale - ) - ) - w2_api_il, w2_api_scale_il, w2_api_residual = ( - fused_moe.preprocess_moe_weights_for_sm90_mixed_gemm_humming( - w2, w2_raw_scale - ) - ) - w1_ref_offset, w1_ref_residual, w1_delta_scale_offsets = ( - _reference_humming_e8m0_weight_scale(w1_raw_scale) - ) - w2_ref_offset, w2_ref_residual, w2_delta_scale_offsets = ( - _reference_humming_e8m0_weight_scale(w2_raw_scale) - ) - torch.testing.assert_close(w1_exp_offset, w1_ref_offset) - torch.testing.assert_close(w2_exp_offset, w2_ref_offset) - torch.testing.assert_close(w1_residual, w1_ref_residual) - torch.testing.assert_close(w2_residual, w2_ref_residual) - expected_max_offset = min(raw_scale_high - raw_scale_low, 12) - assert 1 <= int(w1_exp_offset.min().item()) <= expected_max_offset - assert 1 <= int(w2_exp_offset.min().item()) <= expected_max_offset - assert int(w1_exp_offset.max().item()) == expected_max_offset - assert int(w2_exp_offset.max().item()) == expected_max_offset - clamp_probe = _make_humming_e8m0_weight_scale((1, 64, 16), device) - clamp_offset, clamp_residual, clamp_delta = _reference_humming_e8m0_weight_scale( - clamp_probe - ) - assert int(clamp_offset.max().item()) == 12 - assert int(clamp_delta.max().item()) == 2 - torch.testing.assert_close( - clamp_residual.cpu(), torch.tensor([2.0**-12], dtype=torch.float32) - ) - # Humming mode currently shares the FP8MXFP4 5-slot quant_scales contract; - # slot 2 is reserved for the future post-MMA activation scale and is not - # consumed by the runtime per-token FP8 activation quantization path. - fc2_act_global = torch.ones((), device=device, dtype=torch.float32) - - _assert_humming_payload_rewrite_bit_equal(w1, w1_delta_scale_offsets, w1_processed) - _assert_humming_payload_rewrite_bit_equal(w2, w2_delta_scale_offsets, w2_processed) - - w1_il = fused_moe.interleave_moe_weights_for_sm90_mixed_gemm( - w1_processed, "fp4_fp8" - ) - w2_il = fused_moe.interleave_moe_weights_for_sm90_mixed_gemm( - w2_processed, "fp4_fp8" - ) - w1_scale_il = fused_moe.interleave_moe_scales_for_sm90_mixed_gemm(w1_exp_offset) - w2_scale_il = fused_moe.interleave_moe_scales_for_sm90_mixed_gemm(w2_exp_offset) - if case_name == "small": - torch.testing.assert_close(w1_il, w1_api_il) - torch.testing.assert_close(w2_il, w2_api_il) - torch.testing.assert_close(w1_scale_il, w1_api_scale_il) - torch.testing.assert_close(w2_scale_il, w2_api_scale_il) - torch.testing.assert_close(w1_residual, w1_api_residual) - torch.testing.assert_close(w2_residual, w2_api_residual) - - router_logits = torch.randn(m, e, dtype=output_dtype, device=device) - routing_weights, selected_experts = compute_routing(router_logits, top_k) - # Humming keeps the FP4->FP8 exponent-bias compensation in the epilogue for - # this FP8 x MXFP4 path. Fold the derived residual and the known 2^6 factor - # into the routed-token scale inputs for both GEMMs. - humming_epilogue_compensation = 64.0 - fc1_residual_route_scale = ( - w1_residual[selected_experts.to(torch.long)] * humming_epilogue_compensation - ) - fc2_residual_route_scale = ( - w2_residual[selected_experts.to(torch.long)] * humming_epilogue_compensation - ) - - def make_expert_contiguous_token_scale(route_scale): - return torch.cat( - [route_scale[selected_experts == expert_id] for expert_id in range(e)] - ).contiguous() - - fc1_residual_token_scale = make_expert_contiguous_token_scale( - fc1_residual_route_scale - ) - fc2_residual_token_scale = make_expert_contiguous_token_scale( - fc2_residual_route_scale - ) - quant_scales = [ - w1_scale_il.view(torch.int32), - fc1_residual_token_scale, - fc2_act_global, - w2_scale_il.view(torch.int32), - fc2_residual_token_scale, - ] - - def run_flash(profile_ids=None): - flash_output = torch.zeros(m, k, device=device, dtype=output_dtype) - fused_moe.cutlass_fused_moe( - x, - selected_experts.to(torch.int32), - routing_weights, - w1_il, - w2_il, - output_dtype, - quant_scales=quant_scales, - use_w4_group_scaling=True, - use_wfp4afp8_humming=True, - output=flash_output, - profile_ids=profile_ids, - ) - return flash_output - - w1_ref = _dequant_mxfp4_humming_prescale_on_device(w1_processed, w1_exp_offset) - w2_ref = _dequant_mxfp4_humming_prescale_on_device(w2_processed, w2_exp_offset) - x_ref = x.to(torch.float32) - ref_output = torch.zeros(m, k, dtype=torch.float32, device=device) - print_ref_stats = os.environ.get("FLASHINFER_PRINT_PHASE3_REF_STATS", "0") == "1" - for expert_id in range(e): - mask = selected_experts == expert_id - if not mask.any(): - continue - batch_idx, nth_expert = torch.where(mask) - w3_expert, w1_expert = torch.chunk(w1_ref[expert_id], 2, dim=0) - x_rows = x_ref[batch_idx] - fc1_amax = x_rows.abs().amax(dim=1) - fc1_quant = torch.where( - fc1_amax > 0, - torch.full_like(fc1_amax, 448.0) / fc1_amax, - torch.ones_like(fc1_amax), - ) - x_fp8_tensor = (x_rows * fc1_quant[:, None]).to(torch.float8_e4m3fn) - x_fp8 = x_fp8_tensor.to(torch.float32) - route_fc1_scale = (1.0 / fc1_quant) * fc1_residual_route_scale[ - batch_idx, nth_expert - ] - # FC1 token scale is applied in the GEMM epilogue before activation, so - # both gated branches must be scaled before the SiLU/product. - fc1_w1 = (x_fp8 @ w1_expert.t()) * route_fc1_scale[:, None] - fc1_w3 = (x_fp8 @ w3_expert.t()) * route_fc1_scale[:, None] - fc1 = F.silu(fc1_w1) * fc1_w3 - - fc2_amax = fc1.abs().amax(dim=1) - fc2_quant = torch.where( - fc2_amax > 0, - torch.full_like(fc2_amax, 448.0) / fc2_amax, - torch.ones_like(fc2_amax), - ) - fc1_fp8_tensor = (fc1 * fc2_quant[:, None]).to(torch.float8_e4m3fn) - fc1_fp8 = fc1_fp8_tensor.to(torch.float32) - route_fc2_scale = (1.0 / fc2_quant) * fc2_residual_route_scale[ - batch_idx, nth_expert - ] - fc2 = (fc1_fp8 @ w2_ref[expert_id].t()) * route_fc2_scale[:, None] - ref_output[batch_idx] += routing_weights[batch_idx, nth_expert, None] * fc2 - - torch_ref_rtol, torch_ref_atol = case["torch_ref_tolerance"] - ref_output = ref_output.to(output_dtype) - - def assert_flash_output(profile_label, flash_output): - _assert_close_with_error_stats( - flash_output, - ref_output, - label=f"{case_name}/{profile_label}: FlashInfer vs PyTorch reference", - rtol=torch_ref_rtol, - atol=torch_ref_atol, - max_bad=case["torch_ref_max_bad"], - print_stats=print_ref_stats, - ) - - profile_label = "autotune" if use_autotune else "default" - with autotune(True) if use_autotune else nullcontext(): - assert_flash_output(profile_label, run_flash()) - - # W4A8 Hopper interleaved path. # # Strict-tolerance envelope: h == intermediate_size == 512 with e == 2 only. -# Larger shapes exceed assert_close(rtol=1e-2, atol=2e-1) because of FP8 + INT4 -# accumulation noise. The slightly wider absolute tolerance keeps sparse -# Hopper/toolchain-dependent mismatches inside the test envelope while still -# rejecting broad correctness failures. +# Larger shapes exceed assert_close(rtol=1e-2, atol=1e-1) because of FP8 + INT4 +# accumulation noise — the upstream ``test_moe_w4a8`` above stays inside the +# same envelope for the same reason (verified on H200: e=2/h=2048 and +# e=8/h=512 both fail against a float32 PyTorch reference). W4A8_CORRECTNESS_CONFIGS = [ (1, 512, 2, 2, 512), (4, 512, 2, 2, 512), @@ -3406,10 +2981,10 @@ def _run_w4a8_moe_hopper( w3_scale = ( torch.randn(e, n, k // group_size, dtype=dtype, device=device) * affine_coeff ) - w1_pre_quant_scale = torch.rand(k, dtype=dtype, device=device) * 0.1 + 0.95 - w2_pre_quant_scale = torch.rand(n, dtype=dtype, device=device) * 0.1 + 0.95 - w3_pre_quant_scale = torch.rand(k, dtype=dtype, device=device) * 0.1 + 0.95 - input_scale = torch.rand(1, dtype=torch.float32, device=device) * 0.2 + 0.1 + w1_pre_quant_scale = torch.rand(e, k, dtype=dtype, device=device) * 0.1 + 0.95 + w2_pre_quant_scale = torch.rand(e, n, dtype=dtype, device=device) * 0.1 + 0.95 + w3_pre_quant_scale = torch.rand(e, k, dtype=dtype, device=device) * 0.1 + 0.95 + input_scale = torch.rand(e, 1, dtype=torch.float32, device=device) * 0.2 + 0.1 weight_scale_2 = torch.ones(e, 1, dtype=torch.float32, device=device) fc1_weights = torch.cat([w3_weight, w1_weight], dim=1) @@ -3421,20 +2996,27 @@ def _run_w4a8_moe_hopper( fc2_weights.contiguous().view(torch.uint8), "int4" ) + def _interleave_scales(w, dim): + factor = 4 if dim % 512 == 0 else (2 if dim % 256 == 0 else 1) + s = w.shape + return ( + w.reshape(s[0], s[1], s[2] // factor, factor) + .permute(0, 2, 1, 3) + .reshape(s[0], s[2] // factor, s[1] * factor) + .contiguous() + ) + + w3_w1_scales_int = _interleave_scales(torch.cat([w3_scale, w1_scale], dim=1), k) + w2_scales_int = _interleave_scales(w2_scale, n) # Weight scales: bf16 bit-pattern trick; act scales stay in native dtype. - w3_w1_scales = torch.cat([w3_scale, w1_scale], dim=1) - w3_w1_scales_out = fused_moe.interleave_moe_scales_for_sm90_mixed_gemm( - w3_w1_scales.to(torch.bfloat16).view(dtype), group_size - ) - w2_scales_out = fused_moe.interleave_moe_scales_for_sm90_mixed_gemm( - w2_scale.to(torch.bfloat16).view(dtype), group_size - ) + w3_w1_scales_out = w3_w1_scales_int.to(torch.bfloat16).view(dtype) + w2_scales_out = w2_scales_int.to(torch.bfloat16).view(dtype) w3_w1_input_scale_max = input_scale.max() fc31_act_scale = ( torch.max(w1_pre_quant_scale, w3_pre_quant_scale) / w3_w1_input_scale_max ).to(dtype) - fc2_act_scale = (w2_pre_quant_scale / input_scale).to(dtype) + fc2_act_scale = (w2_pre_quant_scale / input_scale).to(dtype).unsqueeze(-1) fc31_alpha = (weight_scale_2.squeeze(-1) * w3_w1_input_scale_max).float() fc2_alpha = (weight_scale_2.squeeze(-1) * input_scale.squeeze(-1)).float() zero_1 = torch.empty(0, dtype=dtype, device=device) @@ -3501,11 +3083,9 @@ def _run_w4a8_moe_hopper( w31_list.append(torch.cat([w3_dq, w1_dq], dim=0)) w2_list.append(w2_dq) - input_scale_for_ref = torch.full( - (num_experts,), - input_scale.item(), - dtype=input_scale.dtype, - device=input_scale.device, + # Broadcast max over experts; see comment on fc31_act_scale above. + fc1_input_scale_for_ref = torch.full_like( + input_scale.squeeze(-1), w3_w1_input_scale_max.item() ) ref_output = torch_moe_w4a8( num_experts, @@ -3514,14 +3094,14 @@ def _run_w4a8_moe_hopper( torch.stack(w2_list, dim=0), selected_experts, routing_weights, - fc1_input_scale=input_scale_for_ref, - fc2_input_scale=input_scale_for_ref, + fc1_input_scale=fc1_input_scale_for_ref, + fc2_input_scale=input_scale.squeeze(-1), fc1_pre_quant_scale=torch.max(w1_pre_quant_scale, w3_pre_quant_scale), fc2_pre_quant_scale=w2_pre_quant_scale, fc1_weight_scale_2=weight_scale_2.squeeze(-1), fc2_weight_scale_2=weight_scale_2.squeeze(-1), ) - torch.testing.assert_close(ref_output, flash_output, rtol=1e-2, atol=2e-1) + torch.testing.assert_close(ref_output, flash_output, rtol=1e-2, atol=1e-1) @pytest.mark.skipif( diff --git a/version.txt b/version.txt index 194970a6d91..9b79303dbe4 100644 --- a/version.txt +++ b/version.txt @@ -1 +1 @@ -0.6.17rc4 +0.6.17rc5