Skip to content

Commit c44de1a

Browse files
farazkh80yifeizhang-czhangclpeaceh-nvmingyangHao
authored
[TRTLLM-15316][feat] sm107 gemm + quant (#17485)
Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Co-authored-by: Yifei Zhang <yifezhang@nvidia.com> Co-authored-by: Chulian Zhang <chulianz@nvidia.com> Co-authored-by: Peace He <peaceh@nvidia.com> Co-authored-by: Mingyang Hao <mingyangh@nvidia.com> Co-authored-by: Bowen Fu <BowenFu@users.noreply.github.com>
1 parent cb511b5 commit c44de1a

17 files changed

Lines changed: 340 additions & 203 deletions

File tree

‎cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,9 @@
1818
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
1919
set(CMAKE_CUDA_RESOLVE_DEVICE_SYMBOLS ON)
2020

21+
# Common CUDA architectures for FP8/FP16/BF16 GEMM kernels
22+
set(CUTLASS_COMMON_ARCHS 89 90 100f 120f)
23+
2124
# The Python executable will only be defined if building with Torch support. If
2225
# not, we need to find it here.
2326
if(NOT Python3_EXECUTABLE)
@@ -76,15 +79,17 @@ function(process_target target_name enable_hopper enable_blackwell)
7679
if(${enable_blackwell}
7780
AND ("100" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG
7881
OR "103" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG
82+
OR "107" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG
7983
OR "120" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG
8084
OR "121" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG
8185
))
8286

8387
target_compile_options(${target_name}
8488
PRIVATE "-DCUTLASS_ENABLE_GDC_FOR_SM100=1")
85-
# Both 100 and 103 support these kernels
89+
# SM100 family (100, 103, 107) support these kernels
8690
if("100" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG
87-
OR "103" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG)
91+
OR "103" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG
92+
OR "107" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG)
8893
# No kernels should be parsed, unless blackwell is specified. This is a
8994
# build time improvement
9095
target_compile_definitions(${target_name}
@@ -205,12 +210,12 @@ add_cuda_architectures(fpA_intB_gemm_src 89)
205210
add_instantiations(fpA_intB_gemm_src ${INSTANTIATION_GENERATION_DIR}/gemm)
206211

207212
add_library(fb_gemm_src STATIC ${FBGEMM_SRC_CU} ${FBGEMM_CU_INSTANTIATIONS})
208-
set_cuda_architectures(fb_gemm_src 89 90 100f 120f)
213+
set_cuda_architectures(fb_gemm_src ${CUTLASS_COMMON_ARCHS})
209214
# add_instantiations(fb_gemm_src
210215
# ${INSTANTIATION_GENERATION_DIR}/fp8_rowwise_gemm)
211216

212217
add_library(fp8_blockscale_gemm_src STATIC ${FP8_BLOCKSCALE_GEMM_SRC_CU})
213-
set_cuda_architectures(fp8_blockscale_gemm_src 89 90 100f 120f)
218+
set_cuda_architectures(fp8_blockscale_gemm_src ${CUTLASS_COMMON_ARCHS})
214219

215220
set(GEMM_SWIGLU_SM90_SRC_CU
216221
${CMAKE_CURRENT_SOURCE_DIR}/fused_gated_gemm/gemm_swiglu_e4m3.cu)
@@ -264,7 +269,7 @@ if(USING_OSS_CUTLASS_MOE_GEMM)
264269
process_target(_moe_gemm_fp4 false true)
265270

266271
add_library(_moe_gemm_fp8 OBJECT ${MOE_GEMM_SRC_CU_FP8})
267-
set_cuda_architectures(_moe_gemm_fp8 89 90 100f 120f)
272+
set_cuda_architectures(_moe_gemm_fp8 ${CUTLASS_COMMON_ARCHS})
268273
process_target(_moe_gemm_fp8 true true)
269274

270275
add_instantiations(moe_gemm_src ${INSTANTIATION_GENERATION_DIR}/gemm_grouped)

‎cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -382,10 +382,19 @@ std::vector<CutlassGemmConfig> get_candidate_configs_sm100_dynamic_cluster_shape
382382
std::vector<CutlassGemmConfig> candidate_configs;
383383
if ((config & CutlassGemmConfig::FP4_ONLY) != 0)
384384
{
385+
// FP4 block-scaled types only support the TMA epilogue schedule on SM107.
386+
// SM107 uses the shared tile set below; the SM100-only tiles are not enabled for it.
387+
if (sm == 107 && schedule != EpilogueScheduleType::TMA)
388+
{
389+
return {};
390+
}
391+
385392
if (sm == 100)
386393
{
394+
// FP4 block-scaled types only support TMA epilogue schedule
387395
if (schedule != EpilogueScheduleType::TMA)
388396
return {};
397+
389398
candidate_configs.push_back(CutlassGemmConfig{CutlassTileConfigSM100::CtaShape128x64x128B,
390399
MainloopScheduleType::AUTO, schedule, cluster1sm, dynamic_cluster_shape, fallback_cluster_shape, sm});
391400
if (supports_2sm)
@@ -499,6 +508,11 @@ std::vector<CutlassGemmConfig> get_candidate_configs_sm100(
499508
ClusterShape::Undefined, sm},
500509
};
501510
#else
511+
if (tensorrt_llm::common::isSM100Family(sm) && sm != 103 && sm != 107)
512+
{
513+
TLLM_LOG_INFO("Reassigned sm version to 100 for unknown sm version belonging to SM100 family");
514+
sm = 100;
515+
}
502516
if (config & CutlassGemmConfig::GROUPED_GEMM)
503517
{
504518
std::vector<CutlassGemmConfig> candidate_configs;

‎cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_preprocessors.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,7 @@ LayoutDetails getLayoutDetailsForTransform(QuantType quant_type, int arch)
134134
{
135135
return getLayoutDetailsForArch<cutlass::arch::Sm90>(quant_type);
136136
}
137-
else if (arch == 100)
137+
else if (isSM100Family(arch) && arch != 103)
138138
{
139139
return getLayoutDetailsForArch<cutlass::arch::Sm100>(quant_type);
140140
}
@@ -619,7 +619,7 @@ void preprocess_weights_for_mixed_gemm(int8_t* preprocessed_quantized_weight, in
619619
src_buf.swap(dst_buf);
620620
}
621621

622-
if (arch != 100 && arch != 103)
622+
if (!isSM100Family(arch))
623623
{
624624
TLLM_LOG_INFO("add_bias_and_interleave_quantized_tensor_inplace");
625625
add_bias_and_interleave_quantized_tensor_inplace(src_buf.data(), num_elts, quant_type);

‎cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_template.h‎

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -436,7 +436,7 @@ size_t CutlassFp4GemmRunner<T, fp4GemmType>::dispatchToArch(T* D, void const* A,
436436
{
437437
if constexpr (fp4GemmType == FP4GemmType::W4A8_MXFP4_MXFP8)
438438
{
439-
if (mSm == 100 || mSm == 103)
439+
if (tk::isSM100Family(mSm))
440440
{
441441
return dispatchMXFP8xMXFP4GemmCTAShapeSm100<T>(D, A, B, input_sf, weight_sf, global_sf, m, n, k,
442442
batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias);
@@ -449,7 +449,7 @@ size_t CutlassFp4GemmRunner<T, fp4GemmType>::dispatchToArch(T* D, void const* A,
449449
}
450450
else if constexpr (fp4GemmType == FP4GemmType::W8A8_MXFP8_MXFP8)
451451
{
452-
if (mSm == 100 || mSm == 103)
452+
if (tk::isSM100Family(mSm))
453453
{
454454
return dispatchMXFP8xMXFP8GemmCTAShapeSm100<T>(D, A, B, input_sf, weight_sf, global_sf, m, n, k,
455455
batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy);
@@ -462,21 +462,26 @@ size_t CutlassFp4GemmRunner<T, fp4GemmType>::dispatchToArch(T* D, void const* A,
462462
}
463463
else if constexpr (fp4GemmType == FP4GemmType::W4A4_NVFP4_NVFP4)
464464
{
465-
if (mSm == 103)
465+
if (tk::isSM100Family(mSm))
466466
{
467467
#ifdef COMPILE_BLACKWELL_SM103_TMA_GEMMS
468-
return dispatchNVFP4xNVFP4GemmCTAShapeSm10x<cutlass::arch::Sm103, T>(D, A, B, input_sf, weight_sf,
469-
global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias);
468+
if (mSm == 103)
469+
{
470+
return dispatchNVFP4xNVFP4GemmCTAShapeSm10x<cutlass::arch::Sm103, T>(D, A, B, input_sf, weight_sf,
471+
global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias);
472+
}
473+
else
474+
{
475+
return dispatchNVFP4xNVFP4GemmCTAShapeSm10x<cutlass::arch::Sm100, T>(D, A, B, input_sf, weight_sf,
476+
global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias);
477+
}
470478
#else
479+
// SM107, SM100, and other SM100 family members all use the same cutlass::arch::Sm100 kernels (compiled with
480+
// 100f)
471481
return dispatchNVFP4xNVFP4GemmCTAShapeSm10x<cutlass::arch::Sm100, T>(D, A, B, input_sf, weight_sf,
472482
global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias);
473483
#endif
474484
}
475-
else if (mSm == 100)
476-
{
477-
return dispatchNVFP4xNVFP4GemmCTAShapeSm10x<cutlass::arch::Sm100, T>(D, A, B, input_sf, weight_sf,
478-
global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias);
479-
}
480485
else if (mSm == 120 || mSm == 121)
481486
{
482487
return dispatchNVFP4xNVFP4GemmCTAShapeSm120<T>(D, A, B, input_sf, weight_sf, global_sf, m, n, k,
@@ -514,7 +519,7 @@ std::vector<tkc::CutlassGemmConfig> CutlassFp4GemmRunner<T, fp4GemmType>::getCon
514519

515520
std::vector<CutlassGemmConfig> candidateConfigs;
516521

517-
if (mSm == 100 || mSm == 103)
522+
if (tk::isSM100Family(mSm))
518523
{
519524
std::vector<tkc::CutlassTileConfigSM100> tilesSm10x = {
520525
tkc::CutlassTileConfigSM100::CtaShape128x128x256B,

‎cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_rowwise_gemm/fp8_rowwise_gemm_template.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -687,7 +687,7 @@ size_t CutlassFp8RowwiseGemmRunner<T>::dispatchToArch(void* D, void const* A, vo
687687
return dispatchGemmToCutlassSm90<T>(D, A, B, C_bias, quantOption, m, n, k, scale_d0, scale_d1, gemmConfig,
688688
workspace, workspaceBytes, stream, occupancy);
689689
}
690-
else if (mSm == 100 || mSm == 103)
690+
else if (tk::isSM100Family(mSm))
691691
{
692692
return dispatchGemmToCutlassSm100<T>(D, A, B, C_bias, quantOption, m, n, k, scale_d0, scale_d1, gemmConfig,
693693
workspace, workspaceBytes, stream, occupancy);
@@ -759,7 +759,7 @@ std::vector<tkc::CutlassGemmConfig> CutlassFp8RowwiseGemmRunner<T>::getConfigs()
759759
}
760760
}
761761
}
762-
else if (mSm == 100 || mSm == 103)
762+
else if (tk::isSM100Family(mSm))
763763
{
764764
std::vector<tkc::CutlassTileConfigSM100> tilesSm100 = {
765765
tkc::CutlassTileConfigSM100::CtaShape64x32x128B,

‎cpp/tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -465,7 +465,7 @@ void CutlassFpAIntBGemmRunner<ActivationType, WeightType, QuantOp, ScaleZeroType
465465
"arch to build_wheel.py.");
466466
#endif // COMPILE_HOPPER_TMA_GEMMS
467467
}
468-
else if (sm_ == 100 || sm_ == 103)
468+
else if (tk::isSM100Family(sm_))
469469
{
470470
#ifdef COMPILE_BLACKWELL_TMA_GEMMS
471471
cutlass_kernels_oss::sm100_dispatch_gemm_to_cutlass<ActivationType, WeightType, ScaleZeroType, BiasType,

‎cpp/tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1005,7 +1005,8 @@ def has_arch(sm):
10051005
operations = []
10061006
operations += generate_sm120_operations(has_arch(120) or has_arch(121))
10071007
operations += generate_sm103_operations(has_arch(103))
1008-
operations += generate_sm100_operations(has_arch(100) or has_arch(103))
1008+
operations += generate_sm100_operations(
1009+
any(has_arch(sm) for sm in range(100, 110)))
10091010
operations += generate_sm90_operations(has_arch(90))
10101011
operations += generate_sm80_operations(has_arch(80) or has_arch(89))
10111012

‎cpp/tensorrt_llm/kernels/quantization.cu‎

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -189,12 +189,13 @@ void invokeFP4Quantization(int b, int m, int n, T const* input, float const* SFS
189189
////////////////////////////////////////////////////////////////////////////////////////////////////
190190
// MXFP8 Quantization
191191

192-
template <typename T>
192+
template <typename T, int SF_VEC_SIZE, int SF_OUTPUT_VEC_SIZE>
193193
void invokeMxFP8Quantization(int b, int m, int n, int padded_n, T const* input, int64_t* output, int32_t* SFOuput,
194194
QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream)
195195
{
196-
// Fixed SF_VEC_SIZE as 32
197-
static constexpr int SF_VEC_SIZE = 32;
196+
static_assert(SF_VEC_SIZE == 32 || SF_VEC_SIZE == 128, "MXFP8 quantization supports SF vector sizes 32 and 128.");
197+
static_assert(SF_OUTPUT_VEC_SIZE == 32 || SF_OUTPUT_VEC_SIZE == SF_VEC_SIZE,
198+
"MXFP8 output SF vector size must be 32 or match the quantization SF vector size.");
198199

199200
// Grid, Block size.
200201
// Each thread converts 8 values.
@@ -217,8 +218,9 @@ void invokeMxFP8Quantization(int b, int m, int n, int padded_n, T const* input,
217218
config.numAttrs = 1;
218219
config.attrs = attrs;
219220
cudaLaunchKernelEx(&config,
220-
quantize_with_block_size<BlockScaleQuantizationType::FP16_TO_MXFP8, T, SF_VEC_SIZE, true>, b, m, n, padded_n,
221-
input, nullptr, reinterpret_cast<uint32_t*>(output), reinterpret_cast<uint32_t*>(SFOuput), layout);
221+
quantize_with_block_size<BlockScaleQuantizationType::FP16_TO_MXFP8, T, SF_VEC_SIZE, true, SF_OUTPUT_VEC_SIZE>,
222+
b, m, n, padded_n, input, nullptr, reinterpret_cast<uint32_t*>(output), reinterpret_cast<uint32_t*>(SFOuput),
223+
layout);
222224
}
223225

224226
////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -402,8 +404,8 @@ template void invokeFP4Quantization<half, 16>(int b, int m, int n, half const* i
402404
template void invokeFP4Quantization<half, 32>(int b, int m, int n, half const* input, float const* SFScale,
403405
int64_t* output, int32_t* SFOuput, bool useUE8M0, QuantizationSFLayout layout, int multiProcessorCount,
404406
cudaStream_t stream);
405-
template void invokeMxFP8Quantization<half>(int b, int m, int n, int padded_n, half const* input, int64_t* output,
406-
int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream);
407+
template void invokeMxFP8Quantization<half, 32, 32>(int b, int m, int n, int padded_n, half const* input,
408+
int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream);
407409
template void computePerTokenGlobalScaleForFP4Quantization<half>(int b, int m, int n, half const* input,
408410
int const* tokensPerBatch, float* globalScale, int multiProcessorCount, cudaStream_t stream);
409411
#ifdef ENABLE_BF16
@@ -413,8 +415,12 @@ template void invokeFP4Quantization<__nv_bfloat16, 16>(int b, int m, int n, __nv
413415
template void invokeFP4Quantization<__nv_bfloat16, 32>(int b, int m, int n, __nv_bfloat16 const* input,
414416
float const* SFScale, int64_t* output, int32_t* SFOuput, bool useUE8M0, QuantizationSFLayout layout,
415417
int multiProcessorCount, cudaStream_t stream);
416-
template void invokeMxFP8Quantization<__nv_bfloat16>(int b, int m, int n, int padded_n, __nv_bfloat16 const* input,
417-
int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream);
418+
template void invokeMxFP8Quantization<__nv_bfloat16, 32, 32>(int b, int m, int n, int padded_n,
419+
__nv_bfloat16 const* input, int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount,
420+
cudaStream_t stream);
421+
template void invokeMxFP8Quantization<__nv_bfloat16, 128, 32>(int b, int m, int n, int padded_n,
422+
__nv_bfloat16 const* input, int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount,
423+
cudaStream_t stream);
418424
template void computePerTokenGlobalScaleForFP4Quantization<__nv_bfloat16>(int b, int m, int n,
419425
__nv_bfloat16 const* input, int const* tokensPerBatch, float* globalScale, int multiProcessorCount,
420426
cudaStream_t stream);

0 commit comments

Comments
 (0)