[Perf][Kernel] Integrate DeepSelect TopK for the DSA sparse indexer (#56464)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: Roger Wang <hey@rogerw.io>
Co-authored-by: Kimi Code <noreply@moonshot.ai>
Co-authored-by: Roger Wang <hey@rogerw.io>

CMakeLists.txt

@@ -1572,6 +1572,7 @@ endif()
 # For CUDA we also build and ship some external projects.
 if (VLLM_GPU_LANG STREQUAL "CUDA")
     include(cmake/external_projects/deepgemm.cmake)
+    include(cmake/external_projects/deepselect.cmake)
     include(cmake/external_projects/fmha_sm100.cmake)
     include(cmake/external_projects/flashmla.cmake)
     include(cmake/external_projects/flashkda.cmake)

cmake/external_projects/deepselect.cmake

@@ -0,0 +1,97 @@
+include(FetchContent)
+
+# If DEEPSELECT_SRC_DIR is set, DeepSelect is built from that directory
+# instead of downloading.
+# It can be set as an environment variable or passed as a cmake argument.
+# The environment variable takes precedence.
+if (DEFINED ENV{DEEPSELECT_SRC_DIR})
+  set(DEEPSELECT_SRC_DIR $ENV{DEEPSELECT_SRC_DIR})
+endif()
+
+if(DEEPSELECT_SRC_DIR)
+  FetchContent_Declare(
+        deepselect
+        SOURCE_DIR ${DEEPSELECT_SRC_DIR}
+        CONFIGURE_COMMAND ""
+        BUILD_COMMAND ""
+  )
+else()
+  FetchContent_Declare(
+        deepselect
+        GIT_REPOSITORY https://github.com/vllm-project/DeepSelect.git
+        GIT_TAG d96d33afe1fab0d6066da49cdc91e64c2bee65ea
+        GIT_PROGRESS TRUE
+        CONFIGURE_COMMAND ""
+        BUILD_COMMAND ""
+  )
+endif()
+
+FetchContent_MakeAvailable(deepselect)
+message(STATUS "DeepSelect is available at ${deepselect_SOURCE_DIR}")
+
+# DeepSelect kernels require CUDA 12.9+; 10.0f (family) covers SM100 and SM103.
+set(DEEPSELECT_SUPPORT_ARCHS)
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
+    list(APPEND DEEPSELECT_SUPPORT_ARCHS "10.0f")
+endif()
+
+cuda_archs_loose_intersection(DEEPSELECT_ARCHS "${DEEPSELECT_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
+if(DEEPSELECT_ARCHS)
+    message(STATUS "DeepSelect CUDA architectures: ${DEEPSELECT_ARCHS}")
+
+    # All explicit template instantiations must be compiled: api.cpp
+    # references every dispatch combination.
+    file(GLOB DeepSelect_SOURCES CONFIGURE_DEPENDS
+        ${deepselect_SOURCE_DIR}/csrc/cuda_kernels/v3/instantiations/*.cu
+        ${deepselect_SOURCE_DIR}/csrc/cuda_kernels/v3_fp32/instantiations/*.cu
+        ${deepselect_SOURCE_DIR}/csrc/cuda_kernels/v3_cluster/instantiations/*.cu
+    )
+    list(PREPEND DeepSelect_SOURCES ${deepselect_SOURCE_DIR}/csrc/api.cpp)
+
+    set(DeepSelect_INCLUDES
+        ${deepselect_SOURCE_DIR}/csrc
+        ${deepselect_SOURCE_DIR}/csrc/3rdparty/cutlass/include
+        ${deepselect_SOURCE_DIR}/csrc/3rdparty/kerutils/include
+    )
+
+    set_gencode_flags_for_srcs(
+        SRCS "${DeepSelect_SOURCES}"
+        CUDA_ARCHS "${DEEPSELECT_ARCHS}")
+
+    set(VLLM_DEEPSELECT_GPU_FLAGS ${VLLM_GPU_FLAGS})
+    list(APPEND VLLM_DEEPSELECT_GPU_FLAGS
+        "-O3" "--expt-relaxed-constexpr" "--expt-extended-lambda"
+        "--use_fast_math" "--ftz=false"
+        "-U__CUDA_NO_HALF_OPERATORS__" "-U__CUDA_NO_HALF_CONVERSIONS__"
+        "-U__CUDA_NO_HALF2_OPERATORS__" "-U__CUDA_NO_BFLOAT16_CONVERSIONS__")
+
+    define_extension_target(
+        _deepselect_C
+        DESTINATION vllm
+        LANGUAGE ${VLLM_GPU_LANG}
+        SOURCES ${DeepSelect_SOURCES}
+        COMPILE_FLAGS ${VLLM_DEEPSELECT_GPU_FLAGS}
+        ARCHITECTURES ${VLLM_GPU_ARCHES}
+        INCLUDE_DIRECTORIES ${DeepSelect_INCLUDES}
+        USE_SABI 3
+        WITH_SOABI)
+
+    # Only use C-shim APIs available in PyTorch 2.10.
+    # _deepselect_C is abi compatible with PyTorch >= TORCH_TARGET_VERSION.
+    target_compile_definitions(_deepselect_C PRIVATE
+        TORCH_TARGET_VERSION=0x020A000000000000ULL)
+
+    # Needed to use cuda APIs from C-shim
+    if(VLLM_GPU_LANG STREQUAL "CUDA")
+        target_compile_definitions(_deepselect_C PRIVATE USE_CUDA)
+    endif()
+
+    # DeepSelect requires C++20
+    target_compile_options(_deepselect_C PRIVATE
+        $<$<COMPILE_LANGUAGE:CXX>:-std=c++20>
+        $<$<COMPILE_LANGUAGE:CUDA>:-std=c++20>)
+else()
+    message(STATUS "DeepSelect will not compile: unsupported CUDA architecture ${CUDA_ARCHS}")
+    # Create an empty target for setup.py on unsupported systems
+    add_custom_target(_deepselect_C)
+endif()

setup.py

@@ -1015,6 +1015,7 @@ class precompiled_wheel_utils:
                             "vllm/_flashmla_extension_C.abi3.so",
                             "vllm/_flashkda_C.abi3.so",
                             "vllm/_sparse_flashmla_C.abi3.so",
+                            "vllm/_deepselect_C.abi3.so",
                             "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
                             "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
                             "vllm/cumem_allocator.abi3.so",
@@ -1381,6 +1382,9 @@ if _is_cuda():
         ext_modules.append(
             CMakeExtension(name="vllm._flashmla_extension_C", optional=True)
         )
+        # DeepSelect requires CUDA 12.9 or later (SM100a/SM103a only)
+        # Optional since it won't build on unsupported architectures
+        ext_modules.append(CMakeExtension(name="vllm._deepselect_C", optional=True))
     if USE_PRECOMPILED_EXTENSIONS or (
         CUDA_HOME and get_nvcc_cuda_version() >= Version("12.0")
     ):

tests/kernels/test_top_k_per_row.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+
 import numpy as np
 import pytest
 import torch
@@ -21,6 +22,13 @@ def _has_device_capability(major: int) -> bool:
     return current_platform.is_cuda() and current_platform.has_device_capability(major)
 
 
+# DeepSelect is compiled for sm_100a/sm_103a only.
+requires_sm100 = pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="DeepSelect requires SM100a/SM103a",
+)
+
+
 COOPERATIVE_TOPK_BACKEND = pytest.param(
     "cooperative_topk",
     marks=pytest.mark.skipif(
@@ -1315,3 +1323,318 @@ def test_workspace_topk_padded_stride(top_k: int, backend: str) -> None:
                 f"Row {i}: {backend} with padded stride doesn't match. "
                 f"seq_len={sl}, stride={padded_stride}"
             )
+
+
+@requires_sm100
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("batch_size", [1, 8, 64])
+@pytest.mark.parametrize("vocab_size", [4096, 131072])
+@pytest.mark.parametrize("top_k", [512, 2048])
+@torch.inference_mode()
+def test_deep_select_topk(
+    dtype: torch.dtype,
+    batch_size: int,
+    vocab_size: int,
+    top_k: int,
+) -> None:
+    """DeepSelect wrapper vs torch.topk with variable per-row ends."""
+    from vllm.model_executor.layers import indexer_topk
+
+    set_random_seed(0)
+    torch.set_default_device("cuda:0")
+
+    row_ends = torch.randint(
+        1, vocab_size + 1, (batch_size,), dtype=torch.int32, device="cuda"
+    )
+    # Force one short row to exercise the -1 out-of-bounds fill.
+    row_ends[0] = min(top_k - 1, vocab_size)
+    logits = torch.randn(batch_size, vocab_size, dtype=dtype, device="cuda")
+    # Values beyond each row's end must never be read by the kernel: fill
+    # them with NaN (production tails are uninitialized; DeepSelect traps on
+    # NaN with abort_when_nan_found=True, so a wrong read fails loudly).
+    col_idx = torch.arange(vocab_size, device="cuda")
+    logits[col_idx[None, :] >= row_ends[:, None]] = float("nan")
+
+    indices = indexer_topk.deep_select_topk(logits, top_k, end=row_ends)
+    torch.accelerator.synchronize()
+
+    assert indices.shape == (batch_size, top_k)
+    row_starts = torch.zeros(batch_size, dtype=torch.int32, device="cuda")
+    torch_indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda")
+    for i in range(batch_size):
+        row_end = int(row_ends[i])
+        k_i = min(top_k, row_end)
+        idx = logits[i, :row_end].topk(k_i, dim=-1)[1]
+        torch_indices[i, :k_i] = idx
+        assert torch.all(indices[i, k_i:] == -1), (
+            f"Row {i}: expected -1 fill after {k_i} valid indices"
+        )
+        # All selected indices must lie inside [0, row_end).
+        assert torch.all(indices[i, :k_i] < row_end)
+
+    assert compare_top_k_results(
+        logits, indices, torch_indices, row_starts, row_ends, top_k
+    ), "DeepSelect topk results don't match torch.topk"
+
+
+@requires_sm100
+@torch.inference_mode()
+def test_deep_select_topk_preallocated_output() -> None:
+    """DeepSelect writes into a preallocated aligned buffer (no `end`)."""
+    from vllm.model_executor.layers import indexer_topk
+
+    set_random_seed(0)
+    torch.set_default_device("cuda:0")
+
+    batch_size = 8
+    vocab_size = 131072
+    top_k = 2048
+    logits = torch.randn(batch_size, vocab_size, dtype=torch.float32, device="cuda")
+
+    # Mimic vLLM's topk_indices_buffer: slice of a wider contiguous buffer,
+    # whose stride(0) is 32B-aligned.
+    buffer = torch.full((batch_size, 4096), -2, dtype=torch.int32, device="cuda")
+    output_idx = buffer[:, :top_k]
+    result = indexer_topk.deep_select_topk(logits, top_k, output_idx=output_idx)
+    torch.accelerator.synchronize()
+
+    assert result.data_ptr() == output_idx.data_ptr()
+    assert torch.all(buffer[:, top_k:] == -2), "kernel wrote out of bounds"
+
+    row_ends = torch.full((batch_size,), vocab_size, dtype=torch.int32, device="cuda")
+    row_starts = torch.zeros(batch_size, dtype=torch.int32, device="cuda")
+    torch_indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda")
+    for i in range(batch_size):
+        torch_indices[i] = logits[i].topk(top_k, dim=-1)[1]
+
+    assert compare_top_k_results(
+        logits, output_idx, torch_indices, row_starts, row_ends, top_k
+    ), "DeepSelect topk (preallocated output) results don't match torch.topk"
+
+
+def _has_flashinfer_topk() -> bool:
+    if not current_platform.is_cuda():
+        return False
+    # find_spec only: importing flashinfer initializes CUDA at import time.
+    import importlib.util
+
+    return importlib.util.find_spec("flashinfer") is not None
+
+
+def _has_cooperative_topk() -> bool:
+    return _has_device_capability(
+        90
+    ) and not current_platform.is_device_capability_family(120)
+
+
+SPARSE_INDEXER_EXPLICIT_BACKENDS = [
+    pytest.param(
+        "per_row",
+        marks=pytest.mark.skipif(
+            not current_platform.is_cuda(), reason="requires CUDA"
+        ),
+    ),
+    "torch",
+    pytest.param(
+        "persistent",
+        marks=pytest.mark.skipif(
+            not current_platform.is_cuda(), reason="requires CUDA"
+        ),
+    ),
+    pytest.param(
+        "cooperative",
+        marks=pytest.mark.skipif(
+            not _has_cooperative_topk(),
+            reason="cooperative_topk requires SM90+ (non-SM12x)",
+        ),
+    ),
+    pytest.param(
+        "flashinfer",
+        marks=pytest.mark.skipif(
+            not _has_flashinfer_topk(), reason="requires flashinfer top-k"
+        ),
+    ),
+    pytest.param("deep_select", marks=requires_sm100),
+]
+
+
+@pytest.mark.parametrize("next_n", [1, 4])
+@pytest.mark.parametrize("backend", SPARSE_INDEXER_EXPLICIT_BACKENDS)
+@torch.inference_mode()
+def test_sparse_indexer_decode_topk_explicit_backends(
+    backend: str, next_n: int, workspace_init
+) -> None:
+    """Every explicit sparse-indexer decode top-k backend must restrict each
+    row to its seq_len (dirty data past the end must never be selected) and
+    match torch.topk on the valid region."""
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk
+
+    set_random_seed(0)
+    torch.set_default_device("cuda:0")
+
+    top_k = 2048
+    batch_size = 4
+    num_rows = batch_size * next_n
+
+    # Per-row effective lens in the (B, next_n) "native spec decode" form;
+    # cooperative/persistent_topk require lengths.numel() == num_rows.
+    row_ends = torch.randint(4000, 16000, (num_rows,), dtype=torch.int32, device="cuda")
+    seq_lens = row_ends.reshape(batch_size, next_n)
+    max_seq_len = int(row_ends.max())
+    # Pad the stride: multiple of 4 for cooperative_topk (TMA) and
+    # 1024B-aligned (256 fp32) for DeepSelect.
+    stride = (max_seq_len + 255) // 256 * 256
+
+    row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda")
+
+    logits = torch.randn(num_rows, stride, dtype=torch.float32, device="cuda")
+    col_idx = torch.arange(stride, device="cuda")
+    logits[col_idx[None, :] >= row_ends[:, None]] = 1e30  # dirty tail
+
+    indices = torch.full((num_rows, top_k), -2, dtype=torch.int32, device="cuda")
+    cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
+    with set_current_vllm_config(cfg):
+        SparseIndexerTopk()(logits, seq_lens, next_n, indices, top_k, max_seq_len)
+    torch.accelerator.synchronize()
+
+    # k_i == top_k for every row here (row_ends >= 3997 > top_k).
+    assert torch.all(indices >= 0)
+    torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
+    for i in range(num_rows):
+        row_end = int(row_ends[i])
+        torch_indices[i] = logits[i, :row_end].topk(top_k, dim=-1)[1]
+        assert torch.all(indices[i] < row_end), (
+            f"{backend}: row {i} selected indices past its seq_len"
+        )
+
+    assert compare_top_k_results(
+        logits, indices, torch_indices, row_starts, row_ends, top_k
+    ), f"{backend} results don't match torch.topk"
+
+
+@pytest.mark.parametrize(
+    "backend",
+    [
+        "torch",
+        pytest.param(
+            "per_row",
+            marks=pytest.mark.skipif(
+                not current_platform.is_cuda(), reason="requires CUDA"
+            ),
+        ),
+        pytest.param(
+            "flashinfer",
+            marks=pytest.mark.skipif(
+                not _has_flashinfer_topk(), reason="requires flashinfer top-k"
+            ),
+        ),
+        pytest.param("deep_select", marks=requires_sm100),
+    ],
+)
+@torch.inference_mode()
+def test_sparse_indexer_decode_topk_short_seq_lens(
+    backend: str, workspace_init
+) -> None:
+    """seq_len < next_n: derived per-row ends must clamp at 0 (the reference
+    kernels do max(0, ...)); negative ends are OOB for the ragged kernels."""
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk
+
+    set_random_seed(0)
+    next_n = 4
+    top_k = 512
+    vocab = 8192  # 1024B-aligned rows for DeepSelect
+    # 1D (B,) seq_lens: rows of a request share its len and each kernel
+    # derives per-row ends (with clamping) itself.
+    seq_lens = torch.tensor([2, 6], dtype=torch.int32, device="cuda")
+    num_rows = 2 * next_n
+    logits = torch.randn(num_rows, vocab, dtype=torch.float32, device="cuda")
+    indices = torch.full((num_rows, top_k), -2, dtype=torch.int32, device="cuda")
+
+    cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
+    with set_current_vllm_config(cfg):
+        SparseIndexerTopk()(logits, seq_lens, next_n, indices, top_k, vocab)
+    torch.accelerator.synchronize()
+
+    # Derived ends: [2,6] - next_n + 1 + arange(next_n) -> [-1,0,1,2] and
+    # [3,4,5,6], clamped to [0,0,1,2] and [3,4,5,6].
+    for row, end in enumerate([0, 0, 1, 2, 3, 4, 5, 6]):
+        k_i = min(top_k, end)
+        assert torch.all(indices[row, k_i:] == -1), f"{backend}: row {row}"
+        if k_i:
+            ref = logits[row, :end].topk(k_i).indices.sort().values
+            got = indices[row, :k_i].sort().values
+            assert torch.equal(got, ref), f"{backend}: row {row}"
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA")
+@torch.inference_mode()
+def test_sparse_indexer_topk_backend_resolution() -> None:
+    """Auto heuristic chain and fail-fast validation of explicit backends."""
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk
+
+    # 1024B-aligned stride: 131072 * 4 bytes.
+    logits = torch.randn(64, 131072, dtype=torch.float32, device="cuda")
+    # 4-divisible (cooperative TMA ok) but not 1024B-aligned (DeepSelect no):
+    # 131076 * 4 = 524304 bytes.
+    unaligned_logits = torch.randn(128, 131076, dtype=torch.float32, device="cuda")[
+        :, :131072
+    ]
+
+    def resolve(
+        backend: str,
+        t: torch.Tensor = logits,
+        k: int = 2048,
+        num_rows: int = 64,
+    ) -> str:
+        cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
+        with set_current_vllm_config(cfg):
+            return SparseIndexerTopk().resolve_backend(t, k, num_rows)
+
+    is_sm100 = current_platform.is_device_capability_family(100)
+    has_coop = _has_cooperative_topk()
+    has_fi = _has_flashinfer_topk()
+
+    # "auto" is exactly the pre-existing chain: cooperative -> persistent ->
+    # per_row. It must never select the opt-in backends by itself.
+    if has_coop:
+        assert resolve("auto") == "cooperative"
+        assert resolve("auto", k=512, num_rows=8) == "cooperative"
+        # Past cooperative's 64-row limit -> persistent, even where
+        # DeepSelect would be applicable.
+        assert resolve("auto", k=512, num_rows=128) == "persistent"
+        # 16B-aligned (cooperative TMA ok) but not DeepSelect-aligned:
+        # still cooperative at <= 64 rows...
+        assert resolve("auto", unaligned_logits, num_rows=64) == "cooperative"
+    # ...and persistent past the row limit.
+    assert resolve("auto", unaligned_logits, num_rows=128) == "persistent"
+    # Unsupported topk (outside {512, 1024, 2048} for the workspace kernels)
+    # -> per_row fallback.
+    assert resolve("auto", k=5000) == "per_row"
+
+    # Explicit values are returned as-is when constraints are met.
+    assert resolve("per_row") == "per_row"
+    assert resolve("torch") == "torch"
+    assert resolve("persistent") == "persistent"
+    if has_coop:
+        assert resolve("cooperative") == "cooperative"
+    if has_fi:
+        assert resolve("flashinfer") == "flashinfer"
+    if is_sm100:
+        # Explicit deep_select has no row-count requirement.
+        assert resolve("deep_select", num_rows=1) == "deep_select"
+        with pytest.raises(RuntimeError, match="constraints"):
+            resolve("deep_select", unaligned_logits)
+
+    # Fail-fast on unmet constraints.
+    with pytest.raises(RuntimeError, match="num_rows must be <= 64"):
+        resolve("cooperative", num_rows=128)
+    with pytest.raises(RuntimeError, match="topk_tokens must be in"):
+        resolve("persistent", k=3000)
+    with pytest.raises(RuntimeError, match="topk_tokens must be in"):
+        resolve("cooperative", k=3000)
+    if not is_sm100:
+        with pytest.raises(RuntimeError, match="SM100a/SM103a"):
+            resolve("deep_select")

vllm/config/kernel.py

@@ -160,6 +160,16 @@ FLASHINFER_MOE_EP_BACKENDS = frozenset(
 # moe_ep variants.
 MEGA_MOE_BACKENDS = frozenset({"deep_gemm_mega_moe"}) | FLASHINFER_MOE_EP_BACKENDS
 
+SparseIndexerTopkBackend = Literal[
+    "auto",
+    "deep_select",
+    "cooperative",
+    "persistent",
+    "per_row",
+    "flashinfer",
+    "torch",
+]
+
 # Architectures whose model code wires up the flashinfer moe_ep experts. MTP
 # and DSpark draft variants inherit the setting from these target models.
 FLASHINFER_MOE_EP_ARCHITECTURES = frozenset(
@@ -269,6 +279,21 @@ class KernelConfig:
                    running QDQ on activations.
     """
 
+    sparse_indexer_topk_backend: SparseIndexerTopkBackend = "auto"
+    """Backend for the DSA sparse indexer decode top-k kernel. Available options:
+
+    - "auto": The pre-existing chain (cooperative -> persistent -> per_row);
+      the other backends are opt-in
+    - "deep_select": Use DeepSelect kernels (SM100a/SM103a only)
+    - "cooperative": Use vLLM's cooperative_topk kernel
+    - "persistent": Use vLLM's persistent_topk kernel
+    - "per_row": Use vLLM's top_k_per_row_decode kernel
+    - "flashinfer": Use FlashInfer's top_k_ragged_transform kernel
+    - "torch": Use a plain torch.topk implementation (debug reference)
+
+    Explicit values raise RuntimeError when their constraints are not met.
+    """
+
     linear_backend: LinearBackend = "auto"
     """Backend for linear layer GEMM kernels. Available options:
 
@@ -312,6 +337,13 @@ class KernelConfig:
             return value.lower().replace("-", "_")
         return value
 
+    @field_validator("sparse_indexer_topk_backend", mode="before")
+    @classmethod
+    def _normalize_sparse_indexer_topk_backend(cls, value: Any) -> Any:
+        if isinstance(value, str):
+            return value.lower().replace("-", "_")
+        return value
+
     def compute_hash(self) -> str:
         """
         Produces a hash unique to the pass configuration.

vllm/engine/arg_utils.py

@@ -76,7 +76,12 @@ from vllm.config.cache import (
     PrefixCachingHashAlgo,
 )
 from vllm.config.device import Device
-from vllm.config.kernel import IrOpPriorityConfig, LinearBackend, MoEBackend
+from vllm.config.kernel import (
+    IrOpPriorityConfig,
+    LinearBackend,
+    MoEBackend,
+    SparseIndexerTopkBackend,
+)
 from vllm.config.load import SafetensorsLoadStrategy
 from vllm.config.lora import MaxLoRARanks
 from vllm.config.mamba import MambaBackendEnum, MambaSSUAlgorithm
@@ -505,6 +510,9 @@ class EngineArgs:
     enable_ep_weight_filter: bool = ParallelConfig.enable_ep_weight_filter
     moe_backend: MoEBackend = KernelConfig.moe_backend
     linear_backend: LinearBackend = KernelConfig.linear_backend
+    sparse_indexer_topk_backend: SparseIndexerTopkBackend = (
+        KernelConfig.sparse_indexer_topk_backend
+    )
     all2all_backend: All2AllBackend = ParallelConfig.all2all_backend
     enable_elastic_ep: bool = ParallelConfig.enable_elastic_ep
     elastic_ep_max_dp_size: int = ParallelConfig.elastic_ep_max_dp_size
@@ -1670,6 +1678,11 @@ class EngineArgs:
         linear_backend_kwargs = kernel_kwargs["linear_backend"]
         linear_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_")
         kernel_group.add_argument("--linear-backend", **linear_backend_kwargs)
+        sparse_indexer_topk_kwargs = kernel_kwargs["sparse_indexer_topk_backend"]
+        sparse_indexer_topk_kwargs["type"] = lambda s: s.lower().replace("-", "_")
+        kernel_group.add_argument(
+            "--sparse-indexer-topk-backend", **sparse_indexer_topk_kwargs
+        )
 
         # vLLM arguments
         vllm_kwargs = get_kwargs(VllmConfig)
@@ -2548,6 +2561,8 @@ class EngineArgs:
             kernel_config.moe_backend = self.moe_backend
         if self.linear_backend != "auto":
             kernel_config.linear_backend = self.linear_backend
+        if self.sparse_indexer_topk_backend != "auto":
+            kernel_config.sparse_indexer_topk_backend = self.sparse_indexer_topk_backend
 
         # Transfer top-level ir_op_priority into KernelConfig.ir_op_priority
         for op_name, op_priority in asdict(self.ir_op_priority).items():

vllm/model_executor/layers/indexer_topk.py

@@ -0,0 +1,339 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Top-k kernels for the DSA sparse attention indexer."""
+
+import functools
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.config import get_current_vllm_config
+from vllm.platforms import current_platform
+from vllm.utils.flashinfer import has_flashinfer
+from vllm.v1.worker.workspace import current_workspace_manager
+
+RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
+
+# ---------------------------------------------------------------------------
+# DeepSelect (vllm._deepselect_C)
+# ---------------------------------------------------------------------------
+
+# Matches the -1 fill convention used for topk_indices_buffer elsewhere.
+IDX_OOB_FILL_VALUE = -1
+
+try:
+    import vllm._deepselect_C  # noqa: F401  (registers torch.ops.deep_select)
+except ImportError as e:
+    from vllm.logger import init_logger
+
+    init_logger(__name__).warning(
+        "Failed to import the DeepSelect extension (vllm._deepselect_C): %s", e
+    )
+
+
+@functools.lru_cache(maxsize=1)
+def get_deep_select_stride_requirement() -> tuple[int, int]:
+    """Stride alignment requirement (input, output) in bytes."""
+    return torch.ops.deep_select.get_alignment_requirement()
+
+
+def is_deep_select_supported(input: torch.Tensor, topk: int) -> bool:
+    """Whether the kernel accepts this input (dtype/stride/topk constraints)."""
+    return (
+        topk <= 4096
+        and input.shape[1] < 2**23
+        and input.dtype in (torch.float32, torch.bfloat16)
+        and input.stride(1) == 1
+        and input.stride(0)
+        * input.element_size()
+        % get_deep_select_stride_requirement()[0]
+        == 0
+    )
+
+
+def _get_empty_and_aligned_tensor(
+    dim0: int, dim1: int, device: torch.device, dtype: torch.dtype
+) -> torch.Tensor:
+    """Tensor with shape (dim0, dim1) whose stride(0) is 32B-aligned."""
+    output_stride_requirement = get_deep_select_stride_requirement()[1] // (
+        dtype.itemsize
+    )
+    assert output_stride_requirement > 0
+    dim1_rounded = (
+        (dim1 + output_stride_requirement - 1) // output_stride_requirement
+    ) * output_stride_requirement
+    return torch.empty((dim0, dim1_rounded), device=device, dtype=dtype)[:, :dim1]
+
+
+def deep_select_topk(
+    input: torch.Tensor,
+    topk: int,
+    end: torch.Tensor | None = None,
+    output_idx: torch.Tensor | None = None,
+    indices_dtype: torch.dtype = torch.int32,
+) -> torch.Tensor:
+    """Select the top-k indices per row of `input` with DeepSelect.
+
+    Args:
+        input: (num_rows, vocab_size), bf16 or fp32. stride(1) must be 1 and
+            stride(0) must be 1024B-aligned.
+        topk: Number of elements to select per row; must be <= 4096.
+        end: Optional (num_rows,) int32 tensor with the exclusive right
+            boundary of each row. Rows with `end[i] < topk` get their
+            remaining indices filled with -1.
+        output_idx: Optional preallocated (num_rows, topk) output tensor whose
+            stride(0) is 32B-aligned (e.g. a slice of a wider buffer).
+        indices_dtype: Output dtype when `output_idx` is not provided.
+
+    Returns:
+        The (num_rows, topk) indices tensor.
+    """
+    assert input.dim() == 2 and input.stride(1) == 1
+    assert (
+        input.stride(0) * input.element_size() % get_deep_select_stride_requirement()[0]
+        == 0
+    )
+
+    num_rows = input.shape[0]
+    if output_idx is None:
+        output_idx = _get_empty_and_aligned_tensor(
+            num_rows, topk, input.device, indices_dtype
+        )
+    else:
+        assert output_idx.dtype == indices_dtype
+        assert output_idx.shape[0] >= num_rows and output_idx.shape[1] >= topk
+
+    torch.ops.deep_select.topk(
+        input,
+        topk,
+        None,  # begin is not supported
+        end,
+        False,  # sorted_value
+        False,  # sorted_index
+        None,  # output_value
+        output_idx,
+        None,  # output_idx_offset
+        IDX_OOB_FILL_VALUE,
+        float("-inf"),  # value_oob_fill_value
+        False,  # return_value
+        False,  # abort_when_nan_found: dummy/capture passes can feed NaN
+        # logits from uninitialized KV cache; the reference topk kernels do
+        # not trap on NaN, and trapping would abort CUDA graph capture.
+    )
+    return output_idx
+
+
+# ---------------------------------------------------------------------------
+# Backend selection
+# ---------------------------------------------------------------------------
+
+
+@functools.cache
+def get_indexer_topk(backend: str) -> "SparseIndexerTopk":
+    return SparseIndexerTopk(backend)
+
+
+class SparseIndexerTopk(torch.nn.Module):
+    """The sparse indexer's decode top-k stage.
+
+    Selects among the available top-k kernels (see
+    kernel_config.sparse_indexer_topk_backend) and runs the chosen one.
+    """
+
+    def __init__(self, backend: str | None = None) -> None:
+        super().__init__()
+        if backend is None:
+            backend = (
+                get_current_vllm_config().kernel_config.sparse_indexer_topk_backend
+            )
+        self._backend = backend
+        self._is_cuda = current_platform.is_cuda()
+        self._has_deep_select = self._is_cuda and (
+            current_platform.is_device_capability_family(100)
+        )
+        self._has_flashinfer_topk = has_flashinfer()
+        self._cooperative_capable = self._is_cuda and (
+            current_platform.has_device_capability(90)
+            and not current_platform.is_device_capability_family(120)
+        )
+
+    def resolve_backend(
+        self, logits: torch.Tensor, topk_tokens: int, num_rows: int
+    ) -> str:
+        """Resolve the decode top-k implementation from the configured
+        backend ("auto" = the pre-existing chain, or a validated explicit
+        value)."""
+        if self._backend == "auto":
+            return self._resolve_auto(logits, topk_tokens, num_rows)
+
+        failures: list[str] = []
+        if self._backend == "cooperative":
+            failures = self._cooperative_constraints(logits, topk_tokens, num_rows)
+        elif self._backend == "persistent":
+            if not self._is_cuda:
+                failures.append("requires a CUDA platform")
+            if topk_tokens not in (512, 1024, 2048):
+                failures.append(
+                    f"topk_tokens must be in (512, 1024, 2048), got {topk_tokens}"
+                )
+        elif self._backend == "deep_select":
+            if not self._is_cuda:
+                failures.append("requires a CUDA platform")
+            elif not self._has_deep_select:
+                failures.append("requires SM100a/SM103a (10.x device family)")
+            elif not is_deep_select_supported(logits, topk_tokens):
+                failures.append(
+                    f"inputs violate DeepSelect's constraints: dtype={logits.dtype},"
+                    f" stride={logits.stride()}, topk={topk_tokens}"
+                )
+        elif self._backend == "flashinfer":
+            if not self._is_cuda:
+                failures.append("requires a CUDA platform")
+            if not self._has_flashinfer_topk:
+                failures.append(
+                    "flashinfer.topk.top_k_ragged_transform is not importable"
+                )
+            if logits.dtype != torch.float32 or logits.stride(1) != 1:
+                failures.append(
+                    f"requires fp32 logits with stride(1) == 1, got "
+                    f"dtype={logits.dtype}, stride={logits.stride()}"
+                )
+        if failures:
+            raise RuntimeError(
+                f"sparse_indexer_topk_backend='{self._backend}' was requested, but: "
+                + "; ".join(failures)
+            )
+        return self._backend
+
+    def _resolve_auto(
+        self, logits: torch.Tensor, topk_tokens: int, num_rows: int
+    ) -> str:
+        """The pre-existing priority chain: cooperative -> persistent ->
+        per_row. deep_select/flashinfer/torch are opt-in only."""
+        if not self._cooperative_constraints(logits, topk_tokens, num_rows):
+            return "cooperative"
+        if self._is_cuda and topk_tokens in (512, 1024, 2048):
+            return "persistent"
+        return "per_row"
+
+    def _cooperative_constraints(
+        self, logits: torch.Tensor, topk_tokens: int, num_rows: int
+    ) -> list[str]:
+        """Unmet constraints of cooperative_topk (empty when applicable)."""
+        failures = []
+        if not self._is_cuda:
+            failures.append("requires a CUDA platform")
+        if topk_tokens not in (512, 1024, 2048):
+            failures.append(
+                f"topk_tokens must be in (512, 1024, 2048), got {topk_tokens}"
+            )
+        if num_rows > 64:
+            failures.append(f"num_rows must be <= 64, got {num_rows}")
+        if logits.stride(0) % 4 != 0:
+            failures.append(
+                f"logits.stride(0) must be divisible by 4, got {logits.stride(0)}"
+            )
+        if self._is_cuda and not self._cooperative_capable:
+            failures.append("requires SM90+ and is not supported on the SM12x family")
+        return failures
+
+    @staticmethod
+    def _row_ends(seq_lens: torch.Tensor, next_n: int, num_rows: int) -> torch.Tensor:
+        """Per-row exclusive end offsets (int32, (num_rows,)) for top-k
+        kernels that take ragged lengths (DeepSelect, FlashInfer, torch
+        reference).
+
+        seq_lens is (B, next_n) per-row effective lens for native spec decode
+        and (B, 1) otherwise, in which case per-row lens are derived the same
+        way as the other decode top-k kernels.
+        """
+        if seq_lens.numel() == num_rows:
+            row_ends = seq_lens.reshape(-1)
+        else:
+            next_n_offsets = torch.arange(
+                next_n, dtype=torch.int32, device=seq_lens.device
+            )
+            row_ends = (
+                (seq_lens.reshape(-1, 1) - next_n + 1 + next_n_offsets)
+                .clamp_(min=0)
+                .reshape(-1)
+            )
+        assert row_ends.dtype == torch.int32
+        return row_ends
+
+    def forward(
+        self,
+        logits: torch.Tensor,
+        seq_lens: torch.Tensor,
+        next_n: int,
+        topk_indices: torch.Tensor,
+        topk_tokens: int,
+        max_seq_len: int,
+    ) -> None:
+        """Run the resolved decode top-k implementation, writing into
+        topk_indices (int32, -1 fill for rows shorter than topk_tokens)."""
+        backend = self.resolve_backend(logits, topk_tokens, logits.shape[0])
+        if backend == "deep_select":
+            row_ends = self._row_ends(seq_lens, next_n, logits.shape[0])
+            deep_select_topk(logits, topk_tokens, end=row_ends, output_idx=topk_indices)
+        elif backend == "cooperative":
+            (topk_workspace,) = current_workspace_manager().get_simultaneous(
+                ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
+            )
+            torch.ops._C.cooperative_topk(
+                logits,
+                seq_lens,
+                topk_indices,
+                topk_workspace,
+                topk_tokens,
+                max_seq_len,
+            )
+        elif backend == "persistent":
+            (topk_workspace,) = current_workspace_manager().get_simultaneous(
+                ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
+            )
+            torch.ops._C.persistent_topk(
+                logits,
+                seq_lens,
+                topk_indices,
+                topk_workspace,
+                topk_tokens,
+                logits.shape[1],
+            )
+        elif backend == "flashinfer":
+            # Deferred: importing flashinfer initializes CUDA at import time.
+            from flashinfer.topk import top_k_ragged_transform
+
+            row_ends = self._row_ends(seq_lens, next_n, logits.shape[0])
+            offsets = torch.zeros(
+                logits.shape[0], dtype=torch.int32, device=logits.device
+            )
+            # top_k_ragged_transform selects within [0, row_ends[i]) per
+            # row and -1-fills past the row length.
+            indices = top_k_ragged_transform(logits, offsets, row_ends, topk_tokens)
+            topk_indices.copy_(indices)
+        elif backend == "torch":
+            # Debug reference: mask everything past each row's end, then topk.
+            row_ends = self._row_ends(seq_lens, next_n, logits.shape[0])
+            cols = torch.arange(logits.shape[1], device=logits.device)
+            masked = logits.masked_fill(
+                cols.unsqueeze(0) >= row_ends.unsqueeze(1), float("-inf")
+            )
+            indices = masked.topk(topk_tokens, dim=-1).indices
+            in_range = torch.arange(topk_tokens, device=logits.device).unsqueeze(
+                0
+            ) < row_ends.unsqueeze(1)
+            indices = torch.where(in_range, indices, -1)
+            topk_indices.copy_(indices)
+        else:
+            assert backend == "per_row", f"unknown topk backend: {backend}"
+            ops.top_k_per_row_decode(
+                logits,
+                next_n,
+                seq_lens,
+                topk_indices,
+                logits.shape[0],
+                logits.stride(0),
+                logits.stride(1),
+                topk_tokens,
+            )

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -20,6 +20,10 @@ from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
 from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
     select_candidate_blocks as _select_candidate_blocks,
 )
+from vllm.model_executor.layers.indexer_topk import (
+    RADIX_TOPK_WORKSPACE_SIZE,
+    get_indexer_topk,
+)
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     get_fp8_min_max,
 )
@@ -46,8 +50,6 @@ from vllm.v1.worker.workspace import current_workspace_manager
 
 logger = init_logger(__name__)
 
-RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
-
 # MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32.
 MXFP4_BLOCK_SIZE = 32
 
@@ -343,6 +345,7 @@ def sparse_attn_indexer(
     candidate_blocks: torch.Tensor | None = None,
     candidate_block_size: int = 0,
     candidate_write: bool = False,
+    topk_backend: str = "auto",
 ) -> torch.Tensor:
     # careful! this will be None in dummy run
     forward_context = get_forward_context()
@@ -750,56 +753,16 @@ def sparse_attn_indexer(
                 )
         topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
 
-        use_cooperative_topk = (
-            current_platform.is_cuda()
-            and topk_tokens in (512, 1024, 2048)
-            and num_rows <= 64
-            and logits.stride(0) % 4 == 0  # TMA 16-byte alignment
-            and current_platform.has_device_capability(90)
-            and not current_platform.is_device_capability_family(120)
+        # The backend comes from the layer (config is only readable at model
+        # construction); dispatchers are cached per backend.
+        get_indexer_topk(topk_backend)(
+            logits,
+            seq_lens,
+            next_n,
+            topk_indices,
+            topk_tokens,
+            attn_metadata_narrowed.max_seq_len,
         )
-        use_persistent_topk = current_platform.is_cuda() and topk_tokens in (
-            512,
-            1024,
-            2048,
-        )
-        if use_cooperative_topk:
-            workspace_manager = current_workspace_manager()
-            (topk_workspace,) = workspace_manager.get_simultaneous(
-                ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
-            )
-            torch.ops._C.cooperative_topk(
-                logits,
-                seq_lens,
-                topk_indices,
-                topk_workspace,
-                topk_tokens,
-                attn_metadata_narrowed.max_seq_len,
-            )
-        elif use_persistent_topk:
-            workspace_manager = current_workspace_manager()
-            (topk_workspace,) = workspace_manager.get_simultaneous(
-                ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
-            )
-            torch.ops._C.persistent_topk(
-                logits,
-                seq_lens,
-                topk_indices,
-                topk_workspace,
-                topk_tokens,
-                logits.shape[1],
-            )
-        else:
-            ops.top_k_per_row_decode(
-                logits,
-                next_n,
-                seq_lens,
-                topk_indices,
-                num_rows,
-                logits.stride(0),
-                logits.stride(1),
-                topk_tokens,
-            )
 
         if decode_metadata.global_seq_lens is not None:
             _merge_dcp_topk_global(
@@ -851,6 +814,7 @@ def sparse_attn_indexer_fake(
     candidate_blocks: torch.Tensor | None = None,
     candidate_block_size: int = 0,
     candidate_write: bool = False,
+    topk_backend: str = "auto",
 ) -> torch.Tensor:
     return topk_indices_buffer
 
@@ -917,6 +881,7 @@ class SparseAttnIndexer(CustomOp):
         # than threading them through per-step metadata.
         vllm_config = get_current_vllm_config()
         parallel_config = vllm_config.parallel_config
+        self.topk_backend = vllm_config.kernel_config.sparse_indexer_topk_backend
         self._parallel_config = parallel_config
         self.dcp_world_size = parallel_config.decode_context_parallel_size
         self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
@@ -1026,6 +991,7 @@ class SparseAttnIndexer(CustomOp):
             candidate_blocks=self.candidate_blocks,
             candidate_block_size=self.candidate_block_size,
             candidate_write=self.candidate_write,
+            topk_backend=self.topk_backend,
         )
 
     def forward_xpu(