support deepseekv3.2-piecewise-cuda-graph by BJWang-ant · Pull Request #18094 · sgl-project/sglang

BJWang-ant · 2026-02-02T03:50:46Z

Motivation

Modifications

support deepseekv3.2-piecewise-cuda-graph

Accuracy Tests

Benchmarking 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 Process

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

gemini-code-assist · 2026-02-02T03:51:12Z

Summary of Changes

Hello @BJWang-ant, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances support for DeepSeekV3.2 models by integrating them with the piecewise CUDA graph compilation framework. The core objective is to optimize the model's attention and Mixture-of-Experts (MoE) layers for efficient execution on GPUs. This involves adapting data handling, introducing new custom operations for graph compatibility, and refining conditional execution logic to ensure high performance and stability within the CUDA graph environment.

Highlights

Piecewise CUDA Graph Integration: Extended the ForwardContext to include indexer_layers, enabling better management and compilation of DeepSeekV3.2's specific indexing mechanisms within the piecewise CUDA graph framework.
Performance Optimization for MoE Routing: Applied torch.compile(dynamic=True) to the on_select_experts method in expert distribution, which is expected to enhance the performance of Mixture-of-Experts (MoE) routing by leveraging PyTorch's JIT compilation.
Memory and Computation Efficiency in NSA Indexer: Optimized the weights_proj parameter data type in the nsa_indexer from float32 to bfloat16 and ensured proper float casting for projections, leading to reduced memory footprint and potentially faster computations.
Refactored Attention Input Handling: Adjusted how attention inputs are handled in communicator and deepseek_v2 by moving the attn_inputs setup. This change aligns the process with the requirements of piecewise CUDA graph compilation, preventing conflicts and ensuring smooth operation.
Conditional Execution with Custom Operations: Implemented conditional execution paths for self.indexer and mha_one_shot logic within DeepSeekV2 attention. This uses torch.cond and new custom operations (nas_indexer_forward_cus) to ensure compatibility and efficiency when running in CUDA graph extend mode.
New Custom Operations for DeepSeekV3.2 Components: Introduced new custom operations, mla_gdn_with_output and mlp_gdn_with_output, along with their fake implementations. These are crucial for supporting the piecewise CUDA graph compilation for DeepSeekV3.2's Multi-Layer Attention (MLA) and Multi-Layer Perceptron (MLP) components.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request adds support for deepseekv3.2-piecewise-cuda-graph. The changes primarily involve making the model compatible with torch.compile and piecewise CUDA graph execution. This is achieved by replacing dynamic control flow with torch.cond, wrapping model components in custom operators, and adjusting data types and operations for compilation.

Overall, the changes are in the right direction. I've found a critical typo that needs to be fixed and some opportunities for code cleanup and refactoring to improve maintainability. Please see my detailed comments.

gemini-code-assist · 2026-02-02T03:53:21Z

        for layer in self.model.model.layers:
            if hasattr(layer, "self_attn"):
                if hasattr(layer.self_attn, "attn"):
                    self.attention_layers.append(layer.self_attn.attn)
                elif hasattr(layer.self_attn, "attn_mqa"):
                    # For DeepSeek model
                    self.attention_layers.append(layer.self_attn.attn_mqa)
+                elif hasattr(layer.self_attn, "indexer "):


There appears to be a typo in the attribute name check. It's "indexer " with a trailing space, which will likely cause hasattr to fail and prevent indexer_layers from being collected correctly. It should be "indexer".

Suggested change

elif hasattr(layer.self_attn, "indexer "):

elif hasattr(layer.self_attn, "indexer"):

gemini-code-assist · 2026-02-02T03:53:21Z

+        # due to piecewise-cudagraph，so move
+        # if self.qkv_latent_func is not None:
+        #     attn_inputs = AttentionInputs(
+        #         hidden_states, forward_batch, self.qkv_latent_func
+        #     )
+        #     get_attn_tp_context().set_attn_inputs(attn_inputs)


This block of code has been commented out with a note that it was moved. To improve code clarity and remove dead code, it's better to delete this commented-out block entirely.

gemini-code-assist · 2026-02-02T03:53:21Z

@@ -743,7 +743,7 @@ def biased_grouped_topk_gpu(
    experts_per_group = (
        num_experts // num_expert_group if num_expert_group else num_experts
    )
-
+    fused_topk_deepseek = None


The line fused_topk_deepseek = None unconditionally disables the fused_topk_deepseek optimization path below, making the following if block dead code. If this is a permanent change, please remove this line and the unreachable if block. If this is temporary, please add a TODO comment explaining why it's disabled and when it might be re-enabled.

gemini-code-assist · 2026-02-02T03:53:21Z

+        # if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
+        #     # TODO(yuwei): fix the compilation errors for MOE A2A backend
+        #     log_info_on_rank0(
+        #         logger,
+        #         "Disable piecewise CUDA graph due to existing compilation errors",
+        #     )
+        #     return False


This block of code has been commented out. To improve code clarity and remove dead code, it's better to delete this commented-out block entirely.

gemini-code-assist · 2026-02-02T03:53:21Z

+                if (
+                    forward_batch.forward_mode.is_extend()
+                    and get_forward_context() is not None
+                ):
+                    topk_indices = torch.empty(
+                        (hidden_states.shape[0], 2048),
+                        dtype=torch.int32,
+                        device=hidden_states.device,
+                    )
+                    nas_indexer_forward_cus(
+                        layer_id=self.layer_id,
+                        hidden_states=hidden_states,
+                        q_lora=q_lora,
+                        positions=positions,
+                        return_indices=True,
+                        output=topk_indices,
+                    )
+                else:
                    topk_indices = self.indexer(
                        x=hidden_states,
                        q_lora=q_lora,
                        positions=positions,
                        forward_batch=forward_batch,
                        layer_id=self.layer_id,
                    )


The logic for computing topk_indices is duplicated here and in the else block below (lines 1618-1642). Consider refactoring this logic into a helper method to avoid code duplication and improve maintainability.

yiakwy-xpu-ml-framework-team · 2026-02-27T04:55:13Z

Hi @BJWang-ant are you still working on this issue ?

yiakwy-xpu-ml-framework-team · 2026-02-27T05:18:50Z

In most cases from my team, prefill is much faster than decode, hence I guess that there is no need to use piece wise cuda graph for prefill. How do you think about it ?

BJWang-ant · 2026-02-27T06:52:48Z

Hi @BJWang-ant are you still working on this issue ?

yes, I am working it.Most of the work on PCG has been completed, and currently, some tests on accuracy and performance are being carried out normally.

BJWang-ant · 2026-02-27T06:55:04Z

In most cases from my team, prefill is much faster than decode, hence I guess that there is no need to use piece wise cuda graph for prefill. How do you think about it ?

From our some test cases, we can see that PCG does bring some benefits, but as you said, the benefits are not very significant.

Oasis-Git · 2026-03-06T01:17:08Z

@BJWang-ant Hi are u still working on it? we may collaborate on this since the support for nsa is needed now

BJWang-ant · 2026-03-06T01:57:50Z

@BJWang-ant Hi are u still working on it? we may collaborate on this since the support for nsa is needed now

yes, I am working on it. I can first push some code up. Under this commit, it can currently be run through. However, in the PD separation mode, there are still issues with the accuracy. Currently, I'm still trying to identify the cause of the accuracy problem and haven't come up with any good solutions yet.

Oasis-Git · 2026-03-06T02:01:22Z

@BJWang-ant Let's support it on non-pd scenario firstly.

Also if you are in sgl slack let's dm on it.

BJWang-ant · 2026-03-06T02:02:46Z

@BJWang-ant Let's support it on non-pd scenario firstly.

In the non-padding scenario, the accuracy has been achieved.

BJWang-ant · 2026-03-06T02:03:58Z

@BJWang-ant Let's support it on non-pd scenario firstly.

Also if you are in sgl slack let's dm on it.

OK.I contact you

BJWang-ant · 2026-03-07T03:37:37Z

@BJWang-ant Let's support it on non-pd scenario firstly.

Also if you are in sgl slack let's dm on it.

I have rebased the code onto the latest commit of the repository.Perhaps you could work together to look into the issue of accuracy.

BJWang-ant · 2026-03-09T09:07:52Z

Currently, I am conducting an accuracy test using PD separation. I have tested the command, but there are still accuracy issues on the P side that cannot be pinpointed. Moreover, if a function to print tensor data is inserted at any position in the following code, and this parameter is decorated with @register_split_op() and @register_custom_op(mutates_args=[]) then the accuracy becomes normal.

This is my first time developing the functions related to piecewise CUDA graph. I'm not very familiar with it. Do any of you have any suggestions?

update

BJWang-ant · 2026-03-16T11:29:18Z

I attempted to perform PD separation using a non-DP-attention method, but still encountered mismatches. Moreover, in this case, the non-PCG also showed mismatches.

prefill：
NVSHMEM_IB_GID_INDEX=3
TORCH_CUDA_ARCH_LIST="9.0"
NVSHMEM_IB_TRAFFIC_CLASS=184
NVSHMEM_IB_ENABLE_IBGDA=true
MC_SLICE_SIZE=262144
GLOO_SOCKET_IFNAME=eth1
NCCL_SOCKET_IFNAME=eth1
MC_TE_METRIC=true
SGLANG_JIT_DEEPGEMM_COMPILE_WORKERS=32
python3 -m sglang.launch_server
--model-path /deepseek-ai/DeepSeek-V3.2/
--disaggregation-mode prefill
--disaggregation-ib-device mlx5_4,mlx5_5,mlx5_6,mlx5_7,mlx5_8,mlx5_9,mlx5_10,mlx5_11
--dist-init-addr $p0:5757
--nnodes 1
--node-rank $pi
--tp 8
--enable-metrics
--enable-metrics-for-all-schedulers
--enable-expert-distribution-metrics
--decode-log-interval 1
--moe-a2a-backend deepep
--host 0.0.0.0
--port 8000
--trust-remote-code
--moe-dense-tp-size 1
--enable-cache-report
--watchdog-timeout 1000000
--deepep-mode normal
--mem-fraction-static 0.7
--max-running-requests 512
--max-prefill-tokens 8192
--chunked-prefill-size -1
--ep-dispatch-algorithm dynamic
--expert-distribution-recorder-mode stat
--deepep-config /upfs/shishan/configs/deepep.json
--page-size 64
--log-level debug
--context-length 23000

decode：
GLOO_SOCKET_IFNAME=eth1 \
GLOO_SOCKET_TIMEOUT_MS=60000 \
TORCH_CUDA_ARCH_LIST="9.0" \
NCCL_SOCKET_IFNAME=eth1 \
MC_TE_METRIC=true \
NVSHMEM_HCA_LIST=mlx5_4,mlx5_5,mlx5_6,mlx5_7 \
NVSHMEM_IB_GID_INDEX=3 \
NVSHMEM_IB_ENABLE_IBGDA=true \
NVSHMEM_IB_TRAFFIC_CLASS=184 \
NVSHMEM_BOOTSTRAP_UID_SOCK_FAMILY=AF_INET \
NVSHMEM_BOOTSTRAP_UID_SOCK_IFNAME=eth1 \
DEEPEP_DIAGNOSE_INTERVAL=120 \
DEEPEP_DIAGNOSE_LOG_DETAILS=1 \
SGLANG_JIT_DEEPGEMM_COMPILE_WORKERS=32 \
python3 -m sglang.launch_server \
--model-path /deepseek-ai/DeepSeek-V3.2/ \
--disaggregation-mode decode \
--disaggregation-ib-device mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_6,mlx5_7 \
--disaggregation-transfer-backend mooncake \
--dist-init-addr $d0:5757 \
--nnodes 1 \
--node-rank $di \
--tp 8 \
--decode-log-interval 1 \
--enable-metrics \
--enable-metrics-for-all-schedulers \
--host 0.0.0.0 \
--port 8000 \
--trust-remote-code \
--moe-dense-tp-size 1 \
--enable-cache-report \
--disaggregation-mode decode \
--watchdog-timeout 1000000 \
--deepep-mode low_latency \
--mem-fraction-static 0.8 \
--max-running-requests 256 \
--context-length 32000 \
--moe-a2a-backend deepep \
--prefill-round-robin-balance \
--page-size 64 \
--log-level debug \
--speculative-algorithm EAGLE \
--speculative-num-steps 3 \
--speculative-eagle-topk 1 \
--speculative-num-draft-tokens 4 \
--enable-flashinfer-allreduce-fusion \
--speculative-attention-mode decode \
--cuda-graph-max-bs 32 \
--load-balance-method round_robin

weireweire · 2026-03-20T07:57:16Z

@BJWang-ant I see you are using --enable-flashinfer-allreduce-fusion and --moe-dense-tp-size 1. It should be this issue: #19918

Could you try remove --moe-dense-tp-size 1? If the accuracy issue gone, I think you can ignore this accuracy issue and we'll fix it.

cc @nvpohanh @Fridge003

b8zhong · 2026-06-18T18:02:30Z

@BJWang-ant Hello, I thikn this feature is working in main, it's covered by #23351

BJWang-ant requested review from AniZpZ, BBuf, Edwardf0t1, FlamingoPg, Fridge003, HaiShaw, Qiaolin-Yu, Ying1123, ch-wan, fzyzcjy, hebiao064, hnyls2002, ispobock, merrymercy and zhyncs as code owners February 2, 2026 03:50

github-actions Bot added deepseek piecewise-cuda-graph labels Feb 2, 2026

BJWang-ant marked this pull request as draft February 2, 2026 03:51

gemini-code-assist Bot reviewed Feb 2, 2026

View reviewed changes

This was referenced Feb 3, 2026

[Feature - Doing] support deepseekv3.2 Prefill piecewise cuda graph #18148

Closed

[TODO List] Piecewise Cuda Graph Default Enable #18130

Closed

BJWang-ant force-pushed the opt_weight_proj_v0 branch from 2ae08bf to 46101c7 Compare March 6, 2026 02:00

BJWang-ant force-pushed the opt_weight_proj_v0 branch from 46101c7 to cfdf170 Compare March 7, 2026 03:36

BJWang-ant marked this pull request as ready for review March 9, 2026 09:08

BJWang-ant requested a review from xiezhq-hermann as a code owner March 9, 2026 09:08

BJWang-ant mentioned this pull request Mar 9, 2026

[Feature] DeepSeekV3.2 PCG support #20176

Closed

2 tasks

BJWang-ant added 3 commits March 16, 2026 04:26

PCG run through

f74a37e

update

update code

60a9289

pcg capture for no-dp-attention

60c6574

BJWang-ant force-pushed the opt_weight_proj_v0 branch from c7c7048 to 60c6574 Compare March 16, 2026 11:27

BJWang-ant requested a review from yizhang2077 as a code owner March 16, 2026 11:27

Fridge003 mentioned this pull request Apr 3, 2026

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

Closed

40 tasks

Fridge003 self-assigned this Apr 3, 2026

nvpohanh mentioned this pull request Apr 8, 2026

[Tracking] GLM-5 (G)B200 Functional Support and Optimizations #19380

Closed

b8zhong closed this Jun 18, 2026

	elif hasattr(layer.self_attn, "indexer "):
	elif hasattr(layer.self_attn, "indexer"):

Uh oh!

Conversation

BJWang-ant commented Feb 2, 2026

Motivation

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Review Process

Uh oh!

gemini-code-assist Bot commented Feb 2, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

yiakwy-xpu-ml-framework-team commented Feb 27, 2026

Uh oh!

yiakwy-xpu-ml-framework-team commented Feb 27, 2026

Uh oh!

BJWang-ant commented Feb 27, 2026

Uh oh!

BJWang-ant commented Feb 27, 2026

Uh oh!

Oasis-Git commented Mar 6, 2026

Uh oh!

BJWang-ant commented Mar 6, 2026

Uh oh!

Oasis-Git commented Mar 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

BJWang-ant commented Mar 6, 2026

Uh oh!

BJWang-ant commented Mar 6, 2026

Uh oh!

BJWang-ant commented Mar 7, 2026

Uh oh!

BJWang-ant commented Mar 9, 2026

Uh oh!

BJWang-ant commented Mar 16, 2026

Uh oh!

weireweire commented Mar 20, 2026

Uh oh!

b8zhong commented Jun 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Oasis-Git commented Mar 6, 2026 •

edited

Loading