Fix cutlass MLA gets almost zero accuracy

[fzyzcjy]

Fix issues in cuda graph and now it is roughly same as non-cuda-graph or slightly lower. But seems non-cuda-graph still has bug, thus may need to check separately.

test

while true; do python3 benchmark/gsm8k/bench_sglang.py --parallel 10000 --num-questions 1400; done

master + no cudagraph

python3 -m sglang.launch_server --model /dev/shm/DeepSeek-V3-0324 --tp 8 --enable-dp-attention --dp 8 --trust-remote-code --attention-backend cutlass_mla --context-length 4096 --cuda-graph-bs 64 128 --max-running-requests 128 --disable-cuda-graph

Accuracy: 0.917
Accuracy: 0.917
Accuracy: 0.910
Accuracy: 0.920
Accuracy: 0.916

master + cuda graph

python3 -m sglang.launch_server --model /dev/shm/DeepSeek-V3-0324 --tp 8 --enable-dp-attention --dp 8 --trust-remote-code --attention-backend cutlass_mla --context-length 4096 --cuda-graph-bs 64 128 --max-running-requests 128

Accuracy: 0.060

pr + cuda graph

python3 -m sglang.launch_server --model /dev/shm/DeepSeek-V3-0324 --tp 8 --enable-dp-attention --dp 8 --trust-remote-code --attention-backend cutlass_mla --context-length 4096 --cuda-graph-bs 64 128 --max-running-requests 128

Accuracy: 0.906
Accuracy: 0.914
Accuracy: 0.917
Accuracy: 0.913
Accuracy: 0.917

FYI

baseline + triton

python3 -m sglang.launch_server --model /dev/shm/DeepSeek-V3-0324 --tp 8 --enable-dp-attention --dp 8 --trust-remote-code --context-length 4096 --cuda-graph-bs 64 128 --max-running-requests 128

Accuracy: 0.933
Accuracy: 0.933
Accuracy: 0.937
Accuracy: 0.927

Motivation

Modifications

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[gemini-code-assist]

Hello @fzyzcjy, 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!

Summary of Changes

Hello everyone,

Gemini (gemini-code-assist) here to provide a summary of this pull request. This PR aims to address a critical issue where the Cutlass MLA backend was exhibiting almost zero accuracy. Based on the changes, it appears the problem was related to how CUDA graphs were being handled, specifically concerning the management of sequence padding length, KV index buffers, and workspace allocation during graph capture and replay.

The core modifications involve ensuring that buffer sizes (max_seqlen_pad) are determined based on the pre-allocated CUDA graph buffers rather than dynamic runtime values, explicitly initializing the KV index buffer before use in replay, and removing dynamic workspace allocation within the capture/replay functions, suggesting workspace is now handled statically or differently for CUDA graph paths. These changes should stabilize the memory access patterns and state within the captured graph, resolving the accuracy degradation.

Highlights

• Accuracy Fix: The primary goal of this PR is to fix a severe accuracy issue observed when using the Cutlass MLA backend, likely related to CUDA graph execution.
• CUDA Graph Buffer Handling: Modifies the logic for determining max_seqlen_pad during CUDA graph capture to use the size of the pre-allocated KV index buffer, ensuring compatibility with the static nature of CUDA graphs.
• KV Index Buffer Initialization: Adds explicit initialization of the cuda_graph_kv_indices buffer to -1 before populating it during CUDA graph replay, preventing potential use of stale or incorrect indices.
• Workspace Allocation Refinement: Removes dynamic workspace size calculation and allocation within the CUDA graph capture and replay paths, suggesting workspace is managed differently (e.g., statically allocated) for these specific execution modes.

Changelog

• python/sglang/srt/layers/attention/cutlass_mla_backend.py
  • In init_forward_metadata_capture_cuda_graph, changed max_seqlen_pad calculation (L160) to use the shape of self.cuda_graph_kv_indices.
  • Removed dynamic workspace size calculation and allocation (L172-177) in init_forward_metadata_capture_cuda_graph.
  • Added initialization of self.cuda_graph_kv_indices to -1 (L203) in init_forward_metadata_replay_cuda_graph.
  • Removed dynamic workspace size calculation, allocation, and metadata updates (L220-228) in init_forward_metadata_replay_cuda_graph.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 is currently in preview and 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 to provide feedback.

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 addresses an accuracy issue in the cutlass MLA backend, specifically within the CUDA graph capture and replay logic. The changes focus on how sequence padding and KV index metadata are handled for fixed-shape CUDA graphs. The core idea seems to be reusing a captured metadata object and ensuring the underlying data it references is correctly updated for each batch during replay.

Summary of Findings

• Missing max_seqlen_pad calculation in replay function: The calculation of max_seqlen_pad based on the current batch's sequence lengths was removed from init_forward_metadata_replay_cuda_graph, but this value is necessary for the captured metadata object to correctly handle variable sequence lengths during CUDA graph replay. This will likely cause incorrect attention results.
• Corrected workspace handling for CUDA graphs: The removal of workspace allocation within the CUDA graph capture and replay functions is a positive change, aligning with best practices for CUDA graph usage.
• Resetting KV indices buffer: Adding a step to reset the cuda_graph_kv_indices buffer before populating it for each batch during replay is a good practice to prevent stale data issues.

Merge Readiness

This pull request addresses an important accuracy issue and makes positive changes to CUDA graph handling by removing unnecessary workspace allocations within capture/replay. However, there is a critical issue where a required variable (max_seqlen_pad) is no longer calculated in the replay function, which will likely lead to incorrect attention results. This issue must be addressed before merging. I am unable to approve this pull request; please have other reviewers review and approve this code after the critical issue is fixed.

[gemini-code-assist]

Adding this line to reset the cuda_graph_kv_indices buffer before populating it for the current batch is a good practice. It ensures that stale indices from previous, potentially shorter, sequences do not interfere with the current batch's attention calculation during graph replay.