Skip to content

Remove unnecessary metadata_expand.max_seq_len_k operations in fa3 to… #7140

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 2 commits into from
Jun 13, 2025
Merged

Remove unnecessary metadata_expand.max_seq_len_k operations in fa3 to… #7140

merged 2 commits into from
Jun 13, 2025

Conversation

byjiang1996
Copy link
Contributor

@byjiang1996 byjiang1996 commented Jun 13, 2025 •
edited by hebiao064
Loading

Before: 15us + 18us -> After: 0

Motivation

Remove unnecessary metadata_expand.max_seq_len_k operations in fa3 - there is no code reference to metadata_expand.max_seq_len_k variable anymore

Profile scripts:

python3 -m sglang.launch_server --speculative-num-steps 3 --speculative-eagle-topk 2 --speculative-num-draft-tokens 4 --attention-backend fa3 --trust-remote-code --dtype float16 --speculative-algo EAGLE3 --model-path /shared/public/models/meta-llama/Meta-Llama-3.1-8B-Instruct --speculative-draft-model-path /shared/public/elr-models/jamesliu1/sglang-EAGLE3-Llama-3.1-Instruct-8B/e5ed08d66f528a95ce89f5d4fd136a28f6def714 --disable-radix-cache

python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 512 --request-rate 16 --num-prompt 512 --model-path /shared/public/models/meta-llama/Meta-Llama-3.1-8B-Instruct

Before: 50% of init_forward_metadata_replay_cuda_graph steps calculate metadata_expand.max_seq_len_k using metadata_expand.cache_seqlens_int32.max().item() (max() & item() function)
image

After: no more max() or item() calculation (15us + 18us=33us) within init_forward_metadata_replay_cuda_graph for metadata_expand.max_seq_len_k
image

Validated correctness via gsm8k:

python -m sglang.launch_server --model-path /shared/public/models/meta-llama/Meta-Llama-3.1-8B-Instruct --attention-backend fa3

python benchmark/gsm8k/bench_sglang.py --num-shots 8 --num-questions 1319 --parallel 1319 --
data-path /shared/public/data/gsm8k/test.jsonl 

Accuracy: 0.790
Invalid: 0.001
Latency: 16.377 s
Output throughput: 8214.036 token/s

Modifications

Checklist

gemini-code-assist[bot]
gemini-code-assist bot reviewed Jun 13, 2025
View reviewed changes
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Summary of Changes

Hello @byjiang1996, I'm Gemini Code Assist1! 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 focuses on a micro-optimization within the FlashAttention v3 implementation used in the SGLang runtime. It identifies and removes calculations for a metadata field (max_seq_len_k) that is no longer used, particularly in the speculative decoding path. This cleanup aims to reduce computational overhead and improve performance by eliminating dead code execution.

Highlights

  • Performance Optimization: Removed unnecessary calculations for metadata_expand.max_seq_len_k within the FlashAttention v3 backend, specifically in the speculative decode path.
  • Code Cleanup: Eliminated assignments and calculations for a variable (metadata_expand.max_seq_len_k) that is no longer referenced in the code.
  • Expected Acceleration: The removal of these operations is expected to provide a performance improvement of 33+ microseconds.
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 configureGemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist[bot]
gemini-code-assist bot reviewed Jun 13, 2025
View reviewed changes
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

The pull request successfully removes the calculation and assignment of the metadata_expand.max_seq_len_k variable in specific speculative decoding paths (Draft Decode and Target Verify when topk > 1). Based on the code review, this variable was indeed not being used after these assignments, making the removal a valid optimization for performance as described in the motivation. The changes are minimal, targeted, and improve efficiency without impacting correctness. The code generally follows common Python practices, similar to PEP 8.

hebiao064
hebiao064 approved these changes Jun 13, 2025
View reviewed changes
Copy link
Collaborator

@hebiao064 hebiao064 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks! 🚀 🚀 🚀

@zhyncs zhyncs merged commit 22a6b9f into sgl-project:main Jun 13, 2025
39 of 53 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@hebiao064 hebiao064 hebiao064 approved these changes

@merrymercy merrymercy Awaiting requested review from merrymercy merrymercy is a code owner

@Ying1123 Ying1123 Awaiting requested review from Ying1123 Ying1123 is a code owner

@zhyncs zhyncs Awaiting requested review from zhyncs zhyncs is a code owner

@ispobock ispobock Awaiting requested review from ispobock ispobock is a code owner

@HaiShaw HaiShaw Awaiting requested review from HaiShaw HaiShaw is a code owner

@ch-wan ch-wan Awaiting requested review from ch-wan ch-wan is a code owner

@BBuf BBuf Awaiting requested review from BBuf BBuf is a code owner

+1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@byjiang1996 @hebiao064 @zhyncs