Implementation:Sgl project Sglang GEMM Universal Base Compat

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

Overview

Compatibility device-layer wrapper replicating the CUTLASS 2.10 GemmUniversalBase API, enabling mixed-dtype GEMM and SmoothQuant kernels that are incompatible with newer CUTLASS 3.x device APIs.

Description

The GemmUniversalBaseCompat class template in cutlass::gemm::device wraps a GemmKernel_ and provides the standard CUTLASS device-layer interface for launching GEMM kernels. It is adapted from NVIDIA TensorRT-LLM's compatibility layer and replicated from CUTLASS 2.10 (SHA cc85b64).

The class extracts type aliases from the kernel including ElementA, ElementB, ElementC, LayoutA/B/C, ThreadblockShape, ElementAccumulator, EpilogueOutputOp, and ThreadblockSwizzle. Key functionality includes:

• get_grid_shape_ -- computes tiled grid dimensions with split-K support, handling both standard and split-K parallel modes via GemmUniversalMode
• can_implement -- validates that the kernel supports the given problem size
• get_workspace_size -- returns required temporary workspace bytes for split-K reduction
• initialize -- sets up kernel parameters from arguments and workspace pointer
• run -- launches the CUDA kernel with computed grid dimensions

The split-K implementation computes K-dimension tiling by dividing the reduction dimension across multiple threadblocks, with alignment to 16-element boundaries enforced by GemmKernel::kAlignmentA.

Usage

Wrap any CUTLASS 2.x-style GEMM kernel with this class when using it in the SGLang build environment. Required for mixed-dtype GEMM kernels and SmoothQuant quantization kernels that rely on older kernel-level APIs not compatible with CUTLASS 3.x device layers.

Code Reference

Source Location

• Repository: Sgl_project_Sglang
• File: sgl-kernel/csrc/cutlass_extensions/gemm/gemm_universal_base_compat.h
• Lines: 1-356

Signature

namespace cutlass::gemm::device {

template <typename GemmKernel_>
class GemmUniversalBaseCompat {
public:
  using GemmKernel = GemmKernel_;
  using ThreadblockShape = typename GemmKernel::Mma::Shape;

  using ElementA = typename GemmKernel::ElementA;
  using LayoutA = typename GemmKernel::LayoutA;
  using ElementB = typename GemmKernel::ElementB;
  using LayoutB = typename GemmKernel::LayoutB;
  using ElementC = typename GemmKernel::ElementC;
  using LayoutC = typename GemmKernel::LayoutC;
  using ElementAccumulator =
      typename GemmKernel::Mma::Policy::Operator::ElementC;
  using EpilogueOutputOp = typename GemmKernel::EpilogueOutputOp;
  using ThreadblockSwizzle = typename GemmKernel::ThreadblockSwizzle;

  using Arguments = typename GemmKernel::Arguments;

protected:
  typename GemmKernel::Params params_;

  static void get_grid_shape_(
      gemm::GemmCoord& grid_tiled_shape,
      int& gemm_k_size,
      Arguments const& args);

public:
  static Status can_implement(Arguments const& args);
  static size_t get_workspace_size(Arguments const& args);
  Status initialize(Arguments const& args, void* workspace, cudaStream_t);
  Status run(cudaStream_t stream);
};

} // namespace cutlass::gemm::device

Import

#include "cutlass_extensions/gemm/gemm_universal_base_compat.h"

// Underlying dependencies:
#include <cutlass/cutlass.h>
#include <cutlass/device_kernel.h>
#include <cutlass/trace.h>

I/O Contract

Inputs

Name	Type	Required	Description
GemmKernel_	template parameter	Yes	CUTLASS 2.x GEMM kernel type providing element types, layouts, and MMA configuration
args	Arguments	Yes	GEMM problem size, tensor references, epilogue parameters, and split-K configuration
workspace	void*	No	Temporary workspace for split-K parallel reduction (nullptr if not needed)
stream	cudaStream_t	Yes	CUDA stream for kernel launch

Outputs

Name	Type	Description
status	cutlass::Status	Success or error code indicating kernel launch result
output (via args)	TensorRefD	GEMM result written to the output tensor reference provided in arguments

Usage Examples

// Define the GEMM device type using compatibility wrapper
using GemmDevice = cutlass::gemm::device::GemmUniversalBaseCompat<MyGemmKernel>;

// Set up arguments
typename GemmDevice::Arguments args{
    cutlass::gemm::GemmUniversalMode::kGemm,
    {M, N, K},       // problem size
    ref_A, ref_B, ref_C, ref_D,
    epilogue_params
};

// Check feasibility and get workspace
auto status = GemmDevice::can_implement(args);
size_t workspace_size = GemmDevice::get_workspace_size(args);

// Initialize and run
GemmDevice gemm_op;
gemm_op.initialize(args, workspace_ptr, stream);
gemm_op.run(stream);

Related Pages

• Environment:Sgl_project_Sglang_CUDA_GPU_Runtime