[Fix] Enforce Buffer ownership for communication plans

moonep/api.py

@@ -394,6 +394,7 @@ def _create_context(
         'num_sms_dedup': num_sms_dedup,
         'token_padding': token_padding,
         'device': device,
+        '_plan_context_marker': object(),
         'token_padding_extra': token_padding_extra,
         # NVLink shared
         'hidden_buf': hidden_buf,
@@ -585,6 +586,40 @@ def _plan_runtime_tensors(plan: MoonEPCommPlan):
             plan.dup_counts,
         )
 
+    @staticmethod
+    def _validate_plan(
+        ctx: dict,
+        plan: MoonEPCommPlan,
+        operation: str,
+    ) -> None:
+        if not isinstance(plan, MoonEPCommPlan):
+            raise TypeError(f"{operation}: plan must be a MoonEPCommPlan")
+
+        if plan._context_marker is not ctx['_plan_context_marker']:
+            raise ValueError(
+                f"{operation}: plan was created by a different Buffer; "
+                "reuse it only with the Buffer that produced it"
+            )
+
+        expected = {
+            'N': int(ctx['N']),
+            'R': int(ctx['R']),
+            'E': int(ctx['E']),
+            'B': int(ctx['B']),
+            'NvS': int(ctx['NvS']),
+            'K': int(ctx['K']),
+        }
+        mismatches = [
+            f"{name}: expected {value}, got {int(getattr(plan, name))}"
+            for name, value in expected.items()
+            if int(getattr(plan, name)) != value
+        ]
+        if mismatches:
+            raise ValueError(
+                f"{operation}: plan configuration does not match its Buffer "
+                f"({'; '.join(mismatches)})"
+            )
+
     def _run_dispatch_on_current_stream(
         self,
         ctx: dict,
@@ -750,7 +785,7 @@ def dispatch(
         else:
             cu_seqlens = None
             planning_args = None
-            assert isinstance(plan, MoonEPCommPlan)
+            self._validate_plan(ctx, plan, "Buffer.dispatch")
 
         if zero_copy:
             hidden_nvsh = ctx['hidden_buf_local']
@@ -844,7 +879,7 @@ def prefetch_weight(
         """
         ctx = self._require_ctx()
 
-        assert isinstance(plan, MoonEPCommPlan), "Buffer.prefetch_weight: plan is required"
+        self._validate_plan(ctx, plan, "Buffer.prefetch_weight")
         weight_prefetch_args = (full_gate_weight, full_up_weight, full_down_weight)
         assert all(w is not None for w in weight_prefetch_args), \
             "prefetch_weight tensors must be provided together"
@@ -923,7 +958,7 @@ def combine(
         """
         ctx = self._require_ctx()
 
-        assert isinstance(plan, MoonEPCommPlan), "Buffer.combine: plan is required"
+        self._validate_plan(ctx, plan, "Buffer.combine")
 
         assert hidden_nvsh is not None
         assert hidden_nvsh.dtype == torch.bfloat16
@@ -1054,7 +1089,7 @@ def reduce_grad(
         )
         assert all(t is not None for t in grad_reduce_args), \
             "reduce_grad tensors must be provided together"
-        assert isinstance(plan, MoonEPCommPlan), "Buffer.reduce_grad: plan is required"
+        self._validate_plan(ctx, plan, "Buffer.reduce_grad")
 
         if not async_finish:
             self._run_reduce_grad_on_current_stream(

moonep/planning.py

@@ -11,7 +11,7 @@
 """
 
 import functools
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 import torch
 import cuda.bindings.driver as cuda
@@ -48,6 +48,7 @@ class MoonEPCommPlan:
     dup_groups: torch.Tensor
     dup_loffs: torch.Tensor
     dup_counts: torch.Tensor
+    _context_marker: object | None = field(default=None, repr=False, compare=False)
 
     def __post_init__(self) -> None:
         N = int(self.N)
@@ -85,6 +86,7 @@ def clone(self) -> "MoonEPCommPlan":
             B=self.B,
             NvS=self.NvS,
             K=self.K,
+            _context_marker=self._context_marker,
         )
 
 
@@ -1246,6 +1248,7 @@ def allocate_planning_outputs(ctx: dict):
         B=B,
         NvS=NvS,
         K=K,
+        _context_marker=ctx.get('_plan_context_marker'),
     )
     return plan, cu_seqlens
 

tests/test_dispatch.py

@@ -425,6 +425,54 @@ def test_dispatch_saved_plan_hidden_only_reuses_dst_and_skips_weights(dist_env):
     )
 
 
+def test_buffer_plan_provenance_rejects_foreign_and_incompatible_plans(
+    dist_env, monkeypatch
+):
+    from moonep.planning import allocate_planning_outputs
+
+    rank, R = dist_env
+    case = KernelCase(
+        "plan_provenance",
+        S=1,
+        K=1,
+        epn=1,
+        H=8,
+        num_sms=1,
+        B=1,
+        token_padding=1,
+        min_R=2,
+    )
+    ctx_a = init_case(case, R)
+    ctx_b = init_case(case, R)
+    buffer_a = ctx_a["_buffer"]
+    plan_a, _ = allocate_planning_outputs(ctx_a)
+    plan_b, _ = allocate_planning_outputs(ctx_b)
+    hidden = torch.empty(
+        case.S, case.H, dtype=torch.bfloat16, device=f"cuda:{rank}"
+    )
+
+    launches = []
+
+    def record_launch(*args, **kwargs):
+        launches.append((args, kwargs))
+
+    monkeypatch.setattr(buffer_a, "_run_dispatch_on_current_stream", record_launch)
+
+    _, _, cu_seqlens, accepted = buffer_a.dispatch(hidden, plan=plan_a.clone())
+    assert accepted._context_marker is plan_a._context_marker
+    assert cu_seqlens is None
+    assert len(launches) == 1
+
+    with pytest.raises(ValueError, match="created by a different Buffer"):
+        buffer_a.dispatch(hidden, plan=plan_b)
+    assert len(launches) == 1
+
+    incompatible = replace(plan_a, K=plan_a.K + 1)
+    with pytest.raises(ValueError, match=r"K: expected 1, got 2"):
+        buffer_a.dispatch(hidden, plan=incompatible)
+    assert len(launches) == 1
+
+
 @pytest.mark.parametrize("case", case_params(LARGE_DISPATCH_CASES))
 def test_dispatch_large_hidden_stride_spotcheck(dist_env, case):
     rank, R = dist_env