# 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