[None][chore] Refresh PrimTS source onto FlashInfer main (2026-10-30)

3rdparty/vendor_sources.lock.yaml

@@ -14,15 +14,15 @@ vendors:
     digest: sha256-tree-v1:d2d382335ccb68b55078523e63835c28b699607227afb81308cb7bb1d85845ac
   flashinfer-prims-ts:
     url: https://github.com/yuxianq/flashinfer.git
-    branch: trtllm-prims-ts-dev-20261012
-    commit: aa709061230a58858375c513056e7733fa74294a
+    branch: trtllm-prims-ts-dev-20261030
+    commit: d1bed93d69d172fd8f9af96cf14f4f869ac488ef
     source: flashinfer/attention/prims_ts
     destination: tensorrt_llm/_torch/attention/backends/prims_ts
     include:
     - '**/*.py'
     patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch
-    patch_digest: sha256:fc5c23ec3d65c562025a1845f211fdac66bfcd3da156dbb8385f4c4407d20d07
-    digest: sha256-tree-v1:f2f444f7e6d4ad688fbf4e1e9eed11ff7c78d835cfd0af5ef9d7660780b059e0
+    patch_digest: sha256:b41e7a8c0f5eef13fd7b3284402ce8c489a544e590883500735a9b757a29e403
+    digest: sha256-tree-v1:741bf149276a18236e5b91cde228626a001e8c66a71a275d0db4128cd254e0df
   openengine-proto:
     url: https://github.com/ai-dynamo/openengine.git
     tag: v0.1.0

tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py

@@ -672,12 +672,12 @@ def _get_or_plan_context_wrapper(
         v_cache: torch.Tensor,
         *,
         batch_size: int,
-        max_seq_len_q: int,
+        max_q_len: int,
         max_kv_len: int,
         page_size: int,
         mask_type: str,
         window_left: int,
-        sm_scale: float,
+        softmax_scale: float,
         output_dtype: torch.dtype,
     ) -> "BatchPrefillPagedTSWrapper":
         wrapper = self._context_wrappers.get(batch_size)
@@ -693,7 +693,7 @@ def _get_or_plan_context_wrapper(
         wrapper.plan(
             device=q.device,
             batch_size=batch_size,
-            max_seq_len_q=max_seq_len_q,
+            max_q_len=max_q_len,
             max_kv_len=max_kv_len,
             num_qo_heads=int(q.shape[-2]),
             num_kv_heads=int(k_cache.shape[1]),
@@ -705,7 +705,7 @@ def _get_or_plan_context_wrapper(
             page_size=page_size,
             mask_type=mask_type,
             window_left=window_left,
-            sm_scale=sm_scale,
+            softmax_scale=softmax_scale,
             output_scale=1.0,
             # Preprocessing establishes this invariant for every context
             # request admitted by this adapter.
@@ -749,11 +749,11 @@ def _get_or_plan_decode_wrapper(
             head_dim,
             page_size,
             max_kv_len,
-            max_seq_len_q=seq_len_q,
+            max_q_len=seq_len_q,
             packed_query=False,
-            q_data_type=q_dtype,
-            kv_data_type=kv_dtype,
-            o_data_type=output_dtype,
+            q_dtype=q_dtype,
+            kv_dtype=kv_dtype,
+            out_dtype=output_dtype,
             mask_type=mask_type,
             window_left=window_left,
             workspace_buffer=workspace_buffer,
@@ -770,7 +770,7 @@ def _get_or_plan_mla_decode_wrapper(
         kv_lora_rank: int,
         qk_rope_head_dim: int,
         page_size: int,
-        max_seq_len_q: int,
+        max_q_len: int,
         max_kv_len: int,
         q_dtype: torch.dtype,
         kv_dtype: torch.dtype,
@@ -797,11 +797,11 @@ def _get_or_plan_mla_decode_wrapper(
             qk_rope_head_dim,
             page_size,
             max_kv_len,
-            max_seq_len_q=max_seq_len_q,
+            max_q_len=max_q_len,
             packed_query=False,
-            q_data_type=q_dtype,
-            kv_data_type=kv_dtype,
-            o_data_type=output_dtype,
+            q_dtype=q_dtype,
+            kv_dtype=kv_dtype,
+            out_dtype=output_dtype,
             mask_type=mask_type,
             workspace_buffer=workspace_buffer,
         )
@@ -850,13 +850,15 @@ def _forward_without_kv_cache(
         # Per-tensor E4M3 operands as selected by the recipe. The Q and K dequant scales
         # fold into the per-run base-2 softmax scale and the V dequant scale into the
         # kernel output scale.
-        sm_scale = self._get_bmm1_scale(attn)
+        softmax_scale = self._get_bmm1_scale(attn)
         scale_softmax_log2 = None
         if forward_args.qk_fp8_per_tensor:
             q_bshd, q_dequant_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor(q_bshd)
             k_bshd, k_dequant_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor(k_bshd)
             scale_softmax_log2 = (
-                q_dequant_scale.float() * k_dequant_scale.float() * (sm_scale * math.log2(math.e))
+                q_dequant_scale.float()
+                * k_dequant_scale.float()
+                * (softmax_scale * math.log2(math.e))
             ).reshape(1)
         output_scale = None
         if forward_args.v_fp8_per_tensor:
@@ -881,7 +883,7 @@ def _forward_without_kv_cache(
             wrapper.plan(
                 device=q.device,
                 batch_size=batch_size,
-                max_seq_len_q=seq_len,
+                max_q_len=seq_len,
                 max_kv_len=seq_len,
                 num_qo_heads=num_heads,
                 num_kv_heads=num_kv_heads,
@@ -891,7 +893,8 @@ def _forward_without_kv_cache(
                 v_dtype=v_bshd.dtype,
                 out_dtype=out.dtype,
                 mask_type=mask_type,
-                sm_scale=sm_scale,
+                softmax_scale=softmax_scale,
+                cuda_scales=forward_args.qk_fp8_per_tensor or forward_args.v_fp8_per_tensor,
             )
             self._contiguous_context_wrappers[key] = wrapper
         wrapper.run(
@@ -1000,7 +1003,7 @@ def prepare_workspace(
                 int(self.attn.qk_rope_head_dim),
                 int(metadata.tokens_per_block),
                 max_seq_len,
-                max_seq_len_q=seq_len_q,
+                max_q_len=seq_len_q,
                 q_dtype=kernel_dtype,
                 kv_dtype=kernel_dtype,
                 out_dtype=forward_args.output.dtype,
@@ -1177,7 +1180,7 @@ def run_context(self, params: FmhaParams) -> None:
         if kv_page_offset is None:
             raise RuntimeError("PrimTS could not resolve the K-to-V page displacement.")
         k_cache, v_cache = self._standard_kv_views(kv_pool, kv_page_offset)
-        max_seq_len_q = int(meta.max_context_length)
+        max_q_len = int(meta.max_context_length)
         max_kv_len = int(meta.max_seq_len)
         fixed_block_tables = self._get_fixed_block_tables(block_tables, params.batch_size)
         seq_lens_kv = self._get_sequence_lengths(
@@ -1190,12 +1193,12 @@ def run_context(self, params: FmhaParams) -> None:
             k_cache,
             v_cache,
             batch_size=params.batch_size,
-            max_seq_len_q=max_seq_len_q,
+            max_q_len=max_q_len,
             max_kv_len=max_kv_len,
             page_size=params.tokens_per_block,
             mask_type=mask_type,
             window_left=window_left,
-            sm_scale=self._get_bmm1_scale(attn),
+            softmax_scale=self._get_bmm1_scale(attn),
             output_dtype=params.output.dtype,
         )
         wrapper.run(
@@ -1396,8 +1399,8 @@ def run_generation(self, params: FmhaParams) -> None:
             (k_cache, v_cache),
             seq_lens,
             block_tables=fixed_block_tables,
-            bmm1_scale=self._get_bmm1_scale(attn),
-            bmm2_scale=1.0,
+            softmax_scale=self._get_bmm1_scale(attn),
+            output_scale=1.0,
             out=output,
             validate=False,
         )
@@ -1480,7 +1483,7 @@ def run_mla_generation(self, params: FmhaParams) -> None:
             kv_lora_rank=int(attn.kv_lora_rank),
             qk_rope_head_dim=int(attn.qk_rope_head_dim),
             page_size=params.tokens_per_block,
-            max_seq_len_q=seq_len_q,
+            max_q_len=seq_len_q,
             max_kv_len=max_seq_len,
             q_dtype=query.dtype,
             kv_dtype=kv_cache.dtype,
@@ -1491,10 +1494,10 @@ def run_mla_generation(self, params: FmhaParams) -> None:
             query,
             kv_cache,
             block_tables=fixed_block_tables,
-            seq_lens=seq_lens,
+            seq_lens_kv=seq_lens,
             out=output,
-            bmm1_scale=bmm1_scale,
-            bmm2_scale=bmm2_scale,
+            softmax_scale=bmm1_scale,
+            output_scale=bmm2_scale,
             validate=False,
         )
 

tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py

@@ -132,8 +132,8 @@ def strict() -> bool:
 def encode_payload(device_index: int, rows: int, heads: int, plan_kwargs: dict) -> dict[str, Any]:
     """JSON-safe description of one ``BatchSparseMLADecodePagedTSWrapper.plan`` call."""
     kwargs = dict(plan_kwargs)
-    dtype = kwargs.pop("q_data_type")
-    kwargs["q_data_type"] = str(dtype).removeprefix("torch.")
+    dtype = kwargs.pop("q_dtype")
+    kwargs["q_dtype"] = str(dtype).removeprefix("torch.")
     return {"device_index": device_index, "rows": rows, "heads": heads, "plan_kwargs": kwargs}
 
 
@@ -240,7 +240,7 @@ def _plan(mla_decode, payload: dict[str, Any], torch) -> None:
     device = torch.device("cuda", int(payload["device_index"]))
     torch.cuda.set_device(device)
     kwargs = dict(payload["plan_kwargs"])
-    kwargs["q_data_type"] = getattr(torch, kwargs["q_data_type"])
+    kwargs["q_dtype"] = getattr(torch, kwargs["q_dtype"])
     mla_decode.BatchSparseMLADecodePagedTSWrapper().plan(
         device, int(payload["rows"]), int(payload["heads"]), **kwargs
     )

tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py

@@ -538,10 +538,10 @@ def _plan_kwargs(
         plan_kwargs = dict(
             max_topk=int(self.attn.sparse_attention_config.window_size),
             max_extra_topk=extra,
-            max_seq_len_q=1,
-            q_data_type=q_dtype,
+            max_q_len=1,
+            q_dtype=q_dtype,
             kv_layout="NHD",
-            has_sinks=has_sinks,
+            use_attention_sinks=has_sinks,
             return_lse=False,
             assume_valid_prefix=not holes_possible,
         )
@@ -879,7 +879,7 @@ def forward(
             extra_pool,
             softmax_scale=softmax_scale,
             output_scale=output_scale,
-            sinks=sinks,
+            attention_sinks=sinks,
             validate=validate,
             **run_kwargs,
         )

tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py

@@ -41,6 +41,7 @@
     batch_sparse_mla_decode_with_paged_kv_cache,
     get_prims_ts_sparse_mla_decode_workspace_size,
     BatchMLADecodePagedTSWrapper,
+    MLAAttentionTuning,
     batch_mla_decode_with_paged_kv_cache,
     get_prims_ts_batch_mla_decode_workspace_size,
 )
@@ -50,6 +51,13 @@
     q_token_kv_block_sparse_attention_with_paged_kv_cache,
 )
 from .sage import SageAttentionConfig, SageAttentionParams
+from .vc_attention import (
+    VCAttentionConfig,
+    VCAttentionParams,
+)
+from .vc_attention.quantization import (
+    VCAttentionPreprocessor,
+)
 
 __all__ = [
     "BatchSparseMLADecodePagedTSWrapper",
@@ -58,6 +66,9 @@
     "BlockSparseTSWrapper",
     "SageAttentionConfig",
     "SageAttentionParams",
+    "VCAttentionConfig",
+    "VCAttentionParams",
+    "VCAttentionPreprocessor",
     "BlockSparsePagedTSWrapper",
     "block_sparse_attention",
     "block_sparse_attention_with_paged_kv_cache",
@@ -77,6 +88,7 @@
     "get_q_token_kv_block_sparse_workspace_size",
     "q_token_kv_block_sparse_attention_with_paged_kv_cache",
     "BatchMLADecodePagedTSWrapper",
+    "MLAAttentionTuning",
     "batch_mla_decode_with_paged_kv_cache",
     "get_prims_ts_batch_mla_decode_workspace_size",
 ]

tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/common.py

@@ -14,14 +14,56 @@
 
 """Dependency-neutral semantic contract for PrimTS block-sparse attention."""
 
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
+from typing import TypeVar
 
 _FINE_KV_BLOCK_SIZES = (8, 16, 32)
 _PREPARED_KV_ROUTE_SIZE = 128
 _MAX_KV_ATOM_SIZE = 64
 _SIGNED_INT32_MAX = (1 << 31) - 1
 _BLOCK_SPARSE_Q_TILE_SIZES = (8, 16, 32, 64, 128)
 _BLOCK_SPARSE_MAX_HEADS_Q_PER_KV = 32
+# Q and KV block sizes of a dense plan that names none; they select the
+# Q128/KV128 Keeps tile.
+_DENSE_Q_BLOCK_SIZE = 128
+_DENSE_KV_BLOCK_SIZE = 128
+# Sage scale slots of the block-sparse attention adapter, in ABI order.
+SAGE_ADAPTER_SLOTS = ("q_scale", "k_scale", "k_summary_scale", "v_scale", "v_mean")
+
+_T = TypeVar("_T")
+
+
+def sage_adapter_slots(values: Mapping[str, _T | None]) -> tuple[_T | None, ...]:
+    """Arrange Sage values named by scale into the adapter slots.
+
+    An absent name leaves its slot ``None``, which the adapter binds as a null
+    pointer.
+    """
+    return tuple(values.get(name) for name in SAGE_ADAPTER_SLOTS)
+
+
+def _resolve_plan_block_sizes(
+    q_block_size: int | None,
+    kv_block_size: int | None,
+    *,
+    use_block_sparse: bool,
+) -> tuple[int, int]:
+    """Return a plan's Q and KV block sizes.
+
+    A block-sparse plan names both, since they give its routes their meaning;
+    a dense plan may omit either, which then takes the dense default.
+    """
+
+    if use_block_sparse:
+        if q_block_size is None or kv_block_size is None:
+            raise ValueError(
+                "q_block_size and kv_block_size are required by a block-sparse plan"
+            )
+        return q_block_size, kv_block_size
+    return (
+        _DENSE_Q_BLOCK_SIZE if q_block_size is None else q_block_size,
+        _DENSE_KV_BLOCK_SIZE if kv_block_size is None else kv_block_size,
+    )
 
 
 def _num_sparse_pattern_heads(

tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py

@@ -23,8 +23,8 @@
 
 from flashinfer.utils import ceil_div
 
-from .common import _num_sparse_pattern_heads
-from ..sage import sage_adapter_slots, sage_scale_shapes
+from .common import _num_sparse_pattern_heads, sage_adapter_slots
+from ..sage.config import sage_scale_shapes
 from .config import _BlockSparseCompileKey, _make_block_sparse_config
 
 
@@ -59,7 +59,7 @@ def _compile_block_sparse(key: _BlockSparseCompileKey) -> Callable[..., object]:
         fmha_block_sparse_launch,
         fmha_decode_launch,
     )
-    from ..kernels.fmha_decode.sage_prepare import _PrepareSageScales
+    from ..sage.prepare import _PrepareSageScales
 
     config = _make_block_sparse_config(key)
     sparse_format = key.sparse_format
@@ -74,10 +74,7 @@ def _compile_block_sparse(key: _BlockSparseCompileKey) -> Callable[..., object]:
     if config.use_sage_attention:
         sage_scale_layout = config.sage_scale_layout(key.seq_len_kv)
         prepare_sage_scales = _PrepareSageScales(
-            cfg=config,
-            layout=sage_scale_layout,
-            batch_size=key.batch_size,
-            num_kv_heads=key.num_kv_heads,
+            layout=sage_scale_layout, num_kv_heads=key.num_kv_heads
         )
     if key.use_block_sparse:
         prepare_kwargs = {