Implementation:Sgl project Sglang Flash MLA

Knowledge Sources	Sgl_project_Sglang
Domains	Kernel, Attention, MLA
Last Updated	2026-02-10 00:00 GMT

Overview

Python interface for Flash Multi-head Latent Attention (MLA) operations optimized for DeepSeek-style models with paged KV cache support.

Description

The flash_mla.py module provides specialized attention kernels for MLA architectures used in DeepSeek V2/V3 models. get_mla_metadata computes tile scheduling metadata and split information for MLA decode, with distinct code paths for dense FP8 KV cache and general (sparse/non-FP8) cases. flash_mla_with_kvcache performs paged MLA attention with KV cache, supporting optional FP8 KV cache (with descaling factors), optional sparse attention via an indices tensor, and configurable softmax scaling and causal masking. It dispatches to separate C++ ops for FP8 (fwd_kvcache_mla_fp8) and non-FP8 (fwd_kvcache_mla) paths. flash_mla_sparse_fwd handles sparse MLA forward prefill, computing attention over selected KV positions indicated by an indices tensor. The module requires the flashmla_ops CUDA extension (SM90) and CUDA Driver >= 12.4.

Usage

Use these functions for efficient decode and prefill in DeepSeek-style MLA models, where KV cache is stored in compressed latent representation format with paged memory management.

Code Reference

Source Location

• Repository: Sgl_project_Sglang
• File: sgl-kernel/python/sgl_kernel/flash_mla.py
• Lines: 1-156

Signature

def get_mla_metadata(
    cache_seqlens: torch.Tensor,
    num_q_tokens_per_head_k: int,
    num_heads_k: int,
    num_heads_q: Optional[int] = None,
    is_fp8_kvcache: bool = False,
    topk: Optional[int] = None,
) -> Tuple[torch.Tensor, torch.Tensor]: ...

def flash_mla_with_kvcache(
    q: torch.Tensor,
    k_cache: torch.Tensor,
    block_table: torch.Tensor,
    cache_seqlens: torch.Tensor,
    head_dim_v: int,
    tile_scheduler_metadata: torch.Tensor,
    num_splits: torch.Tensor,
    softmax_scale: Optional[float] = None,
    causal: bool = False,
    descale_q: torch.Tensor | None = None,
    descale_k: torch.Tensor | None = None,
    is_fp8_kvcache: bool = False,
    indices: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]: ...

def flash_mla_sparse_fwd(
    q: torch.Tensor,
    kv: torch.Tensor,
    indices: torch.Tensor,
    sm_scale: float,
    d_v: int = 512,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ...

Import

from sgl_kernel.flash_mla import (
    get_mla_metadata,
    flash_mla_with_kvcache,
    flash_mla_sparse_fwd,
)

I/O Contract

Inputs

Name	Type	Required	Description
cache_seqlens	torch.Tensor	Yes	Sequence lengths per batch, shape (batch_size,), int32
num_q_tokens_per_head_k	int	Yes	num_q_tokens_per_q_seq * num_heads_q // num_heads_k
num_heads_k	int	Yes	Number of key heads
num_heads_q	Optional[int]	No	Number of query heads (required for sparse attention)
is_fp8_kvcache	bool	No	Whether KV cache is in FP8 format
topk	Optional[int]	No	If set, enables sparse attention
q	torch.Tensor	Yes	Query tensor, shape (batch_size, seq_len_q, num_heads_q, head_dim)
k_cache	torch.Tensor	Yes	Paged KV cache, shape (num_blocks, page_size, num_heads_k, head_dim)
block_table	torch.Tensor	Yes	Page table mapping, shape (batch_size, max_blocks), int32
head_dim_v	int	Yes	Value head dimension
tile_scheduler_metadata	torch.Tensor	Yes	Scheduling metadata from get_mla_metadata
num_splits	torch.Tensor	Yes	Split info from get_mla_metadata
softmax_scale	Optional[float]	No	QK^T scaling, defaults to 1/sqrt(head_dim)
descale_q	Optional[torch.Tensor]	No	FP8 descaling factors for Q
descale_k	Optional[torch.Tensor]	No	FP8 descaling factors for K
indices	Optional[torch.Tensor]	No	Sparse attention indices, shape (batch_size, seq_len_q, topk)

Outputs

Name	Type	Description
tile_scheduler_metadata	torch.Tensor	Tile scheduling metadata (from get_mla_metadata)
num_splits	torch.Tensor	Split information (from get_mla_metadata)
out	torch.Tensor	Attention output, shape (batch_size, seq_len_q, num_heads_q, head_dim_v)
softmax_lse	torch.Tensor	Log-sum-exp, shape (batch_size, num_heads_q, seq_len_q)
max_logits	torch.Tensor	Maximum logits per position (sparse_fwd only)

Usage Examples

from sgl_kernel.flash_mla import get_mla_metadata, flash_mla_with_kvcache

# Step 1: Compute scheduling metadata
tile_meta, num_splits = get_mla_metadata(
    cache_seqlens, num_q_tokens_per_head_k=1,
    num_heads_k=8
)

# Step 2: Run MLA decode with paged KV cache
out, softmax_lse = flash_mla_with_kvcache(
    q, k_cache, block_table, cache_seqlens,
    head_dim_v=512,
    tile_scheduler_metadata=tile_meta,
    num_splits=num_splits,
    softmax_scale=0.125
)

Related Pages

• Environment:Sgl_project_Sglang_CUDA_GPU_Runtime