Skip to content

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

Description

@morluto

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions