[TRTLLM-15316][feat] sm107 gemm + quant

cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt

@@ -18,6 +18,9 @@
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 set(CMAKE_CUDA_RESOLVE_DEVICE_SYMBOLS ON)
 
+# Common CUDA architectures for FP8/FP16/BF16 GEMM kernels
+set(CUTLASS_COMMON_ARCHS 89 90 100f 120f)
+
 # The Python executable will only be defined if building with Torch support. If
 # not, we need to find it here.
 if(NOT Python3_EXECUTABLE)
@@ -76,15 +79,17 @@ function(process_target target_name enable_hopper enable_blackwell)
   if(${enable_blackwell}
      AND ("100" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG
           OR "103" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG
+          OR "107" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG
           OR "120" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG
           OR "121" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG
          ))
 
     target_compile_options(${target_name}
                            PRIVATE "-DCUTLASS_ENABLE_GDC_FOR_SM100=1")
-    # Both 100 and 103 support these kernels
+    # SM100 family (100, 103, 107) support these kernels
     if("100" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG
-       OR "103" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG)
+       OR "103" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG
+       OR "107" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG)
       # No kernels should be parsed, unless blackwell is specified. This is a
       # build time improvement
       target_compile_definitions(${target_name}
@@ -205,12 +210,12 @@ add_cuda_architectures(fpA_intB_gemm_src 89)
 add_instantiations(fpA_intB_gemm_src ${INSTANTIATION_GENERATION_DIR}/gemm)
 
 add_library(fb_gemm_src STATIC ${FBGEMM_SRC_CU} ${FBGEMM_CU_INSTANTIATIONS})
-set_cuda_architectures(fb_gemm_src 89 90 100f 120f)
+set_cuda_architectures(fb_gemm_src ${CUTLASS_COMMON_ARCHS})
 # add_instantiations(fb_gemm_src
 # ${INSTANTIATION_GENERATION_DIR}/fp8_rowwise_gemm)
 
 add_library(fp8_blockscale_gemm_src STATIC ${FP8_BLOCKSCALE_GEMM_SRC_CU})
-set_cuda_architectures(fp8_blockscale_gemm_src 89 90 100f 120f)
+set_cuda_architectures(fp8_blockscale_gemm_src ${CUTLASS_COMMON_ARCHS})
 
 set(GEMM_SWIGLU_SM90_SRC_CU
     ${CMAKE_CURRENT_SOURCE_DIR}/fused_gated_gemm/gemm_swiglu_e4m3.cu)
@@ -264,7 +269,7 @@ if(USING_OSS_CUTLASS_MOE_GEMM)
   process_target(_moe_gemm_fp4 false true)
 
   add_library(_moe_gemm_fp8 OBJECT ${MOE_GEMM_SRC_CU_FP8})
-  set_cuda_architectures(_moe_gemm_fp8 89 90 100f 120f)
+  set_cuda_architectures(_moe_gemm_fp8 ${CUTLASS_COMMON_ARCHS})
   process_target(_moe_gemm_fp8 true true)
 
   add_instantiations(moe_gemm_src ${INSTANTIATION_GENERATION_DIR}/gemm_grouped)

cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp

@@ -382,10 +382,19 @@ std::vector<CutlassGemmConfig> get_candidate_configs_sm100_dynamic_cluster_shape
     std::vector<CutlassGemmConfig> candidate_configs;
     if ((config & CutlassGemmConfig::FP4_ONLY) != 0)
     {
+        // FP4 block-scaled types only support the TMA epilogue schedule on SM107.
+        // SM107 uses the shared tile set below; the SM100-only tiles are not enabled for it.
+        if (sm == 107 && schedule != EpilogueScheduleType::TMA)
+        {
+            return {};
+        }
+
         if (sm == 100)
         {
+            // FP4 block-scaled types only support TMA epilogue schedule
             if (schedule != EpilogueScheduleType::TMA)
                 return {};
+
             candidate_configs.push_back(CutlassGemmConfig{CutlassTileConfigSM100::CtaShape128x64x128B,
                 MainloopScheduleType::AUTO, schedule, cluster1sm, dynamic_cluster_shape, fallback_cluster_shape, sm});
             if (supports_2sm)
@@ -499,6 +508,11 @@ std::vector<CutlassGemmConfig> get_candidate_configs_sm100(
             ClusterShape::Undefined, sm},
     };
 #else
+    if (tensorrt_llm::common::isSM100Family(sm) && sm != 103 && sm != 107)
+    {
+        TLLM_LOG_INFO("Reassigned sm version to 100 for unknown sm version belonging to SM100 family");
+        sm = 100;
+    }
     if (config & CutlassGemmConfig::GROUPED_GEMM)
     {
         std::vector<CutlassGemmConfig> candidate_configs;

cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_preprocessors.cpp

@@ -134,7 +134,7 @@ LayoutDetails getLayoutDetailsForTransform(QuantType quant_type, int arch)
     {
         return getLayoutDetailsForArch<cutlass::arch::Sm90>(quant_type);
     }
-    else if (arch == 100)
+    else if (isSM100Family(arch) && arch != 103)
     {
         return getLayoutDetailsForArch<cutlass::arch::Sm100>(quant_type);
     }
@@ -619,7 +619,7 @@ void preprocess_weights_for_mixed_gemm(int8_t* preprocessed_quantized_weight, in
         src_buf.swap(dst_buf);
     }
 
-    if (arch != 100 && arch != 103)
+    if (!isSM100Family(arch))
     {
         TLLM_LOG_INFO("add_bias_and_interleave_quantized_tensor_inplace");
         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

@@ -436,7 +436,7 @@ size_t CutlassFp4GemmRunner<T, fp4GemmType>::dispatchToArch(T* D, void const* A,
 {
     if constexpr (fp4GemmType == FP4GemmType::W4A8_MXFP4_MXFP8)
     {
-        if (mSm == 100 || mSm == 103)
+        if (tk::isSM100Family(mSm))
         {
             return dispatchMXFP8xMXFP4GemmCTAShapeSm100<T>(D, A, B, input_sf, weight_sf, global_sf, m, n, k,
                 batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias);
@@ -449,7 +449,7 @@ size_t CutlassFp4GemmRunner<T, fp4GemmType>::dispatchToArch(T* D, void const* A,
     }
     else if constexpr (fp4GemmType == FP4GemmType::W8A8_MXFP8_MXFP8)
     {
-        if (mSm == 100 || mSm == 103)
+        if (tk::isSM100Family(mSm))
         {
             return dispatchMXFP8xMXFP8GemmCTAShapeSm100<T>(D, A, B, input_sf, weight_sf, global_sf, m, n, k,
                 batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy);
@@ -462,21 +462,26 @@ size_t CutlassFp4GemmRunner<T, fp4GemmType>::dispatchToArch(T* D, void const* A,
     }
     else if constexpr (fp4GemmType == FP4GemmType::W4A4_NVFP4_NVFP4)
     {
-        if (mSm == 103)
+        if (tk::isSM100Family(mSm))
         {
 #ifdef COMPILE_BLACKWELL_SM103_TMA_GEMMS
-            return dispatchNVFP4xNVFP4GemmCTAShapeSm10x<cutlass::arch::Sm103, T>(D, A, B, input_sf, weight_sf,
-                global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias);
+            if (mSm == 103)
+            {
+                return dispatchNVFP4xNVFP4GemmCTAShapeSm10x<cutlass::arch::Sm103, T>(D, A, B, input_sf, weight_sf,
+                    global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias);
+            }
+            else
+            {
+                return dispatchNVFP4xNVFP4GemmCTAShapeSm10x<cutlass::arch::Sm100, T>(D, A, B, input_sf, weight_sf,
+                    global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias);
+            }
 #else
+            // SM107, SM100, and other SM100 family members all use the same cutlass::arch::Sm100 kernels (compiled with
+            // 100f)
             return dispatchNVFP4xNVFP4GemmCTAShapeSm10x<cutlass::arch::Sm100, T>(D, A, B, input_sf, weight_sf,
                 global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias);
 #endif
         }
-        else if (mSm == 100)
-        {
-            return dispatchNVFP4xNVFP4GemmCTAShapeSm10x<cutlass::arch::Sm100, T>(D, A, B, input_sf, weight_sf,
-                global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias);
-        }
         else if (mSm == 120 || mSm == 121)
         {
             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
 
     std::vector<CutlassGemmConfig> candidateConfigs;
 
-    if (mSm == 100 || mSm == 103)
+    if (tk::isSM100Family(mSm))
     {
         std::vector<tkc::CutlassTileConfigSM100> tilesSm10x = {
             tkc::CutlassTileConfigSM100::CtaShape128x128x256B,

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

@@ -687,7 +687,7 @@ size_t CutlassFp8RowwiseGemmRunner<T>::dispatchToArch(void* D, void const* A, vo
         return dispatchGemmToCutlassSm90<T>(D, A, B, C_bias, quantOption, m, n, k, scale_d0, scale_d1, gemmConfig,
             workspace, workspaceBytes, stream, occupancy);
     }
-    else if (mSm == 100 || mSm == 103)
+    else if (tk::isSM100Family(mSm))
     {
         return dispatchGemmToCutlassSm100<T>(D, A, B, C_bias, quantOption, m, n, k, scale_d0, scale_d1, gemmConfig,
             workspace, workspaceBytes, stream, occupancy);
@@ -759,7 +759,7 @@ std::vector<tkc::CutlassGemmConfig> CutlassFp8RowwiseGemmRunner<T>::getConfigs()
             }
         }
     }
-    else if (mSm == 100 || mSm == 103)
+    else if (tk::isSM100Family(mSm))
     {
         std::vector<tkc::CutlassTileConfigSM100> tilesSm100 = {
             tkc::CutlassTileConfigSM100::CtaShape64x32x128B,

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

@@ -465,7 +465,7 @@ void CutlassFpAIntBGemmRunner<ActivationType, WeightType, QuantOp, ScaleZeroType
             "arch to build_wheel.py.");
 #endif // COMPILE_HOPPER_TMA_GEMMS
     }
-    else if (sm_ == 100 || sm_ == 103)
+    else if (tk::isSM100Family(sm_))
     {
 #ifdef COMPILE_BLACKWELL_TMA_GEMMS
         cutlass_kernels_oss::sm100_dispatch_gemm_to_cutlass<ActivationType, WeightType, ScaleZeroType, BiasType,

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

@@ -1005,7 +1005,8 @@ def has_arch(sm):
     operations = []
     operations += generate_sm120_operations(has_arch(120) or has_arch(121))
     operations += generate_sm103_operations(has_arch(103))
-    operations += generate_sm100_operations(has_arch(100) or has_arch(103))
+    operations += generate_sm100_operations(
+        any(has_arch(sm) for sm in range(100, 110)))
     operations += generate_sm90_operations(has_arch(90))
     operations += generate_sm80_operations(has_arch(80) or has_arch(89))
 

cpp/tensorrt_llm/kernels/quantization.cu

@@ -189,12 +189,13 @@ void invokeFP4Quantization(int b, int m, int n, T const* input, float const* SFS
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 // MXFP8 Quantization
 
-template <typename T>
+template <typename T, int SF_VEC_SIZE, int SF_OUTPUT_VEC_SIZE>
 void invokeMxFP8Quantization(int b, int m, int n, int padded_n, T const* input, int64_t* output, int32_t* SFOuput,
     QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream)
 {
-    // Fixed SF_VEC_SIZE as 32
-    static constexpr int SF_VEC_SIZE = 32;
+    static_assert(SF_VEC_SIZE == 32 || SF_VEC_SIZE == 128, "MXFP8 quantization supports SF vector sizes 32 and 128.");
+    static_assert(SF_OUTPUT_VEC_SIZE == 32 || SF_OUTPUT_VEC_SIZE == SF_VEC_SIZE,
+        "MXFP8 output SF vector size must be 32 or match the quantization SF vector size.");
 
     // Grid, Block size.
     // Each thread converts 8 values.
@@ -217,8 +218,9 @@ void invokeMxFP8Quantization(int b, int m, int n, int padded_n, T const* input,
     config.numAttrs = 1;
     config.attrs = attrs;
     cudaLaunchKernelEx(&config,
-        quantize_with_block_size<BlockScaleQuantizationType::FP16_TO_MXFP8, T, SF_VEC_SIZE, true>, b, m, n, padded_n,
-        input, nullptr, reinterpret_cast<uint32_t*>(output), reinterpret_cast<uint32_t*>(SFOuput), layout);
+        quantize_with_block_size<BlockScaleQuantizationType::FP16_TO_MXFP8, T, SF_VEC_SIZE, true, SF_OUTPUT_VEC_SIZE>,
+        b, m, n, padded_n, input, nullptr, reinterpret_cast<uint32_t*>(output), reinterpret_cast<uint32_t*>(SFOuput),
+        layout);
 }
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -402,8 +404,8 @@ template void invokeFP4Quantization<half, 16>(int b, int m, int n, half const* i
 template void invokeFP4Quantization<half, 32>(int b, int m, int n, half const* input, float const* SFScale,
     int64_t* output, int32_t* SFOuput, bool useUE8M0, QuantizationSFLayout layout, int multiProcessorCount,
     cudaStream_t stream);
-template void invokeMxFP8Quantization<half>(int b, int m, int n, int padded_n, half const* input, int64_t* output,
-    int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream);
+template void invokeMxFP8Quantization<half, 32, 32>(int b, int m, int n, int padded_n, half const* input,
+    int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream);
 template void computePerTokenGlobalScaleForFP4Quantization<half>(int b, int m, int n, half const* input,
     int const* tokensPerBatch, float* globalScale, int multiProcessorCount, cudaStream_t stream);
 #ifdef ENABLE_BF16
@@ -413,8 +415,12 @@ template void invokeFP4Quantization<__nv_bfloat16, 16>(int b, int m, int n, __nv
 template void invokeFP4Quantization<__nv_bfloat16, 32>(int b, int m, int n, __nv_bfloat16 const* input,
     float const* SFScale, int64_t* output, int32_t* SFOuput, bool useUE8M0, QuantizationSFLayout layout,
     int multiProcessorCount, cudaStream_t stream);
-template void invokeMxFP8Quantization<__nv_bfloat16>(int b, int m, int n, int padded_n, __nv_bfloat16 const* input,
-    int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream);
+template void invokeMxFP8Quantization<__nv_bfloat16, 32, 32>(int b, int m, int n, int padded_n,
+    __nv_bfloat16 const* input, int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount,
+    cudaStream_t stream);
+template void invokeMxFP8Quantization<__nv_bfloat16, 128, 32>(int b, int m, int n, int padded_n,
+    __nv_bfloat16 const* input, int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount,
+    cudaStream_t stream);
 template void computePerTokenGlobalScaleForFP4Quantization<__nv_bfloat16>(int b, int m, int n,
     __nv_bfloat16 const* input, int const* tokensPerBatch, float* globalScale, int multiProcessorCount,
     cudaStream_t stream);