# SPDX-FileCopyrightText: Copyright (c) 2026 MiniMax
# SPDX-License-Identifier: MIT
"""Sparse attention interface.
Current delivery scope:
- head dimension is supported only for D=128
Public API:
sparse_atten_func(...)
sparse_decode_atten_func(...)
SparseDecodePagedAttentionWrapper
Internal forward core:
_sparse_atten_csr_varlen_forward(...)
Preprocessing (external, done once):
q2k_indices [head_kv, total_q, topK] -> sparse_index_utils.build_k2q_csr()
-> k2q_row_ptr [head_kv, total_rows + 1] int32
-> k2q_q_indices [head_kv, total_q * topK] int32
"""
import math
import os
from typing import Optional
import cutlass
import cutlass.cute as cute
import torch
from cutlass import Float32, Int32
from cutlass.cute.runtime import from_dlpack
from .src.sm100.fwd.combine import combine
from .src.sm100.fwd.atten_fwd import SparseAttentionForwardSm100
from .src.sm100.fwd.atten_fwd_nvfp4_kv import SparseAttentionForwardNvfp4KvSm100
from .src.sm100.prepare_scheduler import (
SparseAttentionSchedule,
prepare_sparse_fwd_schedule_and_split,
)
from .src.sm100.decode_schedule import (
DecodeAttentionSchedule,
prepare_decode_schedule,
)
from .src.common.cute_dsl_utils import to_cute_tensor as to_cute_tensor_kvouter
from .src.common.tma_utils import (
create_q_gather4_tma_desc,
)
_compile_cache: dict = {}
_TEMPERATURE_LSE_FAST_PATH_ABS_TOL = 1e-12
_SUPPORTED_SPARSE_TOPK = (4, 8, 16, 32)
_SUPPORTED_FWD_DTYPES = (torch.bfloat16, torch.float8_e4m3fn)
_SUPPORTED_FWD_MMA_DTYPES = (torch.bfloat16, torch.float8_e4m3fn)
_SUPPORTED_DECODE_QHEAD_PER_KV = 16
def _normalize_partial_dtype(partial_dtype: torch.dtype) -> torch.dtype:
supported = {torch.float32, torch.bfloat16, torch.float16, torch.float8_e4m3fn}
if partial_dtype not in supported:
raise TypeError(
"partial_dtype must be one of torch.float32 / torch.bfloat16 / "
"torch.float16 / torch.float8_e4m3fn, "
f"got {partial_dtype}"
)
return partial_dtype
def _normalize_forward_mma_dtype(dtype: Optional[torch.dtype], fallback: torch.dtype, name: str) -> torch.dtype:
dtype = fallback if dtype is None else dtype
if dtype not in _SUPPORTED_FWD_MMA_DTYPES:
raise TypeError(
f"{name} must be one of torch.bfloat16 / torch.float8_e4m3fn, got {dtype}"
)
return dtype
def _resolve_forward_mma_dtypes(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
qk_dtype: Optional[torch.dtype],
pv_dtype: Optional[torch.dtype],
) -> tuple[torch.dtype, torch.dtype]:
qk_dtype = _normalize_forward_mma_dtype(qk_dtype, q.dtype, "qk_dtype")
if pv_dtype is None:
# Preserve the historical fp8 KV-cache path: BF16 Q with FP8 K/V
# stages both K and V as BF16 compute operands.
if (
q.dtype == torch.bfloat16
and k.dtype == torch.float8_e4m3fn
and v.dtype == torch.float8_e4m3fn
):
pv_dtype = torch.bfloat16
else:
pv_dtype = v.dtype
pv_dtype = _normalize_forward_mma_dtype(pv_dtype, pv_dtype, "pv_dtype")
if q.dtype != qk_dtype:
raise ValueError(
"qk_dtype must match q storage dtype; Q fp8->bf16 staging is not supported"
)
if k.dtype != qk_dtype:
if not (k.dtype == torch.float8_e4m3fn and qk_dtype == torch.bfloat16):
raise ValueError(
"unsupported K storage/qk_dtype combination; only fp8 K -> bf16 QK staging is supported"
)
if v.dtype != pv_dtype:
if not (v.dtype == torch.float8_e4m3fn and pv_dtype == torch.bfloat16):
raise ValueError(
"unsupported V storage/pv_dtype combination; only fp8 V -> bf16 PV staging is supported"
)
return qk_dtype, pv_dtype
def _to_cute_tensor_meta(t: torch.Tensor, assumed_align: int = 4):
tensor = from_dlpack(t.detach(), assumed_align=assumed_align, enable_tvm_ffi=True)
return tensor.mark_layout_dynamic(leading_dim=0)
def _torch_dtype_to_cutlass_dtype(dtype: torch.dtype):
if dtype == torch.bfloat16:
return cutlass.BFloat16
if dtype == torch.float16:
return cutlass.Float16
if dtype == torch.float8_e4m3fn:
return cutlass.Float8E4M3FN
raise TypeError(
f"Only torch.bfloat16, torch.float16, torch.float8_e4m3fn supported, got {dtype}"
)
def _prepare_paged_kv_for_tma(k, v, blk_kv: int):
page_size = int(k.shape[2])
if page_size != blk_kv:
raise ValueError(f"Sparse Page Attention requires page_size == blk_kv, got {page_size} vs {blk_kv}")
return k, v
def _validate_cu_seqlens(
cu_seqlens: torch.Tensor,
*,
name: str,
device: torch.device,
) -> None:
if cu_seqlens.device != device:
raise ValueError(f"{name} must be on the same device as q")
if cu_seqlens.dtype != torch.int32:
raise TypeError(f"{name} must be torch.int32")
if cu_seqlens.ndim != 1:
raise ValueError(f"{name} must have shape [B + 1]")
if cu_seqlens.shape[0] < 1:
raise ValueError(f"{name} must have at least one element")
if not cu_seqlens.is_contiguous():
raise ValueError(f"{name} must be contiguous")
def _csr_row_capacity(k2q_row_ptr: torch.Tensor) -> int:
return int(k2q_row_ptr.shape[1] - 1)
def _validate_csr_varlen_inputs(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
k2q_row_ptr: torch.Tensor,
k2q_q_indices: torch.Tensor,
topK: int,
blk_kv: int,
page_table: Optional[torch.Tensor],
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
seqused_k: Optional[torch.Tensor],
) -> tuple[int, int]:
if q.ndim != 3:
raise ValueError("CSR sparse forward requires q to have shape [total_q, Hq, D]")
if q.dtype not in _SUPPORTED_FWD_DTYPES:
raise TypeError(
"CSR sparse forward supports only torch.bfloat16 and "
f"torch.float8_e4m3fn Q/K/V, got {q.dtype}"
)
if q.device != k.device or q.device != v.device:
raise ValueError("q, k, v must be on the same device")
mixed_fp8_kv_bf16_q = (
q.dtype == torch.bfloat16
and k.dtype == torch.float8_e4m3fn
and v.dtype == torch.float8_e4m3fn
)
if not mixed_fp8_kv_bf16_q and (q.dtype != k.dtype or q.dtype != v.dtype):
raise ValueError(
"q, k, v must have the same dtype, except q=bf16 with fp8_e4m3 K/V cache"
)
if q.shape[-1] != k.shape[-1] or q.shape[-1] != v.shape[-1]:
raise ValueError("q, k, v must have the same head dimension")
dim = q.shape[-1]
if dim != 128:
raise NotImplementedError(
f"CSR sparse forward currently supports only D=128, got D={dim}"
)
if page_table is None:
if k.shape[-2] != v.shape[-2] or k.shape[-1] != v.shape[-1]:
raise ValueError("k and v must have the same [Hkv, D] tail dimensions")
head_kv = k.shape[-2]
else:
if k.ndim != 4 or v.ndim != 4:
raise ValueError(
"Sparse Page Attention requires k and v to have shape "
"[num_pages, Hkv, page_size, D]"
)
if k.shape[1] != v.shape[1] or k.shape[-1] != v.shape[-1]:
raise ValueError(
"Sparse Page Attention k and v must have the same Hkv and D"
)
head_kv = k.shape[1]
if (
q.device != k2q_row_ptr.device
or q.device != k2q_q_indices.device
):
raise ValueError("CSR metadata must be on the same device as q")
if (
k2q_row_ptr.dtype != torch.int32
or k2q_q_indices.dtype != torch.int32
):
raise TypeError("CSR metadata tensors must be torch.int32")
if k2q_row_ptr.ndim != 2 or k2q_q_indices.ndim != 2:
raise ValueError("k2q_row_ptr and k2q_q_indices must be rank-2")
_validate_cu_seqlens(cu_seqlens_q, name="cu_seqlens_q", device=q.device)
_validate_cu_seqlens(cu_seqlens_k, name="cu_seqlens_k", device=q.device)
if cu_seqlens_k.shape != cu_seqlens_q.shape:
raise ValueError("cu_seqlens_k must have shape [B + 1] matching cu_seqlens_q")
batch = int(cu_seqlens_q.shape[0] - 1)
total_q = q.shape[0]
head_q = q.shape[1]
if head_q % head_kv != 0:
raise ValueError("q.shape[1] must be divisible by Hkv")
qhead_per_kv = head_q // head_kv
if qhead_per_kv not in (1, 2, 4, 8, 16):
raise NotImplementedError(
"CSR forward is currently supported only for qhead_per_kv in {1, 2, 4, 8, 16}"
)
if k2q_row_ptr.shape[0] != head_kv or k2q_q_indices.shape[0] != head_kv:
raise ValueError("CSR metadata head dimension must match KV head count")
if k2q_q_indices.shape[1] < total_q * topK:
raise ValueError(
f"k2q_q_indices.shape[1] ({k2q_q_indices.shape[1]}) must be >= total_q * topK ({total_q * topK})"
)
if k2q_row_ptr.shape[1] < 1:
raise ValueError("k2q_row_ptr must contain at least one row pointer column")
if page_table is None:
if seqused_k is not None:
raise ValueError("seqused_k is only supported together with page_table")
total_k = k.shape[0]
if k.ndim != 3 or v.ndim != 3:
raise ValueError("Sparse Attention requires k and v to have shape [total_k, Hkv, D]")
if k.shape != (total_k, head_kv, q.shape[-1]) or v.shape != (total_k, head_kv, q.shape[-1]):
raise ValueError("Sparse Attention k and v must match [total_k, Hkv, D]")
else:
if page_table.device != q.device:
raise ValueError("page_table must be on the same device as q")
if page_table.dtype != torch.int32:
raise TypeError("page_table must be torch.int32")
if page_table.ndim != 2 or page_table.shape[0] != batch:
raise ValueError("page_table must have shape [B, max_num_pages_per_seq]")
if page_table.stride(-1) != 1:
raise ValueError("page_table must be contiguous in the last dimension")
if k.ndim != 4 or v.ndim != 4:
raise ValueError(
"Sparse Page Attention requires k and v to have shape "
"[num_pages, Hkv, page_size, D]"
)
if k.shape != v.shape:
raise ValueError(f"k and v must have the same shape, got {k.shape} and {v.shape}")
if k.shape[1] != head_kv or k.shape[3] != q.shape[-1]:
raise ValueError(
"Sparse Page Attention k and v must match "
"[num_pages, Hkv, page_size, D]"
)
page_size = int(k.shape[2])
if page_size != blk_kv:
raise ValueError(
f"Unsupported Sparse Page Attention page_size={page_size} for blk_kv={blk_kv}; "
"require page_size == blk_kv"
)
if seqused_k is not None:
if seqused_k.device != q.device:
raise ValueError("seqused_k must be on the same device as q")
if seqused_k.dtype != torch.int32:
raise TypeError("seqused_k must be torch.int32")
if seqused_k.shape != (batch,):
raise ValueError("seqused_k must have shape [B]")
if not seqused_k.is_contiguous():
raise ValueError("seqused_k must be contiguous")
if topK not in _SUPPORTED_SPARSE_TOPK:
raise ValueError(
f"CSR sparse forward supports topK in {_SUPPORTED_SPARSE_TOPK}, got {topK}"
)
return batch, head_kv
def _validate_csr_varlen_nvfp4_kv_inputs(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
k_scale_128x4: torch.Tensor,
v_scale_128x4: torch.Tensor,
k_global_scale: Optional[torch.Tensor],
v_global_scale: Optional[torch.Tensor],
k2q_row_ptr: torch.Tensor,
k2q_q_indices: torch.Tensor,
topK: int,
blk_kv: int,
page_table: Optional[torch.Tensor],
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
seqused_k: Optional[torch.Tensor],
) -> tuple[int, int]:
if q.ndim != 3:
raise ValueError("KVFP4 CSR sparse forward requires q to have shape [total_q, Hq, D]")
if q.dtype not in (torch.bfloat16, torch.float8_e4m3fn):
raise TypeError(f"KVFP4 CSR sparse forward requires BF16 or FP8 E4M3 q, got {q.dtype}")
if q.shape[-1] != 128:
raise NotImplementedError(
f"KVFP4 CSR sparse forward currently supports only D=128, got {q.shape[-1]}"
)
if k.dtype != torch.uint8 or v.dtype != torch.uint8:
raise TypeError(f"KVFP4 k/v must be torch.uint8, got {k.dtype} and {v.dtype}")
if k_scale_128x4.dtype != torch.uint8 or v_scale_128x4.dtype != torch.uint8:
raise TypeError(
"KVFP4 block scales must be torch.uint8 E4M3 tensors, got "
f"{k_scale_128x4.dtype} and {v_scale_128x4.dtype}"
)
if k_global_scale is not None and k_global_scale.dtype != torch.float32:
raise TypeError("KVFP4 K global scale must be a torch.float32 tensor or None")
if v_global_scale is not None and v_global_scale.dtype != torch.float32:
raise TypeError("KVFP4 V global scale must be a torch.float32 tensor or None")
tensors = (
k,
v,
k_scale_128x4,
v_scale_128x4,
k2q_row_ptr,
k2q_q_indices,
cu_seqlens_q,
cu_seqlens_k,
)
optional_tensors = tuple(t for t in (k_global_scale, v_global_scale) if t is not None)
if any(t.device != q.device for t in tensors + optional_tensors):
raise ValueError("KVFP4 inputs and metadata must be on the same device as q")
if k.shape != v.shape:
raise ValueError(f"KVFP4 k and v must have the same shape, got {k.shape} and {v.shape}")
packed_dim = q.shape[-1] // 2
scale_cols = q.shape[-1] // 16
if k_scale_128x4.ndim != 2 or v_scale_128x4.ndim != 2:
raise ValueError("KVFP4 block scales must be rank-2 128x4 tiled tensors")
if k_scale_128x4.shape[1] < scale_cols or v_scale_128x4.shape[1] < scale_cols:
raise ValueError(
"KVFP4 block scales must have at least D/16 columns; "
f"need {scale_cols}, got {k_scale_128x4.shape[1]} and {v_scale_128x4.shape[1]}"
)
if k_global_scale is not None and k_global_scale.numel() < 1:
raise ValueError("KVFP4 K global scale must contain at least one element")
if v_global_scale is not None and v_global_scale.numel() < 1:
raise ValueError("KVFP4 V global scale must contain at least one element")
if page_table is None:
if seqused_k is not None:
raise ValueError("seqused_k is only supported together with page_table")
if k.ndim != 3:
raise ValueError("KVFP4 Sparse Attention requires k/v shape [total_k, Hkv, D/2]")
if k.shape[-1] != packed_dim:
raise ValueError(f"KVFP4 packed K/V last dimension must be D/2={packed_dim}")
total_k = int(k.shape[0])
head_kv = int(k.shape[1])
required_scale_rows = total_k * head_kv
else:
if k.ndim != 4:
raise ValueError(
"KVFP4 Sparse Page Attention requires k/v shape "
"[num_pages, Hkv, page_size, D/2]"
)
if k.shape[-1] != packed_dim:
raise ValueError(f"KVFP4 packed K/V last dimension must be D/2={packed_dim}")
page_size = int(k.shape[2])
if page_size != int(blk_kv):
raise ValueError(
f"KVFP4 Sparse Page Attention requires page_size == blk_kv, got {page_size} vs {blk_kv}"
)
head_kv = int(k.shape[1])
required_scale_rows = int(k.shape[0]) * head_kv * page_size
if page_table.device != q.device:
raise ValueError("page_table must be on the same device as q")
if page_table.dtype != torch.int32:
raise TypeError("page_table must be torch.int32")
if page_table.ndim != 2:
raise ValueError("page_table must have shape [B, max_num_pages_per_seq]")
if page_table.stride(-1) != 1:
raise ValueError("page_table must be contiguous in the last dimension")
if seqused_k is not None:
if seqused_k.device != q.device:
raise ValueError("seqused_k must be on the same device as q")
if seqused_k.dtype != torch.int32:
raise TypeError("seqused_k must be torch.int32")
if not seqused_k.is_contiguous():
raise ValueError("seqused_k must be contiguous")
padded_scale_rows = ((required_scale_rows + 127) // 128) * 128
padded_scale_cols = ((scale_cols + 3) // 4) * 4
for name, scale in (("k_scale_128x4", k_scale_128x4), ("v_scale_128x4", v_scale_128x4)):
if scale.shape[0] < padded_scale_rows or scale.shape[1] < padded_scale_cols:
raise ValueError(
f"{name} is too small for 128x4 layout: got {tuple(scale.shape)}, "
f"need at least {(padded_scale_rows, padded_scale_cols)}"
)
if k2q_row_ptr.device != q.device or k2q_q_indices.device != q.device:
raise ValueError("CSR metadata must be on the same device as q")
if k2q_row_ptr.dtype != torch.int32 or k2q_q_indices.dtype != torch.int32:
raise TypeError("CSR metadata tensors must be torch.int32")
if k2q_row_ptr.ndim != 2 or k2q_q_indices.ndim != 2:
raise ValueError("k2q_row_ptr and k2q_q_indices must be rank-2")
_validate_cu_seqlens(cu_seqlens_q, name="cu_seqlens_q", device=q.device)
_validate_cu_seqlens(cu_seqlens_k, name="cu_seqlens_k", device=q.device)
if cu_seqlens_k.shape != cu_seqlens_q.shape:
raise ValueError("cu_seqlens_k must have shape [B + 1] matching cu_seqlens_q")
batch = int(cu_seqlens_q.shape[0] - 1)
if page_table is not None and page_table.shape[0] != batch:
raise ValueError("page_table must have shape [B, max_num_pages_per_seq]")
if seqused_k is not None and seqused_k.shape != (batch,):
raise ValueError("seqused_k must have shape [B]")
head_q = int(q.shape[1])
if head_q % head_kv != 0:
raise ValueError("q.shape[1] must be divisible by Hkv")
qhead_per_kv = head_q // head_kv
if qhead_per_kv not in (1, 2, 4, 8, 16):
raise NotImplementedError(
"KVFP4 CSR forward is currently supported only for qhead_per_kv in {1, 2, 4, 8, 16}"
)
if k2q_row_ptr.shape[0] != head_kv or k2q_q_indices.shape[0] != head_kv:
raise ValueError("CSR metadata head dimension must match KV head count")
if k2q_q_indices.shape[1] < q.shape[0] * topK:
raise ValueError(
f"k2q_q_indices.shape[1] ({k2q_q_indices.shape[1]}) must be >= total_q * topK ({q.shape[0] * topK})"
)
if k2q_row_ptr.shape[1] < 1:
raise ValueError("k2q_row_ptr must contain at least one row pointer column")
if topK not in _SUPPORTED_SPARSE_TOPK:
raise ValueError(
f"KVFP4 CSR sparse forward supports topK in {_SUPPORTED_SPARSE_TOPK}, got {topK}"
)
return batch, head_kv
def _validate_sparse_decode_inputs(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
q2k_indices: Optional[torch.Tensor],
*,
page_table: torch.Tensor,
seqused_k: torch.Tensor,
seqlen_q: int,
max_seqlen_k: int,
blk_kv: int,
causal: bool,
) -> tuple[int, int]:
if q.ndim != 3:
raise ValueError("decode attention requires q to have shape [total_q, Hq, D]")
if k.ndim != 4 or v.ndim != 4:
raise ValueError(
"decode attention requires paged k/v with shape [num_pages, Hkv, page_size, D]"
)
if q.device != k.device or q.device != v.device:
raise ValueError("decode q, k, and v must be on the same device")
if q.dtype != torch.float8_e4m3fn or k.dtype != q.dtype or v.dtype != q.dtype:
raise TypeError(
"decode attention currently supports only torch.float8_e4m3fn Q/K/V"
)
if k.shape != v.shape:
raise ValueError(f"decode k and v must have the same shape, got {k.shape} and {v.shape}")
if q.shape[-1] != 128 or k.shape[-1] != 128:
raise NotImplementedError(
f"decode attention currently supports only D=128, got q={q.shape[-1]} k={k.shape[-1]}"
)
if not bool(causal):
raise NotImplementedError("decode attention currently supports only causal=True")
page_size = int(k.shape[2])
if page_size != int(blk_kv):
raise ValueError(f"decode attention requires page_size == blk_kv, got {page_size} vs {blk_kv}")
head_kv = int(k.shape[1])
head_q = int(q.shape[1])
if head_q % head_kv != 0:
raise ValueError("decode q.shape[1] must be divisible by Hkv")
qhead_per_kv = head_q // head_kv
if qhead_per_kv != _SUPPORTED_DECODE_QHEAD_PER_KV:
raise NotImplementedError(
"decode attention currently supports only "
f"qhead_per_kv={_SUPPORTED_DECODE_QHEAD_PER_KV}, got {qhead_per_kv}"
)
if page_table is None:
raise ValueError("decode attention requires page_table")
if page_table.device != q.device:
raise ValueError("decode page_table must be on the same device as q")
if page_table.dtype != torch.int32:
raise TypeError("decode page_table must be torch.int32")
if page_table.ndim != 2:
raise ValueError("decode page_table must have shape [B, max_num_pages_per_seq]")
batch = int(page_table.shape[0])
if page_table.stride(-1) != 1:
raise ValueError("decode page_table must be contiguous in the last dimension")
if seqused_k is None:
raise ValueError("decode attention requires seqused_k")
if seqused_k.device != q.device:
raise ValueError("decode seqused_k must be on the same device as q")
if seqused_k.dtype != torch.int32:
raise TypeError("decode seqused_k must be torch.int32")
if seqused_k.shape != (batch,):
raise ValueError("decode seqused_k must have shape [B]")
if not seqused_k.is_contiguous():
raise ValueError("decode seqused_k must be contiguous")
seqlen_q = int(seqlen_q)
max_seqlen_k = int(max_seqlen_k)
if seqlen_q <= 0 or max_seqlen_k <= 0:
raise ValueError("decode seqlen_q and max_seqlen_k must be positive")
if int(q.shape[0]) != batch * seqlen_q:
raise ValueError("decode q.shape[0] must equal batch * seqlen_q")
if q2k_indices is not None:
if q2k_indices.device != q.device:
raise ValueError("decode q2k_indices must be on the same device as q")
if q2k_indices.dtype != torch.int32:
raise TypeError("decode q2k_indices must be torch.int32")
if q2k_indices.ndim != 3:
raise ValueError("decode q2k_indices must have shape [Hkv, total_q, topK]")
if q2k_indices.shape[0] != head_kv or q2k_indices.shape[1] != q.shape[0]:
raise ValueError("decode q2k_indices must match [Hkv, total_q, topK]")
if not q2k_indices.is_contiguous():
raise ValueError("decode q2k_indices must be contiguous")
return batch, head_kv
def _validate_schedule_common(
schedule: SparseAttentionSchedule,
*,
device: torch.device,
) -> None:
if schedule.scheduler_metadata is None:
raise ValueError("schedule.scheduler_metadata is required")
if schedule.work_count is None:
raise ValueError("schedule.work_count is required")
metadata = schedule.scheduler_metadata
work_count = schedule.work_count
if metadata.device != device or work_count.device != device:
raise ValueError("schedule tensors must be on the same device as q")
if metadata.dtype != torch.int32 or work_count.dtype != torch.int32:
raise TypeError("schedule.scheduler_metadata and schedule.work_count must be torch.int32")
if metadata.ndim != 2 or metadata.shape[1] != 6:
raise ValueError("schedule.scheduler_metadata must have shape [capacity, 6]")
if work_count.shape != (1,):
raise ValueError("schedule.work_count must have shape [1]")
if not metadata.is_contiguous() or not work_count.is_contiguous():
raise ValueError("schedule.scheduler_metadata and schedule.work_count must be contiguous")
def _validate_fwd_schedule(
schedule: SparseAttentionSchedule,
*,
q: torch.Tensor,
k2q_q_indices: torch.Tensor,
head_kv: int,
) -> None:
_validate_schedule_common(schedule, device=q.device)
if schedule.qsplit_indices is None:
raise ValueError("schedule.qsplit_indices is required for forward")
if schedule.split_counts is None:
raise ValueError("schedule.split_counts is required for forward")
qsplit = schedule.qsplit_indices
split_counts = schedule.split_counts
if qsplit.device != q.device or split_counts.device != q.device:
raise ValueError("forward schedule tensors must be on the same device as q")
if qsplit.dtype != torch.int32 or split_counts.dtype != torch.int32:
raise TypeError("schedule.qsplit_indices and schedule.split_counts must be torch.int32")
if qsplit.shape != k2q_q_indices.shape:
raise ValueError("schedule.qsplit_indices shape must match k2q_q_indices")
total_q = q.shape[0]
if split_counts.shape != (total_q, head_kv):
raise ValueError(
"schedule.split_counts must have shape "
f"({total_q}, {head_kv}), got {tuple(split_counts.shape)}"
)
if not qsplit.is_contiguous() or not split_counts.is_contiguous():
raise ValueError("schedule.qsplit_indices and schedule.split_counts must be contiguous")
def sparse_atten_func(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
k2q_row_ptr: torch.Tensor,
k2q_q_indices: torch.Tensor,
topK: int,
*,
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
max_seqlen_q: int,
max_seqlen_k: int,
blk_kv: int = 128,
causal: bool = False,
softmax_scale: Optional[float] = None,
lse_temperature_scale: float = 1.0,
return_temperature_lse: bool = False,
partial_dtype: torch.dtype = torch.bfloat16,
return_softmax_lse: bool = False,
page_table: Optional[torch.Tensor] = None,
seqused_k: Optional[torch.Tensor] = None,
schedule: Optional[SparseAttentionSchedule] = None,
usable_SM_count: int = -1,
qk_dtype: Optional[torch.dtype] = None,
pv_dtype: Optional[torch.dtype] = None,
):
"""Run SM100 CSR block-sparse varlen attention.
This is the public forward-only sparse attention API. It consumes
query-to-key block selections converted to CSR metadata by
``build_k2q_csr`` and supports both dense KV layout and paged KV layout.
Parameters
----------
q : torch.Tensor
Shape ``[total_q, Hq, 128]`` on CUDA. Supported dtypes are BF16 and
FP8 E4M3.
k : torch.Tensor
Dense layout ``[total_k, Hkv, 128]`` or paged layout
``[num_pages, Hkv, blk_kv, 128]``. For BF16 Q with FP8 K/V cache, K
may be FP8 E4M3 while QK compute uses BF16 staging.
v : torch.Tensor
Same layout and head count as ``k``.
k2q_row_ptr : torch.Tensor
CSR row pointers with shape ``[Hkv, total_rows + 1]`` and dtype int32.
k2q_q_indices : torch.Tensor
CSR query indices with shape ``[Hkv, >= total_q * topK]`` and dtype
int32.
topK : int
Number of selected KV blocks per query. Supported values are
``4, 8, 16, 32``.
cu_seqlens_q : torch.Tensor
Shape ``[batch_size + 1]``, dtype int32. Prefix sums of Q lengths.
cu_seqlens_k : torch.Tensor
Shape ``[batch_size + 1]``, dtype int32. Prefix sums of KV lengths.
max_seqlen_q : int
Maximum Q sequence length in the batch.
max_seqlen_k : int
Maximum KV sequence length in the batch.
blk_kv : int, optional
KV block size. Paged KV requires ``k.shape[2] == blk_kv``.
causal : bool, optional
Whether to apply causal masking.
softmax_scale : float, optional
Softmax scale. Defaults to ``1 / sqrt(128)``.
lse_temperature_scale : float, optional
Extra divisor used only for temperature-scaled LSE output.
return_temperature_lse : bool, optional
If True, also return LSE computed with logits scaled by
``softmax_scale / lse_temperature_scale``. Requires
``return_softmax_lse=True``.
partial_dtype : torch.dtype, optional
Accumulation dtype for per-block partial O. Supported values are
FP32, BF16, FP16, and FP8 E4M3.
return_softmax_lse : bool, optional
If True, return ``(out, softmax_lse)`` or
``(out, softmax_lse, temperature_lse)``.
page_table : torch.Tensor, optional
Paged-KV physical page table with shape
``[batch_size, max_num_pages_per_seq]`` and dtype int32.
seqused_k : torch.Tensor, optional
Shape ``[batch_size]``, dtype int32. Effective KV length per request
for paged causal attention.
schedule : SparseAttentionSchedule, optional
Prebuilt sparse forward schedule. If omitted, the schedule is built
during the call.
usable_SM_count : int, optional
Maximum number of SMs used by the scheduler. ``-1`` uses all SMs.
qk_dtype : torch.dtype, optional
Compile-time MMA operand dtype for QK. Defaults to Q storage dtype,
except supported FP8 K/V cache staging modes.
pv_dtype : torch.dtype, optional
Compile-time MMA operand dtype for PV. Defaults to V storage dtype,
except supported FP8 K/V cache staging modes.
Returns
-------
torch.Tensor or tuple[torch.Tensor, torch.Tensor]
Output shape ``[total_q, Hq, 128]`` with BF16 dtype. Optional LSE
outputs have shape ``[total_q, Hq]`` and dtype float32.
Notes
-----
``Hq / Hkv`` must be one of ``1, 2, 4, 8, 16``. Current kernels support
head dimension 128 only.
"""
if softmax_scale is None:
softmax_scale = q.shape[-1] ** -0.5
lse_temperature_scale = float(lse_temperature_scale)
if not math.isfinite(lse_temperature_scale) or lse_temperature_scale <= 0.0:
raise ValueError(
f"lse_temperature_scale must be finite and > 0, got {lse_temperature_scale}"
)
return_temperature_lse = bool(return_temperature_lse)
if return_temperature_lse and not return_softmax_lse:
raise ValueError("return_temperature_lse=True requires return_softmax_lse=True")
partial_dtype = _normalize_partial_dtype(partial_dtype)
qk_dtype, pv_dtype = _resolve_forward_mma_dtypes(q, k, v, qk_dtype, pv_dtype)
if cu_seqlens_q is None or cu_seqlens_k is None:
raise ValueError(
"sparse_atten_func requires CSR varlen metadata: pass cu_seqlens_q and cu_seqlens_k"
)
batch, head_kv = _validate_csr_varlen_inputs(
q,
k,
v,
k2q_row_ptr,
k2q_q_indices,
topK,
blk_kv,
page_table,
cu_seqlens_q,
cu_seqlens_k,
seqused_k,
)
max_seqlen_q = int(max_seqlen_q)
max_seqlen_k = int(max_seqlen_k)
return _sparse_atten_csr_varlen_forward(
q.contiguous(),
k.contiguous(),
v.contiguous(),
k2q_row_ptr.contiguous(),
k2q_q_indices.contiguous(),
int(topK),
int(blk_kv),
bool(causal),
float(softmax_scale),
lse_temperature_scale,
return_temperature_lse,
partial_dtype,
bool(return_softmax_lse),
cu_seqlens_q.contiguous(),
cu_seqlens_k.contiguous(),
None if page_table is None else page_table.contiguous(),
None if seqused_k is None else seqused_k.contiguous(),
schedule,
int(usable_SM_count),
int(batch),
int(head_kv),
int(max_seqlen_q),
int(max_seqlen_k),
qk_dtype,
pv_dtype,
)
def sparse_atten_nvfp4_kv_func(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
k_scale_128x4: torch.Tensor,
v_scale_128x4: torch.Tensor,
k_global_scale: Optional[torch.Tensor],
v_global_scale: Optional[torch.Tensor],
k2q_row_ptr: torch.Tensor,
k2q_q_indices: torch.Tensor,
topK: int,
*,
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
max_seqlen_q: int,
max_seqlen_k: int,
blk_kv: int = 128,
causal: bool = False,
softmax_scale: Optional[float] = None,
lse_temperature_scale: float = 1.0,
return_temperature_lse: bool = False,
partial_dtype: torch.dtype = torch.bfloat16,
return_softmax_lse: bool = False,
page_table: Optional[torch.Tensor] = None,
seqused_k: Optional[torch.Tensor] = None,
schedule: Optional[SparseAttentionSchedule] = None,
):
"""Run SM100 CSR sparse attention with packed NVFP4 K/V.
Parameters
----------
q : torch.Tensor
Shape ``[total_q, Hq, 128]`` on CUDA. Supported dtypes are BF16 and
FP8 E4M3.
k : torch.Tensor
Packed NVFP4 K data. Dense layout is ``[total_k, Hkv, 64]``; paged
layout is ``[num_pages, Hkv, blk_kv, 64]``. Dtype must be uint8
because each byte packs two FP4 values.
v : torch.Tensor
Packed NVFP4 V data with the same shape as ``k``.
k_scale_128x4 : torch.Tensor
K block scales in cuBLAS/cuDNN 128x4 tiled storage. Dtype uint8
containing FP8 E4M3 scale values.
v_scale_128x4 : torch.Tensor
V block scales in the same 128x4 tiled storage.
k_global_scale : torch.Tensor, optional
FP32 tensor/global dequant scale for K. May be ``None``.
v_global_scale : torch.Tensor, optional
FP32 tensor/global dequant scale for V. May be ``None``. The V global
scale is applied in the combine stage.
k2q_row_ptr : torch.Tensor
CSR row pointers with shape ``[Hkv, total_rows + 1]`` and dtype int32.
k2q_q_indices : torch.Tensor
CSR query indices with shape ``[Hkv, >= total_q * topK]`` and dtype
int32.
topK : int
Number of selected KV blocks per query. Supported values are
``4, 8, 16, 32``.
cu_seqlens_q, cu_seqlens_k : torch.Tensor
Shape ``[batch_size + 1]``, dtype int32. Prefix sums of Q and KV
lengths.
max_seqlen_q, max_seqlen_k : int
Maximum Q and KV sequence lengths in the batch.
blk_kv : int, optional
KV block/page size. Paged KV requires ``k.shape[2] == blk_kv``.
causal : bool, optional
Whether to apply causal masking.
softmax_scale : float, optional
Softmax scale. Defaults to ``1 / sqrt(128)``.
lse_temperature_scale : float, optional
Extra divisor used only for temperature-scaled LSE output.
return_temperature_lse : bool, optional
If True, also return temperature-scaled LSE. Requires
``return_softmax_lse=True``.
partial_dtype : torch.dtype, optional
Accumulation dtype for per-block partial O.
return_softmax_lse : bool, optional
If True, return LSE together with the output.
page_table : torch.Tensor, optional
Paged-KV physical page table with shape
``[batch_size, max_num_pages_per_seq]`` and dtype int32.
seqused_k : torch.Tensor, optional
Effective KV length per request for paged causal attention.
schedule : SparseAttentionSchedule, optional
Prebuilt sparse forward schedule.
Returns
-------
torch.Tensor or tuple[torch.Tensor, torch.Tensor]
Output shape ``[total_q, Hq, 128]`` with BF16 dtype. Optional LSE
outputs have shape ``[total_q, Hq]`` and dtype float32.
"""
if softmax_scale is None:
softmax_scale = q.shape[-1] ** -0.5
lse_temperature_scale = float(lse_temperature_scale)
if not math.isfinite(lse_temperature_scale) or lse_temperature_scale <= 0.0:
raise ValueError(
f"lse_temperature_scale must be finite and > 0, got {lse_temperature_scale}"
)
return_temperature_lse = bool(return_temperature_lse)
if return_temperature_lse and not return_softmax_lse:
raise ValueError("return_temperature_lse=True requires return_softmax_lse=True")
partial_dtype = _normalize_partial_dtype(partial_dtype)
if cu_seqlens_q is None or cu_seqlens_k is None:
raise ValueError(
"sparse_atten_nvfp4_kv_func requires CSR varlen metadata: pass cu_seqlens_q and cu_seqlens_k"
)
batch, head_kv = _validate_csr_varlen_nvfp4_kv_inputs(
q,
k,
v,
k_scale_128x4,
v_scale_128x4,
k_global_scale,
v_global_scale,
k2q_row_ptr,
k2q_q_indices,
topK,
blk_kv,
page_table,
cu_seqlens_q,
cu_seqlens_k,
seqused_k,
)
total_q, head_q, dim = q.shape
max_num_kv_blocks = _csr_row_capacity(k2q_row_ptr)
temperature_lse_fast_path = (
return_temperature_lse
and math.isclose(
float(lse_temperature_scale),
1.0,
rel_tol=0.0,
abs_tol=_TEMPERATURE_LSE_FAST_PATH_ABS_TOL,
)
)
kernel_return_temperature_lse = (
return_temperature_lse and not temperature_lse_fast_path
)
O_partial = torch.empty(
topK, total_q, head_q, dim, dtype=partial_dtype, device=q.device
)
LSE_partial = torch.empty(
topK, total_q, head_q, dtype=torch.float32, device=q.device
)
LSE_temperature_partial = (
torch.empty(topK, total_q, head_q, dtype=torch.float32, device=q.device)
if kernel_return_temperature_lse
else None
)
O_out = torch.empty(total_q, head_q, dim, dtype=torch.bfloat16, device=q.device)
LSE_out = torch.empty(total_q, head_q, dtype=torch.float32, device=q.device)
LSE_temperature_out = (
torch.empty_like(LSE_out) if kernel_return_temperature_lse else None
)
if schedule is None:
k2q_qsplit_indices = torch.empty_like(k2q_q_indices)
split_counts = torch.zeros(
(total_q, head_kv),
dtype=torch.int32,
device=q.device,
)
else:
_validate_fwd_schedule(
schedule,
q=q,
k2q_q_indices=k2q_q_indices,
head_kv=head_kv,
)
k2q_qsplit_indices = schedule.qsplit_indices
split_counts = schedule.split_counts
schedule = _call_sparse_forward_sm100_csr_varlen_nvfp4_kv(
q.contiguous(),
k.contiguous(),
v.contiguous(),
k_scale_128x4.contiguous(),
v_scale_128x4.contiguous(),
None if k_global_scale is None else k_global_scale.contiguous(),
None if v_global_scale is None else v_global_scale.contiguous(),
k2q_row_ptr.contiguous(),
k2q_q_indices.contiguous(),
k2q_qsplit_indices.contiguous(),
split_counts.contiguous(),
cu_seqlens_q.contiguous(),
cu_seqlens_k.contiguous(),
None if page_table is None else page_table.contiguous(),
None if seqused_k is None else seqused_k.contiguous(),
O_partial,
LSE_partial,
LSE_temperature_partial,
float(softmax_scale),
lse_temperature_scale,
kernel_return_temperature_lse,
max_num_kv_blocks,
int(blk_kv),
head_kv,
int(max_seqlen_q),
causal=bool(causal),
schedule=schedule,
)
combine(
O_partial,
LSE_partial,
O_out,
LSE_out,
lse_temperature_partial=LSE_temperature_partial,
lse_temperature_out=LSE_temperature_out,
cu_seqlens=cu_seqlens_q,
split_counts=split_counts,
output_scale=v_global_scale,
use_pdl=True,
)
if temperature_lse_fast_path:
LSE_temperature_out = LSE_out
if return_softmax_lse:
if return_temperature_lse:
return O_out, LSE_out, LSE_temperature_out
return O_out, LSE_out
return O_out
def sparse_decode_atten_func(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
q2k_indices: Optional[torch.Tensor] = None,
*,
page_table: torch.Tensor,
seqused_k: torch.Tensor,
seqlen_q: int,
max_seqlen_k: int,
blk_kv: int = 128,
causal: bool = True,
softmax_scale: Optional[float] = None,
return_softmax_lse: bool = False,
schedule: Optional[DecodeAttentionSchedule] = None,
O_partial: Optional[torch.Tensor] = None,
LSE_partial: Optional[torch.Tensor] = None,
):
"""Run forward-only paged FP8 decode attention.
Parameters
----------
q : torch.Tensor
Shape ``[batch_size * seqlen_q, Hq, 128]``. Dtype must be FP8 E4M3.
k : torch.Tensor
Paged K cache with shape ``[num_pages, Hkv, blk_kv, 128]`` and FP8
E4M3 dtype.
v : torch.Tensor
Paged V cache with the same shape and dtype as ``k``.
q2k_indices : torch.Tensor, optional
Sparse selected KV blocks with shape ``[Hkv, total_q, topK]`` and dtype
int32. ``None`` selects the dense all-KV decode path.
page_table : torch.Tensor
Physical page table with shape ``[batch_size, max_num_pages_per_seq]``
and dtype int32.
seqused_k : torch.Tensor
Shape ``[batch_size]``, dtype int32. Effective KV length per request.
seqlen_q : int
Uniform query length per request. Ragged Q lengths should use prefill
or append paths instead.
max_seqlen_k : int
Maximum KV sequence length in the batch.
blk_kv : int, optional
Page size. Must match ``k.shape[2]``.
causal : bool, optional
Whether to apply causal masking. Current decode kernel requires True.
softmax_scale : float, optional
Softmax scale. Defaults to ``1 / sqrt(128)``.
return_softmax_lse : bool, optional
If True, return ``(out, lse)``.
schedule : DecodeAttentionSchedule, optional
Prebuilt decode schedule.
O_partial, LSE_partial : torch.Tensor, optional
Optional split-KV partial workspaces. Normally owned by
``SparseDecodePagedAttentionWrapper``.
Returns
-------
torch.Tensor or tuple[torch.Tensor, torch.Tensor]
BF16 output with shape ``q.shape``. Optional LSE has shape
``[batch_size * seqlen_q, Hq]`` and dtype float32.
"""
if softmax_scale is None:
softmax_scale = q.shape[-1] ** -0.5
batch, head_kv = _validate_sparse_decode_inputs(
q,
k,
v,
q2k_indices,
page_table=page_table,
seqused_k=seqused_k,
seqlen_q=seqlen_q,
max_seqlen_k=max_seqlen_k,
blk_kv=blk_kv,
causal=causal,
)
head_q = int(q.shape[1])
head_dim = int(q.shape[2])
if schedule is None:
schedule = prepare_decode_schedule(
seqused_k=seqused_k.contiguous(),
page_size=int(blk_kv),
seqlen_q=int(seqlen_q),
num_qo_heads=head_q,
num_kv_heads=head_kv,
head_dim=head_dim,
max_seqlen_k=int(max_seqlen_k),
)
if schedule.split_kv:
if O_partial is None:
O_partial = torch.empty(
(schedule.partial_rows, head_q, head_dim),
dtype=torch.float32,
device=q.device,
)
if LSE_partial is None:
LSE_partial = torch.empty(
(schedule.partial_rows, head_q),
dtype=torch.float32,
device=q.device,
)
out = torch.empty(q.shape, dtype=torch.bfloat16, device=q.device)
lse = torch.empty(
q.shape[:2] if (return_softmax_lse or schedule.split_kv) else (1, head_q),
dtype=torch.float32,
device=q.device,
)
_call_sparse_decode_forward_sm100_paged_fp8(
q.contiguous(),
k.contiguous(),
v.contiguous(),
None if q2k_indices is None else q2k_indices.contiguous(),
page_table.contiguous(),
seqused_k.contiguous(),
out,
lse,
schedule,
O_partial,
LSE_partial,
softmax_scale=float(softmax_scale),
seqlen_q=int(seqlen_q),
max_seqlen_k=int(max_seqlen_k),
blk_kv=int(blk_kv),
causal=bool(causal),
return_lse=bool(return_softmax_lse),
)
if return_softmax_lse:
return out, lse
return out
class SparseDecodePagedAttentionWrapper:
"""Plan/run helper for paged FP8 decode attention.
Use this wrapper when the same page table shape and sequence metadata are
reused across multiple decode layers. ``plan`` validates metadata and
allocates persistent schedules/workspaces; ``run`` then launches the decode
kernel with lower per-call overhead than ``sparse_decode_atten_func``.
"""
def __init__(self, *, blk_kv: int = 128, causal: bool = True):
self.blk_kv = int(blk_kv)
self.causal = bool(causal)
self.batch: Optional[int] = None
self.num_qo_heads: Optional[int] = None
self.num_kv_heads: Optional[int] = None
self.head_dim: Optional[int] = None
self.page_table: Optional[torch.Tensor] = None
self.seqused_k: Optional[torch.Tensor] = None
self.q2k_indices: Optional[torch.Tensor] = None
self.seqlen_q: Optional[int] = None
self.max_seqlen_k: Optional[int] = None
self.is_sparse: bool = False
self.decode_schedule: Optional[DecodeAttentionSchedule] = None
self.request_indices: Optional[torch.Tensor] = None
self.qo_tile_indices: Optional[torch.Tensor] = None
self.kv_tile_indices: Optional[torch.Tensor] = None
self.merge_indptr: Optional[torch.Tensor] = None
self.o_indptr: Optional[torch.Tensor] = None
self.block_valid_mask: Optional[torch.Tensor] = None
self.kv_pages: Optional[torch.Tensor] = None
self.split_counts: Optional[torch.Tensor] = None
self.split_kv: bool = False
self.cta_tile_q: int = 0
self.num_q_tiles: int = 0
self.kv_chunk_size_pages: int = 0
self.kv_chunk_size_tokens: int = 0
self.work_count: int = 0
self.padded_work_count: int = 0
self.O_partial: Optional[torch.Tensor] = None
self.LSE_partial: Optional[torch.Tensor] = None
# Cached dummy buffers used in non-split path to satisfy the kernel's
# positional arg signature without per-call torch.empty (saves ~5us
# on every run() for small kv).
self._O_partial_dummy: Optional[torch.Tensor] = None
self._LSE_partial_dummy: Optional[torch.Tensor] = None
# When the caller doesn't ask for LSE, the kernel still needs a valid
# tensor pointer to write to. Cache a small placeholder so run() can
# skip the per-call torch.empty for it as well.
self._lse_dummy: Optional[torch.Tensor] = None
def plan(
self,
*,
page_table: torch.Tensor,
seqused_k: torch.Tensor,
seqlen_q: int,
max_seqlen_k: int,
q2k_indices: Optional[torch.Tensor] = None,
num_qo_heads: Optional[int] = None,
num_kv_heads: Optional[int] = None,
head_dim: Optional[int] = 128,
enable_cuda_graph: bool = False,
max_grid_size: Optional[int] = None,
fixed_split_size: Optional[int] = None,
disable_split_kv: bool = False,
) -> "SparseDecodePagedAttentionWrapper":
"""Prepare decode scheduling metadata and reusable workspaces.
Parameters
----------
page_table : torch.Tensor
Shape ``[batch_size, max_num_pages_per_seq]``, dtype int32. Maps
logical pages to physical KV-cache pages.
seqused_k : torch.Tensor
Shape ``[batch_size]``, dtype int32. Effective KV length per
request.
seqlen_q : int
Uniform query length per request.
max_seqlen_k : int
Maximum KV sequence length in the batch.
q2k_indices : torch.Tensor, optional
Sparse selected KV blocks with shape ``[Hkv, total_q, topK]`` and
dtype int32. ``None`` selects the dense all-KV path.
num_qo_heads : int
Number of Q/O heads.
num_kv_heads : int
Number of KV heads. Current decode kernel requires
``num_qo_heads / num_kv_heads == 16`` at run time.
head_dim : int, optional
Head dimension. Must be 128.
enable_cuda_graph : bool, optional
Build schedule metadata compatible with CUDA graph capture.
max_grid_size : int, optional
Override maximum CTA count used by the scheduler.
fixed_split_size : int, optional
Force a fixed split-KV chunk size in pages.
disable_split_kv : bool, optional
Disable split-KV even for long KV sequences.
Returns
-------
SparseDecodePagedAttentionWrapper
``self``, planned and ready for ``run``.
"""
if page_table.ndim != 2:
raise ValueError("decode plan requires page_table with shape [B, max_num_pages_per_seq]")
if page_table.dtype != torch.int32:
raise TypeError("decode plan requires page_table to be torch.int32")
if seqused_k.dtype != torch.int32:
raise TypeError("decode plan requires seqused_k to be torch.int32")
if not page_table.is_cuda or not seqused_k.is_cuda:
raise ValueError("decode plan requires page_table and seqused_k to be CUDA tensors")
if page_table.device != seqused_k.device:
raise ValueError("decode plan requires page_table and seqused_k on the same device")
if page_table.stride(-1) != 1:
raise ValueError("decode plan requires page_table contiguous in the last dimension")
if seqused_k.shape != (int(page_table.shape[0]),):
raise ValueError("decode plan requires seqused_k with shape [B]")
if q2k_indices is not None and q2k_indices.dtype != torch.int32:
raise TypeError("decode plan requires q2k_indices to be torch.int32")
if int(seqlen_q) <= 0 or int(max_seqlen_k) <= 0:
raise ValueError("decode plan requires positive seqlen_q and max_seqlen_k")
if num_qo_heads is None or num_kv_heads is None or head_dim is None:
raise ValueError("decode plan requires num_qo_heads, num_kv_heads, and head_dim")
if head_dim is not None and int(head_dim) != 128:
raise NotImplementedError("decode plan currently supports only head_dim=128")
if int(num_qo_heads) % int(num_kv_heads) != 0:
raise ValueError("decode plan requires num_qo_heads divisible by num_kv_heads")
self.batch = int(page_table.shape[0])
self.num_qo_heads = None if num_qo_heads is None else int(num_qo_heads)
self.num_kv_heads = None if num_kv_heads is None else int(num_kv_heads)
self.head_dim = None if head_dim is None else int(head_dim)
self.page_table = page_table.contiguous()
self.seqused_k = seqused_k.contiguous()
self.q2k_indices = None if q2k_indices is None else q2k_indices.contiguous()
self.seqlen_q = int(seqlen_q)
self.max_seqlen_k = int(max_seqlen_k)
self.is_sparse = q2k_indices is not None
# max_grid_size is hardcoded to num_sms (1 CTA/SM) inside the C++
# schedule launcher because the decode attn kernel always runs at
# 1 CTA/SM (its register/smem budget saturates the SM). Callers
# can still override via the explicit max_grid_size kwarg.
schedule = prepare_decode_schedule(
seqused_k=self.seqused_k,
page_size=self.blk_kv,
seqlen_q=self.seqlen_q,
num_qo_heads=self.num_qo_heads,
num_kv_heads=self.num_kv_heads,
head_dim=self.head_dim,
max_seqlen_k=self.max_seqlen_k,
enable_cuda_graph=bool(enable_cuda_graph),
max_grid_size=max_grid_size,
fixed_split_size=fixed_split_size,
disable_split_kv=bool(disable_split_kv),
)
self.decode_schedule = schedule
self.request_indices = schedule.request_indices
self.qo_tile_indices = schedule.qo_tile_indices
self.kv_tile_indices = schedule.kv_tile_indices
self.merge_indptr = schedule.merge_indptr
self.o_indptr = schedule.o_indptr
self.block_valid_mask = schedule.block_valid_mask
self.kv_pages = schedule.kv_pages
self.split_counts = schedule.split_counts
self.split_kv = schedule.split_kv
self.cta_tile_q = schedule.cta_tile_q
self.num_q_tiles = schedule.num_q_tiles
self.kv_chunk_size_pages = schedule.kv_chunk_size_pages
self.kv_chunk_size_tokens = schedule.kv_chunk_size_tokens
self.work_count = schedule.work_count
self.padded_work_count = schedule.padded_work_count
if schedule.split_kv:
self.O_partial = torch.empty(
(schedule.partial_rows, self.num_qo_heads, self.head_dim),
dtype=torch.float32,
device=page_table.device,
)
self.LSE_partial = torch.empty(
(schedule.partial_rows, self.num_qo_heads),
dtype=torch.float32,
device=page_table.device,
)
self._O_partial_dummy = None
self._LSE_partial_dummy = None
else:
self.O_partial = None
self.LSE_partial = None
# decode_forward_paged_fp8 always wants non-None partial buffers
# for the kernel's positional arg layout (compile keeps the slot
# alive even when split_kv=False). Allocate once here and reuse.
self._O_partial_dummy = torch.empty(
(1, self.head_dim),
dtype=torch.float32,
device=page_table.device,
)
self._LSE_partial_dummy = torch.empty(
(1, self.num_qo_heads),
dtype=torch.float32,
device=page_table.device,
)
# LSE dummy is shape (1, head_q) — used when caller doesn't request
# LSE and the schedule isn't split-KV (split-KV always writes LSE).
self._lse_dummy = torch.empty(
(1, self.num_qo_heads),
dtype=torch.float32,
device=page_table.device,
)
return self
def run(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
*,
softmax_scale: Optional[float] = None,
return_softmax_lse: bool = False,
out: Optional[torch.Tensor] = None,
lse: Optional[torch.Tensor] = None,
):
"""Launch decode using metadata cached by ``plan``.
Parameters
----------
q : torch.Tensor
Shape ``[batch_size * seqlen_q, Hq, 128]`` and dtype FP8 E4M3.
k : torch.Tensor
Paged K cache with shape ``[num_pages, Hkv, blk_kv, 128]``.
v : torch.Tensor
Paged V cache with the same shape as ``k``.
softmax_scale : float, optional
Softmax scale. Defaults to ``1 / sqrt(128)``.
return_softmax_lse : bool, optional
If True, return ``(out, lse)``.
out : torch.Tensor, optional
Preallocated BF16 output buffer with shape ``q.shape``.
lse : torch.Tensor, optional
Preallocated float32 LSE buffer with shape ``[total_q, Hq]``.
Returns
-------
torch.Tensor or tuple[torch.Tensor, torch.Tensor]
BF16 output, optionally with float32 LSE.
"""
if self.decode_schedule is None:
raise RuntimeError("decode wrapper must be planned before run")
if self.is_sparse:
# Sparse path still goes through the validating wrapper for now;
# only the dense fast path is collapsed.
return sparse_decode_atten_func(
q, k, v, self.q2k_indices,
page_table=self.page_table, seqused_k=self.seqused_k,
seqlen_q=self.seqlen_q, max_seqlen_k=self.max_seqlen_k,
blk_kv=self.blk_kv, causal=self.causal,
softmax_scale=softmax_scale, return_softmax_lse=return_softmax_lse,
schedule=self.decode_schedule,
O_partial=self.O_partial, LSE_partial=self.LSE_partial,
)
if softmax_scale is None:
softmax_scale = q.shape[-1] ** -0.5
if out is None:
out = torch.empty(q.shape, dtype=torch.bfloat16, device=q.device)
if lse is None:
if return_softmax_lse or self.split_kv:
# Real LSE needed — must allocate per-call (shape depends on q).
lse = torch.empty(
q.shape[:2], dtype=torch.float32, device=q.device,
)
else:
# Kernel only needs a valid pointer; reuse cached dummy.
lse = self._lse_dummy
from .src.sm100.fwd_decode import decode_forward_paged_fp8
schedule = self.decode_schedule
decode_forward_paged_fp8(
q, k, v,
self.page_table, self.seqused_k,
out, lse,
schedule.request_indices, schedule.qo_tile_indices,
schedule.kv_tile_indices, schedule.block_valid_mask,
schedule.split_counts, schedule.o_indptr, schedule.merge_indptr,
self.O_partial, self.LSE_partial,
softmax_scale=float(softmax_scale),
seqlen_q=self.seqlen_q,
page_size=self.blk_kv,
kv_chunk_size_pages=int(schedule.kv_chunk_size_pages),
max_split_count=int(schedule.max_split_count),
split_kv=bool(schedule.split_kv),
causal=self.causal,
return_lse=bool(return_softmax_lse),
# cached dummies — avoid per-call torch.empty inside run_decode_attention
O_partial_dummy=self._O_partial_dummy,
LSE_partial_dummy=self._LSE_partial_dummy,
)
if return_softmax_lse:
return out, lse
return out
def _sparse_atten_csr_varlen_forward(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
k2q_row_ptr: torch.Tensor,
k2q_q_indices: torch.Tensor,
topK: int,
blk_kv: int,
causal: bool,
softmax_scale: float,
lse_temperature_scale: float,
return_temperature_lse: bool,
partial_dtype: torch.dtype,
return_softmax_lse: bool,
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
page_table: Optional[torch.Tensor],
seqused_k: Optional[torch.Tensor],
schedule: Optional[SparseAttentionSchedule],
usable_SM_count: int,
batch: int,
head_kv: int,
max_seqlen_q: int,
max_seqlen_k: int,
qk_dtype: torch.dtype,
pv_dtype: torch.dtype,
):
total_q, head_q, dim = q.shape
if head_q % head_kv != 0:
raise ValueError("q.shape[1] must be divisible by head_kv")
max_num_kv_blocks = _csr_row_capacity(k2q_row_ptr)
temperature_lse_fast_path = (
return_temperature_lse
and math.isclose(
float(lse_temperature_scale),
1.0,
rel_tol=0.0,
abs_tol=_TEMPERATURE_LSE_FAST_PATH_ABS_TOL,
)
)
kernel_return_temperature_lse = (
return_temperature_lse and not temperature_lse_fast_path
)
O_partial = torch.empty(
topK, total_q, head_q, dim, dtype=partial_dtype, device=q.device
)
LSE_partial = torch.empty(
topK, total_q, head_q, dtype=torch.float32, device=q.device
)
LSE_temperature_partial = (
torch.empty(topK, total_q, head_q, dtype=torch.float32, device=q.device)
if kernel_return_temperature_lse
else None
)
O_out = torch.empty(total_q, head_q, dim, dtype=torch.bfloat16, device=q.device)
LSE_out = torch.empty(total_q, head_q, dtype=torch.float32, device=q.device)
LSE_temperature_out = (
torch.empty_like(LSE_out) if kernel_return_temperature_lse else None
)
if schedule is None:
k2q_qsplit_indices = torch.empty_like(k2q_q_indices)
split_counts = torch.zeros(
(total_q, head_kv),
dtype=torch.int32,
device=q.device,
)
else:
_validate_fwd_schedule(
schedule,
q=q,
k2q_q_indices=k2q_q_indices,
head_kv=head_kv,
)
k2q_qsplit_indices = schedule.qsplit_indices
split_counts = schedule.split_counts
schedule = _call_sparse_forward_sm100_csr_varlen(
q,
k,
v,
k2q_row_ptr,
k2q_q_indices,
k2q_qsplit_indices,
split_counts,
cu_seqlens_q,
cu_seqlens_k,
page_table,
seqused_k,
O_partial,
LSE_partial,
LSE_temperature_partial,
softmax_scale,
lse_temperature_scale,
kernel_return_temperature_lse,
max_num_kv_blocks,
blk_kv,
head_kv,
max_seqlen_q,
usable_SM_count,
causal=causal,
schedule=schedule,
qk_dtype=qk_dtype,
pv_dtype=pv_dtype,
)
# Sparse Attention and Sparse Page Attention both use the varlen-Q
# combine path; the kernel-written LSE_out is the final contract.
combine(
O_partial,
LSE_partial,
O_out,
LSE_out,
lse_temperature_partial=LSE_temperature_partial,
lse_temperature_out=LSE_temperature_out,
cu_seqlens=cu_seqlens_q,
split_counts=split_counts,
use_pdl=True,
)
if temperature_lse_fast_path:
LSE_temperature_out = LSE_out
if return_softmax_lse:
if return_temperature_lse:
return O_out, LSE_out, LSE_temperature_out
return O_out, LSE_out
return O_out
def _call_sparse_decode_forward_sm100_paged_fp8(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
q2k_indices: Optional[torch.Tensor],
page_table: torch.Tensor,
seqused_k: torch.Tensor,
out: torch.Tensor,
lse: torch.Tensor,
schedule: DecodeAttentionSchedule,
O_partial: Optional[torch.Tensor],
LSE_partial: Optional[torch.Tensor],
*,
softmax_scale: float,
seqlen_q: int,
max_seqlen_k: int,
blk_kv: int,
causal: bool,
return_lse: bool = True,
) -> None:
"""Compile and launch the SM100 paged fp8 decode forward kernel.
Dense decode is selected by ``q2k_indices=None``. Sparse decode will reuse
the same schedule wrapper but needs a separate q2k gather path.
"""
if q2k_indices is not None:
raise NotImplementedError("SM100 paged fp8 sparse decode forward is not implemented yet")
if schedule.cta_tile_q != 128:
raise NotImplementedError(f"decode forward requires cta_tile_q=128, got {schedule.cta_tile_q}")
if schedule.split_kv:
if O_partial is None or LSE_partial is None:
raise ValueError("split decode forward requires O_partial and LSE_partial")
if O_partial.dtype != torch.float32:
raise TypeError(f"O_partial must be torch.float32, got {O_partial.dtype}")
if LSE_partial.dtype != torch.float32:
raise TypeError(f"LSE_partial must be torch.float32, got {LSE_partial.dtype}")
from .src.sm100.fwd_decode import decode_forward_paged_fp8
decode_forward_paged_fp8(
q,
k,
v,
page_table,
seqused_k,
out,
lse,
schedule.request_indices,
schedule.qo_tile_indices,
schedule.kv_tile_indices,
schedule.block_valid_mask,
schedule.split_counts,
schedule.o_indptr,
schedule.merge_indptr,
O_partial,
LSE_partial,
softmax_scale=float(softmax_scale),
seqlen_q=int(seqlen_q),
page_size=int(blk_kv),
kv_chunk_size_pages=int(schedule.kv_chunk_size_pages),
max_split_count=int(schedule.max_split_count),
split_kv=bool(schedule.split_kv),
causal=bool(causal),
return_lse=bool(return_lse),
)
def _call_sparse_forward_sm100_csr_varlen(
q,
k,
v,
k2q_row_ptr,
k2q_q_indices,
k2q_qsplit_indices,
split_counts,
cu_seqlens_q,
cu_seqlens_k,
page_table,
seqused_k,
O_partial,
LSE_partial,
LSE_temperature_partial,
softmax_scale,
lse_temperature_scale,
return_temperature_lse,
max_num_kv_blocks,
blk_kv,
head_kv,
max_seqlen_q,
usable_SM_count=-1,
*,
causal=False,
use_prepare_scheduler=True,
schedule: Optional[SparseAttentionSchedule] = None,
qk_dtype: torch.dtype,
pv_dtype: torch.dtype,
):
"""Compile and launch the SM100 sparse forward K1 kernel on CSR metadata."""
head_dim = q.shape[-1]
dtype = q.dtype
qk_dtype = _normalize_forward_mma_dtype(qk_dtype, q.dtype, "qk_dtype")
pv_dtype = _normalize_forward_mma_dtype(pv_dtype, v.dtype, "pv_dtype")
partial_dtype = O_partial.dtype
return_temperature_lse = bool(return_temperature_lse)
if return_temperature_lse != (LSE_temperature_partial is not None):
raise ValueError(
"return_temperature_lse must match LSE_temperature_partial presence"
)
lse_temperature_scale = float(lse_temperature_scale)
if not math.isfinite(lse_temperature_scale) or lse_temperature_scale <= 0.0:
raise ValueError(
f"lse_temperature_scale must be finite and > 0, got {lse_temperature_scale}"
)
lse_temperature_inv_scale = 1.0 / lse_temperature_scale
n_block_size = int(blk_kv)
head_q = q.shape[1]
qhead_per_kv = head_q // head_kv
paged_kv = page_table is not None
if not bool(use_prepare_scheduler):
raise RuntimeError("sparse forward requires prepare scheduler")
schedule_enabled = k2q_row_ptr.shape[1] > 1
page_size = int(k.shape[2]) if paged_kv else None
if paged_kv:
k_kernel, v_kernel = _prepare_paged_kv_for_tma(k, v, n_block_size)
else:
k_kernel = k
v_kernel = v
O_partial_flat = O_partial.reshape(-1, head_dim).contiguous()
Q_flat = q.reshape(-1, head_dim).contiguous()
Q_gather4_desc = (
create_q_gather4_tma_desc(
Q_flat,
box_x=128 if q.dtype == torch.float8_e4m3fn else 64,
)
if qhead_per_kv in (1, 2, 4)
else None
)
if schedule is None:
schedule = prepare_sparse_fwd_schedule_and_split(
k2q_row_ptr=k2q_row_ptr,
k2q_q_indices=k2q_q_indices,
k2q_qsplit_indices=k2q_qsplit_indices,
split_counts=split_counts,
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_k,
total_q=int(q.shape[0]),
max_seqlen_q=max_seqlen_q,
topk=int(O_partial.shape[0]),
head_kv=head_kv,
qhead_per_kv=qhead_per_kv,
blk_kv=n_block_size,
device=q.device,
enabled=schedule_enabled,
)
use_prepare_scheduler = schedule.enabled
scheduler_metadata = schedule.scheduler_metadata
work_count = schedule.work_count
work_capacity = schedule.work_capacity
if not use_prepare_scheduler or scheduler_metadata is None or work_count is None or work_capacity <= 0:
raise RuntimeError("sparse forward requires a non-empty prepared schedule")
key = (
"sparse_forward_sm100_csr_varlen",
head_dim,
n_block_size,
qhead_per_kv,
dtype,
k.dtype,
v.dtype,
qk_dtype,
pv_dtype,
partial_dtype,
bool(causal),
bool(paged_kv),
bool(use_prepare_scheduler),
page_size,
bool(seqused_k is not None),
bool(return_temperature_lse),
)
if key not in _compile_cache:
from .src.common.aot_cache import try_load_aot, save_aot
loaded = try_load_aot(key)
if loaded is not None:
_compile_cache[key] = loaded
else:
kernel = SparseAttentionForwardSm100(
head_dim=head_dim,
qheadperkv=qhead_per_kv,
n_block_size=n_block_size,
paged_kv=paged_kv,
page_size=page_size,
has_seqused_k=seqused_k is not None,
causal=bool(causal),
use_prepare_scheduler=use_prepare_scheduler,
qk_dtype=_torch_dtype_to_cutlass_dtype(qk_dtype),
pv_dtype=_torch_dtype_to_cutlass_dtype(pv_dtype),
)
_compile_cache[key] = cute.compile(
kernel,
to_cute_tensor_kvouter(k_kernel),
to_cute_tensor_kvouter(v_kernel),
to_cute_tensor_kvouter(k2q_q_indices),
to_cute_tensor_kvouter(k2q_qsplit_indices),
to_cute_tensor_kvouter(k2q_row_ptr),
None if scheduler_metadata is None else to_cute_tensor_kvouter(scheduler_metadata),
None if work_count is None else to_cute_tensor_kvouter(work_count),
to_cute_tensor_kvouter(O_partial_flat),
to_cute_tensor_kvouter(LSE_partial),
None
if LSE_temperature_partial is None
else to_cute_tensor_kvouter(LSE_temperature_partial),
to_cute_tensor_kvouter(Q_flat),
None if Q_gather4_desc is None else to_cute_tensor_kvouter(Q_gather4_desc),
None if page_table is None else to_cute_tensor_kvouter(page_table),
None if seqused_k is None else to_cute_tensor_kvouter(seqused_k),
to_cute_tensor_kvouter(cu_seqlens_q),
to_cute_tensor_kvouter(cu_seqlens_k),
Float32(softmax_scale),
Float32(lse_temperature_inv_scale),
Int32(max_num_kv_blocks),
Int32(head_kv),
Int32(max_seqlen_q),
Int32(work_capacity),
cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
options="--enable-tvm-ffi",
)
save_aot(key, _compile_cache[key])
with torch.cuda.nvtx.range("Fwd_SparseAttn_Sm100_CsrVarlen"):
_compile_cache[key](
k_kernel,
v_kernel,
k2q_q_indices,
k2q_qsplit_indices,
k2q_row_ptr,
scheduler_metadata,
work_count,
O_partial_flat,
LSE_partial,
LSE_temperature_partial,
Q_flat,
Q_gather4_desc,
page_table,
seqused_k,
cu_seqlens_q,
cu_seqlens_k,
softmax_scale,
lse_temperature_inv_scale,
max_num_kv_blocks,
head_kv,
max_seqlen_q,
work_capacity,
)
return schedule
def _call_sparse_forward_sm100_csr_varlen_nvfp4_kv(
q,
k,
v,
k_scale_128x4,
v_scale_128x4,
k_global_scale,
v_global_scale,
k2q_row_ptr,
k2q_q_indices,
k2q_qsplit_indices,
split_counts,
cu_seqlens_q,
cu_seqlens_k,
page_table,
seqused_k,
O_partial,
LSE_partial,
LSE_temperature_partial,
softmax_scale,
lse_temperature_scale,
return_temperature_lse,
max_num_kv_blocks,
blk_kv,
head_kv,
max_seqlen_q,
*,
causal=False,
use_prepare_scheduler=True,
schedule: Optional[SparseAttentionSchedule] = None,
):
"""Compile and launch the SM100 sparse forward K1 kernel with NVFP4 K/V."""
head_dim = q.shape[-1]
dtype = q.dtype
partial_dtype = O_partial.dtype
return_temperature_lse = bool(return_temperature_lse)
if return_temperature_lse != (LSE_temperature_partial is not None):
raise ValueError(
"return_temperature_lse must match LSE_temperature_partial presence"
)
lse_temperature_scale = float(lse_temperature_scale)
if not math.isfinite(lse_temperature_scale) or lse_temperature_scale <= 0.0:
raise ValueError(
f"lse_temperature_scale must be finite and > 0, got {lse_temperature_scale}"
)
lse_temperature_inv_scale = 1.0 / lse_temperature_scale
n_block_size = int(blk_kv)
head_q = q.shape[1]
qhead_per_kv = head_q // head_kv
fp8_pair_dequant = os.environ.get("MINIMAX_KVFP4_FP8_PAIR_DEQUANT", "1") != "0"
k_global_scale_kernel = k_global_scale
# V global scale is linear in the final output. Keep K1 on block-scale-only V
# and apply the tensor scale once in K2 combine.
v_global_scale_kernel = None
has_k_global_scale = k_global_scale_kernel is not None
has_v_global_scale = v_global_scale_kernel is not None
paged_kv = page_table is not None
if not bool(use_prepare_scheduler):
raise RuntimeError("KVFP4 sparse forward requires prepare scheduler")
schedule_enabled = k2q_row_ptr.shape[1] > 1
page_size = int(k.shape[2]) if paged_kv else None
if paged_kv:
_prepare_paged_kv_for_tma(k, v, n_block_size)
k_kernel = k
v_kernel = v
O_partial_flat = O_partial.reshape(-1, head_dim).contiguous()
Q_flat = q.reshape(-1, head_dim).contiguous()
Q_gather4_desc = (
create_q_gather4_tma_desc(
Q_flat,
box_x=128 if q.dtype == torch.float8_e4m3fn else 64,
)
if qhead_per_kv in (1, 2, 4)
else None
)
if schedule is None:
schedule = prepare_sparse_fwd_schedule_and_split(
k2q_row_ptr=k2q_row_ptr,
k2q_q_indices=k2q_q_indices,
k2q_qsplit_indices=k2q_qsplit_indices,
split_counts=split_counts,
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_k,
total_q=int(q.shape[0]),
max_seqlen_q=max_seqlen_q,
topk=int(O_partial.shape[0]),
head_kv=head_kv,
qhead_per_kv=qhead_per_kv,
blk_kv=n_block_size,
device=q.device,
enabled=schedule_enabled,
)
use_prepare_scheduler = schedule.enabled
scheduler_metadata = schedule.scheduler_metadata
work_count = schedule.work_count
work_capacity = schedule.work_capacity
if not use_prepare_scheduler or scheduler_metadata is None or work_count is None or work_capacity <= 0:
raise RuntimeError("KVFP4 sparse forward requires a non-empty prepared schedule")
key = (
"sparse_forward_sm100_csr_varlen_nvfp4_kv",
head_dim,
n_block_size,
qhead_per_kv,
dtype,
partial_dtype,
bool(causal),
bool(paged_kv),
bool(use_prepare_scheduler),
page_size,
bool(seqused_k is not None),
bool(return_temperature_lse),
bool(fp8_pair_dequant),
bool(has_k_global_scale),
bool(has_v_global_scale),
)
if key not in _compile_cache:
from .src.common.aot_cache import try_load_aot, save_aot
loaded = try_load_aot(key)
if loaded is not None:
_compile_cache[key] = loaded
else:
kernel = SparseAttentionForwardNvfp4KvSm100(
head_dim=head_dim,
qheadperkv=qhead_per_kv,
n_block_size=n_block_size,
paged_kv=paged_kv,
page_size=page_size,
has_seqused_k=seqused_k is not None,
causal=bool(causal),
use_prepare_scheduler=use_prepare_scheduler,
fp8_pair_dequant=bool(fp8_pair_dequant),
has_k_global_scale=bool(has_k_global_scale),
has_v_global_scale=bool(has_v_global_scale),
)
_compile_cache[key] = cute.compile(
kernel,
to_cute_tensor_kvouter(k_kernel),
to_cute_tensor_kvouter(v_kernel),
to_cute_tensor_kvouter(k_scale_128x4),
to_cute_tensor_kvouter(v_scale_128x4),
None if k_global_scale_kernel is None else to_cute_tensor_kvouter(k_global_scale_kernel),
None if v_global_scale_kernel is None else to_cute_tensor_kvouter(v_global_scale_kernel),
to_cute_tensor_kvouter(k2q_q_indices),
to_cute_tensor_kvouter(k2q_qsplit_indices),
to_cute_tensor_kvouter(k2q_row_ptr),
None if scheduler_metadata is None else to_cute_tensor_kvouter(scheduler_metadata),
None if work_count is None else to_cute_tensor_kvouter(work_count),
to_cute_tensor_kvouter(O_partial_flat),
to_cute_tensor_kvouter(LSE_partial),
None
if LSE_temperature_partial is None
else to_cute_tensor_kvouter(LSE_temperature_partial),
to_cute_tensor_kvouter(Q_flat),
None if Q_gather4_desc is None else to_cute_tensor_kvouter(Q_gather4_desc),
None if page_table is None else to_cute_tensor_kvouter(page_table),
None if seqused_k is None else to_cute_tensor_kvouter(seqused_k),
to_cute_tensor_kvouter(cu_seqlens_q),
to_cute_tensor_kvouter(cu_seqlens_k),
Float32(softmax_scale),
Float32(lse_temperature_inv_scale),
Int32(max_num_kv_blocks),
Int32(head_kv),
Int32(max_seqlen_q),
Int32(work_capacity),
cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
options="--enable-tvm-ffi",
)
save_aot(key, _compile_cache[key])
with torch.cuda.nvtx.range("Fwd_SparseAttn_Sm100_CsrVarlen_KVFP4"):
_compile_cache[key](
k_kernel,
v_kernel,
k_scale_128x4,
v_scale_128x4,
k_global_scale_kernel,
v_global_scale_kernel,
k2q_q_indices,
k2q_qsplit_indices,
k2q_row_ptr,
scheduler_metadata,
work_count,
O_partial_flat,
LSE_partial,
LSE_temperature_partial,
Q_flat,
Q_gather4_desc,
page_table,
seqused_k,
cu_seqlens_q,
cu_seqlens_k,
softmax_scale,
lse_temperature_inv_scale,
max_num_kv_blocks,
head_kv,
max_seqlen_q,
work_capacity,
)
return schedule
# SPDX-FileCopyrightText: Copyright (c) 2026 MiniMax
# SPDX-License-Identifier: MIT
"""Sparse attention interface.
Current delivery scope:
- head dimension is supported only for D=128
Public API:
sparse_atten_func(...)
sparse_decode_atten_func(...)
SparseDecodePagedAttentionWrapper
Internal forward core:
_sparse_atten_csr_varlen_forward(...)
Preprocessing (external, done once):
q2k_indices [head_kv, total_q, topK] -> sparse_index_utils.build_k2q_csr()
-> k2q_row_ptr [head_kv, total_rows + 1] int32
-> k2q_q_indices [head_kv, total_q * topK] int32
"""
import math
import os
from typing import Optional
import cutlass
import cutlass.cute as cute
import torch
from cutlass import Float32, Int32
from cutlass.cute.runtime import from_dlpack
from .src.sm100.fwd.combine import combine
from .src.sm100.fwd.atten_fwd import SparseAttentionForwardSm100
from .src.sm100.fwd.atten_fwd_nvfp4_kv import SparseAttentionForwardNvfp4KvSm100
from .src.sm100.prepare_scheduler import (
SparseAttentionSchedule,
prepare_sparse_fwd_schedule_and_split,
)
from .src.sm100.decode_schedule import (
DecodeAttentionSchedule,
prepare_decode_schedule,
)
from .src.common.cute_dsl_utils import to_cute_tensor as to_cute_tensor_kvouter
from .src.common.tma_utils import (
create_q_gather4_tma_desc,
)
_compile_cache: dict = {}
_TEMPERATURE_LSE_FAST_PATH_ABS_TOL = 1e-12
_SUPPORTED_SPARSE_TOPK = (4, 8, 16, 32)
_SUPPORTED_FWD_DTYPES = (torch.bfloat16, torch.float8_e4m3fn)
_SUPPORTED_FWD_MMA_DTYPES = (torch.bfloat16, torch.float8_e4m3fn)
_SUPPORTED_DECODE_QHEAD_PER_KV = 16
def _normalize_partial_dtype(partial_dtype: torch.dtype) -> torch.dtype:
supported = {torch.float32, torch.bfloat16, torch.float16, torch.float8_e4m3fn}
if partial_dtype not in supported:
raise TypeError(
"partial_dtype must be one of torch.float32 / torch.bfloat16 / "
"torch.float16 / torch.float8_e4m3fn, "
f"got {partial_dtype}"
)
return partial_dtype
def _normalize_forward_mma_dtype(dtype: Optional[torch.dtype], fallback: torch.dtype, name: str) -> torch.dtype:
dtype = fallback if dtype is None else dtype
if dtype not in _SUPPORTED_FWD_MMA_DTYPES:
raise TypeError(
f"{name} must be one of torch.bfloat16 / torch.float8_e4m3fn, got {dtype}"
)
return dtype
def _resolve_forward_mma_dtypes(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
qk_dtype: Optional[torch.dtype],
pv_dtype: Optional[torch.dtype],
) -> tuple[torch.dtype, torch.dtype]:
qk_dtype = _normalize_forward_mma_dtype(qk_dtype, q.dtype, "qk_dtype")
if pv_dtype is None:
# Preserve the historical fp8 KV-cache path: BF16 Q with FP8 K/V
# stages both K and V as BF16 compute operands.
if (
q.dtype == torch.bfloat16
and k.dtype == torch.float8_e4m3fn
and v.dtype == torch.float8_e4m3fn
):
pv_dtype = torch.bfloat16
else:
pv_dtype = v.dtype
pv_dtype = _normalize_forward_mma_dtype(pv_dtype, pv_dtype, "pv_dtype")
if q.dtype != qk_dtype:
raise ValueError(
"qk_dtype must match q storage dtype; Q fp8->bf16 staging is not supported"