Implementation:Sgl project Sglang SM90 MMA Mixed Input

Knowledge Sources	Sgl_project_Sglang
Domains	CUDA_Kernels, CUTLASS_Extensions, Hopper_GPU, Mixed_Precision_GEMM
Last Updated	2026-02-10 00:00 GMT

Overview

SM90 (Hopper) warp-specialized collective MMA implementation for array-based mixed-precision GEMM, using TMA for global-to-shared memory transfers and GMMA with register-file sourcing for matrix computation.

Description

This file specializes the CollectiveMmaArrayMixedInput template for the MainloopSm90ArrayTmaGmmaWarpSpecializedMixedInput dispatch policy within the cutlass::gemm::collective namespace. The implementation leverages two key Hopper architecture features:

TMA (Tensor Memory Accelerator): Handles asynchronous data movement from global memory to shared memory, with computed transaction byte sizes for operands A, B, and optional scale/zero-point metadata. The utility class MixedGroupedGemmInputUtils provides the byte calculations with 128-byte alignment enforcement required by TMA hardware.

GMMA (Group Matrix Multiply-Accumulate): Performs the actual matrix math using warp-specialized execution. Different warps are assigned producer (data loading) and consumer (compute) roles in a pipelined fashion across multiple stages.

The template accepts 16+ parameters including Stages (pipeline depth), ClusterShape (SM cluster dimensions), TileShape_ (problem tiling), element types for both operands as optional tuples (to support mixed-precision inputs such as FP16 activations with INT4/INT8 weights), stride types, MMA configuration, and copy/layout atoms for both global and shared memory. A fixed GROUP_SIZE of 128 is used for quantization groups.

Nested types include SharedStorage (union of tensor and pipeline barrier storage), TensorStorage (shared memory for A, B operands and scale/zero data), TensorMapStorage (TMA descriptor storage), Arguments (host-side parameters), and Params (device-side parameters).

Usage

This collective is instantiated by the CUTLASS kernel dispatch infrastructure when targeting Hopper GPUs for mixed-precision grouped GEMM operations. It supports quantized MoE inference workloads where expert weights are stored in lower precision formats.

Code Reference

Source Location

• Repository: Sgl_project_Sglang
• File: sgl-kernel/csrc/cutlass_extensions/gemm/collective/sm90_mma_array_tma_gmma_rs_warpspecialized_mixed_input_.hpp
• Lines: 1-1538

Signature

#define GROUP_SIZE 128

namespace cutlass::gemm::collective {

template <
    int Stages,
    class ClusterShape,
    class KernelSchedule_,
    class TileShape_,
    class ElementAOptionalTuple,
    class StrideA_,
    class ElementBOptionalTuple,
    class StrideB_,
    class TiledMma_,
    class GmemTiledCopyA_,
    class SmemLayoutAtomA_,
    class SmemCopyAtomA_,
    class TransformA_,
    class GmemTiledCopyB_,
    class SmemLayoutAtomB_,
    class SmemCopyAtomB_,
    class TransformB_>
struct CollectiveMmaArrayMixedInput<
    MainloopSm90ArrayTmaGmmaWarpSpecializedMixedInput<
        Stages, ClusterShape, KernelSchedule_>,
    TileShape_,
    ElementAOptionalTuple, StrideA_,
    ElementBOptionalTuple, StrideB_,
    TiledMma_,
    GmemTiledCopyA_, SmemLayoutAtomA_, SmemCopyAtomA_, TransformA_,
    GmemTiledCopyB_, SmemLayoutAtomB_, SmemCopyAtomB_, TransformB_> {

  struct SharedStorage { ... };
  struct TensorStorage { ... };
  struct TensorMapStorage { ... };
  struct Arguments { ... };
  struct Params { ... };
};

} // namespace cutlass::gemm::collective

Import

#include "cutlass_extensions/gemm/collective/sm90_mma_array_tma_gmma_rs_warpspecialized_mixed_input_.hpp"

// Key dependencies:
#include "cute/algorithm/gemm.hpp"
#include "cute/arch/cluster_sm90.hpp"
#include "cute/arch/copy_sm90.hpp"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass_extensions/detail/collective/mixed_input_utils.hpp"

I/O Contract

Inputs

Name	Type	Required	Description
ElementAOptionalTuple	template tuple	Yes	Element type(s) for operand A (activations), may include scale info
ElementBOptionalTuple	template tuple	Yes	Element type(s) for operand B (weights), may include quantized type
Stages	int	Yes	Number of pipeline stages for async global-to-shared fetch
ClusterShape	class	Yes	SM cluster shape for cooperative kernel launch
TileShape_	class	Yes	Tile dimensions for the GEMM problem decomposition

Outputs

Name	Type	Description
accumulator	Fragment array	Accumulated MMA results in register file, passed to epilogue

Usage Examples

// This collective is instantiated via CUTLASS dispatch policy
using CollectiveMma = cutlass::gemm::collective::CollectiveMmaArrayMixedInput<
    cutlass::gemm::collective::MainloopSm90ArrayTmaGmmaWarpSpecializedMixedInput<
        3,                    // Stages
        Shape<_1,_1,_1>,     // ClusterShape
        KernelSchedule>,
    TileShape,
    cute::tuple<ElementA, ElementScale>,  // mixed-precision A
    StrideA,
    cute::tuple<ElementB>,                // quantized B
    StrideB,
    TiledMma,
    GmemTiledCopyA, SmemLayoutAtomA, SmemCopyAtomA, TransformA,
    GmemTiledCopyB, SmemLayoutAtomB, SmemCopyAtomB, TransformB>;

Related Pages

• Environment:Sgl_project_Sglang_CUDA_GPU_Runtime