[Feat][Kernel] Add an opt-in deterministic FlashInfer TopK backend

[LopezCastroRoberto]

Motivation

Sparse-attention index selection can be non-deterministic when values tie at the TopK boundary.

This adds an opt-in backend for users who need deterministic index selection, without changing the vLLM native TopK default or its performance characteristics.

This PR does not change default inference behavior and does not claim a model-quality improvement; its purpose is reproducible TopK selection for those who need this functionality.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[jahnclawdmonet]

Ran this PR on two DGX Sparks (GB10, sm_121, drivers 580.159.03 and 580.95.05) at TP=2 over Ray, the same topology and flags as #54521. With FlashInfer 0.6.18 the backend doesn't run on this GPU: the server loads the model and profiles, then dies in vLLM's post-profiling kernel warmup (v1/worker/gpu/warmup.py:345, the prefill execute_model), inside the QSA indexer's ragged prefill top-k, before the first request.

Setup: vllm/vllm-openai:nightly-9ea8f3ffc354901b740f0b31988900897b7221d7 (vLLM 0.28.1rc1.dev516+g9ea8f3ffc, torch 2.13.0+cu130, FlashInfer 0.6.18) with this PR's seven vllm/ files at bccc06c applied over the installed package and nothing else changed in the image; Qwen/Qwen3.8-Flash-Next-FP8 at 236dfdf2; --tensor-parallel-size 2 --distributed-executor-backend ray --gpu-memory-utilization 0.84 --max-model-len 16384 --max-num-seqs 1 --max-num-batched-tokens 16384 --enforce-eager --hf-overrides '{"text_config":{"indexer_budget":8192}}' --no-enable-prefix-caching --no-enable-flashinfer-autotune --dsa-topk-backend flashinfer, tie-break left at small. The flags are picked up (non-default args ... 'dsa_topk_backend': 'flashinfer' in the log; --help=all lists both options).

The failure, identical on both ranks at the same second, four seconds after Warmed up Qwen4Exp QSA sparse attention kernels (frames omitted at the ... marks, caret lines dropped):

  ...
  File ".../vllm/v1/worker/gpu_worker.py", line 800, in compile_or_warm_up_model
    warmup_kernels(self.model_runner, self.execute_model, self.sample_tokens)
  ...
  File ".../vllm/v1/worker/gpu/warmup.py", line 345, in _warmup_kernels
    worker_execute_model(prefill_output)
  ...
  File ".../vllm/v1/worker/gpu/model_runner.py", line 1849, in execute_model
    model_output = self.model(**model_inputs)
  ...
  File ".../vllm/models/qwen4_exp/nvidia/qsa.py", line 374, in _run_qsa
    selected = self.indexer(
  ...
  File ".../vllm/models/qwen4_exp/nvidia/indexer_qsa.py", line 449, in forward
    qsa_select_paged_prefill(
  File ".../vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py", line 665, in qsa_select_paged_prefill
    _topk(
  File ".../vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py", line 488, in _topk
    flashinfer_deterministic_topk(
  File ".../vllm/model_executor/layers/sparse_attn_topk.py", line 50, in flashinfer_deterministic_topk
    indices = top_k_ragged_transform(
  File ".../flashinfer/api_logging.py", line 2333, in _auto_dump_wrapper
    return _inner(*args, **kwargs)
  File ".../flashinfer/topk.py", line 981, in top_k_ragged_transform
    get_topk_module().radix_topk_ragged_transform(
  File ".../flashinfer/topk.py", line 329, in radix_topk_ragged_transform
    module.radix_topk_ragged_transform(
  File "python/tvm_ffi/cython/function.pxi", line 968, in tvm_ffi.core.Function.__call__
  File "<unknown>", line 0, in __tvm_ffi_radix_topk_ragged_transform
  File "/workspace/csrc/topk.cu", line 269, in void radix_topk_ragged_transform(...)
RuntimeError: Check failed: (status == cudaSuccess) is false: TopKRaggedTransform failed with error code operation not supported

Where that comes from in FlashInfer 0.6.18: TopKRaggedTransformDispatch (include/flashinfer/topk.cuh:3672-3681) returns cudaErrorNotSupported before launching anything when the call asks for dsa_graph_safe or a tie-break and either k > FILTERED_TOPK_MAX_K (2048) or CanImplementFilteredTopK() is false. Both modes exist only in the FilteredTopK kernel. CanImplementFilteredTopK() (topk.cuh:3540-3549) requires cudaDevAttrMaxSharedMemoryPerMultiprocessor >= FILTERED_TOPK_SMEM_DYNAMIC, which is 2 x 16384 x 4 = 131072 bytes (topk.cuh:2507-2511). GB10 reports 102400 bytes for that attribute (per-block opt-in 101376), so the check fails for every k, and FlashInfer's own flashinfer.topk.can_implement_filtered_topk() returns False on this device. flashinfer_deterministic_topk always passes dsa_graph_safe=True together with a tie-break, so the call never reaches a kernel here. k is not the trigger here: the QSA indexer passes block_topk = indexer_budget // indexer_compress_ratio (_topk in qsa_indexer.py; the ratio is 4 in this checkpoint's config.json), so 8192 // 4 = 2048, and the check is a strict >. With this ratio an indexer_budget of 8196 or more pushes block_topk to 2049 and returns the same error on any GPU; the DeepSeek indexer in the PR passes topk_tokens undivided, so there the limit sits at 2048 tokens.

Standalone in the same image on the same GPU, top_k_ragged_transform on a [4, 4096] float32 input, k = 2048:

call	result
deterministic=True, tie_break=1, dsa_graph_safe=True (the PR's call)	operation not supported
deterministic=True, tie_break=1, dsa_graph_safe=False	operation not supported
deterministic=True, tie_break=0, dsa_graph_safe=True	operation not supported
deterministic=True, tie_break=0, dsa_graph_safe=False (radix path)	runs; every row's index set equals torch.topk
deterministic=False (default)	runs; every row's index set equals torch.topk
the PR's call with k = 512	operation not supported

So the constructor check in indexer_qsa.py (has_flashinfer(), which checks for the package and for cubins or nvcc, with no device component) passes on a GPU where the backend can't work, and the problem surfaces as an engine death in warmup with nothing pointing at the shared-memory requirement. A guard at construction would need both halves of FlashInfer's condition: flashinfer.topk.can_implement_filtered_topk() (it loads the top-k module), and the k the indexer is going to pass being at most 2048 (FILTERED_TOPK_MAX_K is not exposed in Python). Falling back to the deterministic radix path is the other option, but that path is only selected when tie_break is TopKTieBreak::None (0) and dsa_graph_safe is False (topk.cuh:3577-3584, 2222-2227), so it drops both properties the backend is offered for.

For the native path I have one launch of each: the stock image, and the PR files with --dsa-topk-backend native. Same setup, 8 identical greedy requests per row (temperature 0, max_tokens 300), distinct completions out of 8:

prompt tokens	stock nightly	PR files, --dsa-topk-backend native
5731	1	1
6953	1 (2 requests, then my host-memory gate stopped the row)	1
7785	1	1
8591	5	3
9839	not measured (earlyoom killed the worker on the first request)	6

Both launches show the #54521 pattern on this nightly at TP=2: byte-identical below the 8192 budget, several distinct completions above it. The completions themselves differ between the two launches even on the identical rows. With one launch per arm I can't separate the PR overlay from launch state. I can attach the probe and the logs.

[jschmied]

@jahnclawdmonet's shared-memory finding also explains a failure we reported earlier without a cause, and
we have two measurements since that bear on the PR's premise.

Third GB10, and the mechanism names our earlier symptom. On 2026-09-08 we reported on #55122 that
this backend does not start on sm_121 — TopKRaggedTransform failed with error code operation not supported at engine init, on a single GB10 at TP=1, with the native arm on the same build fine.
We could not say why. @jahnclawdmonet's number accounts for it: this device reports

NVIDIA GB10, sm_121, 48 SMs
shared memory per SM: 102400 bytes      (FlashInfer 0.6.18 requires >= 131072)

So the failure is not TP- or Ray-specific — TP=1 and TP=2 fail for the same structural reason, on three
GB10s across at least three driver versions. Anything gated on 131072 bytes/SM is unavailable on this
whole hardware family, not just misconfigured on it.

On the premise, and we say this about our own PR first. The motivation here is that "index selection
can be non-deterministic when values tie at the TopK boundary". On this workload we can no longer find
that tie. Two runs on Qwen3.8-Flash-Next NVFP4, GB10, TP=1, MTP-3, k = 512 blocks:

1. The boundary does not tie. A census inside _topk over 400 instrumented calls: 87,257 rows, of
   which 6,192 performed a real selection (the rest had fewer visible blocks than the budget, so
   everything was taken). Rows with any value equal to the k-th: 0. Not "tied but resolved
   consistently" — never tied. Exact ties between float32 logits out of a real GEMM are simply rare.

2. The divergence is in the scores, not the selection. Hashing the input scores and the selected
   index set per call across 7 byte-identical greedy requests: on a stock server, all 13 comparable
   prefill calls have different input scores, and zero have identical scores with a different
   selection. With our four determinism fixes on, all 13 are bit-identical on both. Boundary gaps
   (kth − (k+1)th) are min 4.58e-05, median 3.20e-04, so a perturbation of that order is enough to
   reorder the selection — which is @rybruscoe's reading on [Bug]: Qwen3.8-Flash-Next: greedy decoding is non-deterministic from persistent_topk in prefill when prompt length nears indexer_budget (sm121/GB10) #54521, and it puts the cause upstream of
   the selection kernel.

We published exactly this about our own #55122 before writing it here, and we think it applies the
same way: a deterministic tie-break is kernel correctness under ties, not a route to end-to-end
reproducibility on traffic where the boundary never ties. That does not make this PR wrong — your body
already declines to claim a quality or performance improvement, and correctness under a rare condition
is worth having. It does mean a user adopting it for reproducibility on this model would not get it, and
@jahnclawdmonet's native-backend result is consistent with that: still 3–6 distinct completions of 8
above the budget, clean below it.

Bounds, so this is not read as more than it is. Both runs are prefill-only comparisons on one model
on one GPU family. Decode calls could not be aligned at all in the stock arm, because once outputs
diverge the MTP acceptance changes and so does the number of decode steps — that drift is itself a
symptom rather than a nuisance. And a tie rate indistinguishable from zero at ~6×10³ selecting rows is
not zero at 10⁶; a long-context, low-entropy workload might sample the boundary very differently.

TP=1 is the axis this thread does not otherwise have, and the box is free — happy to run any specific
cell you want, including a census on a workload you think should tie.

Data: notes/determinism-investigation.md findings det-190 and det-191.

AI assistance was used in preparing this comment; the measurements are ours and were reviewed before posting.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LopezCastroRoberto.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork