[DeepSeek v3.2] Opt MTP decode cuda batch sizes and nsa implementation

[xu-yfei]

Motivation

1、draft_extend and target_verify use nsa_decode_backend as backend implementation instead of nsa_prefill_backend
In the MTP scenario, the NSA attention implementation selected for prefill, draft_extend, and target_verify is all nsa_prefill_backend, which is typically set to flashmla_sparse. For draft_extend and target_verify scenarios with a small number of tokens, we observed that fa3 delivers better performance. However, when flashmla_sparse is chosen for prefill, there is no interface to additionally specify the NSA attention implementation for draft_extend and target_verify. This PR modifies the backend implementation of draft_extend and target_verify to nsa_decode_backend —in practice, their CUDA Graph initialization and replay already use nsa_decode_impl.

num_draft_tokens=4, TPOT:

batch size	flashmla_sparse	fa3
1	11.87	9.44
2	13.34	13.08
8	15.25	15.87
16	21.25	21.50
32	32.94	33.79

2、Opt cuda graph batch sizes when MTP:
In the Prefill CP scenario, require_gathered_buffer is set to True. Assuming the number of MTP num_draft_tokens is 4, the CUDA Graph tokens in the MTP scenario are calculated as follows:
get_batch_sizes_to_capture returns batch sizes that are multiples of attn_tp_size, i.e., [8, 16, 24, 32, ....].
In capture_one_batch_size, num_tokens = bs * self.num_tokens_per_bs, meaning the actual number of tokens is [32, 64, 96, 128, ...].
When bs=1, the number of tokens is 4 but needs to be padded to 32, resulting in significant performance overhead. In reality, num_tokens only needs to be a multiple of attn_tp_size.
Modification: get_batch_sizes_to_capture is updated to return batch sizes where batch size * num_draft_tokens is divisible by attn_tp_size, i.e., [2, 4, 6, 8, ....].

Modifications

Accuracy Tests

MODEL_PATH=/home/models/DeepSeek-V3.2/

python3 -m sglang.launch_server --model-path $MODEL_PATH --trust-remote-code \
--port 8000 --host 0.0.0.0 --attention-backend  nsa \
--enable-metrics --mem-fraction-static 0.8 --max-running-requests 128 --enable-cache-report --page-size 64 \
--tp-size 8 --moe-dense-tp-size 1 \
--tool-call-parser deepseekv32 \
--reasoning-parser deepseek-v3 \
--moe-dense-tp-size 1 \
--chunked-prefill-size 16384 \
--nsa-decode-backend fa3 \
--enable-nsa-prefill-context-parallel \
--nsa-prefill-cp-mode round-robin-split \
--speculative-algorithm EAGLE \
--speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4

# gsm8k
Accuracy: 0.948
Invalid: 0.000
Latency: 192.326 s
Output throughput: 630.665 token/s

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

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

[gemini-code-assist]

Summary of Changes

Hello @xu-yfei, 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 introduces significant performance and efficiency improvements for the DeepSeek v3.2 model, particularly in scenarios involving multi-token prediction. It refines the selection of attention backend implementations for specific forward modes like draft_extend and target_verify, allowing for the use of more performant algorithms like fa3 where appropriate. Additionally, it optimizes the CUDA graph batch size determination process to prevent excessive token padding, leading to more efficient resource utilization during model execution.

Highlights

• Dynamic NSA Attention Backend Selection: Refactored the NSA attention implementation selection for draft_extend and target_verify operations to dynamically use nsa_decode_backend. This change enables the use of more performant algorithms like fa3 for small token counts in these specific scenarios, improving efficiency.
• Optimized CUDA Graph Batch Sizes for MTP: Optimized the calculation of CUDA graph batch sizes in Multi-Token Prediction (MTP) scenarios. The updated logic ensures that the total number of tokens (batch size multiplied by tokens per batch) is divisible by attn_tp_size, thereby reducing unnecessary padding and improving CUDA graph capture efficiency.

🧠 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.

Footnotes

1. 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. ↩

[xu-yfei]

@Fridge003 Could you please have a review?

[gemini-code-assist]

Code Review

This pull request introduces two performance optimizations for Multi-Token Prediction (MTP) scenarios. First, it switches the NSA attention backend to nsa_decode_backend for draft_extend and target_verify modes, which is more efficient for smaller token counts. Second, it optimizes CUDA graph batch sizes to minimize padding overhead. The changes are well-aligned with the motivations described, and the refactoring in CudaGraphRunner is clean. I have one suggestion to ensure the first optimization is consistently applied to all relevant draft extend modes.

[Fridge003]

/tag-and-rerun-ci