[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>
Jiangyun Zhu committed
Sep 13, 2026 at 15:44 UTC
fa008bdccf10f2f31f84553112a227cf1b8947d7
8 files changed
+829
-52
CMakeLists.txt
+1
@@ -1572,6 +1572,7 @@ endif()
1572
# For CUDA we also build and ship some external projects.
1573
if (VLLM_GPU_LANG STREQUAL "CUDA")
1574
include(cmake/external_projects/deepgemm.cmake)
1575
+ include(cmake/external_projects/deepselect.cmake)
1576
include(cmake/external_projects/fmha_sm100.cmake)
1577
include(cmake/external_projects/flashmla.cmake)
1578
include(cmake/external_projects/flashkda.cmake)
cmake/external_projects/deepselect.cmake
new
+97
@@ -0,0 +1,97 @@
1
+include(FetchContent)
2
+
3
+# If DEEPSELECT_SRC_DIR is set, DeepSelect is built from that directory
4
+# instead of downloading.
5
+# It can be set as an environment variable or passed as a cmake argument.
6
+# The environment variable takes precedence.
7
+if (DEFINED ENV{DEEPSELECT_SRC_DIR})
8
+ set(DEEPSELECT_SRC_DIR $ENV{DEEPSELECT_SRC_DIR})
9
+endif()
10
+
11
+if(DEEPSELECT_SRC_DIR)
12
+ FetchContent_Declare(
13
+ deepselect
14
+ SOURCE_DIR ${DEEPSELECT_SRC_DIR}
15
+ CONFIGURE_COMMAND ""
16
+ BUILD_COMMAND ""
17
+ )
18
+else()
19
+ FetchContent_Declare(
20
+ deepselect
21
+ GIT_REPOSITORY https://github.com/vllm-project/DeepSelect.git
22
+ GIT_TAG d96d33afe1fab0d6066da49cdc91e64c2bee65ea
23
+ GIT_PROGRESS TRUE
24
+ CONFIGURE_COMMAND ""
25
+ BUILD_COMMAND ""
26
+ )
27
+endif()
28
+
29
+FetchContent_MakeAvailable(deepselect)
30
+message(STATUS "DeepSelect is available at ${deepselect_SOURCE_DIR}")
31
+
32
+# DeepSelect kernels require CUDA 12.9+; 10.0f (family) covers SM100 and SM103.
33
+set(DEEPSELECT_SUPPORT_ARCHS)
34
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
35
+ list(APPEND DEEPSELECT_SUPPORT_ARCHS "10.0f")
36
+endif()
37
+
38
+cuda_archs_loose_intersection(DEEPSELECT_ARCHS "${DEEPSELECT_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
39
+if(DEEPSELECT_ARCHS)
40
+ message(STATUS "DeepSelect CUDA architectures: ${DEEPSELECT_ARCHS}")
41
+
42
+ # All explicit template instantiations must be compiled: api.cpp
43
+ # references every dispatch combination.
44
+ file(GLOB DeepSelect_SOURCES CONFIGURE_DEPENDS
45
+ ${deepselect_SOURCE_DIR}/csrc/cuda_kernels/v3/instantiations/*.cu
46
+ ${deepselect_SOURCE_DIR}/csrc/cuda_kernels/v3_fp32/instantiations/*.cu
47
+ ${deepselect_SOURCE_DIR}/csrc/cuda_kernels/v3_cluster/instantiations/*.cu
48
+ )
49
+ list(PREPEND DeepSelect_SOURCES ${deepselect_SOURCE_DIR}/csrc/api.cpp)
50
+
51
+ set(DeepSelect_INCLUDES
52
+ ${deepselect_SOURCE_DIR}/csrc
53
+ ${deepselect_SOURCE_DIR}/csrc/3rdparty/cutlass/include
54
+ ${deepselect_SOURCE_DIR}/csrc/3rdparty/kerutils/include
55
+ )
56
+
57
+ set_gencode_flags_for_srcs(
58
+ SRCS "${DeepSelect_SOURCES}"
59
+ CUDA_ARCHS "${DEEPSELECT_ARCHS}")
60
+
61
+ set(VLLM_DEEPSELECT_GPU_FLAGS ${VLLM_GPU_FLAGS})
62
+ list(APPEND VLLM_DEEPSELECT_GPU_FLAGS
63
+ "-O3" "--expt-relaxed-constexpr" "--expt-extended-lambda"
64
+ "--use_fast_math" "--ftz=false"
65
+ "-U__CUDA_NO_HALF_OPERATORS__" "-U__CUDA_NO_HALF_CONVERSIONS__"
66
+ "-U__CUDA_NO_HALF2_OPERATORS__" "-U__CUDA_NO_BFLOAT16_CONVERSIONS__")
67
+
68
+ define_extension_target(
69
+ _deepselect_C
70
+ DESTINATION vllm
71
+ LANGUAGE ${VLLM_GPU_LANG}
72
+ SOURCES ${DeepSelect_SOURCES}
73
+ COMPILE_FLAGS ${VLLM_DEEPSELECT_GPU_FLAGS}
74
+ ARCHITECTURES ${VLLM_GPU_ARCHES}
75
+ INCLUDE_DIRECTORIES ${DeepSelect_INCLUDES}
76
+ USE_SABI 3
77
+ WITH_SOABI)
78
+
79
+ # Only use C-shim APIs available in PyTorch 2.10.
80
+ # _deepselect_C is abi compatible with PyTorch >= TORCH_TARGET_VERSION.
81
+ target_compile_definitions(_deepselect_C PRIVATE
82
+ TORCH_TARGET_VERSION=0x020A000000000000ULL)
83
+
84
+ # Needed to use cuda APIs from C-shim
85
+ if(VLLM_GPU_LANG STREQUAL "CUDA")
86
+ target_compile_definitions(_deepselect_C PRIVATE USE_CUDA)
87
+ endif()
88
+
89
+ # DeepSelect requires C++20
90
+ target_compile_options(_deepselect_C PRIVATE
91
+ $<$<COMPILE_LANGUAGE:CXX>:-std=c++20>
92
+ $<$<COMPILE_LANGUAGE:CUDA>:-std=c++20>)
93
+else()
94
+ message(STATUS "DeepSelect will not compile: unsupported CUDA architecture ${CUDA_ARCHS}")
95
+ # Create an empty target for setup.py on unsupported systems
96
+ add_custom_target(_deepselect_C)
97
+endif()
setup.py
+4
@@ -1015,6 +1015,7 @@ class precompiled_wheel_utils:
1015
"vllm/_flashmla_extension_C.abi3.so",
1016
"vllm/_flashkda_C.abi3.so",
1017
"vllm/_sparse_flashmla_C.abi3.so",
1018
+ "vllm/_deepselect_C.abi3.so",
1019
"vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
1020
"vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
1021
"vllm/cumem_allocator.abi3.so",
@@ -1381,6 +1382,9 @@ if _is_cuda():
1382
ext_modules.append(
1383
CMakeExtension(name="vllm._flashmla_extension_C", optional=True)
1384
)
1385
+ # DeepSelect requires CUDA 12.9 or later (SM100a/SM103a only)
1386
+ # Optional since it won't build on unsupported architectures
1387
+ ext_modules.append(CMakeExtension(name="vllm._deepselect_C", optional=True))
1388
if USE_PRECOMPILED_EXTENSIONS or (
1389
CUDA_HOME and get_nvcc_cuda_version() >= Version("12.0")
1390
):
tests/kernels/test_top_k_per_row.py
+323
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+
5
import numpy as np
6
import pytest
7
import torch
@@ -21,6 +22,13 @@ def _has_device_capability(major: int) -> bool:
22
return current_platform.is_cuda() and current_platform.has_device_capability(major)
23
24
25
+# DeepSelect is compiled for sm_100a/sm_103a only.
26
+requires_sm100 = pytest.mark.skipif(
27
+ not current_platform.is_device_capability_family(100),
28
+ reason="DeepSelect requires SM100a/SM103a",
29
+)
30
+
31
+
32
COOPERATIVE_TOPK_BACKEND = pytest.param(
33
"cooperative_topk",
34
marks=pytest.mark.skipif(
@@ -1315,3 +1323,318 @@ def test_workspace_topk_padded_stride(top_k: int, backend: str) -> None:
1323
f"Row {i}: {backend} with padded stride doesn't match. "
1324
f"seq_len={sl}, stride={padded_stride}"
1325
)
1326
+
1327
+
1328
+@requires_sm100
1329
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
1330
+@pytest.mark.parametrize("batch_size", [1, 8, 64])
1331
+@pytest.mark.parametrize("vocab_size", [4096, 131072])
1332
+@pytest.mark.parametrize("top_k", [512, 2048])
1333
+@torch.inference_mode()
1334
+def test_deep_select_topk(
1335
+ dtype: torch.dtype,
1336
+ batch_size: int,
1337
+ vocab_size: int,
1338
+ top_k: int,
1339
+) -> None:
1340
+ """DeepSelect wrapper vs torch.topk with variable per-row ends."""
1341
+ from vllm.model_executor.layers import indexer_topk
1342
+
1343
+ set_random_seed(0)
1344
+ torch.set_default_device("cuda:0")
1345
+
1346
+ row_ends = torch.randint(
1347
+ 1, vocab_size + 1, (batch_size,), dtype=torch.int32, device="cuda"
1348
+ )
1349
+ # Force one short row to exercise the -1 out-of-bounds fill.
1350
+ row_ends[0] = min(top_k - 1, vocab_size)
1351
+ logits = torch.randn(batch_size, vocab_size, dtype=dtype, device="cuda")
1352
+ # Values beyond each row's end must never be read by the kernel: fill
1353
+ # them with NaN (production tails are uninitialized; DeepSelect traps on
1354
+ # NaN with abort_when_nan_found=True, so a wrong read fails loudly).
1355
+ col_idx = torch.arange(vocab_size, device="cuda")
1356
+ logits[col_idx[None, :] >= row_ends[:, None]] = float("nan")
1357
+
1358
+ indices = indexer_topk.deep_select_topk(logits, top_k, end=row_ends)
1359
+ torch.accelerator.synchronize()
1360
+
1361
+ assert indices.shape == (batch_size, top_k)
1362
+ row_starts = torch.zeros(batch_size, dtype=torch.int32, device="cuda")
1363
+ torch_indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda")
1364
+ for i in range(batch_size):
1365
+ row_end = int(row_ends[i])
1366
+ k_i = min(top_k, row_end)
1367
+ idx = logits[i, :row_end].topk(k_i, dim=-1)[1]
1368
+ torch_indices[i, :k_i] = idx
1369
+ assert torch.all(indices[i, k_i:] == -1), (
1370
+ f"Row {i}: expected -1 fill after {k_i} valid indices"
1371
+ )
1372
+ # All selected indices must lie inside [0, row_end).
1373
+ assert torch.all(indices[i, :k_i] < row_end)
1374
+
1375
+ assert compare_top_k_results(
1376
+ logits, indices, torch_indices, row_starts, row_ends, top_k
1377
+ ), "DeepSelect topk results don't match torch.topk"
1378
+
1379
+
1380
+@requires_sm100
1381
+@torch.inference_mode()
1382
+def test_deep_select_topk_preallocated_output() -> None:
1383
+ """DeepSelect writes into a preallocated aligned buffer (no `end`)."""
1384
+ from vllm.model_executor.layers import indexer_topk
1385
+
1386
+ set_random_seed(0)
1387
+ torch.set_default_device("cuda:0")
1388
+
1389
+ batch_size = 8
1390
+ vocab_size = 131072
1391
+ top_k = 2048
1392
+ logits = torch.randn(batch_size, vocab_size, dtype=torch.float32, device="cuda")
1393
+
1394
+ # Mimic vLLM's topk_indices_buffer: slice of a wider contiguous buffer,
1395
+ # whose stride(0) is 32B-aligned.
1396
+ buffer = torch.full((batch_size, 4096), -2, dtype=torch.int32, device="cuda")
1397
+ output_idx = buffer[:, :top_k]
1398
+ result = indexer_topk.deep_select_topk(logits, top_k, output_idx=output_idx)
1399
+ torch.accelerator.synchronize()
1400
+
1401
+ assert result.data_ptr() == output_idx.data_ptr()
1402
+ assert torch.all(buffer[:, top_k:] == -2), "kernel wrote out of bounds"
1403
+
1404
+ row_ends = torch.full((batch_size,), vocab_size, dtype=torch.int32, device="cuda")
1405
+ row_starts = torch.zeros(batch_size, dtype=torch.int32, device="cuda")
1406
+ torch_indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda")
1407
+ for i in range(batch_size):
1408
+ torch_indices[i] = logits[i].topk(top_k, dim=-1)[1]
1409
+
1410
+ assert compare_top_k_results(
1411
+ logits, output_idx, torch_indices, row_starts, row_ends, top_k
1412
+ ), "DeepSelect topk (preallocated output) results don't match torch.topk"
1413
+
1414
+
1415
+def _has_flashinfer_topk() -> bool:
1416
+ if not current_platform.is_cuda():
1417
+ return False
1418
+ # find_spec only: importing flashinfer initializes CUDA at import time.
1419
+ import importlib.util
1420
+
1421
+ return importlib.util.find_spec("flashinfer") is not None
1422
+
1423
+
1424
+def _has_cooperative_topk() -> bool:
1425
+ return _has_device_capability(
1426
+ 90
1427
+ ) and not current_platform.is_device_capability_family(120)
1428
+
1429
+
1430
+SPARSE_INDEXER_EXPLICIT_BACKENDS = [
1431
+ pytest.param(
1432
+ "per_row",
1433
+ marks=pytest.mark.skipif(
1434
+ not current_platform.is_cuda(), reason="requires CUDA"
1435
+ ),
1436
+ ),
1437
+ "torch",
1438
+ pytest.param(
1439
+ "persistent",
1440
+ marks=pytest.mark.skipif(
1441
+ not current_platform.is_cuda(), reason="requires CUDA"
1442
+ ),
1443
+ ),
1444
+ pytest.param(
1445
+ "cooperative",
1446
+ marks=pytest.mark.skipif(
1447
+ not _has_cooperative_topk(),
1448
+ reason="cooperative_topk requires SM90+ (non-SM12x)",
1449
+ ),
1450
+ ),
1451
+ pytest.param(
1452
+ "flashinfer",
1453
+ marks=pytest.mark.skipif(
1454
+ not _has_flashinfer_topk(), reason="requires flashinfer top-k"
1455
+ ),
1456
+ ),
1457
+ pytest.param("deep_select", marks=requires_sm100),
1458
+]
1459
+
1460
+
1461
+@pytest.mark.parametrize("next_n", [1, 4])
1462
+@pytest.mark.parametrize("backend", SPARSE_INDEXER_EXPLICIT_BACKENDS)
1463
+@torch.inference_mode()
1464
+def test_sparse_indexer_decode_topk_explicit_backends(
1465
+ backend: str, next_n: int, workspace_init
1466
+) -> None:
1467
+ """Every explicit sparse-indexer decode top-k backend must restrict each
1468
+ row to its seq_len (dirty data past the end must never be selected) and
1469
+ match torch.topk on the valid region."""
1470
+ from vllm.config import VllmConfig, set_current_vllm_config
1471
+ from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk
1472
+
1473
+ set_random_seed(0)
1474
+ torch.set_default_device("cuda:0")
1475
+
1476
+ top_k = 2048
1477
+ batch_size = 4
1478
+ num_rows = batch_size * next_n
1479
+
1480
+ # Per-row effective lens in the (B, next_n) "native spec decode" form;
1481
+ # cooperative/persistent_topk require lengths.numel() == num_rows.
1482
+ row_ends = torch.randint(4000, 16000, (num_rows,), dtype=torch.int32, device="cuda")
1483
+ seq_lens = row_ends.reshape(batch_size, next_n)
1484
+ max_seq_len = int(row_ends.max())
1485
+ # Pad the stride: multiple of 4 for cooperative_topk (TMA) and
1486
+ # 1024B-aligned (256 fp32) for DeepSelect.
1487
+ stride = (max_seq_len + 255) // 256 * 256
1488
+
1489
+ row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda")
1490
+
1491
+ logits = torch.randn(num_rows, stride, dtype=torch.float32, device="cuda")
1492
+ col_idx = torch.arange(stride, device="cuda")
1493
+ logits[col_idx[None, :] >= row_ends[:, None]] = 1e30 # dirty tail
1494
+
1495
+ indices = torch.full((num_rows, top_k), -2, dtype=torch.int32, device="cuda")
1496
+ cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
1497
+ with set_current_vllm_config(cfg):
1498
+ SparseIndexerTopk()(logits, seq_lens, next_n, indices, top_k, max_seq_len)
1499
+ torch.accelerator.synchronize()
1500
+
1501
+ # k_i == top_k for every row here (row_ends >= 3997 > top_k).
1502
+ assert torch.all(indices >= 0)
1503
+ torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
1504
+ for i in range(num_rows):
1505
+ row_end = int(row_ends[i])
1506
+ torch_indices[i] = logits[i, :row_end].topk(top_k, dim=-1)[1]
1507
+ assert torch.all(indices[i] < row_end), (
1508
+ f"{backend}: row {i} selected indices past its seq_len"
1509
+ )
1510
+
1511
+ assert compare_top_k_results(
1512
+ logits, indices, torch_indices, row_starts, row_ends, top_k
1513
+ ), f"{backend} results don't match torch.topk"
1514
+
1515
+
1516
+@pytest.mark.parametrize(
1517
+ "backend",
1518
+ [
1519
+ "torch",
1520
+ pytest.param(
1521
+ "per_row",
1522
+ marks=pytest.mark.skipif(
1523
+ not current_platform.is_cuda(), reason="requires CUDA"
1524
+ ),
1525
+ ),
1526
+ pytest.param(
1527
+ "flashinfer",
1528
+ marks=pytest.mark.skipif(
1529
+ not _has_flashinfer_topk(), reason="requires flashinfer top-k"
1530
+ ),
1531
+ ),
1532
+ pytest.param("deep_select", marks=requires_sm100),
1533
+ ],
1534
+)
1535
+@torch.inference_mode()
1536
+def test_sparse_indexer_decode_topk_short_seq_lens(
1537
+ backend: str, workspace_init
1538
+) -> None:
1539
+ """seq_len < next_n: derived per-row ends must clamp at 0 (the reference
1540
+ kernels do max(0, ...)); negative ends are OOB for the ragged kernels."""
1541
+ from vllm.config import VllmConfig, set_current_vllm_config
1542
+ from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk
1543
+
1544
+ set_random_seed(0)
1545
+ next_n = 4
1546
+ top_k = 512
1547
+ vocab = 8192 # 1024B-aligned rows for DeepSelect
1548
+ # 1D (B,) seq_lens: rows of a request share its len and each kernel
1549
+ # derives per-row ends (with clamping) itself.
1550
+ seq_lens = torch.tensor([2, 6], dtype=torch.int32, device="cuda")
1551
+ num_rows = 2 * next_n
1552
+ logits = torch.randn(num_rows, vocab, dtype=torch.float32, device="cuda")
1553
+ indices = torch.full((num_rows, top_k), -2, dtype=torch.int32, device="cuda")
1554
+
1555
+ cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
1556
+ with set_current_vllm_config(cfg):
1557
+ SparseIndexerTopk()(logits, seq_lens, next_n, indices, top_k, vocab)
1558
+ torch.accelerator.synchronize()
1559
+
1560
+ # Derived ends: [2,6] - next_n + 1 + arange(next_n) -> [-1,0,1,2] and
1561
+ # [3,4,5,6], clamped to [0,0,1,2] and [3,4,5,6].
1562
+ for row, end in enumerate([0, 0, 1, 2, 3, 4, 5, 6]):
1563
+ k_i = min(top_k, end)
1564
+ assert torch.all(indices[row, k_i:] == -1), f"{backend}: row {row}"
1565
+ if k_i:
1566
+ ref = logits[row, :end].topk(k_i).indices.sort().values
1567
+ got = indices[row, :k_i].sort().values
1568
+ assert torch.equal(got, ref), f"{backend}: row {row}"
1569
+
1570
+
1571
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA")
1572
+@torch.inference_mode()
1573
+def test_sparse_indexer_topk_backend_resolution() -> None:
1574
+ """Auto heuristic chain and fail-fast validation of explicit backends."""
1575
+ from vllm.config import VllmConfig, set_current_vllm_config
1576
+ from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk
1577
+
1578
+ # 1024B-aligned stride: 131072 * 4 bytes.
1579
+ logits = torch.randn(64, 131072, dtype=torch.float32, device="cuda")
1580
+ # 4-divisible (cooperative TMA ok) but not 1024B-aligned (DeepSelect no):
1581
+ # 131076 * 4 = 524304 bytes.
1582
+ unaligned_logits = torch.randn(128, 131076, dtype=torch.float32, device="cuda")[
1583
+ :, :131072
1584
+ ]
1585
+
1586
+ def resolve(
1587
+ backend: str,
1588
+ t: torch.Tensor = logits,
1589
+ k: int = 2048,
1590
+ num_rows: int = 64,
1591
+ ) -> str:
1592
+ cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
1593
+ with set_current_vllm_config(cfg):
1594
+ return SparseIndexerTopk().resolve_backend(t, k, num_rows)
1595
+
1596
+ is_sm100 = current_platform.is_device_capability_family(100)
1597
+ has_coop = _has_cooperative_topk()
1598
+ has_fi = _has_flashinfer_topk()
1599
+
1600
+ # "auto" is exactly the pre-existing chain: cooperative -> persistent ->
1601
+ # per_row. It must never select the opt-in backends by itself.
1602
+ if has_coop:
1603
+ assert resolve("auto") == "cooperative"
1604
+ assert resolve("auto", k=512, num_rows=8) == "cooperative"
1605
+ # Past cooperative's 64-row limit -> persistent, even where
1606
+ # DeepSelect would be applicable.
1607
+ assert resolve("auto", k=512, num_rows=128) == "persistent"
1608
+ # 16B-aligned (cooperative TMA ok) but not DeepSelect-aligned:
1609
+ # still cooperative at <= 64 rows...
1610
+ assert resolve("auto", unaligned_logits, num_rows=64) == "cooperative"
1611
+ # ...and persistent past the row limit.
1612
+ assert resolve("auto", unaligned_logits, num_rows=128) == "persistent"
1613
+ # Unsupported topk (outside {512, 1024, 2048} for the workspace kernels)
1614
+ # -> per_row fallback.
1615
+ assert resolve("auto", k=5000) == "per_row"
1616
+
1617
+ # Explicit values are returned as-is when constraints are met.
1618
+ assert resolve("per_row") == "per_row"
1619
+ assert resolve("torch") == "torch"
1620
+ assert resolve("persistent") == "persistent"
1621
+ if has_coop:
1622
+ assert resolve("cooperative") == "cooperative"
1623
+ if has_fi:
1624
+ assert resolve("flashinfer") == "flashinfer"
1625
+ if is_sm100:
1626
+ # Explicit deep_select has no row-count requirement.
1627
+ assert resolve("deep_select", num_rows=1) == "deep_select"
1628
+ with pytest.raises(RuntimeError, match="constraints"):
1629
+ resolve("deep_select", unaligned_logits)
1630
+
1631
+ # Fail-fast on unmet constraints.
1632
+ with pytest.raises(RuntimeError, match="num_rows must be <= 64"):
1633
+ resolve("cooperative", num_rows=128)
1634
+ with pytest.raises(RuntimeError, match="topk_tokens must be in"):
1635
+ resolve("persistent", k=3000)
1636
+ with pytest.raises(RuntimeError, match="topk_tokens must be in"):
1637
+ resolve("cooperative", k=3000)
1638
+ if not is_sm100:
1639
+ with pytest.raises(RuntimeError, match="SM100a/SM103a"):
1640
+ resolve("deep_select")
vllm/config/kernel.py
+32
@@ -160,6 +160,16 @@ FLASHINFER_MOE_EP_BACKENDS = frozenset(
160
# moe_ep variants.
161
MEGA_MOE_BACKENDS = frozenset({"deep_gemm_mega_moe"}) | FLASHINFER_MOE_EP_BACKENDS
162
163
+SparseIndexerTopkBackend = Literal[
164
+ "auto",
165
+ "deep_select",
166
+ "cooperative",
167
+ "persistent",
168
+ "per_row",
169
+ "flashinfer",
170
+ "torch",
171
+]
172
+
173
# Architectures whose model code wires up the flashinfer moe_ep experts. MTP
174
# and DSpark draft variants inherit the setting from these target models.
175
FLASHINFER_MOE_EP_ARCHITECTURES = frozenset(
@@ -269,6 +279,21 @@ class KernelConfig:
279
running QDQ on activations.
280
"""
281
282
+ sparse_indexer_topk_backend: SparseIndexerTopkBackend = "auto"
283
+ """Backend for the DSA sparse indexer decode top-k kernel. Available options:
284
+
285
+ - "auto": The pre-existing chain (cooperative -> persistent -> per_row);
286
+ the other backends are opt-in
287
+ - "deep_select": Use DeepSelect kernels (SM100a/SM103a only)
288
+ - "cooperative": Use vLLM's cooperative_topk kernel
289
+ - "persistent": Use vLLM's persistent_topk kernel
290
+ - "per_row": Use vLLM's top_k_per_row_decode kernel
291
+ - "flashinfer": Use FlashInfer's top_k_ragged_transform kernel
292
+ - "torch": Use a plain torch.topk implementation (debug reference)
293
+
294
+ Explicit values raise RuntimeError when their constraints are not met.
295
+ """
296
+
297
linear_backend: LinearBackend = "auto"
298
"""Backend for linear layer GEMM kernels. Available options:
299
@@ -312,6 +337,13 @@ class KernelConfig:
337
return value.lower().replace("-", "_")
338
return value
339
340
+ @field_validator("sparse_indexer_topk_backend", mode="before")
341
+ @classmethod
342
+ def _normalize_sparse_indexer_topk_backend(cls, value: Any) -> Any:
343
+ if isinstance(value, str):
344
+ return value.lower().replace("-", "_")
345
+ return value
346
+
347
def compute_hash(self) -> str:
348
"""
349
Produces a hash unique to the pass configuration.
vllm/engine/arg_utils.py
+16
-1
@@ -76,7 +76,12 @@ from vllm.config.cache import (
76
PrefixCachingHashAlgo,
77
)
78
from vllm.config.device import Device
79
-from vllm.config.kernel import IrOpPriorityConfig, LinearBackend, MoEBackend
79
+from vllm.config.kernel import (
80
+ IrOpPriorityConfig,
81
+ LinearBackend,
82
+ MoEBackend,
83
+ SparseIndexerTopkBackend,
84
+)
85
from vllm.config.load import SafetensorsLoadStrategy
86
from vllm.config.lora import MaxLoRARanks
87
from vllm.config.mamba import MambaBackendEnum, MambaSSUAlgorithm
@@ -505,6 +510,9 @@ class EngineArgs:
510
enable_ep_weight_filter: bool = ParallelConfig.enable_ep_weight_filter
511
moe_backend: MoEBackend = KernelConfig.moe_backend
512
linear_backend: LinearBackend = KernelConfig.linear_backend
513
+ sparse_indexer_topk_backend: SparseIndexerTopkBackend = (
514
+ KernelConfig.sparse_indexer_topk_backend
515
+ )
516
all2all_backend: All2AllBackend = ParallelConfig.all2all_backend
517
enable_elastic_ep: bool = ParallelConfig.enable_elastic_ep
518
elastic_ep_max_dp_size: int = ParallelConfig.elastic_ep_max_dp_size
@@ -1670,6 +1678,11 @@ class EngineArgs:
1678
linear_backend_kwargs = kernel_kwargs["linear_backend"]
1679
linear_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_")
1680
kernel_group.add_argument("--linear-backend", **linear_backend_kwargs)
1681
+ sparse_indexer_topk_kwargs = kernel_kwargs["sparse_indexer_topk_backend"]
1682
+ sparse_indexer_topk_kwargs["type"] = lambda s: s.lower().replace("-", "_")
1683
+ kernel_group.add_argument(
1684
+ "--sparse-indexer-topk-backend", **sparse_indexer_topk_kwargs
1685
+ )
1686
1687
# vLLM arguments
1688
vllm_kwargs = get_kwargs(VllmConfig)
@@ -2548,6 +2561,8 @@ class EngineArgs:
2561
kernel_config.moe_backend = self.moe_backend
2562
if self.linear_backend != "auto":
2563
kernel_config.linear_backend = self.linear_backend
2564
+ if self.sparse_indexer_topk_backend != "auto":
2565
+ kernel_config.sparse_indexer_topk_backend = self.sparse_indexer_topk_backend
2566
2567
# Transfer top-level ir_op_priority into KernelConfig.ir_op_priority
2568
for op_name, op_priority in asdict(self.ir_op_priority).items():
vllm/model_executor/layers/indexer_topk.py
new
+339
@@ -0,0 +1,339 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Top-k kernels for the DSA sparse attention indexer."""
4
+
5
+import functools
6
+
7
+import torch
8
+
9
+from vllm import _custom_ops as ops
10
+from vllm.config import get_current_vllm_config
11
+from vllm.platforms import current_platform
12
+from vllm.utils.flashinfer import has_flashinfer
13
+from vllm.v1.worker.workspace import current_workspace_manager
14
+
15
+RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
16
+
17
+# ---------------------------------------------------------------------------
18
+# DeepSelect (vllm._deepselect_C)
19
+# ---------------------------------------------------------------------------
20
+
21
+# Matches the -1 fill convention used for topk_indices_buffer elsewhere.
22
+IDX_OOB_FILL_VALUE = -1
23
+
24
+try:
25
+ import vllm._deepselect_C # noqa: F401 (registers torch.ops.deep_select)
26
+except ImportError as e:
27
+ from vllm.logger import init_logger
28
+
29
+ init_logger(__name__).warning(
30
+ "Failed to import the DeepSelect extension (vllm._deepselect_C): %s", e
31
+ )
32
+
33
+
34
+@functools.lru_cache(maxsize=1)
35
+def get_deep_select_stride_requirement() -> tuple[int, int]:
36
+ """Stride alignment requirement (input, output) in bytes."""
37
+ return torch.ops.deep_select.get_alignment_requirement()
38
+
39
+
40
+def is_deep_select_supported(input: torch.Tensor, topk: int) -> bool:
41
+ """Whether the kernel accepts this input (dtype/stride/topk constraints)."""
42
+ return (
43
+ topk <= 4096
44
+ and input.shape[1] < 2**23
45
+ and input.dtype in (torch.float32, torch.bfloat16)
46
+ and input.stride(1) == 1
47
+ and input.stride(0)
48
+ * input.element_size()
49
+ % get_deep_select_stride_requirement()[0]
50
+ == 0
51
+ )
52
+
53
+
54
+def _get_empty_and_aligned_tensor(
55
+ dim0: int, dim1: int, device: torch.device, dtype: torch.dtype
56
+) -> torch.Tensor:
57
+ """Tensor with shape (dim0, dim1) whose stride(0) is 32B-aligned."""
58
+ output_stride_requirement = get_deep_select_stride_requirement()[1] // (
59
+ dtype.itemsize
60
+ )
61
+ assert output_stride_requirement > 0
62
+ dim1_rounded = (
63
+ (dim1 + output_stride_requirement - 1) // output_stride_requirement
64
+ ) * output_stride_requirement
65
+ return torch.empty((dim0, dim1_rounded), device=device, dtype=dtype)[:, :dim1]
66
+
67
+
68
+def deep_select_topk(
69
+ input: torch.Tensor,
70
+ topk: int,
71
+ end: torch.Tensor | None = None,
72
+ output_idx: torch.Tensor | None = None,
73
+ indices_dtype: torch.dtype = torch.int32,
74
+) -> torch.Tensor:
75
+ """Select the top-k indices per row of `input` with DeepSelect.
76
+
77
+ Args:
78
+ input: (num_rows, vocab_size), bf16 or fp32. stride(1) must be 1 and
79
+ stride(0) must be 1024B-aligned.
80
+ topk: Number of elements to select per row; must be <= 4096.
81
+ end: Optional (num_rows,) int32 tensor with the exclusive right
82
+ boundary of each row. Rows with `end[i] < topk` get their
83
+ remaining indices filled with -1.
84
+ output_idx: Optional preallocated (num_rows, topk) output tensor whose
85
+ stride(0) is 32B-aligned (e.g. a slice of a wider buffer).
86
+ indices_dtype: Output dtype when `output_idx` is not provided.
87
+
88
+ Returns:
89
+ The (num_rows, topk) indices tensor.
90
+ """
91
+ assert input.dim() == 2 and input.stride(1) == 1
92
+ assert (
93
+ input.stride(0) * input.element_size() % get_deep_select_stride_requirement()[0]
94
+ == 0
95
+ )
96
+
97
+ num_rows = input.shape[0]
98
+ if output_idx is None:
99
+ output_idx = _get_empty_and_aligned_tensor(
100
+ num_rows, topk, input.device, indices_dtype
101
+ )
102
+ else:
103
+ assert output_idx.dtype == indices_dtype
104
+ assert output_idx.shape[0] >= num_rows and output_idx.shape[1] >= topk
105
+
106
+ torch.ops.deep_select.topk(
107
+ input,
108
+ topk,
109
+ None, # begin is not supported
110
+ end,
111
+ False, # sorted_value
112
+ False, # sorted_index
113
+ None, # output_value
114
+ output_idx,
115
+ None, # output_idx_offset
116
+ IDX_OOB_FILL_VALUE,
117
+ float("-inf"), # value_oob_fill_value
118
+ False, # return_value
119
+ False, # abort_when_nan_found: dummy/capture passes can feed NaN
120
+ # logits from uninitialized KV cache; the reference topk kernels do
121
+ # not trap on NaN, and trapping would abort CUDA graph capture.
122
+ )
123
+ return output_idx
124
+
125
+
126
+# ---------------------------------------------------------------------------
127
+# Backend selection
128
+# ---------------------------------------------------------------------------
129
+
130
+
131
+@functools.cache
132
+def get_indexer_topk(backend: str) -> "SparseIndexerTopk":
133
+ return SparseIndexerTopk(backend)
134
+
135
+
136
+class SparseIndexerTopk(torch.nn.Module):
137
+ """The sparse indexer's decode top-k stage.
138
+
139
+ Selects among the available top-k kernels (see
140
+ kernel_config.sparse_indexer_topk_backend) and runs the chosen one.
141
+ """
142
+
143
+ def __init__(self, backend: str | None = None) -> None:
144
+ super().__init__()
145
+ if backend is None:
146
+ backend = (
147
+ get_current_vllm_config().kernel_config.sparse_indexer_topk_backend
148
+ )
149
+ self._backend = backend
150
+ self._is_cuda = current_platform.is_cuda()
151
+ self._has_deep_select = self._is_cuda and (
152
+ current_platform.is_device_capability_family(100)
153
+ )
154
+ self._has_flashinfer_topk = has_flashinfer()
155
+ self._cooperative_capable = self._is_cuda and (
156
+ current_platform.has_device_capability(90)
157
+ and not current_platform.is_device_capability_family(120)
158
+ )
159
+
160
+ def resolve_backend(
161
+ self, logits: torch.Tensor, topk_tokens: int, num_rows: int
162
+ ) -> str:
163
+ """Resolve the decode top-k implementation from the configured
164
+ backend ("auto" = the pre-existing chain, or a validated explicit
165
+ value)."""
166
+ if self._backend == "auto":
167
+ return self._resolve_auto(logits, topk_tokens, num_rows)
168
+
169
+ failures: list[str] = []
170
+ if self._backend == "cooperative":
171
+ failures = self._cooperative_constraints(logits, topk_tokens, num_rows)
172
+ elif self._backend == "persistent":
173
+ if not self._is_cuda:
174
+ failures.append("requires a CUDA platform")
175
+ if topk_tokens not in (512, 1024, 2048):
176
+ failures.append(
177
+ f"topk_tokens must be in (512, 1024, 2048), got {topk_tokens}"
178
+ )
179
+ elif self._backend == "deep_select":
180
+ if not self._is_cuda:
181
+ failures.append("requires a CUDA platform")
182
+ elif not self._has_deep_select:
183
+ failures.append("requires SM100a/SM103a (10.x device family)")
184
+ elif not is_deep_select_supported(logits, topk_tokens):
185
+ failures.append(
186
+ f"inputs violate DeepSelect's constraints: dtype={logits.dtype},"
187
+ f" stride={logits.stride()}, topk={topk_tokens}"
188
+ )
189
+ elif self._backend == "flashinfer":
190
+ if not self._is_cuda:
191
+ failures.append("requires a CUDA platform")
192
+ if not self._has_flashinfer_topk:
193
+ failures.append(
194
+ "flashinfer.topk.top_k_ragged_transform is not importable"
195
+ )
196
+ if logits.dtype != torch.float32 or logits.stride(1) != 1:
197
+ failures.append(
198
+ f"requires fp32 logits with stride(1) == 1, got "
199
+ f"dtype={logits.dtype}, stride={logits.stride()}"
200
+ )
201
+ if failures:
202
+ raise RuntimeError(
203
+ f"sparse_indexer_topk_backend='{self._backend}' was requested, but: "
204
+ + "; ".join(failures)
205
+ )
206
+ return self._backend
207
+
208
+ def _resolve_auto(
209
+ self, logits: torch.Tensor, topk_tokens: int, num_rows: int
210
+ ) -> str:
211
+ """The pre-existing priority chain: cooperative -> persistent ->
212
+ per_row. deep_select/flashinfer/torch are opt-in only."""
213
+ if not self._cooperative_constraints(logits, topk_tokens, num_rows):
214
+ return "cooperative"
215
+ if self._is_cuda and topk_tokens in (512, 1024, 2048):
216
+ return "persistent"
217
+ return "per_row"
218
+
219
+ def _cooperative_constraints(
220
+ self, logits: torch.Tensor, topk_tokens: int, num_rows: int
221
+ ) -> list[str]:
222
+ """Unmet constraints of cooperative_topk (empty when applicable)."""
223
+ failures = []
224
+ if not self._is_cuda:
225
+ failures.append("requires a CUDA platform")
226
+ if topk_tokens not in (512, 1024, 2048):
227
+ failures.append(
228
+ f"topk_tokens must be in (512, 1024, 2048), got {topk_tokens}"
229
+ )
230
+ if num_rows > 64:
231
+ failures.append(f"num_rows must be <= 64, got {num_rows}")
232
+ if logits.stride(0) % 4 != 0:
233
+ failures.append(
234
+ f"logits.stride(0) must be divisible by 4, got {logits.stride(0)}"
235
+ )
236
+ if self._is_cuda and not self._cooperative_capable:
237
+ failures.append("requires SM90+ and is not supported on the SM12x family")
238
+ return failures
239
+
240
+ @staticmethod
241
+ def _row_ends(seq_lens: torch.Tensor, next_n: int, num_rows: int) -> torch.Tensor:
242
+ """Per-row exclusive end offsets (int32, (num_rows,)) for top-k
243
+ kernels that take ragged lengths (DeepSelect, FlashInfer, torch
244
+ reference).
245
+
246
+ seq_lens is (B, next_n) per-row effective lens for native spec decode
247
+ and (B, 1) otherwise, in which case per-row lens are derived the same
248
+ way as the other decode top-k kernels.
249
+ """
250
+ if seq_lens.numel() == num_rows:
251
+ row_ends = seq_lens.reshape(-1)
252
+ else:
253
+ next_n_offsets = torch.arange(
254
+ next_n, dtype=torch.int32, device=seq_lens.device
255
+ )
256
+ row_ends = (
257
+ (seq_lens.reshape(-1, 1) - next_n + 1 + next_n_offsets)
258
+ .clamp_(min=0)
259
+ .reshape(-1)
260
+ )
261
+ assert row_ends.dtype == torch.int32
262
+ return row_ends
263
+
264
+ def forward(
265
+ self,
266
+ logits: torch.Tensor,
267
+ seq_lens: torch.Tensor,
268
+ next_n: int,
269
+ topk_indices: torch.Tensor,
270
+ topk_tokens: int,
271
+ max_seq_len: int,
272
+ ) -> None:
273
+ """Run the resolved decode top-k implementation, writing into
274
+ topk_indices (int32, -1 fill for rows shorter than topk_tokens)."""
275
+ backend = self.resolve_backend(logits, topk_tokens, logits.shape[0])
276
+ if backend == "deep_select":
277
+ row_ends = self._row_ends(seq_lens, next_n, logits.shape[0])
278
+ deep_select_topk(logits, topk_tokens, end=row_ends, output_idx=topk_indices)
279
+ elif backend == "cooperative":
280
+ (topk_workspace,) = current_workspace_manager().get_simultaneous(
281
+ ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
282
+ )
283
+ torch.ops._C.cooperative_topk(
284
+ logits,
285
+ seq_lens,
286
+ topk_indices,
287
+ topk_workspace,
288
+ topk_tokens,
289
+ max_seq_len,
290
+ )
291
+ elif backend == "persistent":
292
+ (topk_workspace,) = current_workspace_manager().get_simultaneous(
293
+ ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
294
+ )
295
+ torch.ops._C.persistent_topk(
296
+ logits,
297
+ seq_lens,
298
+ topk_indices,
299
+ topk_workspace,
300
+ topk_tokens,
301
+ logits.shape[1],
302
+ )
303
+ elif backend == "flashinfer":
304
+ # Deferred: importing flashinfer initializes CUDA at import time.
305
+ from flashinfer.topk import top_k_ragged_transform
306
+
307
+ row_ends = self._row_ends(seq_lens, next_n, logits.shape[0])
308
+ offsets = torch.zeros(
309
+ logits.shape[0], dtype=torch.int32, device=logits.device
310
+ )
311
+ # top_k_ragged_transform selects within [0, row_ends[i]) per
312
+ # row and -1-fills past the row length.
313
+ indices = top_k_ragged_transform(logits, offsets, row_ends, topk_tokens)
314
+ topk_indices.copy_(indices)
315
+ elif backend == "torch":
316
+ # Debug reference: mask everything past each row's end, then topk.
317
+ row_ends = self._row_ends(seq_lens, next_n, logits.shape[0])
318
+ cols = torch.arange(logits.shape[1], device=logits.device)
319
+ masked = logits.masked_fill(
320
+ cols.unsqueeze(0) >= row_ends.unsqueeze(1), float("-inf")
321
+ )
322
+ indices = masked.topk(topk_tokens, dim=-1).indices
323
+ in_range = torch.arange(topk_tokens, device=logits.device).unsqueeze(
324
+ 0
325
+ ) < row_ends.unsqueeze(1)
326
+ indices = torch.where(in_range, indices, -1)
327
+ topk_indices.copy_(indices)
328
+ else:
329
+ assert backend == "per_row", f"unknown topk backend: {backend}"
330
+ ops.top_k_per_row_decode(
331
+ logits,
332
+ next_n,
333
+ seq_lens,
334
+ topk_indices,
335
+ logits.shape[0],
336
+ logits.stride(0),
337
+ logits.stride(1),
338
+ topk_tokens,
339
+ )
vllm/model_executor/layers/sparse_attn_indexer.py
+17
-51
@@ -20,6 +20,10 @@ from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
20
from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
21
select_candidate_blocks as _select_candidate_blocks,
22
)
23
+from vllm.model_executor.layers.indexer_topk import (
24
+ RADIX_TOPK_WORKSPACE_SIZE,
25
+ get_indexer_topk,
26
+)
27
from vllm.model_executor.layers.quantization.utils.quant_utils import (
28
get_fp8_min_max,
29
)
@@ -46,8 +50,6 @@ from vllm.v1.worker.workspace import current_workspace_manager
50
51
logger = init_logger(__name__)
52
49
-RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
50
-
53
# MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32.
54
MXFP4_BLOCK_SIZE = 32
55
@@ -343,6 +345,7 @@ def sparse_attn_indexer(
345
candidate_blocks: torch.Tensor | None = None,
346
candidate_block_size: int = 0,
347
candidate_write: bool = False,
348
+ topk_backend: str = "auto",
349
) -> torch.Tensor:
350
# careful! this will be None in dummy run
351
forward_context = get_forward_context()
@@ -750,56 +753,16 @@ def sparse_attn_indexer(
753
)
754
topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
755
753
- use_cooperative_topk = (
754
- current_platform.is_cuda()
755
- and topk_tokens in (512, 1024, 2048)
756
- and num_rows <= 64
757
- and logits.stride(0) % 4 == 0 # TMA 16-byte alignment
758
- and current_platform.has_device_capability(90)
759
- and not current_platform.is_device_capability_family(120)
760
- )
761
- use_persistent_topk = current_platform.is_cuda() and topk_tokens in (
762
- 512,
763
- 1024,
764
- 2048,
756
+ # The backend comes from the layer (config is only readable at model
757
+ # construction); dispatchers are cached per backend.
758
+ get_indexer_topk(topk_backend)(
759
+ logits,
760
+ seq_lens,
761
+ next_n,
762
+ topk_indices,
763
+ topk_tokens,
764
+ attn_metadata_narrowed.max_seq_len,
765
)
766
- if use_cooperative_topk:
767
- workspace_manager = current_workspace_manager()
768
- (topk_workspace,) = workspace_manager.get_simultaneous(
769
- ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
770
- )
771
- torch.ops._C.cooperative_topk(
772
- logits,
773
- seq_lens,
774
- topk_indices,
775
- topk_workspace,
776
- topk_tokens,
777
- attn_metadata_narrowed.max_seq_len,
778
- )
779
- elif use_persistent_topk:
780
- workspace_manager = current_workspace_manager()
781
- (topk_workspace,) = workspace_manager.get_simultaneous(
782
- ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
783
- )
784
- torch.ops._C.persistent_topk(
785
- logits,
786
- seq_lens,
787
- topk_indices,
788
- topk_workspace,
789
- topk_tokens,
790
- logits.shape[1],
791
- )
792
- else:
793
- ops.top_k_per_row_decode(
794
- logits,
795
- next_n,
796
- seq_lens,
797
- topk_indices,
798
- num_rows,
799
- logits.stride(0),
800
- logits.stride(1),
801
- topk_tokens,
802
- )
766
767
if decode_metadata.global_seq_lens is not None:
768
_merge_dcp_topk_global(
@@ -851,6 +814,7 @@ def sparse_attn_indexer_fake(
814
candidate_blocks: torch.Tensor | None = None,
815
candidate_block_size: int = 0,
816
candidate_write: bool = False,
817
+ topk_backend: str = "auto",
818
) -> torch.Tensor:
819
return topk_indices_buffer
820
@@ -917,6 +881,7 @@ class SparseAttnIndexer(CustomOp):
881
# than threading them through per-step metadata.
882
vllm_config = get_current_vllm_config()
883
parallel_config = vllm_config.parallel_config
884
+ self.topk_backend = vllm_config.kernel_config.sparse_indexer_topk_backend
885
self._parallel_config = parallel_config
886
self.dcp_world_size = parallel_config.decode_context_parallel_size
887
self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
@@ -1026,6 +991,7 @@ class SparseAttnIndexer(CustomOp):
991
candidate_blocks=self.candidate_blocks,
992
candidate_block_size=self.candidate_block_size,
993
candidate_write=self.candidate_write,
994
+ topk_backend=self.topk_backend,
995
)
996
997
def forward_xpu(