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/context.py

@@ -42,6 +42,21 @@ def _numpy_to_dtype_str(np_dtype) -> str:
     raise ValueError(f"unsupported numpy dtype: {np_dtype!r}")
 
 
+# ADR-0027 D3: weak registry of the currently-active RuntimeContext so
+# module-level helpers (e.g. ``kernbench.tp.parallel_state``) can resolve
+# the ctx without threading it through every call.
+import weakref as _weakref
+
+_ACTIVE_CTX_REF: _weakref.ref | None = None
+
+
+def _get_active_context():
+    """Return the most-recently-entered RuntimeContext, or None."""
+    if _ACTIVE_CTX_REF is None:
+        return None
+    return _ACTIVE_CTX_REF()
+
+
 class _AhbmNamespace:
     """torch.ahbm — per-greenlet SIP device binding (ADR-0024 D10).
 
@@ -89,6 +104,10 @@ class RuntimeContext:
 
     _handles: list[RequestHandle] = field(default_factory=list, init=False)
     _completed: set[RequestHandle] = field(default_factory=set, init=False)
+    # ADR-0027 D0.1: worker-deferred wait queue. When a worker greenlet
+    # calls ctx.wait(h), the handle is appended here and control yields to
+    # main. Main's scheduler drain consumes this list.
+    _pending_worker_waits: list[RequestHandle] = field(default_factory=list, init=False)
     _allocators: dict[tuple[int, int, int], Any] = field(default_factory=dict, init=False)
     _va_allocator: Any = field(default=None, init=False)
     _tensor_counter: int = field(default=0, init=False)
@@ -109,6 +128,9 @@ class RuntimeContext:
         # (PyTorch 2.x portable) namespaces for per-greenlet device binding.
         self.ahbm = _AhbmNamespace()
         self.accelerator = _AcceleratorNamespace(self.ahbm)
+        # ADR-0027 D1.3: torch.multiprocessing.spawn namespace.
+        from kernbench.runtime_api.multiprocessing import _MultiprocessingNamespace
+        self.multiprocessing = _MultiprocessingNamespace(self)
 
     def install_ipcq(
         self,
@@ -160,10 +182,16 @@ class RuntimeContext:
         return plan
 
     def __enter__(self):
+        global _ACTIVE_CTX_REF
+        _ACTIVE_CTX_REF = _weakref.ref(self)
         return self
 
     def __exit__(self, *exc):
+        global _ACTIVE_CTX_REF
         self.cleanup()
+        # Clear active-context registry if we are it.
+        if _ACTIVE_CTX_REF is not None and _ACTIVE_CTX_REF() is self:
+            _ACTIVE_CTX_REF = None
         return False
 
     def submit(self, request: Any) -> RequestHandle:
@@ -178,10 +206,24 @@ class RuntimeContext:
         return handle in self._completed
 
     def wait(self, handle: RequestHandle, *, _meta: dict | None = None) -> Completion:
+        # ADR-0027 D0.2: fast-path for already-completed handles (avoid
+        # redundant worker→main→worker round-trip).
         if handle in self._completed:
             completion, trace = self.engine.get_completion(handle)
             return completion
 
+        # ADR-0027 D0.2: if called from a worker greenlet (parent is main,
+        # not dead), defer the wait to the main scheduler — enqueue and
+        # yield. Main drains env.run, then switches back. On resume the
+        # handle must be in _completed (D0.3 resume invariant).
+        from greenlet import getcurrent
+        g = getcurrent()
+        if g.parent is not None and not g.parent.dead:
+            self._pending_worker_waits.append(handle)
+            g.parent.switch()
+            # Resume: main drained. Fall through to completion/trace assembly.
+
+        # Main context (or single-driver): drive engine directly.
         wait_fn = getattr(self.engine, "wait", None)
         if wait_fn is not None:
             wait_fn(handle)  # type: ignore[misc]
@@ -543,6 +585,21 @@ class RuntimeContext:
                     "sip": shard.sip, "cube": shard.cube, "pe": shard.pe,
                     "nbytes": shard.nbytes,
                 })
+            # ADR-0027: also populate MemoryStore at VA keys so kernels
+            # reading via VA (the common ``tl.load`` path) see the init
+            # data. Phase 1 MemoryWriteMsg writes via PA; kernels read via
+            # VA; Phase 2 DataExecutor reads via the addresses captured in
+            # op_log (VA for tl.load). Without this, zero-init tensors are
+            # invisible to kernels in Phase 2.
+            store = getattr(self.engine, "_memory_store", None)
+            if store is not None and pattern == "zero" and handle.va_base:
+                import numpy as np
+                from kernbench.runtime_api.tensor import _numpy_dtype
+                np_dtype = _numpy_dtype(dtype)
+                for shard in handle.shards:
+                    count = shard.nbytes // itemsize
+                    addr = handle.va_base + shard.offset_bytes
+                    store.write("hbm", addr, np.zeros(count, dtype=np_dtype))
 
         return t