[Dependency] Upgrade FlashInfer version to 0.7.0 (#58069)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>

docker/Dockerfile

@@ -999,7 +999,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.18.post1
+ARG FLASHINFER_VERSION=0.7.0
 RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "${INSTALL_RUBIN_PRERELEASE}" != "true" ]; then \
         uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \

docker/versions.json

@@ -71,7 +71,7 @@
       "default": "true"
     },
     "FLASHINFER_VERSION": {
-      "default": "0.6.18.post1"
+      "default": "0.7.0"
     },
     "GDRCOPY_CUDA_VERSION": {
       "default": "12.8"

requirements/cuda.txt

@@ -14,8 +14,8 @@ PyNvVideoCodec==2.0.4
 # flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from
 # install_requires so the published wheel does not carry an unresolvable pin
 --extra-index-url https://flashinfer.ai/whl/
-flashinfer-python==0.6.18.post1
-flashinfer-cubin==0.6.18.post1
+flashinfer-python==0.7.0
+flashinfer-cubin==0.7.0
 apache-tvm-ffi==0.1.11
 tilelang==0.1.12
 nvidia-cudnn-frontend>=1.19.1

requirements/rubin-prerelease.txt

@@ -16,11 +16,11 @@ torchaudio==2.11.0.dev20260911+cu134
 # INSTALL_RUBIN_PRERELEASE=true, so they must be pinned here instead to
 # still get installed.
 nvidia-cutlass-dsl[cu13]==4.8.0.dev0
-flashinfer-python==0.6.18.post1
+flashinfer-python==0.7.0
 quack-kernels==0.6.5
 
 cuda-python==13.4.1
 cuda-bindings==13.4.1
-flashinfer-jit-cache==0.6.18.post1+cu134
-flashinfer-cubin==0.6.18.post1
+flashinfer-jit-cache==0.7.0+cu134
+flashinfer-cubin==0.7.0
 nixl==1.4.1

tests/kernels/attention/test_flashmla_sparse.py

@@ -509,3 +509,59 @@ def test_flashinfer_mixed_sparse_indices_separates_window_and_padded_width():
     assert sparse_indices.shape == (1, padded_width)
     assert sparse_indices[0].cpu().tolist() == [0, 1, 2, 3, -1, -1, -1, -1]
     assert sparse_lens.cpu().tolist() == [padded_width]
+
+
+@pytest.mark.parametrize("model", ["deepseek_v4", "deepseek_v41"])
+@pytest.mark.parametrize("page_padding", [0, 64])
+def test_flashinfer_sparse_forward_reads_packed_kv_rows(model, page_padding):
+    """Decode and prefill must read storage rows after vLLM remaps packed pages."""
+    import importlib
+
+    from vllm.platforms import current_platform
+    from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
+
+    if not current_platform.is_device_capability_family(100):
+        pytest.skip("Requires FlashInfer TRTLLM sparse MLA on SM100/SM103")
+    mod = importlib.import_module(f"vllm.models.{model}.nvidia.flashinfer_sparse")
+    device = "cuda"
+    block_size = 64
+    storage = torch.full(
+        (3, block_size + page_padding, 512), 99, dtype=torch.bfloat16, device=device
+    )
+    cache = storage[:, :block_size]
+    cache[1].fill_(2)
+    cache[2].fill_(3)
+    attn = object.__new__(mod.DeepseekV4FlashInferMLAAttention)
+    attn.compress_ratio = 1
+    attn.window_size = 128
+    attn.scale = 512**-0.5
+    attn.kv_cache_torch_dtype = torch.bfloat16
+    attn.attn_sink = torch.full((64,), -float("inf"), device=device)
+    attn.topk_indices_buffer = torch.empty((3, 0), dtype=torch.int32, device=device)
+    decode_indices = torch.full((1, 128), -1, dtype=torch.int32, device=device)
+    decode_indices[0, :4] = torch.arange(64, 68, dtype=torch.int32, device=device)
+    query_start = torch.tensor([0, 1, 3], dtype=torch.int32)
+    metadata = DeepseekSparseSWAMetadata(
+        block_table=torch.tensor([[1], [2]], dtype=torch.int32, device=device),
+        slot_mapping=torch.tensor([67, 130, 131], device=device),
+        block_size=block_size,
+        seq_lens=torch.tensor([4, 4], dtype=torch.int32, device=device),
+        query_start_loc=query_start.to(device),
+        query_start_loc_cpu=query_start,
+        token_to_req_indices=torch.tensor([0, 1, 1], dtype=torch.int32, device=device),
+        decode_swa_indices=decode_indices,
+        decode_swa_lens=torch.tensor([4], dtype=torch.int32, device=device),
+        decode_swa_width=128,
+        is_valid_token=torch.ones(3, dtype=torch.bool, device=device),
+        replay_start=torch.zeros(2, dtype=torch.int32, device=device),
+        num_decodes=1,
+        num_prefills=1,
+        num_decode_tokens=1,
+        num_prefill_tokens=2,
+        max_decode_query_len=1,
+    )
+    query = torch.zeros((3, 64, 512), dtype=torch.bfloat16, device=device)
+    output = torch.empty_like(query)
+    attn._forward(query, None, cache, metadata, None, True, output)
+    expected = torch.tensor([2, 3, 3], dtype=output.dtype, device=device)
+    torch.testing.assert_close(output, expected[:, None, None].expand_as(output))

tests/samplers/conftest.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm import envs
+from vllm.platforms import current_platform
+
+
+@pytest.fixture(autouse=True)
+def skip_unsupported_flashinfer_sampler():
+    # CI runs this suite with FlashInfer both explicitly enabled and disabled.
+    if (
+        current_platform.is_cuda()
+        and envs.is_set("VLLM_USE_FLASHINFER_SAMPLER")
+        and envs.VLLM_USE_FLASHINFER_SAMPLER
+        and current_platform.num_compute_units(torch.accelerator.current_device_index())
+        <= 16
+    ):
+        pytest.skip("FlashInfer top-k masking requires more than 16 SMs")

tests/v1/attention/test_flashinfer_mla_sparse_sm90.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""CPU tests for the FlashInfer SM90 sparse MLA backend wiring (no GPU).
+"""Tests for the FlashInfer SM90 sparse MLA backend wiring and index packing.
 
 The FlashInfer wrapper and top-k conversion are replaced by CPU recorders;
 the tests pin the contract between the impl and the kernel API: page_size=1
@@ -65,12 +65,18 @@ class FakeState:
     def __init__(self, width, max_tokens=64):
         self.kv_indices = torch.zeros(max_tokens * width, dtype=torch.int32)
         self.kv_len_arr = torch.zeros(max_tokens, dtype=torch.int32)
+        self.kv_indptr = torch.zeros(max_tokens + 1, dtype=torch.int32)
         self.wrapper = FakeWrapper()
         self.plan_calls = []
 
     def plan(self, num_tokens, kv_lens):
         self.plan_calls.append((num_tokens, kv_lens))
 
+    def pack_indices(self, slots):
+        for row in range(slots.shape[0]):
+            start, end = self.kv_indptr[row : row + 2].tolist()
+            self.kv_indices[start:end] = slots[row, : end - start].clamp(min=0)
+
 
 def make_impl(qk_rope, kv_dtype="fp8_e4m3", num_heads=2, topk_width=TOPK):
     impl = object.__new__(FlashInferMLASparseSM90Impl)
@@ -113,6 +119,9 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype):
         [6, 5, 4] + [-1] * (TOPK - 3),
         [2] + [-1] * (TOPK - 1),
     ]
+    impl.topk_indices_buffer.copy_(torch.tensor(topk_rows, dtype=torch.int32))
+    state.kv_indptr[1:5] = torch.tensor([8, 13, 16, 17], dtype=torch.int32)
+    state.kv_indptr[5:] = 17
     meta = make_batch(rows, topk_rows, [3])
     meta.state = state
     q_nope = torch.randn(rows, impl.num_heads, HEAD)
@@ -133,11 +142,13 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype):
     ref_slots, ref_counts = ref_convert(
         meta.req_id_per_token, meta.block_table, impl.topk_indices_buffer
     )
-    width = TOPK
-    got_slots = state.kv_indices[: rows * width].view(rows, width)
+    offset = 0
     for t in range(rows):
         k = int(ref_counts[t])
-        assert got_slots[t, :k].tolist() == ref_slots[t, :k].tolist()
+        assert (
+            state.kv_indices[offset : offset + k].tolist() == ref_slots[t, :k].tolist()
+        )
+        offset += k
     assert state.plan_calls == []
 
     assert state.wrapper.run_args is not None
@@ -221,13 +232,51 @@ def test_plan_uses_state_params(monkeypatch):
     args, kwargs = wrapper.plan_args
     (qo, kv, indices, kv_len, heads, ckv, kpe, page, causal, scale) = args
     assert qo.tolist() == [0, 1, 2, 3, 3]  # clamp: rows past 3 have no queries
-    assert kv.tolist() == [i * TOPK for i in (0, 1, 2, 3, 3)]
-    assert kv_len.tolist() == [2, 5, 7, TOPK]  # padded row keeps full width
+    assert kv.tolist() == [0, 2, 7, 14, 14]
+    assert kv_len.tolist() == [2, 5, 7, 0]
     assert (heads, ckv, kpe, page, causal) == (4, HEAD, 64, 1, False)
     assert scale == 576**-0.5
     assert kwargs["q_data_type"] == torch.bfloat16
     assert kwargs["kv_data_type"] == torch.bfloat16
 
+    # Replanning a smaller batch must clear the previous rows' lengths.
+    state.plan(1, torch.tensor([TOPK], dtype=torch.int32))
+    assert state._kv_cpu.tolist() == [0, TOPK, TOPK, TOPK, TOPK]
+    assert state._lens_cpu.tolist() == [TOPK, 0, 0, 0]
+    state.plan(0, torch.empty(0, dtype=torch.int32))
+    assert state._kv_cpu.tolist() == [0, 0, 0, 0, 0]
+    assert state._lens_cpu.tolist() == [0, 0, 0, 0]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+@pytest.mark.parametrize("width", [128, 2048, 2176])
+def test_pack_indices_replays_with_updated_offsets(width):
+    """Graph replay packs exact prefixes after row lengths and slots change."""
+    state = sm90_mod._SM90State.__new__(sm90_mod._SM90State)
+    state.kv_indptr = torch.zeros(5, dtype=torch.int32, device="cuda")
+    state.kv_indices = torch.full((4 * width,), -99, dtype=torch.int32, device="cuda")
+    # Exercise a contiguous view with an unaligned starting address.
+    slots = torch.arange(4 * width + 1, dtype=torch.int32, device="cuda")[1:]
+    slots = slots.view(4, width)
+    state.kv_indptr.copy_(torch.tensor([0, 1, 4, 4, 4]))
+    state.pack_indices(slots)
+    torch.accelerator.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        state.pack_indices(slots)
+
+    for lengths in ([1, 3, 0, 0], [width, 0, 17, width - 1], [0, 0, 0, 0]):
+        offsets = torch.tensor([0, *torch.tensor(lengths).cumsum(0).tolist()])
+        state.kv_indptr.copy_(offsets)
+        state.kv_indices.fill_(-99)
+        slots.add_(1)
+        graph.replay()
+        expected = torch.cat(
+            [slots[row, :length] for row, length in enumerate(lengths)]
+        )
+        torch.testing.assert_close(state.kv_indices[: expected.numel()], expected)
+        assert (state.kv_indices[expected.numel() :] == -99).all()
+
 
 def test_kv_lens_host_formula():
     """Per-row host lengths: context == position + 1; capped at

tests/v1/sample/test_topk_topp_sampler.py

@@ -25,7 +25,7 @@ def _flashinfer_topk_topp_supported() -> bool:
 
     Mirrors the gate in `TopKTopPSampler.__init__`: CUDA + flashinfer
     importable + GPU compute capability supported by the FlashInfer
-    backend.
+    backend + more than 16 SMs.
     """
     if not current_platform.is_cuda():
         return False
@@ -38,7 +38,10 @@ def _flashinfer_topk_topp_supported() -> bool:
     capability = current_platform.get_device_capability()
     if capability is None:
         return False
-    return FlashInferBackend.supports_compute_capability(capability)
+    return FlashInferBackend.supports_compute_capability(capability) and (
+        current_platform.num_compute_units(torch.accelerator.current_device_index())
+        > 16
+    )
 
 
 FLASHINFER_TOPK_TOPP_SUPPORTED = _flashinfer_topk_topp_supported()

vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py

@@ -227,6 +227,7 @@ def fi_chunk_gated_delta_rule(
         initial_state=fi_state,
         output_final_state=output_final_state,
         cu_seqlens=cu_seqlens,
+        backend="flashinfer",
     )
     # FlashInfer returns (output, state) when output_final_state=True,
     # or just output when output_final_state=False.

vllm/models/deepseek_v4/common/ops/cache_utils.py

@@ -1234,7 +1234,7 @@ def build_flashinfer_mixed_sparse_indices(
 def _remap_flashinfer_index(values, block_size, block_span):
     # FlashInfer's DSv4 kernel indexes sparse KV by physical token stride, so
     # packed pages (#44577) need block*block_size+off -> block*block_span+off.
-    # TODO: remove once flashinfer-ai/flashinfer#3856 is fixed.
+    # The launcher declares these indices with sparse_indices_are_storage_offsets=True.
     is_valid = values >= 0
     safe_values = tl.where(is_valid, values, 0)
     values = (safe_values // block_size) * block_span

vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py

@@ -521,6 +521,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
                 swa_kv_cache=swa_k_cache,
                 workspace_buffer=workspace,
                 sparse_indices=sparse_indices[:num_decode_tokens],
+                sparse_indices_are_storage_offsets=True,
                 compressed_kv_cache=compressed_kv_cache,
                 sparse_topk_lens=sparse_topk_lens[:num_decode_tokens],
                 seq_lens=seq_lens[:num_decodes],
@@ -546,6 +547,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
                 swa_kv_cache=swa_k_cache,
                 workspace_buffer=workspace,
                 sparse_indices=sparse_indices[num_decode_tokens:num_tokens],
+                sparse_indices_are_storage_offsets=True,
                 compressed_kv_cache=compressed_kv_cache,
                 sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens],
                 seq_lens=seq_lens[num_decodes:num_reqs],

vllm/models/deepseek_v41/common/ops/cache_utils.py

@@ -1266,7 +1266,7 @@ def build_flashinfer_mixed_sparse_indices(
 def _remap_flashinfer_index(values, block_size, block_span):
     # FlashInfer's DSv4 kernel indexes sparse KV by physical token stride, so
     # packed pages (#44577) need block*block_size+off -> block*block_span+off.
-    # TODO: remove once flashinfer-ai/flashinfer#3856 is fixed.
+    # The launcher declares these indices with sparse_indices_are_storage_offsets=True.
     is_valid = values >= 0
     safe_values = tl.where(is_valid, values, 0)
     values = (safe_values // block_size) * block_span

vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py

@@ -560,6 +560,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
                 swa_kv_cache=swa_k_cache,
                 workspace_buffer=workspace,
                 sparse_indices=sparse_indices[:num_decode_tokens],
+                sparse_indices_are_storage_offsets=True,
                 compressed_kv_cache=compressed_kv_cache,
                 sparse_topk_lens=decode_topk_lens,
                 seq_lens=decode_seq_lens,
@@ -585,6 +586,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
                 swa_kv_cache=swa_k_cache,
                 workspace_buffer=workspace,
                 sparse_indices=sparse_indices[num_decode_tokens:num_tokens],
+                sparse_indices_are_storage_offsets=True,
                 compressed_kv_cache=compressed_kv_cache,
                 sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens],
                 seq_lens=seq_lens[num_decodes:num_reqs],

vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py

@@ -43,6 +43,7 @@ from vllm.model_executor.layers.attention.sparse_mla_attention import (
     SparseMLACommonImpl,
 )
 from vllm.platforms.interface import DeviceCapability
+from vllm.triton_utils import tl, triton
 from vllm.utils.flashinfer import has_flashinfer_sm90_nope_mla
 from vllm.v1.attention.backend import (
     AttentionBackend,
@@ -64,6 +65,23 @@ _FP8_KV_DTYPES = ("fp8", "fp8_e4m3")
 _WORKSPACE_BYTES = 128 * 1024 * 1024
 
 
+@triton.jit
+def _pack_topk_indices(
+    slots,
+    indptr,
+    indices,
+    WIDTH: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    start = tl.load(indptr + row)
+    end = tl.load(indptr + row + 1)
+    mask = (cols < end - start) & (cols < WIDTH)
+    values = tl.load(slots + row * WIDTH + cols, mask=mask, other=0)
+    tl.store(indices + start + cols, tl.maximum(values, 0), mask=mask)
+
+
 class FlashInferMLASparseSM90Backend(AttentionBackend):
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
     supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
@@ -192,10 +210,11 @@ class _SM90State:
         self.kv_len_arr = torch.full(
             (max_tokens,), topk_width, dtype=torch.int32, device=device
         )
+        self.kv_indptr = torch.zeros(max_tokens + 1, dtype=torch.int32, device=device)
         self.wrapper = BatchMLAPagedAttentionWrapper(
             self.workspace,
             qo_indptr=torch.zeros(max_tokens + 1, dtype=torch.int32, device=device),
-            kv_indptr=torch.zeros(max_tokens + 1, dtype=torch.int32, device=device),
+            kv_indptr=self.kv_indptr,
             kv_indices=self.kv_indices,
             kv_len_arr=self.kv_len_arr,
             use_cuda_graph=True,
@@ -234,12 +253,13 @@ class _SM90State:
         # fixed (max_tokens+1)-sized buffers with exact-size copy_, so the
         # indptr must always be full-size. Rows past num_tokens are padded
         # empty (qo_indptr flat at num_tokens) — zero-query rows read no q
-        # and schedule no work. Padded rows keep the full width lens; the
-        # value is never dereferenced.
+        # and schedule no work. FlashInfer 0.7 requires packed CSR offsets
+        # matching the exact lengths, including zero-length padded rows.
         torch.clamp(self._arange_cpu, max=num_tokens, out=self._qo_cpu)
-        torch.mul(self._qo_cpu, self.topk_width, out=self._kv_cpu)
-        self._lens_cpu.fill_(self.topk_width)
+        self._lens_cpu.zero_()
         self._lens_cpu[:num_tokens] = kv_lens.to(torch.int32)
+        self._kv_cpu[0] = 0
+        torch.cumsum(self._lens_cpu, dim=0, out=self._kv_cpu[1:])
         self.wrapper.plan(
             self._qo_cpu,
             self._kv_cpu,
@@ -255,6 +275,12 @@ class _SM90State:
             kv_data_type=self.kv_dtype,
         )
 
+    def pack_indices(self, topk_slots: torch.Tensor) -> None:
+        num_tokens, width = topk_slots.shape
+        _pack_topk_indices[(num_tokens, triton.cdiv(width, 256))](
+            topk_slots, self.kv_indptr, self.kv_indices, width, 256
+        )
+
 
 @dataclass
 class FlashInferMLASparseSM90Metadata(FlashInferMLASparseMetadata):
@@ -462,12 +488,8 @@ class FlashInferMLASparseSM90Impl(SparseMLACommonImpl[FlashInferMLASparseSM90Met
         )
         state = attn_metadata.state
         assert state is not None
-        # Refresh top-k rows in graph and clamp masked tails to a valid slot;
-        # per-row lengths are already baked into the host-side plan.
-        width = topk_slots.shape[1]
-        state.kv_indices[: num_tokens * width].copy_(
-            topk_slots.reshape(-1).clamp_(min=0).to(torch.int32)
-        )
+        # Pack valid prefixes at the offsets refreshed by plan() before replay.
+        state.pack_indices(topk_slots)
 
         flat = (
             kv_c_and_k_pe_cache.view(torch.float8_e4m3fn)

vllm/v1/sample/ops/topk_topp_sampler.py

@@ -79,9 +79,9 @@ def flashinfer_sampler_supported() -> bool:
 
     Returns False (with appropriate logging) when ``VLLM_USE_FLASHINFER_SAMPLER``
     is 0, when the platform isn't CUDA, when the GPU's compute capability is
-    unsupported, or when FlashInfer cannot JIT-compile for the current
-    GPU/CUDA toolchain. Raises ``RuntimeError`` if the user explicitly opted in
-    via the env var but FlashInfer is unavailable.
+    unsupported, when the GPU has 16 or fewer SMs, or when FlashInfer cannot
+    JIT-compile for the current GPU/CUDA toolchain. Raises ``RuntimeError`` if
+    the user explicitly opted in via the env var but FlashInfer is unavailable.
 
     Assumes flashinfer is installed, as guaranteed by ``requirements/cuda.txt``;
     otherwise importing the FlashInfer backend below raises ``ImportError``.
@@ -107,6 +107,16 @@ def flashinfer_sampler_supported() -> bool:
         unsupported_reason = (
             f"unsupported compute capability {capability.as_version_str()}"
         )
+    elif (
+        num_sms := current_platform.num_compute_units(
+            torch.accelerator.current_device_index()
+        )
+    ) <= 16:
+        # FlashInfer 0.7+ rejects multi-CTA top-k masking on <=16 SMs because
+        # its cross-CTA software barrier cannot guarantee forward progress.
+        unsupported_reason = (
+            f"top-k masking requires more than 16 SMs; device has {num_sms}"
+        )
     else:
         unsupported_reason = _flashinfer_jit_unsupported_reason(capability)