Skip to content

[BUG] Fix stale attention metadata after Eagle pre-planner pads the batch - #29262

Open
1096125073 wants to merge 1 commit into
sgl-project:mainfrom
1096125073:fix/eagle-dp-attn-stale-metadata
Open

1096125073 wants to merge 1 commit into
sgl-project:mainfrom
1096125073:fix/eagle-dp-attn-stale-metadata

Conversation

@1096125073

@1096125073 1096125073 commented Jun 25, 2026

Copy link
Copy Markdown

Problem

When Eagle speculative decoding + DP-attention (global_num_tokens_cpu is set) are both enabled, flash_mla_with_kvcache raises "num_splits must have shape(b+1)" with the following sequence:

  1. Eagle pre-planner calls init_forward_metadata with the unpadded batch
  2. mark_forward_metadata_ready(False) prevents re-init
  3. prepare_mlp_sync_batch pads batch_size in-place — but the stored forward_metadata still has num_splits shaped for the old unpadded batch
  4. The attention kernel receives q with the padded batch but num_splits with the unpadded shape → crash

Fix

In _forward_raw (model_runner.py), after prepare_mlp_sync_batch changes the batch, detect the staleness and clear forward_metadata_ready to let init_forward_metadata re-run with the final padded batch:

if (
    forward_batch.forward_metadata_ready
    and not forward_batch.forward_metadata_replan_equivalent
    and forward_batch.batch_size != forward_batch.forward_metadata_planned_bs
):
    forward_batch.forward_metadata_ready = False











<!-- pr-states:start -->
---
### CI States

Latest PR Test (Base): <!-- slot:pr-test:start -->:x: [Run #28145748899](https://github.com/sgl-project/sglang/actions/runs/28145748899)<!-- slot:pr-test:end -->
Latest PR Test (Extra): <!-- slot:pr-test-extra:start -->:x: [Run #28145748780](https://github.com/sgl-project/sglang/actions/runs/28145748780)<!-- slot:pr-test-extra:end -->
<!-- pr-states:end -->

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

@1096125073 1096125073 changed the title Fix stale attention metadata after Eagle pre-planner pads the batch [BUG] Fix stale attention metadata after Eagle pre-planner pads the batch Jun 26, 2026
@1096125073

Copy link
Copy Markdown
Author

@merrymercy Hi, could you help to review this fix? Thanks

@1096125073

Copy link
Copy Markdown
Author

/tag-and-rerun-ci

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant