[TRTLLM-17238][feat] Add PrimTS sparse core backend for Qwen3.8 QSA attention

docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md

@@ -25,6 +25,8 @@ The BF16, block-FP8, and NVFP4 checkpoints support the same feature set:
 
 The KV-cache dtype and the GDN recurrent-state dtype are selected independently of the model weight precision, through `kv_cache_config.dtype` and `kv_cache_config.mamba_ssm_cache_dtype`.
 
+With an FP8 KV cache, QSA sparse attention selects the PrimTS kernel on SM100/SM103/SM107 when the inputs qualify, and Triton otherwise (small decode batches, unsupported shapes, or PrimTS unavailable). PrimTS converts the BF16 queries to FP8 (unit scale, clamped to the E4M3 range) and attends against the FP8 cache; Triton upcasts K/V to BF16 instead. The two differ at FP8 rounding level: full GSM8K on B300 at TP1 scores 95.45 with either kernel.
+
 ## Prerequisites
 
 * GPU: NVIDIA Blackwell architecture. The BF16 checkpoint also runs on Hopper.

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

tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 """Module-layer integration for QSA sparse attention."""
 
+import itertools
 import math
 from typing import TYPE_CHECKING, Optional
 
@@ -374,6 +375,7 @@ def select_qsa_paged_tokens(
     visible_blocks: torch.Tensor | None = None,
     context_rows: bool = False,
     mtp_share_slot: int | None = None,
+    expand_blocks: bool = True,
 ) -> torch.Tensor:
     """Select tokens with packed, fixed-width paged scoring.
 
@@ -386,6 +388,9 @@ def select_qsa_paged_tokens(
     steps rebuild their rows from it without running the indexer at all. It is
     the calling layer's row in the shared buffers, and ``None`` turns the reuse
     off.
+
+    ``expand_blocks=False`` returns the ``[rows, block_topk]`` compressed-block
+    IDs for consumers that resolve blocks and the causal tail themselves.
     """
     from .kernels import triton_qsa_mtp_reuse_topk, triton_qsa_paged_index_scores
 
@@ -406,6 +411,8 @@ def select_qsa_paged_tokens(
             visible_blocks=visible_blocks,
             output=top_k_output[: q.shape[0]],
         )
+        if not expand_blocks:
+            return indices
         return expand_qsa_block_indices(
             indices,
             query_positions,
@@ -414,20 +421,42 @@ def select_qsa_paged_tokens(
             token_topk=params.token_topk,
         )
 
-    logits = triton_qsa_paged_index_scores(
-        q=q,
-        index_cache=index_cache,
-        block_table=metadata.qsa_block_table,
-        query_positions=query_positions,
-        request_indices=request_indices,
+    # CUDA radix Top-K scans only [row_starts, row_ends), so score tiles
+    # outside the causal boundary need not be materialized. Preserve fully
+    # initialized logits for the Torch fallback below.
+    only_visible_blocks = top_k is not None and q.is_cuda
+    from .prims_ts_core import use_prims_ts_indexer
+
+    if use_prims_ts_indexer(
+        params,
+        q,
         tokens_per_block=metadata.kv_cache_manager.tokens_per_block,
-        compress_ratio=params.compress_ratio,
-        # CUDA radix Top-K scans only [row_starts, row_ends), so score tiles
-        # outside the causal boundary need not be materialized. Preserve fully
-        # initialized logits for the Torch fallback below.
-        only_visible_blocks=top_k is not None and q.is_cuda,
-        context_rows=context_rows,
-    )
+        is_decode=not context_rows,
+    ):
+        from .prims_ts_indexer import prims_ts_qsa_index_scores
+
+        logits = prims_ts_qsa_index_scores(
+            q,
+            index_cache,
+            metadata.qsa_block_table,
+            query_positions,
+            request_indices,
+            tokens_per_block=metadata.kv_cache_manager.tokens_per_block,
+            compress_ratio=params.compress_ratio,
+            only_visible_blocks=only_visible_blocks,
+        )
+    else:
+        logits = triton_qsa_paged_index_scores(
+            q=q,
+            index_cache=index_cache,
+            block_table=metadata.qsa_block_table,
+            query_positions=query_positions,
+            request_indices=request_indices,
+            tokens_per_block=metadata.kv_cache_manager.tokens_per_block,
+            compress_ratio=params.compress_ratio,
+            only_visible_blocks=only_visible_blocks,
+            context_rows=context_rows,
+        )
     if top_k is not None and logits.is_cuda:
         if top_k_output is None or top_k_row_starts is None:
             raise ValueError("QSA CUDA Top-K requires caller-owned output and row starts")
@@ -464,6 +493,8 @@ def select_qsa_paged_tokens(
             )
     if mtp_share_slot is not None:
         _capture_mtp_shared_topk(metadata, mtp_share_slot, indices, visible_blocks)
+    if not expand_blocks:
+        return indices
     return expand_qsa_block_indices(
         indices,
         query_positions,
@@ -629,7 +660,32 @@ def forward(
                 tokens_per_block,
             )
 
-        if attn_metadata.num_contexts == 0:
+        softmax_scale = 1.0 / (attention.q_scaling * attention.head_dim**0.5)
+        is_decode = attn_metadata.num_contexts == 0
+        if params.sparse_core_backend == "triton":
+            core_backend = "triton"
+        else:
+            from .prims_ts_core import resolve_sparse_core_backend
+
+            core_backend = resolve_sparse_core_backend(
+                params,
+                q=q,
+                k_cache=k_cache,
+                page_size=tokens_per_block,
+                is_decode=is_decode,
+                num_rows=num_tokens,
+            )
+        use_prims_ts = core_backend == "prims_ts"
+        if use_prims_ts:
+            from .prims_ts_core import prims_ts_sparse_gqa
+
+            logger.info_once(
+                "QSA PrimTS sparse core attention is active",
+                key="qsa_prims_ts_sparse_core_active",
+            )
+            max_seq_len_kv = attn_metadata.qsa_block_table.shape[1] * tokens_per_block
+
+        if is_decode:
             index_cache = attn_metadata.kv_cache_manager.get_index_k_buffer(attention.layer_idx)
             if index_cache is None:
                 raise RuntimeError(
@@ -661,32 +717,48 @@ def forward(
                 top_k_row_starts=attn_metadata.qsa_topk_row_starts,
                 visible_blocks=attn_metadata.qsa_visible_blocks[:num_tokens],
                 mtp_share_slot=mtp_share_slot,
+                expand_blocks=not use_prims_ts,
             )
             fuse_output_gate = False
-            if output_gate is not None:
-                from .kernels import can_fuse_qsa_splitk_output_gate
-
-                fuse_output_gate = can_fuse_qsa_splitk_output_gate(
+            if use_prims_ts:
+                output = prims_ts_sparse_gqa(
                     q,
                     k_cache,
+                    v_cache,
+                    attn_metadata.qsa_block_table,
                     selected,
-                    output_gate,
+                    req_idx,
+                    logical,
+                    page_size=tokens_per_block,
+                    max_seq_len_kv=max_seq_len_kv,
+                    softmax_scale=softmax_scale,
+                    is_decode=True,
                 )
-                if fuse_output_gate:
-                    logger.info_once(
-                        "QSA split-K merge fuses the attention output gate",
-                        key="qsa_fused_output_gate_active",
+            else:
+                if output_gate is not None:
+                    from .kernels import can_fuse_qsa_splitk_output_gate
+
+                    fuse_output_gate = can_fuse_qsa_splitk_output_gate(
+                        q,
+                        k_cache,
+                        selected,
+                        output_gate,
                     )
-            output = qsa_sparse_gqa(
-                q=q,
-                k_cache=k_cache,
-                v_cache=v_cache,
-                selected_tokens=selected,
-                request_indices=req_idx,
-                metadata=attn_metadata,
-                softmax_scale=1.0 / (attention.q_scaling * attention.head_dim**0.5),
-                output_gate=output_gate if fuse_output_gate else None,
-            )
+                    if fuse_output_gate:
+                        logger.info_once(
+                            "QSA split-K merge fuses the attention output gate",
+                            key="qsa_fused_output_gate_active",
+                        )
+                output = qsa_sparse_gqa(
+                    q=q,
+                    k_cache=k_cache,
+                    v_cache=v_cache,
+                    selected_tokens=selected,
+                    request_indices=req_idx,
+                    metadata=attn_metadata,
+                    softmax_scale=softmax_scale,
+                    output_gate=output_gate if fuse_output_gate else None,
+                )
             output = output.reshape(num_tokens, -1)
             if output_gate is not None and not fuse_output_gate:
                 output = attention.apply_output_gate(output, output_gate)
@@ -700,11 +772,17 @@ def forward(
             attn_metadata.qsa_block_table.shape[1] * tokens_per_block // params.compress_ratio
         )
         chunk_size = _query_chunk_size(num_tokens, score_columns)
+        if use_prims_ts:
+            from .prims_ts_core import slice_request_boundaries
+
+            cu_seq_lens = list(
+                itertools.accumulate(attn_metadata.seq_lens[: attn_metadata.num_seqs].tolist())
+            )
         output = torch.empty_like(q)
-        for start in range(0, num_tokens, chunk_size):
-            end = min(start + chunk_size, num_tokens)
+
+        def select(start: int, end: int) -> torch.Tensor:
             packed_slice = slice(start, end)
-            selected = select_qsa_paged_tokens(
+            return select_qsa_paged_tokens(
                 q_index[packed_slice],
                 index_cache,
                 logical[packed_slice],
@@ -717,18 +795,54 @@ def forward(
                 top_k_row_starts=attn_metadata.qsa_topk_row_starts[packed_slice],
                 visible_blocks=attn_metadata.qsa_visible_blocks[packed_slice],
                 context_rows=True,
+                expand_blocks=not use_prims_ts,
             )
-            output[packed_slice] = qsa_sparse_gqa(
-                q=q[packed_slice],
-                k_cache=k_cache,
-                v_cache=v_cache,
-                selected_tokens=selected,
-                request_indices=req_idx[packed_slice],
-                metadata=attn_metadata,
-                softmax_scale=1.0 / (attention.q_scaling * attention.head_dim**0.5),
-                query_positions=logical[packed_slice],
-                compress_ratio=params.compress_ratio,
-            )
+
+        bounds = [
+            (start, min(start + chunk_size, num_tokens))
+            for start in range(0, num_tokens, chunk_size)
+        ]
+        if use_prims_ts:
+            # The next chunk's scoring and Top-K are queued before this chunk's
+            # PrimTS launch is prepared on the host, so the GPU has work while
+            # the host runs the wrapper. This relies on select() writing only
+            # its own packed_slice of top_k_output and visible_blocks; a shared
+            # scratch buffer inside select_qsa_paged_tokens would corrupt the
+            # pending chunk.
+            selected_next = select(*bounds[0])
+            for i, (start, end) in enumerate(bounds):
+                selected = selected_next
+                if i + 1 < len(bounds):
+                    selected_next = select(*bounds[i + 1])
+                prims_ts_sparse_gqa(
+                    q[start:end],
+                    k_cache,
+                    v_cache,
+                    attn_metadata.qsa_block_table,
+                    selected,
+                    req_idx[start:end],
+                    logical[start:end],
+                    page_size=tokens_per_block,
+                    max_seq_len_kv=max_seq_len_kv,
+                    softmax_scale=softmax_scale,
+                    is_decode=False,
+                    request_boundaries=slice_request_boundaries(cu_seq_lens, start, end),
+                    out=output[start:end],
+                )
+        else:
+            for start, end in bounds:
+                selected = select(start, end)
+                output[start:end] = qsa_sparse_gqa(
+                    q=q[start:end],
+                    k_cache=k_cache,
+                    v_cache=v_cache,
+                    selected_tokens=selected,
+                    request_indices=req_idx[start:end],
+                    metadata=attn_metadata,
+                    softmax_scale=softmax_scale,
+                    query_positions=logical[start:end],
+                    compress_ratio=params.compress_ratio,
+                )
         output = output.reshape(num_tokens, -1)
         if output_gate is not None:
             output = attention.apply_output_gate(output, output_gate)

tensorrt_llm/_torch/attention/backends/sparse/qsa/params.py

@@ -7,6 +7,9 @@
 from ..params import SparseMetadataParams, SparseParams
 from .constants import QSA_INDEX_KV_HEADS
 
+QSA_SPARSE_CORE_BACKENDS = ("auto", "triton", "prims_ts")
+QSA_INDEXER_BACKENDS = ("follow", *QSA_SPARSE_CORE_BACKENDS)
+
 
 def _validate_selection_geometry(token_topk: int, compress_ratio: int) -> None:
     if token_topk <= 0:
@@ -41,8 +44,24 @@ class QSASparseParams(SparseParams):
     enable_heuristic_topk: bool = False
     # Reuse one indexer selection across the MTP draft loop.
     mtp_index_share: bool = False
+    # Kernel selection is decided at runtime like the other sparse backends:
+    # PrimTS on SM100/SM103/SM107 except for small decode batches, Triton
+    # otherwise. These fields pin one kernel for tests and A/B measurements;
+    # the LLM API does not expose them.
+    sparse_core_backend: str = "auto"
+    indexer_backend: str = "follow"
 
     def __post_init__(self) -> None:
+        if self.sparse_core_backend not in QSA_SPARSE_CORE_BACKENDS:
+            raise ValueError(
+                f"QSA sparse_core_backend must be one of {QSA_SPARSE_CORE_BACKENDS}, "
+                f"got {self.sparse_core_backend!r}"
+            )
+        if self.indexer_backend not in QSA_INDEXER_BACKENDS:
+            raise ValueError(
+                f"QSA indexer_backend must be one of {QSA_INDEXER_BACKENDS}, "
+                f"got {self.indexer_backend!r}"
+            )
         values = {
             "index_n_heads": self.index_n_heads,
             "index_kv_heads": self.index_kv_heads,