Skip to content

fix(qsa): enable trtllm-gen sparse decode on sm_121 (GB10 / DGX Spark) - #36649

Merged
yhyang201 merged 1 commit into
sgl-project:qwen4-main-squashedfrom
cshintov:fix/qsa-trtllm-decode-sm121
Aug 27, 2026
Merged

yhyang201 merged 1 commit into
sgl-project:qwen4-main-squashedfrom
cshintov:fix/qsa-trtllm-decode-sm121

Conversation

@cshintov

@cshintov cshintov commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

Motivation

_resolve_trtllm_sparse_decode() gates the trtllm-gen paged decode kernel on
is_sm100_supported(). GB10 (DGX Spark, sm_121) fails that check, so QSA decode
falls through to the FA4 cute varlen fallback.

As the resolver's own docstring notes, that fallback "runs a prefill-shaped kernel
at decode row counts". On sm_121 it does not merely run slower, it aborts during
CUDA graph capture:

File "flash_attn/cute/flash_fwd.py", line 393, in epilogue
cutlass._mlir._mlir_libs._site_initialize.<locals>.MLIRError: Operation creation failed:
error: unknown: expects `coord` and shape of view are weakly congruent, but got
'!cute.layout<"(?,?):(?{i64 div=8},1)">', '!cute.coord<"(_,_,?)">'

This makes Qwen3.8-Flash-Next unservable on DGX Spark, because
_forward_paged_attention is the only decode path QSA has.

flashinfer.decode.trtllm_batch_decode_with_kv_cache imports fine on sm_121, and
the repo already ships an is_sm121() helper (added for GB10). Admitting sm_121
to the gate selects the working kernel.

Modifications

python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py: widen the gate
from is_sm100_supported() to is_sm100_supported() or is_sm121(). One line;
no behaviour change on any other architecture.

Accuracy Test

Verified on 2x DGX Spark (GB10, sm_121, 128 GB unified each, ConnectX 200 Gb/s)
serving RadixArk/Qwen3.8-Flash-Next-NVFP4 at --tp-size 2 --nnodes 2:

  • before: crash at CUDA graph capture, server never reaches ready
  • after: server reaches ready and returns coherent generations

Benchmarking and Profiling

Same host, --mem-fraction-static 0.70, --ple-offload-embedding,
--moe-runner-backend flashinfer_cutlass, context 65536:

metric value
prefill (49,317-token prompt) ~1318 tok/s
decode, short prompt, no spec decode ~24.8 tok/s

No before-numbers are possible: the server does not start without this change.

Checklist

  • Change is covered by the existing QSA decode path; no new API surface.
  • No effect on sm_90/sm_100 code paths.

CI States

Latest PR Test (Base): ❌ Run #33052279089
Latest PR Test (Extra): ❌ Run #33052278925
Latest PR Test (AMD ROCm 7.2): ❌ Run #33052279077

The QSA decode resolver gated trtllm-gen paged decode on is_sm100_supported(),
so GB10 (sm_121) fell through to the FA4 cute varlen fallback. That fallback
runs a prefill-shaped kernel at decode row counts and aborts in its epilogue:

  MLIRError: expects `coord` and shape of view are weakly congruent, but got
  '!cute.layout<"(?,?):(?{i64 div=8},1)">', '!cute.coord<"(_,_,?)">'

sm_121 runs the same trtllm-gen decode kernel, so admit it via the existing
is_sm121() helper. Verified on 2x DGX Spark serving Qwen3.8-Flash-Next NVFP4
at tp=2: server reaches ready and decodes correctly; prefill ~1318 tok/s.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants