[None][chore] Refresh PrimTS vendor source onto FlashInfer main

3rdparty/vendor_patches/flashinfer-prims-ts.patch

@@ -1,76 +1,79 @@
 diff --git a/block_sparse.py b/block_sparse.py
-index c7e881ca73dfb95314da532912db39fe7b05eb94..e19dc7b1de7fca03c44b8e98217f8d10fb17e26b 100644
+index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea063f9e6e4 100644
 --- a/block_sparse.py
 +++ b/block_sparse.py
 @@ -1,3 +1,4 @@
 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
  # Copyright (c) 2026 by FlashInfer team.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
-@@ -35,12 +36,6 @@ from typing import Literal
+@@ -36,13 +37,7 @@ from typing import Literal
+
  import torch
 
- from flashinfer.api_logging import flashinfer_api
+-from flashinfer.api_logging import flashinfer_experimental_api
 -from flashinfer.trace.templates.attention import (
 -    prims_ts_block_sparse_trace_dispatch,
 -    prims_ts_block_sparse_wrapper_trace_dispatch,
 -    prims_ts_paged_block_sparse_trace_dispatch,
 -    prims_ts_paged_block_sparse_wrapper_trace_dispatch,
 -)
++from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api
 
- from ._block_sparse.common import _validate_contiguous_route_mode
- from ._block_sparse.config import _validate_block_sparse_static_profile
-@@ -239,7 +234,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase):
+ from ._block_sparse.common import (
+     _num_sparse_pattern_heads,
+@@ -301,7 +296,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase):
          # previously published revision intact and runnable.
          self._plan_state = candidate
 
--    @flashinfer_api(trace=prims_ts_block_sparse_wrapper_trace_dispatch)
-+    @flashinfer_api
+-    @flashinfer_experimental_api(trace=prims_ts_block_sparse_wrapper_trace_dispatch)
++    @flashinfer_experimental_api
      def run(
          self,
          q: torch.Tensor,
-@@ -361,7 +356,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase):
+@@ -433,7 +428,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase):
          return self._launch_validated_run(state, run_args, run_stream)
 
 
--@flashinfer_api(trace=prims_ts_block_sparse_trace_dispatch)
-+@flashinfer_api
+-@flashinfer_experimental_api(trace=prims_ts_block_sparse_trace_dispatch)
++@flashinfer_experimental_api
  def block_sparse_attention(
      q: torch.Tensor,
      k: torch.Tensor,
-@@ -607,7 +602,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase):
+@@ -726,9 +721,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase):
              )
          self._plan_state = candidate
 
--    @flashinfer_api(trace=prims_ts_paged_block_sparse_wrapper_trace_dispatch)
-+    @flashinfer_api
+-    @flashinfer_experimental_api(
+-        trace=prims_ts_paged_block_sparse_wrapper_trace_dispatch
+-    )
++    @flashinfer_experimental_api
      def run(
          self,
          q: torch.Tensor,
-@@ -727,7 +722,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase):
+@@ -850,7 +843,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase):
          return self._launch_validated_run(state, run_args, run_stream)
 
 
--@flashinfer_api(trace=prims_ts_paged_block_sparse_trace_dispatch)
-+@flashinfer_api
+-@flashinfer_experimental_api(trace=prims_ts_paged_block_sparse_trace_dispatch)
++@flashinfer_experimental_api
  def block_sparse_attention_with_paged_kv_cache(
      q: torch.Tensor,
      paged_kv_cache: PagedKVCache,
 diff --git a/context.py b/context.py
-index 4ada42762c38c1c7215d2dfbe4582926fd8f9701..678432780332b20a78fb553405ad1e29d692b430 100644
+index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..1163324ec632f2a0938f78a3636fb415450369df 100644
 --- a/context.py
 +++ b/context.py
-@@ -31,8 +31,7 @@ position is ``q + (S_kv - S_q)`` and ``window_left`` is measured from that
- position.
+@@ -46,7 +46,7 @@ from typing import TYPE_CHECKING, Callable, Literal, Optional
+
+ import torch
 
- PrimTS context entry points are intentionally excluded from ``fi_trace`` for
--now; unlike the decode APIs, their ``@flashinfer_api`` decorators do not
--register trace templates.
-+now; their ``@flashinfer_api`` decorators do not register trace templates.
- """
+-from flashinfer.api_logging import flashinfer_experimental_api
++from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api
 
- from dataclasses import dataclass
-@@ -438,7 +437,7 @@ def _validate_device(device: torch.device) -> int:
+
+ if TYPE_CHECKING:
+@@ -488,7 +488,7 @@ def _validate_device(device: torch.device) -> int:
      # Rubin runs through the sm_100f family target; a CuTe DSL older than 4.8
      # cannot emit for it unless CUTE_DSL_ARCH=sm_100f is set before import.
      if capability == (10, 7):
@@ -80,85 +83,111 @@ index 4ada42762c38c1c7215d2dfbe4582926fd8f9701..678432780332b20a78fb553405ad1e29
          require_cute_dsl_arch(device_index)
      return device_index
 diff --git a/decode.py b/decode.py
-index b47093bd298ea9b5514e6b25ad2a34a47ac3bff0..35236fcc7ce98bdee1f51f72dbd1f841cdbce7b2 100644
+index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa6d76a609 100644
 --- a/decode.py
 +++ b/decode.py
 @@ -1,3 +1,4 @@
 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
  # Copyright (c) 2026 by FlashInfer team.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
-@@ -25,11 +26,6 @@ from typing import TYPE_CHECKING, Literal, Optional, Union
+@@ -24,11 +25,7 @@ from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union, cast
+
  import torch
 
- from flashinfer.api_logging import flashinfer_api
+-from flashinfer.api_logging import flashinfer_experimental_api
 -from flashinfer.trace.templates.attention import (
 -    attention_ts_decode_trace_dispatch,
--    prims_ts_decode_trace_dispatch,
 -    prims_ts_decode_wrapper_trace_dispatch,
 -)
++from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api
 
- from ._tensor_aliasing import (
-     _validate_out_does_not_overlap_inputs,
-@@ -702,7 +698,7 @@ def _validate_runtime_device(device: torch.device) -> int:
+ from . import _q_token_kv_block_sparse_policy as _sparse_policy
+ from ._block_sparse.common import _num_sparse_pattern_heads
+@@ -789,7 +786,7 @@ def _validate_runtime_device(device: torch.device) -> None:
      # Rubin runs through the sm_100f family target; a CuTe DSL older than 4.8
      # cannot emit for it unless CUTE_DSL_ARCH=sm_100f is set before import.
      if capability == (10, 7):
 -        from ...cute_dsl.utils import require_cute_dsl_arch
 +        from flashinfer.cute_dsl.utils import require_cute_dsl_arch
 
          require_cute_dsl_arch(device_index)
-     return device_index
-@@ -2145,7 +2141,7 @@ def _validate_decode_run_metadata_values(
-                 )
 
-
--@flashinfer_api(trace=prims_ts_decode_trace_dispatch)
-+@flashinfer_api
- def prims_ts_batch_decode_with_kv_cache(
-     query: torch.Tensor,
-     kv_cache: PagedKVCache,
-@@ -2708,7 +2704,7 @@ class BatchDecodePagedTSWrapper:
+@@ -4046,7 +4043,7 @@ class BatchDecodePagedTSWrapper:
          # previous complete plan revision usable.
          self._plan_state = candidate
 
--    @flashinfer_api(trace=prims_ts_decode_wrapper_trace_dispatch)
-+    @flashinfer_api
+-    @flashinfer_experimental_api(trace=prims_ts_decode_wrapper_trace_dispatch)
++    @flashinfer_experimental_api
      def run(
          self,
          q: torch.Tensor,
-@@ -2893,7 +2889,7 @@ class BatchDecodePagedTSWrapper:
+@@ -4228,7 +4225,7 @@ class BatchDecodePagedTSWrapper:
          )
 
 
--@flashinfer_api(trace=attention_ts_decode_trace_dispatch)
-+@flashinfer_api
+-@flashinfer_experimental_api(trace=attention_ts_decode_trace_dispatch)
++@flashinfer_experimental_api
  def batch_decode_with_paged_kv_cache(
      q: torch.Tensor,
      paged_kv_cache: PagedKVCache,
 diff --git a/mla_decode.py b/mla_decode.py
-index 8b7f19d1cc2baab6309018d46213f509c44d919a..cff2ee7d2ad051c8061ad396b87ee973eebed788 100644
+index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..3d5fb159a3298d1edf72935b43d9a2eef74a1a4b 100644
 --- a/mla_decode.py
 +++ b/mla_decode.py
 @@ -1,3 +1,4 @@
 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
  # Copyright (c) 2026 by FlashInfer team.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
-@@ -25,7 +26,6 @@ from flashinfer.api_logging import flashinfer_api
+@@ -21,10 +22,9 @@ from typing import Any, Literal, Optional, cast
+
+ import torch
+
+-from flashinfer.api_logging import flashinfer_experimental_api
++from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api
  from flashinfer.trace.templates.attention import (
      prims_ts_decode_mla_one_shot_trace_dispatch,
-     prims_ts_decode_mla_trace_dispatch,
 -    prims_ts_decode_mla_wrapper_trace_dispatch,
  )
 
- from ._tensor_aliasing import (
-@@ -1799,7 +1799,7 @@ class BatchMLADecodePagedTSWrapper:
+ from .decode import (
+@@ -1576,7 +1576,7 @@ class BatchMLADecodePagedTSWrapper:
              split_kv=int(dict(policy)["split_kv"]),
          )
 
--    @flashinfer_api(trace=prims_ts_decode_mla_wrapper_trace_dispatch)
-+    @flashinfer_api
+-    @flashinfer_experimental_api(trace=prims_ts_decode_mla_wrapper_trace_dispatch)
++    @flashinfer_experimental_api
      def run(
          self,
          query: torch.Tensor,
+diff --git a/q_token_kv_block_sparse_metadata.py b/q_token_kv_block_sparse_metadata.py
+index 35efc60b7ff3dac6af1c5c5e46631712a171fb72..0c5cf40dbd4e4708b8f74bfa698f3c0788f9f3a4 100644
+--- a/q_token_kv_block_sparse_metadata.py
++++ b/q_token_kv_block_sparse_metadata.py
+@@ -1,3 +1,4 @@
++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+ # Copyright (c) 2026 by FlashInfer team.
+ #
+ # Licensed under the Apache License, Version 2.0 (the "License");
+@@ -40,9 +41,9 @@ from dataclasses import dataclass
+ from typing import Any, Literal, Optional
+
+ import torch
+-from flashinfer.api_logging import flashinfer_experimental_api
++from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api
+
+-from ...utils import device_support_pdl
++from flashinfer.utils import device_support_pdl
+ from ._block_sparse.common import _num_sparse_pattern_heads
+
+ _Q_TOKEN_KV_BLOCK_SPARSE_MIN_BLOCK_SIZE = 4
+@@ -62,7 +63,7 @@ _Q_TOKEN_KV_BLOCK_SPARSE_WORKSPACE_ALIGNMENT = 256
+ def _get_prims_ts_q_token_kv_block_sparse_metadata_module() -> Any:
+     """Build and cache the CUDA C++ QToken-KvBlock-Sparse-Attention metadata implementation."""
+
+-    from ...jit import gen_prims_ts_q_token_kv_block_sparse_metadata_module
++    from flashinfer.jit import gen_prims_ts_q_token_kv_block_sparse_metadata_module
+
+     return gen_prims_ts_q_token_kv_block_sparse_metadata_module().build_and_load()
+

3rdparty/vendor_sources.lock.yaml

@@ -2,15 +2,15 @@ schema_version: 1
 vendors:
   flashinfer-prims-ts:
     url: https://github.com/yuxianq/flashinfer.git
-    branch: trtllm-prims-ts-dev
-    commit: 6e0d64e79f2a87eeb539e2f40ea12d009679a8f5
+    branch: trtllm-prims-ts-dev-20261012
+    commit: ae312f7cc1723b06fda5168105ade0f328ee47d8
     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:ba917f330d8281781ce30ddb89f1772ee8f9017264d948b588b881a92f511e45
-    digest: sha256-tree-v1:fa65b674a9ec6121b4158069d5585cc56df28243d1ac1199eb97003d81135d76
+    patch_digest: sha256:65cc08e20a9cc7bdc190eb6f3e79124986dced74394c1f3250ae2d100a8ee340
+    digest: sha256-tree-v1:e411b6c6679474d362b166e586df4276f52322d75a35c59c782d6266143e560e
   openengine-proto:
     url: https://github.com/ai-dynamo/openengine.git
     tag: v0.1.0

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

@@ -621,7 +621,8 @@ def _get_or_plan_context_wrapper(
             num_kv_heads=int(k_cache.shape[1]),
             head_dim=int(q.shape[-1]),
             q_dtype=q.dtype,
-            kv_dtype=k_cache.dtype,
+            k_dtype=k_cache.dtype,
+            v_dtype=v_cache.dtype,
             out_dtype=output_dtype,
             page_size=page_size,
             mask_type=mask_type,

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

@@ -22,9 +22,13 @@
 )
 from .decode import (
     BatchDecodePagedTSWrapper,
+    PrimsTSBatchDecodePlan,
     batch_decode_with_paged_kv_cache,
     get_prims_ts_batch_decode_workspace_size,
-    prims_ts_batch_decode_with_kv_cache,
+    make_q_token_kv_block_sparse_qo_indptr,
+    prepare_prims_ts_batch_decode_with_kv_cache,
+    suggest_q_token_kv_block_sparse_group_size,
+    validate_q_token_kv_block_sparse_group_size,
 )
 from .context import (
     BatchPrefillPagedTSWrapper,
@@ -36,11 +40,18 @@
     BatchMLADecodePagedTSWrapper,
     batch_mla_decode_with_paged_kv_cache,
     get_prims_ts_batch_mla_decode_workspace_size,
-    prims_ts_batch_mla_decode_with_kv_cache,
 )
+from .q_token_kv_block_sparse_metadata import (
+    QTokenKvBlockSparsePagedTSWrapper,
+    get_q_token_kv_block_sparse_workspace_size,
+    q_token_kv_block_sparse_attention_with_paged_kv_cache,
+)
+from .sage import SageAttentionConfig, SageAttentionParams
 
 __all__ = [
     "BlockSparseTSWrapper",
+    "SageAttentionConfig",
+    "SageAttentionParams",
     "BlockSparsePagedTSWrapper",
     "block_sparse_attention",
     "block_sparse_attention_with_paged_kv_cache",
@@ -49,11 +60,17 @@
     "batch_prefill",
     "batch_prefill_with_paged_kv_cache",
     "BatchDecodePagedTSWrapper",
+    "PrimsTSBatchDecodePlan",
     "batch_decode_with_paged_kv_cache",
     "get_prims_ts_batch_decode_workspace_size",
-    "prims_ts_batch_decode_with_kv_cache",
+    "make_q_token_kv_block_sparse_qo_indptr",
+    "prepare_prims_ts_batch_decode_with_kv_cache",
+    "suggest_q_token_kv_block_sparse_group_size",
+    "validate_q_token_kv_block_sparse_group_size",
+    "QTokenKvBlockSparsePagedTSWrapper",
+    "get_q_token_kv_block_sparse_workspace_size",
+    "q_token_kv_block_sparse_attention_with_paged_kv_cache",
     "BatchMLADecodePagedTSWrapper",
     "batch_mla_decode_with_paged_kv_cache",
     "get_prims_ts_batch_mla_decode_workspace_size",
-    "prims_ts_batch_mla_decode_with_kv_cache",
 ]

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

@@ -14,6 +14,8 @@
 
 """Dependency-neutral semantic contract for PrimTS block-sparse attention."""
 
+from collections.abc import Iterable
+
 _FINE_KV_BLOCK_SIZES = (8, 16, 32)
 _PREPARED_KV_ROUTE_SIZE = 128
 _MAX_KV_ATOM_SIZE = 64
@@ -22,6 +24,16 @@
 _BLOCK_SPARSE_MAX_HEADS_Q_PER_KV = 32
 
 
+def _num_sparse_pattern_heads(
+    num_kv_heads: int, share_pattern_across_kv_heads: bool
+) -> int:
+    """Return metadata head count without changing physical KV-head geometry.
+
+    The caller resolves its API's pattern-sharing default before calling.
+    """
+    return 1 if share_pattern_across_kv_heads else num_kv_heads
+
+
 def _validate_contiguous_route_mode(
     sparse_format: object,
     use_proxy_routes: object,
@@ -36,6 +48,19 @@ def _validate_contiguous_route_mode(
         raise TypeError("use_proxy_routes must be a bool")
 
 
+def _validate_dense_contiguous_plan_inputs(
+    inputs: Iterable[tuple[str, object, object]],
+) -> None:
+    """Reject block-sparse-only inputs given to a dense contiguous plan.
+
+    Each entry is ``(name, value, unused value)``.
+    """
+
+    for name, value, unused in inputs:
+        if value != unused:
+            raise ValueError(f"{name} is unsupported by a dense contiguous plan")
+
+
 def _validate_sparse_q_block_size(value: object) -> int:
     """Return a positive semantic Q block size representable by the ABI."""