Skip to content

[None][fix] Keep one-sided MoE combine workspace offset stable - #1

Closed
roborluo wants to merge 1 commit into
mainfrom
fix/moe-a2a-stable-combine-offset
Closed

roborluo wants to merge 1 commit into
mainfrom
fix/moe-a2a-stable-combine-offset

Conversation

@roborluo

Copy link
Copy Markdown
Owner

Description

The one-sided MoE combine offset previously depended on the current dispatch
payload size. Changes in payload precision or token count could let a later
dispatch overwrite the previous combine payload while a peer was still reading it.

This change:

  • Places combine storage at a fixed, aligned workspace offset.
  • Checks dispatch bounds before native writes.
  • Uses the same offset for direct-to-workspace outputs.

The existing workspace allocation policy is unchanged.

Test Coverage

  • CPU layout tests cover changing dispatch layouts and token counts, overlap
    detection, insufficient capacity, incompatible combine layouts, and lazy
    reservation for both CFT and non-CFT workspaces.
  • Constructor tests verify that failed reservation does not publish shared
    workspace state.
  • A four-rank GPU regression covers mixed dispatch layouts, uneven token counts,
    empty local ranks, eager execution/CUDA graphs, CFT/non-CFT paths,
    low-precision combine, and direct-to-workspace outputs.

Validation on this branch:

  • 22 CPU-isolated tests passed; native CUDA dependencies were isolated/mocked.
  • Ruff lint/format, Python compile, and diff checks passed.
  • The four-rank GPU regression has not been rerun on this base.

Keep the one-sided combine payload at a fixed workspace offset when dispatch precision or token count changes. Validate dispatch bounds before native writes and use the same offset for direct-to-workspace outputs.

Preserve the existing workspace allocation policy and add layout checks and a four-rank mixed-layout regression.

Signed-off-by: bofengluo <bofengl@nvidia.com>
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