recipes/docker-images/mia-tp4/runtime/flash_mla_sm120.py

"""SM120 FlashMLA sparse decode implementation.

On SM120 (Blackwell Desktop / RTX PRO 6000) the flash_mla CUDA kernel
is not available, so this module provides alternative implementations:

- A fused Triton kernel (default, ``SGLANG_SM120_TRITON_FLASHMLA=1``)
- A pure-PyTorch fallback (``SGLANG_SM120_TRITON_FLASHMLA=0``)

The FP8 KV cache uses a page-internal layout where NOPE+ROPE data has
stride (nope_dim + rope_dim*2) per token, and scales are stored in a
separate region at the end of each page.
"""

import logging
import math
from typing import Optional

import torch
import triton
import triton.language as tl

from sglang.srt.environ import envs
from sglang.srt.utils import is_hip

logger = logging.getLogger(__name__)
_is_hip = is_hip()

_GLM_DSA_MODEL_ARCHS = (
    "GlmMoeDsaForCausalLM",
    "GlmMoeDsaForCausalLMNextN",
)

# Page layout constants for DSv4-Flash (MODEL1):
#   nope_dim = 448, rope_dim = 64, quantize_block_size = 64
#   nope_rope_stride = 448 + 64*2 = 576 bytes per token
#   scale_stride = ceil(448/64) + 1 = 8 bytes per token (7 scales + 1 pad)
#   bytes_per_token = 448 + 128 + 8 = 584
#   page_bytes = ceil_div(page_size * 584, 576) * 576

_NOPE_DIM = 448
_ROPE_DIM = 64
_NOPE_ROPE_STRIDE = _NOPE_DIM + _ROPE_DIM * 2  # 576
_TILE_SIZE = 64
_NUM_TILES = _NOPE_DIM // _TILE_SIZE  # 7
_SCALE_STRIDE = _NUM_TILES + 1  # 8 (7 scales + 1 pad)
_D = _NOPE_DIM + _ROPE_DIM  # 512


def _gather_and_dequant(k_cache, indices, page_size):
    """Gather KV entries from the paged buffer using correct page-internal addressing.

    Args:
        k_cache: (num_pages, page_size, 1, bytes_per_token) float8_e4m3fn
                 Non-contiguous view of the raw page buffer.
        indices: (...) int32/int64, token-level indices. -1 = invalid.
        page_size: tokens per page (256)

    Returns:
        kv: (..., _D) bfloat16, dequantized KV vectors
    """
    idx_shape = indices.shape
    flat_idx = indices.reshape(-1)  # (N,)
    N = flat_idx.shape[0]
    device = k_cache.device

    # Page-level addressing
    page_bytes = k_cache.stride(0)  # actual byte stride between pages
    pages = flat_idx // page_size
    offsets = flat_idx % page_size

    # Clamp invalid indices
    safe_pages = pages.clamp(min=0)
    safe_offsets = offsets.clamp(min=0)

    # Access raw buffer as uint8 — use as_strided to get full page view
    num_pages = k_cache.shape[0]
    raw_pages = k_cache.as_strided(
        (num_pages, page_bytes),
        (page_bytes, 1),
    ).view(torch.uint8)  # (num_pages, page_bytes) uint8
    # Note: float8_e4m3fn and uint8 are both 1 byte, view is safe

    # Compute byte offsets within each page
    # NOPE: page[safe_page, safe_offset * 576 + 0:448]
    # ROPE: page[safe_page, safe_offset * 576 + 448:576]
    # SCALES: page[safe_page, page_size * 576 + safe_offset * 8 + 0:7]

    nope_base = safe_offsets * _NOPE_ROPE_STRIDE  # (N,)
    nope_offsets = nope_base.unsqueeze(-1) + torch.arange(
        _NOPE_DIM, device=device, dtype=torch.long
    )  # (N, 448)

    rope_base = nope_base + _NOPE_DIM  # (N,)
    rope_offsets = rope_base.unsqueeze(-1) + torch.arange(
        _ROPE_DIM * 2, device=device, dtype=torch.long
    )  # (N, 128)

    scale_section_offset = page_size * _NOPE_ROPE_STRIDE  # 147456
    scale_base = scale_section_offset + safe_offsets * _SCALE_STRIDE  # (N,)
    scale_offsets = scale_base.unsqueeze(-1) + torch.arange(
        _NUM_TILES, device=device, dtype=torch.long
    )  # (N, 7)

    # Gather bytes per page — use advanced indexing
    # raw_pages[safe_pages, nope_offsets] → (N, 448)
    page_idx_nope = safe_pages.unsqueeze(-1).expand_as(nope_offsets)
    nope_bytes = raw_pages[page_idx_nope, nope_offsets]  # (N, 448) uint8

    page_idx_rope = safe_pages.unsqueeze(-1).expand_as(rope_offsets)
    rope_bytes = raw_pages[page_idx_rope, rope_offsets]  # (N, 128) uint8

    page_idx_scale = safe_pages.unsqueeze(-1).expand_as(scale_offsets)
    scale_bytes = raw_pages[page_idx_scale, scale_offsets]  # (N, 7) uint8

    # Reinterpret dtypes
    nope_fp8 = nope_bytes.view(torch.float8_e4m3fn)  # (N, 448)
    rope_bf16 = rope_bytes.contiguous().view(torch.bfloat16)  # (N, 64)
    scale_e8m0 = scale_bytes.view(torch.float8_e8m0fnu)  # (N, 7)

    # Dequantize: nope_tile * scale_tile → bf16 (vectorized)
    result = torch.empty(N, _D, dtype=torch.bfloat16, device=device)
    result[:, :_NOPE_DIM] = (
        (
            nope_fp8.view(N, _NUM_TILES, _TILE_SIZE).float()
            * scale_e8m0.view(N, _NUM_TILES, 1).float()
        )
        .view(N, _NOPE_DIM)
        .to(torch.bfloat16)
    )
    result[:, _NOPE_DIM:] = rope_bf16

    return result.reshape(*idx_shape, _D)


def _sm120_sparse_decode_fwd(
    q,
    k_cache,
    indices,
    topk_length,
    attn_sink,
    head_dim_v,
    softmax_scale,
    extra_k_cache=None,
    extra_indices=None,
    extra_topk_length=None,
):
    B, s_q, H_q, D_qk = q.shape
    num_pages, page_size, H_k, bpt = k_cache.shape
    topk = indices.shape[-1]

    invalid_mask = indices < 0
    safe_indices = indices.clamp(min=0)

    if topk_length is not None:
        topk_range = torch.arange(topk, device=topk_length.device).view(1, 1, topk)
        invalid_mask = invalid_mask | (topk_range >= topk_length.view(B, 1, 1))

    # Gather and dequantize using page-aware addressing
    gathered_kv = _gather_and_dequant(k_cache, safe_indices, page_size)

    if extra_k_cache is not None and extra_indices is not None:
        extra_topk = extra_indices.shape[-1]
        extra_page_size = extra_k_cache.shape[1]
        extra_invalid = extra_indices < 0
        extra_safe = extra_indices.clamp(min=0)
        if extra_topk_length is not None:
            extra_range = torch.arange(
                extra_topk, device=extra_topk_length.device
            ).view(1, 1, extra_topk)
            extra_invalid = extra_invalid | (
                extra_range >= extra_topk_length.view(B, 1, 1)
            )
        extra_kv = _gather_and_dequant(extra_k_cache, extra_safe, extra_page_size)
        gathered_kv = torch.cat([gathered_kv, extra_kv], dim=2)
        invalid_mask = torch.cat([invalid_mask, extra_invalid], dim=2)

    gathered_kv[invalid_mask] = 0.0

    q_f = q.float()
    kv_f = gathered_kv.float()
    kv_d = kv_f.shape[-1]
    if D_qk != kv_d:
        q_f = q_f[..., :kv_d]

    scores = torch.einsum("bshd,bstd->bsht", q_f, kv_f) * softmax_scale
    scores.masked_fill_(invalid_mask.unsqueeze(2).expand_as(scores), float("-inf"))

    lse = torch.logsumexp(scores, dim=-1)

    if attn_sink is not None:
        lse_for_out = torch.logsumexp(
            torch.stack([lse, attn_sink.view(1, 1, H_q).expand_as(lse)], dim=0), dim=0
        )
    else:
        lse_for_out = lse.clone()

    lonely = lse == float("-inf")
    lse_for_out[lonely] = float("inf")
    weights = torch.exp(scores - lse_for_out.unsqueeze(-1))
    out = torch.einsum("bsht,bstv->bshv", weights, kv_f[..., :head_dim_v])
    out[lonely.unsqueeze(-1).expand_as(out)] = 0.0

    return out.to(torch.bfloat16), lse.permute(0, 2, 1)


# SM120 FlashMLA: default FlashInfer (CUTLASS SM120 sparse MLA decode).
# Override with SGLANG_SM120_FLASHMLA_BACKEND=triton|torch to force fallback.
_sm120_default_backend = envs.SGLANG_SM120_FLASHMLA_BACKEND.get()


SM120_DECODE_MAX_TOKENS = 64


def _flash_mla_sm120_prefill(
    q,
    k_cache,
    indices,
    topk_length,
    attn_sink,
    head_dim_v,
    softmax_scale,
    extra_k_cache,
    extra_indices,
    extra_topk_length,
):
    from flashinfer.mla._sparse_mla_sm120 import _sparse_mla_sm120_paged_attention

    q2 = q.squeeze(1) if q.ndim == 4 else q
    num_tokens, num_heads, _ = q2.shape
    dev = q2.device
    kv_u8 = k_cache.view(torch.uint8) if k_cache.dtype != torch.uint8 else k_cache
    src_pbs = k_cache.shape[1] if k_cache.ndim >= 3 else _PBS_SRC
    idx = indices.squeeze(1) if indices.dim() == 3 else indices
    kv_64 = (
        _split_kv_pages_to_64(kv_u8, src_pbs, touched_indices=idx)
        if src_pbs != _PBS_DST
        else kv_u8
    )
    extra_kv_u8 = (
        extra_k_cache.view(torch.uint8)
        if extra_k_cache is not None and extra_k_cache.dtype != torch.uint8
        else extra_k_cache
    )
    extra_idx = (
        extra_indices.squeeze(1)
        if extra_indices is not None and extra_indices.dim() == 3
        else extra_indices
    )
    # V4.1 ratio-1/2 sources use 256/128-token pages. The prefill
    # specialization needs 64-token pages for both sources, with independent
    # scratch buffers so converting the extra source cannot overwrite SWA.
    if extra_kv_u8 is not None and extra_kv_u8.shape[1] != _PBS_DST:
        extra_kv_u8 = _split_kv_pages_to_64(
            extra_kv_u8, extra_kv_u8.shape[1],
            touched_indices=extra_idx, buffer_tag="extra",
        )
    output = q2.new_empty((num_tokens, num_heads, head_dim_v), dtype=torch.bfloat16)
    out_lse = torch.empty((num_tokens, num_heads), dtype=torch.float32, device=dev)
    _sparse_mla_sm120_paged_attention(
        q2,
        kv_64,
        idx,
        output,
        out_lse,
        softmax_scale,
        topk_length=topk_length,
        attn_sink=attn_sink,
        extra_kv_cache=extra_kv_u8,
        extra_indices=extra_idx,
        extra_topk_length=extra_topk_length,
        mid_out=None,
        mid_lse=None,
    )
    return (output.unsqueeze(1), None)


def flash_mla_with_kvcache_sm120(**kwargs):
    """SM120 FlashMLA sparse decode entry point.

    Dispatches to FlashInfer (default if available), Triton, or PyTorch fallback.
    """
    q = kwargs["q"]
    k_cache = kwargs["k_cache"]
    indices = kwargs["indices"]
    topk_length = kwargs.get("topk_length")
    attn_sink = kwargs.get("attn_sink")
    head_dim_v = kwargs["head_dim_v"]
    softmax_scale = kwargs.get("softmax_scale")
    if softmax_scale is None:
        softmax_scale = q.shape[-1] ** (-0.5)
    extra_k_cache = kwargs.get("extra_k_cache")
    extra_indices = kwargs.get("extra_indices_in_kvcache")
    extra_topk_length = kwargs.get("extra_topk_length")

    if _sm120_default_backend == "flashinfer":
        if q.shape[0] > SM120_DECODE_MAX_TOKENS:
            return _flash_mla_sm120_prefill(
                q,
                k_cache,
                indices,
                topk_length,
                attn_sink,
                head_dim_v,
                softmax_scale,
                extra_k_cache,
                extra_indices,
                extra_topk_length,
            )
        return _flash_mla_flashinfer(
            q,
            k_cache,
            indices,
            topk_length,
            attn_sink,
            head_dim_v,
            softmax_scale,
            extra_k_cache,
            extra_indices,
            extra_topk_length,
        )

    if _sm120_default_backend == "triton":
        from sglang.kernels.ops.attention.flash_mla_sm120_triton import (
            flash_mla_sparse_decode_triton,
        )

        out, lse = flash_mla_sparse_decode_triton(
            q,
            k_cache,
            indices,
            topk_length,
            attn_sink,
            head_dim_v,
            softmax_scale,
            extra_k_cache,
            extra_indices,
            extra_topk_length,
        )
        return (out, lse)

    out, lse = _sm120_sparse_decode_fwd(
        q,
        k_cache,
        indices,
        topk_length,
        attn_sink,
        head_dim_v,
        softmax_scale,
        extra_k_cache,
        extra_indices,
        extra_topk_length,
    )
    return (out, lse)


# --- Page-split utilities: pbs=256 → pbs=64 ---
# SGLang SWA KV cache footer layout per 256-token page:
#   [data: 256 * 576 bytes] [scale: 256 * 8 bytes] [padding]
# FlashInfer decode_dsv4 expects per 64-token page:
#   [data: 64 * 576 bytes] [scale: 64 * 8 bytes] [padding to 37440]
_PBS_SRC = 256  # SGLang physical page size
_PBS_DST = 64  # FlashInfer page_block_size
_NOPE_ROPE_STRIDE = 576  # bytes per token for nope+rope
_SCALE_STRIDE = 8  # bytes per token for scale (7 + 1 pad)
_BYTES_PER_DST_PAGE = (
    _PBS_DST * _NOPE_ROPE_STRIDE + _PBS_DST * _SCALE_STRIDE
)  # 64*576 + 64*8 = 37376 + 512 = 37888
# Padded to 576 alignment

_BYTES_PER_DST_PAGE_PADDED = math.ceil(_BYTES_PER_DST_PAGE / 576) * 576  # 37440


@triton.jit
def _page_split_kernel(
    src_ptr,
    dst_ptr,
    N_pages,
    src_stride0: tl.constexpr,
    dst_stride0: tl.constexpr,
    DATA_PER_SUB: tl.constexpr,  # 64 * 576 = 36864
    SCALE_PER_SUB: tl.constexpr,  # 64 * 8 = 512
    SRC_SCALE_OFF: tl.constexpr,  # 256 * 576 = 147456
    DST_SCALE_OFF: tl.constexpr,  # 64 * 576 = 36864
    RATIO: tl.constexpr,  # 4
    BLOCK_SIZE: tl.constexpr,
    mask_ptr,
    HAS_MASK: tl.constexpr,
):
    """Fused page-split: copy data+scale for all sub-pages in one kernel.

    When HAS_MASK is set, only pages flagged in ``mask_ptr`` (int8, 1=touched)
    are copied; untouched pages are skipped so the kernel no longer rewrites the
    entire KV pool every decode step.
    """
    pid = tl.program_id(0)
    page_idx = pid // RATIO
    sub = pid % RATIO

    if page_idx >= N_pages:
        return

    if HAS_MASK:
        if tl.load(mask_ptr + page_idx) == 0:
            return

    # All layout strides/offsets are 8-byte aligned (asserted at the call
    # site), so copy in u64 lanes instead of single bytes.
    src_u64 = src_ptr.to(tl.pointer_type(tl.uint64))
    dst_u64 = dst_ptr.to(tl.pointer_type(tl.uint64))
    src_base = src_u64 + page_idx * (src_stride0 // 8)
    dst_base = dst_u64 + (page_idx * RATIO + sub) * (dst_stride0 // 8)

    DATA_U64: tl.constexpr = DATA_PER_SUB // 8
    SCALE_U64: tl.constexpr = SCALE_PER_SUB // 8

    # Copy data region: DATA_U64 u64 lanes from src offset sub*DATA_U64
    data_src_off = sub * DATA_U64
    for start in tl.range(0, DATA_U64, BLOCK_SIZE):
        offs = start + tl.arange(0, BLOCK_SIZE)
        mask = offs < DATA_U64
        vals = tl.load(src_base + data_src_off + offs, mask=mask)
        tl.store(dst_base + offs, vals, mask=mask)

    # Copy scale region: SCALE_U64 u64 lanes
    scale_src_off = SRC_SCALE_OFF // 8 + sub * SCALE_U64
    for start in tl.range(0, SCALE_U64, BLOCK_SIZE):
        offs = start + tl.arange(0, BLOCK_SIZE)
        mask = offs < SCALE_U64
        vals = tl.load(src_base + scale_src_off + offs, mask=mask)
        tl.store(dst_base + DST_SCALE_OFF // 8 + offs, vals, mask=mask)


@triton.jit
def _page_mark_kernel(
    indices_ptr,
    mask_ptr,
    N_idx,
    SRC_PBS: tl.constexpr,
    BLOCK: tl.constexpr,
):
    """Mark touched source pages (1 byte each) from token-level indices.

    ``indices`` are token indices into the pbs=SRC_PBS SWA pool; -1 = invalid.
    Each valid token marks ``mask[token // SRC_PBS] = 1``. Concurrent stores of
    the same value 1 are safe (no atomic needed).
    """
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    valid = offs < N_idx
    idx = tl.load(indices_ptr + offs, mask=valid, other=-1)
    keep = valid & (idx >= 0)
    page = idx // SRC_PBS
    tl.store(mask_ptr + page, tl.full((BLOCK,), 1, tl.int8), mask=keep)


def _split_kv_pages_to_64(
    kv_u8: torch.Tensor,
    src_pbs: int,
    touched_indices: Optional[torch.Tensor] = None,
    buffer_tag: str = "primary",
) -> torch.Tensor:
    """Split pbs=N footer-format pages into pbs=64 footer-format pages.

    When ``touched_indices`` (token-level int32 indices into the pbs=src_pbs
    SWA pool, -1 = invalid) is provided, only the source pages that actually
    contain a referenced token are copied. This avoids rewriting the entire KV
    pool on every decode step (only ~2*batch pages are touched vs the full
    pool). The output buffer is persistent and reused across steps; untouched
    dst pages simply retain their (unreferenced) stale data.
    """
    assert src_pbs % _PBS_DST == 0 and src_pbs >= _PBS_DST
    if src_pbs == _PBS_DST:
        return kv_u8

    N = kv_u8.shape[0]
    ratio = src_pbs // _PBS_DST
    num_dst_pages = N * ratio

    from sglang.srt.runtime_context import get_resources

    # Pre-allocated grow-only buffer for page-split output per device.
    dev = kv_u8.device
    buffers = get_resources().buffers
    key = f"flash_mla_sm120_split:{dev}:{buffer_tag}"
    buf = buffers.get(key)
    if buf is None or buf.shape[0] < num_dst_pages:
        # The first allocation can happen under inference mode (autotune), but
        # the buffer is written again during CUDA graph capture outside
        # inference mode, where an inference tensor cannot be mutated.
        with torch.inference_mode(False):
            buf = torch.empty(
                num_dst_pages,
                _BYTES_PER_DST_PAGE_PADDED,
                dtype=torch.uint8,
                device=dev,
            )
        buffers[key] = buf
    out = buf[:num_dst_pages]

    # Get raw 2D view of source
    src_2d = kv_u8
    if src_2d.ndim == 4:
        src_stride0 = src_2d.stride(0)
        src_2d = torch.as_strided(src_2d, (N, src_stride0), (src_stride0, 1))
    else:
        src_stride0 = src_2d.stride(0)

    use_mask = touched_indices is not None and touched_indices.numel() > 0
    mask_ptr = src_2d  # dummy, never dereferenced when HAS_MASK is False
    if use_mask:
        # Persistent per-device int8 mask, zeroed each call (cheap memset,
        # captured cleanly by CUDA graph). 1 = page is referenced this step.
        mkey = f"flash_mla_sm120_mask:{dev}:{buffer_tag}"
        mbuf = buffers.get(mkey)
        if mbuf is None or mbuf.shape[0] < N:
            # The first allocation can happen under inference mode (autotune),
            # but the buffer is zeroed again later during CUDA graph capture
            # outside inference mode -- an inference tensor cannot be mutated
            # there, so force a normal tensor.
            with torch.inference_mode(False):
                mbuf = torch.empty(N, dtype=torch.int8, device=dev)
            buffers[mkey] = mbuf
        mask = mbuf[:N]
        mask.zero_()
        idx_flat = touched_indices.reshape(-1).contiguous()
        if idx_flat.dtype != torch.int32:
            idx_flat = idx_flat.to(torch.int32)
        _MARK_BLOCK = 1024
        _page_mark_kernel[(triton.cdiv(idx_flat.numel(), _MARK_BLOCK),)](
            idx_flat,
            mask,
            idx_flat.numel(),
            src_pbs,  # SRC_PBS
            _MARK_BLOCK,
        )
        mask_ptr = mask

    assert src_stride0 % 8 == 0 and _BYTES_PER_DST_PAGE_PADDED % 8 == 0
    grid = (N * ratio,)
    _page_split_kernel[grid](
        src_2d,
        out,
        N,
        src_stride0,
        _BYTES_PER_DST_PAGE_PADDED,
        _PBS_DST * _NOPE_ROPE_STRIDE,  # DATA_PER_SUB = 36864
        _PBS_DST * _SCALE_STRIDE,  # SCALE_PER_SUB = 512
        src_pbs * _NOPE_ROPE_STRIDE,  # SRC_SCALE_OFF = 147456
        _PBS_DST * _NOPE_ROPE_STRIDE,  # DST_SCALE_OFF = 36864
        ratio,  # RATIO = 4
        1024,  # BLOCK_SIZE
        mask_ptr,
        use_mask,  # HAS_MASK
    )

    bpt = _NOPE_ROPE_STRIDE + _SCALE_STRIDE  # 584
    return out.as_strided(
        (num_dst_pages, _PBS_DST, 1, bpt),
        (_BYTES_PER_DST_PAGE_PADDED, bpt, bpt, 1),
    )


def _flash_mla_flashinfer(
    q,
    k_cache,
    indices,
    topk_length,
    attn_sink,
    head_dim_v,
    softmax_scale,
    extra_k_cache,
    extra_indices,
    extra_topk_length,
):
    """FlashInfer SM120 sparse MLA via the paged-attention dispatcher.

    SGLang SWA pool uses page_size=256 (footer format: 256*576 bytes data + 256*8 bytes scale).
    FlashInfer decode_dsv4 fast path requires page_block_size=64 (footer: 64*576 + 64*8).
    We split 256-token pages into 4 virtual 64-token pages.
    Token indices are invariant under page-split (identity mapping).
    """
    from flashinfer.mla._sparse_mla_sm120 import (
        _DECODE_MAX_TOKENS as _FI_DECODE_MAX_TOKENS,
    )
    from flashinfer.mla._sparse_mla_sm120 import (
        _sparse_mla_sm120_paged_attention,
    )

    B, _, H, D = q.shape  # (batch, 1, num_heads, head_dim)
    dev = q.device

    # Indices: no remapping needed (page-split preserves token addressing).
    idx = indices.squeeze(1) if indices.dim() == 3 else indices

    # --- Page-split: convert pbs=N kv_cache to pbs=64 view ---
    # Only the SWA pages actually referenced by `idx` are copied (the rest of
    # the persistent dst buffer is left untouched and never read).
    kv_u8 = k_cache.view(torch.uint8) if k_cache.dtype != torch.uint8 else k_cache
    src_pbs = k_cache.shape[1] if k_cache.ndim >= 3 else _PBS_SRC
    kv_64 = (
        _split_kv_pages_to_64(kv_u8, src_pbs, touched_indices=idx)
        if src_pbs != _PBS_DST
        else kv_u8
    )

    extra_kv_u8 = (
        extra_k_cache.view(torch.uint8)
        if extra_k_cache is not None and extra_k_cache.dtype != torch.uint8
        else extra_k_cache
    )
    extra_kv_64 = extra_kv_u8

    extra_idx = (
        extra_indices.squeeze(1)
        if extra_indices is not None and extra_indices.dim() == 3
        else extra_indices
    )

    output = torch.empty(B, H, head_dim_v, dtype=torch.bfloat16, device=dev)
    out_lse = torch.empty(B, H, dtype=torch.float32, device=dev)

    # Use split-K for decode-sized batches and paged attention otherwise.
    if B <= _FI_DECODE_MAX_TOKENS:
        topk = idx.shape[-1]
        extra_topk = extra_idx.shape[-1] if extra_idx is not None else 0
        _BI = 64
        num_splits = (topk + _BI - 1) // _BI + (
            (extra_topk + _BI - 1) // _BI if extra_topk > 0 else 0
        )
        mid_out = torch.empty(
            B, H, num_splits, head_dim_v, dtype=torch.bfloat16, device=dev
        )
        mid_lse = torch.empty(B, H, num_splits, dtype=torch.float32, device=dev)
    else:
        mid_out = None
        mid_lse = None

    _sparse_mla_sm120_paged_attention(
        q.squeeze(1) if q.ndim == 4 else q,
        kv_64,
        idx,
        output,
        out_lse,
        softmax_scale,
        d_v=head_dim_v,
        topk_length=topk_length,
        attn_sink=attn_sink,
        extra_kv_cache=extra_kv_64,
        extra_indices=extra_idx,
        extra_topk_length=extra_topk_length,
        mid_out=mid_out,
        mid_lse=mid_lse,
    )

    return (output.unsqueeze(1), None)


def _validate_flashinfer_sparse_mla_backend(
    *,
    model_arch: str,
    device_sm_major: int,
    kv_cache_dtype: torch.dtype,
    prefill_impl: str,
    decode_impl: str,
) -> bool:
    selected = {prefill_impl, decode_impl}
    uses_flashinfer_sparse_mla = "flashinfer_sparse_mla" in selected
    is_glm_sm12_fp8 = (
        model_arch in _GLM_DSA_MODEL_ARCHS
        and device_sm_major == 12
        and kv_cache_dtype == torch.float8_e4m3fn
        and not _is_hip
    )
    if uses_flashinfer_sparse_mla and not is_glm_sm12_fp8:
        raise ValueError(
            "flashinfer_sparse_mla supports only GLM DSA with FP8 KV cache "
            "on NVIDIA SM120/SM121; "
            f"got model_arch={model_arch!r}, sm_major={device_sm_major}, "
            f"kv_cache_dtype={kv_cache_dtype}, prefill_impl={prefill_impl!r}, "
            f"decode_impl={decode_impl!r}."
        )
    if is_glm_sm12_fp8:
        unsupported = selected - {"flashinfer_sparse_mla"}
        if unsupported:
            raise ValueError(
                "GLM DSA with FP8 KV cache on NVIDIA SM120/SM121 supports "
                "only flashinfer_sparse_mla, "
                f"but got {sorted(unsupported)}."
            )
    return uses_flashinfer_sparse_mla


def flashinfer_sparse_mla_forward(
    q: torch.Tensor,
    kv_cache: torch.Tensor,
    indices: torch.Tensor,
    seq_lens: torch.Tensor,
    workspace_buffer: torch.Tensor,
    *,
    page_size: int,
    kv_cache_dim: int,
    qk_nope_head_dim: int,
    kv_lora_rank: int,
    qk_rope_head_dim: int,
    sm_scale: float,
    skip_softmax_threshold_scale_factor: float | None,
) -> torch.Tensor:
    """Run FlashInfer's SM120 sparse MLA kernel on SGLang's packed DSA cache."""
    from flashinfer.mla import trtllm_batch_decode_with_kv_cache_mla

    topk = indices.shape[1]
    result = trtllm_batch_decode_with_kv_cache_mla(
        query=q.unsqueeze(1),
        kv_cache=kv_cache.view(torch.uint8)
        .view(-1, page_size, kv_cache_dim)
        .unsqueeze(1),
        workspace_buffer=workspace_buffer,
        qk_nope_head_dim=qk_nope_head_dim,
        kv_lora_rank=kv_lora_rank,
        qk_rope_head_dim=qk_rope_head_dim,
        block_tables=indices.unsqueeze(1),
        seq_lens=seq_lens,
        max_seq_len=topk,
        sparse_mla_top_k=topk,
        bmm1_scale=float(sm_scale),
        bmm2_scale=1.0,
        kv_scale_format="arbitrary_fp32",
        skip_softmax_threshold_scale_factor=skip_softmax_threshold_scale_factor,
    )
    return result.squeeze(1)