[bugfix] avoid attention padding tokens computation in pcg

[Chen-0210]

Motivation

fix #17330.

When PCG is enabled, the attention metadata is initialized with real_num_tokens, but the input tensor still contains padded tokens. Attention backends such as FlashInfer can not handle this well, which can lead to undefined behavior, including nan value, corrupted outputs (repeated !!!!!), result in abnormally long output lengths.

To fix this, exclude the padded tokens and make PCG more robust.

Modifications

1. sliced the padded tensors to their actual lengths (real_num_tokens) before passing them to attention.
2. the attention is now splited from the graph, mamba cache related tensors are no longer required in pcg. Removed the related code.

Accuracy Tests

python3 -m sglang.launch_server --model Qwen/Qwen3-Next-80B-A3B-Instruct --tp 2 --enable-piecewise-cuda-graph --piecewise-cuda-graph-compiler eager --port 60002 --skip-server-warmup --log-requests --log-requests-level 3 --attention-backend flashinfer --mamba-scheduler-strategy extra_buffer

python3 benchmark/gsm8k/bench_sglang.py --parallel 1319 --num-questions 1319 --host http://127.0.0.1 --port 60002

100%|███████████████████████████████████████████████████████████████████████████████████████████| 1319/1319 [01:43<00:00, 12.75it/s]
Accuracy: 0.938
Invalid: 0.000
Latency: 103.774 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.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.

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 @Chen-0210, 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 implements a bug fix to prevent padding tokens from being included in attention calculations within the Piecewise CUDA Graph (PCG) execution. By introducing a real_num_tokens field to ForwardBatch and modifying attention layers to process only the actual tokens, the change ensures more accurate and efficient attention computations, eliminating the influence of padding on model outputs.

Highlights

• Padding Token Exclusion: Attention mechanisms in radix_attention.py and qwen3_next.py now explicitly ignore padding tokens by slicing input tensors (query, key, value, hidden_states, rope, sinks) based on a new real_num_tokens field, ensuring computations only involve actual data.
• Introduction of real_num_tokens: A new optional integer field, real_num_tokens, has been added to the ForwardBatch class to accurately track the number of non-padding tokens in a batch.
• Piecewise CUDA Graph Runner Refactoring: The piecewise_cuda_graph_runner.py has been updated to correctly pass and utilize real_num_tokens during graph warmup and capture. It also streamlines the handling of out_cache_loc and out_cache_loc_swa by making them dynamically created or passed via ForwardBatch rather than persistent graph runner attributes.

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

[gemini-code-assist]

Code Review

This pull request introduces a bugfix for handling padding tokens within the Piecewise CUDA Graph (PCG) execution path. The core change is the addition of a real_num_tokens field to the ForwardBatch dataclass, which allows distinguishing actual tokens from padding. This field is then utilized in custom attention operations (unified_attention_with_output and gdn_with_output) to correctly slice input tensors, ensuring that only real tokens are processed during attention computation. Consequently, the logic in PiecewiseCudaGraphRunner for handling out_cache_loc has been simplified by removing pre-allocated tensors and passing them directly from the forward_batch.

The changes appear correct and effectively address the padding issue in PCG mode. I have one minor suggestion to improve code clarity by correcting a duplicated comment.

[Oasis-Git]

Since these changes affect a broad range of functionality, the full unit test suite should be carefully validated before approving this PR.

[Oasis-Git]

/tag-run-ci-label

[hebiao064]

why we need to remove these code?

[Chen-0210]

The relevant part is not included in PCG, so it does not require a fixed memory address.

[Chen-0210]

previously disabled CI test_qwen3_next_models_pcg is now enabled and pass.

[zminglei]

previously disabled CI `test_qwen3_next_models_pcg` is now enabled and pass.

Seems previously this CI only has a fixed chance to fail, not fail every time. Could we run the test several times, let's say 10 times to verify it passes every time?

[Chen-0210]

@zminglei
I’ve tested this multiple times on my side and haven’t seen the issue reproduce.

The occasional “all requests invalid” fail was already fixed by another PR #17404 related to mamba cache. this pr mainly fix the rare precision issue related to padding token.

Running the CI multiple times would be great, but I’m not sure if the current CI can accept 10 runs due to time and resource constraints..

[zminglei]

@zminglei I’ve tested this multiple times on my side and haven’t seen the issue reproduce.

The occasional “all requests invalid” fail was already fixed by another PR #17404 related to mamba cache. this pr mainly fix the rare precision issue related to padding token.

Running the CI multiple times would be great, but I’m not sure if the current CI can accept 10 runs due to time and resource constraints..

Thanks, that makes sense, as we've already locally run it multiple times and verified it's fixed, it should be good enough. And for CI it's ok to only run once as others.

[zminglei]

/rerun-failed-ci again

[Chen-0210]

@zminglei It seems the re-run didn’t take effect

[Chen-0210]

/rerun-failed-ci again

[Chen-0210]

/rerun-failed-ci again