[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>
Wei Zhao committed
Sep 24, 2026 at 01:45 UTC
44bf011f447b15ae762dbbf4d9dba2752be1ca8b
15 files changed
+197
-31
docker/Dockerfile
+1
-1
@@ -999,7 +999,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
999
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
1000
# https://docs.flashinfer.ai/installation.html
1001
# From versions.json: .flashinfer.version
1002
-ARG FLASHINFER_VERSION=0.6.18.post1
1002
+ARG FLASHINFER_VERSION=0.7.0
1003
RUN --mount=type=cache,target=/opt/uv/cache \
1004
if [ "${INSTALL_RUBIN_PRERELEASE}" != "true" ]; then \
1005
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
docker/versions.json
+1
-1
@@ -71,7 +71,7 @@
71
"default": "true"
72
},
73
"FLASHINFER_VERSION": {
74
- "default": "0.6.18.post1"
74
+ "default": "0.7.0"
75
},
76
"GDRCOPY_CUDA_VERSION": {
77
"default": "12.8"
requirements/cuda.txt
+2
-2
@@ -14,8 +14,8 @@ PyNvVideoCodec==2.0.4
14
# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from
15
# install_requires so the published wheel does not carry an unresolvable pin
16
--extra-index-url https://flashinfer.ai/whl/
17
-flashinfer-python==0.6.18.post1
18
-flashinfer-cubin==0.6.18.post1
17
+flashinfer-python==0.7.0
18
+flashinfer-cubin==0.7.0
19
apache-tvm-ffi==0.1.11
20
tilelang==0.1.12
21
nvidia-cudnn-frontend>=1.19.1
requirements/rubin-prerelease.txt
+3
-3
@@ -16,11 +16,11 @@ torchaudio==2.11.0.dev20260911+cu134
16
# INSTALL_RUBIN_PRERELEASE=true, so they must be pinned here instead to
17
# still get installed.
18
nvidia-cutlass-dsl[cu13]==4.8.0.dev0
19
-flashinfer-python==0.6.18.post1
19
+flashinfer-python==0.7.0
20
quack-kernels==0.6.5
21
22
cuda-python==13.4.1
23
cuda-bindings==13.4.1
24
-flashinfer-jit-cache==0.6.18.post1+cu134
25
-flashinfer-cubin==0.6.18.post1
24
+flashinfer-jit-cache==0.7.0+cu134
25
+flashinfer-cubin==0.7.0
26
nixl==1.4.1
tests/kernels/attention/test_flashmla_sparse.py
+56
@@ -509,3 +509,59 @@ def test_flashinfer_mixed_sparse_indices_separates_window_and_padded_width():
509
assert sparse_indices.shape == (1, padded_width)
510
assert sparse_indices[0].cpu().tolist() == [0, 1, 2, 3, -1, -1, -1, -1]
511
assert sparse_lens.cpu().tolist() == [padded_width]
512
+
513
+
514
+@pytest.mark.parametrize("model", ["deepseek_v4", "deepseek_v41"])
515
+@pytest.mark.parametrize("page_padding", [0, 64])
516
+def test_flashinfer_sparse_forward_reads_packed_kv_rows(model, page_padding):
517
+ """Decode and prefill must read storage rows after vLLM remaps packed pages."""
518
+ import importlib
519
+
520
+ from vllm.platforms import current_platform
521
+ from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
522
+
523
+ if not current_platform.is_device_capability_family(100):
524
+ pytest.skip("Requires FlashInfer TRTLLM sparse MLA on SM100/SM103")
525
+ mod = importlib.import_module(f"vllm.models.{model}.nvidia.flashinfer_sparse")
526
+ device = "cuda"
527
+ block_size = 64
528
+ storage = torch.full(
529
+ (3, block_size + page_padding, 512), 99, dtype=torch.bfloat16, device=device
530
+ )
531
+ cache = storage[:, :block_size]
532
+ cache[1].fill_(2)
533
+ cache[2].fill_(3)
534
+ attn = object.__new__(mod.DeepseekV4FlashInferMLAAttention)
535
+ attn.compress_ratio = 1
536
+ attn.window_size = 128
537
+ attn.scale = 512**-0.5
538
+ attn.kv_cache_torch_dtype = torch.bfloat16
539
+ attn.attn_sink = torch.full((64,), -float("inf"), device=device)
540
+ attn.topk_indices_buffer = torch.empty((3, 0), dtype=torch.int32, device=device)
541
+ decode_indices = torch.full((1, 128), -1, dtype=torch.int32, device=device)
542
+ decode_indices[0, :4] = torch.arange(64, 68, dtype=torch.int32, device=device)
543
+ query_start = torch.tensor([0, 1, 3], dtype=torch.int32)
544
+ metadata = DeepseekSparseSWAMetadata(
545
+ block_table=torch.tensor([[1], [2]], dtype=torch.int32, device=device),
546
+ slot_mapping=torch.tensor([67, 130, 131], device=device),
547
+ block_size=block_size,
548
+ seq_lens=torch.tensor([4, 4], dtype=torch.int32, device=device),
549
+ query_start_loc=query_start.to(device),
550
+ query_start_loc_cpu=query_start,
551
+ token_to_req_indices=torch.tensor([0, 1, 1], dtype=torch.int32, device=device),
552
+ decode_swa_indices=decode_indices,
553
+ decode_swa_lens=torch.tensor([4], dtype=torch.int32, device=device),
554
+ decode_swa_width=128,
555
+ is_valid_token=torch.ones(3, dtype=torch.bool, device=device),
556
+ replay_start=torch.zeros(2, dtype=torch.int32, device=device),
557
+ num_decodes=1,
558
+ num_prefills=1,
559
+ num_decode_tokens=1,
560
+ num_prefill_tokens=2,
561
+ max_decode_query_len=1,
562
+ )
563
+ query = torch.zeros((3, 64, 512), dtype=torch.bfloat16, device=device)
564
+ output = torch.empty_like(query)
565
+ attn._forward(query, None, cache, metadata, None, True, output)
566
+ expected = torch.tensor([2, 3, 3], dtype=output.dtype, device=device)
567
+ torch.testing.assert_close(output, expected[:, None, None].expand_as(output))
tests/samplers/conftest.py
new
+21
@@ -0,0 +1,21 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm import envs
8
+from vllm.platforms import current_platform
9
+
10
+
11
+@pytest.fixture(autouse=True)
12
+def skip_unsupported_flashinfer_sampler():
13
+ # CI runs this suite with FlashInfer both explicitly enabled and disabled.
14
+ if (
15
+ current_platform.is_cuda()
16
+ and envs.is_set("VLLM_USE_FLASHINFER_SAMPLER")
17
+ and envs.VLLM_USE_FLASHINFER_SAMPLER
18
+ and current_platform.num_compute_units(torch.accelerator.current_device_index())
19
+ <= 16
20
+ ):
21
+ pytest.skip("FlashInfer top-k masking requires more than 16 SMs")
tests/v1/attention/test_flashinfer_mla_sparse_sm90.py
+55
-6
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""CPU tests for the FlashInfer SM90 sparse MLA backend wiring (no GPU).
3
+"""Tests for the FlashInfer SM90 sparse MLA backend wiring and index packing.
4
5
The FlashInfer wrapper and top-k conversion are replaced by CPU recorders;
6
the tests pin the contract between the impl and the kernel API: page_size=1
@@ -65,12 +65,18 @@ class FakeState:
65
def __init__(self, width, max_tokens=64):
66
self.kv_indices = torch.zeros(max_tokens * width, dtype=torch.int32)
67
self.kv_len_arr = torch.zeros(max_tokens, dtype=torch.int32)
68
+ self.kv_indptr = torch.zeros(max_tokens + 1, dtype=torch.int32)
69
self.wrapper = FakeWrapper()
70
self.plan_calls = []
71
72
def plan(self, num_tokens, kv_lens):
73
self.plan_calls.append((num_tokens, kv_lens))
74
75
+ def pack_indices(self, slots):
76
+ for row in range(slots.shape[0]):
77
+ start, end = self.kv_indptr[row : row + 2].tolist()
78
+ self.kv_indices[start:end] = slots[row, : end - start].clamp(min=0)
79
+
80
81
def make_impl(qk_rope, kv_dtype="fp8_e4m3", num_heads=2, topk_width=TOPK):
82
impl = object.__new__(FlashInferMLASparseSM90Impl)
@@ -113,6 +119,9 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype):
119
[6, 5, 4] + [-1] * (TOPK - 3),
120
[2] + [-1] * (TOPK - 1),
121
]
122
+ impl.topk_indices_buffer.copy_(torch.tensor(topk_rows, dtype=torch.int32))
123
+ state.kv_indptr[1:5] = torch.tensor([8, 13, 16, 17], dtype=torch.int32)
124
+ state.kv_indptr[5:] = 17
125
meta = make_batch(rows, topk_rows, [3])
126
meta.state = state
127
q_nope = torch.randn(rows, impl.num_heads, HEAD)
@@ -133,11 +142,13 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype):
142
ref_slots, ref_counts = ref_convert(
143
meta.req_id_per_token, meta.block_table, impl.topk_indices_buffer
144
)
136
- width = TOPK
137
- got_slots = state.kv_indices[: rows * width].view(rows, width)
145
+ offset = 0
146
for t in range(rows):
147
k = int(ref_counts[t])
140
- assert got_slots[t, :k].tolist() == ref_slots[t, :k].tolist()
148
+ assert (
149
+ state.kv_indices[offset : offset + k].tolist() == ref_slots[t, :k].tolist()
150
+ )
151
+ offset += k
152
assert state.plan_calls == []
153
154
assert state.wrapper.run_args is not None
@@ -221,13 +232,51 @@ def test_plan_uses_state_params(monkeypatch):
232
args, kwargs = wrapper.plan_args
233
(qo, kv, indices, kv_len, heads, ckv, kpe, page, causal, scale) = args
234
assert qo.tolist() == [0, 1, 2, 3, 3] # clamp: rows past 3 have no queries
224
- assert kv.tolist() == [i * TOPK for i in (0, 1, 2, 3, 3)]
225
- assert kv_len.tolist() == [2, 5, 7, TOPK] # padded row keeps full width
235
+ assert kv.tolist() == [0, 2, 7, 14, 14]
236
+ assert kv_len.tolist() == [2, 5, 7, 0]
237
assert (heads, ckv, kpe, page, causal) == (4, HEAD, 64, 1, False)
238
assert scale == 576**-0.5
239
assert kwargs["q_data_type"] == torch.bfloat16
240
assert kwargs["kv_data_type"] == torch.bfloat16
241
242
+ # Replanning a smaller batch must clear the previous rows' lengths.
243
+ state.plan(1, torch.tensor([TOPK], dtype=torch.int32))
244
+ assert state._kv_cpu.tolist() == [0, TOPK, TOPK, TOPK, TOPK]
245
+ assert state._lens_cpu.tolist() == [TOPK, 0, 0, 0]
246
+ state.plan(0, torch.empty(0, dtype=torch.int32))
247
+ assert state._kv_cpu.tolist() == [0, 0, 0, 0, 0]
248
+ assert state._lens_cpu.tolist() == [0, 0, 0, 0]
249
+
250
+
251
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
252
+@pytest.mark.parametrize("width", [128, 2048, 2176])
253
+def test_pack_indices_replays_with_updated_offsets(width):
254
+ """Graph replay packs exact prefixes after row lengths and slots change."""
255
+ state = sm90_mod._SM90State.__new__(sm90_mod._SM90State)
256
+ state.kv_indptr = torch.zeros(5, dtype=torch.int32, device="cuda")
257
+ state.kv_indices = torch.full((4 * width,), -99, dtype=torch.int32, device="cuda")
258
+ # Exercise a contiguous view with an unaligned starting address.
259
+ slots = torch.arange(4 * width + 1, dtype=torch.int32, device="cuda")[1:]
260
+ slots = slots.view(4, width)
261
+ state.kv_indptr.copy_(torch.tensor([0, 1, 4, 4, 4]))
262
+ state.pack_indices(slots)
263
+ torch.accelerator.synchronize()
264
+ graph = torch.cuda.CUDAGraph()
265
+ with torch.cuda.graph(graph):
266
+ state.pack_indices(slots)
267
+
268
+ for lengths in ([1, 3, 0, 0], [width, 0, 17, width - 1], [0, 0, 0, 0]):
269
+ offsets = torch.tensor([0, *torch.tensor(lengths).cumsum(0).tolist()])
270
+ state.kv_indptr.copy_(offsets)
271
+ state.kv_indices.fill_(-99)
272
+ slots.add_(1)
273
+ graph.replay()
274
+ expected = torch.cat(
275
+ [slots[row, :length] for row, length in enumerate(lengths)]
276
+ )
277
+ torch.testing.assert_close(state.kv_indices[: expected.numel()], expected)
278
+ assert (state.kv_indices[expected.numel() :] == -99).all()
279
+
280
281
def test_kv_lens_host_formula():
282
"""Per-row host lengths: context == position + 1; capped at
tests/v1/sample/test_topk_topp_sampler.py
+5
-2
@@ -25,7 +25,7 @@ def _flashinfer_topk_topp_supported() -> bool:
25
26
Mirrors the gate in `TopKTopPSampler.__init__`: CUDA + flashinfer
27
importable + GPU compute capability supported by the FlashInfer
28
- backend.
28
+ backend + more than 16 SMs.
29
"""
30
if not current_platform.is_cuda():
31
return False
@@ -38,7 +38,10 @@ def _flashinfer_topk_topp_supported() -> bool:
38
capability = current_platform.get_device_capability()
39
if capability is None:
40
return False
41
- return FlashInferBackend.supports_compute_capability(capability)
41
+ return FlashInferBackend.supports_compute_capability(capability) and (
42
+ current_platform.num_compute_units(torch.accelerator.current_device_index())
43
+ > 16
44
+ )
45
46
47
FLASHINFER_TOPK_TOPP_SUPPORTED = _flashinfer_topk_topp_supported()
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+1
@@ -227,6 +227,7 @@ def fi_chunk_gated_delta_rule(
227
initial_state=fi_state,
228
output_final_state=output_final_state,
229
cu_seqlens=cu_seqlens,
230
+ backend="flashinfer",
231
)
232
# FlashInfer returns (output, state) when output_final_state=True,
233
# or just output when output_final_state=False.
vllm/models/deepseek_v4/common/ops/cache_utils.py
+1
-1
@@ -1234,7 +1234,7 @@ def build_flashinfer_mixed_sparse_indices(
1234
def _remap_flashinfer_index(values, block_size, block_span):
1235
# FlashInfer's DSv4 kernel indexes sparse KV by physical token stride, so
1236
# packed pages (#44577) need block*block_size+off -> block*block_span+off.
1237
- # TODO: remove once flashinfer-ai/flashinfer#3856 is fixed.
1237
+ # The launcher declares these indices with sparse_indices_are_storage_offsets=True.
1238
is_valid = values >= 0
1239
safe_values = tl.where(is_valid, values, 0)
1240
values = (safe_values // block_size) * block_span
vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
+2
@@ -521,6 +521,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
521
swa_kv_cache=swa_k_cache,
522
workspace_buffer=workspace,
523
sparse_indices=sparse_indices[:num_decode_tokens],
524
+ sparse_indices_are_storage_offsets=True,
525
compressed_kv_cache=compressed_kv_cache,
526
sparse_topk_lens=sparse_topk_lens[:num_decode_tokens],
527
seq_lens=seq_lens[:num_decodes],
@@ -546,6 +547,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
547
swa_kv_cache=swa_k_cache,
548
workspace_buffer=workspace,
549
sparse_indices=sparse_indices[num_decode_tokens:num_tokens],
550
+ sparse_indices_are_storage_offsets=True,
551
compressed_kv_cache=compressed_kv_cache,
552
sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens],
553
seq_lens=seq_lens[num_decodes:num_reqs],
vllm/models/deepseek_v41/common/ops/cache_utils.py
+1
-1
@@ -1266,7 +1266,7 @@ def build_flashinfer_mixed_sparse_indices(
1266
def _remap_flashinfer_index(values, block_size, block_span):
1267
# FlashInfer's DSv4 kernel indexes sparse KV by physical token stride, so
1268
# packed pages (#44577) need block*block_size+off -> block*block_span+off.
1269
- # TODO: remove once flashinfer-ai/flashinfer#3856 is fixed.
1269
+ # The launcher declares these indices with sparse_indices_are_storage_offsets=True.
1270
is_valid = values >= 0
1271
safe_values = tl.where(is_valid, values, 0)
1272
values = (safe_values // block_size) * block_span
vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py
+2
@@ -560,6 +560,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
560
swa_kv_cache=swa_k_cache,
561
workspace_buffer=workspace,
562
sparse_indices=sparse_indices[:num_decode_tokens],
563
+ sparse_indices_are_storage_offsets=True,
564
compressed_kv_cache=compressed_kv_cache,
565
sparse_topk_lens=decode_topk_lens,
566
seq_lens=decode_seq_lens,
@@ -585,6 +586,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
586
swa_kv_cache=swa_k_cache,
587
workspace_buffer=workspace,
588
sparse_indices=sparse_indices[num_decode_tokens:num_tokens],
589
+ sparse_indices_are_storage_offsets=True,
590
compressed_kv_cache=compressed_kv_cache,
591
sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens],
592
seq_lens=seq_lens[num_decodes:num_reqs],
vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py
+33
-11
@@ -43,6 +43,7 @@ from vllm.model_executor.layers.attention.sparse_mla_attention import (
43
SparseMLACommonImpl,
44
)
45
from vllm.platforms.interface import DeviceCapability
46
+from vllm.triton_utils import tl, triton
47
from vllm.utils.flashinfer import has_flashinfer_sm90_nope_mla
48
from vllm.v1.attention.backend import (
49
AttentionBackend,
@@ -64,6 +65,23 @@ _FP8_KV_DTYPES = ("fp8", "fp8_e4m3")
65
_WORKSPACE_BYTES = 128 * 1024 * 1024
66
67
68
+@triton.jit
69
+def _pack_topk_indices(
70
+ slots,
71
+ indptr,
72
+ indices,
73
+ WIDTH: tl.constexpr,
74
+ BLOCK_SIZE: tl.constexpr,
75
+):
76
+ row = tl.program_id(0).to(tl.int64)
77
+ cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
78
+ start = tl.load(indptr + row)
79
+ end = tl.load(indptr + row + 1)
80
+ mask = (cols < end - start) & (cols < WIDTH)
81
+ values = tl.load(slots + row * WIDTH + cols, mask=mask, other=0)
82
+ tl.store(indices + start + cols, tl.maximum(values, 0), mask=mask)
83
+
84
+
85
class FlashInferMLASparseSM90Backend(AttentionBackend):
86
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
87
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
@@ -192,10 +210,11 @@ class _SM90State:
210
self.kv_len_arr = torch.full(
211
(max_tokens,), topk_width, dtype=torch.int32, device=device
212
)
213
+ self.kv_indptr = torch.zeros(max_tokens + 1, dtype=torch.int32, device=device)
214
self.wrapper = BatchMLAPagedAttentionWrapper(
215
self.workspace,
216
qo_indptr=torch.zeros(max_tokens + 1, dtype=torch.int32, device=device),
198
- kv_indptr=torch.zeros(max_tokens + 1, dtype=torch.int32, device=device),
217
+ kv_indptr=self.kv_indptr,
218
kv_indices=self.kv_indices,
219
kv_len_arr=self.kv_len_arr,
220
use_cuda_graph=True,
@@ -234,12 +253,13 @@ class _SM90State:
253
# fixed (max_tokens+1)-sized buffers with exact-size copy_, so the
254
# indptr must always be full-size. Rows past num_tokens are padded
255
# empty (qo_indptr flat at num_tokens) — zero-query rows read no q
237
- # and schedule no work. Padded rows keep the full width lens; the
238
- # value is never dereferenced.
256
+ # and schedule no work. FlashInfer 0.7 requires packed CSR offsets
257
+ # matching the exact lengths, including zero-length padded rows.
258
torch.clamp(self._arange_cpu, max=num_tokens, out=self._qo_cpu)
240
- torch.mul(self._qo_cpu, self.topk_width, out=self._kv_cpu)
241
- self._lens_cpu.fill_(self.topk_width)
259
+ self._lens_cpu.zero_()
260
self._lens_cpu[:num_tokens] = kv_lens.to(torch.int32)
261
+ self._kv_cpu[0] = 0
262
+ torch.cumsum(self._lens_cpu, dim=0, out=self._kv_cpu[1:])
263
self.wrapper.plan(
264
self._qo_cpu,
265
self._kv_cpu,
@@ -255,6 +275,12 @@ class _SM90State:
275
kv_data_type=self.kv_dtype,
276
)
277
278
+ def pack_indices(self, topk_slots: torch.Tensor) -> None:
279
+ num_tokens, width = topk_slots.shape
280
+ _pack_topk_indices[(num_tokens, triton.cdiv(width, 256))](
281
+ topk_slots, self.kv_indptr, self.kv_indices, width, 256
282
+ )
283
+
284
285
@dataclass
286
class FlashInferMLASparseSM90Metadata(FlashInferMLASparseMetadata):
@@ -462,12 +488,8 @@ class FlashInferMLASparseSM90Impl(SparseMLACommonImpl[FlashInferMLASparseSM90Met
488
)
489
state = attn_metadata.state
490
assert state is not None
465
- # Refresh top-k rows in graph and clamp masked tails to a valid slot;
466
- # per-row lengths are already baked into the host-side plan.
467
- width = topk_slots.shape[1]
468
- state.kv_indices[: num_tokens * width].copy_(
469
- topk_slots.reshape(-1).clamp_(min=0).to(torch.int32)
470
- )
491
+ # Pack valid prefixes at the offsets refreshed by plan() before replay.
492
+ state.pack_indices(topk_slots)
493
494
flat = (
495
kv_c_and_k_pe_cache.view(torch.float8_e4m3fn)
vllm/v1/sample/ops/topk_topp_sampler.py
+13
-3
@@ -79,9 +79,9 @@ def flashinfer_sampler_supported() -> bool:
79
80
Returns False (with appropriate logging) when ``VLLM_USE_FLASHINFER_SAMPLER``
81
is 0, when the platform isn't CUDA, when the GPU's compute capability is
82
- unsupported, or when FlashInfer cannot JIT-compile for the current
83
- GPU/CUDA toolchain. Raises ``RuntimeError`` if the user explicitly opted in
84
- via the env var but FlashInfer is unavailable.
82
+ unsupported, when the GPU has 16 or fewer SMs, or when FlashInfer cannot
83
+ JIT-compile for the current GPU/CUDA toolchain. Raises ``RuntimeError`` if
84
+ the user explicitly opted in via the env var but FlashInfer is unavailable.
85
86
Assumes flashinfer is installed, as guaranteed by ``requirements/cuda.txt``;
87
otherwise importing the FlashInfer backend below raises ``ImportError``.
@@ -107,6 +107,16 @@ def flashinfer_sampler_supported() -> bool:
107
unsupported_reason = (
108
f"unsupported compute capability {capability.as_version_str()}"
109
)
110
+ elif (
111
+ num_sms := current_platform.num_compute_units(
112
+ torch.accelerator.current_device_index()
113
+ )
114
+ ) <= 16:
115
+ # FlashInfer 0.7+ rejects multi-CTA top-k masking on <=16 SMs because
116
+ # its cross-CTA software barrier cannot guarantee forward progress.
117
+ unsupported_reason = (
118
+ f"top-k masking requires more than 16 SMs; device has {num_sms}"
119
+ )
120
else:
121
unsupported_reason = _flashinfer_jit_unsupported_reason(capability)
122