Implementation:Sgl project Sglang SM100 MLA TMA WarpSpecialized

Knowledge Sources	Sgl_project_Sglang
Domains	CUDA Kernels, Attention, CUTLASS, Blackwell GPU
Last Updated	2026-02-10 00:00 GMT

Overview

The core warp-specialized Flash Multi-Head Attention (FMHA) kernel for Multi-head Latent Attention (MLA) on NVIDIA SM100 Blackwell GPUs, implementing concurrent data loading, matrix multiplication, and softmax computation through specialized warp roles.

Description

sm100_fmha_mla_tma_warpspecialized.hpp is the largest and most performance-critical kernel in the CUTLASS MLA implementation at over 2000 lines. It implements the Sm100FmhaMlaKernelTmaWarpspecialized struct within the cutlass::fmha::kernel namespace.

Template Parameters:

• TileShape: Defines tile dimensions as Shape<TileShapeH(128), TileShapeS, Shape<TileShapeL, TileShapeR>> for head dim, sequence tile, and latent/rope head dimensions
• Element_: Input element type (e.g., half_t, bfloat16_t)
• ElementAcc_: Accumulator type (typically float)
• ElementOut_: Output element type
• ElementLSE_: Log-sum-exp element type
• TileScheduler: Controls work distribution across CTAs
• kIsCpAsync: Toggle between CP.ASYNC and TMA data loading paths

Warp Specialization Architecture:

The kernel uses 256 threads organized into specialized warp groups via the WarpRole enum and kWarpAssignment bitmask:

• kLoad (role 0x2): 1-2 warps responsible for TMA or CP.ASYNC data loading from global to shared memory
• kMma (role 0x1): 1 warp performing matrix multiply-accumulate (UMMA) operations
• kCompute (role 0x3): 4 warps executing online softmax and output epilogue
• kLoadPageTable (role 0x4): 1 warp (CP.ASYNC mode) loading page table entries for paged KV cache

Two-SM Cluster Mode:

The kernel operates exclusively in 2-SM mode (kIs2Sm = true), forming thread block clusters of shape (2, 1, 1) using the SM100 cluster launch mechanism.

Pipeline Architecture:

Five concurrent pipelines coordinate data flow between warps:

• PipelineLoadQK: Load warp to MMA warp for Q/K tiles
• PipelineLoadPV: Load warp to MMA warp for P/V tiles (reuses same pipeline type)
• PipelineS: MMA warp to Compute warps for S=Q@K results
• PipelineP: Compute warps to MMA warp for P (softmax output) tiles
• PipelineO: MMA warp to Compute warps for O accumulator rescaling
• PipelinePT: Page table loader to Load warps for paged attention indices

MLA-Specific Design:

The attention computation is split into latent and rope components:

• Q@K is computed in two phases: Q_latent @ C_latent + Q_rope @ K_rope
• V loading uses a transposed layout (tma_load_c_latent_transpose) to read the value projection directly from the latent cache
• This eliminates the need for a separate V cache, reducing memory by exploiting MLA's weight-tied architecture

Shared Memory Layout:

The SharedStorage struct contains:

• PipelineStorage: Barrier storage for all 5 pipeline types
• TensorStorage: Shared memory tiles for Q, K/C, V/C (K and V share the same union), P, page table, and exchange buffers
• tmem_base_ptr: Tensor Memory (TMEM) base pointer for SM100 accumulator storage

TMEM Allocation:

The kernel uses TMEM (Tensor Memory) for accumulator storage with explicit allocation:

• kS0, kS1: Two S buffers for double-buffered QK results
• kO0: Output accumulator split into kSizeAccO-sized chunks for iterative PV accumulation

Key Methods:

• operator(): Main kernel entry point dispatching to warp role-specific loops
• load_tma / load_cpasync: Two data loading paths with TMA descriptor-based or CP.ASYNC-based loads
• load_page_table: Loads paged attention page indices into shared memory
• mma: Performs Q@K and P@V matrix multiplications with pipeline synchronization
• softmax: Online softmax with TMEM load, max reduction, exp2 computation, and quantization to half precision
• rescale: Correction factor application for online softmax running max updates
• epilogue: Final output writing with linear combination scaling and optional LSE output
• compute: Orchestrates softmax, rescale, and epilogue across K tiles

Paged Attention Support:

The kernel supports paged KV cache through:

• A Gather struct that remaps linear indices through a page table stored in shared memory
• ComposedLayout tensors that combine the gather operation with standard memory layouts
• Page sizes constrained to powers of 2 (CP.ASYNC path) or exactly TileShapeS (TMA path)

Usage

This kernel is not launched directly. It is instantiated as the template parameter to the cutlass::fmha::device::MLA class, which handles workspace allocation, parameter conversion, and cluster launch.

Code Reference

Source Location

• Repository: Sgl_project_Sglang
• File: sgl-kernel/csrc/attention/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp
• Lines: 1-2019

Signature

namespace cutlass::fmha::kernel {

template<
    class TileShape,      // Shape<_128, TileS, Shape<TileL, TileR>>
    class Element_,       // Input type (half_t, bfloat16_t)
    class ElementAcc_,    // Accumulator type (float)
    class ElementOut_,    // Output type
    class ElementLSE_,    // LSE type
    class TileScheduler,  // Work distribution policy
    bool kIsCpAsync = false
>
struct Sm100FmhaMlaKernelTmaWarpspecialized {

    // Constants
    static const bool kIs2Sm = true;
    static const int MaxThreadsPerBlock = 256;
    using ArchTag = cutlass::arch::Sm100;
    using ClusterShape = Shape<_2, _1, _1>;

    // Warp roles
    enum class WarpRole { kMma = 0x1, kLoad = 0x2, kCompute = 0x3, kLoadPageTable = 0x4, kEmpty = 0x0 };

    // Nested types
    struct MainloopArguments { ... };
    struct EpilogueArguments { ... };
    struct Arguments { ProblemShape problem_shape; MainloopArguments mainloop; EpilogueArguments epilogue; ... };
    struct Params { ... };
    struct SharedStorage { PipelineStorage pipelines; TensorStorage tensors; uint32_t tmem_base_ptr; };

    // Static methods
    static Params to_underlying_arguments(Arguments const& args, void* workspace);
    static size_t get_workspace_size(Arguments const& args);
    static dim3 get_grid_shape(Params const& params);
    static dim3 get_block_shape();
    static bool can_implement(Arguments const& args);

    // Kernel entry point
    CUTLASS_DEVICE void operator()(Params const& params, char* smem_raw);

    // Warp role methods
    void load_tma(...);
    void load_cpasync(...);
    void load_page_table(...);
    void mma(...);
    void softmax(...);
    void rescale(...);
    void compute(...);
    void epilogue(...);
};

} // namespace cutlass::fmha::kernel

Import

#include "cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp"

I/O Contract

Inputs

Name	Type	Required	Description
problem_shape	Shape<H=128, K, Shape<L, R>, B>	Yes	MLA problem dimensions: heads, seq_len, (latent_dim, rope_dim), batch
ptr_q_latent	Element*	Yes	Query latent projection tensor [H, L, B]
ptr_q_rope	Element*	Yes	Query rope projection tensor [H, R, B]
ptr_c_latent	Element*	Yes	Compressed KV latent cache [K, L, B]
ptr_k_rope	Element*	Yes	K rope cache [K, R, B]
softmax_scale	float	Yes	Attention softmax scaling factor (1/sqrt(d))
ptr_page_table	int*	No	Page table for paged KV cache
page_size	int	No	Page size (power of 2 for CP.ASYNC, TileShapeS for TMA)
ptr_seq	int*	No	Per-batch sequence lengths
split_kv	int	Yes	Number of KV splits for parallelism
ptr_o	ElementOut*	Yes	Output attention tensor
ptr_lse	ElementLSE*	No	Log-sum-exp output for reduction

Outputs

Name	Type	Description
ptr_o / ptr_o_acc	ElementOut* / ElementAcc*	Output attention or partial accumulator (when split_kv > 1)
ptr_lse / ptr_lse_acc	ElementLSE*	Log-sum-exp values or partial LSE (when split_kv > 1)

Usage Examples

Kernel Type Instantiation

using MlaKernel = cutlass::fmha::kernel::Sm100FmhaMlaKernelTmaWarpspecialized<
    cute::Shape<cute::_128, cute::_512, cute::Shape<cute::_512, cute::_64>>,
    cutlass::half_t,    // Element
    float,              // ElementAcc
    cutlass::half_t,    // ElementOut
    float,              // ElementLSE
    PersistentTileScheduler,
    false               // kIsCpAsync = false (use TMA)
>;

Problem Shape for DeepSeek MLA

// DeepSeek-V2 style MLA:
// H=128 heads, K=seq_len, D_latent=512, D_rope=64, B=batch_size
auto problem_shape = cute::make_shape(128, seq_len, cute::make_shape(512, 64), batch_size);

Warp Role Assignment

// kWarpAssignment (TMA mode) = 0x0021'3333
// Warp 0: kCompute (0x3)
// Warp 1: kCompute (0x3)
// Warp 2: kCompute (0x3)
// Warp 3: kCompute (0x3)
// Warp 4: kMma     (0x1)
// Warp 5: kLoad    (0x2)
// Warp 6: kEmpty   (0x0)
// Warp 7: kEmpty   (0x0)

Related Pages

• Environment:Sgl_project_Sglang_CUDA