diff --git a/docker/Dockerfile b/docker/Dockerfile index 895968601f..26b6156881 100644 --- a/docker/Dockerfile +++ b/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} \ diff --git a/docker/versions.json b/docker/versions.json index 7a5b9f44ab..ec23b95186 100644 --- a/docker/versions.json +++ b/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" diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 59e5c4dcc8..2bbf1e244d 100644 --- a/requirements/cuda.txt +++ b/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 diff --git a/requirements/rubin-prerelease.txt b/requirements/rubin-prerelease.txt index 8b9f7d3155..d02fb38884 100644 --- a/requirements/rubin-prerelease.txt +++ b/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 diff --git a/tests/kernels/attention/test_flashmla_sparse.py b/tests/kernels/attention/test_flashmla_sparse.py index bfcd6ae10a..7c1a455094 100644 --- a/tests/kernels/attention/test_flashmla_sparse.py +++ b/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)) diff --git a/tests/samplers/conftest.py b/tests/samplers/conftest.py new file mode 100644 index 0000000000..4585e98898 --- /dev/null +++ b/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") diff --git a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py index 8dc082b08b..72b2de05c2 100644 --- a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py +++ b/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 diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 8ff16e5599..e40d3048c1 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/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() diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 094dd74e78..9d5bcc89da 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/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. diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index 12fdb04711..0cc784e348 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/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 diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index c3b7092fbe..bb9e42349a 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/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], diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index 02fb0d5539..ad03c32945 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/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 diff --git a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index 949f383324..423849dcf1 100644 --- a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py +++ b/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], diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py index d39b7343a4..27fd633601 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py +++ b/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) diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 2f051b6768..69ec368309 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/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)