diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index d6f80daccc16..5ea705b71b48 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -1,5 +1,5 @@ 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 @@ @@ -7,22 +7,19 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0 # 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,7 +60,7 @@ 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 @@ @@ -91,20 +68,16 @@ index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa # 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,26 +86,8 @@ 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 @@ @@ -140,7 +95,7 @@ index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..3d5fb159a3298d1edf72935b43d9a2ee # 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 @@ -148,19 +103,7 @@ index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..3d5fb159a3298d1edf72935b43d9a2ee +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 diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 2e914da8f052..b46f81988cfb 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/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 diff --git a/ATTRIBUTIONS-Python.md b/ATTRIBUTIONS-Python.md index 11c934c2ee82..8d5c798c531b 100644 --- a/ATTRIBUTIONS-Python.md +++ b/ATTRIBUTIONS-Python.md @@ -5261,7 +5261,7 @@ For more information, please refer to - `Tracker`: https://github.com/tox-dev/py-filelock/issues -## flashinfer-python (0.6.18) +## flashinfer-python (0.7.1) ### Licenses License: `Apache-2.0` diff --git a/requirements.txt b/requirements.txt index 0e3070694b06..19524a1d3dff 100644 --- a/requirements.txt +++ b/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 diff --git a/security_scanning/pyproject.toml b/security_scanning/pyproject.toml index e9e6550ac15b..36531007708b 100644 --- a/security_scanning/pyproject.toml +++ b/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)", diff --git a/tensorrt_llm/_torch/attention/backends/flashinfer.py b/tensorrt_llm/_torch/attention/backends/flashinfer.py index e3373666692c..33ad107f7a28 100644 --- a/tensorrt_llm/_torch/attention/backends/flashinfer.py +++ b/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( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py b/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py index bb56724e986d..855a17627a6b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py +++ b/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, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py index 3fcec5efdc2a..0f0be88df5f3 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py +++ b/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, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py index 3d5fb159a329..7996bee06a74 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py +++ b/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, diff --git a/tests/unittest/_torch/attention/test_flashinfer_attention.py b/tests/unittest/_torch/attention/test_flashinfer_attention.py index 94e154517ebf..695e7be55859 100644 --- a/tests/unittest/_torch/attention/test_flashinfer_attention.py +++ b/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,