Skip to content

Simplify prepare_extend_after_decode and remove unnecessary padding #6987

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 3 commits into from
Jun 9, 2025
Merged

Simplify prepare_extend_after_decode and remove unnecessary padding #6987

merged 3 commits into from
Jun 9, 2025

Conversation

merrymercy
Copy link
Contributor

@merrymercy merrymercy commented Jun 9, 2025 •
edited
Loading

  • Remove redundant padding in prepare_extend_after_decode
  • Remove some redundant device sync

before

acc_length=2.88
speed=310.63 token/s

after

acc_length=2.88
speed=322.92 token/s

gemini-code-assist[bot]
gemini-code-assist bot reviewed Jun 9, 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.

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

Summary of Changes

Hello team, gemini-code-assist here with a summary of this pull request titled "Simplify prepare_extend_after_decode".

The primary goal of this PR is to simplify the logic within the prepare_extend_after_decode function, which is crucial for the speculative decoding process, particularly with the EAGLE draft model. This simplification involves removing the input padding logic from this function and refactoring the associated Triton kernel (create_extend_spec_info) to handle cumulative sums internally, making the Python code cleaner.

Additionally, the pull request refactors the model_capture_mode context manager, moving it from a class method to a standalone function for better reusability and cleaner code in the CUDA graph runners. Minor changes include reordering conditions in a batch scheduling check and reordering options in the attention backend argument list.

Highlights

  • Speculative Decoding Simplification: The prepare_extend_after_decode function in eagle_utils.py has been significantly simplified by removing the input padding logic. This responsibility is likely handled elsewhere or deemed unnecessary in this specific context.
  • Triton Kernel Refactoring: The create_extend_spec_info Triton kernel has been updated to calculate the cumulative sum of accept_lens internally using Triton's tl.sum and masking, removing the need for a pre-calculated cumulative sum argument.
  • CUDA Graph Context Manager Refactor: The model_capture_mode context manager has been moved from being a method of CudaGraphRunner to a standalone function in cuda_graph_runner.py. This standalone function is now used consistently across the CUDA graph runners (CudaGraphRunner, eagle_draft_cuda_graph_runner.py, eagle_draft_extend_cuda_graph_runner.py).
  • CUDA Graph Replay Condition Update: In eagle_draft_extend_cuda_graph_runner.py, the condition for resetting state during CUDA graph replay in the replay method has been changed from checking bs != raw_bs to bs * self.num_tokens_per_bs != num_tokens. This adjusts when the state is reset based on the total number of tokens rather than just the batch size.

Changelog

Click here to see the changelog
  • python/sglang/srt/managers/schedule_batch.py
    • Reordered conditions in the if statement checking which attention backends require seq_lens_cpu.
  • python/sglang/srt/model_executor/cuda_graph_runner.py
    • Moved the model_capture_mode context manager from a class method to a standalone function (lines 59-66).
    • Updated __init__ to use the standalone model_capture_mode function (line 304).
  • python/sglang/srt/server_args.py
    • Reordered the list of choices for the --attention-backend command-line argument.
  • python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
    • Imported the standalone model_capture_mode function (line 13).
    • Updated __init__ to use the standalone model_capture_mode function (line 84).
  • python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
    • Imported the standalone model_capture_mode function (line 14).
    • Updated __init__ to use the standalone model_capture_mode function (line 91).
    • Changed the condition for resetting state in the replay method from bs != raw_bs to bs * self.num_tokens_per_bs != num_tokens (line 203).
  • python/sglang/srt/speculative/eagle_utils.py
    • Imported ForwardMode (line 25).
    • Refactored prepare_extend_after_decode to remove input padding logic and simplify batch attribute assignments (lines 92-113).
    • Refactored the create_extend_spec_info Triton kernel to remove the accept_len_cum argument and calculate cumulative sums internally (lines 569-590).
  • python/sglang/srt/speculative/eagle_worker.py
    • Removed setting batch.forward_mode, batch.spec_info.capture_hidden_mode, and batch.return_logprob as this is now handled within prepare_extend_after_decode (lines 686, 693, 694 removed).
    • Added logic to calculate forward_batch.seq_lens_sum (lines 696-699).
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[bot]
gemini-code-assist bot suggested changes Jun 9, 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

This pull request introduces a simplification to prepare_extend_after_decode and refactors model_capture_mode into a global utility, which is a positive step towards cleaner code.

However, I've found some critical issues in python/sglang/srt/speculative/eagle_utils.py concerning tensor sizing and the data prepared for the draft model's forward pass. Additionally, there's a potential issue regarding the handling of CUDA graph inputs after the removal of explicit padding logic, and an efficiency concern with the modified Triton kernel.

These issues need to be addressed to ensure correctness and performance.

Summary of Findings

  • Critical: Tensor Sizing in prepare_extend_after_decode: self.positions is likely initialized with an incorrect (smaller) size, potentially leading to out-of-bounds writes in the Triton kernel.
  • Critical: Incorrect input_ids for Draft Extend: batch.input_ids for the draft extend forward pass is not correctly assembled. It currently seems to only contain bonus tokens from the previous step, not the full sequence of accepted tokens corresponding to the computed positions.
  • High: CUDA Graph Input Padding: Explicit padding logic was removed from prepare_extend_after_decode. It's unclear if the fixed-shape input requirements for CUDA graphs (especially for input_ids and positions in EAGLEDraftExtendCudaGraphRunner) are still fully met, as uninitialized trailing data in tensors could cause issues.
  • Medium: Triton Kernel Efficiency: The create_extend_spec_info Triton kernel now computes cumulative sums internally per kernel launch, which might be less efficient for large batch sizes compared to using a precomputed sum.
  • Refactoring of model_capture_mode: The model_capture_mode has been refactored into a global context manager, which improves code organization and reusability. This is a good change.

Merge Readiness

The pull request has some positive refactoring, particularly with model_capture_mode. However, there are critical issues identified in prepare_extend_after_decode related to tensor sizing and data preparation that must be addressed before merging. Additionally, clarification is needed on how CUDA graph input padding is handled after the recent changes. I recommend addressing these critical and high-severity issues. I am not authorized to approve pull requests; please ensure further review and approval by authorized maintainers after these points are resolved.

@merrymercy merrymercy force-pushed the lianmin/sync-eagle-oss branch from dc8fe31 to bd30277 Compare June 9, 2025 11:34
@merrymercy merrymercy force-pushed the lianmin/sync-eagle-oss branch from 4aee40e to a9d53a3 Compare June 9, 2025 12:34
@merrymercy merrymercy merged commit dc0705a into main Jun 9, 2025
60 of 73 checks passed
@merrymercy merrymercy deleted the lianmin/sync-eagle-oss branch June 9, 2025 23:39
almaslof pushed a commit to mpashkovskii/sglang that referenced this pull request Jun 11, 2025
@merrymercy @almaslof
Simplify prepare_extend_after_decode (sgl-project#6987)
8368a97
jianan-gu pushed a commit to jianan-gu/sglang that referenced this pull request Jun 12, 2025
@merrymercy @jianan-gu
Simplify prepare_extend_after_decode (sgl-project#6987)
f492ede
@merrymercy merrymercy changed the title Simplify prepare_extend_after_decode Simplify prepare_extend_after_decode and remove unnecessary padding Jun 15, 2025
walker-ai pushed a commit to walker-ai/sglang that referenced this pull request Jul 8, 2025
@yudian0504
PullRequest: 52 sgl_20250610_sync_tag047
a19e2e2 
Merge branch 'sgl_20250610_sync_tag047 of git@code.alipay.com:Theta/SGLang.git into main

https://code.alipay.com/Theta/SGLang/pull_requests/52


Reviewed-by: 剑川 <jianchuan.gys@antgroup.com>


* [Bugfix] Fix slice operation when chunk size mismatch (sgl-project#6697)
* [Bugfix] Fix ChatCompletion endpoint of mini_lb when stream is set (sgl-project#6703)
* [CI] Fix setup of disaggregation with different tp (sgl-project#6706)
* [PD] Remove Unnecessary Exception Handling for FastQueue.get() (sgl-project#6712)
* Fuse routed_scaling_factor in DeepSeek (sgl-project#6710)
* Overlap two kernels in DeepSeek with communication (sgl-project#6711)
* Minor refactor two-batch overlap (sgl-project#6682)
* Speed up when having padding tokens two-batch overlap (sgl-project#6668)
* [Feature] Support Flashinfer fp8 blockwise GEMM kernel on Blackwell (sgl-project#6479)
* Fix LoRA bench (sgl-project#6719)
* temp
* Fix PP for Qwen3 MoE (sgl-project#6709)
* [feat] triton kernel for get_last_loc (sgl-project#6676)
* [fix] more mem for draft_extend cuda_graph (sgl-project#6726)
* [PD] bug fix:  Update status if nixl receiver send a a dummy req. (sgl-project#6720)
* Tune memory arguments on B200 (sgl-project#6718)
* Add DeepSeek-R1-0528 function call chat template (sgl-project#6725)
* refactor(tool call): Fix BaseFormatDetector tool_index issue and refactor `parse_streaming_increment` (sgl-project#6715)
* Add draft extend CUDA graph for Triton backend (sgl-project#6705)
* refactor apply_w8a8_block_fp8_linear in fp (sgl-project#6545)
* [PD] Support completion endpoint (sgl-project#6729)
* PD Rust LB (PO2) (sgl-project#6437)
* Super tiny enable sole usage of expert distribution metrics and update doc (sgl-project#6680)
* Support picking variants of EPLB algorithms (sgl-project#6728)
* Support tuning DeepEP configs (sgl-project#6742)
* [test] add ut and bm for get_last_loc (sgl-project#6746)
* Fix mem_fraction_static for AMD CI (sgl-project#6748)
* [fix][RL] Fix DeepSeekV3ForCausalLM.post_load_weights for multiple update weight (sgl-project#6265)
* Improve EPLB logical to physical dispatch map (sgl-project#6727)
* Update DeepSeek-R1-0528 function call chat template (sgl-project#6765)
* [PD] Optimize time out logic and add env var doc for mooncake (sgl-project#6761)
* Fix aiohttp 'Chunk too big' in bench_serving (sgl-project#6737)
* Support sliding window in triton backend (sgl-project#6509)
* Fix shared experts fusion error (sgl-project#6289)
* Fix one bug in the grouped-gemm triton kernel (sgl-project#6772)
* update llama4 chat template and pythonic parser (sgl-project#6679)
* feat(tool call): Enhance Llama32Detector for improved JSON parsing in non-stream (sgl-project#6784)
* Support token-level quantization for EP MoE (sgl-project#6782)
* Temporarily lower mmlu threshold for triton sliding window backend (sgl-project#6785)
* ci: relax test_function_call_required (sgl-project#6786)
* Add intel_amx backend for Radix Attention for CPU (sgl-project#6408)
* Fix incorrect LoRA weight loading for fused gate_up_proj (sgl-project#6734)
* fix(PD-disaggregation): Can not get local ip (sgl-project#6792)
* [FIX] mmmu bench serving result display error (sgl-project#6525) (sgl-project#6791)
* Bump torch to 2.7.0 (sgl-project#6788)
* chore: bump sgl-kernel v0.1.5 (sgl-project#6794)
* Improve profiler and integrate profiler in bench_one_batch_server (sgl-project#6787)
* chore: upgrade sgl-kernel v0.1.5 (sgl-project#6795)
* [Minor] Always append newline after image token when parsing chat message (sgl-project#6797)
* Update CI tests for Llama4 models (sgl-project#6421)
* [Feat] Enable PDL automatically on Hopper architecture (sgl-project#5981)
* chore: update blackwell docker (sgl-project#6800)
* misc: cache is_hopper_arch (sgl-project#6799)
* Remove contiguous before Flashinfer groupwise fp8 gemm (sgl-project#6804)
* Correctly abort the failed grammar requests & Improve the handling of abort (sgl-project#6803)
* [EP] Add cuda kernel for moe_ep_pre_reorder (sgl-project#6699)
* Add draft extend CUDA graph for flashinfer backend  (sgl-project#6805)
* Refactor CustomOp to avoid confusing bugs (sgl-project#5382)
* Tiny log prefill time (sgl-project#6780)
* Tiny fix EPLB assertion about rebalancing period and recorder window size (sgl-project#6813)
* Add simple utility to dump tensors for debugging (sgl-project#6815)
* Fix profiles do not have consistent names (sgl-project#6811)
* Speed up rebalancing when using non-static dispatch algorithms (sgl-project#6812)
* [1/2] Add Kernel support for Cutlass based Fused FP4 MoE (sgl-project#6093)
* [Router] Fix k8s Service Discovery (sgl-project#6766)
* Add CPU optimized kernels for topk and rope fusions  (sgl-project#6456)
* fix new_page_count_next_decode (sgl-project#6671)
* Fix wrong weight reference in dynamic EPLB (sgl-project#6818)
* Minor add metrics to expert location updater (sgl-project#6816)
* [Refactor] Rename `n_share_experts_fusion` as `num_fused_shared_experts` (sgl-project#6735)
* [FEAT] Add transformers backend support  (sgl-project#5929)
* [fix] recover auto-dispatch for rmsnorm and rope (sgl-project#6745)
* fix ep_moe_reorder kernel bugs (sgl-project#6858)
* [Refactor] Multimodal data processing for VLM (sgl-project#6659)
* Decoder-only Scoring API (sgl-project#6460)
* feat: add dp-rank to KV events (sgl-project#6852)
* Set `num_fused_shared_experts` as `num_shared_experts` when shared_experts fusion is not disabled (sgl-project#6736)
* Fix one missing arg in DeepEP (sgl-project#6878)
* Support LoRA in TestOpenAIVisionServer and fix fused kv_proj loading bug. (sgl-project#6861)
* support 1 shot allreduce  in 1-node and 2-node using mscclpp (sgl-project#6277)
* Fix Qwen3MoE missing token padding optimization (sgl-project#6820)
* Tiny update error hints (sgl-project#6846)
* Support layerwise rebalancing experts (sgl-project#6851)
* Tiny allow profiler API to auto create directory (sgl-project#6865)
* Support Blackwell DeepEP docker images (sgl-project#6868)
* [EP] Add cuda kernel for moe_ep_post_reorder (sgl-project#6837)
* [theta]merge 0605
* oai: fix openAI client error with single request via batch api (sgl-project#6170)
* [PD] Fix potential perf spike caused by tracker gc and optimize doc (sgl-project#6764)
* Use deepgemm instead of triton for fused_qkv_a_proj_with_mqa (sgl-project#6890)
* [CUTLASS-FP4-MOE]  Introduce CutlassMoEParams class for easy initialization of Cutlass Grouped Gems Metadata (sgl-project#6887)
* bugfix(OAI): Fix image_data processing for jinja chat templates (sgl-project#6877)
* [CPU] enable CI for PRs, add Dockerfile and auto build task (sgl-project#6458)
* AITER backend extension and workload optimizations (sgl-project#6838)
* [theta]merge
* [theta]merge
* [Feature] Support Flashinfer fmha on Blackwell (sgl-project#6930)
* Fix a bug in abort & Improve docstrings for abort (sgl-project#6931)
* Tiny support customize DeepEP max dispatch tokens per rank (sgl-project#6934)
* Sync the changes on cuda graph runners (sgl-project#6932)
* [PD] Optimize transfer queue forward logic for dummy rank (sgl-project#6922)
* [Refactor] image data process in bench_serving (sgl-project#6879)
* [fix] logical_to_all_physical_map index 256 is out of bounds in EP parallel. (sgl-project#6767)
* Add triton fused moe kernel config for E=257 on B200 (sgl-project#6939)
* [sgl-kernel] update deepgemm (sgl-project#6942)
* chore: bump sgl-kernel v0.1.6 (sgl-project#6943)
* Minor compile fused topk (sgl-project#6944)
* [Bugfix] pipeline parallelism and Eagle Qwen2 (sgl-project#6910)
* Tiny re-introduce profile id logging (sgl-project#6912)
* Add triton version as a fused_moe_triton config search key to avoid performace decrease in different Triton version (sgl-project#5955)
* reduce torch.zeros overhead in moe align block size kernel (sgl-project#6369)
* chore: upgrade sgl-kernel v0.1.6 (sgl-project#6945)
* add fbgemm moe grouped gemm kernel benchmark (sgl-project#6924)
* [Docker] Add docker file for SGL Router (sgl-project#6915)
* Disabling mixed chunked prefill when eagle is enabled (sgl-project#6874)
* Add canary for EPLB rebalancing (sgl-project#6895)
* Refactor global_server_args_dict (sgl-project#6866)
* Fuse routed scaling factor in topk_reduce kernel (sgl-project#6220)
* Update server timeout time in AMD CI. (sgl-project#6953)
* [misc] add is_cpu() (sgl-project#6950)
* Add H20 fused MoE kernel tuning configs for DeepSeek-R1/V3 (sgl-project#6885)
* Add a CUDA kernel for fusing mapping and weighted sum for MoE. (sgl-project#6916)
* chore: bump sgl-kernel v0.1.6.post1 (sgl-project#6955)
* chore: upgrade sgl-kernel v0.1.6.post1 (sgl-project#6957)
* [DeepseekR1-FP4] Add Support for nvidia/DeepSeekR1-FP4 model (sgl-project#6853)
* Revert "Fuse routed scaling factor in topk_reduce kernel (sgl-project#6220)" (sgl-project#6968)
* [AMD] Add more tests to per-commit-amd (sgl-project#6926)
* chore: bump sgl-kernel v0.1.7 (sgl-project#6963)
* Slightly improve the sampler to skip unnecessary steps (sgl-project#6956)
* rebase h20 fused_moe config (sgl-project#6966)
* Fix CI and triton moe Configs (sgl-project#6974)
* Remove unnecessary kernels of num_token_non_padded (sgl-project#6965)
* Extend cuda graph capture bs for B200 (sgl-project#6937)
* Fuse routed scaling factor in deepseek (sgl-project#6970)
* Sync cuda graph runners (sgl-project#6976)
* Fix draft extend ut stability with flush cache (sgl-project#6979)
* Fix triton sliding window test case (sgl-project#6981)
* Fix expert distribution dumping causes OOM (sgl-project#6967)
* Minor remove one kernel for DeepSeek (sgl-project#6977)
* [perf][sgl-kernel] extend cutlass_mla_decode to support num_head < 128 (sgl-project#6929)
* Enable more unit tests for AMD CI. (sgl-project#6983)
* Use torch.compile to fuse flash attention decode metadata preparation (sgl-project#6973)
* Eliminate stream sync to speed up LoRA batch init  (sgl-project#6960)
* support qwen3 emebedding (sgl-project#6990)
* Fix torch profiler bugs for bench_offline_throughput.py (sgl-project#6557)
* chore: upgrade flashinfer v0.2.6.post1 jit (sgl-project#6958)
* cleanup tmp dir (sgl-project#7007)
* chore: update pr test xeon (sgl-project#7008)
* Fix cutlass MLA gets almost zero accuracy (sgl-project#6998)
* Update amd nightly models CI. (sgl-project#6992)
* feat: add direct routing strategy to DP worker (sgl-project#6884)
* Fallback to lower triton version for unfound fused moe configs (sgl-project#7013)
* Fix torchvision version for Blackwell (sgl-project#7015)
* Simplify prepare_extend_after_decode (sgl-project#6987)
* Migrate to assertEqual (sgl-project#6741)
* Fix torch version in blackwell dockerfile (sgl-project#7017)
* chore: update pr test xeon (sgl-project#7018)
* Update default settings for blackwell (sgl-project#7023)
* Support both approximate and exact expert distribution collection (sgl-project#6964)
* Add decode req pool (sgl-project#6980)
* [theta]merge 0610
* [theta]merge 0610
* [CI] Add CI workflow for sgl-router docker build (sgl-project#7027)
* Fix fused_moe triton configs (sgl-project#7029)
* CPU: map changes from developing branch in sgl-kernel (sgl-project#6833)
* chore: bump v0.4.7 (sgl-project#7038)
* Update README.md (sgl-project#7040)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

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

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

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

@hnyls2002 hnyls2002 Awaiting requested review from hnyls2002 hnyls2002 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

@xiezhq-hermann xiezhq-hermann Awaiting requested review from xiezhq-hermann xiezhq-hermann is a code owner

@ByronHsu ByronHsu Awaiting requested review from ByronHsu

@HaiShaw HaiShaw Awaiting requested review from HaiShaw

@ch-wan ch-wan Awaiting requested review from ch-wan

@BBuf BBuf Awaiting requested review from BBuf

1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] requested changes

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

Successfully merging this pull request may close these issues.
2 participants
@merrymercy @zhyncs