From b3412dce0c36b147f805ce9a556652ae680fbba3 Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Tue, 28 Jul 2026 10:50:55 +0000 Subject: [PATCH] Reject communication plans from incompatible buffers --- moonep/api.py | 43 +++++++++++++++++++++++++++++++++---- moonep/planning.py | 5 ++++- tests/test_dispatch.py | 48 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 91 insertions(+), 5 deletions(-) diff --git a/moonep/api.py b/moonep/api.py index 8d0cbc7..370a132 100644 --- a/moonep/api.py +++ b/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( diff --git a/moonep/planning.py b/moonep/planning.py index 23e67e6..d230a1f 100644 --- a/moonep/planning.py +++ b/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 diff --git a/tests/test_dispatch.py b/tests/test_dispatch.py index b9e3397..f502f7c 100644 --- a/tests/test_dispatch.py +++ b/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