Skip to content

Support async scheduling for MoE and expert parallelism - #18

Closed
lmcafee-nvidia wants to merge 2 commits into
defer-req-res-mamba-single-bankfrom
defer-req-res-moe
Closed

Support async scheduling for MoE and expert parallelism#18
lmcafee-nvidia wants to merge 2 commits into
defer-req-res-mamba-single-bankfrom
defer-req-res-moe

Conversation

@lmcafee-nvidia

@lmcafee-nvidia lmcafee-nvidia commented Jul 10, 2026

Copy link
Copy Markdown
Owner

Summary

Enable serial and overlap async scheduling for MoE models and expert-parallel execution.

  • Inherit standalone primer steps from the single-forward parent and single-bank Mamba state from the Mamba parent.
  • Remove the MoE and expert-parallel compatibility rejections from async scheduling.
  • Advance the existing expert-parallel work-consensus cadence when a primer-only controller result immediately re-enters the controller.
  • Retain the existing restriction on MoE routing replay.

Architecture

This PR contains only the MoE/EP integration that remains after the generic one-forward-per-step behavior was moved into the single-forward parent. A primer still returns from the controller and immediately continues without normal engine bookkeeping. When expert parallelism is active, the engine accounts for that continuation in the same consensus cadence used by ordinary coordinator iterations, keeping all participating ranks on the same forward slot.

The model, controller scheduling flow, and token dispatcher are otherwise unchanged.

@lmcafee-nvidia
lmcafee-nvidia force-pushed the defer-req-res-moe branch 2 times, most recently from f23c25a to ffc7459 Compare July 14, 2026 17:01
@lmcafee-nvidia
lmcafee-nvidia changed the base branch from defer-req-res-mamba to defer-req-res-mamba-single-bank July 14, 2026 21:04
@lmcafee-nvidia
lmcafee-nvidia force-pushed the defer-req-res-mamba-single-bank branch from d326910 to ee6de13 Compare July 14, 2026 21:23
Signed-off-by: Lawrence McAfee <lmcafee@nvidia.com>
Signed-off-by: Lawrence McAfee <lmcafee@nvidia.com>
@lmcafee-nvidia
lmcafee-nvidia force-pushed the defer-req-res-mamba-single-bank branch from ee6de13 to f89da81 Compare July 17, 2026 18:34
@lmcafee-nvidia

Copy link
Copy Markdown
Owner Author

Closing this archived implementation in favor of the active cumulative stack. Its MTP, EP, Mamba admission, primer, prefill, and step-order architecture is consolidated in #26 and #27.

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant