[https://nvbugs/6646140][fix] Update flashinfer-python from 0.6.18 to 0.7.1

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

@@ -1,28 +1,25 @@
 diff --git a/block_sparse.py b/block_sparse.py
-index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea063f9e6e4 100644
+index d4652b8b68958ee3dab29d062a7cd082ea604c47..855a17627a6ba340e7c96ca94389aa0e4fb1a74c 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");
-@@ -36,13 +37,7 @@ from typing import Literal
+@@ -36,10 +37,8 @@ from typing import Literal
 
  import torch
 
 -from flashinfer.api_logging import flashinfer_experimental_api
--from flashinfer.trace.templates.attention import (
++from flashinfer.api_logging import flashinfer_api as 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 (
-     _num_sparse_pattern_heads,
-@@ -301,7 +296,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase):
+     prims_ts_paged_block_sparse_trace_dispatch,
+     prims_ts_paged_block_sparse_wrapper_trace_dispatch,
+ )
+@@ -301,7 +300,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase):
          # previously published revision intact and runnable.
          self._plan_state = candidate
 
@@ -31,7 +28,7 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0
      def run(
          self,
          q: torch.Tensor,
-@@ -433,7 +428,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase):
+@@ -433,7 +432,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase):
          return self._launch_validated_run(state, run_args, run_stream)
 
 
@@ -40,26 +37,6 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0
  def block_sparse_attention(
      q: torch.Tensor,
      k: torch.Tensor,
-@@ -726,9 +721,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase):
-             )
-         self._plan_state = candidate
-
--    @flashinfer_experimental_api(
--        trace=prims_ts_paged_block_sparse_wrapper_trace_dispatch
--    )
-+    @flashinfer_experimental_api
-     def run(
-         self,
-         q: torch.Tensor,
-@@ -850,7 +843,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase):
-         return self._launch_validated_run(state, run_args, run_stream)
-
-
--@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 dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..1163324ec632f2a0938f78a3636fb415450369df 100644
 --- a/context.py
@@ -83,28 +60,24 @@ index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..1163324ec632f2a0938f78a3636fb415
          require_cute_dsl_arch(device_index)
      return device_index
 diff --git a/decode.py b/decode.py
-index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa6d76a609 100644
+index 41479a386d72c27eb0df76107318dc09256aff06..0f0be88df5f31619fd4b3c5628073121000252c5 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");
-@@ -24,11 +25,7 @@ from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union, cast
+@@ -24,7 +25,7 @@ from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union, cast
 
  import torch
 
 -from flashinfer.api_logging import flashinfer_experimental_api
--from flashinfer.trace.templates.attention import (
--    attention_ts_decode_trace_dispatch,
--    prims_ts_decode_wrapper_trace_dispatch,
--)
 +from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api
-
- 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:
+ from flashinfer.trace.templates.attention import (
+     attention_ts_decode_trace_dispatch,
+     prims_ts_decode_wrapper_trace_dispatch,
+@@ -789,7 +790,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):
@@ -113,54 +86,24 @@ index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa
 
          require_cute_dsl_arch(device_index)
 
-@@ -4046,7 +4043,7 @@ class BatchDecodePagedTSWrapper:
-         # previous complete plan revision usable.
-         self._plan_state = candidate
-
--    @flashinfer_experimental_api(trace=prims_ts_decode_wrapper_trace_dispatch)
-+    @flashinfer_experimental_api
-     def run(
-         self,
-         q: torch.Tensor,
-@@ -4228,7 +4225,7 @@ class BatchDecodePagedTSWrapper:
-         )
-
-
--@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 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..3d5fb159a3298d1edf72935b43d9a2eef74a1a4b 100644
+index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..7996bee06a74ceb3a4bc8c53ce90ea2ad7fe449a 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");
-@@ -21,10 +22,9 @@ from typing import Any, Literal, Optional, cast
+@@ -21,7 +22,7 @@ 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_wrapper_trace_dispatch,
- )
-
- from .decode import (
-@@ -1576,7 +1576,7 @@ class BatchMLADecodePagedTSWrapper:
-             split_kv=int(dict(policy)["split_kv"]),
-         )
-
--    @flashinfer_experimental_api(trace=prims_ts_decode_mla_wrapper_trace_dispatch)
-+    @flashinfer_experimental_api
-     def run(
-         self,
-         query: torch.Tensor,
+     prims_ts_decode_mla_wrapper_trace_dispatch,
 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

3rdparty/vendor_sources.lock.yaml

@@ -9,8 +9,8 @@ vendors:
     include:
     - '**/*.py'
     patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch
-    patch_digest: sha256:65cc08e20a9cc7bdc190eb6f3e79124986dced74394c1f3250ae2d100a8ee340
-    digest: sha256-tree-v1:e411b6c6679474d362b166e586df4276f52322d75a35c59c782d6266143e560e
+    patch_digest: sha256:635a5b7445863b024055ba3404c28d52baf1b37bf701a40102afb3e99c121e07
+    digest: sha256-tree-v1:f6f1b4d9ba26138193de1215a1b93abea7868eca7b641a6f75b6d422e32088b0
   openengine-proto:
     url: https://github.com/ai-dynamo/openengine.git
     tag: v0.1.0

ATTRIBUTIONS-Python.md

@@ -5261,7 +5261,7 @@ For more information, please refer to <http://unlicense.org>
   - `Tracker`: https://github.com/tox-dev/py-filelock/issues
 
 
-## flashinfer-python (0.6.18)
+## flashinfer-python (0.7.1)
 
 ### Licenses
 License: `Apache-2.0`

requirements.txt

@@ -68,7 +68,7 @@ peft>=0.18.1,<0.19.0
 patchelf
 einops
 fla-core==0.5.2 # the Kimi Linear / Kimi K3 KDA layers import fla.modules and fla.ops when they load
-flashinfer-python==0.6.18
+flashinfer-python==0.7.1
 xgrammar==0.1.32
 llguidance==0.7.29
 jsonschema
@@ -106,7 +106,7 @@ cuda-core
 llist
 cuda-tile>=1.0.1
 nvidia-cuda-tileiras>=13.1,<13.2
-nvidia-cudnn-frontend>=1.27.0,<1.28
+nvidia-cudnn-frontend>=1.31.0,<1.32
 etcd-sdk-python==0.0.7
 # Required by the checked-in OpenEngine gRPC bindings.
 grpcio>=1.67.1,<2

security_scanning/pyproject.toml

@@ -56,7 +56,7 @@ dependencies = [
     "patchelf (>=0.19.1.0,<0.20.0.0)",
     "einops (>=0.8.2,<0.9.0)",
     "fla-core (==0.5.2)",
-    "flashinfer-python (==0.6.18)",
+    "flashinfer-python (==0.7.1)",
     "xgrammar (==0.1.32)",
     "llguidance (==0.7.29)",
     "jsonschema (>=4.26.0,<5.0.0)",

tensorrt_llm/_torch/attention/backends/flashinfer.py

@@ -1344,15 +1344,6 @@ def _plan_ragged_no_kv(
             max_query_tokens_per_sequence = cast(int,
                                                  q_and_kv_seqlens.max().item())
 
-            if flashinfer_backend == "cudnn":
-                # cuDNN ragged prefill uses *element* offsets in qo/kv indptr, not token indptr.
-                key_value_indptr = key_value_indptr.clone(
-                )  # unalias query_output_indptr
-                query_output_indptr[1:].mul_(num_query_output_heads *
-                                             attention_head_dim)
-                key_value_indptr[1:].mul_(num_key_value_heads *
-                                          attention_head_dim)
-
         if (multi_item_params := plan_params.multi_item_params) is not None:
             if flashinfer_backend not in ["fa2", "fa3"]:
                 raise ValueError(

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

@@ -38,6 +38,10 @@
 import torch
 
 from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api
+from flashinfer.trace.templates.attention import (
+    prims_ts_paged_block_sparse_trace_dispatch,
+    prims_ts_paged_block_sparse_wrapper_trace_dispatch,
+)
 
 from ._block_sparse.common import (
     _num_sparse_pattern_heads,
@@ -721,7 +725,9 @@ def plan(
             )
         self._plan_state = candidate
 
-    @flashinfer_experimental_api
+    @flashinfer_experimental_api(
+        trace=prims_ts_paged_block_sparse_wrapper_trace_dispatch
+    )
     def run(
         self,
         q: torch.Tensor,
@@ -843,7 +849,7 @@ def run(
         return self._launch_validated_run(state, run_args, run_stream)
 
 
-@flashinfer_experimental_api
+@flashinfer_experimental_api(trace=prims_ts_paged_block_sparse_trace_dispatch)
 def block_sparse_attention_with_paged_kv_cache(
     q: torch.Tensor,
     paged_kv_cache: PagedKVCache,

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

@@ -26,6 +26,10 @@
 import torch
 
 from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api
+from flashinfer.trace.templates.attention import (
+    attention_ts_decode_trace_dispatch,
+    prims_ts_decode_wrapper_trace_dispatch,
+)
 
 from . import _q_token_kv_block_sparse_policy as _sparse_policy
 from ._block_sparse.common import _num_sparse_pattern_heads
@@ -4043,7 +4047,7 @@ def plan(
         # previous complete plan revision usable.
         self._plan_state = candidate
 
-    @flashinfer_experimental_api
+    @flashinfer_experimental_api(trace=prims_ts_decode_wrapper_trace_dispatch)
     def run(
         self,
         q: torch.Tensor,
@@ -4225,7 +4229,7 @@ def run(
         )
 
 
-@flashinfer_experimental_api
+@flashinfer_experimental_api(trace=attention_ts_decode_trace_dispatch)
 def batch_decode_with_paged_kv_cache(
     q: torch.Tensor,
     paged_kv_cache: PagedKVCache,

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

@@ -25,6 +25,7 @@
 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_wrapper_trace_dispatch,
 )
 
 from .decode import (
@@ -1576,7 +1577,7 @@ def plan(
             split_kv=int(dict(policy)["split_kv"]),
         )
 
-    @flashinfer_experimental_api
+    @flashinfer_experimental_api(trace=prims_ts_decode_mla_wrapper_trace_dispatch)
     def run(
         self,
         query: torch.Tensor,

tests/unittest/_torch/attention/test_flashinfer_attention.py

@@ -604,6 +604,9 @@ def plan(self, **kwargs):
         self.assertIn("kv_indptr", wrapper.plan_kwargs)
         self.assertNotIn("v_indptr", wrapper.plan_kwargs)
         self.assertNotIn("o_indptr", wrapper.plan_kwargs)
+        # FlashInfer's cudnn backend takes token-unit indptrs like every other backend.
+        self.assertEqual(wrapper.plan_kwargs["qo_indptr"].tolist(), [0, 3, 8])
+        self.assertEqual(wrapper.plan_kwargs["kv_indptr"].tolist(), [0, 3, 8])
 
     @parameterized.expand([
         Scenario(num_layers=1,