Skip to content

[Fix] Enable graph capture and MSCCL++ for attention TP groups - #31629

Merged
merrymercy merged 4 commits into
sgl-project:mainfrom
caiocbr:caiorocha/attn_tp
Jul 27, 2026
Merged

[Fix] Enable graph capture and MSCCL++ for attention TP groups#31629
merrymercy merged 4 commits into
sgl-project:mainfrom
caiocbr:caiorocha/attn_tp

Conversation

@caiocbr

@caiocbr caiocbr commented Jul 18, 2026 •
edited by github-actions Bot
Loading

Copy link
Copy Markdown
Contributor

Motivation

When attention tensor parallelism uses a process group distinct from the main
tensor-parallel group, the attention TP group is not currently included in the
CUDA graph capture context.

As a result, specialized communicators associated with the attention TP group
are not enabled during graph capture. In addition, MSCCL++ is explicitly
disabled when creating a distinct attention TP group, even when MSCCL++ is
enabled globally.

This prevents attention TP collectives from using the configured MSCCL++
backend and may cause them to fall back to torch.distributed.

Modifications

  • Include _ATTN_TP in the process groups entered by the shared graph-capture
    context.
  • Reuse the existing identity-based deduplication so no additional context is
    entered when _ATTN_TP aliases _TP.
  • Allow distinct attention TP groups to inherit MSCCL++ all-reduce
    instead of disabling it explicitly.

CI States

Latest PR Test (Base): ⏳ Run #30310861373
Latest PR Test (Extra): ❌ Run #30310861267

@gemini-code-assist

gemini-code-assist Bot commented Jul 18, 2026

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!

@ch-wan ch-wan self-assigned this Jul 19, 2026
@ch-wan

ch-wan commented Jul 19, 2026

Copy link
Copy Markdown
Collaborator

/tag-and-rerun-ci

@github-actions github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Jul 19, 2026
@caiocbr

caiocbr commented Jul 23, 2026

Copy link
Copy Markdown
Contributor Author

/tag-and-rerun-ci

@merrymercy
merrymercy merged commit 5a46e16 into sgl-project:main Jul 27, 2026
82 of 119 checks passed
saturn-acc pushed a commit to saturn-acc/sglang that referenced this pull request Aug 16, 2026
@caiocbr @ch-wan @saturn-acc
[Fix] Enable graph capture and MSCCL++ for attention TP groups (sgl-p…
2d0c5a8 
…roject#31629)

Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
Atituiset pushed a commit to Atituiset/sglang that referenced this pull request Sep 10, 2026
@caiocbr @ch-wan
[Fix] Enable graph capture and MSCCL++ for attention TP groups (sgl-p…
bb79dfd 
…roject#31629)

Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
Richardyu114 added a commit to sammysun0711/sglang that referenced this pull request Oct 19, 2026
@Richardyu114
fix(distributed): enable AITER collectives for attention TP
975bc21 
Let attention-TP groups inherit custom-allreduce and MSCCL++ policies
and participate in graph capture with existing group deduplication.

Allow eligible attention-TP groups to use the existing AITER
reduce-scatter and all-gather paths without changing PyNccl selection.

Add regression tests for policy inheritance and aliased capture groups.

Adapted from sgl-project#31629 and sgl-project#35330.

Signed-off-by: Richardyu114 <zhentyu@amd.com>
Richardyu114 added a commit to sammysun0711/sglang that referenced this pull request Nov 2, 2026
@Richardyu114
fix(distributed): enable custom collectives for attention TP groups
97d1347 
Let attention TP groups inherit custom-allreduce and MSCCL++ policies
and participate in graph capture. Deduplicate aliased groups and retain
the existing PyNccl selection. Test policy inheritance and capture scope.

Selectively adapt the attention-group changes from the upstream PRs.

Adapted-from: sgl-project#31629
Adapted-from: sgl-project#35330
Signed-off-by: Richardyu114 <zhentyu@amd.com>
@Richardyu114 Richardyu114 mentioned this pull request Nov 2, 2026
Merged
Richardyu114 added a commit to sammysun0711/sglang that referenced this pull request Nov 3, 2026
@Richardyu114
fix(distributed): enable custom collectives for attention TP groups
547c87b 
Let attention TP groups inherit custom-allreduce and MSCCL++ policies
and participate in graph capture. Deduplicate aliased groups and retain
the existing PyNccl selection. Test policy inheritance and capture scope.

Selectively adapt the attention-group changes from the upstream PRs.

Adapted-from: sgl-project#31629
Adapted-from: sgl-project#35330
Signed-off-by: Richardyu114 <zhentyu@amd.com>
sammysun0711 pushed a commit to sammysun0711/sglang that referenced this pull request Nov 3, 2026
@Richardyu114 @hanming-lu @silencejade
[Feat][EP] Add MiMo-V2.6-Flash DP2+EP8+DP-ATTN prefill support in MI3…
e3c10a4 
…00X (#39)

* fix(rocm): always mask padded MoE expert IDs

Mask padded top-k IDs to a valid expert even when EPLB is disabled,
while retaining zero routing weights for padded rows. Add a regression
for invalid padded IDs without an expert-location remap.

Adapted-from: sgl-project#28378
Upstream-commit: 6309fb9
Signed-off-by: Richardyu114 <zhentyu@amd.com>

* perf(moriep): skip unused expert counts in captured graphs

Do not capture LocalExpertCount for the no-op recorder. Keep the
count kernel for configured recorders so recording can start after
capture. Cover inactive, active and capture-time recording policies.

Adapted-from: sgl-project#30302
Upstream-commit: 9ddea8d
Signed-off-by: Richardyu114 <zhentyu@amd.com>

* feat(mimo): route MiMo MoE through MORI expert parallelism

Recognize MORI in the MiMo A2A initialization and forward-path gates.
Use the existing dispatcher-based forward with padded-token metadata
and avoid the tensor-parallel post-expert reduction path.

Add a regression for MORI forward-path selection.

Signed-off-by: Richardyu114 <zhentyu@amd.com>

* fix(distributed): enable custom collectives for attention TP groups

Let attention TP groups inherit custom-allreduce and MSCCL++ policies
and participate in graph capture. Deduplicate aliased groups and retain
the existing PyNccl selection. Test policy inheritance and capture scope.

Selectively adapt the attention-group changes from the upstream PRs.

Adapted-from: sgl-project#31629
Adapted-from: sgl-project#35330
Signed-off-by: Richardyu114 <zhentyu@amd.com>

* perf(rocm): backport AITER reduce-scatter and capture handling

Use AITER for eligible reduce-scatter calls with ROCm defaults. Support
unregistered capture inputs for memory-saver and an opt-in AG/RS copy-in
path, retaining existing collective fallbacks.

Adapt the backport to keep RS inside the registered collective boundary
and check per-rank alignment, output shape/dtype/device and deterministic
policy. Add CPU dispatch tests and a four-GPU collective smoke test.

Only the collective portions are included; gfx950 EP decode kernels and
pre-dispatch token hints are not part of this change.

Adapted-from: sgl-project#29103
Adapted-from: sgl-project#34203
Adapted-from: sgl-project#43077
Signed-off-by: Richardyu114 <zhentyu@amd.com>

* fix(bench): wait for idle before flushing the server cache

Backport the post-warmup cache-flush race fix. Add --flush-cache-timeout
with a 60-second default for SGLang and retain backend-specific endpoints.
Raise on HTTP failure before issuing measured requests. Cover deferred
flush completion, CLI validation and benchmark abort on flush failure.

Adapted-from: sgl-project#33663
Upstream-commit: 8a1637a
Signed-off-by: Richardyu114 <zhentyu@amd.com>

* eval(mimo): prepare the Flash DP2 attention-TP4 EP8 launcher

Configure global TP8/EP8 with DP2, MORI normal mode, dense TP1 and a DP
LM head. Size and validate dispatch capacity for prefill and DFlash blocks.

Default PERFORMANCE_MODE=1 to fake routing plus PD prefill; setting it
to 0 selects real routing and ordinary serving. Keep TBO and fused QKV
quantization disabled and leave collective workaround flags unset.

The launcher prepares the target topology; DFlash DP-attention support
is still pending and the existing admission guard is unchanged.

Signed-off-by: Richardyu114 <zhentyu@amd.com>

* eval(mimo) refine dp2_attn-tp4_ep8 prefill server secript

Signed-off-by: Richardyu114 <zhentyu@amd.com>

* perf(mimo): enable tuned MORI EP8 launch configs by default

Add gfx942 tuning rules for FP8 dispatch and BF16 combine at
2048, 4096, and 8192 tokens per rank.

Default the MiMo-V2.6-Flash DP2/attention-TP4/EP8 launcher to AUTO
mode with the bundled JSON, while preserving environment overrides.
Validate config readability and log the active config path.

(cherry picked from commit 1c6b707)
Signed-off-by: Richardyu114 <zhentyu@amd.com>

* perf(mimo): pack FP8 QKV activations and scales into one all-gather

Write fused QKV quantization outputs into a shared FP32 bit-view
buffer and reuse the existing AITER/RCCL collective path. Keep
activations as row-strided views and transpose scales once for
CK/CKTile GEMMs.

Enable the path through SGLANG_MIMO_FUSED_RMS_QKV_QUANT and default
it on in the EP8 launcher. Add layout and 8-GPU graph replay tests.

Validate DFlash GSM8K accuracy and improve input throughput by 4.86%
at global chunk 32K, input 64K, and concurrency 16 on MI300X.

Signed-off-by: Richardyu114 <zhentyu@amd.com>

* feat(eval): enable prefill delayer in MiMo Flash DP2 launcher

- Enable coordinated prefill admission by default with a configurable
  maximum of 100 scheduler passes and an opt-out switch.
- Validate prefill delayer settings and log the resolved configuration.
- Expose stable router-weight fusion as an environment override,
  disabled by default.

Signed-off-by: Richardyu114 <zhentyu@amd.com>

* Negotiate PrefillDelayer only after KV-budget admission checks (sgl-project#31835)

Signed-off-by: Richardyu114 <zhentyu@amd.com>

* [Bugfix] Fix prefill suspension caused by delayed negotiate_should_allow_prefill invocation (sgl-project#32389)

Signed-off-by: Richardyu114 <zhentyu@amd.com>

* fix(ci): run MiMo EP regression tests through pytest

Add executable pytest entry points for the packed FP8 and MORI tests.
Replace the partial parallel-state runner so direct CI invocation
executes all cases and propagates failures.

Signed-off-by: Richardyu114 <zhentyu@amd.com>

---------

Signed-off-by: Richardyu114 <zhentyu@amd.com>
Co-authored-by: Hanming Lu <69857889+hanming-lu@users.noreply.github.com>
Co-authored-by: silencejade <222910030+silencejade@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@merrymercy merrymercy merrymercy approved these changes
@ch-wan ch-wan ch-wan approved these changes
@yizhang2077 yizhang2077 Awaiting requested review from yizhang2077 yizhang2077 is a code owner

Assignees

@ch-wan ch-wan

Labels

run-ci CI: run the baseline test suite on this PR

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@caiocbr @ch-wan @merrymercy