Implementation:Sgl project Sglang GEMM With Epilogue Visitor

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

Overview

CUTLASS GEMM kernel template supporting custom epilogue visitors for fused post-processing operations such as per-row/per-column scaling, adapted from NVIDIA TensorRT-LLM.

Description

The GemmWithEpilogueVisitor struct template in cutlass::gemm::kernel combines three components:

• Mma_ -- a threadblock-scoped matrix multiply-accumulate operation providing iterators for A and B operands, the MMA policy, and architectural tags
• Epilogue_ -- an epilogue with a visitor pattern (Epilogue::Visitor) that defines custom post-processing callbacks
• ThreadblockSwizzle_ -- a threadblock swizzling function for workload distribution

The struct extracts comprehensive type information: element types and layouts for A, B, and C matrices; ElementCompute from the epilogue visitor; alpha scale references in both row-major (LayoutAlphaCol) and column-major (LayoutAlphaRow) layouts; warp counts, alignment requirements, and pipeline stage count. It derives OperatorClass, ThreadblockShape, WarpShape, InstructionShape, and ArchTag from the MMA configuration.

The nested Arguments struct holds the GEMM problem size (as GemmCoord), tensor references for A, B, C, and D, per-row and per-column alpha scale references, the epilogue visitor's own arguments, and batch information including batch counts and strides.

The nested Params struct manages runtime kernel parameters and provides conversion from Arguments via an initializer static method.

Usage

Use this kernel when GEMM output requires fused custom post-processing that follows the visitor pattern, such as quantized inference with per-row/per-column dequantization scaling. It is combined with GemmUniversalBaseCompat for the device-layer launch interface.

Code Reference

Source Location

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

Signature

namespace cutlass::gemm::kernel {

template <
    typename Mma_,                // Threadblock-scoped MMA
    typename Epilogue_,           // Epilogue with visitor
    typename ThreadblockSwizzle_  // Threadblock swizzling function
>
struct GemmWithEpilogueVisitor {
  using Mma = Mma_;
  using Epilogue = Epilogue_;
  using EpilogueVisitor = typename Epilogue::Visitor;
  using ThreadblockSwizzle = ThreadblockSwizzle_;

  using ElementA = typename Mma::IteratorA::Element;
  using LayoutA = typename Mma::IteratorA::Layout;
  using ElementB = typename Mma::IteratorB::Element;
  using LayoutB = typename Mma::IteratorB::Layout;
  using ElementCompute = typename EpilogueVisitor::ElementCompute;
  using ElementC = typename EpilogueVisitor::ElementOutput;

  using OperatorClass = typename Mma::Operator::OperatorClass;
  using ThreadblockShape = typename Mma::Shape;
  using WarpShape = typename Mma::Operator::Shape;
  using InstructionShape =
      typename Mma::Policy::Operator::InstructionShape;

  static int const kStages = Mma::kStages;
  static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
  static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
  static int const kAlignmentC = EpilogueVisitor::kElementsPerAccess;

  struct Arguments { ... };
  struct Params { ... };
};

} // namespace cutlass::gemm::kernel

Import

#include "cutlass_extensions/gemm/gemm_with_epilogue_visitor.h"

// Underlying dependencies:
#include <cutlass/complex.h>
#include <cutlass/cutlass.h>
#include <cutlass/fast_math.h>
#include <cutlass/matrix_coord.h>
#include <cutlass/trace.h>

I/O Contract

Inputs

Name	Type	Required	Description
problem_size	GemmCoord	Yes	GEMM dimensions (M, N, K)
ref_A	TensorRefA	Yes	Reference to input matrix A (activations)
ref_B	TensorRefB	Yes	Reference to input matrix B (weights)
ref_C	TensorRefC	Yes	Reference to bias/source matrix C
ref_alpha_row	TensorRefAlphaRow	Yes	Per-row scale factors (column-major layout)
ref_alpha_col	TensorRefAlphaCol	Yes	Per-column scale factors (row-major layout)
epilogue_visitor_args	EpilogueVisitor::Arguments	Yes	Epilogue visitor-specific arguments
batch_count	int	No	Number of batches for batched GEMM (default 1)

Outputs

Name	Type	Description
ref_D	TensorRefC	GEMM output with applied epilogue visitor scaling, written in-place

Usage Examples

// Define the GEMM kernel with epilogue visitor
using GemmKernel = cutlass::gemm::kernel::GemmWithEpilogueVisitor<
    MmaType,
    EpilogueType,
    ThreadblockSwizzleType>;

// Set up arguments
typename GemmKernel::Arguments args{
    {M, N, K},            // problem_size
    ref_A, ref_B, ref_C,
    ref_D,
    ref_alpha_row,
    ref_alpha_col,
    epilogue_visitor_args,
    batch_count
};

// Launch via GemmUniversalBaseCompat
using GemmDevice = cutlass::gemm::device::GemmUniversalBaseCompat<GemmKernel>;
GemmDevice gemm_op;
gemm_op.initialize(args, workspace, stream);
gemm_op.run(stream);

Related Pages

• Environment:Sgl_project_Sglang_CUDA_GPU_Runtime