Skip to content

minimal_async_ep: enable eager EP overlap - #3955

Draft
sanketpurandare wants to merge 6 commits into
gh/sanketpurandare/6/basefrom
gh/sanketpurandare/6/head
Draft

minimal_async_ep: enable eager EP overlap#3955
sanketpurandare wants to merge 6 commits into
gh/sanketpurandare/6/basefrom
gh/sanketpurandare/6/head

Conversation

@sanketpurandare

@sanketpurandare sanketpurandare commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Summary

Enable eager two-chunk EP overlap for MinimalAsyncEP. Allocate independent symmetric-memory buffer sets, give dispatch and combine distinct pools sized to their own receive capacities, select stable buffer ownership while tracing each chunk, and teach the shared EP scheduler to recognize MinimalAsyncEP launch and wait markers.

Force rematerialization inside MinimalAsyncEP MoE regions so receive-buffer views are never saved across backward. The existing one-set path remains unchanged when overlap is disabled.

Why

Launch/wait separation is not sufficient when adjacent chunks can reuse static receive storage. Buffer ownership, separate exchange pools, and rematerialization are part of the correctness contract when communication moves across expert compute. Exact pool sizing also avoids allocating combine storage at the larger dispatch bound.

Test Plan

  • Run real two-rank dispatch, expert computation, combine, reduction, and backward through two independent buffer sets.
  • Compare outputs and gradients with a local reference.
  • Verify dispatch/combine pool capacities and storage separation.
  • Verify bounded CTA configuration and prep/copy/wait stream placement.
  • Run the two-rank dispatcher, full-Inductor waits, and CUDA graph wrapping in H100 CI.

## Summary

Enable eager two-chunk EP overlap for MinimalAsyncEP. Allocate independent
symmetric-memory buffer sets, select stable buffer ownership while tracing each
chunk, and teach the shared EP scheduler to recognize MinimalAsyncEP launch and
wait markers.

Force rematerialization inside MinimalAsyncEP MoE regions so receive-buffer
views are never saved across backward. The existing one-buffer path remains
unchanged when overlap is disabled.

## Why

Launch/wait separation is not sufficient when adjacent chunks can reuse the
same static receive storage. Buffer ownership and rematerialization are part of
the correctness contract when communication moves across expert compute.

## Test Plan

- Run real two-rank dispatch, expert computation, combine, reduction, and
  backward through two independent buffer sets.
- Compare outputs and gradients with a local reference.
- Verify bounded CTA configuration and prep/copy/wait stream placement.
- Run the two-rank dispatcher, full-Inductor waits, and CUDA graph wrapping in
  H100 CI.

[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/8gpu CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant