SM90 native BF16/FP8 GEMM kernels, fused-SiLU quantization, and linear test harness

lmdeploy/turbomind/builders/ffn.py

@@ -20,6 +20,9 @@
     'fuse_w1w3',
 ]
 
+_SM90_FP8_FUSED_SILU_BLOCK = 128
+_SM90_BF16_FUSED_SILU_BLOCK = 64
+
 # ---------------------------------------------------------------------------
 # @transform_output_dim / @transform_input_dim helpers
 # ---------------------------------------------------------------------------
@@ -31,6 +34,24 @@ def _interleave_w1w3(w1: torch.Tensor, w3: torch.Tensor) -> torch.Tensor:
     return torch.stack([w1, w3], dim=-1).reshape(w1.shape[:-1] + (-1,)).contiguous()
 
 
+@transform_output_dim
+def _block_pack_w1w3(w1: torch.Tensor, w3: torch.Tensor, *,
+                     groups: int) -> torch.Tensor:
+    """Interleave gate/up by logical output groups.
+
+    The elements per group self-adapt for each Linear tensor kind. For FP8, one logical 128-wide weight group
+    corresponds to one scale element.
+    """
+    assert groups > 0, f'groups must be positive, got {groups}'
+    assert w1.shape[-1] % groups == 0, (
+        f'output dim {w1.shape[-1]} not divisible by groups {groups}')
+    assert w3.shape[-1] == w1.shape[-1], (
+        f'w1/w3 output dims differ: {w1.shape[-1]} != {w3.shape[-1]}')
+    w1g = w1.unflatten(-1, (groups, -1))
+    w3g = w3.unflatten(-1, (groups, -1))
+    return torch.stack([w1g, w3g], dim=-2).flatten(-3, -1).contiguous()
+
+
 @transform_output_dim
 def _chunk_w1w3(w1: torch.Tensor, w3: torch.Tensor, *,
                 tp: int) -> torch.Tensor:
@@ -49,61 +70,67 @@ def _chunk_w1w3(w1: torch.Tensor, w3: torch.Tensor, *,
 # ---------------------------------------------------------------------------
 
 
+def _is_sm90() -> bool:
+    """Whether the current CUDA device is Hopper SM90."""
+    return (torch.cuda.is_available()
+            and torch.cuda.get_device_capability() == (9, 0))
+
+
 def _should_fuse_silu(w1_linear: Linear, act_type: str, is_moe: bool = False) -> bool:
-    """Determine if fused SiLU (interleave) should be used for w1+w3 fusion.
+    """Determine if fused SiLU should be used for w1+w3 fusion.
 
     Gold standard condition (from GEMM kernel constraints — trust it):
 
-    act_type == SiLU && (int4 || mxfp4 || fp8 || moe) && !(fp8 && SM90)
+    act_type == SiLU && (int4 || mxfp4 || fp8 || moe)
+
+    Packing depends on format: FP8 uses [g128|u128|...] and SM90 BF16 uses
+    [g64|u64|...]; int4/mxfp4 use element interleave. Controllable via
+    config.fuse_silu after commit.
     """
     if act_type not in ('', 'silu', 'SiLU'):
         return False
 
-    # Dense bf16/fp16 without MoE -> chunk, not interleave
-    weight = w1_linear.tensors.get('weight')
-    is_quantized = weight is not None and weight.element_size() < 2
+    # SM90 BF16 dense uses the native fused-SiLU kernel. Other dense
+    # bf16/fp16 paths keep chunk layout and apply activation separately.
+    weight = w1_linear.tensors['weight']
+    is_quantized = weight.element_size() < 2
     if not is_quantized and not is_moe:
-        return False
-
-    # FP8 on SM90 -> chunk
-    fmt = w1_linear.weight_format
-    if fmt is not None and fmt.name == 'fp8':
-        if torch.cuda.is_available():
-            cap = torch.cuda.get_device_capability()
-            if cap == (9, 0):
-                return False
+        return weight.dtype == torch.bfloat16 and _is_sm90()
 
     # SM100+ grouped bf16 MoE: CublasGroupedKernel has no fused GatedSilu
-    if is_moe:
-        weight = w1_linear.tensors.get('weight')
-        if weight is not None and weight.dtype == torch.bfloat16:
-            if torch.cuda.is_available():
-                cap = torch.cuda.get_device_capability()
-                if cap >= (10, 0):
-                    return False
+    if (is_moe and weight.dtype == torch.bfloat16
+            and torch.cuda.is_available()
+            and torch.cuda.get_device_capability() >= (10, 0)):
+        return False
 
     return True
 
 
-def _can_fuse_w1w3(w1: Linear, tp: int) -> bool:
+def _fused_silu_block(w1: Linear) -> int | None:
+    """Return the gate/up block width required by the fused SiLU kernel."""
+    if w1.weight_format.name == 'fp8' and _is_sm90():
+        return _SM90_FP8_FUSED_SILU_BLOCK
+    if _is_sm90() and w1.tensors['weight'].dtype == torch.bfloat16:
+        return _SM90_BF16_FUSED_SILU_BLOCK
+    return None
+
+
+def _can_fuse_w1w3(w1: Linear, tp: int, *,
+                    pack_block: int | None = None) -> bool:
     """Check whether w1+w3 fusion is safe for the given TP.
 
-    Fusion (interleave or chunk) concatenates w1 and w3 along the output dim.
-    For block-quantized formats (e.g. FP8 with block_out=128), the fused
-    scale count ``2 * cdiv(N/tp, block_out)`` must equal
-    ``cdiv(2*N/tp, block_out)``.  This holds iff ``(N/tp) % block_out == 0``.
-    When it doesn't, the fused module's C++ allocation won't match the
-    concatenated scales and we must commit w1/w3 separately.
+    Fusion (interleave, block-pack, or chunk) concatenates w1 and w3 along the
+    output dim. For block-quantized formats, the local output must align to
+    ``block_out`` so concatenated scale counts match the fused allocation.
+    Native SM90 fused SiLU additionally requires format-specific gate/up groups.
+    Misaligned projections are committed separately.
     """
-    if tp <= 1:
-        return True
-    fmt = w1.weight_format
-    if fmt is None or fmt.block_out is None:
-        return True
-    w = w1.tensors.get('weight')
-    if w is None:
-        return True
-    return (w.size(-1) // tp) % fmt.block_out == 0
+    w = w1.tensors['weight']
+    if w.size(-1) % tp != 0:
+        return False
+    block_out = w1.weight_format.block_out or 1
+    required_block = math.lcm(block_out, pack_block or 1)
+    return (w.size(-1) // tp) % required_block == 0
 
 
 def fuse_w1w3(
@@ -124,10 +151,15 @@ def fuse_w1w3(
     block-scale alignment check in ``_can_fuse_w1w3``.
     """
     fused_silu = _should_fuse_silu(w1, act_type, is_moe)
-    can_fuse = _can_fuse_w1w3(w1, tp)
+    pack_block = _fused_silu_block(w1) if fused_silu else None
+    can_fuse = _can_fuse_w1w3(w1, tp, pack_block=pack_block)
 
     if can_fuse:
-        if fused_silu:
+        if pack_block is not None:
+            inter_size = w1.tensors['weight'].shape[-1]
+            w1w3 = _block_pack_w1w3(
+                w1, w3, groups=inter_size // pack_block)
+        elif fused_silu:
             w1w3 = _interleave_w1w3(w1, w3)
         else:
             w1w3 = _chunk_w1w3(w1, w3, tp=tp)

src/turbomind/core/module.h

@@ -177,6 +177,13 @@ class Param {
         return slot_ ? *slot_ : Tensor{};
     }
 
+    /// Replace the parameter slot (e.g. after quant alloc or MakeStridedPtrs).
+    void set(Tensor t)
+    {
+        TM_CHECK(slot_ != nullptr);
+        *slot_ = std::move(t);
+    }
+
     explicit operator bool() const
     {
         return slot_ && static_cast<bool>(*slot_);

src/turbomind/kernels/gemm/CMakeLists.txt

@@ -19,27 +19,102 @@ set(GEMM2_KERNELS_SM90
     tma.cu
     kernel/sm90_16816_4.cu
     kernel/sm90_16816_8.cu
-    kernel/sm90_16816_16.cu
     kernel/sm90_64n32_8.cu
+    kernel/sm90_64n16_16.cu
 )
 
-set(GEMM2_ARCH_90_ENABLED FALSE)
-set(_sm90_archs "${CMAKE_CUDA_ARCHITECTURES}")
-list(FILTER _sm90_archs INCLUDE REGEX "^90")
-if(_sm90_archs)
-    set(GEMM2_ARCH_90_ENABLED TRUE)
-else()
+# Bucket the selected archs by kernel family (ranges mirror arch.h).
+set(GEMM2_SM70_ARCHS "")
+set(GEMM2_SM75_ARCHS "")
+set(GEMM2_SM80_ARCHS "")
+set(GEMM2_SM90_ARCHS "")
+foreach(arch IN LISTS CMAKE_CUDA_ARCHITECTURES)
+    if(arch MATCHES "^([0-9]+)(a)?(-real)?$")
+        set(arch_number "${CMAKE_MATCH_1}")
+        if(arch_number GREATER_EQUAL 70 AND arch_number LESS 75)
+            list(APPEND GEMM2_SM70_ARCHS "${arch}")
+        elseif(arch_number GREATER_EQUAL 75 AND arch_number LESS 80)
+            list(APPEND GEMM2_SM75_ARCHS "${arch}")
+        elseif(arch_number GREATER_EQUAL 80 AND arch_number LESS 90)
+            list(APPEND GEMM2_SM80_ARCHS "${arch}")
+        elseif(arch_number EQUAL 90)
+            list(APPEND GEMM2_SM90_ARCHS "${arch}")
+        elseif(arch_number GREATER_EQUAL 100)
+            # SM100/SM120 run the SM80 (s16816 mma.sync) kernels for now.
+            list(APPEND GEMM2_SM80_ARCHS "${arch}")
+        endif()
+    endif()
+endforeach()
+
+if(NOT GEMM2_SM90_ARCHS)
     # When building for SM100+ without explicit SM90, still compile SM90 CUTLASS
     # kernels so the fat binary can run MoE models on H100 (CUTLASS fused path).
     set(_sm100_archs "${CMAKE_CUDA_ARCHITECTURES}")
     list(FILTER _sm100_archs INCLUDE REGEX "^100")
     if(_sm100_archs)
-        set(GEMM2_ARCH_90_ENABLED TRUE)
-        set(_sm90_archs "90")
+        set(GEMM2_SM90_ARCHS "90")
         message(STATUS "GEMM: auto-enabling SM90 CUTLASS kernels for H100 backward compatibility")
     endif()
 endif()
 
+message(STATUS "GEMM SM70 CUDA archs: ${GEMM2_SM70_ARCHS}")
+message(STATUS "GEMM SM75 CUDA archs: ${GEMM2_SM75_ARCHS}")
+message(STATUS "GEMM SM80 CUDA archs: ${GEMM2_SM80_ARCHS}")
+message(STATUS "GEMM SM90 CUDA archs: ${GEMM2_SM90_ARCHS}")
+
+function(configure_gemm2_kernel_target target)
+    set_property(TARGET ${target} PROPERTY POSITION_INDEPENDENT_CODE ON)
+    set_property(TARGET ${target} PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON)
+    target_compile_options(${target} PRIVATE
+            $<$<COMPILE_LANGUAGE:CUDA>:
+                    -Xptxas=-v
+                    --generate-line-info
+                    --threads 16>
+    )
+    target_link_libraries(${target} PRIVATE parser nvidia::cutlass::cutlass CUDA::cuda_driver)
+endfunction()
+
+# cublas.cu registers for every arch.
+add_library(gemm2_kernels STATIC cublas.cu)
+configure_gemm2_kernel_target(gemm2_kernels)
+
+set(GEMM2_KERNEL_TARGETS gemm2_kernels)
+
+if(GEMM2_SM70_ARCHS)
+    add_library(gemm2_kernels_sm70 STATIC ${GEMM2_KERNELS_SM70})
+    set_property(TARGET gemm2_kernels_sm70 PROPERTY
+            CUDA_ARCHITECTURES "${GEMM2_SM70_ARCHS}")
+    configure_gemm2_kernel_target(gemm2_kernels_sm70)
+    list(APPEND GEMM2_KERNEL_TARGETS gemm2_kernels_sm70)
+endif()
+
+if(GEMM2_SM75_ARCHS)
+    add_library(gemm2_kernels_sm75 STATIC ${GEMM2_KERNELS_SM75})
+    set_property(TARGET gemm2_kernels_sm75 PROPERTY
+            CUDA_ARCHITECTURES "${GEMM2_SM75_ARCHS}")
+    configure_gemm2_kernel_target(gemm2_kernels_sm75)
+    list(APPEND GEMM2_KERNEL_TARGETS gemm2_kernels_sm75)
+endif()
+
+if(GEMM2_SM80_ARCHS)
+    add_library(gemm2_kernels_sm80 STATIC ${GEMM2_KERNELS_SM80})
+    set_property(TARGET gemm2_kernels_sm80 PROPERTY
+            CUDA_ARCHITECTURES "${GEMM2_SM80_ARCHS}")
+    configure_gemm2_kernel_target(gemm2_kernels_sm80)
+    list(APPEND GEMM2_KERNEL_TARGETS gemm2_kernels_sm80)
+endif()
+
+if(GEMM2_SM90_ARCHS)
+    # SM90 kernels only compile for 90/90a; avoid building them for sm_100.
+    add_library(gemm2_sm90 STATIC ${GEMM2_KERNELS_SM90})
+    set_target_properties(gemm2_sm90 PROPERTIES
+        CUDA_ARCHITECTURES "${GEMM2_SM90_ARCHS}"
+    )
+    target_compile_definitions(gemm2_sm90 PRIVATE -DCUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED)
+    configure_gemm2_kernel_target(gemm2_sm90)
+    list(APPEND GEMM2_KERNEL_TARGETS gemm2_sm90)
+endif()
+
 add_library(gemm2
         gemm.cu
         kernel.cu
@@ -55,23 +130,22 @@ add_library(gemm2
         tuner/sampler.cu
         tuner/stopping_criterion.cc
         tuner/params.cc
-        ${GEMM2_KERNELS_SM70}
-        ${GEMM2_KERNELS_SM75}
-        ${GEMM2_KERNELS_SM80}
-        cublas.cu
         moe_utils_v2.cu
         test/test_utils.cu
 )
 
 target_link_libraries(gemm2 PRIVATE parser nvidia::cutlass::cutlass CUDA::cuda_driver)
 
+# Retain file-scope Registrar objects (self-registration). Same pattern as attention_kernels.
+foreach(target IN LISTS GEMM2_KERNEL_TARGETS)
+    target_link_libraries(gemm2 PUBLIC $<LINK_LIBRARY:WHOLE_ARCHIVE,${target}>)
+endforeach()
+
 # cublasGemmGroupedBatchedEx (CUDA 12.5+): grouped batched GEMM for MoE on SM100
-set(_has_sm100 FALSE)
 set(_archs_100 "${CMAKE_CUDA_ARCHITECTURES}")
 list(FILTER _archs_100 INCLUDE REGEX "^100")
 if(_archs_100 AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL "12.5")
-    set(_has_sm100 TRUE)
-    target_compile_definitions(gemm2 PRIVATE ENABLE_CUBLAS_GROUPED=1)
+    target_compile_definitions(gemm2_kernels PRIVATE ENABLE_CUBLAS_GROUPED=1)
     message(STATUS "GEMM: ENABLE_CUBLAS_GROUPED=1 (cublasGemmGroupedBatchedEx for MoE on SM100)")
 endif()
 
@@ -83,56 +157,3 @@ target_compile_options(gemm2 PRIVATE
 )
 set_property(TARGET gemm2 PROPERTY POSITION_INDEPENDENT_CODE ON)
 set_property(TARGET gemm2 PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON)
-
-if(GEMM2_ARCH_90_ENABLED)
-    # SM90 kernels only compile for 90/90a; avoid building them for sm_100.
-    add_library(gemm2_sm90 STATIC ${GEMM2_KERNELS_SM90})
-    set_target_properties(gemm2_sm90 PROPERTIES
-        CUDA_ARCHITECTURES "${_sm90_archs}"
-        POSITION_INDEPENDENT_CODE ON
-        CUDA_RESOLVE_DEVICE_SYMBOLS ON
-    )
-    target_compile_definitions(gemm2_sm90 PRIVATE -DCUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED)
-    target_compile_options(gemm2_sm90 PRIVATE
-        $<$<COMPILE_LANGUAGE:CUDA>:
-            -Xptxas=-v
-            --generate-line-info
-            --threads 16>
-    )
-    target_link_libraries(gemm2_sm90 PRIVATE parser nvidia::cutlass::cutlass CUDA::cuda_driver)
-    target_link_libraries(gemm2 PRIVATE gemm2_sm90)
-
-    target_compile_definitions(gemm2 PRIVATE GEMM2_ARCH_90_ENABLED)
-endif()
-
-if (BUILD_TEST)
-        # add_executable(test_gemm_v2
-        #         test/test_gemm_v2.cc
-        #         ../../models/llama/LlamaLinear.cu
-        #         test/reference.cu)
-        # target_link_libraries(test_gemm_v2 PRIVATE gemm2 core cublas quantization_kernels gpt_kernels)
-
-        # add_executable(test_moe_utils test/test_moe_utils.cu test/test_utils.cu)
-        # target_link_libraries(test_moe_utils PRIVATE gemm2 core cublas)
-
-        # if (NOT MSVC)
-        #         FetchContent_Declare(
-        #         repo-nvbench
-        #         GIT_REPOSITORY https://github.com/NVIDIA/nvbench.git
-        #         GIT_TAG        d8dced8a64d9ce305add92fa6d274fd49b569b7e
-        #         )
-
-        #         set(NVBench_ENABLE_EXAMPLES OFF)
-        #         set(NVBench_ENABLE_TESTING OFF)
-        #         set(BUILD_SHARED_LIBS OFF)
-
-        #         FetchContent_MakeAvailable(repo-nvbench)
-
-        #         add_executable(gemm_bench
-        #                 test/gemm_bench.cu
-        #                 # test/test_utils.cu
-        #                 test/quantization.cu
-        #                 test/reference.cu)
-        #         target_link_libraries(gemm_bench PRIVATE gemm2 core nvbench::nvbench cublas)
-        # endif ()
-endif ()

src/turbomind/kernels/gemm/arch.h

@@ -35,7 +35,7 @@ struct Sm100: Arch<1000, 1200> {
     static constexpr int value = 1000;
 };
 
-// SM12.x (e.g. sm_120): use same CUTLASS SM90 kernel family as pre-PR Sm90+ range
+// SM12.x (e.g. sm_120): no native kernels; falls back to the SM80 s16816 family
 struct Sm120: Arch<1200, 1300> {
     static constexpr int value = 1200;
 };
@@ -50,9 +50,9 @@ inline bool is_arch_compatible(int karch, int darch)
         case 750:
             return Sm75::is_compatible(darch);
         case 800:
-            return Sm80::is_compatible(darch);
+            return Sm80::is_compatible(darch) || Sm100::is_compatible(darch) || Sm120::is_compatible(darch);
         case 900:
-            return Sm90::is_compatible(darch) || Sm120::is_compatible(darch);
+            return Sm90::is_compatible(darch);
         case 1000:
             return Sm100::is_compatible(darch);
         case 1200: