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>

tests/test_tp_parallel_state.py

@@ -0,0 +1,85 @@
+"""ADR-0027 T1: TP parallel_state (D3).
+
+Phase 1: ``kernbench.tp`` module does not exist yet — tests fail at import.
+Phase 2 (D2/D3) lands the package and these pass.
+"""
+from __future__ import annotations
+
+import pytest
+
+
+def _make_ctx(topology):
+    from kernbench.runtime_api.context import RuntimeContext
+    from kernbench.runtime_api.types import DeviceSelector
+    from kernbench.sim_engine.engine import GraphEngine
+
+    engine = GraphEngine(topology.topology_obj, enable_data=True)
+    return RuntimeContext(
+        engine=engine,
+        target_device=DeviceSelector("all"),
+        correlation_id="test_t1",
+        spec=topology.topology_obj.spec,
+    )
+
+
+def test_tp_package_importable():
+    """D2: kernbench.tp must be importable."""
+    import kernbench.tp as tp
+    assert hasattr(tp, "initialize_model_parallel")
+    assert hasattr(tp, "get_tensor_model_parallel_world_size")
+    assert hasattr(tp, "get_tensor_model_parallel_rank")
+
+
+def test_initialize_model_parallel_matches_world_size(topology, tmp_path, monkeypatch):
+    """D3: TP size must equal dist world_size; otherwise NotImplementedError."""
+    import kernbench.tp as tp
+
+    with _make_ctx(topology) as ctx:
+        ctx.distributed.init_process_group(backend="ahbm")
+        ws = ctx.distributed.get_world_size()
+
+        tp.initialize_model_parallel(ws)
+        assert tp.get_tensor_model_parallel_world_size() == ws
+
+
+def test_initialize_mismatched_ws_raises(topology):
+    """D3: calling with tp_size != world_size raises NotImplementedError."""
+    import kernbench.tp as tp
+
+    with _make_ctx(topology) as ctx:
+        ctx.distributed.init_process_group(backend="ahbm")
+        ws = ctx.distributed.get_world_size()
+
+        with pytest.raises(NotImplementedError):
+            tp.initialize_model_parallel(ws + 1)
+
+
+def test_get_tp_rank_is_greenlet_local(topology):
+    """D3: get_tensor_model_parallel_rank returns greenlet-local rank
+    (delegates to torch.distributed.get_rank, ADR-0024 D9)."""
+    import kernbench.tp as tp
+
+    with _make_ctx(topology) as ctx:
+        ctx.distributed.init_process_group(backend="ahbm")
+        ws = ctx.distributed.get_world_size()
+        tp.initialize_model_parallel(ws)
+
+        observed: list[int] = []
+
+        def _worker(rank: int):
+            observed.append(tp.get_tensor_model_parallel_rank())
+
+        ctx.multiprocessing.spawn(_worker, args=(), nprocs=ws)
+
+        assert sorted(observed) == list(range(ws))
+
+
+def test_get_world_size_before_init_raises():
+    """D3: uninitialised TP group → accessing world_size fails informatively."""
+    from kernbench.tp import parallel_state
+
+    # Reset internal state if previous tests (or parallel workers) left it set.
+    parallel_state._reset_for_tests()
+
+    with pytest.raises((RuntimeError, AssertionError, TypeError)):
+        _ = parallel_state.get_tensor_model_parallel_world_size() + 0