diff --git a/CMakeLists.txt b/CMakeLists.txt index 018138f54c..fbacfd9091 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/cmake/external_projects/deepselect.cmake b/cmake/external_projects/deepselect.cmake new file mode 100644 index 0000000000..e2a5592e61 --- /dev/null +++ b/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 + $<$:-std=c++20> + $<$:-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() diff --git a/setup.py b/setup.py index 81ff499be4..13692f3d20 100644 --- a/setup.py +++ b/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") ): diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index a1e27d2e50..c629ae15de 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/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") diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index c16e4fc169..d44e83bf67 100644 --- a/vllm/config/kernel.py +++ b/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. diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index ae6f14b1df..b420e84f54 100644 --- a/vllm/engine/arg_utils.py +++ b/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(): diff --git a/vllm/model_executor/layers/indexer_topk.py b/vllm/model_executor/layers/indexer_topk.py new file mode 100644 index 0000000000..a3b9da4e1e --- /dev/null +++ b/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, + ) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index d663c54de3..32902901c7 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/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(