Implementation:Sgl project Sglang CUTLASS Mixed Input Utils
| Knowledge Sources | |
|---|---|
| Domains | CUDA_Kernels, CUTLASS_Extensions, Mixed_Precision_GEMM |
| Last Updated | 2026-02-10 00:00 GMT |
Overview
Utility class for computing shared memory layouts, TMA transaction bytes, and scale/zero-point allocation for mixed-precision grouped GEMM operations in the CUTLASS framework.
Description
The MixedGroupedGemmInputUtils class template resides in the cutlass::gemm::collective::detail namespace and provides static constexpr helper methods for mixed-precision GEMM infrastructure. It extracts type information from a Collective type parameter (including SmemLayoutA, SmemLayoutB, SmemLayoutScale, ConversionMode, element types for scales and zero-points, etc.) and computes:
Shared memory element counts: elements_per_smem_scale() returns the number of scale elements needed in shared memory based on the conversion mode (zero for DirectConvert, cosize_v<SmemLayoutScale> when scales are used). elements_per_smem_zero() similarly returns zero-point element counts only for the ConvertAndScaleWithZero mode.
TMA transaction bytes: compute_tma_transaction_bytes_mk() and compute_tma_transaction_bytes_nk() calculate the byte count for TMA (Tensor Memory Accelerator) transfers for operands A and B, based on shared memory layout dimensions and element sizes. compute_tma_transaction_bytes_extra_mk() adds scale and zero-point TMA bytes with 128-byte alignment checks required by the TMA hardware.
The three ConversionMode values supported are DirectConvert (type conversion only), ConvertAndScale (conversion plus scale application), and ConvertAndScaleWithZero (conversion with scale and zero-point correction).
Usage
This utility is used internally by the SM90 warp-specialized mixed-input MMA collective to compute memory layout parameters. It is not called directly by user code but is instantiated as part of the CUTLASS template metaprogramming pipeline for Hopper GPU mixed-precision GEMM.
Code Reference
Source Location
- Repository: Sgl_project_Sglang
- File: sgl-kernel/csrc/cutlass_extensions/detail/collective/mixed_input_utils.hpp
- Lines: 1-482
Signature
namespace cutlass::gemm::collective::detail {
template <class Collective>
struct MixedGroupedGemmInputUtils {
// Shared memory element counts
static constexpr auto elements_per_smem_scale();
static constexpr auto elements_per_smem_zero();
// TMA transaction byte calculations
static constexpr uint32_t compute_tma_transaction_bytes_mk();
static constexpr uint32_t compute_tma_transaction_bytes_nk();
static constexpr uint32_t compute_tma_transaction_bytes_extra_mk();
};
} // namespace cutlass::gemm::collective::detail
Import
#include "cutlass_extensions/detail/collective/mixed_input_utils.hpp"
// Underlying dependencies:
#include "cute/arch/copy_sm90.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
#include "cute/util/type_traits.hpp"
#include "cutlass/cutlass.h"
#include "cutlass/numeric_conversion.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| Collective | template parameter | Yes | CUTLASS collective type providing SmemLayout, ConversionMode, element types, and schedule information |
Outputs
| Name | Type | Description |
|---|---|---|
| elements_per_smem_scale() | constexpr int | Number of scale elements to allocate in shared memory |
| elements_per_smem_zero() | constexpr int | Number of zero-point elements to allocate in shared memory |
| compute_tma_transaction_bytes_mk() | constexpr uint32_t | TMA transaction bytes for MK-dimension operand |
| compute_tma_transaction_bytes_nk() | constexpr uint32_t | TMA transaction bytes for NK-dimension operand |
Usage Examples
// Used internally by CollectiveMmaArrayMixedInput specialization
using Utils = detail::MixedGroupedGemmInputUtils<CollectiveType>;
// Compute shared memory requirements
constexpr auto scale_elems = Utils::elements_per_smem_scale();
constexpr auto zero_elems = Utils::elements_per_smem_zero();
// Compute TMA transaction sizes for pipeline setup
constexpr uint32_t tma_bytes_mk = Utils::compute_tma_transaction_bytes_mk();
constexpr uint32_t tma_bytes_nk = Utils::compute_tma_transaction_bytes_nk();