diff --git a/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md index 4cb49d003240..f8225bac5010 100644 --- a/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md @@ -25,6 +25,8 @@ The BF16, block-FP8, and NVFP4 checkpoints support the same feature set: The KV-cache dtype and the GDN recurrent-state dtype are selected independently of the model weight precision, through `kv_cache_config.dtype` and `kv_cache_config.mamba_ssm_cache_dtype`. +With an FP8 KV cache, QSA sparse attention selects the PrimTS kernel on SM100/SM103/SM107 when the inputs qualify, and Triton otherwise (small decode batches, unsupported shapes, or PrimTS unavailable). PrimTS converts the BF16 queries to FP8 (unit scale, clamped to the E4M3 range) and attends against the FP8 cache; Triton upcasts K/V to BF16 instead. The two differ at FP8 rounding level: full GSM8K on B300 at TP1 scores 95.45 with either kernel. + ## Prerequisites * GPU: NVIDIA Blackwell architecture. The BF16 checkpoint also runs on Hopper. diff --git a/docs/source/release-notes.md b/docs/source/release-notes.md index 26022b9cf621..543f98573a61 100644 --- a/docs/source/release-notes.md +++ b/docs/source/release-notes.md @@ -12,6 +12,7 @@ All published functionality in the Release Notes has been fully tested and verif - **First-launch compile cost.** The kernels are JIT-compiled by the CuTe DSL and cached on disk under the FlashInfer workspace. The first start on a host (or after a FlashInfer / CuTe DSL upgrade) compiles the kernel variants the deployment's generation row counts select during `LLM()` startup, about 5 to 10 minutes on B200 for DeepSeek-V4; later starts load the cache. Under `trtllm-llmapi-launch` and the default `LLM()` MPI session every rank has its own FlashInfer workspace (`TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=1`), so each rank compiles its own copy; set `TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0` to share one cache between the ranks of a host. - **Warm-up controls.** `TLLM_PRIMTS_WARM_WORKERS` sets the compile subprocesses per rank (default 4, reduced so that the ranks on a host, counted from the model's parallel mapping, start at most 16 together; `0` compiles in-process), `TLLM_PRIMTS_WARM_TIMEOUT` bounds the wait for them (default 1800 s), and `TLLM_PRIMTS_WARM_STRICT=1` turns a worker failure into an error instead of an in-process retry. Each worker imports torch and the CuTe DSL and opens a CUDA context, so lower the worker count on nodes with little host memory. - **`FLASHINFER_DISABLE_JIT=1`** now also refuses to compile these kernels (with the on-disk CuTe DSL cache enabled, its default): startup then fails on a variant missing from the cache instead of falling back to TRTLLM-Gen, so a deployment that sets it must pre-populate the cache (one start without the variable) or opt out of the library. +- **Qwen3.8-Flash-Next QSA attention on SM100/SM103/SM107 now runs on PrimTS kernels by default.** The sparse core uses FlashInfer's QToken-KvBlock-Sparse kernel and the indexer scorer a TRT-LLM CuTe DSL kernel. Triton remains the fallback for small decode batches, other GPUs and ineligible inputs. Prefill is 1.7 to 4.2x faster on the core and 1.2 to 4.9x on the indexer (Rubin and B300). - Disaggregated serving supports the Responses API: `trtllm-serve disaggregated` serves `/v1/responses` with context-first and generation-first scheduling, and `/v1/models`. `previous_response_id` is rejected and `store` is ignored there, since no response history is kept. See [Disaggregated Serving](features/disagg-serving.md#responses-api). ### Infrastructure Changes diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py index c1c127aa35df..fc050b98e520 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Module-layer integration for QSA sparse attention.""" +import itertools import math from typing import TYPE_CHECKING, Optional @@ -374,6 +375,7 @@ def select_qsa_paged_tokens( visible_blocks: torch.Tensor | None = None, context_rows: bool = False, mtp_share_slot: int | None = None, + expand_blocks: bool = True, ) -> torch.Tensor: """Select tokens with packed, fixed-width paged scoring. @@ -386,6 +388,9 @@ def select_qsa_paged_tokens( steps rebuild their rows from it without running the indexer at all. It is the calling layer's row in the shared buffers, and ``None`` turns the reuse off. + + ``expand_blocks=False`` returns the ``[rows, block_topk]`` compressed-block + IDs for consumers that resolve blocks and the causal tail themselves. """ from .kernels import triton_qsa_mtp_reuse_topk, triton_qsa_paged_index_scores @@ -406,6 +411,8 @@ def select_qsa_paged_tokens( visible_blocks=visible_blocks, output=top_k_output[: q.shape[0]], ) + if not expand_blocks: + return indices return expand_qsa_block_indices( indices, query_positions, @@ -414,20 +421,42 @@ def select_qsa_paged_tokens( token_topk=params.token_topk, ) - logits = triton_qsa_paged_index_scores( - q=q, - index_cache=index_cache, - block_table=metadata.qsa_block_table, - query_positions=query_positions, - request_indices=request_indices, + # CUDA radix Top-K scans only [row_starts, row_ends), so score tiles + # outside the causal boundary need not be materialized. Preserve fully + # initialized logits for the Torch fallback below. + only_visible_blocks = top_k is not None and q.is_cuda + from .prims_ts_core import use_prims_ts_indexer + + if use_prims_ts_indexer( + params, + q, tokens_per_block=metadata.kv_cache_manager.tokens_per_block, - compress_ratio=params.compress_ratio, - # CUDA radix Top-K scans only [row_starts, row_ends), so score tiles - # outside the causal boundary need not be materialized. Preserve fully - # initialized logits for the Torch fallback below. - only_visible_blocks=top_k is not None and q.is_cuda, - context_rows=context_rows, - ) + is_decode=not context_rows, + ): + from .prims_ts_indexer import prims_ts_qsa_index_scores + + logits = prims_ts_qsa_index_scores( + q, + index_cache, + metadata.qsa_block_table, + query_positions, + request_indices, + tokens_per_block=metadata.kv_cache_manager.tokens_per_block, + compress_ratio=params.compress_ratio, + only_visible_blocks=only_visible_blocks, + ) + else: + logits = triton_qsa_paged_index_scores( + q=q, + index_cache=index_cache, + block_table=metadata.qsa_block_table, + query_positions=query_positions, + request_indices=request_indices, + tokens_per_block=metadata.kv_cache_manager.tokens_per_block, + compress_ratio=params.compress_ratio, + only_visible_blocks=only_visible_blocks, + context_rows=context_rows, + ) if top_k is not None and logits.is_cuda: if top_k_output is None or top_k_row_starts is None: raise ValueError("QSA CUDA Top-K requires caller-owned output and row starts") @@ -464,6 +493,8 @@ def select_qsa_paged_tokens( ) if mtp_share_slot is not None: _capture_mtp_shared_topk(metadata, mtp_share_slot, indices, visible_blocks) + if not expand_blocks: + return indices return expand_qsa_block_indices( indices, query_positions, @@ -629,7 +660,32 @@ def forward( tokens_per_block, ) - if attn_metadata.num_contexts == 0: + softmax_scale = 1.0 / (attention.q_scaling * attention.head_dim**0.5) + is_decode = attn_metadata.num_contexts == 0 + if params.sparse_core_backend == "triton": + core_backend = "triton" + else: + from .prims_ts_core import resolve_sparse_core_backend + + core_backend = resolve_sparse_core_backend( + params, + q=q, + k_cache=k_cache, + page_size=tokens_per_block, + is_decode=is_decode, + num_rows=num_tokens, + ) + use_prims_ts = core_backend == "prims_ts" + if use_prims_ts: + from .prims_ts_core import prims_ts_sparse_gqa + + logger.info_once( + "QSA PrimTS sparse core attention is active", + key="qsa_prims_ts_sparse_core_active", + ) + max_seq_len_kv = attn_metadata.qsa_block_table.shape[1] * tokens_per_block + + if is_decode: index_cache = attn_metadata.kv_cache_manager.get_index_k_buffer(attention.layer_idx) if index_cache is None: raise RuntimeError( @@ -661,32 +717,48 @@ def forward( top_k_row_starts=attn_metadata.qsa_topk_row_starts, visible_blocks=attn_metadata.qsa_visible_blocks[:num_tokens], mtp_share_slot=mtp_share_slot, + expand_blocks=not use_prims_ts, ) fuse_output_gate = False - if output_gate is not None: - from .kernels import can_fuse_qsa_splitk_output_gate - - fuse_output_gate = can_fuse_qsa_splitk_output_gate( + if use_prims_ts: + output = prims_ts_sparse_gqa( q, k_cache, + v_cache, + attn_metadata.qsa_block_table, selected, - output_gate, + req_idx, + logical, + page_size=tokens_per_block, + max_seq_len_kv=max_seq_len_kv, + softmax_scale=softmax_scale, + is_decode=True, ) - if fuse_output_gate: - logger.info_once( - "QSA split-K merge fuses the attention output gate", - key="qsa_fused_output_gate_active", + else: + if output_gate is not None: + from .kernels import can_fuse_qsa_splitk_output_gate + + fuse_output_gate = can_fuse_qsa_splitk_output_gate( + q, + k_cache, + selected, + output_gate, ) - output = qsa_sparse_gqa( - q=q, - k_cache=k_cache, - v_cache=v_cache, - selected_tokens=selected, - request_indices=req_idx, - metadata=attn_metadata, - softmax_scale=1.0 / (attention.q_scaling * attention.head_dim**0.5), - output_gate=output_gate if fuse_output_gate else None, - ) + if fuse_output_gate: + logger.info_once( + "QSA split-K merge fuses the attention output gate", + key="qsa_fused_output_gate_active", + ) + output = qsa_sparse_gqa( + q=q, + k_cache=k_cache, + v_cache=v_cache, + selected_tokens=selected, + request_indices=req_idx, + metadata=attn_metadata, + softmax_scale=softmax_scale, + output_gate=output_gate if fuse_output_gate else None, + ) output = output.reshape(num_tokens, -1) if output_gate is not None and not fuse_output_gate: output = attention.apply_output_gate(output, output_gate) @@ -700,11 +772,17 @@ def forward( attn_metadata.qsa_block_table.shape[1] * tokens_per_block // params.compress_ratio ) chunk_size = _query_chunk_size(num_tokens, score_columns) + if use_prims_ts: + from .prims_ts_core import slice_request_boundaries + + cu_seq_lens = list( + itertools.accumulate(attn_metadata.seq_lens[: attn_metadata.num_seqs].tolist()) + ) output = torch.empty_like(q) - for start in range(0, num_tokens, chunk_size): - end = min(start + chunk_size, num_tokens) + + def select(start: int, end: int) -> torch.Tensor: packed_slice = slice(start, end) - selected = select_qsa_paged_tokens( + return select_qsa_paged_tokens( q_index[packed_slice], index_cache, logical[packed_slice], @@ -717,18 +795,54 @@ def forward( top_k_row_starts=attn_metadata.qsa_topk_row_starts[packed_slice], visible_blocks=attn_metadata.qsa_visible_blocks[packed_slice], context_rows=True, + expand_blocks=not use_prims_ts, ) - output[packed_slice] = qsa_sparse_gqa( - q=q[packed_slice], - k_cache=k_cache, - v_cache=v_cache, - selected_tokens=selected, - request_indices=req_idx[packed_slice], - metadata=attn_metadata, - softmax_scale=1.0 / (attention.q_scaling * attention.head_dim**0.5), - query_positions=logical[packed_slice], - compress_ratio=params.compress_ratio, - ) + + bounds = [ + (start, min(start + chunk_size, num_tokens)) + for start in range(0, num_tokens, chunk_size) + ] + if use_prims_ts: + # The next chunk's scoring and Top-K are queued before this chunk's + # PrimTS launch is prepared on the host, so the GPU has work while + # the host runs the wrapper. This relies on select() writing only + # its own packed_slice of top_k_output and visible_blocks; a shared + # scratch buffer inside select_qsa_paged_tokens would corrupt the + # pending chunk. + selected_next = select(*bounds[0]) + for i, (start, end) in enumerate(bounds): + selected = selected_next + if i + 1 < len(bounds): + selected_next = select(*bounds[i + 1]) + prims_ts_sparse_gqa( + q[start:end], + k_cache, + v_cache, + attn_metadata.qsa_block_table, + selected, + req_idx[start:end], + logical[start:end], + page_size=tokens_per_block, + max_seq_len_kv=max_seq_len_kv, + softmax_scale=softmax_scale, + is_decode=False, + request_boundaries=slice_request_boundaries(cu_seq_lens, start, end), + out=output[start:end], + ) + else: + for start, end in bounds: + selected = select(start, end) + output[start:end] = qsa_sparse_gqa( + q=q[start:end], + k_cache=k_cache, + v_cache=v_cache, + selected_tokens=selected, + request_indices=req_idx[start:end], + metadata=attn_metadata, + softmax_scale=softmax_scale, + query_positions=logical[start:end], + compress_ratio=params.compress_ratio, + ) output = output.reshape(num_tokens, -1) if output_gate is not None: output = attention.apply_output_gate(output, output_gate) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py index cc1e4c7ce596..203ef42a564f 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py @@ -7,6 +7,9 @@ from ..params import SparseMetadataParams, SparseParams from .constants import QSA_INDEX_KV_HEADS +QSA_SPARSE_CORE_BACKENDS = ("auto", "triton", "prims_ts") +QSA_INDEXER_BACKENDS = ("follow", *QSA_SPARSE_CORE_BACKENDS) + def _validate_selection_geometry(token_topk: int, compress_ratio: int) -> None: if token_topk <= 0: @@ -41,8 +44,24 @@ class QSASparseParams(SparseParams): enable_heuristic_topk: bool = False # Reuse one indexer selection across the MTP draft loop. mtp_index_share: bool = False + # Kernel selection is decided at runtime like the other sparse backends: + # PrimTS on SM100/SM103/SM107 except for small decode batches, Triton + # otherwise. These fields pin one kernel for tests and A/B measurements; + # the LLM API does not expose them. + sparse_core_backend: str = "auto" + indexer_backend: str = "follow" def __post_init__(self) -> None: + if self.sparse_core_backend not in QSA_SPARSE_CORE_BACKENDS: + raise ValueError( + f"QSA sparse_core_backend must be one of {QSA_SPARSE_CORE_BACKENDS}, " + f"got {self.sparse_core_backend!r}" + ) + if self.indexer_backend not in QSA_INDEXER_BACKENDS: + raise ValueError( + f"QSA indexer_backend must be one of {QSA_INDEXER_BACKENDS}, " + f"got {self.indexer_backend!r}" + ) values = { "index_n_heads": self.index_n_heads, "index_kv_heads": self.index_kv_heads, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py new file mode 100644 index 000000000000..c2c33e919f59 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py @@ -0,0 +1,459 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""PrimTS QToken-KvBlock-Sparse core attention for QSA.""" + +import bisect +import functools +import threading + +import torch + +from tensorrt_llm.logger import logger + +from .params import QSASparseParams + +_PREFILL_GROUP_SIZE = 4 +# Prefill plans are capacity plans over a route count: the route offsets are +# padded with trailing empty routes to a multiple of this, so the request +# composition of a chunk does not create a plan (the wrapper requires +# routes <= rows, so chunks shorter than the bucket plan per row count). +_PREFILL_ROUTE_BUCKET = 256 +# The wrapper re-prepares its launch, with a host readback of the route +# offsets, whenever the q/out shape changes. Full chunks share one shape; +# tail chunks are padded to this many rows so each padded shape keeps its +# own prepared plan instead of re-preparing on every tail. Decode rows are +# never padded. +_PREFILL_ROW_BUCKET = 64 +# Query rows (requests x query tokens per request, so MTP counts draft tokens) +# below which Triton is faster under CUDA graphs (Rubin SM107 and B300, 32K +# context, TP2): the core ties at one row, the indexer until ~32 rows. +_MIN_DECODE_ROWS = 2 +_MIN_INDEXER_DECODE_ROWS = 32 +_KV_BLOCK_SIZE = 4 +_SUPPORTED_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) +_SUPPORTED_Q_DTYPES = (torch.float16, torch.bfloat16) +_SUPPORTED_KV_DTYPES = (torch.float16, torch.bfloat16, torch.float8_e4m3fn) +_FP8_MAX = 448.0 + + +@functools.cache +def _prims_ts_api(): + try: + from ...prims_ts import ( + QTokenKvBlockSparsePagedTSWrapper, + get_q_token_kv_block_sparse_workspace_size, + make_q_token_kv_block_sparse_qo_indptr, + ) + except ImportError: + return None + return ( + QTokenKvBlockSparsePagedTSWrapper, + get_q_token_kv_block_sparse_workspace_size, + make_q_token_kv_block_sparse_qo_indptr, + ) + + +def prims_ts_available() -> bool: + return _prims_ts_api() is not None + + +def _unsupported_reason( + q: torch.Tensor, + k_cache: torch.Tensor, + page_size: int, +) -> str | None: + if not q.is_cuda: + return "requires CUDA tensors" + if torch.cuda.get_device_capability(q.device) not in _SUPPORTED_CAPABILITIES: + return "requires an SM100, SM103 or SM107 GPU" + if not prims_ts_available(): + return "the vendored PrimTS package is not importable" + if q.dtype not in _SUPPORTED_Q_DTYPES or k_cache.dtype not in _SUPPORTED_KV_DTYPES: + return f"unsupported dtypes q={q.dtype} kv={k_cache.dtype}" + if k_cache.dtype != q.dtype and k_cache.dtype != torch.float8_e4m3fn: + return f"mismatched dtypes q={q.dtype} kv={k_cache.dtype}" + if page_size % 4 != 0: + return f"page size {page_size} is not a multiple of the sparse block size 4" + if q.shape[-1] not in (64, 128, 256): + return f"unsupported head_dim {q.shape[-1]}" + return None + + +@functools.cache +def _prims_ts_indexer_available() -> bool: + try: + from . import prims_ts_indexer # noqa: F401 + except ImportError: + return False + return True + + +def use_prims_ts_indexer( + params: QSASparseParams, + q: torch.Tensor, + *, + tokens_per_block: int | None = None, + is_decode: bool = False, +) -> bool: + """Whether the indexer scoring runs on the PrimTS kernel instead of Triton. + + Follows ``sparse_core_backend`` unless ``indexer_backend`` is set: + ``triton`` never, ``prims_ts`` always (raising if unsupported), ``auto`` + when the device and inputs qualify and decode batches are large enough. + """ + backend = params.indexer_backend + if backend == "follow": + backend = params.sparse_core_backend + if backend == "triton": + return False + if backend == "auto" and is_decode and q.shape[0] < _MIN_INDEXER_DECODE_ROWS: + return False + reason = None + if not q.is_cuda: + reason = "requires CUDA tensors" + elif torch.cuda.get_device_capability(q.device) not in _SUPPORTED_CAPABILITIES: + reason = "requires an SM100, SM103 or SM107 GPU" + elif q.dtype != torch.bfloat16: + reason = f"requires BF16 index queries, got {q.dtype}" + elif tokens_per_block is not None and ( + tokens_per_block % params.compress_ratio != 0 + or 128 % (tokens_per_block // params.compress_ratio) != 0 + ): + reason = f"unsupported page size {tokens_per_block}" + elif not _prims_ts_indexer_available(): + reason = "the CUTLASS task-scheduling DSL is not importable" + if reason is None: + return True + if backend == "prims_ts": + raise RuntimeError(f"QSA PrimTS indexer {reason}") + logger.warning_once( + f"QSA PrimTS indexer unavailable ({reason}); using the Triton scorer", + key="qsa_prims_ts_indexer_fallback", + ) + return False + + +def resolve_sparse_core_backend( + params: QSASparseParams, + *, + q: torch.Tensor, + k_cache: torch.Tensor, + page_size: int, + is_decode: bool, + num_rows: int, +) -> str: + backend = params.sparse_core_backend + if backend == "triton": + return "triton" + reason = _unsupported_reason(q, k_cache, page_size) + if backend == "prims_ts": + if reason is not None: + raise RuntimeError(f"QSA sparse_core_backend=prims_ts {reason}") + return "prims_ts" + if reason is not None: + logger.warning_once( + f"QSA PrimTS sparse core unavailable ({reason}); using Triton", + key="qsa_prims_ts_unavailable", + ) + return "triton" + # Small decode batches are launch-bound; PrimTS pays warp-specialization + # barrier costs there and Triton stays faster. + if is_decode and num_rows < _MIN_DECODE_ROWS: + return "triton" + return "prims_ts" + + +def _request_boundaries(request_indices: torch.Tensor) -> tuple[int, ...]: + req = request_indices.to(torch.int64).cpu() + change = torch.nonzero(req[1:] != req[:-1]).flatten() + 1 + return (0, *change.tolist(), req.numel()) + + +def _pad_rows(t: torch.Tensor, rows: int, fill: float) -> torch.Tensor: + padded = t.new_full((rows, *t.shape[1:]), fill) + padded[: t.shape[0]] = t + return padded + + +def slice_request_boundaries(cu_seq_lens: list[int], start: int, end: int) -> tuple[int, ...]: + """Request boundaries of packed rows ``[start, end)`` from host cumulative lengths.""" + inner = cu_seq_lens[ + bisect.bisect_right(cu_seq_lens, start) : bisect.bisect_left(cu_seq_lens, end) + ] + return (0, *(offset - start for offset in inner), end - start) + + +def _run_abi(run_args) -> tuple: + return tuple((name, tuple(t.shape), tuple(t.stride()), t.dtype) for name, t in run_args) + + +class _PrimsTSSparseCore: + def __init__(self) -> None: + self._plans: dict[tuple, object] = {} + self._last_abi: dict[tuple, tuple] = {} + self._workspace: torch.Tensor | None = None + # Graphs captured before a growth keep the old buffer's address. + self._retired_workspaces: list[torch.Tensor] = [] + self._qo_indptr: dict[tuple, torch.Tensor] = {} + self._qo_indptr_key: tuple | None = None + self._qo_indptr_buffer_key: tuple | None = None + self._lock = threading.Lock() + + def _workspace_for(self, nbytes: int, device: torch.device) -> torch.Tensor: + ws = self._workspace + if ws is None or ws.device != device or ws.numel() < nbytes: + # Plans bind typed views into the workspace; growth invalidates them. + self._plans.clear() + if ws is not None: + self._retired_workspaces.append(ws) + ws = torch.zeros(nbytes, dtype=torch.uint8, device=device) + self._workspace = ws + return ws + + def _route_offsets( + self, + make_qo_indptr, + request_boundaries: tuple[int, ...], + rows: int, + device: torch.device, + ) -> torch.Tensor: + # The FlashInfer wrapper re-prepares, with host readbacks, whenever the + # qo_indptr storage changes, so offsets live in one persistent buffer + # per route count. The buffer is rewritten, stream-ordered, when the + # layout differs from the previous call: once per chunk per layer with + # several prefill chunks, never within a chunk's layers. + layout = (device.index, rows, request_boundaries) + if layout == self._qo_indptr_key: + return self._qo_indptr[self._qo_indptr_buffer_key] + host = make_qo_indptr( + torch.tensor(request_boundaries, dtype=torch.int32), + rows, + group_size=_PREFILL_GROUP_SIZE, + device="cpu", + ) + routes = min(rows, -(-(host.numel() - 1) // _PREFILL_ROUTE_BUCKET) * _PREFILL_ROUTE_BUCKET) + host = torch.cat([host, host[-1:].expand(routes + 1 - host.numel())]) + buffer_key = (device.index, host.numel()) + buffer = self._qo_indptr.get(buffer_key) + if buffer is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("QSA PrimTS route offsets must be allocated before capture") + buffer = torch.empty(host.numel(), dtype=torch.int32, device=device) + self._qo_indptr[buffer_key] = buffer + buffer.copy_(host, non_blocking=True) + self._qo_indptr_key = layout + self._qo_indptr_buffer_key = buffer_key + return buffer + + def run( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + block_ids: torch.Tensor, + request_indices: torch.Tensor, + query_positions: torch.Tensor, + *, + page_size: int, + max_seq_len_kv: int, + softmax_scale: float, + is_decode: bool, + request_boundaries: tuple[int, ...] | None = None, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + wrapper_cls, workspace_size, make_qo_indptr = _prims_ts_api() + rows, num_qo_heads, head_dim = q.shape + num_kv_heads = k_cache.shape[1] + block_topk = block_ids.shape[-1] + out_dtype = q.dtype + + if k_cache.dtype == torch.float8_e4m3fn and q.dtype != torch.float8_e4m3fn: + # The sparse route has no BF16-Q/FP8-KV kernel, so Q is cast to FP8 + # with unit scale: the normed Q of this model sits well inside the + # E4M3 range, and a per-call scale would need a host readback that + # CUDA-graph capture forbids. This differs numerically from the + # Triton path, which upcasts K/V. + q = q.clamp(-_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn) + + request_indices = request_indices.to(torch.int32).contiguous() + q = q.contiguous() + padded_rows = rows if is_decode else -(-rows // _PREFILL_ROW_BUCKET) * _PREFILL_ROW_BUCKET + if padded_rows != rows: + # Padding rows select no blocks and are sliced off the output. + q = _pad_rows(q, padded_rows, 0) + block_ids = _pad_rows(block_ids, padded_rows, -1) + request_indices = _pad_rows(request_indices, padded_rows, 0) + query_positions = _pad_rows(query_positions, padded_rows, 0) + if is_decode: + route_q = q.view(rows, 1, 1, num_qo_heads, head_dim) + qo_indptr = None + group_size = 1 + num_routes = rows + split_kv = True + else: + route_q = q + if request_boundaries is None: + # Host readback; callers with host sequence lengths pass the + # boundaries to keep the stream running. + request_boundaries = _request_boundaries(request_indices) + with self._lock: + qo_indptr = self._route_offsets( + make_qo_indptr, request_boundaries, padded_rows, q.device + ) + group_size = _PREFILL_GROUP_SIZE + num_routes = qo_indptr.numel() - 1 + split_kv = False + + # The wrapper retains K/V storage, and every QSA layer has its own cache + # pool view, so plans are per layer buffer. + key = ( + q.device.index, + k_cache.data_ptr(), + v_cache.data_ptr(), + padded_rows, + num_routes, + group_size, + is_decode, + max_seq_len_kv, + q.dtype, + k_cache.dtype, + out_dtype, + split_kv, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + block_topk, + ) + result = out + if out is None or padded_rows != rows: + out = torch.empty( + (padded_rows, num_qo_heads, head_dim), dtype=out_dtype, device=q.device + ) + with self._lock: + wrapper = self._plans.get(key) + if wrapper is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("QSA PrimTS plan must be prepared before CUDA-graph capture") + nbytes = workspace_size( + route_q, + k_cache, + block_table, + block_topk=block_topk, + max_seq_len_kv=max_seq_len_kv, + o_data_type=out_dtype, + qo_indptr=qo_indptr, + group_size=group_size, + kv_block_size=_KV_BLOCK_SIZE, + split_kv=split_kv, + ) + workspace = self._workspace_for(nbytes, q.device) + wrapper = wrapper_cls() + wrapper.plan( + num_routes, + num_qo_heads, + num_kv_heads, + head_dim, + _KV_BLOCK_SIZE, + page_size, + block_topk, + max_seq_len_kv, + device=q.device, + workspace_buffer=workspace, + group_size=group_size, + use_packed_q=not is_decode, + split_kv=split_kv, + q_data_type=q.dtype, + kv_data_type=k_cache.dtype, + o_data_type=out_dtype, + ) + self._plans[key] = wrapper + run_out = out.view_as(route_q) if is_decode else out + run_args = ( + ("q", route_q), + ("block_ids", block_ids), + ("block_table", block_table), + ("request_indices", request_indices), + ("query_positions", query_positions), + ("out", run_out), + ) + try: + wrapper.run( + route_q, + (k_cache, v_cache), + block_table, + block_ids, + request_indices, + query_positions, + qo_indptr=qo_indptr, + sm_scale=softmax_scale, + out=run_out, + ) + except RuntimeError as error: + previous = self._last_abi.get(key) + changed = ( + [(a, b) for a, b in zip(previous, _run_abi(run_args)) if a != b] + if previous + else [("no eager run for this plan key", None)] + ) + raise RuntimeError(f"{error}; ABI changes vs first eager run: {changed}") from error + if key not in self._last_abi: + self._last_abi[key] = _run_abi(run_args) + if result is None: + return out[:rows] + if result is not out: + result.copy_(out[:rows]) + return result + + +_CORE = _PrimsTSSparseCore() + + +def prims_ts_sparse_gqa( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + block_ids: torch.Tensor, + request_indices: torch.Tensor, + query_positions: torch.Tensor, + *, + page_size: int, + max_seq_len_kv: int, + softmax_scale: float, + is_decode: bool, + request_boundaries: tuple[int, ...] | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Sparse GQA over HND paged K/V from compact indexer block IDs. + + ``block_ids`` is the ``[rows, block_topk]`` Int32 Top-K output; the causal + tail is added by the kernel. Decode rows use the fixed ``[rows,1,1,Hq,D]`` + layout with split-KV; prefill packs rows into request-safe groups of four, + using host ``request_boundaries`` (row offsets of request starts plus the + row count) when given. ``out``, when given, receives the result in place. + """ + return _CORE.run( + q, + k_cache, + v_cache, + block_table, + block_ids, + request_indices, + query_positions, + page_size=page_size, + max_seq_len_kv=max_seq_len_kv, + softmax_scale=softmax_scale, + is_decode=is_decode, + request_boundaries=request_boundaries, + out=out, + ) + + +__all__ = [ + "prims_ts_available", + "prims_ts_sparse_gqa", + "resolve_sparse_core_backend", + "slice_request_boundaries", +] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/__init__.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/__init__.py new file mode 100644 index 000000000000..9521385b632f --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .api import prepare_qsa_index_scores, prims_ts_qsa_index_scores + +__all__ = ["prepare_qsa_index_scores", "prims_ts_qsa_index_scores"] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/api.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/api.py new file mode 100644 index 000000000000..33cbd4b3d33f --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/api.py @@ -0,0 +1,230 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Host entry for the paged task-scheduled QSA indexer scorer.""" + +import functools + +import torch + +from . import kernel as _k + +_COMPILE_OPTIONS = "--enable-tvm-ffi --opt-level 3" +# One CTA fits per SM; a CTA's prologue (page-id staging, Q load, ring fill) +# costs about as much as this many block-tile iterations. +_CTA_FIXED_ITERS = 2 + + +@functools.cache +def _num_sms(device_index: int) -> int: + return torch.cuda.get_device_properties(device_index).multi_processor_count + + +@functools.lru_cache(maxsize=1024) +def _pick_n_splits(tiles_bound: int, n_max: int, min_splits: int, num_sms: int) -> int: + """Split each tile's block range so whole CTA waves do the least serial work.""" + best, best_cost = min_splits, None + for n_splits in range(max(1, min_splits), n_max + 1): + waves = -(-tiles_bound * n_splits // num_sms) + iters = -(-n_max // n_splits) + cost = waves * (iters + _CTA_FIXED_ITERS) + if best_cost is None or cost < best_cost: + best, best_cost = n_splits, cost + if waves * _CTA_FIXED_ITERS >= best_cost: + break + return best + + +def _aligned(t: torch.Tensor) -> torch.Tensor: + # Row slices of the caller's packed metadata land on 4 B boundaries; the + # kernel's vector loads and TMA need 16 B. + t = t.contiguous() + return t if t.data_ptr() % 16 == 0 else t.clone() + + +def _paged_view(index_cache: torch.Tensor) -> torch.Tensor: + # The KV pool may coalesce the index cache with K/V, so the page stride is + # arbitrary; the kernel's TMA map takes it as a dimension. Pages must stay + # dense inside and 16 B aligned; anything else is copied once. + if ( + index_cache.stride()[1:] == (_k.HEAD_DIM, _k.HEAD_DIM, 1) + and index_cache.stride(0) * index_cache.element_size() % 16 == 0 + and index_cache.data_ptr() % 16 == 0 + ): + return index_cache + return index_cache.contiguous() + + +def _ptr(t, dtype): + import cutlass + from cutlass.cute.runtime import make_ptr + + return make_ptr(dtype, t.data_ptr(), cutlass.AddressSpace.gmem, assumed_align=16) + + +@functools.cache +def _compiled(blocks_per_page: int, compress_ratio: int): + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import make_fake_stream + + dev = torch.device("cuda", torch.cuda.current_device()) + i32 = torch.empty(16, dtype=torch.int32, device=dev) + f32 = torch.empty(16, dtype=torch.float32, device=dev) + bf16 = torch.empty(16, dtype=torch.bfloat16, device=dev) + return cute.compile( + _k.qsa_index_scores_host, + _ptr(bf16, cutlass.BFloat16), + _ptr(bf16, cutlass.BFloat16), + _ptr(i32, cutlass.Int32), + _ptr(i32, cutlass.Int32), + _ptr(i32, cutlass.Int32), + _ptr(i32, cutlass.Int32), + _ptr(f32, cutlass.Float32), + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Float32(1.0), + cutlass.Int32(1), + blocks_per_page, + compress_ratio, + make_fake_stream(), + options=_COMPILE_OPTIONS, + ) + + +def prepare_qsa_index_scores( + q: torch.Tensor, + index_cache: torch.Tensor, + block_table: torch.Tensor, + query_positions: torch.Tensor, + request_indices: torch.Tensor, + *, + tokens_per_block: int, + compress_ratio: int, + only_visible_blocks: bool = False, +): + """Build device metadata and return ``(launch, scores)``. + + Rows must be sorted by request. ``launch()`` runs the kernel on the current + stream without host synchronization and is CUDA-graph capturable; the + metadata built here depends only on request boundaries, so a captured + launch stays valid while those are unchanged. + """ + rows, num_heads, head_dim = q.shape + if (num_heads, head_dim) != (_k.NUM_HEADS, _k.HEAD_DIM): + raise ValueError("prims_ts QSA scorer supports exactly 4 heads x 128 dims") + if compress_ratio < 2 or tokens_per_block % compress_ratio != 0: + raise ValueError("prims_ts QSA scorer requires compress_ratio >= 2 dividing the page size") + blocks_per_page = tokens_per_block // compress_ratio + if blocks_per_page > _k.TILE_N or _k.TILE_N % blocks_per_page != 0: + raise ValueError("compressed blocks per page must divide 128") + if ( + index_cache.ndim != 4 + or index_cache.shape[1] != tokens_per_block + or index_cache.shape[2] != 1 + ): + raise ValueError("index_cache must be [pages, tokens_per_block, 1, 128]") + if q.dtype != torch.bfloat16 or index_cache.dtype != torch.bfloat16: + raise ValueError("prims_ts QSA scorer requires BF16 q and index cache") + if block_table.dtype != torch.int32 or block_table.ndim != 2: + raise ValueError("block_table must be int32 [requests, max_pages]") + if request_indices.shape != (rows,) or query_positions.shape != (rows,): + raise ValueError("request_indices and query_positions must have one entry per row") + if request_indices.dtype != torch.int32: + raise ValueError("request_indices must be int32") + device = q.device + num_requests = block_table.shape[0] + max_blocks = block_table.shape[1] * blocks_per_page + tile_meta = torch.empty(2 * num_requests, dtype=torch.int32, device=device) + request_indices = _aligned(request_indices) + tiles_bound = (rows + _k.QUERY_ROWS_PER_TILE - 1) // _k.QUERY_ROWS_PER_TILE + num_requests + n_max = max(1, (max_blocks + _k.TILE_N - 1) // _k.TILE_N) + pages_per_tile = _k.TILE_N // blocks_per_page + min_splits = (n_max * pages_per_tile + _k.PAGE_ID_CAPACITY - 1) // _k.PAGE_ID_CAPACITY + num_sms = _num_sms(device.index if device.index is not None else torch.cuda.current_device()) + n_splits = _pick_n_splits(tiles_bound, n_max, min_splits, num_sms) + grid_size = tiles_bound * n_splits + scores = torch.empty(rows, max_blocks, dtype=torch.float32, device=device) + positions = _aligned( + query_positions if query_positions.dtype == torch.int32 else query_positions.to(torch.int32) + ) + fn = _compiled(blocks_per_page, compress_ratio) + import cuda.bindings.driver as cuda_driver + + block_table = _aligned(block_table) + q = _aligned(q) + index_cache = _paged_view(index_cache) + # Raw pointers go into the compiled call; these tensors own the storage. + keepalive = ( + q, + index_cache, + block_table, + request_indices, + tile_meta, + positions, + scores, + ) + # The TVM-FFI entry takes raw addresses and Python scalars; building + # cutlass pointer and numeric objects here would cost ~11 us per call. + args = ( + q.data_ptr(), + index_cache.data_ptr(), + block_table.data_ptr(), + request_indices.data_ptr(), + tile_meta.data_ptr(), + positions.data_ptr(), + scores.data_ptr(), + num_requests, + block_table.shape[1], + max_blocks, + index_cache.shape[0], + index_cache.stride(0), + rows, + n_splits, + head_dim**-0.5, + grid_size, + ) + + def launch(): + if not only_visible_blocks: + scores.fill_(float("-inf")) + fn(*args, cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream)) + return scores + + launch.keepalive = keepalive + return launch, scores + + +def prims_ts_qsa_index_scores( + q: torch.Tensor, + index_cache: torch.Tensor, + block_table: torch.Tensor, + query_positions: torch.Tensor, + request_indices: torch.Tensor, + *, + tokens_per_block: int, + compress_ratio: int, + only_visible_blocks: bool = False, +) -> torch.Tensor: + """Score compressed keys for each query row; same contract as the Triton scorer. + + Rows must be ordered by request. Columns at or past each row's causal block + bound hold ``-inf`` unless ``only_visible_blocks``, which leaves them + unspecified. + """ + launch, _ = prepare_qsa_index_scores( + q, + index_cache, + block_table, + query_positions, + request_indices, + tokens_per_block=tokens_per_block, + compress_ratio=compress_ratio, + only_visible_blocks=only_visible_blocks, + ) + return launch() diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/kernel.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/kernel.py new file mode 100644 index 000000000000..7a15367b0bec --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/kernel.py @@ -0,0 +1,1046 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Task-scheduled QSA indexer scoring over the paged index cache. + +A CTA owns one tile of 32 consecutive query rows of one request (x 4 indexer +heads interleaved into the MMA M axis), holds its Q tile in SMEM and streams a +range of 128-block key tiles through a TmaUmma ring: + + SmemQResource, SmemKResource -> TmemSResource -> GmemScoreResource + +Tiles are derived on device from the request-sorted rows: request ``r`` owns +rows ``[req_start[r], req_start[r+1])`` and tile ``t`` of the launch maps to +request ``r`` with ``tile_offsets[r-1] <= t < tile_offsets[r]``. The load task +TMA-fills Q once and one 128-block x 128-dim K-bar stage per iteration (a box +per 8-row core group, both 64-dim halves adjacent, resolved through the block +table); the MMA task accumulates S in a two-stage TMEM ring; the store task +applies scale and ReLU, sums the four head rows of each query with warp +shuffles and writes the fp32 scores straight from registers. +""" + +from dataclasses import dataclass, field + +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +import cutlass.pipeline as pipeline +from cutlass.experimental import primitives as prims +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import ( + SmemAllocation, + SmemAllocator, + TmemAllocation, + TmemAllocator, +) +from cutlass.experimental.task_scheduling.resources import ( + MemoryResource, + PipelineConfig, + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) +from cutlass.experimental.task_scheduling.schedule_builder import domain_loop, schedule +from cutlass.experimental.task_scheduling.task import Task +from cutlass.experimental.task_scheduling.task_manager import TaskManager + +from tensorrt_llm._torch.attention.backends.prims_ts.kernels.placeholder_helpers import ( + _placeholder_smem_array, +) + +io_dtype = cutlass.BFloat16 +acc_dtype = cutlass.Float32 + +HEAD_DIM = 128 +NUM_HEADS = 4 +QUERY_ROWS_PER_TILE = 32 +TILE_M = QUERY_ROWS_PER_TILE * NUM_HEADS +TILE_N = 128 +STAGE_K = 64 +Q_HALVES = HEAD_DIM // STAGE_K +MMA_K = 16 +SMEM_STAGES = 5 +# SW128 K-major core group: 8 rows x 128 B. A K stage holds all 128 dims of +# 128 blocks; when a page holds at least 8 compressed rows each TMA box +# brings one 8-row group with both 64-dim halves adjacent (16 boxes per +# stage), otherwise each half of the stage is filled page by page. +GROUP_ROWS = 8 +TMEM_STAGES = 2 +TMEM_COLS = TILE_N * TMEM_STAGES +T2R_COLS = 32 +NUM_SUBTILES = TILE_N // T2R_COLS +COLS_PER_LANE = T2R_COLS // NUM_HEADS +SEARCH_STEPS = 16 +# Named barriers below 12 are reserved for the task-scheduling runtime; 12 +# and 13 bracket the task run. +NEG_INF = float("-inf") +# Page ids for the CTA's whole block-tile range are staged in SMEM once; the +# host bounds n_splits so a CTA never needs more than this many. +PAGE_ID_CAPACITY = 4096 + +_Q_STAGE_ELEMS = TILE_M * STAGE_K +_Q_ELEMS = TILE_M * HEAD_DIM +_K_STAGE_ELEMS = TILE_N * HEAD_DIM +_K_GROUP_ELEMS = GROUP_ROWS * HEAD_DIM +_K_HALF_ELEMS = TILE_N * STAGE_K +_Q_BYTES = _Q_ELEMS * io_dtype.width // 8 +_K_STAGE_BYTES = _K_STAGE_ELEMS * io_dtype.width // 8 + + +@cute.jit +def _tile_coords( + tile_offsets: cute.Tensor, + req_start: cute.Tensor, + positions: cute.Tensor, + num_rows: cutlass.Int32, + n_splits: cutlass.Int32, + compress_ratio: cutlass.Constexpr[int], +): + """Resolve this CTA's (request, first row, row count, block-tile range). + + Launch tiles past the real tile count repeat the last tile; their stores + are idempotent. + """ + bx, _, _ = cute.arch.block_idx() + num_requests = cutlass.Int32(tile_offsets.shape[0]) + total_tiles = cutlass.Int32(tile_offsets[num_requests - 1]) + tile = bx // n_splits + if tile >= total_tiles: + tile = total_tiles - 1 + split = bx % n_splits + lo = cutlass.Int32(0) + hi = num_requests + for _ in cutlass.range_constexpr(SEARCH_STEPS): + mid = (lo + hi) >> 1 + probe = mid + if probe > num_requests - 1: + probe = num_requests - 1 + if cutlass.Int32(tile_offsets[probe]) > tile: + hi = mid + else: + lo = mid + 1 + request = lo + if request > num_requests - 1: + request = num_requests - 1 + tiles_before = cutlass.Int32(0) + if request > 0: + tiles_before = cutlass.Int32(tile_offsets[request - 1]) + row_begin = cutlass.Int32(req_start[request]) + row_end = num_rows + if request + 1 < num_requests: + row_end = cutlass.Int32(req_start[request + 1]) + first_row = row_begin + (tile - tiles_before) * QUERY_ROWS_PER_TILE + tile_rows = row_end - first_row + if tile_rows > QUERY_ROWS_PER_TILE: + tile_rows = QUERY_ROWS_PER_TILE + if tile_rows < 0: + tile_rows = cutlass.Int32(0) + lane = cute.arch.lane_idx() + visible = cutlass.Int32(0) + if lane < tile_rows: + visible = (cutlass.Int32(positions[first_row + lane]) + 1) // compress_ratio + tile_visible = cutlass.Int32(cute.arch.warp_redux_sync(visible, "max")) + n_count = (tile_visible + TILE_N - 1) // TILE_N + chunk = (n_count + n_splits - 1) // n_splits + n_begin = split * chunk + n_end = n_begin + chunk + if n_end > n_count: + n_end = n_count + n_iters = n_end - n_begin + if n_iters < 1: + n_iters = cutlass.Int32(1) + n_begin = n_count - 1 + if n_begin < 0: + n_begin = cutlass.Int32(0) + return request, first_row, tile_rows, n_begin, n_iters + + +@dataclass +class SmemQResource(MemoryResource): + """The CTA's Q tile, TMA-loaded once and held for every K tile.""" + + tma_desc_q: cutlass.Pointer = field(init=False, default=None) + first_row: cutlass.Int32 = field(init=False, default=None) + smem_q: cutlass.Array = field(init=False, default=None) + desc_q_base: cutlass.Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + _alloc_q: cutlass.Constexpr = field(init=False, default=None) + + def __init__( + self, tma_desc_q: cutlass.Pointer, first_row: cutlass.Int32, **kwargs: object + ) -> None: + super().__init__(**kwargs) + self.tma_desc_q = tma_desc_q + self.first_row = first_row + self.desc_q_base = TaskLocalVariable(dtype=cutlass.Int64, default=cutlass.Int64(0)) + self.smem_q = _placeholder_smem_array(io_dtype, _Q_ELEMS) + self._alloc_q = SmemAllocation("smem_q", _Q_BYTES, alignment=1024) + + def get_smem_requirements(self): + return [self._alloc_q] + + @cute.jit + def _init_state(self, stage_info: StageInfo) -> None: + self.smem_q = cutlass.Array( + stage_info.context.smem_base.data_ptr() + self._alloc_q.offset, + dtype=io_dtype, + shape=(_Q_ELEMS,), + addrspace=3, + ) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + self._init_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_mma_state(self, stage_info: StageInfo) -> None: + self._init_state(stage_info) + + @producer_work + @cute.jit + def tma_load(self, stage_info: StageInfo) -> None: + if prims.elect_sync(): + for half in cutlass.range_constexpr(Q_HALVES): + prims.cp_async_bulk_tensor_shared_cta_global( + self.smem_q.subview(_Q_STAGE_ELEMS * half), + self.tma_desc_q, + (cutlass.Int32(half * STAGE_K), self.first_row * NUM_HEADS), + stage_info.barrier, + ) + + @consumer_work(returns=desc_q_base) + @cute.jit + def build_descriptor(self, stage_info: StageInfo) -> cutlass.Int64: + return prims.Tcgen05SmemDesc.build( + self.smem_q, leading_byte_offset=16, stride_byte_offset=1024, layout=2 + ) + + +@dataclass +class SmemKResource(MemoryResource): + """Ring of K-bar operand tiles: one 128-block x 64-dim stage per TMA group.""" + + tma_desc_k: cutlass.Pointer = field(init=False, default=None) + block_table: cute.Tensor = field(init=False, default=None) + blocks_per_page: cutlass.Constexpr[int] = field(init=False, default=None) + compress_ratio: cutlass.Constexpr[int] = field(init=False, default=None) + merged_halves: cutlass.Constexpr[bool] = field(init=False, default=None) + stride_byte_offset: cutlass.Constexpr[int] = field(init=False, default=None) + half_bytes: cutlass.Constexpr[int] = field(init=False, default=None) + smem_k: cutlass.Array = field(init=False, default=None) + smem_pages: cutlass.Array = field(init=False, default=None) + request: cutlass.Int32 = field(init=False, default=None) + n_begin: cutlass.Int32 = field(init=False, default=None) + n_iters: cutlass.Int32 = field(init=False, default=None) + desc_k_base: cutlass.Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + _alloc_k: cutlass.Constexpr = field(init=False, default=None) + _alloc_pages: cutlass.Constexpr = field(init=False, default=None) + + def __init__( + self, + tma_desc_k: cutlass.Pointer, + block_table: cute.Tensor, + request: cutlass.Int32, + n_begin: cutlass.Int32, + n_iters: cutlass.Int32, + blocks_per_page: int, + compress_ratio: int, + **kwargs: object, + ) -> None: + super().__init__(**kwargs) + self.tma_desc_k = tma_desc_k + self.block_table = block_table + self.blocks_per_page = blocks_per_page + self.compress_ratio = compress_ratio + self.merged_halves = blocks_per_page >= GROUP_ROWS + self.stride_byte_offset = ( + GROUP_ROWS * HEAD_DIM * 2 if self.merged_halves else GROUP_ROWS * STAGE_K * 2 + ) + self.half_bytes = GROUP_ROWS * STAGE_K * 2 if self.merged_halves else _K_HALF_ELEMS * 2 + self.request = request + self.n_begin = n_begin + self.n_iters = n_iters + self.desc_k_base = TaskLocalVariable(dtype=cutlass.Int64, default=cutlass.Int64(0)) + # Placeholders keep the field structure stable across the warp-gated + # task dispatch; the real views are bound in init work. + self.smem_k = _placeholder_smem_array(io_dtype, _K_STAGE_ELEMS * SMEM_STAGES) + self.smem_pages = _placeholder_smem_array(cutlass.Int32, PAGE_ID_CAPACITY) + self._alloc_k = SmemAllocation("smem_k", _K_STAGE_BYTES * SMEM_STAGES, alignment=1024) + self._alloc_pages = SmemAllocation("smem_pages", PAGE_ID_CAPACITY * 4, alignment=16) + + def get_smem_requirements(self): + return [self._alloc_k, self._alloc_pages] + + @cute.jit + def _init_state(self, stage_info: StageInfo) -> None: + smem_base = stage_info.context.smem_base + self.smem_k = cutlass.Array( + smem_base.data_ptr() + self._alloc_k.offset, + dtype=io_dtype, + shape=(_K_STAGE_ELEMS * SMEM_STAGES,), + addrspace=3, + ) + self.smem_pages = cutlass.Array( + smem_base.data_ptr() + self._alloc_pages.offset, + dtype=cutlass.Int32, + shape=(PAGE_ID_CAPACITY,), + addrspace=3, + ) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + self._init_state(stage_info) + # Stage the physical page of every block tile this CTA will visit. + # Pages past the table or unallocated (negative) entries are masked + # in the epilogue; clamp so the TMA stays in bounds. + pages_per_tile = TILE_N // self.blocks_per_page + max_pages = cutlass.Int32(self.block_table.shape[1]) + total = self.n_iters * pages_per_tile + lane = cute.arch.lane_idx() + for i in cutlass.range(lane, total, 32): + logical_page = self.n_begin * pages_per_tile + i + if logical_page >= max_pages: + logical_page = max_pages - 1 + physical_page = cutlass.Int32(self.block_table[(self.request, logical_page)]) + if physical_page < 0: + physical_page = cutlass.Int32(0) + self.smem_pages[i] = physical_page + cute.arch.sync_warp() + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_mma_state(self, stage_info: StageInfo) -> None: + self._init_state(stage_info) + + @consumer_work(returns=desc_k_base) + @cute.jit + def build_descriptor(self, stage_info: StageInfo) -> cutlass.Int64: + sk = self.smem_k.subview(_K_STAGE_ELEMS * stage_info.stage_idx) + return prims.Tcgen05SmemDesc.build( + sk, + leading_byte_offset=16, + stride_byte_offset=self.stride_byte_offset, + layout=2, + ) + + @producer_work + @cute.jit + def tma_load(self, stage_info: StageInfo) -> None: + sk = self.smem_k.subview(_K_STAGE_ELEMS * stage_info.stage_idx) + pages_per_tile = TILE_N // self.blocks_per_page + page_base = stage_info.loop_offset * pages_per_tile + lane = cute.arch.lane_idx() + if cutlass.const_expr(self.merged_halves): + groups_per_page = self.blocks_per_page // GROUP_ROWS + if lane < TILE_N // GROUP_ROWS: + physical_page = cutlass.Int32(self.smem_pages[page_base + lane // groups_per_page]) + row0 = cutlass.Int32((lane % groups_per_page) * GROUP_ROWS) + prims.cp_async_bulk_tensor_shared_cta_global( + sk.subview(_K_GROUP_ELEMS * lane), + self.tma_desc_k, + (cutlass.Int32(0), row0, cutlass.Int32(0), physical_page), + stage_info.barrier, + ) + else: + if lane < pages_per_tile: + physical_page = cutlass.Int32(self.smem_pages[page_base + lane]) + for half in cutlass.range_constexpr(Q_HALVES): + prims.cp_async_bulk_tensor_shared_cta_global( + sk.subview(_K_HALF_ELEMS * half + self.blocks_per_page * STAGE_K * lane), + self.tma_desc_k, + ( + cutlass.Int32(half * STAGE_K), + cutlass.Int32(self.compress_ratio - 1), + cutlass.Int32(0), + physical_page, + ), + stage_info.barrier, + ) + + +@dataclass +class TmemSResource(MemoryResource): + """FP32 score accumulator ring in TMEM between the MMA and store tasks.""" + + tmem_raw_addr: cutlass.Int32 = field(init=False, default=None) + _alloc_acc: cutlass.Constexpr = field(init=False, default=None) + t2r_rmem: cutlass.Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def __init__(self, **kwargs: object) -> None: + super().__init__(**kwargs) + self._alloc_acc = TmemAllocation("tmem_s", TMEM_COLS) + self.tmem_raw_addr = cutlass.Int32(0) + self.t2r_rmem = TaskLocalVariable( + dtype=cutlass.Float32, + default=cutlass.full([T2R_COLS], 0.0, cutlass.Float32), + ) + + def get_tmem_requirements(self): + return [self._alloc_acc] + + @cute.jit + def _init_tmem_state(self, stage_info: StageInfo) -> None: + self.tmem_raw_addr = cutlass.Int32(stage_info.context.tmem_ptr_i32.load()) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_accumulator_state(self, stage_info: StageInfo) -> None: + self._init_tmem_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_store_state(self, stage_info: StageInfo) -> None: + self._init_tmem_state(stage_info) + + @consumer_work(returns=t2r_rmem) + @cute.jit + def load_subtile( + self, stage_info: StageInfo, *, subtile_idx: cutlass.Constexpr[int] + ) -> cutlass.Float32: + # Routed by variable name, so one 32-column load per call; the four + # round trips per tile are hidden behind the loader, not the limiter. + base_col = ( + (self.tmem_raw_addr & 0xFFFF) + stage_info.stage_idx * TILE_N + subtile_idx * T2R_COLS + ) + tmem_ptr = prims.make_tmem_ptr_from_warp_row_col( + self.tmem_raw_addr, cute.arch.warp_idx(), base_col, cutlass.Float32 + ) + t2r_rmem = prims.tcgen05_ld("32x32b", tmem_ptr, num=T2R_COLS) + cute.arch.fence_view_async_tmem_load() + return t2r_rmem + + @producer_work + @cute.jit + def mma( + self, + stage_info: StageInfo, + *, + desc_q_base: cutlass.Int64, + desc_k_base: cutlass.Int64, + k_half_bytes: cutlass.Constexpr[int], + ) -> None: + idesc = prims.Tcgen05InstrDesc.build( + a_dtype=io_dtype, + b_dtype=io_dtype, + c_dtype=cutlass.Float32, + n_dim=TILE_N, + m_dim=TILE_M, + ) + tmem_ptr = prims.make_tmem_ptr( + self.tmem_raw_addr + stage_info.stage_idx * TILE_N, acc_dtype + ) + tmem_acc = cutlass.Array(tmem_ptr.data_ptr(), dtype=cutlass.Int32, addrspace=6) + blocks_per_half = STAGE_K // MMA_K + for k_block in cutlass.range_constexpr(HEAD_DIM // MMA_K): + half = k_block // blocks_per_half + increment = (MMA_K * io_dtype.width // 8 * (k_block % blocks_per_half)) >> 4 + q_offset = ((_Q_STAGE_ELEMS * io_dtype.width // 8) * half) >> 4 + k_offset = (k_half_bytes * half) >> 4 + accumulate = cutlass.Boolean(k_block > 0) + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, + prims.CTAGroup.CTA_1, + tmem_acc, + desc_q_base + q_offset + increment, + desc_k_base + k_offset + increment, + idesc, + accumulate, + ) + + +@dataclass +class GmemScoreResource(MemoryResource): + """Score sink. + + Relu, head-sum, scale and causal mask straight from registers into the + fp32 score rows with 16-byte vector stores. + """ + + scores: cute.Tensor = field(init=False, default=None) + positions: cute.Tensor = field(init=False, default=None) + scale: cutlass.Float32 = field(init=False, default=None) + compress_ratio: cutlass.Constexpr[int] = field(init=False, default=None) + first_row: cutlass.Int32 = field(init=False, default=None) + tile_rows: cutlass.Int32 = field(init=False, default=None) + n_begin: cutlass.Int32 = field(init=False, default=None) + visible: cutlass.Int32 = field(init=False, default=None) + row_ptr: cutlass.Int64 = field(init=False, default=None) + + def __init__( + self, + scores: cute.Tensor, + positions: cute.Tensor, + scale: cutlass.Float32, + first_row: cutlass.Int32, + tile_rows: cutlass.Int32, + n_begin: cutlass.Int32, + compress_ratio: int, + **kwargs: object, + ) -> None: + super().__init__(**kwargs) + self.scores = scores + self.positions = positions + self.scale = scale + self.compress_ratio = compress_ratio + self.first_row = first_row + self.tile_rows = tile_rows + self.n_begin = n_begin + self.visible = cutlass.Int32(0) + self.row_ptr = cutlass.Int64(0) + + def get_smem_requirements(self): + return [] + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_store_state(self, stage_info: StageInfo) -> None: + tx, _, _ = cute.arch.thread_idx() + tile_row = tx // NUM_HEADS + self.visible = cutlass.Int32(0) + if tile_row < self.tile_rows: + self.visible = ( + cutlass.Int32(self.positions[self.first_row + tile_row]) + 1 + ) // self.compress_ratio + num_cols = cutlass.Int64(self.scores.shape[1]) + self.row_ptr = self.scores.iterator.toint() + ( + cutlass.Int64(self.first_row + tile_row) * num_cols * (acc_dtype.width // 8) + ) + + @producer_work + @cute.jit + def stage( + self, + stage_info: StageInfo, + *, + t2r_rmem: cutlass.Float32, + subtile_idx: cutlass.Constexpr[int], + ) -> None: + tx, _, _ = cute.arch.thread_idx() + lane = tx % 32 + upper_half = (lane & 1) != 0 + upper_quarter = (lane & 2) != 0 + # The four head rows of one query sit in four adjacent lanes. Reduce- + # scatter over the two lane bits so each lane ends with 8 fully summed + # columns: lanes with bit0 set keep columns 16..31, then bit1 selects + # the upper 8 of those 16. + half = cutlass.Array(cutlass.Float32, 16, space=cutlass.AddressSpace.rmem) + for i in cutlass.range_constexpr(16): + lo = t2r_rmem[i] + hi = t2r_rmem[16 + i] + if lo < cutlass.Float32(0.0): + lo = cutlass.Float32(0.0) + if hi < cutlass.Float32(0.0): + hi = cutlass.Float32(0.0) + send = hi + keep = lo + if upper_half: + send = lo + keep = hi + recv = cutlass.Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=send, + offset=1, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ) + half[i] = keep + recv + quarter = cutlass.Array(cutlass.Float32, COLS_PER_LANE, space=cutlass.AddressSpace.rmem) + n_tile = self.n_begin + stage_info.loop_offset + col0 = n_tile * TILE_N + subtile_idx * T2R_COLS + (lane & 1) * 16 + (lane & 2) * 4 + for i in cutlass.range_constexpr(COLS_PER_LANE): + lo = half[i] + hi = half[COLS_PER_LANE + i] + send = hi + keep = lo + if upper_quarter: + send = lo + keep = hi + recv = cutlass.Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=send, + offset=2, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ) + value = (keep + recv) * self.scale + if col0 + i >= self.visible: + value = cutlass.Float32(NEG_INF) + quarter[i] = value + num_cols = cutlass.Int32(self.scores.shape[1]) + tile_row = tx // NUM_HEADS + if tile_row < self.tile_rows and col0 < num_cols: + dst = cutlass.inttoptr( + self.row_ptr + cutlass.Int64(col0) * (acc_dtype.width // 8), + mem_space=1, + dtype=cutlass.Float32, + ) + if col0 + COLS_PER_LANE <= num_cols: + dst.store(quarter.data_ptr().load(count=4, alignment=16), alignment=16) + (dst + 4).store((quarter.data_ptr() + 4).load(count=4, alignment=16), alignment=16) + else: + for i in cutlass.range_constexpr(COLS_PER_LANE): + if col0 + i < num_cols: + (dst + i).store(quarter[i]) + + +@cute.jit +def create_smem_q_resource(tma_desc_q, first_row) -> SmemQResource: + cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=1, + num_bytes=_Q_BYTES, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + cta_layout_vmnk=(1, 1, 1, 1), + ) + return SmemQResource( + tma_desc_q=tma_desc_q, first_row=first_row, pipeline_config=cfg, name="SmemQ" + ) + + +@cute.jit +def create_smem_k_resource( + tma_desc_k, block_table, request, n_begin, n_iters, blocks_per_page, compress_ratio +) -> SmemKResource: + cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=SMEM_STAGES, + num_bytes=_K_STAGE_BYTES, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + cta_layout_vmnk=(1, 1, 1, 1), + ) + return SmemKResource( + tma_desc_k=tma_desc_k, + block_table=block_table, + request=request, + n_begin=n_begin, + n_iters=n_iters, + blocks_per_page=blocks_per_page, + compress_ratio=compress_ratio, + pipeline_config=cfg, + name="SmemK", + ) + + +@cute.jit +def create_tmem_s_resource(num_epilogue_warps: int) -> TmemSResource: + cfg = PipelineConfig.create_umma_async_pipeline_cfg( + num_stages=TMEM_STAGES, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, size=num_epilogue_warps * 32 + ), + cta_layout_vmnk=(1, 1, 1, 1), + ) + return TmemSResource(pipeline_config=cfg, name="TmemS") + + +@cute.jit +def create_load_task(smem_q, smem_k, n_iters, num_load_warps: int) -> Task: + @schedule + def load_schedule(q: MemoryResource, k: MemoryResource) -> None: + q.init_load_state() + k.init_load_state() + q.try_acquire() + q.acquire() + q.tma_load() + q.commit() + with domain_loop(0, n_iters, 1, unroll=1): + k.try_acquire() + k.acquire() + k.tma_load() + k.commit() + + return Task( + src_resources=[], + dst_resources=[smem_q, smem_k], + warp_idx=4, + num_warps=num_load_warps, + schedule=load_schedule(smem_q, smem_k), + num_registers=40, + name="LoadTask", + ) + + +@cute.jit +def create_mma_task(smem_q, smem_k, tmem_s, n_iters, num_mma_warps: int) -> Task: + k_half_bytes = smem_k.half_bytes + + @schedule + def mma_schedule(q: MemoryResource, k: MemoryResource, s: MemoryResource) -> None: + q.init_mma_state() + k.init_mma_state() + s.init_accumulator_state() + q.try_wait() + q.wait() + desc_q_base = q.build_descriptor() + with domain_loop(0, n_iters, 1, unroll=1): + s.try_acquire() + s.acquire() + k.try_wait() + k.wait() + desc_k_base = k.build_descriptor() + s.mma( + desc_q_base=desc_q_base, + desc_k_base=desc_k_base, + k_half_bytes=k_half_bytes, + ) + k.release() + s.commit() + q.release() + + return Task( + src_resources=[smem_q, smem_k], + dst_resources=[tmem_s], + warp_idx=5, + num_warps=num_mma_warps, + schedule=mma_schedule(smem_q, smem_k, tmem_s), + num_registers=40, + name="MmaTask", + ) + + +@cute.jit +def create_store_task(tmem_s, gmem_scores, n_iters, num_epilogue_warps: int) -> Task: + @schedule + def store_schedule(s: MemoryResource, out: MemoryResource) -> None: + s.init_store_state() + out.init_store_state() + with domain_loop(0, n_iters, 1, unroll=1): + s.try_wait() + s.wait() + for subtile_idx in cutlass.range_constexpr(NUM_SUBTILES): + t2r_rmem = s.load_subtile(subtile_idx=subtile_idx) + out.stage(t2r_rmem=t2r_rmem, subtile_idx=subtile_idx) + s.release() + + return Task( + src_resources=[tmem_s], + dst_resources=[gmem_scores], + warp_idx=0, + num_warps=num_epilogue_warps, + schedule=store_schedule(tmem_s, gmem_scores), + num_registers=168, + name="StoreTask", + ) + + +@cute.jit +def create_padding_task(n_iters, num_padding_warps: int, first_warp: int) -> Task: + @schedule + def padding_schedule() -> None: + with domain_loop(0, n_iters, 1, unroll=1): + pass + + return Task( + src_resources=[], + dst_resources=[], + warp_idx=first_warp, + num_warps=num_padding_warps, + schedule=padding_schedule(), + num_registers=40, + name="PaddingTask", + ) + + +@cute.jit +def _build_pipeline( + tma_desc_q, + tma_desc_k, + block_table, + positions, + scores, + scale, + request, + first_row, + tile_rows, + n_begin, + n_iters, + blocks_per_page, + compress_ratio, +): + num_epilogue_warps = 4 + num_load_warps = 1 + num_mma_warps = 1 + + allocator = SmemAllocator() + smem_q = create_smem_q_resource(tma_desc_q, first_row) + smem_k = create_smem_k_resource( + tma_desc_k, + block_table, + request, + n_begin, + n_iters, + blocks_per_page, + compress_ratio, + ) + tmem_s = create_tmem_s_resource(num_epilogue_warps) + gmem_scores = GmemScoreResource( + scores=scores, + positions=positions, + scale=scale, + first_row=first_row, + tile_rows=tile_rows, + n_begin=n_begin, + compress_ratio=compress_ratio, + name="GmemScores", + ) + allocator.add_resource(smem_q) + allocator.add_resource(smem_k) + allocator.add_tmem_ptr(SmemAllocation("tmem_ptr_i32", dtype=cutlass.Int32, alignment=4)) + allocator.compute_layout() + allocator.allocate() + tmem_ptr_i32 = allocator.get(allocator.tmem_ptr_alloc) + tmem_allocator = TmemAllocator() + tmem_allocator.add_resource(tmem_s) + tmem_allocator.compute_layout() + + total_warps = num_epilogue_warps + num_load_warps + num_mma_warps + num_padding_warps = (total_warps + 3) // 4 * 4 - total_warps + tasks = [ + create_load_task(smem_q, smem_k, n_iters, num_load_warps), + create_mma_task(smem_q, smem_k, tmem_s, n_iters, num_mma_warps), + create_store_task(tmem_s, gmem_scores, n_iters, num_epilogue_warps), + ] + if cutlass.const_expr(num_padding_warps > 0): + tasks.append(create_padding_task(n_iters, num_padding_warps, total_warps)) + task_manager = TaskManager( + tasks=tasks, + resource_dependency_graph={ + tmem_s: [smem_q, smem_k], + gmem_scores: [tmem_s], + }, + smem_allocator=allocator, + tmem_allocator=tmem_allocator, + tmem_ptr_i32=tmem_ptr_i32, + verbose=False, + exhaustive_deadlock_race_check=False, + ) + return task_manager, tmem_ptr_i32 + + +@cute.jit +def _run(task_manager, warp_idx, tmem_ptr_i32) -> None: + task_manager.setup_resources_and_tasks() + prims.fence_mbarrier_init() + prims.barrier_cta_sync(0) + if warp_idx == 0: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + prims.barrier_cta_sync(12) + task_manager.run() + prims.barrier_cta_sync(13) + if warp_idx == 0: + tmem_ptr = prims.make_tmem_ptr(tmem_ptr_i32.load(), acc_dtype) + prims.tcgen05_dealloc(tmem_ptr, TMEM_COLS) + + +@cute.kernel +def qsa_index_scores_kernel( + tma_desc_q: cutlass.GridConstant[cuda.TensorMap], + tma_desc_k: cutlass.GridConstant[cuda.TensorMap], + block_table: cute.Tensor, + tile_offsets: cute.Tensor, + req_start: cute.Tensor, + positions: cute.Tensor, + scores: cute.Tensor, + num_rows: cutlass.Int32, + n_splits: cutlass.Int32, + scale: cutlass.Float32, + blocks_per_page: cutlass.Constexpr[int], + compress_ratio: cutlass.Constexpr[int], +) -> None: + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 4: + prims.prefetch_tensormap(tma_desc_q.get_ptr()) + prims.prefetch_tensormap(tma_desc_k.get_ptr()) + request, first_row, tile_rows, n_begin, n_iters = _tile_coords( + tile_offsets, req_start, positions, num_rows, n_splits, compress_ratio + ) + task_manager, tmem_ptr_i32 = _build_pipeline( + tma_desc_q.get_ptr(), + tma_desc_k.get_ptr(), + block_table, + positions, + scores, + scale, + request, + first_row, + tile_rows, + n_begin, + n_iters, + blocks_per_page, + compress_ratio, + ) + _run(task_manager, warp_idx, tmem_ptr_i32) + + +META_THREADS = 128 + + +@cute.kernel +def qsa_index_tiles_kernel( + request_indices: cute.Tensor, + tile_meta: cute.Tensor, + num_rows: cutlass.Int32, +) -> None: + """Fill ``tile_meta = [req_start[0..R), tile_offsets[0..R))`` for request-sorted rows. + + ``req_start[r]`` is the first row of request ``r`` (``num_rows`` for an + empty request) and ``tile_offsets`` is the inclusive scan of each + request's 32-row tile count. + """ + tx, _, _ = cute.arch.thread_idx() + num_requests = cutlass.Int32(tile_meta.shape[0]) // 2 + for r in cutlass.range(tx, num_requests, META_THREADS): + lo = cutlass.Int32(0) + hi = num_rows + for _ in cutlass.range_constexpr(32): + if lo < hi: + mid = (lo + hi) >> 1 + if cutlass.Int32(request_indices[mid]) < r: + lo = mid + 1 + else: + hi = mid + tile_meta[r] = lo + cute.arch.barrier() + for r in cutlass.range(tx, num_requests, META_THREADS): + row_begin = cutlass.Int32(tile_meta[r]) + row_end = num_rows + if r + 1 < num_requests: + row_end = cutlass.Int32(tile_meta[r + 1]) + tile_meta[num_requests + r] = ( + row_end - row_begin + QUERY_ROWS_PER_TILE - 1 + ) // QUERY_ROWS_PER_TILE + cute.arch.barrier() + if tx < 32: + carry = cutlass.Int32(0) + for base in cutlass.range(0, num_requests, 32): + idx = base + tx + value = cutlass.Int32(0) + if idx < num_requests: + value = cutlass.Int32(tile_meta[num_requests + idx]) + for step in cutlass.range_constexpr(5): + other = cutlass.Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=value, + offset=1 << step, + mask_and_clamp=0, + kind=prims.Shfl.UP, + ) + ) + if tx >= (1 << step): + value = value + other + total = cutlass.Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=value, + offset=31, + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + if idx < num_requests: + tile_meta[num_requests + idx] = value + carry + carry = carry + total + + +@cute.jit +def qsa_index_scores_host( + q_ptr: cute.Pointer, + index_ptr: cute.Pointer, + block_table_ptr: cute.Pointer, + request_indices_ptr: cute.Pointer, + tile_meta_ptr: cute.Pointer, + positions_ptr: cute.Pointer, + scores_ptr: cute.Pointer, + num_requests: cutlass.Int32, + max_pages: cutlass.Int32, + max_blocks: cutlass.Int32, + num_pages: cutlass.Int32, + page_stride: cutlass.Int32, + num_rows: cutlass.Int32, + n_splits: cutlass.Int32, + scale: cutlass.Float32, + grid_size: cutlass.Int32, + blocks_per_page: cutlass.Constexpr[int], + compress_ratio: cutlass.Constexpr[int], + stream, +) -> None: + q_rows = cute.make_tensor( + q_ptr, cute.make_layout((num_rows * NUM_HEADS, HEAD_DIM), stride=(HEAD_DIM, 1)) + ) + block_table = cute.make_tensor( + block_table_ptr, + cute.make_layout((num_requests, max_pages), stride=(max_pages, 1)), + ) + request_indices = cute.make_tensor(request_indices_ptr, cute.make_layout((num_rows,))) + tile_meta = cute.make_tensor(tile_meta_ptr, cute.make_layout((2 * num_requests,))) + req_start = cute.make_tensor(tile_meta_ptr, cute.make_layout((num_requests,))) + tile_offsets = cute.make_tensor(tile_meta_ptr + num_requests, cute.make_layout((num_requests,))) + positions = cute.make_tensor(positions_ptr, cute.make_layout((num_rows,))) + scores = cute.make_tensor( + scores_ptr, cute.make_layout((num_rows, max_blocks), stride=(max_blocks, 1)) + ) + qsa_index_tiles_kernel(request_indices, tile_meta, num_rows).launch( + grid=(1, 1, 1), block=[META_THREADS, 1, 1], stream=stream + ) + tma_desc_q = cuda.create_tensor_map_tiled_from_view( + q_rows, + box_dims=(TILE_M, STAGE_K), + stride_order=(1, 0), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + if cutlass.const_expr(blocks_per_page >= GROUP_ROWS): + # Compressed rows sit at token slot compress_ratio-1 of every block; + # the two 64-dim halves are separate map dims so one box lands them + # adjacently in SMEM. + index_groups = cute.make_tensor( + index_ptr + (compress_ratio - 1) * HEAD_DIM, + cute.make_layout( + (blocks_per_page, num_pages, Q_HALVES, STAGE_K), + stride=(compress_ratio * HEAD_DIM, page_stride, STAGE_K, 1), + ), + ) + tma_desc_k = cuda.create_tensor_map_tiled_from_view( + index_groups, + box_dims=(GROUP_ROWS, 1, Q_HALVES, STAGE_K), + stride_order=(3, 0, 2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + else: + index_groups = cute.make_tensor( + index_ptr, + cute.make_layout( + (blocks_per_page, num_pages, compress_ratio, HEAD_DIM), + stride=(compress_ratio * HEAD_DIM, page_stride, HEAD_DIM, 1), + ), + ) + tma_desc_k = cuda.create_tensor_map_tiled_from_view( + index_groups, + box_dims=(blocks_per_page, 1, 1, STAGE_K), + stride_order=(3, 2, 0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + qsa_index_scores_kernel( + tma_desc_q, + tma_desc_k, + block_table, + tile_offsets, + req_start, + positions, + scores, + num_rows, + n_splits, + scale, + blocks_per_page, + compress_ratio, + ).launch(grid=(grid_size, 1, 1), block=[256, 1, 1], stream=stream) diff --git a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_prims_ts_core.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_prims_ts_core.py new file mode 100644 index 000000000000..5340abc0fc65 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_prims_ts_core.py @@ -0,0 +1,480 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import itertools +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.qsa.module import qsa_sparse_gqa_reference +from tensorrt_llm._torch.attention.backends.sparse.qsa.params import QSASparseParams +from tensorrt_llm._torch.attention.backends.sparse.qsa.prims_ts_core import ( + prims_ts_available, + prims_ts_sparse_gqa, + slice_request_boundaries, +) + +NUM_QO_HEADS = 12 +NUM_KV_HEADS = 1 +HEAD_DIM = 256 +PAGE_SIZE = 32 +COMPRESS_RATIO = 4 +BLOCK_TOPK = 64 +CONTEXT = 2048 +PAGES_PER_REQUEST = CONTEXT // PAGE_SIZE + + +def _params(**overrides) -> QSASparseParams: + kwargs = dict( + index_n_heads=4, + index_kv_heads=1, + index_head_dim=128, + token_topk=BLOCK_TOPK * COMPRESS_RATIO, + compress_ratio=COMPRESS_RATIO, + ) + kwargs.update(overrides) + return QSASparseParams(**kwargs) + + +def test_sparse_core_backend_defaults_to_auto() -> None: + assert _params().sparse_core_backend == "auto" + assert _params().indexer_backend == "follow" + assert _params(sparse_core_backend="prims_ts").sparse_core_backend == "prims_ts" + with pytest.raises(ValueError, match="sparse_core_backend"): + _params(sparse_core_backend="bogus") + with pytest.raises(ValueError, match="indexer_backend"): + _params(indexer_backend="bogus") + + +def _requires_prims_ts() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + if torch.cuda.get_device_capability() not in ((10, 0), (10, 3), (10, 7)): + pytest.skip("PrimTS QToken sparse attention requires SM100/SM103/SM107") + # The vendored package ships with TRT-LLM, so a missing kernel on a + # supported GPU is a failure, not a skip. + assert prims_ts_available() + + +def _make_case(query_lens: list[int], seed: int, kv_dtype: torch.dtype = torch.bfloat16): + torch.manual_seed(seed) + device = torch.device("cuda") + num_requests = len(query_lens) + rows = sum(query_lens) + num_pages = num_requests * PAGES_PER_REQUEST + k_cache = torch.randn( + num_pages, NUM_KV_HEADS, PAGE_SIZE, HEAD_DIM, dtype=torch.bfloat16, device=device + ).to(kv_dtype) + v_cache = torch.randn( + num_pages, NUM_KV_HEADS, PAGE_SIZE, HEAD_DIM, dtype=torch.bfloat16, device=device + ).to(kv_dtype) + block_table = torch.randperm(num_pages, dtype=torch.int32, device=device).view( + num_requests, PAGES_PER_REQUEST + ) + q = torch.randn(rows, NUM_QO_HEADS, HEAD_DIM, dtype=torch.bfloat16, device=device) + request_indices = torch.repeat_interleave( + torch.arange(num_requests, dtype=torch.int32), + torch.tensor(query_lens), + ) + positions = torch.cat( + [torch.arange(CONTEXT - length, CONTEXT, dtype=torch.int32) for length in query_lens] + ) + block_ids = torch.full((rows, BLOCK_TOPK), -1, dtype=torch.int32) + selected_tokens = torch.full( + (rows, BLOCK_TOPK * COMPRESS_RATIO + COMPRESS_RATIO - 1), -1, dtype=torch.int32 + ) + for row in range(rows): + position = int(positions[row]) + visible = (position + 1) // COMPRESS_RATIO + chosen = torch.randperm(visible)[: min(BLOCK_TOPK, visible)] + block_ids[row, : chosen.numel()] = chosen.to(torch.int32) + tokens = (chosen[:, None] * COMPRESS_RATIO + torch.arange(COMPRESS_RATIO)).flatten() + tail = torch.arange(visible * COMPRESS_RATIO, position + 1) + tokens = torch.cat((tokens, tail)) + selected_tokens[row, : tokens.numel()] = tokens.to(torch.int32) + metadata = SimpleNamespace( + qsa_block_table=block_table, + kv_cache_manager=SimpleNamespace(tokens_per_block=PAGE_SIZE), + ) + return ( + q, + k_cache, + v_cache, + block_table, + block_ids.to(device), + request_indices.to(device), + positions.to(device), + selected_tokens.to(device), + metadata, + ) + + +@pytest.mark.parametrize( + ("query_lens", "is_decode"), + [ + ([37, 91], False), + ([1] * 8, True), + ], + ids=["packed_prefill", "fixed_decode"], +) +def test_prims_ts_sparse_gqa_matches_reference(query_lens, is_decode) -> None: + _requires_prims_ts() + ( + q, + k_cache, + v_cache, + block_table, + block_ids, + request_indices, + positions, + selected_tokens, + metadata, + ) = _make_case(query_lens, seed=1234 + len(query_lens)) + softmax_scale = HEAD_DIM**-0.5 + + actual = prims_ts_sparse_gqa( + q, + k_cache, + v_cache, + block_table, + block_ids, + request_indices, + positions, + page_size=PAGE_SIZE, + max_seq_len_kv=CONTEXT, + softmax_scale=softmax_scale, + is_decode=is_decode, + ) + expected = qsa_sparse_gqa_reference( + q=q, + k_cache=k_cache, + v_cache=v_cache, + selected_tokens=selected_tokens, + request_indices=request_indices, + metadata=metadata, + softmax_scale=softmax_scale, + ) + torch.cuda.synchronize() + assert actual.shape == expected.shape + torch.testing.assert_close(actual.float(), expected.float(), atol=2e-2, rtol=2e-2) + + +def test_slice_request_boundaries() -> None: + cu_seq_lens = [37, 128, 131, 200] + assert slice_request_boundaries(cu_seq_lens, 0, 200) == (0, 37, 128, 131, 200) + assert slice_request_boundaries(cu_seq_lens, 0, 100) == (0, 37, 100) + assert slice_request_boundaries(cu_seq_lens, 37, 128) == (0, 91) + assert slice_request_boundaries(cu_seq_lens, 100, 140) == (0, 28, 31, 40) + assert slice_request_boundaries(cu_seq_lens, 131, 200) == (0, 69) + + +def test_prims_ts_sparse_gqa_host_boundaries_match_readback() -> None: + _requires_prims_ts() + query_lens = [37, 91, 5] + q, k_cache, v_cache, block_table, block_ids, request_indices, positions, _, _ = _make_case( + query_lens, seed=99 + ) + kwargs = dict( + page_size=PAGE_SIZE, max_seq_len_kv=CONTEXT, softmax_scale=HEAD_DIM**-0.5, is_decode=False + ) + readback = prims_ts_sparse_gqa( + q, k_cache, v_cache, block_table, block_ids, request_indices, positions, **kwargs + ) + cu_seq_lens = list(itertools.accumulate(query_lens)) + host = prims_ts_sparse_gqa( + q, + k_cache, + v_cache, + block_table, + block_ids, + request_indices, + positions, + request_boundaries=slice_request_boundaries(cu_seq_lens, 0, sum(query_lens)), + **kwargs, + ) + torch.cuda.synchronize() + assert torch.equal(readback, host) + + +def test_prefill_rows_share_plans_within_bucket() -> None: + _requires_prims_ts() + from tensorrt_llm._torch.attention.backends.sparse.qsa.prims_ts_core import _CORE + + kwargs = dict( + page_size=PAGE_SIZE, max_seq_len_kv=CONTEXT, softmax_scale=HEAD_DIM**-0.5, is_decode=False + ) + q, k_cache, v_cache, block_table, block_ids, request_indices, positions, _, _ = _make_case( + [300, 200], seed=5 + ) + prims_ts_sparse_gqa( + q, k_cache, v_cache, block_table, block_ids, request_indices, positions, **kwargs + ) + plans = len(_CORE._plans) + # 500 and 510 rows both pad to 512 rows and 256 routes. + q2, _, _, block_table2, block_ids2, request_indices2, positions2, selected2, metadata = ( + _make_case([260, 250], seed=6) + ) + out = prims_ts_sparse_gqa( + q2, k_cache, v_cache, block_table2, block_ids2, request_indices2, positions2, **kwargs + ) + assert len(_CORE._plans) == plans + assert out.shape[0] == 510 + expected = qsa_sparse_gqa_reference( + q=q2, + k_cache=k_cache, + v_cache=v_cache, + selected_tokens=selected2, + request_indices=request_indices2, + metadata=metadata, + softmax_scale=HEAD_DIM**-0.5, + ) + torch.cuda.synchronize() + torch.testing.assert_close(out, expected, atol=2e-2, rtol=2e-2) + + +def test_route_offsets_follow_layout_changes_without_replanning() -> None: + _requires_prims_ts() + import tensorrt_llm._torch.attention.backends.prims_ts.q_token_kv_block_sparse_metadata as fi_metadata + + kwargs = dict( + page_size=PAGE_SIZE, max_seq_len_kv=CONTEXT, softmax_scale=HEAD_DIM**-0.5, is_decode=False + ) + # Same row count and route count, different request boundaries. + layouts = ([36, 92], [64, 64], [36, 92]) + cases = [_make_case(lens, seed=40 + i) for i, lens in enumerate(layouts)] + k_cache, v_cache, block_table = cases[0][1], cases[0][2], cases[0][3] + + def run(case, lens): + q, _, _, _, block_ids, request_indices, positions, _, _ = case + cu_seq_lens = list(itertools.accumulate(lens)) + return prims_ts_sparse_gqa( + q, + k_cache, + v_cache, + block_table, + block_ids, + request_indices, + positions, + request_boundaries=slice_request_boundaries(cu_seq_lens, 0, sum(lens)), + **kwargs, + ) + + expected = [ + prims_ts_sparse_gqa( + case[0], k_cache, v_cache, block_table, case[4], case[5], case[6], **kwargs + ).clone() + for case in cases + ] + prepares = 0 + original = fi_metadata._prepare_q_token_kv_block_sparse_attention + + def counted(*args, **kw): + nonlocal prepares + prepares += 1 + return original(*args, **kw) + + fi_metadata._prepare_q_token_kv_block_sparse_attention = counted + try: + outputs = [run(case, lens).clone() for case, lens in zip(cases, layouts)] + # Negative control: a different route count is a new plan and must + # re-prepare, which shows the counter is live. + control = _make_case([38, 92], seed=50) + run(control, [38, 92]) + finally: + fi_metadata._prepare_q_token_kv_block_sparse_attention = original + torch.cuda.synchronize() + assert prepares == 1 + for out, ref in zip(outputs, expected): + assert torch.equal(out, ref) + + +def test_auto_gates_small_decode_batches() -> None: + _requires_prims_ts() + from tensorrt_llm._torch.attention.backends.sparse.qsa.prims_ts_core import ( + resolve_sparse_core_backend, + use_prims_ts_indexer, + ) + + device = torch.device("cuda") + k_cache = torch.empty(1, NUM_KV_HEADS, PAGE_SIZE, HEAD_DIM, dtype=torch.bfloat16, device=device) + + def core(backend, rows, is_decode): + q = torch.empty(rows, NUM_QO_HEADS, HEAD_DIM, dtype=torch.bfloat16, device=device) + return resolve_sparse_core_backend( + _params(sparse_core_backend=backend), + q=q, + k_cache=k_cache, + page_size=PAGE_SIZE, + is_decode=is_decode, + num_rows=rows, + ) + + def indexer(backend, rows, is_decode): + q = torch.empty(rows, 4, 128, dtype=torch.bfloat16, device=device) + return use_prims_ts_indexer(_params(sparse_core_backend=backend), q, is_decode=is_decode) + + assert core("auto", 1, True) == "triton" + assert core("auto", 2, True) == "prims_ts" + assert core("auto", 1, False) == "prims_ts" + assert core("prims_ts", 1, True) == "prims_ts" + assert not indexer("auto", 31, True) + assert indexer("auto", 32, True) + assert indexer("auto", 1, False) + assert indexer("prims_ts", 1, True) + assert not indexer("triton", 4096, False) + q64 = torch.empty(64, 4, 128, dtype=torch.bfloat16, device=device) + q31 = torch.empty(31, 4, 128, dtype=torch.bfloat16, device=device) + assert not use_prims_ts_indexer( + _params(sparse_core_backend="prims_ts", indexer_backend="triton"), q64 + ) + assert use_prims_ts_indexer( + _params(sparse_core_backend="triton", indexer_backend="prims_ts"), q31, is_decode=True + ) + assert not use_prims_ts_indexer( + _params(sparse_core_backend="prims_ts", indexer_backend="auto"), q31, is_decode=True + ) + assert not use_prims_ts_indexer(_params(), q64, tokens_per_block=48) + assert use_prims_ts_indexer(_params(), q64, tokens_per_block=64) + + +def test_prefill_writes_into_caller_output() -> None: + _requires_prims_ts() + kwargs = dict( + page_size=PAGE_SIZE, max_seq_len_kv=CONTEXT, softmax_scale=HEAD_DIM**-0.5, is_decode=False + ) + for query_lens in ([128, 128], [37, 91]): + q, k_cache, v_cache, block_table, block_ids, request_indices, positions, _, _ = _make_case( + query_lens, seed=77 + ) + expected = prims_ts_sparse_gqa( + q, k_cache, v_cache, block_table, block_ids, request_indices, positions, **kwargs + ).clone() + out = torch.full_like(q, float("nan")) + result = prims_ts_sparse_gqa( + q, + k_cache, + v_cache, + block_table, + block_ids, + request_indices, + positions, + out=out, + **kwargs, + ) + torch.cuda.synchronize() + assert result.data_ptr() == out.data_ptr() + assert torch.equal(out, expected) + + +@pytest.mark.parametrize("is_decode", [False, True], ids=["packed_prefill", "fixed_decode"]) +def test_prims_ts_sparse_gqa_fp8_kv_matches_reference(is_decode) -> None: + _requires_prims_ts() + query_lens = [1] * 8 if is_decode else [37, 91] + q, k_cache, v_cache, block_table, block_ids, request_indices, positions, selected, metadata = ( + _make_case(query_lens, seed=4321, kv_dtype=torch.float8_e4m3fn) + ) + softmax_scale = HEAD_DIM**-0.5 + actual = prims_ts_sparse_gqa( + q, + k_cache, + v_cache, + block_table, + block_ids, + request_indices, + positions, + page_size=PAGE_SIZE, + max_seq_len_kv=CONTEXT, + softmax_scale=softmax_scale, + is_decode=is_decode, + ) + expected = qsa_sparse_gqa_reference( + q=q, + k_cache=k_cache.to(torch.bfloat16), + v_cache=v_cache.to(torch.bfloat16), + selected_tokens=selected, + request_indices=request_indices, + metadata=metadata, + softmax_scale=softmax_scale, + ) + torch.cuda.synchronize() + # FP8 Q on the PrimTS side vs BF16 Q upcast on the reference side. + torch.testing.assert_close(actual.float(), expected.float(), atol=5e-2, rtol=5e-2) + + +@pytest.mark.parametrize("kv_dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"]) +def test_decode_graph_capture_replay(kv_dtype) -> None: + _requires_prims_ts() + q, k_cache, v_cache, block_table, block_ids, request_indices, positions, _, _ = _make_case( + [1] * 4, seed=99, kv_dtype=kv_dtype + ) + kwargs = dict( + page_size=PAGE_SIZE, max_seq_len_kv=CONTEXT, softmax_scale=HEAD_DIM**-0.5, is_decode=True + ) + out = torch.empty_like(q) + # Eager run prepares the plan; capture must not need the host afterwards. + prims_ts_sparse_gqa( + q, k_cache, v_cache, block_table, block_ids, request_indices, positions, out=out, **kwargs + ) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + prims_ts_sparse_gqa( + q, + k_cache, + v_cache, + block_table, + block_ids, + request_indices, + positions, + out=out, + **kwargs, + ) + q.copy_(torch.randn_like(q, dtype=torch.float32).to(q.dtype)) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + eager = prims_ts_sparse_gqa( + q, k_cache, v_cache, block_table, block_ids, request_indices, positions, **kwargs + ) + torch.cuda.synchronize() + assert torch.equal(out, eager) + + +def test_decode_rows_get_exact_plans() -> None: + _requires_prims_ts() + from tensorrt_llm._torch.attention.backends.sparse.qsa.prims_ts_core import _CORE + + kwargs = dict( + page_size=PAGE_SIZE, max_seq_len_kv=CONTEXT, softmax_scale=HEAD_DIM**-0.5, is_decode=True + ) + q, k_cache, v_cache, block_table, block_ids, request_indices, positions, selected, metadata = ( + _make_case([1] * 5, seed=7) + ) + prims_ts_sparse_gqa( + q[:3], + k_cache, + v_cache, + block_table, + block_ids[:3], + request_indices[:3], + positions[:3], + **kwargs, + ) + plans = len(_CORE._plans) + out = prims_ts_sparse_gqa( + q, k_cache, v_cache, block_table, block_ids, request_indices, positions, **kwargs + ) + assert len(_CORE._plans) == plans + 1 + assert out.shape[0] == 5 + expected = qsa_sparse_gqa_reference( + q=q, + k_cache=k_cache, + v_cache=v_cache, + selected_tokens=selected, + request_indices=request_indices, + metadata=metadata, + softmax_scale=HEAD_DIM**-0.5, + ) + torch.cuda.synchronize() + torch.testing.assert_close(out, expected, atol=2e-2, rtol=2e-2) diff --git a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_prims_ts_indexer.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_prims_ts_indexer.py new file mode 100644 index 000000000000..433c8a851eac --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_prims_ts_indexer.py @@ -0,0 +1,264 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import pytest +import torch + +HEAD_DIM = 128 +NUM_HEADS = 4 +RATIO = 4 +# BF16 operands are exact; the fp32 head-sum order differs between the UMMA +# accumulator and the einsum reference, so the tolerance scales with |score|. +REL_TOL = 2e-5 +ABS_TOL = 1e-4 + + +def _reference(q, index_cache, block_table, positions, request_indices, tokens_per_block): + max_blocks = block_table.shape[1] * tokens_per_block // RATIO + blocks = torch.arange(max_blocks, device=q.device) + anchors = blocks * RATIO + RATIO - 1 + pages = block_table[request_indices.long()][:, anchors // tokens_per_block] + keys = index_cache[pages, anchors % tokens_per_block, 0] + scores = torch.einsum("rhd,rbd->rhb", q.float(), keys.float()) + scores = scores.clamp_min(0).sum(dim=1) * HEAD_DIM**-0.5 + visible = blocks[None, :] < ((positions + 1) // RATIO)[:, None] + return scores.masked_fill(~visible, float("-inf")), visible + + +def _case(request_lengths, ctx, tokens_per_block, seed, *, q_scale=1.0, positions=None): + torch.manual_seed(seed) + device = torch.device("cuda") + pages_per_request = (ctx + tokens_per_block - 1) // tokens_per_block + num_requests = len(request_lengths) + num_pages = num_requests * pages_per_request + block_table = torch.randperm(num_pages, device=device).to(torch.int32).view(num_requests, -1) + index_cache = torch.randn( + num_pages, tokens_per_block, 1, HEAD_DIM, dtype=torch.bfloat16, device=device + ) + rows = sum(request_lengths) + q = (torch.randn(rows, NUM_HEADS, HEAD_DIM, device=device) * q_scale).to(torch.bfloat16) + request_indices = torch.repeat_interleave( + torch.arange(num_requests, device=device), + torch.tensor(request_lengths, device=device), + ).to(torch.int32) + if positions is None: + positions = torch.cat([torch.arange(ctx - n, ctx, device=device) for n in request_lengths]) + positions = positions.to(torch.int64) + return q, index_cache, block_table, positions, request_indices + + +def _requires_prims_ts(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + if torch.cuda.get_device_capability() not in ((10, 0), (10, 3), (10, 7)): + pytest.skip("PrimTS QSA scorer requires SM100/SM103/SM107") + # The vendored PrimTS kernels need this DSL in the CI image, so a missing + # module on a supported GPU is a failure, not a skip. + import cutlass.experimental.task_scheduling # noqa: F401 + + +def _scorer(): + from tensorrt_llm._torch.attention.backends.sparse.qsa.prims_ts_indexer import ( + prims_ts_qsa_index_scores, + ) + + return prims_ts_qsa_index_scores + + +def _check(out, ref, visible, only_visible): + assert out.shape == ref.shape + assert out.dtype == torch.float32 + if not only_visible: + assert torch.equal(torch.isfinite(out), visible) + assert torch.all(out[~visible] == float("-inf")) + got, want = out[visible], ref[visible] + assert torch.isfinite(got).all() + err = (got - want).abs().max().item() if got.numel() else 0.0 + bound = ABS_TOL + REL_TOL * (want.abs().max().item() if want.numel() else 0.0) + assert err <= bound, f"max abs err {err:.3e} > {bound:.3e}" + return err + + +def _run(q, cache, table, pos, req, tokens_per_block, only_visible): + out = _scorer()( + q, + cache, + table, + pos, + req, + tokens_per_block=tokens_per_block, + compress_ratio=RATIO, + only_visible_blocks=only_visible, + ) + ref, visible = _reference(q, cache, table, pos, req, tokens_per_block) + return out, ref, visible, _check(out, ref, visible, only_visible) + + +@pytest.mark.parametrize("tokens_per_block", (16, 32, 64)) +@pytest.mark.parametrize("only_visible", (False, True)) +@pytest.mark.parametrize( + "request_lengths, ctx", + [ + ((1,), 2048), + ((5, 37), 2048), + ((40, 24, 3, 61), 8192), + ((1024,), 32768 + 36), + ((4096,), 32768), + ], +) +def test_prefill_rows(request_lengths, ctx, tokens_per_block, only_visible): + _requires_prims_ts() + q, cache, table, pos, req = _case(request_lengths, ctx, tokens_per_block, 7, q_scale=8.0) + _run(q, cache, table, pos, req, tokens_per_block, only_visible) + + +@pytest.mark.parametrize("tokens_per_block", (16, 32, 64)) +@pytest.mark.parametrize("num_requests", (1, 8, 64)) +def test_decode_rows(num_requests, tokens_per_block): + _requires_prims_ts() + ctx = 8192 + q, cache, table, _, req = _case([1] * num_requests, ctx, tokens_per_block, 11, q_scale=4.0) + device = q.device + # Rows at the first positions have no complete block, rows inside a partial + # block or at a page boundary exercise the tail of the visible range. + specials = torch.tensor( + [ + 0, + 1, + 2, + 3, + 4, + 7, + tokens_per_block - 1, + tokens_per_block, + tokens_per_block * 3 - 1, + ctx - 1, + ], + device=device, + ) + pos = torch.randint(0, ctx, (num_requests,), device=device) + pos[: specials.numel()] = specials[:num_requests] + for only_visible in (False, True): + _run(q, cache, table, pos.to(torch.int64), req, tokens_per_block, only_visible) + + +def test_prefill_mixed_positions_within_request(): + _requires_prims_ts() + q, cache, table, pos, req = _case((100, 60), 8192, 32, 5) + pos = pos.clone() + pos[10:20] = 3 + pos[130] = 0 + _run(q, cache, table, pos, req, 32, False) + + +@pytest.mark.parametrize("seed", (1, 2, 3)) +def test_prefill_seeds_and_large_scores(seed): + _requires_prims_ts() + q, cache, table, pos, req = _case((513, 7), 16384, 32, seed, q_scale=32.0) + _run(q, cache, table, pos, req, 32, False) + + +def test_prefill_128k(): + _requires_prims_ts() + q, cache, table, pos, req = _case((2048,), 131072, 32, 13) + _run(q, cache, table, pos, req, 32, True) + + +def test_matches_triton_scorer(): + _requires_prims_ts() + pytest.importorskip("tensorrt_llm") + from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( + triton_qsa_paged_index_scores, + ) + + for lengths, ctx, tpb in ( + ((300, 4), 8192, 32), + ((1,) * 16, 32768, 64), + ((2000,), 16384, 16), + ): + q, cache, table, pos, req = _case(lengths, ctx, tpb, 21, q_scale=8.0) + ours = _scorer()( + q, + cache, + table, + pos, + req, + tokens_per_block=tpb, + compress_ratio=RATIO, + only_visible_blocks=True, + ) + theirs = triton_qsa_paged_index_scores( + q=q, + index_cache=cache, + block_table=table, + query_positions=pos, + request_indices=req, + tokens_per_block=tpb, + compress_ratio=RATIO, + only_visible_blocks=True, + ) + visible = ( + torch.arange(ours.shape[1], device=q.device)[None, :] < ((pos + 1) // RATIO)[:, None] + ) + err = (ours[visible] - theirs[visible]).abs().max().item() + assert err <= ABS_TOL + REL_TOL * theirs[visible].abs().max().item(), err + + +def test_deterministic(): + _requires_prims_ts() + q, cache, table, pos, req = _case((257, 31), 8192, 32, 17) + a = _scorer()(q, cache, table, pos, req, tokens_per_block=32, compress_ratio=RATIO) + b = _scorer()(q, cache, table, pos, req, tokens_per_block=32, compress_ratio=RATIO) + assert torch.equal(a, b) + + +def test_prepared_launch_graph_replay(): + _requires_prims_ts() + from tensorrt_llm._torch.attention.backends.sparse.qsa.prims_ts_indexer import ( + prepare_qsa_index_scores, + ) + + q, cache, table, pos, req = _case((64, 64), 8192, 32, 19) + launch, scores = prepare_qsa_index_scores( + q, cache, table, pos, req, tokens_per_block=32, compress_ratio=RATIO + ) + launch() + torch.cuda.synchronize() + ref, visible = _reference(q, cache, table, pos, req, 32) + _check(scores, ref, visible, False) + q.copy_(torch.randn_like(q, dtype=torch.float32).to(torch.bfloat16)) + scores.fill_(float("nan")) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch() + graph.replay() + torch.cuda.synchronize() + ref, visible = _reference(q, cache, table, pos, req, 32) + _check(scores, ref, visible, False) + + +@pytest.mark.parametrize("tokens_per_block", (16, 32)) +def test_strided_index_cache_view_is_not_copied(tokens_per_block): + _requires_prims_ts() + from tensorrt_llm._torch.attention.backends.sparse.qsa.prims_ts_indexer import ( + prepare_qsa_index_scores, + ) + + q, cache, table, pos, req = _case((200, 56), 4096, tokens_per_block, 23) + # The KV pool coalesces the index cache with K/V, so production views have + # a page stride larger than one page. + pool = torch.zeros( + (cache.shape[0], 5, *cache.shape[1:]), dtype=cache.dtype, device=cache.device + ) + view = pool[:, 2] + view.copy_(cache) + assert not view.is_contiguous() + launch, scores = prepare_qsa_index_scores( + q, view, table, pos, req, tokens_per_block=tokens_per_block, compress_ratio=RATIO + ) + assert any(t.data_ptr() == view.data_ptr() for t in launch.keepalive) + launch() + dense = _scorer()( + q, cache, table, pos, req, tokens_per_block=tokens_per_block, compress_ratio=RATIO + ) + torch.cuda.synchronize() + assert torch.equal(scores, dense)