Skip to content

[Fix] Enforce Buffer ownership for communication plans - #12

Open
morluto wants to merge 1 commit into
MoonshotAI:masterfrom
morluto:codex/plan-provenance
Open

[Fix] Enforce Buffer ownership for communication plans#12
morluto wants to merge 1 commit into
MoonshotAI:masterfrom
morluto:codex/plan-provenance

Conversation

@morluto

@morluto morluto commented Jul 28, 2026

Copy link
Copy Markdown

Fixes #8.

Summary

Tie saved communication plans to the Buffer context that produced them and
reject incompatible plan reuse before synchronization or CUDA launch.

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

Problem

MoonEPCommPlan contains routing and layout decisions interpreted relative to
its originating Buffer:

  • dst uses destination_rank * NvS + local_offset;
  • duplicate offsets address rank-local NvS storage;
  • expert-copy decisions depend on R, E, and B;
  • combine iteration depends on N and K;
  • process-group ordering, device ownership, VMM mappings, and barriers belong
    to one context.

Existing checks validate only parts of this geometry.

A concrete incompatible pair is:

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 NvS check accepts this pair because both values are six. Buffer
B can then use its six-entry combine iteration domain while reading Buffer A's
two-entry dst.

A foreign prefetch plan can similarly contain an expert ID valid under its
original E but outside the consuming Buffer's expert tensor.

Change

Each Buffer context now owns a private identity marker. Planning stores that
marker in the returned plan, and MoonEPCommPlan.clone() preserves it.

Before accepting a saved plan, the Buffer validates:

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

The validation applies to:

  • saved-plan dispatch;
  • prefetch_weight;
  • combine;
  • reduce_grad.

Regression coverage

The regression covers three cases:

  • a cloned plan is accepted by its originating Buffer;
  • a plan from a second Buffer is rejected before launch;
  • a same-owner plan with a modified signature is rejected before launch.

The final kernel call is replaced with a recorder because the behavior under
test is the synchronous pre-launch ownership boundary.

Compatibility

The provenance field is private, appended with a default, and excluded from
dataclass repr and equality.

Plans returned by a Buffer and their clones retain their current behavior.
Plans manually constructed, deserialized into another context, or transferred
between Buffer instances are now rejected.

There is no kernel ABI, tensor-layout, or planning-algorithm change.

Validation

  • python3 -m py_compile moonep/planning.py moonep/api.py tests/test_dispatch.py
    — passed.
  • CPU-backed validator harness — same-context clone accepted; foreign and
    signature-mutated plans rejected.
  • Added focused regression coverage for the public dispatch boundary.

@morluto morluto changed the title Reject communication plans from incompatible buffers [Fix] Enforce Buffer ownership for communication plans Jul 28, 2026
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.

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

1 participant