MoonshotAI / MoonshotAI/MoonEP

[Bug] Communication plans can be reused across incompatible Buffer contexts

Open
#8 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
1.1k
Forks
134
PR merge metrics
No merged PRs in 30d

Description

Problem

MoonEPCommPlan stores routing and layout decisions produced for one Buffer
context, but public plan-consuming methods do not verify that a supplied plan
belongs to that Buffer.

The plan contains:

  • dst values encoded as destination_rank * NvS + local_offset;
  • duplicate offsets into rank-local NvS storage;
  • expert-copy decisions indexed by rank and prefetch slot;
  • zero-fill ranges for the planned expert/prefetch layout.

Buffer.dispatch, combine, prefetch_weight, and reduce_grad can receive a
plan from another Buffer. Existing checks cover only parts of the required
geometry.

Concrete combine trigger

Buffer A: R=2, E=4, B=1, S=2, K=1, token_padding=2
          N=2, NvS=6, len(plan.dst)=2

Buffer B: R=2, E=4, B=1, S=6, K=1, token_padding=1
          N=6, NvS=6

The existing combine prologue can accept this pair because NvS agrees. The
kernel then uses Buffer B's N=6 iteration domain while reading Buffer A's
two-entry dst. Equal NvS therefore does not prove that the plan's iteration
domain is compatible.

A foreign prefetch plan has a second trigger: a plan produced with E=8 may
contain expert ID 7. Reusing it with an E=4 Buffer can pass the existing
experts_to_copy shape checks when B agrees, but the expert ID is outside the
consumer's four-expert tensor.

These are source-derived trigger paths.

Required invariant

Buffer A --planning--> Plan(marker A, signature A)
                              |
                 +------------+------------+
                 |                         |
             Buffer A                  Buffer B
          marker matches            marker differs
                 |                         |
              accepted             reject before CUDA

Before synchronization or kernel launch:

plan origin == consuming Buffer
plan {N,R,E,B,NvS,K} == Buffer {N,R,E,B,NvS,K}

Context ownership also establishes process-group rank ordering, device
ownership, VMM layout, and barrier-resource identity. A dimensional signature
alone cannot establish those facts.

Expected behavior

  • Plans remain reusable by the Buffer that created them.
  • plan.clone() preserves ownership and remains reusable.
  • A plan created by another Buffer raises a clear exception before CUDA work.
  • A same-owner plan with an incompatible stored signature is rejected.
  • Kernel ABI, tensor layout, and planning semantics remain unchanged.

Compatibility

MoonEPCommPlan is public. A provenance field should therefore be private,
appended with a construction default, and excluded from repr/equality.
Manually constructed or deserialized plans without live Buffer ownership
should be rejected explicitly rather than treated as portable.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with MoonEPCommPlan construction and clone behavior, then trace the public Buffer.dispatch, combine, prefetch_weight, and reduce_grad entry points. Verify that ownership and the stored {N,R,E,B,NvS,K} signature are checked before synchronization or kernel launch, while same-owner plans and clones remain usable and foreign or invalid plans fail clearly.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.