Repository navigation
[SM120] Support GLM DSA serving on RTX PRO 6000 - #29580
ronhuafeng wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the flashinfer_sparse_mla backend as the default prefill and decode attention mechanism for SM120 architectures, adding fallback implementations (such as a chunked PyTorch fallback for FP8 paged MQA logits and a Triton-based sparse MLA fallback) to support environments with older wheels. It also adds logic to disable shared expert fusion for GLM ModelOpt FP4 checkpoints when shared experts are excluded from quantization, and disables CUDA graphs by default during SM120 fallbacks to prevent excessive memory allocation. The review feedback highlights several critical issues, including multiple missing imports (Phase, log_info_on_rank0, envs, and QuantizationConfig) that would cause runtime NameError exceptions, a potential IndexError when parsing empty architecture lists, unreachable dead code in _forward_trtllm, and a performance optimization to increase pages_per_chunk to reduce loop overhead during eager execution.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| def _disable_cuda_graph_for_sm120_dsa_fallback(self): | ||
| """Keep the SM120 DSA torch fallback on eager execution by default. | ||
|
|
||
| The fallback is selected when older DeepGEMM/FlashInfer wheels do not | ||
| have native sm_120 paged-MQA/TRTLLM-MLA kernels. It is functional under | ||
| eager execution, while decode CUDA graph capture uses an over-allocated | ||
| static page table that makes the PyTorch fallback allocate/capture | ||
| impractically large tensors. | ||
| """ |
There was a problem hiding this comment.
The Phase enum is used in this method (e.g., Phase.PREFILL and Phase.DECODE), but it might not be imported in server_args.py. To prevent a potential NameError on startup, we should import Phase from sglang.srt.model_executor.cuda_graph_config.
| def _disable_cuda_graph_for_sm120_dsa_fallback(self): | |
| """Keep the SM120 DSA torch fallback on eager execution by default. | |
| The fallback is selected when older DeepGEMM/FlashInfer wheels do not | |
| have native sm_120 paged-MQA/TRTLLM-MLA kernels. It is functional under | |
| eager execution, while decode CUDA graph capture uses an over-allocated | |
| static page table that makes the PyTorch fallback allocate/capture | |
| impractically large tensors. | |
| """ | |
| def _disable_cuda_graph_for_sm120_dsa_fallback(self): | |
| """Keep the SM120 DSA torch fallback on eager execution by default. | |
| The fallback is selected when older DeepGEMM/FlashInfer wheels do not | |
| have native sm_120 paged-MQA/TRTLLM-MLA kernels. It is functional under | |
| eager execution, while decode CUDA graph capture uses an over-allocated | |
| static page table that makes the PyTorch fallback allocate/capture | |
| impractically large tensors. | |
| """ | |
| from sglang.srt.model_executor.cuda_graph_config import Phase |
There was a problem hiding this comment.
Addressed in the current head: Phase is imported from sglang.srt.model_executor.cuda_graph_config.
| log_info_on_rank0( | ||
| logger, | ||
| "GLM ModelOpt FP4 checkpoint excludes shared experts from " | ||
| "quantization, so they cannot be fused into routed FP4 expert slots. " | ||
| "Shared experts fusion optimization is disabled.", | ||
| ) |
There was a problem hiding this comment.
The utility function log_info_on_rank0 is used here but might not be imported in glm4_moe.py. To prevent a potential NameError at runtime, we should import it.
| log_info_on_rank0( | |
| logger, | |
| "GLM ModelOpt FP4 checkpoint excludes shared experts from " | |
| "quantization, so they cannot be fused into routed FP4 expert slots. " | |
| "Shared experts fusion optimization is disabled.", | |
| ) | |
| from sglang.srt.utils import log_info_on_rank0 | |
| log_info_on_rank0( | |
| logger, | |
| "GLM ModelOpt FP4 checkpoint excludes shared experts from " | |
| "quantization, so they cannot be fused into routed FP4 expert slots. " | |
| "Shared experts fusion optimization is disabled.", | |
| ) |
There was a problem hiding this comment.
Addressed in the current head: log_info_on_rank0 is imported from sglang.srt.utils before this call site.
| is_sm120_supported, | ||
| ) |
There was a problem hiding this comment.
The envs module is used in dsa_indexer.py (e.g., envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get()), but it is not imported in this file. This will cause a NameError during execution. We should import envs from sglang.srt.utils.
| is_sm120_supported, | |
| ) | |
| is_sm120_supported, | |
| ) | |
| from sglang.srt.utils import envs |
There was a problem hiding this comment.
Addressed in the current head: envs is imported from sglang.srt.environ in dsa_indexer.py.
|
|
||
| kv_value = kv_value_raw.contiguous().view(dtype=FP8_DTYPE).to(torch.float32) | ||
| kv_value = kv_value.view(batch_size, max_padded_seq, head_dim) | ||
| pages_per_chunk = 4 |
There was a problem hiding this comment.
The pages_per_chunk is set to a very small constant value of 4. While this keeps the memory footprint extremely low during CUDA graph warmup, it introduces significant Python loop overhead and kernel launch latency during eager execution (e.g., for a sequence length of 2048 and block size of 16, the loop will run 32 times, launching hundreds of PyTorch kernels).
Consider increasing pages_per_chunk to a larger value like 64 or 128. This still keeps the memory footprint very small (a few tens of megabytes) while dramatically reducing the loop iterations and CPU overhead in the latency-sensitive decode path.
| pages_per_chunk = 4 | |
| pages_per_chunk = 64 |
There was a problem hiding this comment.
This thread is outdated against the current head. The pages_per_chunk = 4 loop is no longer present in python/sglang/srt/layers/attention/dsv4/indexer.py; the current implementation gathers the page table slice and runs the batched computation directly, so there is no chunk-size constant to increase here.
I am leaving this as no-code-change rather than adding a new tuning knob that is not used by the current code.
| model_arch = getattr( | ||
| model_runner.model_config.hf_config, "architectures", [None] | ||
| )[0] |
There was a problem hiding this comment.
If architectures is an empty list [] in the model's Hugging Face config, calling [0] on it will raise an IndexError: list index out of range and crash the server startup. It is safer to handle empty lists defensively.
architectures = getattr(
model_runner.model_config.hf_config, "architectures", []
)
model_arch = architectures[0] if architectures else NoneThere was a problem hiding this comment.
Good catch. Updated the current head to handle an empty architectures list defensively:
architectures = getattr(model_runner.model_config.hf_config, "architectures", [])
model_arch = architectures[0] if architectures else NoneThis keeps the GLM path unchanged when the architecture is present and falls back to "auto" instead of raising IndexError when the list is empty.
Local checks: python3 -m py_compile python/sglang/srt/layers/attention/dsa_backend.py and git diff --check passed.
| def _modelopt_fp4_excludes_shared_experts( | ||
| quant_config: Optional[QuantizationConfig], | ||
| ) -> bool: |
There was a problem hiding this comment.
Using QuantizationConfig as a type annotation can raise a NameError at import time if QuantizationConfig is not imported in glm4_moe.py. Since type annotations are optional, we can remove it or use a string literal to be safe.
| def _modelopt_fp4_excludes_shared_experts( | |
| quant_config: Optional[QuantizationConfig], | |
| ) -> bool: | |
| def _modelopt_fp4_excludes_shared_experts( | |
| quant_config, | |
| ) -> bool: |
There was a problem hiding this comment.
Addressed in the current head: QuantizationConfig is imported, so this annotation is defined at import time.
| extra_kwargs = {} | ||
| if flashinfer_backend == "sparse" and _flashinfer_trtllm_mla_accepts_kwarg( | ||
| "kv_scale_format" | ||
| ): | ||
| extra_kwargs["kv_scale_format"] = self.flashinfer_sparse_mla_kv_scale_format |
There was a problem hiding this comment.
Since _forward_trtllm returns early on line 2496 when flashinfer_backend == "sparse", the check if flashinfer_backend == "sparse" on line 2514 is unreachable dead code. Consequently, extra_kwargs will always be empty, and backend in trtllm_batch_decode_with_kv_cache_mla will always be "trtllm-gen". We can simplify this and remove the unreachable code.
| extra_kwargs = {} | |
| if flashinfer_backend == "sparse" and _flashinfer_trtllm_mla_accepts_kwarg( | |
| "kv_scale_format" | |
| ): | |
| extra_kwargs["kv_scale_format"] = self.flashinfer_sparse_mla_kv_scale_format | |
| extra_kwargs = {} |
There was a problem hiding this comment.
This thread appears outdated. In the current head, the sparse path still populates extra_kwargs["kv_scale_format"] and passes it into the backend call, so this simplification does not apply to the current code.
5312378 to
66cb38c
Compare
|
Follow-up pushed in
I kept Validation after this follow-up:
|
8af780a to
de6c667
Compare
de6c667 to
1a5f115
Compare
1a5f115 to
35450ad
Compare
|
Independently validated on a separate 8× RTX PRO 6000 Blackwell Server (SM120) box, base image 1. End-to-end correctness ✅ All auto-selects work as described: 2. Decode CUDA graph: With the default Workaround: (Tested against this PR's diff at 3. Throughput (TP=8, 2K/512 fixed-len,
Eager TPOT is flat ~90 ms (host overhead floor). With 4. Minor: The nvidia checkpoint ships Happy to re-test if you push a decode-capture fix. |
35450ad to
cf638d6
Compare
|
Thanks for the independent run; I pushed follow-up commit Changes:
This should address the Follow-up validation in this workspace:
I did not rerun the SM120 runtime here, so an independent retest on the |
which version do you use to get vllm b12x performance data? in this repo rtx6kpro,it seems that vllm is way faster than sglang |
|
@zhangvia Fair point — should have been clearer about what that comparison is and isn't. The 442.7 baseline is
The ctx0 vs 2K/512 workload difference alone accounts for a large chunk of 703 vs 443 (prefill steps drag down output-token throughput in the bench-serve metric). The allreduce config is the other suspect — my run inherited So the +29% is strictly "same box, same 2K/512 bench-serve workload, my vLLM config vs my sglang config" — not a claim that sglang beats the tuned dark-devotion ctx0 numbers. I don't have an apples-to-apples ctx0 number for sglang, and my vLLM baseline is clearly not the best-known vLLM config for this card. Thanks for the pointer to the v12 doc. |
|
Reran the exact SM120 runtime validation on current head Environment summary:
Checks before server launch:
Server launch used the default GLM command from the PR body, without workaround flags: python3 -m sglang.launch_server \
--model-path nvidia/GLM-5.2-NVFP4 \
--tensor-parallel-size 8 \
--quantization modelopt_fp4 \
--tool-call-parser glm47 \
--reasoning-parser glm45 \
--trust-remote-code \
--mem-fraction-static 0.8 \
--host 0.0.0.0 \
--port 30000Result:
I also updated the PR body so the validation gaps no longer say the latest |
|
Retested Health 290 s, GSM8K-style 5/5, bench reproduces the earlier manual-flag run within <0.5%:
So the cu13-image hang is covered by the new default. LGTM from the SM120 runtime side. |
cf638d6 to
330483a
Compare
|
Thanks for the retest on I pushed amended head
I also updated the PR body with this root cause and validation. Validation for this follow-up:
Exact SM120 runtime checks:
So the default remains |
330483a to
79e42d2
Compare
79e42d2 to
005ffc6
Compare
|
One more follow-up pushed in Change from the previous head:
Validation for the latest head:
Exact SM120 runtime default launch, with no manual
I also updated the PR body to make the current behavior explicit: default |
|
This is a spam PR |
|
@Fridge003 Sorry for the noise here, and sorry for taking maintainer time on this. I understand why this came across as a broad parallel PR to #26928. I should have searched more carefully for existing PRs related to this issue at the beginning, found #26928 earlier, and coordinated there before opening this as a ready standalone PR and requesting many reviewers. A bit of context on how this happened: I started from #29562, where the first blocker looked like a GLM-5.2-NVFP4 checkpoint/loading issue on SM120. As I debugged it on exact RTX PRO 6000 hardware, each fix exposed the next blocker: shared-expert fusion, then the SM120 DSA backend path, then the FlashInfer/DeepGEMM wheel gap, CUDA graph memory behavior, and finally the One important part of that history is the dependency setup. I found that the released FlashInfer / DeepGEMM wheels did not yet contain the SM120 support needed for this path, so I added a temporary source-build path pinned to upstream GitHub commit hashes. I realize that pinning dependency commits in a PR build path is not ideal and should not be the long-term shape for SGLang. My original intent was narrower: to give the #29562 reporter a reproducible way to try the fix in their own environment before official wheels are ready, while documenting the release gap. In hindsight, I let a focused bug investigation grow into a broad PR. My intent was not to replace or compete with #26928, but bundling all of these pieces together, especially with a default backend change, temporary source-pinned dependencies, and many review requests, created unnecessary review burden and made the relationship to the existing SM120 work unclear. I’ll keep this PR closed and won’t reopen it as-is. If any part of it is useful, I’ll split it into much smaller GLM-5.2-specific follow-ups or issue validation work around #26928. |
Motivation
Fixes #29562.
This PR adds a native SM120 serving path for GLM DSA / GLM-5.2-NVFP4 on RTX PRO
6000 Blackwell. With the default GLM launch, the patched path starts the TP8
server, captures CUDA graphs, passes
/health, and completes a minimal/generaterequest on exact SM120 hardware.The reporter path had multiple independent blockers:
mlp.shared_expertsfrom FP4quantization, but shared-expert fusion remapped those unquantized tensors
into routed packed FP4 expert slots. That caused the initial
3072vs6144loader mismatch.trtllmDSAbackends for fp8 KV on SM120. The underlying FlashInfer TRTLLM-gen MLA path
reports unsupported architecture on SM120.
deep_gemm.get_paged_mqa_logits_metadata, which is unsupported on SM120 inthe currently pinned
sgl-deep-gemmwheel.max_bs=512default OOMed during graph capture on 96GB RTX PRO 6000. Anindependent
dev-cu13run also foundfulldecode graph capture can hang.The current patch keeps SGLang's default
fulldecode graph backend forconsistency with the existing decode path, caps decode
max_bsat 256 bydefault for this SM120 GLM DSA case, and still lets users explicitly select
tc_piecewiseas an alternative decode CUDA graph mode.index_topk_pattern: null; the DSA IndexShare layerlayout is encoded in
indexer_types, so SGLang needs to derive the same F/Ssharing pattern from checkpoint metadata instead of requiring a manual model
config override.
Modifications
glm4_moe.py: detect ModelOpt FP4 checkpoints where shared experts are notquantized and automatically disable GLM shared-expert fusion for that case.
model_config.py: derive the DSA top-k sharing pattern fromcheckpoint-native
indexer_typeswhenindex_topk_patternis absent,preserving explicit pattern precedence and the existing frequency/offset
fallback.
server_args.py: addflashinfer_sparse_mlaas a DSA backend choice andselect it by default for fp8 DSA on SM120. For GLM DSA on SM120, preserve the
default
fulldecode CUDA graph backend, cap the default decode graphmax_bsto 256, and use the derived DSA top-k sharing pattern for the TBOguard. Hopper and SM100 paths keep their existing default backend selections.
dsa_backend.py: call FlashInfer sparse MLA with the SM120 contract requiredby the native kernels: BF16 query, packed uint8 KV, page size 64, int32 sparse
page tables, and
kv_scale_format="arbitrary_fp32"for GLM.forward_mla.py: keep TRTLLM-style FP8 RoPE fusion limited to the TRTLLMbackend so the sparse MLA path uses its own input contract.
Dockerfile: add optional repo+commit source build args for FlashInfer andDeepGEMM so the image can use the validated dependency commits before wheels
are released.
Validated source commits:
5f2bdc41f9ffecef9d8ed590e688e7c0f108504f(includes merged SM120 sparse MLA support from
feat(attention): add SM120 sparse MLA kernels flashinfer-ai/flashinfer#3395)
2073ddb2814892014c33ef4cd1c7d4c148baf1fe(SM120 support merge)
Released wheel gap checked on SM120:
flashinfer-python==0.6.13,flashinfer-cubin==0.6.13, andflashinfer-jit-cache==0.6.13+cu130import, but the releasedtrtllm_batch_decode_with_kv_cache_mlasignature still lackskv_scale_format, so this GLM sparse MLA path needs a FlashInfer sourcebuild until the next wheel release.
sgl-deep-gemmis still at 0.1.3, so the DeepGEMM SM120 metadata path needsthe source commit above until a new SGLang DeepGEMM wheel is published.
The Dockerfile source override path pins the validated repos by commit hash:
I also validated this repo+commit install sequence on the exact SM120 base
image: DeepGEMM was installed from the pinned GitHub commit as
deep_gemm 2.5.0+2073ddb; FlashInfer was installed from the pinned GitHubcommit with
FLASHINFER_CUDA_ARCH_LIST=12.0f, includingflashinfer-python,flashinfer-cubin, andflashinfer-jit-cache0.6.13. The final probe printedhas_kv_scale_format True,deepgemm_metadata_ok Tensor (189, 2) torch.int32 cuda, andrepo_commit_install_ok.Accuracy Tests
This PR changes the serving path for a previously unsupported SM120 GLM DSA
configuration. I did not run a full downstream accuracy benchmark such as GSM8K.
The issue-specific accuracy-relevant validation was functional generation on the
exact target hardware:
/health/generatecompleted with HTTP 200 and 16 completion tokensSM120 selection path plus existing non-SM120 routing behavior
FlashInfer sparse MLA correctness for the SM120 kernels is covered upstream by
FlashInfer's SM120 sparse MLA test suite in
flashinfer-ai/flashinfer#3395.
Speed Tests and Profiling
I did not run a throughput benchmark or profiler sweep for this patch. The
performance-relevant validation was that the SM120 path keeps CUDA graphs
enabled:
max_bs=512OOMed during graph capture on 96GB RTX PRO6000
dev-cu13validation foundfulldecode graph capture can hang,while
tc_piecewisedecode capture worksfulldecode graph backend for consistencywith the existing decode path and caps GLM DSA SM120 decode graph
max_bsto256
tc_piecewiseremains available as an alternative decode CUDA graphmode
Tests
Exact SM120 hardware:
580.159.03, CUDA13.0host /13.0.1containerlmsysorg/sglang:dev-glm52-nvfp42.11.0+cu130Baseline:
the
3072vs6144mismatch--disable-shared-experts-fusiongot past load/KV allocation, then failed indeep_gemm.get_paged_mqa_logits_metadata(...): Unsupported architecture--dsa-prefill-backend trtllm --dsa-decode-backend trtllmfailed with
TllmGenFmhaRunner ... Unsupported architecturePatched dependency probe:
Patched default launch:
Result:
prefill=flashinfer_sparse_mla,decode=flashinfer_sparse_mla368768tokens/rank,21.11 GB/rank)/healthpassed/generatecompleted with HTTP 200 and 16 completion tokensLatest SM120 rerun after the
tc_piecewisedecode graph follow-up, on PR headcf638d6e6b0b1138128196d2349da18460b06d44:has_kv_scale_format TrueandDeepGEMM metadata support
99 passed, 18 warnings, 10 subtests passedflashinfer_sparse_mlacuda_graph_config.decode.backenddefaulted totc_piecewisewithmax_bs=256GlmMoeDsaForCausalLMwithmodelopt_fp4/NVFP4368768tokens/rank,21.11 GB/rank)backend=tc_piecewisebackend=tc_piecewiseandbs=[1, 2, 4, ..., 256]/healthpassed/generatecompleted with HTTP 200 and 16 completion tokensUnsupported architecture,TllmGenFmhaRunner,Traceback,OutOfMemory,RuntimeError,ValueError,3072, and6144Local checks:
Follow-up after independent
dev-cu13validation of the native path:PYTHONPATH=python python3 -m pytest \ test/registered/unit/server_args/test_server_args.py \ test/registered/unit/models/test_glm4_moe.py \ test/registered/unit/models/test_deepseek_mla_forward.py -q # 99 passed, 18 warnings, 10 subtests passedSM120 runtime rerun for the earlier
tc_piecewisedefault follow-up:Result: selected
flashinfer_sparse_mla, defaulted decode CUDA graph backend totc_piecewise, completed prefill and decode CUDA graph capture, passed/health, returned HTTP 200 for minimal/generate, and had a clean failuregrep.
Latest follow-up for GLM-5.2
indexer_typesmetadata and defaultfulldecodegraph validation, on PR head
005ffc65dc0f72d42a1436c31ef129ed98cccc58:index_topk_pattern: nulland 78indexer_typesentries describing thefull/sharedlayer layout97 passed, 18 warnings, 10 subtests passedflashinfer_sparse_mlacuda_graph_config.decode.backendremainedfullwithmax_bs=256backend=tc_piecewisebackend=fullandbs=[1, 2, 4, ..., 256]/healthpassed/generatecompleted with HTTP 200 and 16 completion tokensindex_topk_patternoverride or manual decode backend flag was usedfull, explicitfull, explicittc_piecewise,and explicit disabled decode CUDA graph behavior
fullchecks also passed at--cuda-graph-max-bs-decode 32and at the PR's automatic
max_bs=256capGitHub Actions:
lintpassed on headcf638d6e6b0b1138128196d2349da18460b06d44005ffc65dc0f72d42a1436c31ef129ed98cccc58needs thenormal GitHub check rerun
pr-gatefailures due missing maintainerrun-ciauthorization/label, not test job failures
Validation gaps:
SM120 selection path plus existing non-SM120 routing behavior.
sgl-deep-gemmwheels still need follow-up releases.Until then, the Dockerfile source-pinned build args above are the validated
way for the SGLang image to carry this path.
Checklist
git diff --checkand GitHublintpassed.indexer_typespattern derivation, and non-SM120 routing.Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #28357834679
Latest PR Test (Extra): ❌ Run #28357834575
CI States
Latest PR Test (Base): ❌ Run #28352062088
Latest PR Test (Extra): ❌ Run #28352061942