"""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 functools import lru_cache
from typing import FrozenSet, Optional, Tuple

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)


@lru_cache(maxsize=1)
def _flashinfer_dsv4_decode_capabilities() -> Tuple[int, FrozenSet[int]]:
    """Read the installed FlashInfer DSV4 decode capabilities once."""
    try:
        from flashinfer.mla._sparse_mla_sm120 import (
            _DECODE_DSV4_DISPATCH,
            _DECODE_MAX_TOKENS,
        )
    except (AttributeError, ImportError):
        return 0, frozenset()

    return int(_DECODE_MAX_TOKENS), frozenset(
        heads for heads, _ in _DECODE_DSV4_DISPATCH
    )


def flashinfer_dsv4_decode_supports_num_heads(num_heads: int, num_tokens: int) -> bool:
    """Return whether FlashInfer supports this DSV4 decode head count.

    Keep this capability check fail-closed because SGLang can be used with a
    locally installed FlashInfer even though the release dependency is pinned.
    The padded 64-head decode path remains the safe fallback for older builds.
    Prefill head selection is handled separately by the caller.
    """
    decode_max_tokens, supported_heads = _flashinfer_dsv4_decode_capabilities()
    return num_tokens <= decode_max_tokens and num_heads in supported_heads


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)