ADR-0027: Megatron TP API + worker-wait generalization + mp.spawn

Implements ADR-0027 Phase 2 end-to-end. All 559 tests pass (was 523 +
1 xfail; ring_default_ws strict-xfail is now resolved).

D0 — Worker-wait generalization (context.py):
- _pending_worker_waits queue on RuntimeContext.
- ctx.wait(h) in worker context defers to main via g.parent.switch().
  Fast-path for already-completed handles.
- Worker API is unchanged: tensor deploy, launch, etc. still look
  synchronous; they're transparently cooperatively scheduled.
- Solves ADR-0024 Phase B kernel-greenlet orphan bug (env.run now
  only ever drives from main; kernel _parent is always main).

D0.5 — Host-read barrier (tensor.py):
- Explicit _HOST_READ_BARRIERS registry (T5.g closed-set via code
  review, not reflection-magic).
- numpy/data/__getitem__/__repr__ drain pending worker-waits before
  host-observable read.
- copy_: source-side barrier via source.numpy(). Target-side write
  barrier is intentionally NOT applied — global pending target barrier
  prematurely drains cross-rank collectives → deadlock.
- Collective pending is excluded from barrier drain condition
  (collective is cross-rank; its own yield in all_reduce covers the
  invariant naturally).

D1 — torch.multiprocessing.spawn (runtime_api/multiprocessing.py):
- API signature parity with real PyTorch spawn; execution is
  cooperative greenlet scheduler (process isolation etc. are explicit
  non-goals per D1.0).
- _drain_pending drains worker-waits then collectives in one barrier,
  loop-until-empty.
- Round-based exception handling with SystemExit sibling abort +
  SpawnException(errors) wrapping root-cause ranks.
- RuntimeContext attaches ctx.multiprocessing in __post_init__.
- benches/ccl_allreduce.py hand-rolled loop collapses to one
  torch.multiprocessing.spawn call.

D2–D6 — kernbench.tp package:
- parallel_state: initialize_model_parallel, get_*_rank,
  get_*_world_size, with weak active-ctx registry in context.py.
- layers: ColumnParallelLinear, RowParallelLinear (shape-only
  primitives — fp16 gemm via tl.load + tl.dot + tl.store).
- kernels: _gemm_kernel used by TP layers (self-contained; no bench
  dependency).
- primitives / mappings stubs per D6/D8.

Data-path fixes (surfaced by TP gemm + all_reduce sequence):
- sim_engine/op_log.py: dma_write snapshot is skipped for TCM
  sources (PE scratch is repopulated by Phase 2 math/gemm replay —
  capturing Phase-1-time snapshot picked up STALE data from prior
  kernel's output aliased at the same scratch addr, causing the later
  kernel's dma_write to overwrite Phase 2 result with stale value).
- sim_engine/op_log.py + sim_engine/data_executor.py: per-operand
  space recorded on GemmCmd and composite gemm records so HBM-resident
  operands (tl.load output) don't default to TCM during replay.
- runtime_api/context.py: ctx.zeros writes zero-init to MemoryStore
  at VA keys so kernels reading via VA see deterministic init even
  without explicit copy_().

Tests (Phase 1 + Phase 2):
- test_worker_wait_drain (T3): orphan invariant + resume + multi-rank
  drain + idempotency + exception propagation.
- test_mp_spawn (T4): spawn shape + bind + SpawnException scope.
- test_host_read_barrier (T5): barrier contract per entry-point +
  closed-set registry check.
- test_tp_parallel_state (T1): initialize + rank lookup.
- test_tp_layers (T2): shape + deterministic numerical correctness
  (concat-matmul equality for RowParallel, not mean-only).
- test_tp_mlp (T6): full 2-layer MLP with deterministic weight
  numerical match + rank-consistency post all-reduce.
- test_ccl_allreduce_matrix: ring_default_ws xfail removed (T7).

Regression: 523 pre + 35 new + 1 ex-xfail = 559 passed, 1 intentional
skip (T3.e historical failure documentation).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

src/kernbench/tp/__init__.py

@@ -0,0 +1,21 @@
+"""kernbench.tp — Megatron-style Tensor Parallelism (ADR-0027).
+
+Public API re-exports.
+"""
+from kernbench.tp.layers import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+)
+from kernbench.tp.parallel_state import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    initialize_model_parallel,
+)
+
+__all__ = [
+    "ColumnParallelLinear",
+    "RowParallelLinear",
+    "get_tensor_model_parallel_rank",
+    "get_tensor_model_parallel_world_size",
+    "initialize_model_parallel",
+]

src/kernbench/tp/kernels.py

@@ -0,0 +1,23 @@
+"""Kernel used by ``kernbench.tp`` layers (ADR-0027 D4/D5).
+
+Intentionally self-contained inside the ``tp`` package — the ``tp`` package
+must not import from ``benches/``. Future work: move to a shared
+``kernbench.kernels`` module so benches and TP can share.
+"""
+from __future__ import annotations
+
+
+def _gemm_kernel(a_ptr, b_ptr, out_ptr, M, K, N, tl, DTYPE: str = "f16") -> None:
+    """Single-PE GEMM: out = a @ b via load → dot → store.
+
+    Uses the ``tl.load + tl.dot + tl.store`` path. Unlike ``tl.composite``
+    (which is absorbed by the PE scheduler into TileTokens that don't reach
+    the op_log), this path emits explicit ``DmaReadCmd`` / ``GemmCmd`` /
+    ``DmaWriteCmd`` records, which DataExecutor replays numerically in
+    Phase 2.
+    """
+    M, K, N = int(M), int(K), int(N)
+    a = tl.load(int(a_ptr), shape=(M, K), dtype=DTYPE)
+    b = tl.load(int(b_ptr), shape=(K, N), dtype=DTYPE)
+    out = tl.dot(a, b)
+    tl.store(int(out_ptr), out)

src/kernbench/tp/layers.py

@@ -0,0 +1,150 @@
+"""Megatron-style parallel layers (ADR-0027 D4/D5).
+
+- ``ColumnParallelLinear``: weight's out_features axis split across TP ranks.
+  forward(x) is local gemm; no collective.
+- ``RowParallelLinear``: weight's in_features axis split across TP ranks.
+  forward(x) ends with ``dist.all_reduce`` to sum partial products.
+
+Both layers use the intra-device ``DPPolicy`` (ADR-0026). TP shard
+ownership is determined by ``torch.ahbm.set_device(rank)`` (ADR-0024 D10).
+
+Yield-safety contract (ADR-0027 D4/D5): every forward path contains at
+least one ``ctx.wait`` (via ``torch.launch``) or one collective; this
+keeps the scheduler loop making progress.
+"""
+from __future__ import annotations
+
+from typing import Any
+
+from kernbench.policy.placement.dp import DPPolicy
+from kernbench.tp.kernels import _gemm_kernel
+from kernbench.tp.parallel_state import (
+    get_tensor_model_parallel_world_size,
+)
+
+
+class ColumnParallelLinear:
+    """Weight's K (out_features) axis distributed across TP ranks.
+
+    forward(x):
+        x: (M, N) — full-replicated across ranks
+        W_k: (N, K / world_size) — this rank's slice (on its SIP)
+        y_k = x @ W_k → (M, K / world_size)
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = False,
+        dtype: str = "f16",
+        torch: Any = None,
+    ) -> None:
+        if torch is None:
+            raise TypeError("ColumnParallelLinear requires torch=<RuntimeContext>")
+        ws = get_tensor_model_parallel_world_size()
+        if out_features % ws != 0:
+            raise ValueError(
+                f"out_features ({out_features}) must be divisible by TP world "
+                f"size ({ws})"
+            )
+        self.in_features = in_features
+        self.out_features = out_features
+        self.k_local = out_features // ws
+        self.dtype = dtype
+        self._torch = torch
+        # Per-rank weight slice. ``set_device(rank)`` (ADR-0024 D10) places
+        # it on SIP ``rank``. Intra-SIP layout comes from DPPolicy (ADR-0026).
+        self.weight = torch.zeros(
+            (in_features, self.k_local),
+            dtype=dtype,
+            dp=DPPolicy(cube="replicate", pe="replicate",
+                        num_cubes=1, num_pes=1),
+            name="col_parallel_w",
+        )
+        # Bias omitted in initial scope (ADR-0027 D9).
+        self.bias = None
+        if bias:
+            raise NotImplementedError(
+                "bias=True is deferred (ADR-0027 D9 initial scope)"
+            )
+
+    def forward(self, x):
+        M = int(x.shape[0])
+        out = self._torch.empty(
+            (M, self.k_local),
+            dtype=x.dtype,
+            dp=DPPolicy(cube="replicate", pe="replicate",
+                        num_cubes=1, num_pes=1),
+            name="col_parallel_out",
+        )
+        self._torch.launch(
+            "col_parallel_gemm",
+            _gemm_kernel,
+            x, self.weight, out,
+            M, self.in_features, self.k_local,
+        )
+        return out
+
+
+class RowParallelLinear:
+    """Weight's N (in_features) axis distributed across TP ranks.
+
+    forward(x):
+        x: (M, N / world_size) — rank-local slice (ColumnParallel output)
+        W_k: (N / world_size, K) — this rank's slice
+        y_k = x @ W_k → (M, K) — partial sum
+        y = all_reduce(y_k, op="sum") → (M, K) on every rank
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = False,
+        dtype: str = "f16",
+        torch: Any = None,
+    ) -> None:
+        if torch is None:
+            raise TypeError("RowParallelLinear requires torch=<RuntimeContext>")
+        ws = get_tensor_model_parallel_world_size()
+        if in_features % ws != 0:
+            raise ValueError(
+                f"in_features ({in_features}) must be divisible by TP world "
+                f"size ({ws})"
+            )
+        self.in_features = in_features
+        self.out_features = out_features
+        self.n_local = in_features // ws
+        self.dtype = dtype
+        self._torch = torch
+        self.weight = torch.zeros(
+            (self.n_local, out_features),
+            dtype=dtype,
+            dp=DPPolicy(cube="replicate", pe="replicate",
+                        num_cubes=1, num_pes=1),
+            name="row_parallel_w",
+        )
+        self.bias = None
+        if bias:
+            raise NotImplementedError(
+                "bias=True is deferred (ADR-0027 D9 initial scope)"
+            )
+
+    def forward(self, x):
+        M = int(x.shape[0])
+        y_partial = self._torch.empty(
+            (M, self.out_features),
+            dtype=x.dtype,
+            dp=DPPolicy(cube="replicate", pe="replicate",
+                        num_cubes=1, num_pes=1),
+            name="row_parallel_partial",
+        )
+        self._torch.launch(
+            "row_parallel_gemm",
+            _gemm_kernel,
+            x, self.weight, y_partial,
+            M, self.n_local, self.out_features,
+        )
+        self._torch.distributed.all_reduce(y_partial, op="sum")
+        return y_partial

src/kernbench/tp/mappings.py

@@ -0,0 +1,5 @@
+"""Forward/backward mappings stub (ADR-0027 — future backward work).
+
+Inference-only initial scope. Backward hooks land when training simulation
+arrives.
+"""

src/kernbench/tp/parallel_state.py

@@ -0,0 +1,83 @@
+"""TP group state (ADR-0027 D3).
+
+Single global TP group. Initial scope: TP size == world_size (pure TP;
+mixed DP+TP is future work).
+"""
+from __future__ import annotations
+
+_TP_WORLD_SIZE: int | None = None
+
+
+def initialize_model_parallel(tensor_model_parallel_size: int) -> None:
+    """Initialize the TP process group.
+
+    Must be called after ``torch.distributed.init_process_group``.
+    Only ``tensor_model_parallel_size == world_size`` is supported in the
+    initial scope.
+    """
+    global _TP_WORLD_SIZE
+    # Import here to avoid cycle when tp is imported before a ctx exists.
+    _ws = _current_world_size()
+    if tensor_model_parallel_size != _ws:
+        raise NotImplementedError(
+            f"Only TP == world_size supported; got TP={tensor_model_parallel_size}, "
+            f"world_size={_ws}"
+        )
+    _TP_WORLD_SIZE = tensor_model_parallel_size
+
+
+def get_tensor_model_parallel_world_size() -> int:
+    """Return the TP group's world size.
+
+    Raises if not initialised — callers must call
+    :func:`initialize_model_parallel` first.
+    """
+    if _TP_WORLD_SIZE is None:
+        raise RuntimeError(
+            "TP group not initialised; call initialize_model_parallel() first"
+        )
+    return _TP_WORLD_SIZE
+
+
+def get_tensor_model_parallel_rank() -> int:
+    """Return this worker's rank within the TP group.
+
+    Delegates to the greenlet-local rank registered by the spawn launcher
+    (ADR-0024 D9 via ``torch.distributed.get_rank``).
+    """
+    # Resolve via the global torch.distributed facade on the active ctx.
+    return _current_rank()
+
+
+def _reset_for_tests() -> None:
+    """Clear _TP_WORLD_SIZE so ordering-sensitive tests can re-init."""
+    global _TP_WORLD_SIZE
+    _TP_WORLD_SIZE = None
+
+
+# ── helpers (resolve current ctx) ────────────────────────────────────
+
+
+def _current_ctx():
+    """Best-effort resolution of the currently-active RuntimeContext.
+
+    In KernBench, the ``ctx`` is passed as the ``torch`` positional in
+    bench/worker code. Since parallel_state is a module-global helper,
+    we look it up via a weak registry maintained by RuntimeContext.
+    """
+    from kernbench.runtime_api.context import _get_active_context
+    ctx = _get_active_context()
+    if ctx is None:
+        raise RuntimeError(
+            "No active RuntimeContext; kernbench.tp requires one "
+            "(call init_process_group / spawn under a live ctx)"
+        )
+    return ctx
+
+
+def _current_world_size() -> int:
+    return _current_ctx().distributed.get_world_size()
+
+
+def _current_rank() -> int:
+    return _current_ctx().distributed.get_rank()

src/kernbench/tp/primitives.py

@@ -0,0 +1,34 @@
+"""TP primitive ops (ADR-0027 D6).
+
+``copy_to_tp_region`` / ``reduce_from_tp_region`` are forward-only in the
+initial scope (backward pass is future work). ``scatter`` / ``gather`` are
+not implemented — they require an all-gather kernel that is not yet
+available in KernBench (see ADR-0027 D9).
+"""
+from __future__ import annotations
+
+from typing import Any
+
+
+def copy_to_tp_region(x: Any) -> Any:
+    """Forward: identity. Backward: all-reduce. (Training is future.)"""
+    return x
+
+
+def reduce_from_tp_region(x: Any, torch: Any) -> Any:
+    """Forward: all-reduce. Backward: identity."""
+    torch.distributed.all_reduce(x, op="sum")
+    return x
+
+
+def scatter_to_tp_region(x: Any) -> Any:
+    raise NotImplementedError(
+        "scatter_to_tp_region deferred — caller should create the sharded "
+        "tensor directly (ADR-0027 D9)"
+    )
+
+
+def gather_from_tp_region(x: Any) -> Any:
+    raise NotImplementedError(
+        "gather_from_tp_region deferred — requires all-gather kernel (ADR-0027 D9)"
+    )