[Bug]: DSA persistent_topk crashes on GB10/sm121 at long context -- FilteredTopK fallback needs >=128KB smem but sharedMemPerBlockOptin is 101376B; fix by honoring optin smem on the cooperative path

[SignoreUnito]

Your current environment

• GPU: NVIDIA GB10 (DGX Spark), Blackwell, compute capability sm121 (12.1)
• Device props: multiProcessorCount = 48; sharedMemPerBlock (default) = 49152 (48 KB); sharedMemPerBlockOptin = 101376 B (~99 KB). For contrast, Hopper / datacenter-Blackwell optin smem is ~227 KB.
• vLLM: jasl fork build (sm12x / Blackwell DSA + sparse-MLA), build tag ~v20260613.dev331.
• Model: GLM-5.2 (model_type=glm_moe_dsa, 744B MoE, NVFP4/modelopt_fp4), 78 layers, DSA with index_topk=2048. Maps to the DeepSeek-V2/V3 MLA family in vLLM.
• Topology: 8x GB10, pipeline parallelism (PP shards the MLA KV cache by layers).

Note: collect_env.py output omitted -- this is a kernel launch-logic bug, fully determined by the device properties above (SM count and sharedMemPerBlockOptin), and is independent of the exact driver/torch build.

🐛 Describe the bug

Summary

On GB10 (DGX Spark, sm121), serving a DSA model (glm_moe_dsa / DeepSeek-V3.2/V4 MLA family, index_topk=2048) at long context aborts in the sparse-attention indexer's top-k launcher. The cooperative radix path in csrc/libtorch_stable/topk.cu (launch_persistent_topk) selects shared memory based only on num_rows. For the long-sequence (max_seq_len > RADIX_THRESHOLD) cooperative path this under-sizes the smem chunk, which inflates ctas_per_group past the cooperative resident cap, so the launcher takes the FilteredTopK fallback -- and that fallback STD_TORCH_CHECKs for >= 128*1024 bytes of smem, which GB10's sharedMemPerBlockOptin (101376 B, ~99 KB) cannot satisfy. Net effect: DSA long-context inference is impossible on GB10 with the default smem caps.

This logic is the same in upstream vllm-project/vllm main and in this jasl fork, so the bug is inherited from upstream. The public GB10 DSA recipes route around it by using a Triton top-k indexer instead of the CUDA persistent_topk kernel, which is why this CUDA-path failure has not surfaced before.

This is downstream of vllm-project#45317 / vllm-project#46055 (the sm121 backend-selection gate). With a backend selected, this top-k smem crash is the next wall on long context.

Crash

persistent_topk would oversubscribe and the FilteredTopK fallback requires
>=128KB smem per block (have 101376). total_ctas=N > num_sms*occupancy=48
(TopK=2048, ...)

Root cause

In csrc/libtorch_stable/topk.cu, launch_persistent_topk:

1. max_smem_per_block = device_prop->sharedMemPerBlockOptin; (101376 on GB10).

2. effective_max_smem is keyed only on num_rows:

   • num_rows <= 4 -> min(max_smem_per_block, kSmemMedium) (kSmemMedium = 35968)
   • num_rows <= 8 -> min(max_smem_per_block, 48*1024)
   • else -> max_smem_per_block

   The small caps are a decode-path occupancy optimization.

3. needs_cooperative = (max_seq_len > RADIX_THRESHOLD) (RADIX_THRESHOLD = 32768) is computed, but is used only for the oversubscribe / resident-CTA guard -- it does not influence effective_max_smem.

4. max_chunk_elements = (effective_max_smem - kFixedSmemLarge) / 4, and ctas_per_group = ceil(seq_len / max_chunk_elements).

For small num_rows (decode) plus a long sequence, effective_max_smem is capped at ~36 KB / 48 KB. That shrinks max_chunk_elements, so ctas_per_group exceeds the cooperative resident cap (num_sms * occupancy = 48 on GB10). The launcher then falls back to FilteredTopK, which checks:

STD_TORCH_CHECK(max_smem_per_block >= 128 * 1024, "persistent_topk would oversubscribe ...");

GB10 has only 101376 B of optin smem (< 131072), so the check fails and the run aborts.

Minimal fix

On the long-sequence cooperative path (max_seq_len > RADIX_THRESHOLD), use the full opt-in smem regardless of num_rows. The small num_rows caps should apply only to short, non-cooperative sequences.

const bool needs_cooperative =
    static_cast<uint32_t>(max_seq_len) > P::RADIX_THRESHOLD;

int effective_max_smem;
if (needs_cooperative) {
  // Long-sequence cooperative radix path: chunk sizing depends on smem.
  // Use full opt-in smem so ctas_per_group stays within the resident cap;
  // the small num_rows caps are a decode occupancy tweak only relevant to
  // short (non-cooperative) sequences.
  effective_max_smem = max_smem_per_block;
} else if (num_rows <= 4) {
  effective_max_smem = std::min(max_smem_per_block, (int)P::kSmemMedium);
} else if (num_rows <= 8) {
  effective_max_smem = std::min(max_smem_per_block, 48 * 1024);
} else {
  effective_max_smem = max_smem_per_block;
}

With full optin on the long path: max_chunk_elements ~= (101376 - kFixedSmemLarge) / 4 ~= 24824, so at 1,048,576 context ctas_per_group = ceil(1048576 / 24824) = 43 <= 48. The cooperative launch fits and the >= 128KB FilteredTopK fallback is never reached.

Framing it as "honor sharedMemPerBlockOptin on the cooperative max_seq_len > RADIX_THRESHOLD path instead of assuming 128KB" makes it benefit any Blackwell part with sub-128KB optin smem, not just GB10.

Verification

With this fix, GLM-5.2 (glm_moe_dsa, 744B NVFP4) serves at max_model_len = 1048576 (native 1M) on 8x GB10 via pipeline parallelism:

• GPU KV cache: ~2.4M tokens (PP=8); ~1.2M tokens (TP=2 x PP=4).
• /health = 200; coherent generation observed.
• Single-stream throughput ~3-4.4 tok/s (pipeline-bubble bound at PP=8).

The PP-for-KV-sharding part is standard vLLM; the change here is only the GB10 smem-cap handling on the cooperative top-k path. These serving numbers are from a single local bring-up and have not been independently reproduced.

Caveats / open questions

• This is a host-side launch-logic workaround. Forcing full opt-in smem on the long path reduces occupancy for that path and trades some throughput; a maintainer may prefer a cleaner solution (e.g. sizing effective_max_smem from the actually-selected path, computing ctas_per_group against the cooperative budget directly, or a dedicated sub-128KB FilteredTopK fallback).
• Numbers above are for GB10 specifically (48 SM, 101376 B optin). The ceil(1048576/24824)=43<=48 headroom is tight; other sm12x parts with different SM counts / optin smem should be checked.
• On >=128KB-optin parts the intent is no behavior change (the else branch already returns full optin for the cooperative path), but that should be confirmed.

Happy to open a PR with the above branch, and to cross-link vllm-project#45317 / vllm-project#46055 (backend selection) and the persistent_topk hardening lineage.

---

Reported by Leo Sakaguchi (LifeLabo K.K.). A detailed walkthrough of this diagnosis and the fix -- with the GB10 numbers -- is on YouTube: https://www.youtube.com/@The-Pantheons

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[jasl]

Which model do you want to run?
There's another issue #20 wanna run GLM-5.1 NVFP4 (REAP)
And I'm thinking about GLM-5.2-NVFP4 and GLM-5.2-MXFP4