[Fix] dp-attention: Align prefill graph buckets and decode padding across DP ranks by ispobock · Pull Request #46138 · sgl-project/sglang

ispobock · 2026-10-23T17:39:17Z

Summary

Two DP-attention padding changes:

With breakable prefill CUDA graphs, every DP rank pads to the same captured token bucket. An idle rank keeps the padded shape for the collectives; a backend that opts in also skips its attention metadata there.
Decode picks SUM_LEN padding, one TP all-reduce of the summed rows, when the gather spans several ranks, every rank has tokens, the summed buffer fits the CustomAllReduceV2 buffer, and there is no MoE all-to-all or speculative decoding. Otherwise the existing cost-based choice applies.

Changes to existing code

forward_batch_info: global_num_tokens is rounded up to the smallest captured prefill bucket on every rank.
DpPaddingMode: decode selects SUM_LEN under the condition above, both eagerly and in get_default_mode_in_cuda_graph, which now receives the captured per-rank token counts.
AttentionBackend.skips_prefill_metadata_on_idle_dp_rank (default False): when a backend sets it, PrefillCudaGraphRunner skips init_forward_metadata and the read snapshot on a rank with no real tokens. No existing backend sets it, so their replay is unchanged.

The prefill change affects every model that uses DP attention with breakable prefill graphs; the decode padding change affects DP-attention decode without idle ranks.

Verification

test_dp_attention, test_dsa_backend_dp_padding (an idle rank still selects MAX_LEN), test_runtime_context, test_prefill_cuda_graph_padding, test_prefill_cuda_graph_runner, test_prefill_cuda_graph_runner_helpers, test_prefill_shared_read_done
TODO: DP-attention serving on an existing DP model, with breakable prefill graphs and with decode CUDA graphs.

Cut byte-for-byte from #45925 (MiniMax-M3.1-Flash support) at 10e655a.

CI States

Latest PR Test (Base): ✅ Run #40160842644
Latest PR Test (Extra): ✅ Run #40160842552
Latest PR Test (AMD ROCm 10): ⏳ Run #40160842664

ispobock · 2026-10-23T17:46:18Z

/rerun-test test/registered/unit/layers/test_dp_attention.py test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py

github-actions · 2026-10-23T17:46:43Z

🚀 ubuntu-latest (5 tests): ❌ View workflow run

cd test/ && python3 registered/unit/layers/test_dp_attention.py
cd test/ && python3 registered/unit/layers/attention/test_dsa_backend_dp_padding.py
cd test/ && python3 registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py
cd test/ && python3 registered/unit/model_executor/test_prefill_cuda_graph_runner.py
cd test/ && python3 registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py

Co-authored-by: Chunan Zeng <zcnrex@gmail.com>

Co-authored-by: Chunan Zeng <zcnrex@gmail.com> Co-authored-by: Denji-kk <51020465+Denji-kk@users.noreply.github.com>

github-actions · 2026-10-23T17:53:33Z

🚀 ubuntu-latest (5 tests): ❌ View workflow run

cd test/ && python3 registered/unit/layers/test_dp_attention.py
cd test/ && python3 registered/unit/layers/attention/test_dsa_backend_dp_padding.py
cd test/ && python3 registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py
cd test/ && python3 registered/unit/model_executor/test_prefill_cuda_graph_runner.py
cd test/ && python3 registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py

ispobock · 2026-10-23T17:59:40Z

/rerun-test test/registered/unit/layers/test_dp_attention.py test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py

github-actions · 2026-10-23T18:00:04Z

🚀 ubuntu-latest (5 tests): ❌ View workflow run

cd test/ && python3 registered/unit/layers/test_dp_attention.py
cd test/ && python3 registered/unit/layers/attention/test_dsa_backend_dp_padding.py
cd test/ && python3 registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py
cd test/ && python3 registered/unit/model_executor/test_prefill_cuda_graph_runner.py
cd test/ && python3 registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py

ispobock · 2026-10-23T18:16:05Z

/rerun-test test/registered/unit/layers/test_dp_attention.py test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py

github-actions · 2026-10-23T18:16:30Z

🚀 ubuntu-latest (5 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/test_dp_attention.py
cd test/ && python3 registered/unit/layers/attention/test_dsa_backend_dp_padding.py
cd test/ && python3 registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py
cd test/ && python3 registered/unit/model_executor/test_prefill_cuda_graph_runner.py
cd test/ && python3 registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py

ispobock · 2026-10-24T09:58:40Z

/rerun-test test/registered/unit/layers/test_dp_attention.py test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py test/registered/unit/test_runtime_context.py test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py

github-actions · 2026-10-24T09:59:00Z

🚀 ubuntu-latest (7 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/test_dp_attention.py
cd test/ && python3 registered/unit/layers/attention/test_dsa_backend_dp_padding.py
cd test/ && python3 registered/unit/test_runtime_context.py
cd test/ && python3 registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py
cd test/ && python3 registered/unit/model_executor/test_prefill_cuda_graph_runner.py
cd test/ && python3 registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py
cd test/ && python3 registered/unit/model_executor/runner/test_prefill_shared_read_done.py

mintlify · 2026-10-25T03:30:30Z

Preview deployment for your docs. Learn more about Mintlify Previews.

Project	Status	Preview	Updated
sglang-doc	🟢 Ready	View Preview	Oct 25, 2026, 4:13 AM

💡 Tip: Enable Automations to automatically generate PRs for you.

# Conflicts: # python/sglang/srt/layers/layer_boundary/fusions/custom_all_reduce_fusion.py # test/registered/unit/layer_boundary/test_custom_all_reduce_fusion.py

ispobock · 2026-10-25T04:13:39Z

/run-full-ci

ispobock · 2026-10-25T08:50:13Z

cc @yuxwei @ch-wan

ispobock requested review from BBuf, Edwardf0t1, Fridge003, HaiShaw, Ying1123, ch-wan, hnyls2002 and merrymercy as code owners October 23, 2026 17:39

ispobock mentioned this pull request Oct 23, 2026

[Feat] msa: Add MSA-v1 decode and NVFP4 KV cache kernels #46139

Merged

This comment has been minimized.

Sign in to view

ispobock and others added 2 commits October 24, 2026 01:51

fuse tp all-reduce, residual and rmsnorm on push plane

d0ec8c6

Co-authored-by: Chunan Zeng <zcnrex@gmail.com>

pad dp ranks to one prefill graph bucket

193e54b

Co-authored-by: Chunan Zeng <zcnrex@gmail.com> Co-authored-by: Denji-kk <51020465+Denji-kk@users.noreply.github.com>

ispobock force-pushed the ar-residual-norm-fusion branch from 9343589 to d0ec8c6 Compare October 23, 2026 17:52

ispobock requested review from Alisehen, AniZpZ, DarkSharpness, FlamingoPg, HydraQYH, JustinTong0323, OrangeRedeng, b8zhong, celve, fzyzcjy, sogalin, wisclmy0611, yuan-luo and zijiexia as code owners October 23, 2026 17:52

ispobock force-pushed the dp-prefill-graph-bucket branch from 7e9662d to 193e54b Compare October 23, 2026 17:52

ispobock added 2 commits October 24, 2026 01:58

add test main entry

c83a69b

Merge branch 'ar-residual-norm-fusion' into dp-prefill-graph-bucket

6a8fb28

keep max_len choice with idle dp ranks

29859e0

ispobock added 2 commits October 24, 2026 17:33

skip sum_len check for one-rank gathers

dc80859

opt in to skipping metadata on idle dp ranks

2480c8d

ispobock requested review from Qiaolin-Yu and hebiao064 as code owners October 24, 2026 09:58

This was referenced Oct 24, 2026

[Perf] Qwen3.8: Allow breakable prefill CUDA graphs for image batches eppaneamd/sglang#2

Closed

[Perf] Qwen3.8: Allow breakable prefill CUDA graphs for image batches #46329

Open

Base automatically changed from ar-residual-norm-fusion to main October 25, 2026 03:30

Merge branch 'main' into dp-prefill-graph-bucket

acf3199

# Conflicts: # python/sglang/srt/layers/layer_boundary/fusions/custom_all_reduce_fusion.py # test/registered/unit/layer_boundary/test_custom_all_reduce_fusion.py

mintlify Bot deployed to staging - docs October 25, 2026 03:35

View deployment

mintlify Bot deployed to staging - docs October 25, 2026 04:13

View deployment

github-actions Bot added run-ci CI: run the baseline test suite on this PR run-ci-extra CI: also run the extra suite (requires run-ci) labels Oct 25, 2026

ispobock merged commit 9b230c1 into main Oct 25, 2026
304 of 331 checks passed

ispobock deleted the dp-prefill-graph-bucket branch October 25, 2026 08:50

arathi-hlab mentioned this pull request Oct 25, 2026

[Fix] [Intel] dp-attention: Skip the CustomAllReduceV2 import on XPU #46417

Merged

1 of 2 tasks

Juhyun-Kim-Memphis mentioned this pull request Oct 27, 2026

[INF-674] Backport upstream DP-attention, EPLB-drafter and Hopper DeepSelect fixes onto v0.5.22 DeepSeek-V4.1 vessl-ai/sglang#89

Merged

ispobock mentioned this pull request Oct 27, 2026

[Tracking] MiniMax-M3.1-Flash support #46869

Open

13 tasks done

Conversation

ispobock commented Oct 23, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes to existing code

Verification

CI States

Uh oh!

This comment has been minimized.

ispobock commented Oct 23, 2026

Uh oh!

github-actions Bot commented Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ispobock commented Oct 23, 2026

Uh oh!

github-actions Bot commented Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ispobock commented Oct 23, 2026

Uh oh!

github-actions Bot commented Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ispobock commented Oct 24, 2026

Uh oh!

github-actions Bot commented Oct 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mintlify Bot commented Oct 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ispobock commented Oct 25, 2026

Uh oh!

ispobock commented Oct 25, 2026

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

ispobock commented Oct 23, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Oct 23, 2026 •

edited

Loading

github-actions Bot commented Oct 23, 2026 •

edited

Loading

github-actions Bot commented Oct 23, 2026 •

edited

Loading

github-actions Bot commented Oct 23, 2026 •

edited

Loading

github-actions Bot commented Oct 24, 2026 •

edited

Loading

mintlify Bot commented Oct 25, 2026 •

edited

Loading