Support piecewise CUDA graph with NSA by nvjullin · Pull Request #23351 · sgl-project/sglang

nvjullin · 2026-04-21T08:15:41Z

Motivation

GLM-5/DSV3.2 currently doesn't allow piecewise CUDA graph due to incompatibilities in NSA attention backend and NSA indexer. This commit fixes the incompatibilities on the conditions

is not context parallel
is on CUDA

_store_index_k_cache and _get_topk_ragged are excluded from CUDA graph for now, will revisit whether they can be in CUDA graph in a future PR.

Benchmark results

Benchmark command

python3 bench_serving/benchmark_serving.py --backend sglang --host localhost --port 8080 --model zai-org/GLM-5-FP8 --num-prompts 1024 --ignore-eos --max-concurrency 128 --random-input-len 1024 --random-output-len 1024 --random-range-ratio 0.8 --use-chat-template --dataset-name random --save-result --result-filename results/concurrency_128.json

Modifications

In NSA code path:

Remove/fix torch.compile incompatible constructs
Correctly handling padding rows

Accuracy Tests

python3 -m sglang.test.run_eval --base-url http://localhost:8080 --model zai-org/GLM-5-FP8 --num-threads 128 --eval-name gpqa --max-tokens 128000 --num-examples 198 --repeat 8 --top-p 0.95 --temperature 1.0 --thinking-mode glm-45

Repeat: 8, mean: 0.862
Scores: ['0.848', '0.874', '0.874', '0.859', '0.848', '0.843', '0.869', '0.879']

Speed Tests and Profiling

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review and Merge Process

Ping Merge Oncalls to start the process. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): 🚫 Run #26311925143
Latest PR Test (Extra): 🚫 Run #26311925079

gemini-code-assist

Code Review

This pull request enables piecewise CUDA graph support for Native Sparse Attention (NSA) models by registering custom operators for Hadamard transforms, LayerNorm, and NSA indexing. It updates the context management and graph runners to handle NSA indexers and adjusts tensor slicing to accommodate padding in CUDA graphs. Feedback focused on correcting the output tensor dimensions in the TRT-LLM MLA decode path and adding necessary metadata checks and tensor slicing in the NSA indexer custom operations to prevent runtime errors.

gemini-code-assist · 2026-04-21T08:21:06Z

+        metadata = forward_batch.attn_backend.get_indexer_metadata(
+            layer_id, forward_batch
+        )


In piecewise CUDA graph mode, metadata can be None if the attention backend decides to skip the batch (e.g., if there are no tokens to process). We should check for None to avoid potential crashes in subsequent calls like _get_topk_ragged.

Suggested change

metadata = forward_batch.attn_backend.get_indexer_metadata(

layer_id, forward_batch

)

metadata = forward_batch.attn_backend.get_indexer_metadata(

layer_id, forward_batch

)

if metadata is None:

return

Fridge003 · 2026-04-21T22:45:36Z

            ), "FP8/MXFP8 Cutlass MoE is only supported with ep_size == 1"

-        # TODO(yuwei): Fix piecewise cuda graph support for bypassed topk MoE backends.
-        # Exception: GptOssForCausalLM wraps the entire MoE block in its own


Rather than removing this if block, we can add an exemption for dsa models here

The problem here is torch.compile compatibility instead of PCG compatibility. The fix was in sglang before the removal of the block, and it seems general to all models. Do you still think we need to guard it?

https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/moe/fused_moe_triton/layer.py#L978-L1003

https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/moe/fused_moe_triton/layer.py#L1176-L1200

I think we still need the guard: the custom op fixes Dynamo traceability, but it does not guarantee every flashinfer_trtllm bypassed-topk kernel is safe under PCG capture/replay. I opened a PR to restore the guard and add the exemption only for DSA models / already-supported paths.

#26718

Fridge003 · 2026-04-21T23:16:47Z

We can add a PCG test for GLM-5-fp4, registered as stage-c-test-4-gpu-b200

Fridge003 · 2026-04-21T23:19:54Z

For the gpqa accuracy test, please make sure temperature is 1.0, top_p is 0.95, --thinking-mode glm-45 is appended.
0.813 looks lower than expectation.

We also need another accuracy test for AIME2025, with the instruction here https://docs.sglang.io/docs/basic_usage/deepseek_v32#accuracy-test-with-aime-2025.
The number of turns can be larger than 4, since AIME2025 fluctuate a lot. The larger the better

nvjullin · 2026-04-24T10:13:59Z

We can add a PCG test for GLM-5-fp4, registered as stage-c-test-4-gpu-b200

Added gsm8k test. I understand it doesn't catch all indexer bugs, but it runs in ~10min. Let me know if you want a different test.

nvpohanh · 2026-05-22T01:50:08Z

  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/layers/utils/multi_platform.py", line 83, in forward
    return self._forward_method(*args, **kwargs)
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/layers/attention/dsa/dsa_indexer.py", line 1509, in forward_cuda
    k_cache_and_topk_result(
  File "/usr/local/lib/python3.10/dist-packages/torch/_ops.py", line 1269, in __call__
    return self._op(*args, **kwargs)
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/layers/attention/dsa/dsa_indexer.py", line 122, in k_cache_and_topk_result
    metadata = forward_batch.attn_backend.get_indexer_metadata(
AttributeError: 'ForwardBatch' object has no attribute 'attn_backend'

https://github.com/sgl-project/sglang/actions/runs/26256885165/job/77291831227?pr=23351
@nvjullin seeing this failure

nvjullin · 2026-05-22T02:58:22Z

Failure caused by bad merge from #25983. Fixed.

nvpohanh · 2026-05-22T05:03:31Z

[2026-05-22 04:54:32 TP6] Piecewise CUDA Graph failed with error: Tensor match failed for Tensor<8>[strides=<1>, dtype=int64, device=cuda:6] at /actions-runner/_work/sglang/sglang/python/sglang/jit_kernel/csrc/dsa/fused_store_index_cache.cuh:107
- Root cause: Size mismatch for shape#0('num_tokens'): expected 6 but got 8
Piecewise CUDA Graph is enabled by default as an experimental feature.
To work around this error, add --disable-piecewise-cuda-graph to your launch command.
Please report this issue at https://github.com/sgl-project/sglang/issues/new/choose
[2026-05-22 04:54:32 TP6] Scheduler hit an exception: Traceback (most recent call last):
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/managers/scheduler.py", line 3811, in run_scheduler_process
    scheduler.run_event_loop()
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/managers/scheduler.py", line 1507, in run_event_loop
    dispatch_event_loop(self)
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/managers/scheduler.py", line 3680, in dispatch_event_loop
    scheduler.event_loop_overlap()
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/managers/scheduler.py", line 1567, in event_loop_overlap
    batch_result = self.run_batch(batch)
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/managers/scheduler.py", line 2865, in run_batch
    batch_result = self.model_worker.forward_batch_generation(
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/managers/tp_worker.py", line 472, in forward_batch_generation
    out = self.model_runner.forward(
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/model_executor/model_runner.py", line 3182, in forward
    output = self._forward_raw(
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/model_executor/model_runner.py", line 3314, in _forward_raw
    ret, can_run_graph = self.forward_extend(
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/model_executor/model_runner.py", line 3075, in forward_extend
    ret = self.piecewise_cuda_graph_runner.replay(forward_batch, **kwargs)
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py", line 805, in replay
    output = self.model_runner.model.forward(
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/models/deepseek_v2.py", line 2524, in forward
    hidden_states = self.model(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1790, in _call_impl
    return forward_call(*args, **kwargs)
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/compilation/compile.py", line 195, in trampoline
    return compiled_callable(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py", line 1024, in compile_wrapper
    return fn(*args, **kwargs)
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/models/deepseek_v2.py", line 2216, in forward
    def forward(
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py", line 472, in __call__
    return super().__call__(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1790, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py", line 1263, in _fn
    return fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 949, in call_wrapped
    return self._wrapped_call(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 461, in __call__
    raise e
  File "/usr/local/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 447, in __call__
    return super(self.cls, obj).__call__(*args, **kwargs)  # type: ignore[misc]
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1790, in _call_impl
    return forward_call(*args, **kwargs)
  File "<eval_with_key>.1122", line 1227, in forward
    submod_1 = self.submod_1(getitem, s72, getitem_1, getitem_2, getitem_3);  getitem = getitem_1 = getitem_2 = submod_1 = None
  File "/usr/local/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 949, in call_wrapped
    return self._wrapped_call(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 461, in __call__
    raise e
  File "/usr/local/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 447, in __call__
    return super(self.cls, obj).__call__(*args, **kwargs)  # type: ignore[misc]
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1790, in _call_impl
    return forward_call(*args, **kwargs)
  File "<eval_with_key>.1124", line 5, in forward
    k_cache_and_topk_result = torch.ops.sglang.k_cache_and_topk_result(layer_id = 0, key = key_1, q_fp8 = y, weights = weights, topk_result = topk_result);  key_1 = y = weights = topk_result = k_cache_and_topk_result = None
  File "/usr/local/lib/python3.10/dist-packages/torch/_ops.py", line 1269, in __call__
    return self._op(*args, **kwargs)
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/layers/attention/dsa/dsa_indexer.py", line 127, in k_cache_and_topk_result
    indexer._store_index_k_cache(
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/layers/attention/dsa/dsa_indexer.py", line 1200, in _store_index_k_cache
    fused_store_index_k_cache(
  File "/actions-runner/_work/sglang/sglang/python/sglang/jit_kernel/fused_store_index_cache.py", line 105, in fused_store_index_k_cache
    module.fused_store_index_k_cache(key, index_k_with_scale, out_cache_loc)
  File "python/tvm_ffi/cython/function.pxi", line 929, in tvm_ffi.core.Function.__call__
RuntimeError: Tensor match failed for Tensor<8>[strides=<1>, dtype=int64, device=cuda:6] at /actions-runner/_work/sglang/sglang/python/sglang/jit_kernel/csrc/dsa/fused_store_index_cache.cuh:107
- Root cause: Size mismatch for shape#0('num_tokens'): expected 6 but got 8

https://github.com/sgl-project/sglang/actions/runs/26265766989/job/77308576707?pr=23351
@nvjullin Looks like real failures caused by this change.

nvjullin · 2026-05-22T07:15:35Z

Fixed out_cache_loc, same bad merge from above. Quick accuracy test shows it's correct.

nvpohanh · 2026-05-22T07:22:11Z

/tag-and-rerun-ci

Fridge003 · 2026-05-22T20:09:59Z

/rerun-test test_pcg_glm5_fp4.py

github-actions · 2026-05-22T20:10:19Z

🚀 4-gpu-b200 (1 test): ❌ View workflow run

cd test/ && python3 registered/piecewise_cuda_graph/test_pcg_glm5_fp4.py

Fridge003 · 2026-05-22T21:05:42Z

Other tests are OK https://github.com/sgl-project/sglang/pull/23351/checks?sha=632b381033e3ea1faf5d0561a12ff6b6f1b4fa0a
Now fixing the newly added pcg test

Fridge003 · 2026-05-22T21:05:46Z

/rerun-test test_pcg_glm5_fp4.py

github-actions · 2026-05-22T21:06:06Z

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/piecewise_cuda_graph/test_pcg_glm5_fp4.py

Port of sgl-project#23351 ("Support piecewise CUDA graph with NSA"), translated across this fork's dsa->nsa rename (sgl-project#25014). - nsa_indexer: add k_cache_and_topk_result / logits_head_gate_pcg custom ops; PCG guards in forward_cuda (metadata fetch, delayed K-cache store, gate, seq_lens_cpu access); optional topk_result in _get_topk_ragged; out_cache_loc param in _store_index_k_cache; CP/loop-path asserts; is_compiling rope guard. - nsa_backend: force use_mha=False under PCG; drop _forward_trtllm squeeze(1). - piecewise_context_manager / model_runner / piecewise_cuda_graph_runner: thread nsa_indexers through the forward context. - model_config: un-gate NSA models from the piecewise disable list. - hadamard / layernorm: wrap in register_custom_op with fake impls. - radix_attention: pass cos_sin_cache/is_neox/llama_4_scaling/topk_indices through the split-op schema. - deepseek_v2: skip moe_output_buffer_ctx while compiling. - server_args: disable PCG under NSA prefill-CP; remove flashinfer-MoE PCG-disable block. - test: add test_pcg_glm5_fp4.py (GLM-5-NVFP4, TP=4, B200). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

github-actions Bot added piecewise-cuda-graph jit-kernel labels Apr 21, 2026

nvjullin changed the title ~~Pcg nsa pr~~ Support piecewise CUDA graph with NSA Apr 21, 2026

gemini-code-assist Bot reviewed Apr 21, 2026

View reviewed changes

amd-bot mentioned this pull request Apr 21, 2026

[CI Monitor] Daily Report - 2026-04-21 bingxche/sglang-ci-bot#42

Open

Fridge003 reviewed Apr 21, 2026

View reviewed changes

Fridge003 mentioned this pull request Apr 22, 2026

[Roadmap] DeepSeek v3.2 (GLM 5) Optimization #15025

Closed

40 tasks

nvjullin force-pushed the pcg-nsa-pr branch from a2b963a to f2bf494 Compare April 22, 2026 07:47

amd-bot mentioned this pull request Apr 23, 2026

[CI Monitor] Daily Report - 2026-04-23 bingxche/sglang-ci-bot#44

Open

nvjullin marked this pull request as ready for review April 24, 2026 10:14

nvjullin requested review from 1am9trash, BBuf, DarkSharpness, Edwardf0t1, HaiShaw, Oasis-Git, Qiaolin-Yu, Ying1123, celve, ch-wan, hebiao064, hlu1, hnyls2002, hubertlu-tw, ispobock, kkHuang-amd and merrymercy as code owners April 24, 2026 10:14

Fridge003 reviewed May 21, 2026

View reviewed changes

Comment thread test/registered/piecewise_cuda_graph/test_pcg_glm5_fp4.py Outdated

Update test/registered/piecewise_cuda_graph/test_pcg_glm5_fp4.py

5b1c2c7

Fridge003 added the run-ci-extra label May 21, 2026

fixed get_attn_backend()

ac82409

fixed out_cache_loc

632b381

Update max_tokens value in test_pcg_glm5_fp4.pyfix test

e0fb050

Increase max_tokens from 1024 to 4096

82e6f83

Fridge003 merged commit cadfa2d into sgl-project:main May 22, 2026
57 of 126 checks passed

b8zhong mentioned this pull request May 27, 2026

Remove DeepGEMM for indexer GEMM in piecewise NSA path #26494

Merged

nvjullin deleted the pcg-nsa-pr branch May 29, 2026 08:50

mmangkad mentioned this pull request May 29, 2026

Fix PCG guard for flashinfer TRTLLM MoE #26718

Closed

mattteochen mentioned this pull request Jun 2, 2026

[BCG][GLM5] perf: BCG support and prefill enhancements #27053

Merged

5 tasks

mqhc2020 pushed a commit to mqhc2020/sglang that referenced this pull request Jun 2, 2026

Support piecewise CUDA graph with NSA (sgl-project#23351)

d36097b

janbernloehr mentioned this pull request Jun 14, 2026

[Bug] flashinfer_trtllm BF16 MoE: piecewise CUDA graph capture causes illegal memory access at trtllm_fused_moe_dev_kernel.cu:991 (regression in lmsysorg/sglang:dev between 0.0.0.dev1+g0c8049d9b and 0.0.0.dev1+g208397aff) #26715

Closed

5 tasks

elvischenv mentioned this pull request Jun 15, 2026

Support Confidential Computing (CC) for Qwen3.5-397B-A17B-FP8 on B300 #28251

Closed

b8zhong mentioned this pull request Jun 18, 2026

support deepseekv3.2-piecewise-cuda-graph #18094

Closed

5 tasks

nvjullin mentioned this pull request Jun 30, 2026

Disable dsr1 prefill cudagraphs by default #28053

Merged

5 tasks

Uh oh!

Conversation

nvjullin commented Apr 21, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

Review and Merge Process

CI States

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist Bot Apr 21, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Fridge003 Apr 21, 2026

Choose a reason for hiding this comment

Uh oh!

nvjullin Apr 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mmangkad May 29, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Fridge003 commented Apr 21, 2026

Uh oh!

Fridge003 commented Apr 21, 2026

Uh oh!

nvjullin commented Apr 24, 2026

Uh oh!

Uh oh!

nvpohanh commented May 22, 2026

Uh oh!

nvjullin commented May 22, 2026

Uh oh!

nvpohanh commented May 22, 2026

Uh oh!

nvjullin commented May 22, 2026

Uh oh!

nvpohanh commented May 22, 2026

Uh oh!

Fridge003 commented May 22, 2026

Uh oh!

github-actions Bot commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Fridge003 commented May 22, 2026

Uh oh!

Fridge003 commented May 22, 2026

Uh oh!

github-actions Bot commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

nvjullin commented Apr 21, 2026 •

edited by github-actions Bot

Loading

nvjullin Apr 22, 2026 •

edited

Loading

github-actions Bot commented May 22, 2026 •

edited

Loading

github-actions Bot commented May 22, 2026 •

edited

Loading