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/runtime_api/multiprocessing.py

@@ -0,0 +1,152 @@
+"""``torch.multiprocessing.spawn``-compatible namespace (ADR-0027 D1).
+
+Real-PyTorch API *signature* parity only — execution model is a cooperative
+greenlet scheduler in a single Python process (D1.0). Non-goals: process
+isolation, independent address space, failure isolation, OS-level scheduler
+fairness, mp.Queue/Lock.
+
+Attached to ``RuntimeContext`` as ``ctx.multiprocessing`` in
+``__post_init__`` (D1.3).
+"""
+from __future__ import annotations
+
+from typing import Any, Callable
+
+
+class SpawnException(RuntimeError):
+    """Raised from ``_MultiprocessingNamespace.spawn`` on worker failure.
+
+    ``errors`` contains only root-cause ranks — the rank(s) whose body
+    raised. Sibling greenlets terminated via ``throw(SystemExit)`` during
+    cleanup are NOT recorded (SystemExit does not satisfy ``except
+    Exception`` in the entry wrapper).
+    """
+
+    def __init__(self, errors: dict[int, Exception]):
+        self.errors = errors
+        first = next(iter(errors.items()), None)
+        msg = (
+            f"spawn failed on ranks {sorted(errors.keys())}"
+            + (
+                f": rank {first[0]} raised {first[1]!r}"
+                if first is not None
+                else ""
+            )
+        )
+        super().__init__(msg)
+
+
+def _drain_pending(ctx: Any) -> None:
+    """Drain worker-wait + collective-pending queues in main context (D0.4/D0.5).
+
+    Loop-until-empty: runs until both queues are simultaneously empty. Safe
+    under the current model where main-context ``ctx.wait`` never re-enqueues
+    (D0.5 main-context non-reentrance invariant); also safe under future
+    extensions where drain can add sub-handles (SimPy causality gives finite
+    depth).
+    """
+    distributed = getattr(ctx, "distributed", None)
+    backend = getattr(distributed, "_backend", None) if distributed else None
+
+    def _collective_nonempty() -> bool:
+        if backend is None:
+            return False
+        pending = getattr(backend, "_pending_collective_handles", None)
+        return bool(pending)
+
+    while ctx._pending_worker_waits or _collective_nonempty():
+        # (a) Worker-driven waits (D0.1). FIFO.
+        while ctx._pending_worker_waits:
+            h = ctx._pending_worker_waits.pop(0)
+            if h not in ctx._completed:
+                wait_fn = getattr(ctx.engine, "wait", None)
+                if wait_fn is not None:
+                    wait_fn(h)
+                # Populate _completed so fast-path in ctx.wait short-circuits
+                # on the return leg.
+                ctx._completed.add(h)
+        # (b) Collective backend queue (ADR-0024 D7 + D0.4-(2)).
+        if backend is not None:
+            pending_list = getattr(backend, "_pending_collective_handles", None)
+            if pending_list is not None:
+                while pending_list:
+                    h, _sip_id, meta = pending_list.pop(0)
+                    # Main context: ctx.wait drives engine directly and does
+                    # NOT re-enqueue (D0.5 invariant).
+                    ctx.wait(h, _meta=meta)
+
+
+class _MultiprocessingNamespace:
+    """torch.multiprocessing-compat facade bound to a RuntimeContext."""
+
+    def __init__(self, ctx: Any) -> None:
+        self._ctx = ctx
+
+    def spawn(
+        self,
+        fn: Callable,
+        args: tuple = (),
+        nprocs: int = 1,
+        join: bool = True,
+    ) -> None:
+        """Spawn ``nprocs`` worker greenlets, each calling ``fn(rank, *args)``.
+
+        Mirrors ``torch.multiprocessing.spawn`` signature (minus ``daemon``).
+        Runs the D0.4 round-robin scheduler loop until all workers finish,
+        draining pending queues between rounds.
+        """
+        from greenlet import greenlet
+
+        ctx = self._ctx
+        dist = ctx.distributed
+        gs: list = []
+        errors: dict[int, Exception] = {}
+
+        for rank in range(nprocs):
+            def _entry(r: int = rank) -> None:
+                try:
+                    fn(r, *args)
+                except Exception as e:
+                    errors[r] = e
+                    raise
+
+            g = greenlet(_entry)
+            if dist is not None and hasattr(dist, "_bind_rank"):
+                dist._bind_rank(g, rank)
+            gs.append(g)
+
+        try:
+            while True:
+                alive = [g for g in gs if not g.dead]
+                if not alive:
+                    break
+                for g in alive:
+                    if not g.dead:
+                        g.switch()
+                _drain_pending(ctx)
+        except Exception as outer:
+            # D0.4-(4) sibling cleanup. Abort live greenlets, clear state.
+            for other in gs:
+                if not other.dead:
+                    try:
+                        other.throw(SystemExit)
+                    except BaseException:
+                        # SystemExit inherits BaseException; greenlet.throw
+                        # re-raises in caller if target doesn't catch it.
+                        # Silent — we're already in cleanup.
+                        pass
+            backend = getattr(dist, "_backend", None)
+            if backend is not None:
+                if hasattr(backend, "_barrier") and hasattr(backend._barrier, "reset"):
+                    try:
+                        backend._barrier.reset()
+                    except Exception:
+                        pass
+                pending_collective = getattr(
+                    backend, "_pending_collective_handles", None,
+                )
+                if pending_collective is not None:
+                    pending_collective.clear()
+            ctx._pending_worker_waits.clear()
+            raise SpawnException(errors) from outer
+        # join=True: we already waited for all workers above.