From b7c5e7085cea4da359aebc4f0876df225dc86ebf Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Tue, 20 Oct 2026 23:11:50 +0000 Subject: [PATCH 01/18] [None][feat] Add PrimTS sparse core backend for Qwen3.8 QSA attention Route the QSA sparse core attention to FlashInfer's PrimTS QToken-KvBlock-Sparse kernel behind QSASparseParams.sparse_core_backend (triton | prims_ts | auto, env TLLM_QSA_SPARSE_CORE). PrimTS consumes the indexer's compact block IDs directly, so the Triton token expansion is skipped on that path; the output gate stays on apply_output_gate. Plans are cached per geometry and per layer K/V buffer over one shared workspace, and FP8 caches quantize Q per call with the scale folded into sm_scale. Adds a unit test against qsa_sparse_gqa_reference and a standalone A/B microbenchmark of the two sparse cores. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../attention/backends/sparse/qsa/module.py | 107 ++++- .../attention/backends/sparse/qsa/params.py | 16 + .../backends/sparse/qsa/prims_ts_core.py | 312 +++++++++++++ .../qsa/bench_qsa_sparse_core.py | 442 ++++++++++++++++++ .../sparse/qsa/test_qsa_prims_ts_core.py | 162 +++++++ 5 files changed, 1018 insertions(+), 21 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py create mode 100644 tests/microbenchmarks/qsa/bench_qsa_sparse_core.py create mode 100644 tests/unittest/_torch/attention/sparse/qsa/test_qsa_prims_ts_core.py diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py index c1c127aa35df..e2c2c398a930 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py @@ -374,6 +374,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 +387,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 +410,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, @@ -464,6 +470,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 +637,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 +694,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) @@ -717,7 +766,23 @@ 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, ) + if use_prims_ts: + output[packed_slice] = prims_ts_sparse_gqa( + q[packed_slice], + k_cache, + v_cache, + attn_metadata.qsa_block_table, + selected, + req_idx[packed_slice], + logical[packed_slice], + page_size=tokens_per_block, + max_seq_len_kv=max_seq_len_kv, + softmax_scale=softmax_scale, + is_decode=False, + ) + continue output[packed_slice] = qsa_sparse_gqa( q=q[packed_slice], k_cache=k_cache, @@ -725,7 +790,7 @@ def forward( selected_tokens=selected, request_indices=req_idx[packed_slice], metadata=attn_metadata, - softmax_scale=1.0 / (attention.q_scaling * attention.head_dim**0.5), + softmax_scale=softmax_scale, query_positions=logical[packed_slice], compress_ratio=params.compress_ratio, ) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py index cc1e4c7ce596..3bd6dd932843 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py @@ -2,11 +2,15 @@ # SPDX-License-Identifier: Apache-2.0 """Internal sparse-attention parameters for QSA.""" +import os from dataclasses import dataclass from ..params import SparseMetadataParams, SparseParams from .constants import QSA_INDEX_KV_HEADS +QSA_SPARSE_CORE_BACKENDS = ("auto", "triton", "prims_ts") +QSA_SPARSE_CORE_ENV = "TLLM_QSA_SPARSE_CORE" + def _validate_selection_geometry(token_topk: int, compress_ratio: int) -> None: if token_topk <= 0: @@ -41,8 +45,20 @@ class QSASparseParams(SparseParams): enable_heuristic_topk: bool = False # Reuse one indexer selection across the MTP draft loop. mtp_index_share: bool = False + # Sparse core attention kernel: "triton", "prims_ts", or "auto" (PrimTS + # for prefill and large decode batches on SM100/SM103/SM107, else Triton). + # TLLM_QSA_SPARSE_CORE overrides the configured value. + sparse_core_backend: str = "triton" def __post_init__(self) -> None: + env_backend = os.environ.get(QSA_SPARSE_CORE_ENV) + if env_backend: + self.sparse_core_backend = env_backend.strip().lower() + 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}" + ) 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..ed8563786c26 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py @@ -0,0 +1,312 @@ +# 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 functools +import threading + +import torch + +from tensorrt_llm.logger import logger + +from .params import QSASparseParams + +_PREFILL_GROUP_SIZE = 4 +_MIN_DECODE_ROWS = 64 +_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 flashinfer.attention.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 "flashinfer.attention.prims_ts 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 + + +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) -> torch.Tensor: + req = request_indices.to(torch.int64).cpu() + change = torch.nonzero(req[1:] != req[:-1]).flatten() + 1 + return torch.cat( + ( + torch.zeros(1, dtype=torch.int32), + change.to(torch.int32), + torch.tensor([req.numel()], dtype=torch.int32), + ) + ) + + +class _PrimsTSSparseCore: + def __init__(self) -> None: + self._plans: dict[tuple, object] = {} + self._last_abi: dict[tuple, tuple] = {} + self._workspace: torch.Tensor | 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() + ws = torch.zeros(nbytes, dtype=torch.uint8, device=device) + self._workspace = ws + return ws + + 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, + ) -> 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 + + sm_scale = softmax_scale + 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 quantized per + # call with a dynamic per-tensor scale folded into sm_scale. This + # differs numerically from the Triton path, which upcasts K/V, and + # the scale readback synchronizes the host. + q_scale = (q.float().abs().amax() / _FP8_MAX).clamp_min(1e-12) + q = (q.float() / q_scale).to(torch.float8_e4m3fn) + sm_scale = softmax_scale * float(q_scale) + + request_indices = request_indices.to(torch.int32).contiguous() + q = q.contiguous() + 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 + qo_indptr = make_qo_indptr( + _request_boundaries(request_indices), + rows, + group_size=_PREFILL_GROUP_SIZE, + device=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(), + 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, + ) + out = torch.empty((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=4, + 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, + 4, + 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 + abi = tuple( + (name, tuple(t.shape), tuple(t.stride()), t.dtype) + for name, t in ( + ("q", route_q), + ("block_ids", block_ids), + ("block_table", block_table), + ("request_indices", request_indices), + ("query_positions", query_positions), + ("out", run_out), + ) + ) + (("k_ptr", k_cache.data_ptr()), ("v_ptr", v_cache.data_ptr())) + try: + wrapper.run( + route_q, + (k_cache, v_cache), + block_table, + block_ids, + request_indices, + query_positions, + qo_indptr=qo_indptr, + sm_scale=sm_scale, + out=run_out, + ) + except RuntimeError as error: + previous = self._last_abi.get(key) + changed = ( + [(a, b) for a, b in zip(previous or (), abi) if a != b] + if previous + else [("no eager run for this plan key", None)] + ) + raise RuntimeError(f"{error}; ABI changes vs last eager run: {changed}") from error + self._last_abi[key] = abi + return out + + +_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, +) -> 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. + """ + 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, + ) + + +__all__ = [ + "prims_ts_available", + "prims_ts_sparse_gqa", + "resolve_sparse_core_backend", +] diff --git a/tests/microbenchmarks/qsa/bench_qsa_sparse_core.py b/tests/microbenchmarks/qsa/bench_qsa_sparse_core.py new file mode 100644 index 000000000000..674d2db4fe6f --- /dev/null +++ b/tests/microbenchmarks/qsa/bench_qsa_sparse_core.py @@ -0,0 +1,442 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A/B latency of the QSA sparse core: TRT-LLM Triton vs FlashInfer PrimTS. + +Qwen3.8-Flash-Next geometry (D256, 4-token KV blocks, top-k 512 blocks) over an +HND paged cache. Triton time includes the block-to-token expansion the Triton +kernel needs; PrimTS consumes block IDs directly. Both outputs are checked +against the Torch reference on sampled rows before timing. + + python bench_qsa_sparse_core.py --mode prefill --ctx 32768 --chunk 4096 + python bench_qsa_sparse_core.py --mode decode --ctx 32768 --batch 64 --sq 4 + python bench_qsa_sparse_core.py --preset quick --out results.md +""" + +import argparse +import json +import math +import statistics +import time +from types import SimpleNamespace + +import torch + +HEAD_DIM = 256 +KV_BLOCK = 4 +TOKEN_TOPK = 2048 +BLOCK_TOPK = TOKEN_TOPK // KV_BLOCK +PREFILL_GROUP = 4 +FP8_MAX = 448.0 + +QUICK_PRESET = [ + dict(mode="prefill", ctx=32768, chunk=4096, batch=1, sq=1), + dict(mode="decode", ctx=32768, chunk=0, batch=8, sq=1), + dict(mode="decode", ctx=32768, chunk=0, batch=64, sq=1), + dict(mode="decode", ctx=32768, chunk=0, batch=64, sq=4), +] + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--mode", choices=("prefill", "decode"), default="prefill") + p.add_argument("--tp", type=int, choices=(1, 2), default=2, help="1: Hq/Hkv 24/2, 2: 12/1") + p.add_argument("--ctx", type=int, default=32768) + p.add_argument("--chunk", type=int, default=4096, help="prefill query rows per request") + p.add_argument("--batch", type=int, default=1) + p.add_argument("--sq", type=int, default=1, help="decode query tokens per request (1 or MTP+1)") + p.add_argument("--kv-dtype", choices=("bf16", "fp8"), default="bf16") + p.add_argument("--page", type=int, default=32) + p.add_argument("--pattern", choices=("random", "sorted"), default="random") + p.add_argument("--warmup", type=int, default=20) + p.add_argument("--iters", type=int, default=100) + p.add_argument("--l2-mb", type=int, default=256) + p.add_argument("--check-rows", type=int, default=8) + p.add_argument("--no-graph", action="store_true") + p.add_argument("--skip-triton", action="store_true") + p.add_argument("--skip-prims", action="store_true") + p.add_argument("--preset", choices=("quick",), default=None) + p.add_argument("--seed", type=int, default=1729) + p.add_argument("--out", default=None, help="markdown file to append rows to") + p.add_argument("--json", default=None, help="json file to append rows to") + return p.parse_args() + + +def heads_for_tp(tp): + return (24, 2) if tp == 1 else (12, 1) + + +def make_cache(num_requests, ctx, hkv, page, kv_dtype, device, gen): + pages_per_req = math.ceil(ctx / page) + num_pages = num_requests * pages_per_req + k = ( + torch.randn( + num_pages, hkv, page, HEAD_DIM, dtype=torch.bfloat16, device=device, generator=gen + ) + * 0.25 + ) + v = ( + torch.randn( + num_pages, hkv, page, HEAD_DIM, dtype=torch.bfloat16, device=device, generator=gen + ) + * 0.25 + ) + perm = torch.randperm(num_pages, device=device, generator=gen).to(torch.int32) + block_table = perm.view(num_requests, pages_per_req).contiguous() + return k.to(kv_dtype), v.to(kv_dtype), block_table + + +def make_block_ids(positions, pattern, device, gen): + rows = positions.numel() + visible = (positions + 1) // KV_BLOCK + max_vis = int(visible.max()) + keys = torch.rand(rows, max_vis, device=device, generator=gen) + col = torch.arange(max_vis, device=device) + keys.masked_fill_(col[None, :] >= visible[:, None], 2.0) + order = keys.argsort(dim=1)[:, :BLOCK_TOPK] + count = visible.clamp(max=BLOCK_TOPK) + valid = torch.arange(BLOCK_TOPK, device=device)[None, :] < count[:, None] + ids = order.to(torch.int32) + if pattern == "sorted": + ids = ( + torch.where(valid, ids, torch.full_like(ids, torch.iinfo(torch.int32).max)) + .sort(dim=1) + .values + ) + return torch.where(valid, ids, torch.full_like(ids, -1)).contiguous() + + +def build_case(args, device, gen): + hq, hkv = heads_for_tp(args.tp) + kv_dtype = torch.bfloat16 if args.kv_dtype == "bf16" else torch.float8_e4m3fn + if args.mode == "prefill": + rows_per_req = args.chunk + else: + rows_per_req = args.sq + B = args.batch + rows = B * rows_per_req + q = torch.randn(rows, hq, HEAD_DIM, dtype=torch.bfloat16, device=device, generator=gen) * 0.25 + k, v, block_table = make_cache(B, args.ctx, hkv, args.page, kv_dtype, device, gen) + req = torch.arange(B, dtype=torch.int32, device=device).repeat_interleave(rows_per_req) + pos = torch.arange(args.ctx - rows_per_req, args.ctx, dtype=torch.int64, device=device).repeat( + B + ) + seq_lens = torch.full((rows,), args.ctx, dtype=torch.int32, device=device) + ids = make_block_ids(pos, args.pattern, device, gen) + scale = HEAD_DIM**-0.5 + return SimpleNamespace( + q=q, + k=k, + v=v, + block_table=block_table, + req=req, + pos=pos, + seq_lens=seq_lens, + ids=ids, + scale=scale, + hq=hq, + hkv=hkv, + rows=rows, + rows_per_req=rows_per_req, + B=B, + kv_dtype=kv_dtype, + ) + + +def triton_runner(case, args): + from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( + triton_expand_qsa_block_indices, + triton_qsa_paged_sparse_gqa, + ) + + def run(): + tokens = triton_expand_qsa_block_indices( + case.ids, case.pos, case.seq_lens, compress_ratio=KV_BLOCK, token_topk=TOKEN_TOPK + ) + return triton_qsa_paged_sparse_gqa( + q=case.q, + k_cache=case.k, + v_cache=case.v, + block_table=case.block_table, + selected_tokens=tokens, + request_indices=case.req, + tokens_per_block=args.page, + softmax_scale=case.scale, + query_positions=case.pos, + compress_ratio=KV_BLOCK, + ) + + return run + + +def prims_runner(case, args, device): + from flashinfer.attention.prims_ts import ( + QTokenKvBlockSparsePagedTSWrapper, + get_q_token_kv_block_sparse_workspace_size, + make_q_token_kv_block_sparse_qo_indptr, + ) + + q = case.q + sm_scale = case.scale + if case.kv_dtype == torch.float8_e4m3fn: + q_scale = (q.float().abs().amax() / FP8_MAX).clamp_min(1e-8).item() + q = (q.float() / q_scale).to(torch.float8_e4m3fn) + sm_scale = case.scale * q_scale + out = torch.empty(case.rows, case.hq, HEAD_DIM, dtype=torch.bfloat16, device=device) + packed = args.mode == "prefill" + if packed: + group = PREFILL_GROUP + starts = torch.arange(0, case.rows + 1, case.rows_per_req, dtype=torch.int32) + qo_indptr = make_q_token_kv_block_sparse_qo_indptr( + starts, case.rows, group_size=group, device=device + ) + routes = qo_indptr.numel() - 1 + q_view, out_view, split_kv = q, out, False + else: + group = case.rows_per_req + qo_indptr = None + routes = case.B + q_view = q.view(case.B, 1, group, case.hq, HEAD_DIM) + out_view = out.view_as(q_view) + split_kv = True + ws_bytes = get_q_token_kv_block_sparse_workspace_size( + q_view, + case.k, + case.block_table, + block_topk=BLOCK_TOPK, + max_seq_len_kv=args.ctx, + o_data_type=out.dtype, + qo_indptr=qo_indptr, + group_size=group, + kv_block_size=KV_BLOCK, + split_kv=split_kv, + ) + workspace = torch.zeros(ws_bytes, dtype=torch.uint8, device=device) + wrapper = QTokenKvBlockSparsePagedTSWrapper() + wrapper.plan( + routes, + case.hq, + case.hkv, + HEAD_DIM, + KV_BLOCK, + args.page, + BLOCK_TOPK, + args.ctx, + device=device, + workspace_buffer=workspace, + group_size=group, + use_packed_q=packed, + split_kv=split_kv, + q_data_type=q.dtype, + kv_data_type=case.k.dtype, + o_data_type=out.dtype, + ) + + def run(): + wrapper.run( + q_view, + (case.k, case.v), + case.block_table, + case.ids, + case.req, + case.pos, + qo_indptr=qo_indptr, + sm_scale=sm_scale, + out=out_view, + ) + return out + + return run, dict(group_size=group, routes=routes, workspace_mb=ws_bytes / 2**20) + + +def reference(case, args, rows_idx): + from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( + triton_expand_qsa_block_indices, + ) + from tensorrt_llm._torch.attention.backends.sparse.qsa.module import qsa_sparse_gqa_reference + + tokens = triton_expand_qsa_block_indices( + case.ids[rows_idx], + case.pos[rows_idx], + case.seq_lens[rows_idx], + compress_ratio=KV_BLOCK, + token_topk=TOKEN_TOPK, + ) + meta = SimpleNamespace( + qsa_block_table=case.block_table, + kv_cache_manager=SimpleNamespace(tokens_per_block=args.page), + ) + return qsa_sparse_gqa_reference( + q=case.q[rows_idx], + k_cache=case.k.to(torch.bfloat16), + v_cache=case.v.to(torch.bfloat16), + selected_tokens=tokens, + request_indices=case.req[rows_idx], + metadata=meta, + softmax_scale=case.scale, + ) + + +def check(name, out, ref, rows_idx, tol): + got = out[rows_idx].float() + err = (got - ref.float()).abs().max().item() + ok = torch.allclose(got, ref.float(), atol=tol, rtol=tol) + print(f" [{name}] max_abs_err={err:.4e} {'OK' if ok else 'MISMATCH'} (tol {tol})") + return ok, err + + +def time_fn(fn, args, device): + for _ in range(3): + fn() + torch.cuda.synchronize() + if args.no_graph: + launch = fn + else: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fn() + launch = graph.replay + scrub = torch.empty(args.l2_mb << 20, dtype=torch.uint8, device=device) + samples = [] + for i in range(args.warmup + args.iters): + scrub.random_() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + launch() + end.record() + end.synchronize() + if i >= args.warmup: + samples.append(start.elapsed_time(end) * 1000.0) + return samples + + +def summarize(samples): + s = sorted(samples) + return dict(median_us=statistics.median(s), p95_us=s[int(0.95 * (len(s) - 1))], n=len(s)) + + +def run_case(args, device): + gen = torch.Generator(device=device).manual_seed(args.seed) + case = build_case(args, device, gen) + hq, hkv = case.hq, case.hkv + label = ( + f"{args.mode} tp{args.tp} Hq/Hkv={hq}/{hkv} ctx={args.ctx} " + + (f"chunk={args.chunk} " if args.mode == "prefill" else f"bs={args.batch} sq={args.sq} ") + + f"kv={args.kv_dtype} page={args.page} pattern={args.pattern}" + ) + print(f"== {label}") + rows_idx = torch.linspace( + 0, case.rows - 1, steps=min(args.check_rows, case.rows), device=device + ).long() + ref = reference(case, args, rows_idx) + tol = 2e-2 if args.kv_dtype == "bf16" else 5e-2 + row = dict( + label=label, + mode=args.mode, + tp=args.tp, + ctx=args.ctx, + chunk=args.chunk, + batch=args.batch, + sq=args.sq, + kv_dtype=args.kv_dtype, + page=args.page, + pattern=args.pattern, + rows=case.rows, + ) + runners = [] + if not args.skip_triton: + t0 = time.time() + fn = triton_runner(case, args) + ok, err = check("triton", fn(), ref, rows_idx, tol) + row.update(triton_ok=ok, triton_err=err) + runners.append(("triton", fn)) + print(f" triton ready in {time.time() - t0:.1f}s") + if not args.skip_prims: + t0 = time.time() + fn, info = prims_runner(case, args, device) + ok, err = check("prims_ts", fn(), ref, rows_idx, tol) + row.update(prims_ok=ok, prims_err=err, **info) + runners.append(("prims_ts", fn)) + print( + f" prims_ts ready in {time.time() - t0:.1f}s (G={info['group_size']}, routes={info['routes']})" + ) + samples = {name: [] for name, _ in runners} + for order in (runners, list(reversed(runners))): + for name, fn in order: + samples[name] += time_fn(fn, args, device) + for name in samples: + row[name] = summarize(samples[name]) + print( + f" {name:8s} median {row[name]['median_us']:9.1f} us p95 {row[name]['p95_us']:9.1f} us" + ) + if "triton" in row and "prims_ts" in row: + row["speedup"] = row["triton"]["median_us"] / row["prims_ts"]["median_us"] + print(f" speedup (triton/prims_ts) = {row['speedup']:.3f}x") + return row + + +def md_row(r): + t = f"{r['triton']['median_us']:.1f}" if "triton" in r else "-" + p = f"{r['prims_ts']['median_us']:.1f}" if "prims_ts" in r else "-" + s = f"{r['speedup']:.2f}x" if "speedup" in r else "-" + okc = "".join( + f"{n}:{'ok' if r.get(k) else 'BAD'} " + for n, k in (("T", "triton_ok"), ("P", "prims_ok")) + if k in r + ) + return f"| {r['label']} | {r['rows']} | {t} | {p} | {s} | {okc.strip()} |" + + +def main(): + args = parse_args() + device = torch.device("cuda", torch.cuda.current_device()) + print("device:", torch.cuda.get_device_name(device), torch.cuda.get_device_capability(device)) + cases = [] + if args.preset == "quick": + for over in QUICK_PRESET: + a = argparse.Namespace(**vars(args)) + for k, v in over.items(): + setattr(a, k, v) + cases.append(a) + else: + cases.append(args) + rows = [] + for a in cases: + try: + rows.append(run_case(a, device)) + except Exception as exc: # keep going so one failing config does not hide the rest + print(f" FAILED: {type(exc).__name__}: {exc}") + rows.append( + dict( + label=f"{a.mode} ctx={a.ctx} bs={a.batch} sq={a.sq} chunk={a.chunk}", + rows=0, + error=f"{type(exc).__name__}: {exc}", + ) + ) + torch.cuda.empty_cache() + header = "| case | rows | Triton us | PrimTS us | Triton/PrimTS | check |\n|---|---:|---:|---:|---:|---|" + print(header) + for r in rows: + print( + md_row(r) + if "error" not in r + else f"| {r['label']} | - | error: {r['error'][:80]} | | | |" + ) + if args.out: + with open(args.out, "a") as f: + f.write( + f"\n### {torch.cuda.get_device_name(device)} {time.strftime('%Y-%m-%d %H:%M')}\n{header}\n" + ) + for r in rows: + f.write( + (md_row(r) if "error" not in r else f"| {r['label']} | - | error | | | |") + + "\n" + ) + if args.json: + with open(args.json, "a") as f: + for r in rows: + f.write(json.dumps(r) + "\n") + + +if __name__ == "__main__": + main() 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..9078cdfcff60 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_prims_ts_core.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +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 ( + QSA_SPARSE_CORE_ENV, + QSASparseParams, +) +from tensorrt_llm._torch.attention.backends.sparse.qsa.prims_ts_core import ( + prims_ts_available, + prims_ts_sparse_gqa, +) + +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_triton(monkeypatch) -> None: + monkeypatch.delenv(QSA_SPARSE_CORE_ENV, raising=False) + assert _params().sparse_core_backend == "triton" + assert _params(sparse_core_backend="prims_ts").sparse_core_backend == "prims_ts" + + +def test_sparse_core_backend_env_override(monkeypatch) -> None: + monkeypatch.setenv(QSA_SPARSE_CORE_ENV, "Auto") + assert _params().sparse_core_backend == "auto" + monkeypatch.setenv(QSA_SPARSE_CORE_ENV, "bogus") + with pytest.raises(ValueError, match="sparse_core_backend"): + _params() + + +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") + if not prims_ts_available(): + pytest.skip("flashinfer.attention.prims_ts is not installed") + + +def _make_case(query_lens: list[int], seed: int): + 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 + ) + v_cache = torch.randn_like(k_cache) + 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) From f42f37559108a6110f7404ef3cbbb1ac830f5c23 Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Wed, 21 Oct 2026 12:55:49 +0000 Subject: [PATCH 02/18] [None][feat] Add PrimTS QSA indexer scorer behind the sparse core switch Score compressed keys with a task-scheduled tcgen05 kernel instead of the Triton scorer when sparse_core_backend selects PrimTS. The kernel streams the paged index cache with one TMA box per 8-row core group, reduces the four indexer heads in registers and stores fp32 scores directly; tile metadata is built on device so the hot path has no host synchronization. Contract matches triton_qsa_paged_index_scores, including only_visible_blocks, so TopK consumes the logits unchanged. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../attention/backends/sparse/qsa/module.py | 45 +- .../backends/sparse/qsa/prims_ts_core.py | 27 + .../sparse/qsa/prims_ts_indexer/__init__.py | 17 + .../sparse/qsa/prims_ts_indexer/api.py | 218 ++++ .../sparse/qsa/prims_ts_indexer/kernel.py | 1052 +++++++++++++++++ .../sparse/qsa/test_qsa_prims_ts_indexer.py | 233 ++++ 6 files changed, 1578 insertions(+), 14 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/__init__.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/api.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/kernel.py create mode 100644 tests/unittest/_torch/attention/sparse/qsa/test_qsa_prims_ts_indexer.py diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py index e2c2c398a930..5f4881e55315 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py @@ -420,20 +420,37 @@ 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, - 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, - ) + # 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): + 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") 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 index ed8563786c26..b686f7b9bee8 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py @@ -62,6 +62,33 @@ def _unsupported_reason( return None +def use_prims_ts_indexer(params: QSASparseParams, q: torch.Tensor) -> bool: + """Whether the indexer scoring runs on the PrimTS kernel instead of Triton. + + Follows ``sparse_core_backend``: ``triton`` never, ``prims_ts`` always + (raising if unsupported), ``auto`` when the device and inputs qualify. + """ + backend = params.sparse_core_backend + if backend == "triton": + 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}" + 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, *, 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..1c0a4892ee68 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +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..a60a679ab36c --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/api.py @@ -0,0 +1,218 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""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 _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.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 = request_indices.contiguous() + 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 = ( + 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 + import cutlass + + block_table = block_table.contiguous() + q = q.contiguous() + index_cache = index_cache.contiguous() + # Raw pointers go into the compiled call; these tensors own the storage. + keepalive = ( + q, + index_cache, + block_table, + request_indices, + tile_meta, + positions, + scores, + ) + args = ( + _ptr(q, cutlass.BFloat16), + _ptr(index_cache, cutlass.BFloat16), + _ptr(block_table, cutlass.Int32), + _ptr(request_indices, cutlass.Int32), + _ptr(tile_meta, cutlass.Int32), + _ptr(positions, cutlass.Int32), + _ptr(scores, cutlass.Float32), + cutlass.Int32(num_requests), + cutlass.Int32(block_table.shape[1]), + cutlass.Int32(max_blocks), + cutlass.Int32(index_cache.shape[0] * blocks_per_page), + cutlass.Int32(rows), + cutlass.Int32(n_splits), + cutlass.Float32(head_dim**-0.5), + cutlass.Int32(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..f2c60e1c3026 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/kernel.py @@ -0,0 +1,1052 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""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 = physical_page * self.blocks_per_page + (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)), + 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), + physical_page * self.blocks_per_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, + block_rows: 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( + (block_rows, Q_HALVES, STAGE_K), + stride=(compress_ratio * HEAD_DIM, STAGE_K, 1), + ), + ) + tma_desc_k = cuda.create_tensor_map_tiled_from_view( + index_groups, + box_dims=(GROUP_ROWS, Q_HALVES, STAGE_K), + stride_order=(2, 0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + else: + index_groups = cute.make_tensor( + index_ptr, + cute.make_layout( + (block_rows, compress_ratio, HEAD_DIM), + stride=(compress_ratio * HEAD_DIM, HEAD_DIM, 1), + ), + ) + tma_desc_k = cuda.create_tensor_map_tiled_from_view( + index_groups, + box_dims=(blocks_per_page, 1, STAGE_K), + stride_order=(2, 1, 0), + 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_indexer.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_prims_ts_indexer.py new file mode 100644 index 000000000000..e024bd03d0c1 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_prims_ts_indexer.py @@ -0,0 +1,233 @@ +# 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()[0] != 10: + pytest.skip("Blackwell/Rubin tcgen05 required") + + +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) From 921da4c919722d15cbcee5223a504d74abfbbb26 Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:11:48 +0000 Subject: [PATCH 03/18] [None][fix] Align sliced QSA indexer inputs before the PrimTS scorer The prefill chunk loop passes row slices of packed metadata that start on 4 B boundaries; the scorer's raw-pointer path needs 16 B, so clone misaligned inputs on the host. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../backends/sparse/qsa/prims_ts_indexer/api.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) 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 index a60a679ab36c..f84cc52a7605 100644 --- 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 @@ -46,6 +46,13 @@ def _pick_n_splits(tiles_bound: int, n_max: int, min_splits: int, num_sms: int) 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 _ptr(t, dtype): import cutlass from cutlass.cute.runtime import make_ptr @@ -131,7 +138,7 @@ def prepare_qsa_index_scores( 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 = request_indices.contiguous() + 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 @@ -140,16 +147,16 @@ def prepare_qsa_index_scores( 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 = ( + 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 import cutlass - block_table = block_table.contiguous() - q = q.contiguous() - index_cache = index_cache.contiguous() + block_table = _aligned(block_table) + q = _aligned(q) + index_cache = _aligned(index_cache) # Raw pointers go into the compiled call; these tensors own the storage. keepalive = ( q, From 7b0ed867c909367ef2299055427b4a8ba172c722 Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:17:28 +0000 Subject: [PATCH 04/18] [None][fix] Remove per-call host costs in the QSA PrimTS path The index-K cache view from the V2 pool is coalesced with K/V (page stride larger than one page), so the PrimTS scorer's alignment helper copied the whole cache on every call (~0.8 ms per layer per step). The kernel's TMA map now takes the page stride as its own dimension and the view is passed through. The sparse core no longer reads request indices back to the host to build qo_indptr; the hook passes boundaries derived from the host seq_lens. Prefill rows are padded to a 256-row bucket so plan-cache misses (~40 ms each) stop recurring on every mixed prefill/decode iteration. Also wires the TLLM_QSA_INDEXER override that selects the indexer scorer independently of the sparse core. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../attention/backends/sparse/qsa/module.py | 7 ++ .../backends/sparse/qsa/prims_ts_core.py | 54 ++++++++++++-- .../sparse/qsa/prims_ts_indexer/api.py | 19 ++++- .../sparse/qsa/prims_ts_indexer/kernel.py | 33 +++++---- .../sparse/qsa/test_qsa_prims_ts_core.py | 73 +++++++++++++++++++ .../sparse/qsa/test_qsa_prims_ts_indexer.py | 28 +++++++ 6 files changed, 192 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py index 5f4881e55315..9380c6bb280a 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py @@ -766,6 +766,12 @@ 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 = torch.cumsum( + attn_metadata.seq_lens[: attn_metadata.num_seqs].to(torch.int64), dim=0 + ) output = torch.empty_like(q) for start in range(0, num_tokens, chunk_size): end = min(start + chunk_size, num_tokens) @@ -798,6 +804,7 @@ def forward( 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), ) continue output[packed_slice] = qsa_sparse_gqa( 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 index b686f7b9bee8..e54e729b0af1 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py @@ -3,6 +3,7 @@ """PrimTS QToken-KvBlock-Sparse core attention for QSA.""" import functools +import os import threading import torch @@ -12,6 +13,9 @@ from .params import QSASparseParams _PREFILL_GROUP_SIZE = 4 +# A plan miss costs ~40 ms per layer; padding prefill rows to a bucket bounds +# the plan keys to max_num_tokens / bucket per layer. +_PREFILL_ROW_BUCKET = 256 _MIN_DECODE_ROWS = 64 _SUPPORTED_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) _SUPPORTED_Q_DTYPES = (torch.float16, torch.bfloat16) @@ -65,10 +69,11 @@ def _unsupported_reason( def use_prims_ts_indexer(params: QSASparseParams, q: torch.Tensor) -> bool: """Whether the indexer scoring runs on the PrimTS kernel instead of Triton. - Follows ``sparse_core_backend``: ``triton`` never, ``prims_ts`` always - (raising if unsupported), ``auto`` when the device and inputs qualify. + Follows ``sparse_core_backend`` unless ``TLLM_QSA_INDEXER`` overrides it: + ``triton`` never, ``prims_ts`` always (raising if unsupported), ``auto`` + when the device and inputs qualify. """ - backend = params.sparse_core_backend + backend = os.environ.get("TLLM_QSA_INDEXER", params.sparse_core_backend) if backend == "triton": return False reason = None @@ -131,6 +136,24 @@ def _request_boundaries(request_indices: torch.Tensor) -> torch.Tensor: ) +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: torch.Tensor, start: int, end: int) -> torch.Tensor: + """Request boundaries of packed rows ``[start, end)`` from host cumulative lengths.""" + inner = cu_seq_lens[(cu_seq_lens > start) & (cu_seq_lens < end)] - start + return torch.cat( + ( + torch.zeros(1, dtype=torch.int32), + inner.to(torch.int32), + torch.tensor([end - start], dtype=torch.int32), + ) + ) + + class _PrimsTSSparseCore: def __init__(self) -> None: self._plans: dict[tuple, object] = {} @@ -161,6 +184,7 @@ def run( max_seq_len_kv: int, softmax_scale: float, is_decode: bool, + request_boundaries: torch.Tensor | None = None, ) -> torch.Tensor: wrapper_cls, workspace_size, make_qo_indptr = _prims_ts_api() rows, num_qo_heads, head_dim = q.shape @@ -180,6 +204,16 @@ def run( request_indices = request_indices.to(torch.int32).contiguous() q = q.contiguous() + num_rows = rows + if not is_decode: + rows = -(-rows // _PREFILL_ROW_BUCKET) * _PREFILL_ROW_BUCKET + if rows != num_rows: + # Padding rows attend to nothing but token 0 of request 0; the + # suffix gets its own routes and is sliced off the output. + q = _pad_rows(q, rows, 0) + block_ids = _pad_rows(block_ids, rows, -1) + request_indices = _pad_rows(request_indices, rows, 0) + query_positions = _pad_rows(query_positions, rows, 0) if is_decode: route_q = q.view(rows, 1, 1, num_qo_heads, head_dim) qo_indptr = None @@ -188,8 +222,12 @@ def run( 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) qo_indptr = make_qo_indptr( - _request_boundaries(request_indices), + request_boundaries, rows, group_size=_PREFILL_GROUP_SIZE, device=q.device, @@ -291,7 +329,7 @@ def run( ) raise RuntimeError(f"{error}; ABI changes vs last eager run: {changed}") from error self._last_abi[key] = abi - return out + return out[:num_rows] _CORE = _PrimsTSSparseCore() @@ -310,12 +348,14 @@ def prims_ts_sparse_gqa( max_seq_len_kv: int, softmax_scale: float, is_decode: bool, + request_boundaries: 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. + layout with split-KV; prefill packs rows into request-safe groups of four, + using host ``request_boundaries`` (Int32 offsets into the rows) when given. """ return _CORE.run( q, @@ -329,6 +369,7 @@ def prims_ts_sparse_gqa( max_seq_len_kv=max_seq_len_kv, softmax_scale=softmax_scale, is_decode=is_decode, + request_boundaries=request_boundaries, ) @@ -336,4 +377,5 @@ def prims_ts_sparse_gqa( "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/api.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_indexer/api.py index f84cc52a7605..a9e6130f1de7 100644 --- 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 @@ -53,6 +53,19 @@ def _aligned(t: torch.Tensor) -> torch.Tensor: 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 @@ -85,6 +98,7 @@ def _compiled(blocks_per_page: int, compress_ratio: int): cutlass.Int32(1), cutlass.Int32(1), cutlass.Int32(1), + cutlass.Int32(1), cutlass.Float32(1.0), cutlass.Int32(1), blocks_per_page, @@ -156,7 +170,7 @@ def prepare_qsa_index_scores( block_table = _aligned(block_table) q = _aligned(q) - index_cache = _aligned(index_cache) + index_cache = _paged_view(index_cache) # Raw pointers go into the compiled call; these tensors own the storage. keepalive = ( q, @@ -178,7 +192,8 @@ def prepare_qsa_index_scores( cutlass.Int32(num_requests), cutlass.Int32(block_table.shape[1]), cutlass.Int32(max_blocks), - cutlass.Int32(index_cache.shape[0] * blocks_per_page), + cutlass.Int32(index_cache.shape[0]), + cutlass.Int32(index_cache.stride(0)), cutlass.Int32(rows), cutlass.Int32(n_splits), cutlass.Float32(head_dim**-0.5), 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 index f2c60e1c3026..694ab28c2342 100644 --- 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 @@ -348,11 +348,11 @@ def tma_load(self, stage_info: StageInfo) -> None: 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 = physical_page * self.blocks_per_page + (lane % groups_per_page) * GROUP_ROWS + 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)), + (cutlass.Int32(0), row0, cutlass.Int32(0), physical_page), stage_info.barrier, ) else: @@ -365,7 +365,8 @@ def tma_load(self, stage_info: StageInfo) -> None: ( cutlass.Int32(half * STAGE_K), cutlass.Int32(self.compress_ratio - 1), - physical_page * self.blocks_per_page, + cutlass.Int32(0), + physical_page, ), stage_info.barrier, ) @@ -464,8 +465,11 @@ def mma( @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.""" + """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) @@ -972,7 +976,8 @@ def qsa_index_scores_host( num_requests: cutlass.Int32, max_pages: cutlass.Int32, max_blocks: cutlass.Int32, - block_rows: cutlass.Int32, + num_pages: cutlass.Int32, + page_stride: cutlass.Int32, num_rows: cutlass.Int32, n_splits: cutlass.Int32, scale: cutlass.Float32, @@ -1012,28 +1017,28 @@ def qsa_index_scores_host( index_groups = cute.make_tensor( index_ptr + (compress_ratio - 1) * HEAD_DIM, cute.make_layout( - (block_rows, Q_HALVES, STAGE_K), - stride=(compress_ratio * HEAD_DIM, STAGE_K, 1), + (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, Q_HALVES, STAGE_K), - stride_order=(2, 0, 1), + 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( - (block_rows, compress_ratio, HEAD_DIM), - stride=(compress_ratio * HEAD_DIM, HEAD_DIM, 1), + (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, STAGE_K), - stride_order=(2, 1, 0), + box_dims=(blocks_per_page, 1, 1, STAGE_K), + stride_order=(3, 2, 0, 1), swizzle=cuda.TensorMapSwizzle.s128b, ) qsa_index_scores_kernel( 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 index 9078cdfcff60..87d286f735e7 100644 --- 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 @@ -14,6 +14,7 @@ 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 @@ -160,3 +161,75 @@ def test_prims_ts_sparse_gqa_matches_reference(query_lens, is_decode) -> None: 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 = torch.tensor([37, 128, 131, 200], dtype=torch.int64) + assert slice_request_boundaries(cu_seq_lens, 0, 200).tolist() == [0, 37, 128, 131, 200] + assert slice_request_boundaries(cu_seq_lens, 0, 100).tolist() == [0, 37, 100] + assert slice_request_boundaries(cu_seq_lens, 37, 128).tolist() == [0, 91] + assert slice_request_boundaries(cu_seq_lens, 100, 140).tolist() == [0, 28, 31, 40] + assert slice_request_boundaries(cu_seq_lens, 131, 200).dtype == torch.int32 + + +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 = torch.cumsum(torch.tensor(query_lens, dtype=torch.int64), dim=0) + 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( + [37, 91], seed=5 + ) + prims_ts_sparse_gqa( + q, k_cache, v_cache, block_table, block_ids, request_indices, positions, **kwargs + ) + plans = len(_CORE._plans) + q2, _, _, block_table2, block_ids2, request_indices2, positions2, selected2, metadata = ( + _make_case([100, 30], 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] == 130 + 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) 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 index e024bd03d0c1..4d6e9fecf98d 100644 --- 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 @@ -231,3 +231,31 @@ def test_prepared_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) From 3ceef11944e18afac4777e05163c4d8837ffc3bb Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Wed, 21 Oct 2026 19:56:08 +0000 Subject: [PATCH 05/18] [None][perf] Remove host stalls from the QSA PrimTS prefill path The FlashInfer QToken wrapper re-prepares its launch, with host readbacks, whenever the qo_indptr storage changes, and a fresh qo_indptr was built for every layer. Route offsets now live in a persistent buffer per route count, shared by all QSA layers of a forward and rewritten only when the request layout changes. Prefill chunks are aligned to the 256-row plan bucket so full chunks need no padding, the core writes into the caller's output slice, request boundaries stay host ints, and the indexer passes raw addresses and scalars to its TVM-FFI entry. auto now gates decode by measured Rubin crossovers: the core runs PrimTS from 2 decode rows (was 64), the indexer from 32 (no gate before). Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../attention/backends/sparse/qsa/module.py | 13 +- .../backends/sparse/qsa/prims_ts_core.py | 137 +++++++++++------ .../sparse/qsa/prims_ts_indexer/api.py | 35 ++--- .../sparse/qsa/test_qsa_prims_ts_core.py | 143 +++++++++++++++++- 4 files changed, 254 insertions(+), 74 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py index 9380c6bb280a..8bbf43d2724e 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 @@ -426,7 +427,7 @@ def select_qsa_paged_tokens( 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): + if use_prims_ts_indexer(params, q, is_decode=not context_rows): from .prims_ts_indexer import prims_ts_qsa_index_scores logits = prims_ts_qsa_index_scores( @@ -767,10 +768,11 @@ def forward( ) chunk_size = _query_chunk_size(num_tokens, score_columns) if use_prims_ts: - from .prims_ts_core import slice_request_boundaries + from .prims_ts_core import align_prefill_chunk, slice_request_boundaries - cu_seq_lens = torch.cumsum( - attn_metadata.seq_lens[: attn_metadata.num_seqs].to(torch.int64), dim=0 + chunk_size = align_prefill_chunk(chunk_size) + 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): @@ -792,7 +794,7 @@ def forward( expand_blocks=not use_prims_ts, ) if use_prims_ts: - output[packed_slice] = prims_ts_sparse_gqa( + prims_ts_sparse_gqa( q[packed_slice], k_cache, v_cache, @@ -805,6 +807,7 @@ def forward( softmax_scale=softmax_scale, is_decode=False, request_boundaries=slice_request_boundaries(cu_seq_lens, start, end), + out=output[packed_slice], ) continue output[packed_slice] = qsa_sparse_gqa( 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 index e54e729b0af1..101472bc8ee8 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """PrimTS QToken-KvBlock-Sparse core attention for QSA.""" +import bisect import functools import os import threading @@ -16,7 +17,10 @@ # A plan miss costs ~40 ms per layer; padding prefill rows to a bucket bounds # the plan keys to max_num_tokens / bucket per layer. _PREFILL_ROW_BUCKET = 256 -_MIN_DECODE_ROWS = 64 +# Decode rows below which Triton is faster under CUDA graphs (Rubin SM107, +# 32K context, TP2): the core ties at one row, the indexer until ~32 rows. +_MIN_DECODE_ROWS = 2 +_MIN_INDEXER_DECODE_ROWS = 32 _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) @@ -66,16 +70,20 @@ def _unsupported_reason( return None -def use_prims_ts_indexer(params: QSASparseParams, q: torch.Tensor) -> bool: +def use_prims_ts_indexer( + params: QSASparseParams, q: torch.Tensor, *, is_decode: bool = False +) -> bool: """Whether the indexer scoring runs on the PrimTS kernel instead of Triton. Follows ``sparse_core_backend`` unless ``TLLM_QSA_INDEXER`` overrides it: ``triton`` never, ``prims_ts`` always (raising if unsupported), ``auto`` - when the device and inputs qualify. + when the device and inputs qualify and decode batches are large enough. """ backend = os.environ.get("TLLM_QSA_INDEXER", 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" @@ -124,16 +132,10 @@ def resolve_sparse_core_backend( return "prims_ts" -def _request_boundaries(request_indices: torch.Tensor) -> torch.Tensor: +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 torch.cat( - ( - torch.zeros(1, dtype=torch.int32), - change.to(torch.int32), - torch.tensor([req.numel()], dtype=torch.int32), - ) - ) + return (0, *change.tolist(), req.numel()) def _pad_rows(t: torch.Tensor, rows: int, fill: float) -> torch.Tensor: @@ -142,16 +144,21 @@ def _pad_rows(t: torch.Tensor, rows: int, fill: float) -> torch.Tensor: return padded -def slice_request_boundaries(cu_seq_lens: torch.Tensor, start: int, end: int) -> torch.Tensor: +def align_prefill_chunk(rows: int) -> int: + """Largest bucket multiple not above ``rows``, so full chunks need no padding.""" + return max(rows // _PREFILL_ROW_BUCKET * _PREFILL_ROW_BUCKET, min(rows, _PREFILL_ROW_BUCKET)) + + +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[(cu_seq_lens > start) & (cu_seq_lens < end)] - start - return torch.cat( - ( - torch.zeros(1, dtype=torch.int32), - inner.to(torch.int32), - torch.tensor([end - start], dtype=torch.int32), - ) - ) + 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: @@ -159,6 +166,9 @@ def __init__(self) -> None: self._plans: dict[tuple, object] = {} self._last_abi: dict[tuple, tuple] = {} self._workspace: torch.Tensor | None = None + 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: @@ -170,6 +180,38 @@ def _workspace_for(self, nbytes: int, device: torch.device) -> torch.Tensor: 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. All QSA layers of a forward share one + # request layout, so offsets live in a persistent buffer per route count + # that is rewritten, stream-ordered, only when the layout changes. + 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", + ) + 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, @@ -184,7 +226,8 @@ def run( max_seq_len_kv: int, softmax_scale: float, is_decode: bool, - request_boundaries: torch.Tensor | None = None, + 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 @@ -226,12 +269,8 @@ def run( # Host readback; callers with host sequence lengths pass the # boundaries to keep the stream running. request_boundaries = _request_boundaries(request_indices) - qo_indptr = make_qo_indptr( - request_boundaries, - rows, - group_size=_PREFILL_GROUP_SIZE, - device=q.device, - ) + with self._lock: + qo_indptr = self._route_offsets(make_qo_indptr, request_boundaries, rows, q.device) group_size = _PREFILL_GROUP_SIZE num_routes = qo_indptr.numel() - 1 split_kv = False @@ -257,7 +296,9 @@ def run( page_size, block_topk, ) - out = torch.empty((rows, num_qo_heads, head_dim), dtype=out_dtype, device=q.device) + result = out + if out is None or rows != num_rows: + out = torch.empty((rows, num_qo_heads, head_dim), dtype=out_dtype, device=q.device) with self._lock: wrapper = self._plans.get(key) if wrapper is None: @@ -297,17 +338,14 @@ def run( ) self._plans[key] = wrapper run_out = out.view_as(route_q) if is_decode else out - abi = tuple( - (name, tuple(t.shape), tuple(t.stride()), t.dtype) - for name, t in ( - ("q", route_q), - ("block_ids", block_ids), - ("block_table", block_table), - ("request_indices", request_indices), - ("query_positions", query_positions), - ("out", run_out), - ) - ) + (("k_ptr", k_cache.data_ptr()), ("v_ptr", v_cache.data_ptr())) + 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, @@ -323,13 +361,18 @@ def run( except RuntimeError as error: previous = self._last_abi.get(key) changed = ( - [(a, b) for a, b in zip(previous or (), abi) if a != b] + [(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 last eager run: {changed}") from error - self._last_abi[key] = abi - return out[:num_rows] + 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[:num_rows] + if result is not out: + result.copy_(out[:num_rows]) + return result _CORE = _PrimsTSSparseCore() @@ -348,14 +391,16 @@ def prims_ts_sparse_gqa( max_seq_len_kv: int, softmax_scale: float, is_decode: bool, - request_boundaries: torch.Tensor | None = None, + 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`` (Int32 offsets into the rows) when given. + 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, @@ -370,12 +415,14 @@ def prims_ts_sparse_gqa( softmax_scale=softmax_scale, is_decode=is_decode, request_boundaries=request_boundaries, + out=out, ) __all__ = [ "prims_ts_available", "prims_ts_sparse_gqa", + "align_prefill_chunk", "resolve_sparse_core_backend", "slice_request_boundaries", ] 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 index a9e6130f1de7..c114e42300f3 100644 --- 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 @@ -166,7 +166,6 @@ def prepare_qsa_index_scores( ) fn = _compiled(blocks_per_page, compress_ratio) import cuda.bindings.driver as cuda_driver - import cutlass block_table = _aligned(block_table) q = _aligned(q) @@ -181,23 +180,25 @@ def prepare_qsa_index_scores( 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 = ( - _ptr(q, cutlass.BFloat16), - _ptr(index_cache, cutlass.BFloat16), - _ptr(block_table, cutlass.Int32), - _ptr(request_indices, cutlass.Int32), - _ptr(tile_meta, cutlass.Int32), - _ptr(positions, cutlass.Int32), - _ptr(scores, cutlass.Float32), - cutlass.Int32(num_requests), - cutlass.Int32(block_table.shape[1]), - cutlass.Int32(max_blocks), - cutlass.Int32(index_cache.shape[0]), - cutlass.Int32(index_cache.stride(0)), - cutlass.Int32(rows), - cutlass.Int32(n_splits), - cutlass.Float32(head_dim**-0.5), - cutlass.Int32(grid_size), + 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(): 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 index 87d286f735e7..41d11f4431b2 100644 --- 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 @@ -1,6 +1,7 @@ # 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 @@ -12,6 +13,7 @@ QSASparseParams, ) from tensorrt_llm._torch.attention.backends.sparse.qsa.prims_ts_core import ( + align_prefill_chunk, prims_ts_available, prims_ts_sparse_gqa, slice_request_boundaries, @@ -164,12 +166,12 @@ def test_prims_ts_sparse_gqa_matches_reference(query_lens, is_decode) -> None: def test_slice_request_boundaries() -> None: - cu_seq_lens = torch.tensor([37, 128, 131, 200], dtype=torch.int64) - assert slice_request_boundaries(cu_seq_lens, 0, 200).tolist() == [0, 37, 128, 131, 200] - assert slice_request_boundaries(cu_seq_lens, 0, 100).tolist() == [0, 37, 100] - assert slice_request_boundaries(cu_seq_lens, 37, 128).tolist() == [0, 91] - assert slice_request_boundaries(cu_seq_lens, 100, 140).tolist() == [0, 28, 31, 40] - assert slice_request_boundaries(cu_seq_lens, 131, 200).dtype == torch.int32 + 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: @@ -184,7 +186,7 @@ def test_prims_ts_sparse_gqa_host_boundaries_match_readback() -> None: readback = prims_ts_sparse_gqa( q, k_cache, v_cache, block_table, block_ids, request_indices, positions, **kwargs ) - cu_seq_lens = torch.cumsum(torch.tensor(query_lens, dtype=torch.int64), dim=0) + cu_seq_lens = list(itertools.accumulate(query_lens)) host = prims_ts_sparse_gqa( q, k_cache, @@ -233,3 +235,130 @@ def test_prefill_rows_share_plans_within_bucket() -> None: ) 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 flashinfer.attention.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)] + finally: + fi_metadata._prepare_q_token_kv_block_sparse_attention = original + torch.cuda.synchronize() + assert prepares == 0 + for out, ref in zip(outputs, expected): + assert torch.equal(out, ref) + + +def test_auto_gates_small_decode_batches(monkeypatch) -> None: + _requires_prims_ts() + from tensorrt_llm._torch.attention.backends.sparse.qsa.prims_ts_core import ( + resolve_sparse_core_backend, + use_prims_ts_indexer, + ) + + monkeypatch.delenv("TLLM_QSA_SPARSE_CORE", raising=False) + monkeypatch.delenv("TLLM_QSA_INDEXER", raising=False) + 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) + + +def test_align_prefill_chunk() -> None: + assert align_prefill_chunk(4048) == 3840 + assert align_prefill_chunk(8192) == 8192 + assert align_prefill_chunk(256) == 256 + assert align_prefill_chunk(255) == 255 + assert align_prefill_chunk(96) == 96 + + +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) From 4005b4b28714ac2b918b4757738e3f0fbee8cce3 Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:47:19 +0000 Subject: [PATCH 06/18] [None][perf] Overlap QSA PrimTS prefill chunks with the next chunk's scoring Queue the next chunk's indexer scoring and Top-K before preparing the current chunk's PrimTS launch on the host, so the GPU keeps working while the FlashInfer wrapper runs. Chunks write disjoint Top-K rows, so the reorder changes no results. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../attention/backends/sparse/qsa/module.py | 55 ++++++++++++------- 1 file changed, 35 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py index 8bbf43d2724e..6960fc2bccf7 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py @@ -775,10 +775,10 @@ def forward( 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], @@ -793,34 +793,49 @@ def forward( context_rows=True, expand_blocks=not use_prims_ts, ) - if use_prims_ts: + + 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. Chunks write disjoint Top-K rows. + 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[packed_slice], + q[start:end], k_cache, v_cache, attn_metadata.qsa_block_table, selected, - req_idx[packed_slice], - logical[packed_slice], + 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[packed_slice], + 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, ) - continue - 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=softmax_scale, - query_positions=logical[packed_slice], - 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) From cb42e22609f40b9f0b3428ee98ffe3d982cae21d Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:39:34 +0000 Subject: [PATCH 07/18] [None][feat] Select the QSA PrimTS backend from the LLM API, default auto Expose QSASparseAttentionConfig.sparse_core_backend (auto | triton | prims_ts) and drop the TLLM_QSA_SPARSE_CORE / TLLM_QSA_INDEXER environment overrides. The default is auto: PrimTS on SM100/SM103/SM107 with Triton kept for small decode batches (core below 2 rows, indexer below 32). QSASparseParams.indexer_backend stays internal for A/B measurements. Regenerates the LLM-args golden manifest and the trtllm-serve schemas for the new field. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../attention/backends/sparse/qsa/params.py | 20 +++++++------ .../backends/sparse/qsa/prims_ts_core.py | 7 +++-- tensorrt_llm/llmapi/llm_args.py | 9 ++++++ .../schemas/trtllm-serve-config.schema.json | 11 +++++++ .../trtllm-serve-disagg-config.schema.json | 11 +++++++ .../usage/llm_args_golden_manifest.json | 10 +++++++ .../sparse/qsa/test_qsa_prims_ts_core.py | 29 ++++++++----------- 7 files changed, 68 insertions(+), 29 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py index 3bd6dd932843..fb780df7e48f 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py @@ -2,14 +2,13 @@ # SPDX-License-Identifier: Apache-2.0 """Internal sparse-attention parameters for QSA.""" -import os from dataclasses import dataclass from ..params import SparseMetadataParams, SparseParams from .constants import QSA_INDEX_KV_HEADS QSA_SPARSE_CORE_BACKENDS = ("auto", "triton", "prims_ts") -QSA_SPARSE_CORE_ENV = "TLLM_QSA_SPARSE_CORE" +QSA_INDEXER_BACKENDS = ("follow", *QSA_SPARSE_CORE_BACKENDS) def _validate_selection_geometry(token_topk: int, compress_ratio: int) -> None: @@ -45,20 +44,23 @@ class QSASparseParams(SparseParams): enable_heuristic_topk: bool = False # Reuse one indexer selection across the MTP draft loop. mtp_index_share: bool = False - # Sparse core attention kernel: "triton", "prims_ts", or "auto" (PrimTS - # for prefill and large decode batches on SM100/SM103/SM107, else Triton). - # TLLM_QSA_SPARSE_CORE overrides the configured value. - sparse_core_backend: str = "triton" + # Sparse core attention kernel: "auto" (PrimTS on SM100/SM103/SM107 except + # for small decode batches, else Triton), "triton", or "prims_ts". + sparse_core_backend: str = "auto" + # Indexer scorer kernel; "follow" uses the sparse core selection. + indexer_backend: str = "follow" def __post_init__(self) -> None: - env_backend = os.environ.get(QSA_SPARSE_CORE_ENV) - if env_backend: - self.sparse_core_backend = env_backend.strip().lower() 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 index 101472bc8ee8..8c6000daff2b 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py @@ -4,7 +4,6 @@ import bisect import functools -import os import threading import torch @@ -75,11 +74,13 @@ def use_prims_ts_indexer( ) -> bool: """Whether the indexer scoring runs on the PrimTS kernel instead of Triton. - Follows ``sparse_core_backend`` unless ``TLLM_QSA_INDEXER`` overrides it: + 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 = os.environ.get("TLLM_QSA_INDEXER", params.sparse_core_backend) + 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: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 55a23ac1a87b..4028f541d9ba 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -967,6 +967,14 @@ class QSASparseAttentionConfig(SeqLenAwareSparseAttentionConfig): "recompute; compressed groups that complete during the loop are " "appended at lookup. When omitted, the checkpoint config supplies the " "value, defaulting to off.") + sparse_core_backend: Literal["auto", "triton", "prims_ts"] = Field( + default="auto", + description= + "Kernel for the QSA sparse core attention and indexer scorer. 'auto' " + "runs the PrimTS (CUTLASS task-scheduling) kernels on SM100/SM103/" + "SM107 and keeps Triton for small decode batches, 'triton' always " + "uses the Triton kernels, 'prims_ts' always uses PrimTS and raises " + "when the device does not support it.") # Index projection dimensions, compression, and selection budget are part # of the checkpoint contract rather than serving-time tuning knobs. _resolved_params: Optional["QSASparseParams"] = PrivateAttr(default=None) @@ -1030,6 +1038,7 @@ def to_sparse_params(self, **kwargs: object) -> "QSASparseParams": seq_len_threshold=seq_len_threshold, enable_heuristic_topk=self.enable_heuristic_topk, mtp_index_share=self._mtp_index_share(pretrained_config), + sparse_core_backend=self.sparse_core_backend, ) def _mtp_index_share(self, pretrained_config: object) -> bool: diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index f5f9184c5557..e9c70b3ac5ee 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -5083,6 +5083,17 @@ "default": null, "description": "The sequence length threshold separating dense and QSA attention. When omitted, it is resolved to token_topk after checkpoint geometry is loaded; an explicit value below token_topk is raised to token_topk because it cannot reduce attention work.", "title": "Seq Len Threshold" + }, + "sparse_core_backend": { + "default": "auto", + "description": "Kernel for the QSA sparse core attention and indexer scorer. 'auto' runs the PrimTS (CUTLASS task-scheduling) kernels on SM100/SM103/SM107 and keeps Triton for small decode batches, 'triton' always uses the Triton kernels, 'prims_ts' always uses PrimTS and raises when the device does not support it.", + "enum": [ + "auto", + "triton", + "prims_ts" + ], + "title": "Sparse Core Backend", + "type": "string" } }, "title": "QSASparseAttentionConfig", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index c524c9fe14d4..9004fd71098b 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -6633,6 +6633,17 @@ "default": null, "description": "The sequence length threshold separating dense and QSA attention. When omitted, it is resolved to token_topk after checkpoint geometry is loaded; an explicit value below token_topk is raised to token_topk because it cannot reduce attention work.", "title": "Seq Len Threshold" + }, + "sparse_core_backend": { + "default": "auto", + "description": "Kernel for the QSA sparse core attention and indexer scorer. 'auto' runs the PrimTS (CUTLASS task-scheduling) kernels on SM100/SM103/SM107 and keeps Triton for small decode batches, 'triton' always uses the Triton kernels, 'prims_ts' always uses PrimTS and raises when the device does not support it.", + "enum": [ + "auto", + "triton", + "prims_ts" + ], + "title": "Sparse Core Backend", + "type": "string" } }, "title": "QSASparseAttentionConfig", diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index e58f42598f2d..06d10b1826f1 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1525,6 +1525,16 @@ "kind": "value", "path": "sparse_attention_config.sparse_block_size" }, + { + "allowed_values": [ + "auto", + "triton", + "prims_ts" + ], + "capture_policy": "literal", + "kind": "categorical", + "path": "sparse_attention_config.sparse_core_backend" + }, { "capture_policy": "bool", "kind": "value", 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 index 41d11f4431b2..5da8352b4394 100644 --- 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 @@ -8,10 +8,7 @@ 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 ( - QSA_SPARSE_CORE_ENV, - QSASparseParams, -) +from tensorrt_llm._torch.attention.backends.sparse.qsa.params import QSASparseParams from tensorrt_llm._torch.attention.backends.sparse.qsa.prims_ts_core import ( align_prefill_chunk, prims_ts_available, @@ -41,18 +38,14 @@ def _params(**overrides) -> QSASparseParams: return QSASparseParams(**kwargs) -def test_sparse_core_backend_defaults_to_triton(monkeypatch) -> None: - monkeypatch.delenv(QSA_SPARSE_CORE_ENV, raising=False) - assert _params().sparse_core_backend == "triton" - assert _params(sparse_core_backend="prims_ts").sparse_core_backend == "prims_ts" - - -def test_sparse_core_backend_env_override(monkeypatch) -> None: - monkeypatch.setenv(QSA_SPARSE_CORE_ENV, "Auto") +def test_sparse_core_backend_defaults_to_auto() -> None: assert _params().sparse_core_backend == "auto" - monkeypatch.setenv(QSA_SPARSE_CORE_ENV, "bogus") + 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() + _params(sparse_core_backend="bogus") + with pytest.raises(ValueError, match="indexer_backend"): + _params(indexer_backend="bogus") def _requires_prims_ts() -> None: @@ -289,15 +282,13 @@ def counted(*args, **kw): assert torch.equal(out, ref) -def test_auto_gates_small_decode_batches(monkeypatch) -> None: +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, ) - monkeypatch.delenv("TLLM_QSA_SPARSE_CORE", raising=False) - monkeypatch.delenv("TLLM_QSA_INDEXER", raising=False) device = torch.device("cuda") k_cache = torch.empty(1, NUM_KV_HEADS, PAGE_SIZE, HEAD_DIM, dtype=torch.bfloat16, device=device) @@ -325,6 +316,10 @@ def indexer(backend, rows, is_decode): assert indexer("auto", 1, False) assert indexer("prims_ts", 1, True) assert not indexer("triton", 4096, False) + assert not use_prims_ts_indexer( + _params(sparse_core_backend="prims_ts", indexer_backend="triton"), + torch.empty(64, 4, 128, dtype=torch.bfloat16, device=device), + ) def test_align_prefill_chunk() -> None: From 94f7e07f3e6a02ad1ee05fff7c6809290a31f418 Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Thu, 22 Oct 2026 21:15:16 +0000 Subject: [PATCH 08/18] [None][refactor] Select QSA kernels at runtime without a user-facing switch Match the other sparse backends (glm_kpool, deepseek_v4): the PrimTS core and indexer are chosen in code from device, dtype and decode batch size, with Triton as the fallback. Drop the QSASparseAttentionConfig field and the regenerated manifest and schemas; QSASparseParams keeps internal sparse_core_backend / indexer_backend fields for tests and A/B runs. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../_torch/attention/backends/sparse/qsa/params.py | 7 ++++--- tensorrt_llm/llmapi/llm_args.py | 9 --------- tensorrt_llm/schemas/trtllm-serve-config.schema.json | 11 ----------- .../schemas/trtllm-serve-disagg-config.schema.json | 11 ----------- tensorrt_llm/usage/llm_args_golden_manifest.json | 10 ---------- 5 files changed, 4 insertions(+), 44 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py index fb780df7e48f..203ef42a564f 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py @@ -44,10 +44,11 @@ class QSASparseParams(SparseParams): enable_heuristic_topk: bool = False # Reuse one indexer selection across the MTP draft loop. mtp_index_share: bool = False - # Sparse core attention kernel: "auto" (PrimTS on SM100/SM103/SM107 except - # for small decode batches, else Triton), "triton", or "prims_ts". + # 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 scorer kernel; "follow" uses the sparse core selection. indexer_backend: str = "follow" def __post_init__(self) -> None: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 4028f541d9ba..55a23ac1a87b 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -967,14 +967,6 @@ class QSASparseAttentionConfig(SeqLenAwareSparseAttentionConfig): "recompute; compressed groups that complete during the loop are " "appended at lookup. When omitted, the checkpoint config supplies the " "value, defaulting to off.") - sparse_core_backend: Literal["auto", "triton", "prims_ts"] = Field( - default="auto", - description= - "Kernel for the QSA sparse core attention and indexer scorer. 'auto' " - "runs the PrimTS (CUTLASS task-scheduling) kernels on SM100/SM103/" - "SM107 and keeps Triton for small decode batches, 'triton' always " - "uses the Triton kernels, 'prims_ts' always uses PrimTS and raises " - "when the device does not support it.") # Index projection dimensions, compression, and selection budget are part # of the checkpoint contract rather than serving-time tuning knobs. _resolved_params: Optional["QSASparseParams"] = PrivateAttr(default=None) @@ -1038,7 +1030,6 @@ def to_sparse_params(self, **kwargs: object) -> "QSASparseParams": seq_len_threshold=seq_len_threshold, enable_heuristic_topk=self.enable_heuristic_topk, mtp_index_share=self._mtp_index_share(pretrained_config), - sparse_core_backend=self.sparse_core_backend, ) def _mtp_index_share(self, pretrained_config: object) -> bool: diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index e9c70b3ac5ee..f5f9184c5557 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -5083,17 +5083,6 @@ "default": null, "description": "The sequence length threshold separating dense and QSA attention. When omitted, it is resolved to token_topk after checkpoint geometry is loaded; an explicit value below token_topk is raised to token_topk because it cannot reduce attention work.", "title": "Seq Len Threshold" - }, - "sparse_core_backend": { - "default": "auto", - "description": "Kernel for the QSA sparse core attention and indexer scorer. 'auto' runs the PrimTS (CUTLASS task-scheduling) kernels on SM100/SM103/SM107 and keeps Triton for small decode batches, 'triton' always uses the Triton kernels, 'prims_ts' always uses PrimTS and raises when the device does not support it.", - "enum": [ - "auto", - "triton", - "prims_ts" - ], - "title": "Sparse Core Backend", - "type": "string" } }, "title": "QSASparseAttentionConfig", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 9004fd71098b..c524c9fe14d4 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -6633,17 +6633,6 @@ "default": null, "description": "The sequence length threshold separating dense and QSA attention. When omitted, it is resolved to token_topk after checkpoint geometry is loaded; an explicit value below token_topk is raised to token_topk because it cannot reduce attention work.", "title": "Seq Len Threshold" - }, - "sparse_core_backend": { - "default": "auto", - "description": "Kernel for the QSA sparse core attention and indexer scorer. 'auto' runs the PrimTS (CUTLASS task-scheduling) kernels on SM100/SM103/SM107 and keeps Triton for small decode batches, 'triton' always uses the Triton kernels, 'prims_ts' always uses PrimTS and raises when the device does not support it.", - "enum": [ - "auto", - "triton", - "prims_ts" - ], - "title": "Sparse Core Backend", - "type": "string" } }, "title": "QSASparseAttentionConfig", diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 06d10b1826f1..e58f42598f2d 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1525,16 +1525,6 @@ "kind": "value", "path": "sparse_attention_config.sparse_block_size" }, - { - "allowed_values": [ - "auto", - "triton", - "prims_ts" - ], - "capture_policy": "literal", - "kind": "categorical", - "path": "sparse_attention_config.sparse_core_backend" - }, { "capture_policy": "bool", "kind": "value", From ee49a8b3a8808045dacae8f8a889db40bbc30479 Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Fri, 23 Oct 2026 15:28:22 +0000 Subject: [PATCH 09/18] [None][fix] Address review: vendored PrimTS import, static FP8 Q scale, decode buckets Import the QToken sparse wrapper from the vendored backends/prims_ts package like the other PrimTS consumers, using its seq_len_q plan API. Cast Q to FP8 with unit scale for FP8 caches instead of a per-call dynamic scale, which needed a host readback that CUDA-graph capture forbids. Bucket decode rows to 8 so eager decode batches reuse plans. Gate the PrimTS indexer on page geometry and DSL importability, pass kv_block_size by name, and give the indexer files the NVIDIA header. Tests: FP8 K/V prefill and decode against the reference, decode graph capture and replay, decode plan buckets, independent indexer overrides; the core tests fail instead of skipping when the vendored kernel is missing on a supported GPU. The microbenchmark exits non-zero on a reference mismatch. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../attention/backends/sparse/qsa/module.py | 7 +- .../backends/sparse/qsa/prims_ts_core.py | 98 +++++++----- .../sparse/qsa/prims_ts_indexer/__init__.py | 15 +- .../sparse/qsa/prims_ts_indexer/api.py | 15 +- .../sparse/qsa/prims_ts_indexer/kernel.py | 15 +- .../qsa/bench_qsa_sparse_core.py | 13 +- .../sparse/qsa/test_qsa_prims_ts_core.py | 139 +++++++++++++++++- .../sparse/qsa/test_qsa_prims_ts_indexer.py | 5 +- 8 files changed, 218 insertions(+), 89 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py index 6960fc2bccf7..bf28b600726f 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py @@ -427,7 +427,12 @@ def select_qsa_paged_tokens( 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, is_decode=not context_rows): + if use_prims_ts_indexer( + params, + q, + tokens_per_block=metadata.kv_cache_manager.tokens_per_block, + is_decode=not context_rows, + ): from .prims_ts_indexer import prims_ts_qsa_index_scores logits = prims_ts_qsa_index_scores( 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 index 8c6000daff2b..ec5a40b6dded 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py @@ -16,10 +16,14 @@ # A plan miss costs ~40 ms per layer; padding prefill rows to a bucket bounds # the plan keys to max_num_tokens / bucket per layer. _PREFILL_ROW_BUCKET = 256 -# Decode rows below which Triton is faster under CUDA graphs (Rubin SM107, -# 32K context, TP2): the core ties at one row, the indexer until ~32 rows. +# Decode plans are keyed on the row count too; eager decode batches vary. +_DECODE_ROW_BUCKET = 8 +# 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) @@ -29,7 +33,7 @@ @functools.cache def _prims_ts_api(): try: - from flashinfer.attention.prims_ts import ( + from ...prims_ts import ( QTokenKvBlockSparsePagedTSWrapper, get_q_token_kv_block_sparse_workspace_size, make_q_token_kv_block_sparse_qo_indptr, @@ -57,7 +61,7 @@ def _unsupported_reason( 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 "flashinfer.attention.prims_ts is not importable" + 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: @@ -69,8 +73,21 @@ def _unsupported_reason( 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, *, is_decode: bool = False + 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. @@ -92,6 +109,13 @@ def use_prims_ts_indexer( 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": @@ -236,33 +260,30 @@ def run( block_topk = block_ids.shape[-1] out_dtype = q.dtype - sm_scale = softmax_scale 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 quantized per - # call with a dynamic per-tensor scale folded into sm_scale. This - # differs numerically from the Triton path, which upcasts K/V, and - # the scale readback synchronizes the host. - q_scale = (q.float().abs().amax() / _FP8_MAX).clamp_min(1e-12) - q = (q.float() / q_scale).to(torch.float8_e4m3fn) - sm_scale = softmax_scale * float(q_scale) + # 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() - num_rows = rows - if not is_decode: - rows = -(-rows // _PREFILL_ROW_BUCKET) * _PREFILL_ROW_BUCKET - if rows != num_rows: - # Padding rows attend to nothing but token 0 of request 0; the - # suffix gets its own routes and is sliced off the output. - q = _pad_rows(q, rows, 0) - block_ids = _pad_rows(block_ids, rows, -1) - request_indices = _pad_rows(request_indices, rows, 0) - query_positions = _pad_rows(query_positions, rows, 0) + bucket = _DECODE_ROW_BUCKET if is_decode else _PREFILL_ROW_BUCKET + padded_rows = -(-rows // bucket) * bucket + if padded_rows != rows: + # Padding rows attend to nothing but token 0 of request 0 and are + # sliced off the output; in packed prefill they form inert routes. + 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) + route_q = q.view(padded_rows, 1, 1, num_qo_heads, head_dim) qo_indptr = None group_size = 1 - num_routes = rows + num_routes = padded_rows split_kv = True else: route_q = q @@ -271,7 +292,9 @@ def run( # 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, rows, q.device) + 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 @@ -282,7 +305,7 @@ def run( q.device.index, k_cache.data_ptr(), v_cache.data_ptr(), - rows, + padded_rows, num_routes, group_size, is_decode, @@ -298,8 +321,10 @@ def run( block_topk, ) result = out - if out is None or rows != num_rows: - out = torch.empty((rows, num_qo_heads, head_dim), dtype=out_dtype, device=q.device) + 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: @@ -313,24 +338,25 @@ def run( max_seq_len_kv=max_seq_len_kv, o_data_type=out_dtype, qo_indptr=qo_indptr, - group_size=group_size, - kv_block_size=4, + seq_len_q=None if is_decode else group_size, + kv_block_size=_KV_BLOCK_SIZE, split_kv=split_kv, ) workspace = self._workspace_for(nbytes, q.device) wrapper = wrapper_cls() + # The vendored API names the route width seq_len_q. wrapper.plan( num_routes, + group_size, num_qo_heads, num_kv_heads, head_dim, - 4, + _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, @@ -356,7 +382,7 @@ def run( request_indices, query_positions, qo_indptr=qo_indptr, - sm_scale=sm_scale, + sm_scale=softmax_scale, out=run_out, ) except RuntimeError as error: @@ -370,9 +396,9 @@ def run( if key not in self._last_abi: self._last_abi[key] = _run_abi(run_args) if result is None: - return out[:num_rows] + return out[:rows] if result is not out: - result.copy_(out[:num_rows]) + result.copy_(out[:rows]) return result 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 index 1c0a4892ee68..9521385b632f 100644 --- 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 @@ -1,16 +1,5 @@ -# Copyright (c) 2026 by FlashInfer team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. +# 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 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 index c114e42300f3..33cbd4b3d33f 100644 --- 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 @@ -1,16 +1,5 @@ -# Copyright (c) 2026 by FlashInfer team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. +# 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.""" 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 index 694ab28c2342..7a15367b0bec 100644 --- 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 @@ -1,16 +1,5 @@ -# Copyright (c) 2026 by FlashInfer team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. +# 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. diff --git a/tests/microbenchmarks/qsa/bench_qsa_sparse_core.py b/tests/microbenchmarks/qsa/bench_qsa_sparse_core.py index 674d2db4fe6f..8869d4b914a0 100644 --- a/tests/microbenchmarks/qsa/bench_qsa_sparse_core.py +++ b/tests/microbenchmarks/qsa/bench_qsa_sparse_core.py @@ -16,6 +16,7 @@ import json import math import statistics +import sys import time from types import SimpleNamespace @@ -169,7 +170,7 @@ def run(): def prims_runner(case, args, device): - from flashinfer.attention.prims_ts import ( + from tensorrt_llm._torch.attention.backends.prims_ts import ( QTokenKvBlockSparsePagedTSWrapper, get_q_token_kv_block_sparse_workspace_size, make_q_token_kv_block_sparse_qo_indptr, @@ -206,7 +207,7 @@ def prims_runner(case, args, device): max_seq_len_kv=args.ctx, o_data_type=out.dtype, qo_indptr=qo_indptr, - group_size=group, + seq_len_q=group if packed else None, kv_block_size=KV_BLOCK, split_kv=split_kv, ) @@ -214,6 +215,7 @@ def prims_runner(case, args, device): wrapper = QTokenKvBlockSparsePagedTSWrapper() wrapper.plan( routes, + group, case.hq, case.hkv, HEAD_DIM, @@ -223,7 +225,6 @@ def prims_runner(case, args, device): args.ctx, device=device, workspace_buffer=workspace, - group_size=group, use_packed_q=packed, split_kv=split_kv, q_data_type=q.dtype, @@ -436,7 +437,11 @@ def main(): with open(args.json, "a") as f: for r in rows: f.write(json.dumps(r) + "\n") + bad = [r["label"] for r in rows if "error" in r or not (r["triton_ok"] and r["prims_ok"])] + if bad: + print("FAILED cases:", ", ".join(bad)) + return 1 if bad else 0 if __name__ == "__main__": - main() + sys.exit(main()) 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 index 5da8352b4394..2479cd339df7 100644 --- 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 @@ -53,11 +53,12 @@ def _requires_prims_ts() -> None: 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") - if not prims_ts_available(): - pytest.skip("flashinfer.attention.prims_ts is not installed") + # 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): +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) @@ -65,8 +66,10 @@ def _make_case(query_lens: list[int], seed: int): 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 - ) - v_cache = torch.randn_like(k_cache) + ).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 ) @@ -316,10 +319,19 @@ def indexer(backend, rows, is_decode): 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="triton"), - torch.empty(64, 4, 128, dtype=torch.bfloat16, device=device), + _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_align_prefill_chunk() -> None: @@ -357,3 +369,116 @@ def test_prefill_writes_into_caller_output() -> None: 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_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=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 + 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 index 4d6e9fecf98d..5da690f9ece8 100644 --- 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 @@ -49,8 +49,9 @@ def _case(request_lengths, ctx, tokens_per_block, seed, *, q_scale=1.0, position def _requires_prims_ts(): if not torch.cuda.is_available(): pytest.skip("CUDA required") - if torch.cuda.get_device_capability()[0] != 10: - pytest.skip("Blackwell/Rubin tcgen05 required") + if torch.cuda.get_device_capability() not in ((10, 0), (10, 3), (10, 7)): + pytest.skip("PrimTS QSA scorer requires SM100/SM103/SM107") + pytest.importorskip("cutlass.experimental.task_scheduling") def _scorer(): From 70ae6fc207c070b4dfa9751b23f65fa45a30007e Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Fri, 23 Oct 2026 15:38:02 +0000 Subject: [PATCH 10/18] [None][chore] Drop the QSA sparse-core microbenchmark from the tree Local A/B tooling; the measured tables live in the PR description. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../qsa/bench_qsa_sparse_core.py | 447 ------------------ 1 file changed, 447 deletions(-) delete mode 100644 tests/microbenchmarks/qsa/bench_qsa_sparse_core.py diff --git a/tests/microbenchmarks/qsa/bench_qsa_sparse_core.py b/tests/microbenchmarks/qsa/bench_qsa_sparse_core.py deleted file mode 100644 index 8869d4b914a0..000000000000 --- a/tests/microbenchmarks/qsa/bench_qsa_sparse_core.py +++ /dev/null @@ -1,447 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""A/B latency of the QSA sparse core: TRT-LLM Triton vs FlashInfer PrimTS. - -Qwen3.8-Flash-Next geometry (D256, 4-token KV blocks, top-k 512 blocks) over an -HND paged cache. Triton time includes the block-to-token expansion the Triton -kernel needs; PrimTS consumes block IDs directly. Both outputs are checked -against the Torch reference on sampled rows before timing. - - python bench_qsa_sparse_core.py --mode prefill --ctx 32768 --chunk 4096 - python bench_qsa_sparse_core.py --mode decode --ctx 32768 --batch 64 --sq 4 - python bench_qsa_sparse_core.py --preset quick --out results.md -""" - -import argparse -import json -import math -import statistics -import sys -import time -from types import SimpleNamespace - -import torch - -HEAD_DIM = 256 -KV_BLOCK = 4 -TOKEN_TOPK = 2048 -BLOCK_TOPK = TOKEN_TOPK // KV_BLOCK -PREFILL_GROUP = 4 -FP8_MAX = 448.0 - -QUICK_PRESET = [ - dict(mode="prefill", ctx=32768, chunk=4096, batch=1, sq=1), - dict(mode="decode", ctx=32768, chunk=0, batch=8, sq=1), - dict(mode="decode", ctx=32768, chunk=0, batch=64, sq=1), - dict(mode="decode", ctx=32768, chunk=0, batch=64, sq=4), -] - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__) - p.add_argument("--mode", choices=("prefill", "decode"), default="prefill") - p.add_argument("--tp", type=int, choices=(1, 2), default=2, help="1: Hq/Hkv 24/2, 2: 12/1") - p.add_argument("--ctx", type=int, default=32768) - p.add_argument("--chunk", type=int, default=4096, help="prefill query rows per request") - p.add_argument("--batch", type=int, default=1) - p.add_argument("--sq", type=int, default=1, help="decode query tokens per request (1 or MTP+1)") - p.add_argument("--kv-dtype", choices=("bf16", "fp8"), default="bf16") - p.add_argument("--page", type=int, default=32) - p.add_argument("--pattern", choices=("random", "sorted"), default="random") - p.add_argument("--warmup", type=int, default=20) - p.add_argument("--iters", type=int, default=100) - p.add_argument("--l2-mb", type=int, default=256) - p.add_argument("--check-rows", type=int, default=8) - p.add_argument("--no-graph", action="store_true") - p.add_argument("--skip-triton", action="store_true") - p.add_argument("--skip-prims", action="store_true") - p.add_argument("--preset", choices=("quick",), default=None) - p.add_argument("--seed", type=int, default=1729) - p.add_argument("--out", default=None, help="markdown file to append rows to") - p.add_argument("--json", default=None, help="json file to append rows to") - return p.parse_args() - - -def heads_for_tp(tp): - return (24, 2) if tp == 1 else (12, 1) - - -def make_cache(num_requests, ctx, hkv, page, kv_dtype, device, gen): - pages_per_req = math.ceil(ctx / page) - num_pages = num_requests * pages_per_req - k = ( - torch.randn( - num_pages, hkv, page, HEAD_DIM, dtype=torch.bfloat16, device=device, generator=gen - ) - * 0.25 - ) - v = ( - torch.randn( - num_pages, hkv, page, HEAD_DIM, dtype=torch.bfloat16, device=device, generator=gen - ) - * 0.25 - ) - perm = torch.randperm(num_pages, device=device, generator=gen).to(torch.int32) - block_table = perm.view(num_requests, pages_per_req).contiguous() - return k.to(kv_dtype), v.to(kv_dtype), block_table - - -def make_block_ids(positions, pattern, device, gen): - rows = positions.numel() - visible = (positions + 1) // KV_BLOCK - max_vis = int(visible.max()) - keys = torch.rand(rows, max_vis, device=device, generator=gen) - col = torch.arange(max_vis, device=device) - keys.masked_fill_(col[None, :] >= visible[:, None], 2.0) - order = keys.argsort(dim=1)[:, :BLOCK_TOPK] - count = visible.clamp(max=BLOCK_TOPK) - valid = torch.arange(BLOCK_TOPK, device=device)[None, :] < count[:, None] - ids = order.to(torch.int32) - if pattern == "sorted": - ids = ( - torch.where(valid, ids, torch.full_like(ids, torch.iinfo(torch.int32).max)) - .sort(dim=1) - .values - ) - return torch.where(valid, ids, torch.full_like(ids, -1)).contiguous() - - -def build_case(args, device, gen): - hq, hkv = heads_for_tp(args.tp) - kv_dtype = torch.bfloat16 if args.kv_dtype == "bf16" else torch.float8_e4m3fn - if args.mode == "prefill": - rows_per_req = args.chunk - else: - rows_per_req = args.sq - B = args.batch - rows = B * rows_per_req - q = torch.randn(rows, hq, HEAD_DIM, dtype=torch.bfloat16, device=device, generator=gen) * 0.25 - k, v, block_table = make_cache(B, args.ctx, hkv, args.page, kv_dtype, device, gen) - req = torch.arange(B, dtype=torch.int32, device=device).repeat_interleave(rows_per_req) - pos = torch.arange(args.ctx - rows_per_req, args.ctx, dtype=torch.int64, device=device).repeat( - B - ) - seq_lens = torch.full((rows,), args.ctx, dtype=torch.int32, device=device) - ids = make_block_ids(pos, args.pattern, device, gen) - scale = HEAD_DIM**-0.5 - return SimpleNamespace( - q=q, - k=k, - v=v, - block_table=block_table, - req=req, - pos=pos, - seq_lens=seq_lens, - ids=ids, - scale=scale, - hq=hq, - hkv=hkv, - rows=rows, - rows_per_req=rows_per_req, - B=B, - kv_dtype=kv_dtype, - ) - - -def triton_runner(case, args): - from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( - triton_expand_qsa_block_indices, - triton_qsa_paged_sparse_gqa, - ) - - def run(): - tokens = triton_expand_qsa_block_indices( - case.ids, case.pos, case.seq_lens, compress_ratio=KV_BLOCK, token_topk=TOKEN_TOPK - ) - return triton_qsa_paged_sparse_gqa( - q=case.q, - k_cache=case.k, - v_cache=case.v, - block_table=case.block_table, - selected_tokens=tokens, - request_indices=case.req, - tokens_per_block=args.page, - softmax_scale=case.scale, - query_positions=case.pos, - compress_ratio=KV_BLOCK, - ) - - return run - - -def prims_runner(case, args, device): - from tensorrt_llm._torch.attention.backends.prims_ts import ( - QTokenKvBlockSparsePagedTSWrapper, - get_q_token_kv_block_sparse_workspace_size, - make_q_token_kv_block_sparse_qo_indptr, - ) - - q = case.q - sm_scale = case.scale - if case.kv_dtype == torch.float8_e4m3fn: - q_scale = (q.float().abs().amax() / FP8_MAX).clamp_min(1e-8).item() - q = (q.float() / q_scale).to(torch.float8_e4m3fn) - sm_scale = case.scale * q_scale - out = torch.empty(case.rows, case.hq, HEAD_DIM, dtype=torch.bfloat16, device=device) - packed = args.mode == "prefill" - if packed: - group = PREFILL_GROUP - starts = torch.arange(0, case.rows + 1, case.rows_per_req, dtype=torch.int32) - qo_indptr = make_q_token_kv_block_sparse_qo_indptr( - starts, case.rows, group_size=group, device=device - ) - routes = qo_indptr.numel() - 1 - q_view, out_view, split_kv = q, out, False - else: - group = case.rows_per_req - qo_indptr = None - routes = case.B - q_view = q.view(case.B, 1, group, case.hq, HEAD_DIM) - out_view = out.view_as(q_view) - split_kv = True - ws_bytes = get_q_token_kv_block_sparse_workspace_size( - q_view, - case.k, - case.block_table, - block_topk=BLOCK_TOPK, - max_seq_len_kv=args.ctx, - o_data_type=out.dtype, - qo_indptr=qo_indptr, - seq_len_q=group if packed else None, - kv_block_size=KV_BLOCK, - split_kv=split_kv, - ) - workspace = torch.zeros(ws_bytes, dtype=torch.uint8, device=device) - wrapper = QTokenKvBlockSparsePagedTSWrapper() - wrapper.plan( - routes, - group, - case.hq, - case.hkv, - HEAD_DIM, - KV_BLOCK, - args.page, - BLOCK_TOPK, - args.ctx, - device=device, - workspace_buffer=workspace, - use_packed_q=packed, - split_kv=split_kv, - q_data_type=q.dtype, - kv_data_type=case.k.dtype, - o_data_type=out.dtype, - ) - - def run(): - wrapper.run( - q_view, - (case.k, case.v), - case.block_table, - case.ids, - case.req, - case.pos, - qo_indptr=qo_indptr, - sm_scale=sm_scale, - out=out_view, - ) - return out - - return run, dict(group_size=group, routes=routes, workspace_mb=ws_bytes / 2**20) - - -def reference(case, args, rows_idx): - from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( - triton_expand_qsa_block_indices, - ) - from tensorrt_llm._torch.attention.backends.sparse.qsa.module import qsa_sparse_gqa_reference - - tokens = triton_expand_qsa_block_indices( - case.ids[rows_idx], - case.pos[rows_idx], - case.seq_lens[rows_idx], - compress_ratio=KV_BLOCK, - token_topk=TOKEN_TOPK, - ) - meta = SimpleNamespace( - qsa_block_table=case.block_table, - kv_cache_manager=SimpleNamespace(tokens_per_block=args.page), - ) - return qsa_sparse_gqa_reference( - q=case.q[rows_idx], - k_cache=case.k.to(torch.bfloat16), - v_cache=case.v.to(torch.bfloat16), - selected_tokens=tokens, - request_indices=case.req[rows_idx], - metadata=meta, - softmax_scale=case.scale, - ) - - -def check(name, out, ref, rows_idx, tol): - got = out[rows_idx].float() - err = (got - ref.float()).abs().max().item() - ok = torch.allclose(got, ref.float(), atol=tol, rtol=tol) - print(f" [{name}] max_abs_err={err:.4e} {'OK' if ok else 'MISMATCH'} (tol {tol})") - return ok, err - - -def time_fn(fn, args, device): - for _ in range(3): - fn() - torch.cuda.synchronize() - if args.no_graph: - launch = fn - else: - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - fn() - launch = graph.replay - scrub = torch.empty(args.l2_mb << 20, dtype=torch.uint8, device=device) - samples = [] - for i in range(args.warmup + args.iters): - scrub.random_() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - launch() - end.record() - end.synchronize() - if i >= args.warmup: - samples.append(start.elapsed_time(end) * 1000.0) - return samples - - -def summarize(samples): - s = sorted(samples) - return dict(median_us=statistics.median(s), p95_us=s[int(0.95 * (len(s) - 1))], n=len(s)) - - -def run_case(args, device): - gen = torch.Generator(device=device).manual_seed(args.seed) - case = build_case(args, device, gen) - hq, hkv = case.hq, case.hkv - label = ( - f"{args.mode} tp{args.tp} Hq/Hkv={hq}/{hkv} ctx={args.ctx} " - + (f"chunk={args.chunk} " if args.mode == "prefill" else f"bs={args.batch} sq={args.sq} ") - + f"kv={args.kv_dtype} page={args.page} pattern={args.pattern}" - ) - print(f"== {label}") - rows_idx = torch.linspace( - 0, case.rows - 1, steps=min(args.check_rows, case.rows), device=device - ).long() - ref = reference(case, args, rows_idx) - tol = 2e-2 if args.kv_dtype == "bf16" else 5e-2 - row = dict( - label=label, - mode=args.mode, - tp=args.tp, - ctx=args.ctx, - chunk=args.chunk, - batch=args.batch, - sq=args.sq, - kv_dtype=args.kv_dtype, - page=args.page, - pattern=args.pattern, - rows=case.rows, - ) - runners = [] - if not args.skip_triton: - t0 = time.time() - fn = triton_runner(case, args) - ok, err = check("triton", fn(), ref, rows_idx, tol) - row.update(triton_ok=ok, triton_err=err) - runners.append(("triton", fn)) - print(f" triton ready in {time.time() - t0:.1f}s") - if not args.skip_prims: - t0 = time.time() - fn, info = prims_runner(case, args, device) - ok, err = check("prims_ts", fn(), ref, rows_idx, tol) - row.update(prims_ok=ok, prims_err=err, **info) - runners.append(("prims_ts", fn)) - print( - f" prims_ts ready in {time.time() - t0:.1f}s (G={info['group_size']}, routes={info['routes']})" - ) - samples = {name: [] for name, _ in runners} - for order in (runners, list(reversed(runners))): - for name, fn in order: - samples[name] += time_fn(fn, args, device) - for name in samples: - row[name] = summarize(samples[name]) - print( - f" {name:8s} median {row[name]['median_us']:9.1f} us p95 {row[name]['p95_us']:9.1f} us" - ) - if "triton" in row and "prims_ts" in row: - row["speedup"] = row["triton"]["median_us"] / row["prims_ts"]["median_us"] - print(f" speedup (triton/prims_ts) = {row['speedup']:.3f}x") - return row - - -def md_row(r): - t = f"{r['triton']['median_us']:.1f}" if "triton" in r else "-" - p = f"{r['prims_ts']['median_us']:.1f}" if "prims_ts" in r else "-" - s = f"{r['speedup']:.2f}x" if "speedup" in r else "-" - okc = "".join( - f"{n}:{'ok' if r.get(k) else 'BAD'} " - for n, k in (("T", "triton_ok"), ("P", "prims_ok")) - if k in r - ) - return f"| {r['label']} | {r['rows']} | {t} | {p} | {s} | {okc.strip()} |" - - -def main(): - args = parse_args() - device = torch.device("cuda", torch.cuda.current_device()) - print("device:", torch.cuda.get_device_name(device), torch.cuda.get_device_capability(device)) - cases = [] - if args.preset == "quick": - for over in QUICK_PRESET: - a = argparse.Namespace(**vars(args)) - for k, v in over.items(): - setattr(a, k, v) - cases.append(a) - else: - cases.append(args) - rows = [] - for a in cases: - try: - rows.append(run_case(a, device)) - except Exception as exc: # keep going so one failing config does not hide the rest - print(f" FAILED: {type(exc).__name__}: {exc}") - rows.append( - dict( - label=f"{a.mode} ctx={a.ctx} bs={a.batch} sq={a.sq} chunk={a.chunk}", - rows=0, - error=f"{type(exc).__name__}: {exc}", - ) - ) - torch.cuda.empty_cache() - header = "| case | rows | Triton us | PrimTS us | Triton/PrimTS | check |\n|---|---:|---:|---:|---:|---|" - print(header) - for r in rows: - print( - md_row(r) - if "error" not in r - else f"| {r['label']} | - | error: {r['error'][:80]} | | | |" - ) - if args.out: - with open(args.out, "a") as f: - f.write( - f"\n### {torch.cuda.get_device_name(device)} {time.strftime('%Y-%m-%d %H:%M')}\n{header}\n" - ) - for r in rows: - f.write( - (md_row(r) if "error" not in r else f"| {r['label']} | - | error | | | |") - + "\n" - ) - if args.json: - with open(args.json, "a") as f: - for r in rows: - f.write(json.dumps(r) + "\n") - bad = [r["label"] for r in rows if "error" in r or not (r["triton_ok"] and r["prims_ok"])] - if bad: - print("FAILED cases:", ", ".join(bad)) - return 1 if bad else 0 - - -if __name__ == "__main__": - sys.exit(main()) From 371816f7729de417ebcfa0a978ce18a8e139160a Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Fri, 23 Oct 2026 16:56:33 +0000 Subject: [PATCH 11/18] [None][doc] Note the FP8 KV cache numerics of the QSA PrimTS kernel in the Qwen3.8 guide Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../deployment-guide-for-qwen3.8-flash-next-on-trtllm.md | 2 ++ 1 file changed, 2 insertions(+) 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 22d0d2b6da2d..a315f6bb8473 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, the QSA sparse attention on SM100/SM103/SM107 runs the PrimTS kernel with FP8 queries (unit scale, clamped to the E4M3 range) against the FP8 cache; the Triton path 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. From e78d2dc63c83f77e4d852d1ed581b640a92731bc Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Fri, 23 Oct 2026 17:17:34 +0000 Subject: [PATCH 12/18] [None][doc] Qualify when QSA selects PrimTS with an FP8 KV cache Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../deployment-guide-for-qwen3.8-flash-next-on-trtllm.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 a315f6bb8473..c1caf8fea3fa 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,7 +25,7 @@ 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, the QSA sparse attention on SM100/SM103/SM107 runs the PrimTS kernel with FP8 queries (unit scale, clamped to the E4M3 range) against the FP8 cache; the Triton path 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. +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 From 645b4501f69a266161e78bb0f4c2b25fc017890f Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Sat, 31 Oct 2026 20:20:23 +0000 Subject: [PATCH 13/18] [None][chore] Use the group_size plan API of the vendored PrimTS QToken wrapper #20911 moved the vendored prims_ts tree to aa709061, where FlashInfer #5351 renamed the QToken wrapper's route width from seq_len_q to group_size and accepts it for fixed decode layouts as well. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../_torch/attention/backends/sparse/qsa/prims_ts_core.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) 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 index ec5a40b6dded..96c7a87c824f 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py @@ -338,16 +338,14 @@ def run( max_seq_len_kv=max_seq_len_kv, o_data_type=out_dtype, qo_indptr=qo_indptr, - seq_len_q=None if is_decode else group_size, + group_size=group_size, kv_block_size=_KV_BLOCK_SIZE, split_kv=split_kv, ) workspace = self._workspace_for(nbytes, q.device) wrapper = wrapper_cls() - # The vendored API names the route width seq_len_q. wrapper.plan( num_routes, - group_size, num_qo_heads, num_kv_heads, head_dim, @@ -357,6 +355,7 @@ def run( 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, From 95f6b2bacdac5e36228e4d73ab908fa1a617e2b5 Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Tue, 3 Nov 2026 17:00:04 +0000 Subject: [PATCH 14/18] [None][refactor] Drop row padding from the QSA PrimTS core Prefill routes are padded with trailing empty routes to a 256-route capacity instead of padding q, block ids, request indices and positions to a row bucket and copying the output back; the wrapper runs fewer rows than planned when the route count matches (requires routes <= rows, so chunks under 256 rows plan per row count). Decode plans are per row count: with the vendored occupancy cache a plan miss is ~1-3 ms per layer, not ~40 ms, so the 8-row bucket no longer pays for itself. B300 e2e (TP2, 32K/256, concurrency 16) is unchanged: 1.262 -> 1.260 req/s, TTFT 2613 -> 2595 ms, TPOT 39.28 -> 39.44 ms. Tests: the no-replanning test patches the vendored prepare function (the wheel's was never called) and adds a re-prepare negative control; the indexer tests fail instead of skipping when the task-scheduling DSL is missing on a supported GPU. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../attention/backends/sparse/qsa/module.py | 3 +- .../backends/sparse/qsa/prims_ts_core.py | 54 +++++-------------- .../sparse/qsa/test_qsa_prims_ts_core.py | 29 +++++----- .../sparse/qsa/test_qsa_prims_ts_indexer.py | 4 +- 4 files changed, 30 insertions(+), 60 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py index bf28b600726f..20e613d96e2d 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py @@ -773,9 +773,8 @@ def forward( ) chunk_size = _query_chunk_size(num_tokens, score_columns) if use_prims_ts: - from .prims_ts_core import align_prefill_chunk, slice_request_boundaries + from .prims_ts_core import slice_request_boundaries - chunk_size = align_prefill_chunk(chunk_size) cu_seq_lens = list( itertools.accumulate(attn_metadata.seq_lens[: attn_metadata.num_seqs].tolist()) ) 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 index 96c7a87c824f..a777d6dd2eef 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py @@ -13,11 +13,12 @@ from .params import QSASparseParams _PREFILL_GROUP_SIZE = 4 -# A plan miss costs ~40 ms per layer; padding prefill rows to a bucket bounds -# the plan keys to max_num_tokens / bucket per layer. -_PREFILL_ROW_BUCKET = 256 -# Decode plans are keyed on the row count too; eager decode batches vary. -_DECODE_ROW_BUCKET = 8 +# Prefill plans are capacity plans over a route count: the route offsets are +# padded with trailing empty routes to a multiple of this, so a chunk's row +# count does not create a plan (the wrapper requires routes <= rows, so +# chunks shorter than the bucket plan per row count). Decode plans are per +# row count. +_PREFILL_ROUTE_BUCKET = 256 # 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. @@ -163,17 +164,6 @@ def _request_boundaries(request_indices: torch.Tensor) -> tuple[int, ...]: 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 align_prefill_chunk(rows: int) -> int: - """Largest bucket multiple not above ``rows``, so full chunks need no padding.""" - return max(rows // _PREFILL_ROW_BUCKET * _PREFILL_ROW_BUCKET, min(rows, _PREFILL_ROW_BUCKET)) - - 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[ @@ -225,6 +215,8 @@ def _route_offsets( 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: @@ -270,20 +262,11 @@ def run( request_indices = request_indices.to(torch.int32).contiguous() q = q.contiguous() - bucket = _DECODE_ROW_BUCKET if is_decode else _PREFILL_ROW_BUCKET - padded_rows = -(-rows // bucket) * bucket - if padded_rows != rows: - # Padding rows attend to nothing but token 0 of request 0 and are - # sliced off the output; in packed prefill they form inert routes. - 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(padded_rows, 1, 1, num_qo_heads, head_dim) + route_q = q.view(rows, 1, 1, num_qo_heads, head_dim) qo_indptr = None group_size = 1 - num_routes = padded_rows + num_routes = rows split_kv = True else: route_q = q @@ -293,7 +276,7 @@ def run( request_boundaries = _request_boundaries(request_indices) with self._lock: qo_indptr = self._route_offsets( - make_qo_indptr, request_boundaries, padded_rows, q.device + make_qo_indptr, request_boundaries, rows, q.device ) group_size = _PREFILL_GROUP_SIZE num_routes = qo_indptr.numel() - 1 @@ -305,7 +288,6 @@ def run( q.device.index, k_cache.data_ptr(), v_cache.data_ptr(), - padded_rows, num_routes, group_size, is_decode, @@ -320,11 +302,8 @@ def run( 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 - ) + if out is None: + out = torch.empty((rows, num_qo_heads, head_dim), dtype=out_dtype, device=q.device) with self._lock: wrapper = self._plans.get(key) if wrapper is None: @@ -394,11 +373,7 @@ def run( 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 + return out _CORE = _PrimsTSSparseCore() @@ -448,7 +423,6 @@ def prims_ts_sparse_gqa( __all__ = [ "prims_ts_available", "prims_ts_sparse_gqa", - "align_prefill_chunk", "resolve_sparse_core_backend", "slice_request_boundaries", ] 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 index 2479cd339df7..c3cf859c5303 100644 --- 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 @@ -10,8 +10,7 @@ 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 ( - align_prefill_chunk, - prims_ts_available, + prims_ts_available, prims_ts_sparse_gqa, slice_request_boundaries, ) @@ -206,20 +205,20 @@ def test_prefill_rows_share_plans_within_bucket() -> None: 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( - [37, 91], seed=5 + [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) q2, _, _, block_table2, block_ids2, request_indices2, positions2, selected2, metadata = ( - _make_case([100, 30], seed=6) + _make_case([350, 210], 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] == 130 + assert out.shape[0] == 560 expected = qsa_sparse_gqa_reference( q=q2, k_cache=k_cache, @@ -235,7 +234,7 @@ def test_prefill_rows_share_plans_within_bucket() -> None: def test_route_offsets_follow_layout_changes_without_replanning() -> None: _requires_prims_ts() - import flashinfer.attention.prims_ts.q_token_kv_block_sparse_metadata as fi_metadata + 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 @@ -277,10 +276,14 @@ def counted(*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 == 0 + assert prepares == 1 for out, ref in zip(outputs, expected): assert torch.equal(out, ref) @@ -334,14 +337,6 @@ def indexer(backend, rows, is_decode): assert use_prims_ts_indexer(_params(), q64, tokens_per_block=64) -def test_align_prefill_chunk() -> None: - assert align_prefill_chunk(4048) == 3840 - assert align_prefill_chunk(8192) == 8192 - assert align_prefill_chunk(256) == 256 - assert align_prefill_chunk(255) == 255 - assert align_prefill_chunk(96) == 96 - - def test_prefill_writes_into_caller_output() -> None: _requires_prims_ts() kwargs = dict( @@ -445,7 +440,7 @@ def test_decode_graph_capture_replay(kv_dtype) -> None: assert torch.equal(out, eager) -def test_decode_rows_share_plans_within_bucket() -> None: +def test_decode_rows_get_exact_plans() -> None: _requires_prims_ts() from tensorrt_llm._torch.attention.backends.sparse.qsa.prims_ts_core import _CORE @@ -469,7 +464,7 @@ def test_decode_rows_share_plans_within_bucket() -> None: out = prims_ts_sparse_gqa( q, k_cache, v_cache, block_table, block_ids, request_indices, positions, **kwargs ) - assert len(_CORE._plans) == plans + assert len(_CORE._plans) == plans + 1 assert out.shape[0] == 5 expected = qsa_sparse_gqa_reference( q=q, 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 index 5da690f9ece8..433c8a851eac 100644 --- 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 @@ -51,7 +51,9 @@ def _requires_prims_ts(): 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") - pytest.importorskip("cutlass.experimental.task_scheduling") + # 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 9ff5790d0af44d03b62bacacc46f50540ae3d18a Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Tue, 3 Nov 2026 17:49:29 +0000 Subject: [PATCH 15/18] [None][doc] Release note for QSA PrimTS kernels on by default Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- docs/source/release-notes.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/source/release-notes.md b/docs/source/release-notes.md index 26022b9cf621..621565770dde 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 From f33a4f7700c82bcbb1ab6038d48373a6c32c284a Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Tue, 3 Nov 2026 17:50:03 +0000 Subject: [PATCH 16/18] [None][doc] Split the QSA release note sentence Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- docs/source/release-notes.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/release-notes.md b/docs/source/release-notes.md index 621565770dde..543f98573a61 100644 --- a/docs/source/release-notes.md +++ b/docs/source/release-notes.md @@ -12,7 +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). +- **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 From 1b157347fcb856623f6fd60fbde87345934bb2f3 Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Tue, 3 Nov 2026 21:22:20 +0000 Subject: [PATCH 17/18] [None][perf] Pad QSA PrimTS prefill tails to 64 rows and retire grown workspaces The vendored wrapper re-prepares its launch, with a host readback of the route offsets, whenever the q/out shape changes (prepared_key includes the tensor shapes), so with varying chunk row counts every prefill forward re-prepared every QSA layer, ~0.6 ms plus a stream drain each. Prefill tails are padded to 64-row buckets again and plans are keyed on the padded shape as well as the route capacity, so a repeated shape keeps its prepared plan; full chunks and decode rows are not padded. B300 TP2 e2e with varied ISL (32K +- 3K, OSL 256, concurrency 16): 1.340 -> 1.450 req/s, TTFT 2550 -> 2061 ms, TPOT 36.6 -> 35.0 ms. Workspaces superseded by growth stay referenced for the process lifetime: decode graphs captured before the growth still address the old buffer. Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../attention/backends/sparse/qsa/module.py | 5 +- .../backends/sparse/qsa/prims_ts_core.py | 53 +++++++++++++++---- .../sparse/qsa/test_qsa_prims_ts_core.py | 5 +- 3 files changed, 49 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py index 20e613d96e2d..fc050b98e520 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py @@ -805,7 +805,10 @@ def select(start: int, end: int) -> torch.Tensor: 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. Chunks write disjoint Top-K rows. + # 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 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 index a777d6dd2eef..c2c33e919f59 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/prims_ts_core.py @@ -14,11 +14,16 @@ _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 a chunk's row -# count does not create a plan (the wrapper requires routes <= rows, so -# chunks shorter than the bucket plan per row count). Decode plans are per -# row count. +# 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. @@ -164,6 +169,12 @@ def _request_boundaries(request_indices: torch.Tensor) -> tuple[int, ...]: 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[ @@ -181,6 +192,8 @@ 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 @@ -191,6 +204,8 @@ def _workspace_for(self, nbytes: int, device: torch.device) -> torch.Tensor: 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 @@ -203,9 +218,10 @@ def _route_offsets( device: torch.device, ) -> torch.Tensor: # The FlashInfer wrapper re-prepares, with host readbacks, whenever the - # qo_indptr storage changes. All QSA layers of a forward share one - # request layout, so offsets live in a persistent buffer per route count - # that is rewritten, stream-ordered, only when the layout changes. + # 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] @@ -262,6 +278,13 @@ def run( 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 @@ -276,7 +299,7 @@ def run( request_boundaries = _request_boundaries(request_indices) with self._lock: qo_indptr = self._route_offsets( - make_qo_indptr, request_boundaries, rows, q.device + make_qo_indptr, request_boundaries, padded_rows, q.device ) group_size = _PREFILL_GROUP_SIZE num_routes = qo_indptr.numel() - 1 @@ -288,6 +311,7 @@ def run( q.device.index, k_cache.data_ptr(), v_cache.data_ptr(), + padded_rows, num_routes, group_size, is_decode, @@ -302,8 +326,11 @@ def run( page_size, block_topk, ) - if out is None: - out = torch.empty((rows, num_qo_heads, head_dim), dtype=out_dtype, device=q.device) + 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: @@ -373,7 +400,11 @@ def run( 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) - return out + if result is None: + return out[:rows] + if result is not out: + result.copy_(out[:rows]) + return result _CORE = _PrimsTSSparseCore() 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 index c3cf859c5303..b210082e7bda 100644 --- 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 @@ -211,14 +211,15 @@ def test_prefill_rows_share_plans_within_bucket() -> None: 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([350, 210], seed=6) + _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] == 560 + assert out.shape[0] == 510 expected = qsa_sparse_gqa_reference( q=q2, k_cache=k_cache, From 7480919cffcbff7b252200f20524b5e850faff27 Mon Sep 17 00:00:00 2001 From: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> Date: Wed, 4 Nov 2026 03:11:31 +0000 Subject: [PATCH 18/18] [None][chore] Apply ruff-format to the QSA core test imports Signed-off-by: Faraz Khoubsirat <58580514+farazkh80@users.noreply.github.com> --- .../_torch/attention/sparse/qsa/test_qsa_prims_ts_core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index b210082e7bda..5340abc0fc65 100644 --- 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 @@ -10,7 +10,7 @@ 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_available, prims_ts_sparse_gqa, slice_request_boundaries, )