[Fix] Enforce Buffer ownership for communication plans - #12
Open
morluto wants to merge 1 commit into
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #8.
Summary
Tie saved communication plans to the Buffer context that produced them and
reject incompatible plan reuse before synchronization or CUDA launch.
Problem
MoonEPCommPlancontains routing and layout decisions interpreted relative toits originating Buffer:
dstusesdestination_rank * NvS + local_offset;NvSstorage;R,E, andB;NandK;to one context.
Existing checks validate only parts of this geometry.
A concrete incompatible pair is:
The existing
NvScheck accepts this pair because both values are six. BufferB 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
Ebut 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:
The validation applies to:
prefetch_weight;combine;reduce_grad.Regression coverage
The regression covers three cases:
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.
signature-mutated plans rejected.