diff --git a/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md b/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md new file mode 100644 index 00000000..eda802fd --- /dev/null +++ b/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md @@ -0,0 +1,231 @@ +# P1 mHC + RMSNorm Start Kit (`P1-S0`) + +The start kit unblocks every P1 sub-issue (`P1-D1`…`P1-D6`): it freezes the +layer contract, provides a bit-exact FP32 oracle for the six WS1 operators, +generates seeded golden fixtures, and ships one acceptance command that any +backend PR can run independently. + +Issue: [DSV4][P1/7] mHC 与 RMSNorm 确定性前向/反向 (#2). + +## The eight sub-issues of #2 + +`P1-N` is a development-order label; **GitHub issue numbers stay authoritative +for links**. Five WS1 operator tasks, each delivering forward *and* backward +together, plus three WS2 parallelism tasks. + +| Label | Issue | Stage | Scope | +| --- | --- | --- | --- | +| `P1-S0` | this PR | — | Start kit: contract, oracle, fixtures, provider stub, acceptance command | +| `P1-1` | #14 | WS1 | `hc_split_sinkhorn` — controller mapping + 20 Sinkhorn rounds (fwd + bwd) | +| `P1-2` | #15 | WS1 | `fp32_gemm_rms` — FP32 controller projection + controller RMS (fwd + bwd). **Also carries the fixed-K / batch-invariant GEMM reference + bit-equivalence harness** that P2/P3/P5/P7 consume | +| `P1-3` | #16 | WS1 | `mhc_post` — sublayer output written back into the four residual streams (fwd + bwd) | +| `P1-4` | #17 | WS1 | `mhc_pre` / `h_aggregate` — mHC entry and four-stream aggregation (fwd + composite bwd) | +| `P1-5` | #18 | WS1 | `rmsnorm_residual` — single-stream RMSNorm + residual fork (fwd + bwd) | +| `P1-6` | #19 | WS2 | `fp32_gemm_rms` TP/SP contract — the full `K=16384` controller dot and RMS statistic may not silently become local-K | +| `P1-7` | #20 | WS2 | Rank-local invariance of the three mHC operators under TP/SP/CP/DP/PP | +| `P1-8` | #21 | WS2 | `rmsnorm_residual` TP+SP semantics and cross-rank `dgamma` reduction | + +``` +Foundation v1 -> P1-S0 -> {P1-1 || P1-2 || P1-3 || P1-4 || P1-5} -> P1-R0 + | + +-> WS2: P1-6, P1-7, P1-8 +``` + +Every WS1 task depends only on `P1-S0`, never on the others. `P1-4`'s composite +backward calls `hc_split_sinkhorn_bwd` and `fp32_gemm_rms_bwd`, but it codes +against the oracle for those, so it does not wait on `P1-1` or `P1-2`. + +For the three WS2 tasks, the `placement` field on `LayerContract` and +`check_capability` are the attachment points: a provider that has not declared +a placement already fails closed on it. + +## What is in the kit + +| Module | Contents | +| --- | --- | +| `rl_engine/mhc/reduction.py` | The two pinned reduction trees. Defines the golden bytes for every P1 accumulation. | +| `rl_engine/mhc/contract.py` | `LayerContract`, `ResidualBatch`, `ControllerParams`, `NormParams`, `GradBoundary`, fingerprints. | +| `rl_engine/mhc/oracle.py` | FP32 reference for the six operators plus the full block forward/backward composition. | +| `rl_engine/mhc/provider.py` | `MHCProvider` protocol, `ReferenceProvider` (oracle-backed), `StubProvider` (fail-closed), `check_capability`. | +| `rl_engine/mhc/fixtures.py` | Seeded fixture cases and the golden-hash manifest (`tests/fixtures/p1/golden_hashes.json`, the CI anchor). | +| `rl_engine/mhc/trace.py` | Boundary hashes + `first_divergence` (P1-local stand-in for `TraceEnvelope`). | +| `scripts/check_p1.py` | The acceptance command. | + +## The reduction trees (the heart of the contract) + +Everything in `oracle.py` reduces through `reduction.py`; no `torch.sum`, +`matmul`, `mean` or `einsum` appears anywhere in the operator bodies. Two +trees, and only two: + +- **Long reductions** — the `K = 4·D` controller dot, the `D`-wide + sum-of-squares, the token-major parameter gradients — use a **single FP32 + accumulator walking ascending indices left to right**, with every multiply + and add rounding separately. This is the order the repository's existing + `reduce_rows_fp32` left fold already uses, so a P1 kernel and the WS1 VJP + path agree by construction. +- **4-element stream reductions** — the four mHC residual streams, and the + row/column sums of the 4×4 Sinkhorn matrix — use the balanced tree + **`(a0+a1)+(a2+a3)`** pinned by #2. + +Banned downstream: Split-K, Stream-K, atomic partial accumulation, and any +order that varies with batch size, token count, SM count or any other runtime +condition. `tests/test_p1_reduction.py` pins both trees with cases where the +alternatives visibly disagree. + +## What Megatron actually does — and why it is not the byte reference + +#2 says to prefer Megatron's implementation over vLLM's. Reading the source +(`megatron/core/transformer/hyper_connection.py`, +`megatron/core/fusions/fused_mhc_kernels.py`, both on `dev`) shows what that +can and cannot mean. + +**Megatron's own two paths do not agree with each other.** `config.use_fused_mhc` +selects between a native path and a fused Triton/cuTile path: + +| | native | fused | +| --- | --- | --- | +| controller RMS | `norm = x.norm(-1)`; `r = 1/(norm/sqrt(K) + eps)` | `r_val = sqrt(sum_sq / K)`; `1/(r_val + eps)` | +| controller GEMM | `torch.matmul`, FP32 | `ct.mma(..., tfloat32)` — **TF32** | +| K reduction | whatever Inductor emits (`@torch.compile`) | **`split_k = 16` when `K >= 16384`**, and also runtime-autotuned | +| 4-stream mix | `torch.bmm` (cuBLAS) | Triton kernel | +| softmax | `torch.softmax` (`exp`) | `tl.exp2(x * log2e)` | + +The `sqrt(s)/sqrt(K)` vs `sqrt(s/K)` split is ~1 ulp; TF32 in the fused MMA is +~1e-3 relative. The authors are aware of this class of difference and manage it +as a tolerance budget — there is a comment in the fused kernel reading *"Square +in fp32: a bf16 square/reduction loses ~2e-3 relative on the RMS scale, which +native (fp32) does not."* + +Three of #2's explicit bans are violated by Megatron's fused path **at exactly +the production shape**: Split-K is on (`K = 16384` selects `split_k = 16`), +TF32 is on, and the reduction order is autotuned per machine per run. Measured +separately: eager `.sum(-1)` on CUDA is **not batch-invariant** — rows `0:8` +produce different bytes inside a 64-row batch than in an 8-row batch. On CPU it +happens to be invariant. + +None of that is a defect in Megatron. It is a correct set of performance +choices for a framework that never promised bit-exactness. But it means: + +> **"参考 Megatron" can only mean its formulas and constants — never its +> reduction order.** Reproducing Megatron's bytes is not a goal that can be +> held, because Megatron does not reproduce its own. + +What the source *does* settle, and this kit adopts verbatim: the affine +association `h = r * proj * alpha_ + bias`; the Sinkhorn schedule and its axis +convention; the max-shift before `exp`; the FP32 controller path +(`mark_keep_in_fp32` on the projection weight, alphas and bias); the three +scalar alphas; and the module decomposition that leaves the sublayer outside. + +## Frozen numeric contract (recap of #2 + decisions made here) + +From the issue: + +1. All multiplies and reduction accumulations are FP32. +2. Each operator performs **exactly one** FP32→BF16 downcast, at its output: + `mhc_pre`'s aggregated hidden, `rmsnorm_residual`'s normalized row, and + `mhc_post`'s `R_new`. No intermediate BF16 cast anywhere. +3. `PRE`, `POST`, `C`, the controller projection `P` and the RMS scale `r` + stay FP32 while travelling between operators. +4. Controller arithmetic is `h = ((r * P) * alpha) + bias`; + `PRE[i] = sigmoid(h[i]) + 1e-6`; `POST[i] = 2·sigmoid(h[4+i])`; + `L = h[8:24].reshape(4,4)`. +5. `hc_mult = 4`, layout `PRE[0:4] POST[4:8] COMB[8:24]`, `sinkhorn_iters = 20`, + `eps = 1e-6`; `sum + eps` may not be replaced by a `clamp`. +6. `rmsnorm_residual` is `rsqrt(mean(x²) + eps)`; the controller RMS is + `1/(sqrt(mean(x²)) + eps)`. The two forms are never interchangeable. + +Decisions this kit had to freeze (flagged for review on #2; changing any of +them means regenerating the manifest and bumping the schema/profile id): + +| # | Decision | Rationale | +| --- | --- | --- | +| D1 | **`q = sqrt(s) / sqrt(K)`, not `sqrt(s / K)`**, with `s` from our own fixed-tree sum-of-squares (*not* `torch.norm`, which is not bit-equal to `sqrt(x.square().sum())` on either device). Backward reuses the saved `q`. | #2 states both forms because **Megatron states both**: `native_proj_rms` computes `norm/sqrt(K)` while the fused kernel computes `sqrt(s/K)` under a comment that says `norm/sqrt(K)`. The native path is the semantic reference, so it wins. | +| D2 | **Sinkhorn `sum_row(M)[i] = Σ_j M[i,j]` (row sums, broadcast along j) and `sum_col(M)[j] = Σ_i M[i,j]`.** Schedule is literal: `softmax_row(L)+eps`, one column normalize, then 19×(row, column) — 20 column normalizations, 39 in total. | The issue names the steps but not the axis convention; this is the reading that makes rows/columns each sum to 1. | +| D3 | **`softmax_row` subtracts the row max before `exp`**, with the max taken on the same balanced 4-way tree. | Unshifted `exp` overflows on the saturating-logit fixture; the shift has to be pinned rather than left to the kernel. | +| D4 | **Sinkhorn backward walks the recorded 39 normalizations in reverse, one VJP per step.** No fixed-point / implicit-differentiation shortcut, no fused simplification. | #2 forbids "mathematically equivalent but differently associated" forms. Cross-checked against autograd in `test_p1_oracle.py`. | +| D5 | **Numeric profile `oracle-fp32-mhc-v1`**: FP32, the two trees above, mul-then-add (**no FMA fusion**). A strict CUDA kernel matches with `__fmul_rn`/`__fadd_rn` or registers its own profile. | Byte-equality needs the rounding points pinned, not just the order. | +| D6 | **`alpha` is three learnable FP32 scalars** (`alpha_pre`, `alpha_post`, `alpha_res`) broadcast over the PRE / POST / COMB segments; `bias` is a `[24]` vector. Backward returns three scalars, each a pinned segment fold on top of the token fold. | Matches Megatron exactly (`alpha_pre/post/res` are `nn.Parameter(torch.full((1,), ...))`, cat-expanded in `_compute_h`). Forward is identical to a `[24]` gain, but backward is not — a `[24]` `dAlpha` would leave the segment reduction order unpinned. | +| D7 | **The transformer sublayer is external to P1.** `y_sublayer` enters the block as data and `d_normalized` / `d_residual` enter the backward as data; `dy_sublayer` is an output boundary. | This is what makes P1 acceptance runnable with no P2–P7 code in the loop, exactly as the issue requires. | +| D8 | **`unfused` is canonical and the default.** For `rmsnorm_residual`, #18 says to prefer TE's `TEFusedResidualRMSNorm` first and self-write only if TE fails the deterministic contract — **it fails**: it refuses to expose the pre-normalization intermediate (it raises on any forward hook), so the fork and the norm cannot be hashed as separate boundaries. That is exactly the escape hatch #18 provides, so v1 is unfused. `fused-pre-norm` stays supported for an engine that physically cannot expose the intermediate. Fusion itself is not banned — changing the reduction layout or moving a downcast point is. | Train/infer byte-equality needs every boundary hashable on its own, so a divergence localizes to one operator instead of one megakernel. A fused fast path can be swapped back through the provider hook once proven byte-equal. | +| D9 | **`trainability='mixer-frozen'` returns `None` for `d_controller_weight`/`d_alpha`/`d_bias`**, not zeros. | #2: a stop-grad mixer must not leak `dMixWeight`. `None` cannot be silently summed into an optimizer; a zero tensor can. | +| D10 | **Gradients are returned FP32** (the accumulator dtype); rounding to BF16 happens only at an outer block edge. | Consistent with "FP32 reductions, BF16 boundaries". | + +## Byte-equality scope + +Strict byte-equality is required **between Megatron training and Miles +inference on the same numeric profile and device**. The committed manifest +anchors the CPU x86 oracle; `scripts/check_p1.py` recomputes the oracle on the +provider's device, so transcendentals (`sigmoid`, `exp`, `rsqrt`) never cross +devices inside a strict comparison. Hardware without equivalent capability +must register its own profile with an explicit tolerance — never silently +relax. + +The acceptance command additionally re-runs each fixture row on its own and +checks that the bytes do not move: **same row, different batch / padding / +stride ⇒ identical output**, which is acceptance criterion 2 of #2. + +## How a sub-issue PR uses the kit + +1. Subclass `ReferenceProvider`, override only the operators your PR delivers + (everything else stays on the oracle), and set `name` / `numeric_profile`: + + ```python + from rl_engine.mhc.provider import ReferenceProvider + + class MyCudaProvider(ReferenceProvider): + name = "my-cuda" + numeric_profile = "cuda-ffma-strict-v1" + + def mhc_post_fwd(self, r_old, y, c, post): + return my_cuda_kernel(r_old, y, c, post) + ``` + +2. Run `python scripts/check_p1.py --provider your.module:YourProvider + [--device cuda]`. Every boundary must be byte-equal; exit code 1 otherwise. +3. Ship the check output and your `provenance()` in the PR description. + +Fixture cases: `one_row`, `packed_t16`, `packed_t7_odd`, `fused_pre_norm`, +`mixer_frozen`, plus operator edge cases `sinkhorn_edges` (saturating +sigmoids, tied logits, a zero row, and a magnitude that makes the `sum + eps` +guard load-bearing) and `rms_edges` (zero row, subnormal-ish and large +magnitudes, exact powers of two). + +Fixture geometry is a scaled-down layer (`hidden = 128` ⇒ `K = 512`) so the +serial oracle stays CPU-cheap. `hc_mult`, `controller_n`, `sinkhorn_iters` and +both epsilons are the real production constants; the full DSv4 geometry +(`hidden = 4096`, `K = 16384`, `N = 24`) is pinned separately by +`LayerContract.assert_production()`. + +Regenerate the manifest after an intentional contract change: + +```bash +python -m rl_engine.mhc.fixtures --write-manifest +``` + +## Open questions for review + +- **Miles must also run unfused** (D8). #2 records that Miles/XoRL fuse the + mHC pre-mix and RMSNorm into a single launch. If that kernel keeps the same + reduction layout and downcast points it can stay and still be byte-equal; + if it does not, the inference side has to unfuse for v1. **This decision + moves cost onto the Miles side and needs their sign-off.** +- **Sinkhorn iteration detail** — #2 notes that TogetherAI never published the + per-round detail and that the checkpoint plus the Miles implementation are + the only reference. Megatron's two implementations agree with each other and + with the issue text, and the kit matches both; if Miles differs, D2/D4 change + and the manifest is regenerated. +- **Upstream nit worth raising (not a bug report)** — `fused_mhc_kernels.py` + computes `sqrt(sum_sq / K)` under a comment reading + `# 2. Compute r = norm / sqrt(K)`. Which is the intent? A one-line answer + confirms D1 from the other direction. +- **`dAlpha` consumers** — the three scalar gradients are DP-reduced by the + outer DDP, but any consumer that currently expects a `[24]` `dAlpha` needs + updating (D6). + +## Non-goals of the kit + +No CUDA/Triton kernels, no TE/Megatron/vLLM injection, no RoPE (that is P2's, +per #1), no attention or MoE (#4, #8/#10), and no WS2 TP/SP/CP/PP gates. The +`placement` field and the WS2 notes in the operator docstrings mark where those +gates will attach; `check_capability` already fails closed on any placement a +provider has not declared. diff --git a/rl_engine/mhc/__init__.py b/rl_engine/mhc/__init__.py new file mode 100644 index 00000000..0fb440ba --- /dev/null +++ b/rl_engine/mhc/__init__.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P1 start kit: mHC + RMSNorm deterministic forward/backward contracts (issue #2).""" + +from rl_engine.mhc.contract import ( + COMB_SLICE, + FUSION_MODES, + MHC_EPS, + ORACLE_PROFILE, + POST_SLICE, + PRE_SLICE, + PROD_CONTROLLER_N, + PROD_HIDDEN, + RMSNORM_EPS, + SCHEMA_VERSION, + SINKHORN_ITERS, + TRAINABILITY_MODES, + ControllerParams, + GradBoundary, + LayerContract, + NormParams, + ResidualBatch, + tensor_sha256, +) +from rl_engine.mhc.provider import ( + MHCProvider, + ReferenceProvider, + StubProvider, + check_capability, + resolve_provider, +) +from rl_engine.mhc.reduction import ( + HC_MULT, + STREAM4_TREE, + fixed_dot, + fixed_sum, + fixed_sumsq, + stream4_sum, +) +from rl_engine.mhc.trace import MHCTrace, first_divergence + +__all__ = [ + "COMB_SLICE", + "FUSION_MODES", + "HC_MULT", + "MHC_EPS", + "ORACLE_PROFILE", + "POST_SLICE", + "PRE_SLICE", + "PROD_CONTROLLER_N", + "PROD_HIDDEN", + "RMSNORM_EPS", + "SCHEMA_VERSION", + "SINKHORN_ITERS", + "STREAM4_TREE", + "TRAINABILITY_MODES", + "ControllerParams", + "GradBoundary", + "LayerContract", + "MHCProvider", + "MHCTrace", + "NormParams", + "ReferenceProvider", + "ResidualBatch", + "StubProvider", + "check_capability", + "first_divergence", + "fixed_dot", + "fixed_sum", + "fixed_sumsq", + "resolve_provider", + "stream4_sum", + "tensor_sha256", +] diff --git a/rl_engine/mhc/contract.py b/rl_engine/mhc/contract.py new file mode 100644 index 00000000..1ceab82e --- /dev/null +++ b/rl_engine/mhc/contract.py @@ -0,0 +1,388 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Data contracts for P1: LayerContract, ResidualBatch, controller/norm params.""" + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass, field, replace +from typing import Any + +import torch + +from rl_engine.mhc.reduction import HC_MULT + +SCHEMA_VERSION = "p1-mhc-layer-v1" + +# The oracle's numeric profile: FP32 math everywhere, long reductions as a +# serial ascending left fold, 4-way reductions as the pinned (a0+a1)+(a2+a3) +# tree, mul-then-add rounding (no FMA fusion). Kernel backends declare theirs. +ORACLE_PROFILE = "oracle-fp32-mhc-v1" + +# --- frozen production constants (issue #2) -------------------------------- +PROD_HIDDEN = 4096 # D +PROD_CONTROLLER_N = 24 # PRE 4 + POST 4 + COMB 16 +SINKHORN_ITERS = 20 +MHC_EPS = 1e-6 # controller / Sinkhorn epsilon +RMSNORM_EPS = 1e-6 + +# Controller output layout, frozen: PRE[0:4], POST[4:8], COMB[8:24]. +PRE_SLICE = slice(0, 4) +POST_SLICE = slice(4, 8) +COMB_SLICE = slice(8, 24) + +PLACEMENTS = ("replicated", "tp-sharded", "cp-token-sharded") +ROW_GEOMETRIES = ("one-row", "packed") + +# Modes the kit defines. Anything else must fail closed (acceptance 4). +# +# ``unfused`` is CANONICAL and the default. Train/infer byte-equality is the +# goal, and the unfused decomposition is the only one where every operator +# boundary can be hashed and compared on its own, so a divergence localizes to +# one operator instead of one megakernel. ``fused-pre-norm`` exists for an +# inference engine that physically cannot expose the pre-normalization +# intermediate (TE's ``TEFusedResidualRMSNorm`` refuses to: it raises if you +# hook it). Fusion is not forbidden -- what is forbidden is a fused kernel that +# changes the reduction layout or moves a downcast point. Such a kernel must +# register its own numeric profile rather than claim to be the same operator. +FUSION_MODES = ("unfused", "fused-pre-norm") +CANONICAL_FUSION_MODE = "unfused" +TRAINABILITY_MODES = ("full", "mixer-frozen") + + +def tensor_bytes(t: torch.Tensor) -> bytes: + """Raw little-endian bytes of a tensor, independent of layout.""" + flat = t.detach().contiguous().flatten() + if flat.numel() == 0: + return b"" + return flat.view(torch.uint8).cpu().numpy().tobytes() + + +def tensor_sha256(t: torch.Tensor) -> str: + return hashlib.sha256(tensor_bytes(t)).hexdigest() + + +@dataclass(frozen=True) +class LayerContract: + """Frozen shape/constant contract for one mHC + RMSNorm decoder boundary. + + ``hidden`` and ``controller_n`` are parameters so the start-kit fixtures + can run a scaled-down layer on CPU, but :meth:`assert_production` pins the + real DSv4 numbers. ``hc_mult``, ``sinkhorn_iters`` and both epsilons are + *not* negotiable -- changing any of them is a schema bump. + """ + + hidden: int = PROD_HIDDEN + hc_mult: int = HC_MULT + controller_n: int = PROD_CONTROLLER_N + sinkhorn_iters: int = SINKHORN_ITERS + mhc_eps: float = MHC_EPS + rmsnorm_eps: float = RMSNORM_EPS + layer_index: int = 0 + placement: str = "replicated" + fusion_mode: str = "unfused" + trainability: str = "full" + schema_version: str = SCHEMA_VERSION + numeric_profile: str = ORACLE_PROFILE + + @property + def flat_k(self) -> int: + """Controller GEMM K: the flattened four-stream residual width.""" + return self.hc_mult * self.hidden + + def validate(self) -> None: + if self.schema_version != SCHEMA_VERSION: + raise ValueError(f"schema {self.schema_version!r} != {SCHEMA_VERSION!r}") + if self.hc_mult != HC_MULT: + raise ValueError(f"hc_mult is frozen at {HC_MULT}, got {self.hc_mult}") + if self.controller_n != 2 * self.hc_mult + self.hc_mult**2: + raise ValueError( + f"controller_n {self.controller_n} != PRE+POST+COMB " + f"({2 * self.hc_mult + self.hc_mult ** 2})" + ) + if self.sinkhorn_iters != SINKHORN_ITERS: + raise ValueError(f"sinkhorn_iters is frozen at {SINKHORN_ITERS}") + if self.mhc_eps != MHC_EPS or self.rmsnorm_eps != RMSNORM_EPS: + raise ValueError("eps values are frozen at 1e-6") + if self.placement not in PLACEMENTS: + raise ValueError(f"unknown placement {self.placement!r}") + if self.fusion_mode not in FUSION_MODES: + raise ValueError(f"unknown fusion_mode {self.fusion_mode!r}; want {FUSION_MODES}") + if self.trainability not in TRAINABILITY_MODES: + raise ValueError( + f"unknown trainability {self.trainability!r}; want {TRAINABILITY_MODES}" + ) + if self.hidden <= 0: + raise ValueError("hidden must be positive") + + def assert_production(self) -> None: + """Fail unless this is the real DSv4 layer geometry.""" + if (self.hidden, self.controller_n, self.flat_k) != ( + PROD_HIDDEN, + PROD_CONTROLLER_N, + PROD_HIDDEN * HC_MULT, + ): + raise ValueError( + f"not the production contract: hidden={self.hidden} " + f"controller_n={self.controller_n} K={self.flat_k}" + ) + + def fingerprint(self) -> str: + h = hashlib.sha256() + for value in ( + self.hidden, + self.hc_mult, + self.controller_n, + self.sinkhorn_iters, + repr(self.mhc_eps), + repr(self.rmsnorm_eps), + self.fusion_mode, + self.trainability, + self.schema_version, + ): + h.update(str(value).encode()) + return h.hexdigest() + + +@dataclass(frozen=True) +class ControllerParams: + """Weights of the mHC controller projection (`fp32_gemm_rms` + affine). + + All FP32: issue #2 pins the controller path to FP32 end to end, with no + intermediate BF16 cast between the projection, the RMS scale and the + Sinkhorn split. + """ + + weight: torch.Tensor # FP32 [controller_n, K] + alpha_pre: torch.Tensor # FP32 scalar [1] + alpha_post: torch.Tensor # FP32 scalar [1] + alpha_res: torch.Tensor # FP32 scalar [1] + bias: torch.Tensor # FP32 [controller_n] + + @property + def alphas(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return (self.alpha_pre, self.alpha_post, self.alpha_res) + + def expanded_alpha(self, contract: LayerContract) -> torch.Tensor: + """``cat([alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(n*n)])``. + + Megatron holds three learnable *scalars* and broadcasts them across the + PRE / POST / COMB segments (``hyper_connection.py`` ``_compute_h``); it + does not hold 24 independent gains. Forward is identical either way, + but backward is not: ``dAlpha`` is three numbers, each a reduction over + its segment, so the segment reduction has to be pinned like any other. + """ + n = contract.hc_mult + return torch.cat( + [ + self.alpha_pre.expand(n), + self.alpha_post.expand(n), + self.alpha_res.expand(n * n), + ], + dim=-1, + ) + + def validate(self, contract: LayerContract) -> None: + named = ( + ("weight", self.weight), + ("alpha_pre", self.alpha_pre), + ("alpha_post", self.alpha_post), + ("alpha_res", self.alpha_res), + ("bias", self.bias), + ) + for name, t in named: + if t.dtype != torch.float32: + raise TypeError(f"controller {name} must be FP32, got {t.dtype}") + n, k = contract.controller_n, contract.flat_k + if tuple(self.weight.shape) != (n, k): + raise ValueError(f"controller weight {tuple(self.weight.shape)} != {(n, k)}") + if tuple(self.bias.shape) != (n,): + raise ValueError(f"controller bias {tuple(self.bias.shape)} != {(n,)}") + for name, t in named[1:4]: + if tuple(t.shape) != (1,): + raise ValueError(f"controller {name} must be a scalar [1], got {tuple(t.shape)}") + + def fingerprint(self) -> str: + h = hashlib.sha256() + for t in (self.weight, self.alpha_pre, self.alpha_post, self.alpha_res, self.bias): + h.update(tensor_bytes(t)) + return h.hexdigest() + + def to(self, device: torch.device | str) -> "ControllerParams": + return ControllerParams( + self.weight.to(device), + self.alpha_pre.to(device), + self.alpha_post.to(device), + self.alpha_res.to(device), + self.bias.to(device), + ) + + +@dataclass(frozen=True) +class NormParams: + """RMSNorm gain. BF16 storage, promoted to FP32 inside the operator.""" + + gamma: torch.Tensor # BF16 [hidden] + + def validate(self, contract: LayerContract) -> None: + if self.gamma.dtype != torch.bfloat16: + raise TypeError(f"gamma must be BF16, got {self.gamma.dtype}") + if tuple(self.gamma.shape) != (contract.hidden,): + raise ValueError(f"gamma shape {tuple(self.gamma.shape)} != {(contract.hidden,)}") + + def fingerprint(self) -> str: + return hashlib.sha256(tensor_bytes(self.gamma)).hexdigest() + + def to(self, device: torch.device | str) -> "NormParams": + return NormParams(self.gamma.to(device)) + + +@dataclass(frozen=True) +class ResidualBatch: + """One mHC block boundary: the four-stream residual plus the sublayer edge. + + ``r_old`` is the versioned four-way residual identity entering the block. + ``token_id`` and ``layer_index`` carry the global-token / absolute-layer + identity required by the Foundation ``SemanticTensor``; they are passed + through untouched and only participate in fingerprints. + + The transformer sublayer (attention / dense FFN / MoE) is *external* to + P1: this batch carries the two tensors that cross that boundary -- + ``y_sublayer`` (its BF16 output, consumed by ``mhc_post``) and, for the + backward direction, the incoming ``d_normalized`` / ``d_residual`` in + :class:`GradBoundary`. That is what lets P1 be developed and accepted with + no P2-P7 code in the loop. + """ + + r_old: torch.Tensor # BF16 [T, 4, hidden] + y_sublayer: torch.Tensor # BF16 [T, hidden] + controller: ControllerParams + norm: NormParams + token_id: torch.Tensor # int64 [T] global token identity + contract: LayerContract = field(default_factory=LayerContract) + row_geometry: str = "packed" + weight_fingerprint: str = "" + + @property + def tokens(self) -> int: + return int(self.r_old.shape[0]) + + @property + def hidden(self) -> int: + return int(self.r_old.shape[2]) + + def validate(self) -> None: + self.contract.validate() + if self.row_geometry not in ROW_GEOMETRIES: + raise ValueError(f"row_geometry {self.row_geometry!r} not in {ROW_GEOMETRIES}") + if self.r_old.dtype != torch.bfloat16: + raise TypeError(f"r_old must be BF16, got {self.r_old.dtype}") + if self.y_sublayer.dtype != torch.bfloat16: + raise TypeError(f"y_sublayer must be BF16, got {self.y_sublayer.dtype}") + if self.token_id.dtype != torch.int64: + raise TypeError(f"token_id must be int64, got {self.token_id.dtype}") + t, streams, hidden = self.r_old.shape + if streams != self.contract.hc_mult: + raise ValueError(f"r_old has {streams} streams, contract says {self.contract.hc_mult}") + if hidden != self.contract.hidden: + raise ValueError(f"r_old hidden {hidden} != contract {self.contract.hidden}") + if tuple(self.y_sublayer.shape) != (t, hidden): + raise ValueError(f"y_sublayer shape {tuple(self.y_sublayer.shape)} != {(t, hidden)}") + if tuple(self.token_id.shape) != (t,): + raise ValueError(f"token_id shape {tuple(self.token_id.shape)} != {(t,)}") + if self.row_geometry == "one-row" and t != 1: + raise ValueError("row_geometry 'one-row' requires exactly one token") + self.controller.validate(self.contract) + self.norm.validate(self.contract) + expected = self.compute_weight_fingerprint() + if self.weight_fingerprint and self.weight_fingerprint != expected: + raise ValueError("weight_fingerprint mismatch: checkpoint bytes were modified") + + def compute_weight_fingerprint(self) -> str: + h = hashlib.sha256() + h.update(self.contract.fingerprint().encode()) + h.update(self.controller.fingerprint().encode()) + h.update(self.norm.fingerprint().encode()) + return h.hexdigest() + + def sealed(self) -> "ResidualBatch": + """Return a copy with ``weight_fingerprint`` filled in.""" + return replace(self, weight_fingerprint=self.compute_weight_fingerprint()) + + def to(self, device: torch.device | str) -> "ResidualBatch": + return ResidualBatch( + r_old=self.r_old.to(device), + y_sublayer=self.y_sublayer.to(device), + controller=self.controller.to(device), + norm=self.norm.to(device), + token_id=self.token_id.to(device), + contract=self.contract, + row_geometry=self.row_geometry, + weight_fingerprint=self.weight_fingerprint, + ) + + +@dataclass(frozen=True) +class GradBoundary: + """The three incoming gradients at the P1 block's outer edges. + + - ``d_r_new``: from the next mHC block (or the loss), BF16 [T, 4, hidden]. + - ``d_normalized``: the sublayer's gradient w.r.t. the RMSNorm output. + - ``d_residual``: the sublayer's gradient w.r.t. the *unnormalized* hidden + forked by ``rmsnorm_residual``. Zero when nothing consumes the fork. + + Supplying the two sublayer-side gradients as data (instead of computing + them) is what keeps P1 independent of P2-P7: ``dy`` produced by + ``mhc_post_bwd`` is an *output* boundary that the sublayer owner consumes. + """ + + d_r_new: torch.Tensor # BF16 [T, 4, hidden] + d_normalized: torch.Tensor # BF16 [T, hidden] + d_residual: torch.Tensor # BF16 [T, hidden] + metadata: dict[str, Any] = field(default_factory=dict) + + def validate(self, batch: ResidualBatch) -> None: + t, hidden = batch.tokens, batch.hidden + expect = { + "d_r_new": (t, batch.contract.hc_mult, hidden), + "d_normalized": (t, hidden), + "d_residual": (t, hidden), + } + for name, shape in expect.items(): + got = getattr(self, name) + if got.dtype != torch.bfloat16: + raise TypeError(f"{name} must be BF16, got {got.dtype}") + if tuple(got.shape) != shape: + raise ValueError(f"{name} shape {tuple(got.shape)} != {shape}") + + def to(self, device: torch.device | str) -> "GradBoundary": + return GradBoundary( + d_r_new=self.d_r_new.to(device), + d_normalized=self.d_normalized.to(device), + d_residual=self.d_residual.to(device), + metadata=dict(self.metadata), + ) + + +__all__ = [ + "CANONICAL_FUSION_MODE", + "COMB_SLICE", + "FUSION_MODES", + "MHC_EPS", + "ORACLE_PROFILE", + "POST_SLICE", + "PRE_SLICE", + "PROD_CONTROLLER_N", + "PROD_HIDDEN", + "RMSNORM_EPS", + "SCHEMA_VERSION", + "SINKHORN_ITERS", + "TRAINABILITY_MODES", + "ControllerParams", + "GradBoundary", + "LayerContract", + "NormParams", + "ResidualBatch", + "tensor_bytes", + "tensor_sha256", +] diff --git a/rl_engine/mhc/fixtures.py b/rl_engine/mhc/fixtures.py new file mode 100644 index 00000000..d017698f --- /dev/null +++ b/rl_engine/mhc/fixtures.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Seeded P1 fixtures and the golden-hash manifest (start-kit acceptance data). + +Fixtures are regenerated deterministically from seeds; the committed manifest +``tests/fixtures/p1/golden_hashes.json`` anchors the golden bytes in CI. If a +torch upgrade ever changes RNG or libm behavior, the manifest test fails +loudly instead of the goldens drifting silently. + +Fixture geometry is a scaled-down layer (``hidden=128`` -> ``K=512``) so the +serial oracle stays CPU-cheap; ``hc_mult``, ``controller_n``, +``sinkhorn_iters`` and both epsilons are the real production constants. +``LayerContract.assert_production()`` pins the full DSv4 geometry separately. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from dataclasses import replace +from pathlib import Path +from typing import Any + +import torch + +from rl_engine.mhc import oracle +from rl_engine.mhc.contract import ( + ORACLE_PROFILE, + SCHEMA_VERSION, + ControllerParams, + GradBoundary, + LayerContract, + NormParams, + ResidualBatch, + tensor_sha256, +) +from rl_engine.mhc.trace import MHCTrace + +FIXTURE_HIDDEN = 128 +BASE_SEED = 2026 + +DEFAULT_MANIFEST_PATH = Path("tests/fixtures/p1/golden_hashes.json") + +BLOCK_CASES: dict[str, dict[str, Any]] = { + "one_row": {"tokens": 1, "geometry": "one-row"}, + "packed_t16": {"tokens": 16}, + "packed_t7_odd": {"tokens": 7}, + "fused_pre_norm": {"tokens": 16, "fusion_mode": "fused-pre-norm"}, + "mixer_frozen": {"tokens": 16, "trainability": "mixer-frozen"}, +} + + +def _seed_for(name: str) -> int: + digest = hashlib.sha256(name.encode()).digest() + return BASE_SEED + int.from_bytes(digest[:4], "little") + + +def _gen(name: str) -> torch.Generator: + g = torch.Generator(device="cpu") + g.manual_seed(_seed_for(name)) + return g + + +def _randn(g: torch.Generator, *shape: int, scale: float = 1.0) -> torch.Tensor: + return torch.randn(*shape, generator=g, dtype=torch.float32) * scale + + +def make_contract(name: str) -> LayerContract: + spec = BLOCK_CASES[name] + return LayerContract( + hidden=FIXTURE_HIDDEN, + layer_index=spec.get("layer_index", 3), + fusion_mode=spec.get("fusion_mode", "unfused"), + trainability=spec.get("trainability", "full"), + ) + + +def make_batch(name: str) -> ResidualBatch: + spec = BLOCK_CASES[name] + contract = make_contract(name) + g = _gen(name) + t, d, n, k = spec["tokens"], contract.hidden, contract.controller_n, contract.flat_k + batch = ResidualBatch( + r_old=_randn(g, t, contract.hc_mult, d).to(torch.bfloat16), + y_sublayer=_randn(g, t, d).to(torch.bfloat16), + controller=ControllerParams( + weight=_randn(g, n, k, scale=1.0 / float(k) ** 0.5), + alpha_pre=_randn(g, 1, scale=0.5), + alpha_post=_randn(g, 1, scale=0.5), + alpha_res=_randn(g, 1, scale=0.5), + bias=_randn(g, n, scale=0.1), + ), + norm=NormParams(gamma=(1.0 + _randn(g, d, scale=0.05)).to(torch.bfloat16)), + token_id=torch.arange(t, dtype=torch.int64) + 1000, + contract=contract, + row_geometry=spec.get("geometry", "packed"), + ).sealed() + batch.validate() + return batch + + +def make_grads(name: str, batch: ResidualBatch) -> GradBoundary: + g = _gen(name + ".grad") + t, d, s = batch.tokens, batch.hidden, batch.contract.hc_mult + grads = GradBoundary( + d_r_new=_randn(g, t, s, d).to(torch.bfloat16), + d_normalized=_randn(g, t, d).to(torch.bfloat16), + d_residual=_randn(g, t, d, scale=0.25).to(torch.bfloat16), + ) + grads.validate(batch) + return grads + + +def make_sinkhorn_edge_inputs() -> torch.Tensor: + """Edge inputs for P1-D1: saturating sigmoids, tied logits, degenerate rows. + + Row 0 pushes both sigmoid legs to the flat ends; row 1 makes every COMB + logit identical (a uniform Sinkhorn fixed point); row 2 is all zeros; row 3 + gives one row of the 4x4 a huge value so the ``sum + eps`` guard is the + only thing keeping the normalize finite -- the case where swapping in a + ``clamp`` would change bytes. + """ + rows = [ + torch.tensor([-30.0, -8.0, 8.0, 30.0] * 6, dtype=torch.float32), + torch.cat([torch.zeros(8), torch.full((16,), 0.75)]), + torch.zeros(24, dtype=torch.float32), + torch.cat( + [ + torch.tensor([0.5, -0.5, 2.0, -2.0, 1.0, -1.0, 3.0, -3.0]), + torch.tensor([40.0, -40.0, 0.0, 0.0] + [0.0] * 12), + ] + ), + ] + return torch.stack(rows) + + +def make_rms_edge_inputs() -> torch.Tensor: + """Edge inputs for P1-D5: a zero row (eps is the only guard), tiny and large + magnitudes, and an exact-power-of-two row where ``rsqrt`` and ``1/sqrt`` + are most likely to agree by accident.""" + rows = [ + torch.zeros(FIXTURE_HIDDEN), + torch.full((FIXTURE_HIDDEN,), 2.0**-12), + torch.full((FIXTURE_HIDDEN,), 4.0), + torch.linspace(-8.0, 8.0, FIXTURE_HIDDEN), + ] + return torch.stack(rows).to(torch.bfloat16) + + +def _run_case(name: str) -> dict[str, str]: + batch = make_batch(name) + trace = MHCTrace(numeric_profile=ORACLE_PROFILE) + r_new, saved = oracle.mhc_block_forward(batch, trace) + grads = make_grads(name, batch) + out = oracle.mhc_block_backward(batch, saved, grads, trace) + hashes = trace.hashes() + for key, grad in out.items(): + if grad is not None: + hashes[f"grad.{key}"] = tensor_sha256(grad) + hashes["r_new"] = tensor_sha256(r_new) + return hashes + + +def golden_manifest() -> dict[str, Any]: + """Recompute every golden hash from seeds with the FP32 oracle.""" + cases: dict[str, dict[str, str]] = {name: _run_case(name) for name in BLOCK_CASES} + + contract = LayerContract(hidden=FIXTURE_HIDDEN) + h = make_sinkhorn_edge_inputs() + pre, post, c, saved = oracle.hc_split_sinkhorn_fwd(h, contract) + g = _gen("sinkhorn_edges.grad") + dh = oracle.hc_split_sinkhorn_bwd(_randn(g, 4, 4), _randn(g, 4, 4), _randn(g, 4, 4, 4), saved) + cases["sinkhorn_edges"] = { + "pre": tensor_sha256(pre), + "post": tensor_sha256(post), + "c": tensor_sha256(c), + "grad.dh": tensor_sha256(dh), + } + + x = make_rms_edge_inputs() + gamma = (1.0 + _randn(_gen("rms_edges.gamma"), FIXTURE_HIDDEN, scale=0.05)).to(torch.bfloat16) + y, residual, rsaved = oracle.rmsnorm_residual_fwd(x, gamma, contract.rmsnorm_eps) + ge = _gen("rms_edges.grad") + dx, dgamma = oracle.rmsnorm_residual_bwd( + _randn(ge, *x.shape).to(torch.bfloat16), + _randn(ge, *x.shape, scale=0.25).to(torch.bfloat16), + x, + gamma, + rsaved, + ) + cases["rms_edges"] = { + "y": tensor_sha256(y), + "residual": tensor_sha256(residual), + "grad.dx": tensor_sha256(dx), + "grad.dgamma": tensor_sha256(dgamma), + } + + return { + "schema_version": SCHEMA_VERSION, + "numeric_profile": ORACLE_PROFILE, + "fixture_hidden": FIXTURE_HIDDEN, + "cases": cases, + } + + +def write_manifest(path: Path = DEFAULT_MANIFEST_PATH) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(golden_manifest(), indent=2, sort_keys=True) + "\n") + return path + + +def load_manifest(path: Path = DEFAULT_MANIFEST_PATH) -> dict[str, Any]: + return json.loads(path.read_text()) + + +def slice_batch(batch: ResidualBatch, start: int, stop: int) -> ResidualBatch: + """A token sub-range of a batch, used by the batch-invariance tests.""" + return replace( + batch, + r_old=batch.r_old[start:stop], + y_sublayer=batch.y_sublayer[start:stop], + token_id=batch.token_id[start:stop], + row_geometry="one-row" if stop - start == 1 else "packed", + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description="P1 golden-hash manifest tool") + parser.add_argument("--write-manifest", action="store_true") + parser.add_argument("--path", type=Path, default=DEFAULT_MANIFEST_PATH) + args = parser.parse_args() + if args.write_manifest: + print(f"wrote {write_manifest(args.path)}") + else: + print(json.dumps(golden_manifest(), indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/rl_engine/mhc/oracle.py b/rl_engine/mhc/oracle.py new file mode 100644 index 00000000..1437830c --- /dev/null +++ b/rl_engine/mhc/oracle.py @@ -0,0 +1,691 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""FP32 oracle for the six P1 operators (P1-D1..P1-D6; issue #2). + +Numeric profile ``oracle-fp32-mhc-v1``: + +- Every multiply and accumulate is FP32. Long reductions use the serial + ascending left fold, 4-way reductions the pinned ``(a0+a1)+(a2+a3)`` tree + (see :mod:`rl_engine.mhc.reduction`); nothing here calls ``torch.sum``, + ``matmul``, ``mean`` or ``einsum``. +- Every multiply and add rounds separately (mul-then-add, no FMA fusion). A + strict CUDA kernel reproduces this with ``__fmul_rn`` / ``__fadd_rn`` or + registers its own numeric profile. +- ``PRE``/``POST``/``C``, the controller projection ``P`` and the RMS scale + ``r`` stay FP32 across operator boundaries -- never cast to BF16 in transit. +- Each operator performs exactly one FP32->BF16 downcast, at its output. + ``mhc_pre`` downcasts the aggregated hidden; ``rmsnorm_residual`` downcasts + the normalized row; ``mhc_post`` downcasts ``R_new``. Nothing else. +- Gradients are returned FP32 (the accumulator dtype) and round to BF16 only + when they cross an outer block edge. + +The oracle favors auditability over speed; use the start-kit fixture sizes. +""" + +from __future__ import annotations + +import math +import sys +from typing import Any + +import torch + +from rl_engine.mhc.contract import ( + COMB_SLICE, + POST_SLICE, + PRE_SLICE, + GradBoundary, + LayerContract, + ResidualBatch, +) +from rl_engine.mhc.reduction import ( + HC_MULT, + fixed_dot, + fixed_sum, + fixed_sumsq, + stream4_max, + stream4_sum, + stream4_sum_dim, +) +from rl_engine.mhc.trace import MHCTrace + + +def _f32(t: torch.Tensor) -> torch.Tensor: + return t.to(torch.float32) + + +# --------------------------------------------------------------------------- +# 6. fixed_k_gemm -- P1-D6 reference (also the core of fp32_gemm_rms) +# --------------------------------------------------------------------------- + + +def fixed_k_gemm_fwd(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """``X @ W.T`` with one FP32 accumulator and K walked left to right. + + ``x``: [M, K], ``w``: [N, K] -> FP32 [M, N]. The batch-invariant GEMM + reference P1 owns for P2/P3/P5/P7 (issue #2, ``P1-D6``): a single unsplit + K pass, no Split-K / Stream-K / atomic partial merge, and a single cast at + the output performed by the caller. + """ + return fixed_dot(x, w) + + +def fixed_k_gemm_bwd( + dy: torch.Tensor, x: torch.Tensor, w: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Returns ``(dX, dW)`` FP32. ``dX = dY @ W`` (over N), ``dW = dY.T @ X`` (over M).""" + dx = fixed_dot(dy, w.t()) + dw = fixed_dot(dy.t(), x.t()) + return dx, dw + + +# --------------------------------------------------------------------------- +# 1. hc_split_sinkhorn -- P1-D1 +# --------------------------------------------------------------------------- + + +def _softmax_row(logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Row-wise softmax over the last (j) axis of [T, 4, 4], max-shifted. + + Frozen: subtract the row max before ``exp`` (both the max and the + denominator use the pinned 4-way tree), then divide. Returns ``(S, S)`` + where the second value is what backward needs. + """ + mx = stream4_max(logits, dim=2) # [T, 4] + shifted = logits - mx.unsqueeze(2) + e = torch.exp(shifted) + den = stream4_sum_dim(e, dim=2) # [T, 4] + s = e / den.unsqueeze(2) + return s, s + + +def _row_normalize(m: torch.Tensor, eps: float) -> tuple[torch.Tensor, torch.Tensor]: + """``M / (sum_row(M) + eps)``; ``sum_row(M)[i] = sum_j M[i, j]``.""" + rs = stream4_sum_dim(m, dim=2) + eps # [T, 4] + return m / rs.unsqueeze(2), rs + + +def _col_normalize(m: torch.Tensor, eps: float) -> tuple[torch.Tensor, torch.Tensor]: + """``M / (sum_col(M) + eps)``; ``sum_col(M)[j] = sum_i M[i, j]``.""" + cs = stream4_sum_dim(m, dim=1) + eps # [T, 4] + return m / cs.unsqueeze(1), cs + + +def _row_normalize_bwd(dn: torch.Tensor, m: torch.Tensor, rs_e: torch.Tensor) -> torch.Tensor: + g = stream4_sum_dim(dn * m, dim=2) # [T, 4] + return dn / rs_e.unsqueeze(2) - (g / (rs_e * rs_e)).unsqueeze(2) + + +def _col_normalize_bwd(dn: torch.Tensor, m: torch.Tensor, cs_e: torch.Tensor) -> torch.Tensor: + g = stream4_sum_dim(dn * m, dim=1) # [T, 4] + return dn / cs_e.unsqueeze(1) - (g / (cs_e * cs_e)).unsqueeze(1) + + +def hc_split_sinkhorn_fwd( + h: torch.Tensor, contract: LayerContract | None = None +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]: + """Split the controller row into ``(PRE, POST, C)``. + + ``h``: FP32 [T, 24]. Layout is frozen as ``PRE[0:4], POST[4:8], + COMB[8:24]``. Returns ``(pre, post, c, saved)`` with + + - ``PRE[i] = sigmoid(h[i]) + 1e-6`` -> FP32 [T, 4] + - ``POST[i] = 2 * sigmoid(h[4 + i])`` -> FP32 [T, 4] + - ``C`` = Sinkhorn-Knopp of ``L = h[8:24].reshape(4, 4)`` -> FP32 [T, 4, 4] + + The Sinkhorn schedule is literal: ``M = softmax_row(L) + eps``, one column + normalize, then 19 rounds of (row normalize, column normalize). Every + intermediate is saved so backward can walk the *same* 39 normalizations in + reverse rather than a mathematically equivalent shortcut. + """ + contract = contract or LayerContract() + eps = contract.mhc_eps + h32 = _f32(h) + if h32.shape[1] != contract.controller_n: + raise ValueError(f"h has {h32.shape[1]} controller values, want {contract.controller_n}") + + sig_pre = torch.sigmoid(h32[:, PRE_SLICE]) + pre = sig_pre + eps + sig_post = torch.sigmoid(h32[:, POST_SLICE]) + post = 2.0 * sig_post + + logits = h32[:, COMB_SLICE].reshape(-1, HC_MULT, HC_MULT) + s, _ = _softmax_row(logits) + m = s + eps + + steps: list[tuple[str, torch.Tensor, torch.Tensor]] = [] # (kind, M_in, denom_eps) + m_next, cs = _col_normalize(m, eps) + steps.append(("col", m, cs)) + m = m_next + for _ in range(contract.sinkhorn_iters - 1): + m_next, rs = _row_normalize(m, eps) + steps.append(("row", m, rs)) + m = m_next + m_next, cs = _col_normalize(m, eps) + steps.append(("col", m, cs)) + m = m_next + + saved = {"sig_pre": sig_pre, "sig_post": sig_post, "softmax": s, "steps": steps, "eps": eps} + return pre, post, m, saved + + +def hc_split_sinkhorn_bwd( + dpre: torch.Tensor, + dpost: torch.Tensor, + dc: torch.Tensor, + saved: dict[str, Any], +) -> torch.Tensor: + """Backward of ``hc_split_sinkhorn``. Returns ``dh`` FP32 [T, 24]. + + Walks the recorded normalization steps in reverse, one VJP per step, so + the association order matches the forward graph exactly (issue #2 forbids + a simplified but differently-associated form). The sigmoid legs and the + softmax leg then fold back into the same [T, 24] row. + """ + dm = _f32(dc) + for kind, m_in, denom_e in reversed(saved["steps"]): + if kind == "row": + dm = _row_normalize_bwd(dm, m_in, denom_e) + else: + dm = _col_normalize_bwd(dm, m_in, denom_e) + + # M0 = softmax(L) + eps -> dS = dM0 + s = saved["softmax"] + inner = stream4_sum_dim(dm * s, dim=2) # [T, 4] + dlogits = s * (dm - inner.unsqueeze(2)) + + sig_pre, sig_post = saved["sig_pre"], saved["sig_post"] + dh_pre = _f32(dpre) * (sig_pre * (1.0 - sig_pre)) + dh_post = _f32(dpost) * (2.0 * (sig_post * (1.0 - sig_post))) + return torch.cat([dh_pre, dh_post, dlogits.reshape(dlogits.shape[0], HC_MULT * HC_MULT)], dim=1) + + +# --------------------------------------------------------------------------- +# 2. fp32_gemm_rms -- P1-D2 (controller projection + controller RMS scale) +# --------------------------------------------------------------------------- + + +def fp32_gemm_rms_fwd( + x_flat: torch.Tensor, weight: torch.Tensor, eps: float +) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Controller projection and RMS scale from the same flattened residual. + + ``x_flat``: FP32 [T, K], ``weight``: FP32 [N, K]. Returns ``(P, r, saved)``: + + - ``P[t, n] = sum_k X[t, k] * W[n, k]`` (fixed-K, ascending, unsplit) + - ``s = sum_k X[t, k]^2``; ``norm = sqrt(s)``; ``q = norm / sqrt(K)``; + ``r = 1 / (q + eps)`` + + Note this is the *controller* RMS: ``1 / (sqrt(mean(X^2)) + eps)``, and + deliberately **not** the ``rsqrt(mean + eps)`` of ``rmsnorm_residual``. + The two must never be interchanged (issue #2 acceptance). + """ + x32 = _f32(x_flat) + p = fixed_k_gemm_fwd(x32, _f32(weight)) + s = fixed_sumsq(x32, dim=1) # [T] + norm = torch.sqrt(s) + q = norm / math.sqrt(float(x32.shape[1])) + r = 1.0 / (q + eps) + return p, r, {"s": s, "norm": norm, "q": q, "r": r, "k": int(x32.shape[1])} + + +def fp32_gemm_rms_bwd( + dp: torch.Tensor, + dr: torch.Tensor, + x_flat: torch.Tensor, + weight: torch.Tensor, + saved: dict[str, Any], +) -> tuple[torch.Tensor, torch.Tensor]: + """Returns ``(dX, dW)`` FP32; ``dX`` sums the GEMM leg and the RMS leg. + + ``dX_gemm = dP @ W``, ``dW = dP.T @ X``. For the RMS leg, with ``K``, + ``q = sqrt(s)/sqrt(K)`` and ``r = 1/(q + eps)``: + ``dX_rms[k] = g_r * ((-r^2 * X[k]) / (K * q))``. + """ + x32 = _f32(x_flat) + dx_gemm, dw = fixed_k_gemm_bwd(_f32(dp), x32, _f32(weight)) + r, q, k = saved["r"], saved["q"], float(saved["k"]) + neg_r2 = -(r * r) + denom = k * q + dx_rms = _f32(dr).unsqueeze(1) * ((neg_r2.unsqueeze(1) * x32) / denom.unsqueeze(1)) + return dx_gemm + dx_rms, dw + + +# --------------------------------------------------------------------------- +# 4a. h_aggregate -- the PRE-weighted four-stream merge inside mhc_pre +# --------------------------------------------------------------------------- + + +def h_aggregate_fwd(pre: torch.Tensor, r_old: torch.Tensor) -> torch.Tensor: + """``H = (PRE0*R0 + PRE1*R1) + (PRE2*R2 + PRE3*R3)`` -> BF16 [T, D]. + + ``pre``: FP32 [T, 4], ``r_old``: BF16 [T, 4, D]. The four streams are + promoted to FP32 before any arithmetic; the single FP32->BF16 downcast at + the end is the only rounding point in this operator. + """ + r32 = _f32(r_old) + parts = [pre[:, i].unsqueeze(1) * r32[:, i, :] for i in range(HC_MULT)] + return stream4_sum(parts).to(torch.bfloat16) + + +def h_aggregate_bwd( + dh: torch.Tensor, pre: torch.Tensor, r_old: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Returns ``(dR_from_aggregate [T, 4, D], dPRE [T, 4])``, both FP32. + + ``dR[i, d] = PRE[i] * dH[d]``; ``dPRE[i] = sum_d dH[d] * R[i, d]`` over the + full hidden width with the serial ascending fold. + """ + dh32 = _f32(dh) + r32 = _f32(r_old) + dr = torch.stack([pre[:, i].unsqueeze(1) * dh32 for i in range(HC_MULT)], dim=1) + dpre = torch.stack([fixed_sum(dh32 * r32[:, i, :], dim=1) for i in range(HC_MULT)], dim=1) + return dr, dpre + + +# --------------------------------------------------------------------------- +# 4b. mhc_pre -- P1-D4 (composite: gemm_rms -> affine -> sinkhorn -> aggregate) +# --------------------------------------------------------------------------- + + +def _controller_affine( + p: torch.Tensor, r: torch.Tensor, alpha: torch.Tensor, bias: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """``h = ((r * P) * alpha) + bias``. Returns ``(h, m)`` with ``m = r * P``. + + ``alpha`` is the [24] broadcast of the three learnable scalars; the + association ``((r * P) * alpha) + bias`` is Megatron's + ``h = r * proj * alpha_ + self.bias``. + """ + m = r.unsqueeze(1) * p + return (m * alpha.unsqueeze(0)) + bias.unsqueeze(0), m + + +def mhc_pre_fwd( + batch: ResidualBatch, ops: Any = None +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]: + """mHC entry: ``R_old -> (hidden BF16, PRE, POST, C)``. + + Steps, in the frozen order: + ``X_flat = reshape(R, [T, 4D])`` -> ``P, r = fp32_gemm_rms(X_flat)`` -> + ``h = ((r * P) * alpha) + bias`` -> ``PRE, POST, C = hc_split_sinkhorn(h)`` + -> ``H = h_aggregate(PRE, R)``. + """ + ops = ops if ops is not None else sys.modules[__name__] + contract = batch.contract + x_flat = _f32(batch.r_old).reshape(batch.tokens, contract.flat_k) + p, r, gemm_saved = ops.fp32_gemm_rms_fwd(x_flat, batch.controller.weight, contract.mhc_eps) + alpha = batch.controller.expanded_alpha(contract) + h, m = _controller_affine(p, r, alpha, batch.controller.bias) + pre, post, c, sink_saved = ops.hc_split_sinkhorn_fwd(h, contract) + hidden = ops.h_aggregate_fwd(pre, batch.r_old) + saved = { + "x_flat": x_flat, + "p": p, + "r": r, + "m": m, + "h": h, + "pre": pre, + "post": post, + "c": c, + "gemm": gemm_saved, + "sinkhorn": sink_saved, + } + return hidden, pre, post, c, saved + + +def mhc_pre_bwd( + dhidden: torch.Tensor, + dpost: torch.Tensor, + dc: torch.Tensor, + batch: ResidualBatch, + saved: dict[str, Any], + ops: Any = None, +) -> dict[str, torch.Tensor | None]: + """Composite backward: aggregate -> sinkhorn -> affine -> gemm_rms. + + ``dhidden`` is the gradient of the aggregated hidden; ``dpost``/``dc`` + arrive from :func:`mhc_post_bwd`. Returns ``d_r_old`` (aggregate leg plus + controller leg, summed in FP32) and the controller parameter gradients. + Under ``trainability='mixer-frozen'`` the controller parameter gradients + are ``None`` -- a stop-grad mixer must not leak ``dMixWeight``. + """ + ops = ops if ops is not None else sys.modules[__name__] + contract = batch.contract + dr_aggregate, dpre = ops.h_aggregate_bwd(dhidden, saved["pre"], batch.r_old) + dh = ops.hc_split_sinkhorn_bwd(dpre, dpost, dc, saved["sinkhorn"]) + + alpha = batch.controller.expanded_alpha(contract) + m, p, r = saved["m"], saved["p"], saved["r"] + dbias = fixed_sum(dh, dim=0) + # dAlpha is three scalars, not 24: each learnable alpha is broadcast over + # its segment, so its gradient is a reduction over that segment on top of + # the token reduction. Both folds are the pinned ascending order. + dalpha_per_n = fixed_sum(dh * m, dim=0) # [24] + dalpha_pre = fixed_sum(dalpha_per_n[PRE_SLICE], dim=0).reshape(1) + dalpha_post = fixed_sum(dalpha_per_n[POST_SLICE], dim=0).reshape(1) + dalpha_res = fixed_sum(dalpha_per_n[COMB_SLICE], dim=0).reshape(1) + dm = dh * alpha.unsqueeze(0) + dp = dm * r.unsqueeze(1) + dr_scale = fixed_sum(dm * p, dim=1) # [T], over the 24 controller values + + dx_flat, dweight = ops.fp32_gemm_rms_bwd( + dp, dr_scale, saved["x_flat"], batch.controller.weight, saved["gemm"] + ) + dr_controller = dx_flat.reshape(batch.tokens, contract.hc_mult, contract.hidden) + d_r_old = dr_aggregate + dr_controller + + frozen = contract.trainability == "mixer-frozen" + return { + "d_r_old": d_r_old, + "d_controller_weight": None if frozen else dweight, + "d_alpha_pre": None if frozen else dalpha_pre, + "d_alpha_post": None if frozen else dalpha_post, + "d_alpha_res": None if frozen else dalpha_res, + "d_bias": None if frozen else dbias, + } + + +# --------------------------------------------------------------------------- +# 3. mhc_post -- P1-D3 +# --------------------------------------------------------------------------- + + +def mhc_post_fwd( + r_old: torch.Tensor, y: torch.Tensor, c: torch.Tensor, post: torch.Tensor +) -> torch.Tensor: + """``R_new[j, d] = (C[0,j]R0 + C[1,j]R1) + (C[2,j]R2 + C[3,j]R3) + POST[j]*y[d]``. + + ``r_old``: BF16 [T, 4, D], ``y``: BF16 [T, D], ``c``: FP32 [T, 4, 4], + ``post``: FP32 [T, 4] -> BF16 [T, 4, D]. One FP32->BF16 downcast, at the + output. + """ + r32 = _f32(r_old) + y32 = _f32(y) + columns = [] + for j in range(HC_MULT): + parts = [c[:, i, j].unsqueeze(1) * r32[:, i, :] for i in range(HC_MULT)] + old_mix = stream4_sum(parts) + columns.append(old_mix + post[:, j].unsqueeze(1) * y32) + return torch.stack(columns, dim=1).to(torch.bfloat16) + + +def mhc_post_bwd( + dr_new: torch.Tensor, + r_old: torch.Tensor, + y: torch.Tensor, + c: torch.Tensor, + post: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Returns ``(dR_old, dy, dC, dPOST)`` FP32. + + ``dR_old[i,d] = sum_j C[i,j] G[j,d]`` and ``dy[d] = sum_j POST[j] G[j,d]`` + use the pinned 4-way tree; ``dC[i,j] = sum_d R_old[i,d] G[j,d]`` and + ``dPOST[j] = sum_d y[d] G[j,d]`` use the serial ascending fold over the + hidden width. + """ + g = _f32(dr_new) + r32 = _f32(r_old) + y32 = _f32(y) + dr_old = torch.stack( + [ + stream4_sum([c[:, i, j].unsqueeze(1) * g[:, j, :] for j in range(HC_MULT)]) + for i in range(HC_MULT) + ], + dim=1, + ) + dy = stream4_sum([post[:, j].unsqueeze(1) * g[:, j, :] for j in range(HC_MULT)]) + dc = torch.stack( + [ + torch.stack( + [fixed_sum(r32[:, i, :] * g[:, j, :], dim=1) for j in range(HC_MULT)], dim=1 + ) + for i in range(HC_MULT) + ], + dim=1, + ) + dpost = torch.stack([fixed_sum(y32 * g[:, j, :], dim=1) for j in range(HC_MULT)], dim=1) + return dr_old, dy, dc, dpost + + +# --------------------------------------------------------------------------- +# 5. rmsnorm_residual -- P1-D5 +# --------------------------------------------------------------------------- + + +def rmsnorm_residual_fwd( + x: torch.Tensor, gamma: torch.Tensor, eps: float +) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """RMSNorm with a fork of the *unnormalized* input as the residual branch. + + This is **not** ``x += residual`` followed by a norm. ``x``: BF16 [T, D], + ``gamma``: BF16 [D]. Returns ``(y BF16, residual BF16, saved)``: + + ``s = sum_d FP32(x[d])^2``; ``m = s / D``; ``r = rsqrt(m + eps)``; + ``y[d] = (FP32(x[d]) * r) * FP32(gamma[d])``. + + ``rsqrt`` is mandatory here -- mixing in ``1 / sqrt(...)`` changes bytes, + and the controller RMS in :func:`fp32_gemm_rms_fwd` uses that other form + on purpose. The residual fork keeps the original BF16 bytes untouched. + + P1-5 (#18) says to prefer TE's ``TEFusedResidualRMSNorm`` first and to + self-write only when TE's reduction/dtype fails the deterministic contract. + It fails: that module fuses the fork and the norm and refuses to expose the + intermediate (it raises on any forward hook), so the two boundaries cannot + be hashed separately and a divergence cannot be localized. v1 is therefore + unfused; a TE fast path can be swapped back through the provider hook once + it is proven byte-equal. + """ + x32 = _f32(x) + d = x32.shape[1] + s = fixed_sumsq(x32, dim=1) + m = s / float(d) + r = torch.rsqrt(m + eps) + y = (x32 * r.unsqueeze(1)) * _f32(gamma).unsqueeze(0) + residual = x.clone() + return y.to(torch.bfloat16), residual, {"r": r, "x32": x32, "d": d} + + +def rmsnorm_residual_bwd( + dy: torch.Tensor, + d_residual: torch.Tensor, + x: torch.Tensor, + gamma: torch.Tensor, + saved: dict[str, Any], +) -> tuple[torch.Tensor, torch.Tensor]: + """Returns ``(dX, dGamma)`` FP32. + + ``u[d] = dy[d]*gamma[d]``; ``q = sum_j u[j]*x[j]``; + ``dx_norm[d] = (r*u[d]) - (((x[d]*r^3)*q)/D)``; + ``dgamma[d] = sum_t (dy[t,d]*x[t,d])*r[t]``. + The input also feeds the residual fork, so ``dX = dx_norm + d_residual``. + """ + x32, r, d = saved["x32"], saved["r"], saved["d"] + dy32 = _f32(dy) + u = dy32 * _f32(gamma).unsqueeze(0) + q = fixed_sum(u * x32, dim=1) # [T] + r3 = (r * r) * r + dx_norm = (r.unsqueeze(1) * u) - (((x32 * r3.unsqueeze(1)) * q.unsqueeze(1)) / float(d)) + dgamma = fixed_sum((dy32 * x32) * r.unsqueeze(1), dim=0) + return dx_norm + _f32(d_residual), dgamma + + +# --------------------------------------------------------------------------- +# Fused pre+norm boundary (issue #2: fused/unfused equivalence case) +# --------------------------------------------------------------------------- + + +def mhc_pre_rmsnorm_fused_fwd( + batch: ResidualBatch, ops: Any = None +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]: + """The Miles/XoRL fused boundary: one call for pre-mix + normalize. + + Miles fuses the four-stream merge and the normalization into a single + launch, so M0/M1 have no separate boundary there. The oracle defines the + fused joint boundary as *exactly* the unfused composition, which makes + "fused equals unfused" a testable claim rather than an assumption. A + kernel whose fused residual store changes the reduction layout must say + so explicitly (register a different numeric profile) -- it may not present + itself as the same kernel. + """ + ops = ops if ops is not None else sys.modules[__name__] + hidden, pre, post, c, saved = ops.mhc_pre_fwd(batch, ops=ops) + normalized, residual, norm_saved = ops.rmsnorm_residual_fwd( + hidden, batch.norm.gamma, batch.contract.rmsnorm_eps + ) + saved = {**saved, "hidden": hidden, "norm": norm_saved} + return hidden, normalized, residual, post, c, {**saved, "pre": pre} + + +# --------------------------------------------------------------------------- +# Block composition (the full P1 forward/backward chain) +# --------------------------------------------------------------------------- + +SUPPORTED_FUSION = ("unfused", "fused-pre-norm") +SUPPORTED_TRAINABILITY = ("full", "mixer-frozen") + + +def _check_modes(contract: LayerContract) -> None: + """Fail-closed on any fusion / trainability mode the kit does not define.""" + if contract.fusion_mode not in SUPPORTED_FUSION: + raise NotImplementedError( + f"fusion_mode {contract.fusion_mode!r} is not defined by the P1 contract; " + f"supported: {SUPPORTED_FUSION} (fail-closed, issue #2 acceptance 4)" + ) + if contract.trainability not in SUPPORTED_TRAINABILITY: + raise NotImplementedError( + f"trainability {contract.trainability!r} is not defined by the P1 contract; " + f"supported: {SUPPORTED_TRAINABILITY} (fail-closed, issue #2 acceptance 4)" + ) + + +def mhc_block_forward( + batch: ResidualBatch, trace: MHCTrace | None = None, ops: Any = None +) -> tuple[torch.Tensor, dict[str, Any]]: + """Full P1 forward: ``R_old -> mhc_pre -> rmsnorm_residual -> [sublayer] -> mhc_post``. + + The transformer sublayer is external: its output arrives as + ``batch.y_sublayer``. Returns ``(R_new BF16, saved)``. + """ + ops = ops if ops is not None else sys.modules[__name__] + batch.validate() + _check_modes(batch.contract) + + if batch.contract.fusion_mode == "fused-pre-norm": + hidden, normalized, residual, post, c, saved = ops.mhc_pre_rmsnorm_fused_fwd(batch, ops=ops) + pre = saved["pre"] + norm_saved = saved["norm"] + else: + hidden, pre, post, c, saved = ops.mhc_pre_fwd(batch, ops=ops) + normalized, residual, norm_saved = ops.rmsnorm_residual_fwd( + hidden, batch.norm.gamma, batch.contract.rmsnorm_eps + ) + + r_new = ops.mhc_post_fwd(batch.r_old, batch.y_sublayer, c, post) + + if trace is not None: + trace.note("reduction_tree", "long=serial-ascending-left-fold; stream4=(a0+a1)+(a2+a3)") + trace.note("fma", "mul-then-add, no fusion") + trace.note("rsqrt", "rmsnorm=rsqrt(mean+eps); controller=1/(sqrt(mean)+eps)") + trace.note("downcast_points", "mhc_pre.hidden, rmsnorm.normalized, mhc_post.r_new") + trace.note("fusion_mode", batch.contract.fusion_mode) + trace.note("trainability", batch.contract.trainability) + trace.note("weight_fingerprint", batch.compute_weight_fingerprint()) + trace.record("controller.p", saved["p"]) + trace.record("controller.r", saved["r"]) + trace.record("controller.h", saved["h"]) + trace.record("split.pre", pre) + trace.record("split.post", post) + trace.record("split.c", c) + trace.record("pre.hidden", hidden) + trace.record("norm.normalized", normalized) + trace.record("norm.residual", residual) + trace.record("post.r_new", r_new) + + saved = { + **saved, + "hidden": hidden, + "normalized": normalized, + "residual": residual, + "norm": norm_saved, + "post": post, + "c": c, + "pre": pre, + } + return r_new, saved + + +def mhc_block_backward( + batch: ResidualBatch, + saved: dict[str, Any], + grads: GradBoundary, + trace: MHCTrace | None = None, + ops: Any = None, +) -> dict[str, torch.Tensor | None]: + """Full P1 backward. Returns ``dStream[0..3]`` (as ``d_r_old``), ``dy``, + ``dX``/``dResidual`` at the norm edge, ``dGamma`` and the controller + parameter gradients. + + ``grads.d_normalized`` and ``grads.d_residual`` come from the sublayer + owner: P1 never differentiates attention/FFN/MoE. ``dy_sublayer`` is an + output boundary handed back to that owner. + """ + ops = ops if ops is not None else sys.modules[__name__] + grads.validate(batch) + _check_modes(batch.contract) + + dr_old_post, dy, dc, dpost = ops.mhc_post_bwd( + grads.d_r_new, batch.r_old, batch.y_sublayer, saved["c"], saved["post"] + ) + dhidden, dgamma = ops.rmsnorm_residual_bwd( + grads.d_normalized, grads.d_residual, saved["hidden"], batch.norm.gamma, saved["norm"] + ) + pre_grads = ops.mhc_pre_bwd(dhidden, dpost, dc, batch, saved, ops=ops) + d_r_old = pre_grads["d_r_old"] + dr_old_post + + out: dict[str, torch.Tensor | None] = { + "d_r_old": d_r_old, + "dy_sublayer": dy, + "d_hidden": dhidden, + "d_gamma": dgamma, + "d_c": dc, + "d_post": dpost, + "d_controller_weight": pre_grads["d_controller_weight"], + "d_alpha_pre": pre_grads["d_alpha_pre"], + "d_alpha_post": pre_grads["d_alpha_post"], + "d_alpha_res": pre_grads["d_alpha_res"], + "d_bias": pre_grads["d_bias"], + } + if trace is not None: + trace.record("bwd.dy_sublayer", dy) + trace.record("bwd.d_c", dc) + trace.record("bwd.d_post", dpost) + trace.record("bwd.d_hidden", dhidden) + trace.record("bwd.d_gamma", dgamma) + trace.record("bwd.d_r_old", d_r_old) + for i in range(HC_MULT): + trace.record(f"bwd.d_stream{i}", d_r_old[:, i, :]) + return out + + +__all__ = [ + "SUPPORTED_FUSION", + "SUPPORTED_TRAINABILITY", + "fixed_k_gemm_bwd", + "fixed_k_gemm_fwd", + "fp32_gemm_rms_bwd", + "fp32_gemm_rms_fwd", + "h_aggregate_bwd", + "h_aggregate_fwd", + "hc_split_sinkhorn_bwd", + "hc_split_sinkhorn_fwd", + "mhc_block_backward", + "mhc_block_forward", + "mhc_post_bwd", + "mhc_post_fwd", + "mhc_pre_bwd", + "mhc_pre_fwd", + "mhc_pre_rmsnorm_fused_fwd", + "rmsnorm_residual_bwd", + "rmsnorm_residual_fwd", +] diff --git a/rl_engine/mhc/provider.py b/rl_engine/mhc/provider.py new file mode 100644 index 00000000..bf189cdc --- /dev/null +++ b/rl_engine/mhc/provider.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P1 provider interface (``P1-D4``) plus reference and fail-closed stub. + +A provider implements the six WS1 operators. Sub-issue owners (D1-D6) +subclass :class:`ReferenceProvider` and override only the methods their PR +delivers; every other method stays on the oracle, so each PR can run the full +acceptance command independently and land without waiting on the others. + +Fail-closed contract (issue #2): a provider must raise on an input it does not +support instead of silently falling back to another implementation, and +``provenance()`` must report the backend that actually ran. When an external +core (TE / Megatron) fails strict bytes, the dispatcher switches to the +RL-Kernel core through :meth:`capabilities` -- never silently. +""" + +from __future__ import annotations + +import importlib +from typing import Any, Protocol, runtime_checkable + +import torch + +from rl_engine.mhc import oracle +from rl_engine.mhc.contract import ORACLE_PROFILE, LayerContract, ResidualBatch + + +@runtime_checkable +class MHCProvider(Protocol): + """The six P1 WS1 operators. See :mod:`rl_engine.mhc.oracle` for semantics.""" + + name: str + numeric_profile: str + + def capabilities(self) -> dict[str, Any]: ... + + def provenance(self) -> dict[str, Any]: ... + + def hc_split_sinkhorn_fwd( + self, h: torch.Tensor, contract: LayerContract | None = None + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]: ... + + def hc_split_sinkhorn_bwd( + self, + dpre: torch.Tensor, + dpost: torch.Tensor, + dc: torch.Tensor, + saved: dict[str, Any], + ) -> torch.Tensor: ... + + def fp32_gemm_rms_fwd( + self, x_flat: torch.Tensor, weight: torch.Tensor, eps: float + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: ... + + def fp32_gemm_rms_bwd( + self, + dp: torch.Tensor, + dr: torch.Tensor, + x_flat: torch.Tensor, + weight: torch.Tensor, + saved: dict[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor]: ... + + def mhc_post_fwd( + self, r_old: torch.Tensor, y: torch.Tensor, c: torch.Tensor, post: torch.Tensor + ) -> torch.Tensor: ... + + def mhc_post_bwd( + self, + dr_new: torch.Tensor, + r_old: torch.Tensor, + y: torch.Tensor, + c: torch.Tensor, + post: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ... + + def h_aggregate_fwd(self, pre: torch.Tensor, r_old: torch.Tensor) -> torch.Tensor: ... + + def h_aggregate_bwd( + self, dh: torch.Tensor, pre: torch.Tensor, r_old: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: ... + + def rmsnorm_residual_fwd( + self, x: torch.Tensor, gamma: torch.Tensor, eps: float + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: ... + + def rmsnorm_residual_bwd( + self, + dy: torch.Tensor, + d_residual: torch.Tensor, + x: torch.Tensor, + gamma: torch.Tensor, + saved: dict[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor]: ... + + def fixed_k_gemm_fwd(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: ... + + def fixed_k_gemm_bwd( + self, dy: torch.Tensor, x: torch.Tensor, w: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: ... + + def mhc_pre_fwd( + self, batch: ResidualBatch, ops: Any = None + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]: ... + + def mhc_pre_bwd( + self, + dhidden: torch.Tensor, + dpost: torch.Tensor, + dc: torch.Tensor, + batch: ResidualBatch, + saved: dict[str, Any], + ops: Any = None, + ) -> dict[str, torch.Tensor | None]: ... + + def mhc_pre_rmsnorm_fused_fwd( + self, batch: ResidualBatch, ops: Any = None + ) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any] + ]: ... + + +class ReferenceProvider: + """Binds the FP32 oracle. Always passes acceptance; defines the golden bytes.""" + + name = "reference" + numeric_profile = ORACLE_PROFILE + + def capabilities(self) -> dict[str, Any]: + return { + "backend": "pytorch-oracle", + "geometry": ["one-row", "packed"], + "devices": ["cpu", "cuda"], + "fusion_modes": list(oracle.SUPPORTED_FUSION), + "trainability": list(oracle.SUPPORTED_TRAINABILITY), + "placements": ["replicated"], + } + + def provenance(self) -> dict[str, Any]: + return { + "requested_backend": self.name, + "actual_backend": self.name, + "numeric_profile": self.numeric_profile, + "torch_version": torch.__version__, + } + + hc_split_sinkhorn_fwd = staticmethod(oracle.hc_split_sinkhorn_fwd) + hc_split_sinkhorn_bwd = staticmethod(oracle.hc_split_sinkhorn_bwd) + fp32_gemm_rms_fwd = staticmethod(oracle.fp32_gemm_rms_fwd) + fp32_gemm_rms_bwd = staticmethod(oracle.fp32_gemm_rms_bwd) + mhc_post_fwd = staticmethod(oracle.mhc_post_fwd) + mhc_post_bwd = staticmethod(oracle.mhc_post_bwd) + h_aggregate_fwd = staticmethod(oracle.h_aggregate_fwd) + h_aggregate_bwd = staticmethod(oracle.h_aggregate_bwd) + rmsnorm_residual_fwd = staticmethod(oracle.rmsnorm_residual_fwd) + rmsnorm_residual_bwd = staticmethod(oracle.rmsnorm_residual_bwd) + fixed_k_gemm_fwd = staticmethod(oracle.fixed_k_gemm_fwd) + fixed_k_gemm_bwd = staticmethod(oracle.fixed_k_gemm_bwd) + mhc_pre_fwd = staticmethod(oracle.mhc_pre_fwd) + mhc_pre_bwd = staticmethod(oracle.mhc_pre_bwd) + mhc_pre_rmsnorm_fused_fwd = staticmethod(oracle.mhc_pre_rmsnorm_fused_fwd) + + +class StubProvider(ReferenceProvider): + """Fail-closed placeholder: every operator raises until a backend claims it. + + Deliberately NOT a fallback to the oracle -- issue #2 forbids silent + fallback, so an unimplemented operator must be loud. + """ + + name = "stub" + numeric_profile = "unimplemented" + + @staticmethod + def _todo(task: str) -> NotImplementedError: + return NotImplementedError( + f"P1 operator not implemented; claim it on {task} " + "(fail-closed: no silent fallback to the oracle)" + ) + + def hc_split_sinkhorn_fwd(self, h, contract=None): + raise self._todo("P1-1 (#14)") + + def hc_split_sinkhorn_bwd(self, dpre, dpost, dc, saved): + raise self._todo("P1-1 (#14)") + + def fp32_gemm_rms_fwd(self, x_flat, weight, eps): + raise self._todo("P1-2 (#15)") + + def fp32_gemm_rms_bwd(self, dp, dr, x_flat, weight, saved): + raise self._todo("P1-2 (#15)") + + def mhc_post_fwd(self, r_old, y, c, post): + raise self._todo("P1-3 (#16)") + + def mhc_post_bwd(self, dr_new, r_old, y, c, post): + raise self._todo("P1-3 (#16)") + + def h_aggregate_fwd(self, pre, r_old): + raise self._todo("P1-4 (#17)") + + def h_aggregate_bwd(self, dh, pre, r_old): + raise self._todo("P1-4 (#17)") + + def rmsnorm_residual_fwd(self, x, gamma, eps): + raise self._todo("P1-5 (#18)") + + def rmsnorm_residual_bwd(self, dy, d_residual, x, gamma, saved): + raise self._todo("P1-5 (#18)") + + def fixed_k_gemm_fwd(self, x, w): + raise self._todo("P1-2 (#15)") + + def fixed_k_gemm_bwd(self, dy, x, w): + raise self._todo("P1-2 (#15)") + + +def resolve_provider(spec: str) -> MHCProvider: + """Instantiate a provider from ``"module.path:ClassName"`` (or an alias).""" + aliases = { + "reference": "rl_engine.mhc.provider:ReferenceProvider", + "stub": "rl_engine.mhc.provider:StubProvider", + } + spec = aliases.get(spec, spec) + if ":" not in spec: + raise ValueError(f"provider spec {spec!r} must look like 'module.path:ClassName'") + module_name, class_name = spec.split(":", 1) + cls = getattr(importlib.import_module(module_name), class_name) + instance = cls() + if not isinstance(instance, MHCProvider): + raise TypeError(f"{spec} does not implement the MHCProvider protocol") + return instance + + +def check_capability(provider: MHCProvider, contract: LayerContract) -> None: + """Fail closed when the contract asks for something the provider lacks. + + This is the hook ``P1-D4`` uses at the TE/Megatron/Miles dispatch point: + an unsupported fusion mode, trainability mode or placement must raise + here, never degrade into a different-but-similar implementation. + """ + caps = provider.capabilities() + for key, label, want in ( + ("fusion_modes", "fusion mode", contract.fusion_mode), + ("trainability", "trainability mode", contract.trainability), + ("placements", "placement", contract.placement), + ): + supported = caps.get(key) + if supported is not None and want not in supported: + raise NotImplementedError( + f"provider {provider.name!r} does not support {label} {want!r}; " + f"supported: {supported} (fail-closed)" + ) + + +__all__ = [ + "MHCProvider", + "ReferenceProvider", + "StubProvider", + "check_capability", + "resolve_provider", +] diff --git a/rl_engine/mhc/reduction.py b/rl_engine/mhc/reduction.py new file mode 100644 index 00000000..a12185a4 --- /dev/null +++ b/rl_engine/mhc/reduction.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Frozen reduction primitives for the P1 mHC/RMSNorm contract (issue #2). + +Every reduction inside the P1 operators goes through this module. Nothing in +`oracle.py` is allowed to call `torch.sum`, `torch.matmul`, `mean`, `einsum` +or any other library reduction: those pick their own tree and would silently +un-freeze the arithmetic. + +Two trees are pinned, and only two: + +- **Long reductions** (K = 4 * D controller dot, D-wide sum-of-squares, + token-major parameter grads) use :func:`fixed_sum` / :func:`fixed_dot`: + a *single FP32 accumulator, left-to-right in ascending index order*, with + every multiply and add rounding separately (mul-then-add, no FMA fusion). + This is the same order the repository's existing `reduce_rows_fp32` + left-fold uses, so a P1 kernel and the WS1 VJP path agree by construction. +- **4-element stream reductions** (the four mHC residual streams, the 4x4 + Sinkhorn row/column sums) use :func:`stream4_sum`: the balanced tree + ``(a0 + a1) + (a2 + a3)`` pinned by issue #2. + +Banned everywhere downstream, per issue #2: Split-K, Stream-K, atomic partial +accumulation, and any reduction order that varies with batch size, token +count, SM count or any other runtime condition. A kernel either reproduces +these bytes (`__fmul_rn` / `__fadd_rn`) or registers its own numeric profile +-- never silently. +""" + +from __future__ import annotations + +import torch + +# The pinned 4-way tree, written out so the docstring and the code cannot drift. +STREAM4_TREE = "(a0+a1)+(a2+a3)" +LONG_REDUCTION_ORDER = "serial-ascending-left-fold" + +HC_MULT = 4 + + +def fixed_sum(x: torch.Tensor, dim: int = -1) -> torch.Tensor: + """Left-fold ``x`` along ``dim`` with a single FP32 accumulator. + + Ascending index order, one add at a time. Equivalent in value to + ``x.sum(dim)`` but with the addition order pinned. + """ + x32 = x.to(torch.float32) + x32 = x32.movedim(dim, 0) + acc = torch.zeros(x32.shape[1:], dtype=torch.float32, device=x32.device) + for i in range(x32.shape[0]): + acc = acc + x32[i] + return acc + + +def fixed_dot(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """``a @ b.T`` with FP32 mul-then-add in ascending-k order. + + ``a``: [M, K], ``b``: [N, K] (any float dtype) -> FP32 [M, N]. This is the + P1-D6 fixed-K GEMM reference: one FP32 accumulator per output element, K + walked left to right in a single unsplit pass, and a single cast at the + output (performed by the caller). Split-K / Stream-K / atomic partial + merges would all change these bytes. + """ + a32 = a.to(torch.float32) + b32 = b.to(torch.float32) + m, k = a32.shape + n, kb = b32.shape + if kb != k: + raise ValueError(f"fixed_dot K mismatch: {k} vs {kb}") + acc = torch.zeros(m, n, dtype=torch.float32, device=a32.device) + for kk in range(k): + acc = acc + a32[:, kk].unsqueeze(1) * b32[:, kk].unsqueeze(0) + return acc + + +def fixed_sumsq(x: torch.Tensor, dim: int = -1) -> torch.Tensor: + """``sum(x**2)`` along ``dim``, FP32, ascending order, mul-then-add. + + The square rounds before it is accumulated, matching a kernel that does + ``acc = __fadd_rn(acc, __fmul_rn(v, v))``. + """ + x32 = x.to(torch.float32) + x32 = x32.movedim(dim, 0) + acc = torch.zeros(x32.shape[1:], dtype=torch.float32, device=x32.device) + for i in range(x32.shape[0]): + acc = acc + x32[i] * x32[i] + return acc + + +def stream4_sum(parts: list[torch.Tensor] | tuple[torch.Tensor, ...]) -> torch.Tensor: + """The pinned 4-element tree ``(a0 + a1) + (a2 + a3)`` in FP32. + + Used for every 4-way reduction in P1: the four residual streams in + ``mhc_pre`` / ``mhc_post``, and the row/column sums of the 4x4 Sinkhorn + matrix. Note this is deliberately *not* the ascending left fold -- a + balanced tree is what the 4-stream kernels can actually emit, and issue #2 + pins it explicitly. + """ + if len(parts) != HC_MULT: + raise ValueError(f"stream4_sum expects {HC_MULT} parts, got {len(parts)}") + a0, a1, a2, a3 = (p.to(torch.float32) for p in parts) + return (a0 + a1) + (a2 + a3) + + +def stream4_sum_dim(x: torch.Tensor, dim: int) -> torch.Tensor: + """:func:`stream4_sum` over a length-4 axis of a tensor.""" + if x.shape[dim] != HC_MULT: + raise ValueError(f"stream4_sum_dim needs a length-{HC_MULT} axis, got {x.shape[dim]}") + moved = x.movedim(dim, 0) + return stream4_sum([moved[0], moved[1], moved[2], moved[3]]) + + +def stream4_max(x: torch.Tensor, dim: int) -> torch.Tensor: + """Max over a length-4 axis with the same balanced tree as :func:`stream4_sum`. + + ``max(max(a0, a1), max(a2, a3))`` -- pinned so the Sinkhorn softmax shift + cannot drift between a tree reduction and a serial scan. (Max is + associative in IEEE arithmetic, so this fixes the code, not the value.) + """ + if x.shape[dim] != HC_MULT: + raise ValueError(f"stream4_max needs a length-{HC_MULT} axis, got {x.shape[dim]}") + m = x.movedim(dim, 0).to(torch.float32) + return torch.maximum(torch.maximum(m[0], m[1]), torch.maximum(m[2], m[3])) + + +__all__ = [ + "HC_MULT", + "LONG_REDUCTION_ORDER", + "STREAM4_TREE", + "fixed_dot", + "fixed_sum", + "fixed_sumsq", + "stream4_max", + "stream4_sum", + "stream4_sum_dim", +] diff --git a/rl_engine/mhc/trace.py b/rl_engine/mhc/trace.py new file mode 100644 index 00000000..aa62c20e --- /dev/null +++ b/rl_engine/mhc/trace.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Boundary trace for P1 (first-divergence localization). + +Every operator boundary records (name, dtype, shape, sha256 of raw bytes). +This is the P1-local stand-in for the Foundation ``TraceEnvelope``; the +``notes`` map carries the arithmetic provenance issue #2 requires a trace to +record -- reduction tree, rsqrt vs 1/sqrt, FMA policy, rounding points and +addition order. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + +from rl_engine.mhc.contract import tensor_sha256 + + +@dataclass(frozen=True) +class BoundaryRecord: + name: str + dtype: str + shape: tuple[int, ...] + sha256: str + + def to_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "dtype": self.dtype, + "shape": list(self.shape), + "sha256": self.sha256, + } + + +@dataclass +class MHCTrace: + """Ordered boundary hashes plus arithmetic provenance for one P1 run.""" + + numeric_profile: str + records: list[BoundaryRecord] = field(default_factory=list) + notes: dict[str, str] = field(default_factory=dict) + + def record(self, name: str, tensor: torch.Tensor) -> None: + self.records.append( + BoundaryRecord( + name=name, + dtype=str(tensor.dtype).removeprefix("torch."), + shape=tuple(tensor.shape), + sha256=tensor_sha256(tensor), + ) + ) + + def note(self, key: str, value: str) -> None: + self.notes[key] = value + + def hashes(self) -> dict[str, str]: + return {r.name: r.sha256 for r in self.records} + + def to_dict(self) -> dict[str, Any]: + return { + "numeric_profile": self.numeric_profile, + "records": [r.to_dict() for r in self.records], + "notes": dict(self.notes), + } + + +def first_divergence(a: MHCTrace, b: MHCTrace) -> str | None: + """Name of the first boundary whose hash differs, or None if identical.""" + for ra, rb in zip(a.records, b.records, strict=False): + if ra.name != rb.name: + return ra.name + if ra.sha256 != rb.sha256: + return ra.name + if len(a.records) != len(b.records): + return "" + return None + + +__all__ = ["BoundaryRecord", "MHCTrace", "first_divergence"] diff --git a/scripts/check_p1.py b/scripts/check_p1.py new file mode 100755 index 00000000..871f1c59 --- /dev/null +++ b/scripts/check_p1.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P1 start-kit acceptance command (issue #2, ``P1-S0``). + +Runs a provider's operators through the frozen mHC + RMSNorm block and +compares every operator boundary byte-for-byte against the FP32 oracle +executed on the same device. Any mismatching strict boundary fails the run. + +Examples: + python scripts/check_p1.py + python scripts/check_p1.py --provider mypkg.p1:CudaMHCProvider --device cuda + python scripts/check_p1.py --cases packed_t16,fused_pre_norm --json out.json +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import sys +from typing import Any + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.mhc import fixtures, oracle # noqa: E402 +from rl_engine.mhc.contract import tensor_sha256 # noqa: E402 +from rl_engine.mhc.provider import ( # noqa: E402 + MHCProvider, + check_capability, + resolve_provider, +) +from rl_engine.mhc.trace import MHCTrace # noqa: E402 + + +def _compare(golden: dict[str, str], candidate: dict[str, str]) -> list[dict[str, Any]]: + rows = [] + for name, want in golden.items(): + got = candidate.get(name) + rows.append( + { + "boundary": name, + "ok": got == want, + "golden": want[:12], + "got": (got or "")[:12], + } + ) + return rows + + +def _hashes(trace: MHCTrace, grads: dict[str, Any], r_new: Any) -> dict[str, str]: + out = trace.hashes() + for key, grad in grads.items(): + if grad is not None: + out[f"grad.{key}"] = tensor_sha256(grad) + out["r_new"] = tensor_sha256(r_new) + return out + + +def _run_block(provider: MHCProvider, name: str, device: str) -> list[dict[str, Any]]: + batch = fixtures.make_batch(name).to(device) + check_capability(provider, batch.contract) + grads = fixtures.make_grads(name, fixtures.make_batch(name)).to(device) + + gold_trace = MHCTrace(numeric_profile="oracle") + r_gold, saved_gold = oracle.mhc_block_forward(batch, gold_trace) + grads_gold = oracle.mhc_block_backward(batch, saved_gold, grads, gold_trace) + + cand_trace = MHCTrace(numeric_profile=provider.numeric_profile) + r_cand, saved_cand = oracle.mhc_block_forward(batch, cand_trace, ops=provider) + grads_cand = oracle.mhc_block_backward(batch, saved_cand, grads, cand_trace, ops=provider) + + return _compare( + _hashes(gold_trace, grads_gold, r_gold), _hashes(cand_trace, grads_cand, r_cand) + ) + + +def _run_batch_invariance(provider: MHCProvider, name: str, device: str) -> list[dict[str, Any]]: + """Same row, different batch: bytes must not move (issue #2 acceptance).""" + batch = fixtures.make_batch(name).to(device) + full, _ = oracle.mhc_block_forward(batch, ops=provider) + rows = [] + for row in range(batch.tokens): + single = fixtures.slice_batch(batch, row, row + 1) + one, _ = oracle.mhc_block_forward(single, ops=provider) + ok = bool((one[0] == full[row]).all()) + rows.append( + { + "boundary": f"batch_invariance.row{row}", + "ok": ok, + "golden": "packed", + "got": "one-row" if ok else "DIVERGED", + } + ) + return rows + + +def main() -> int: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "--provider", default="reference", help="'reference', 'stub', or module.path:ClassName" + ) + parser.add_argument("--cases", default=None, help="comma-separated case names (default: all)") + parser.add_argument("--device", default="cpu") + parser.add_argument("--json", dest="json_path", default=None, help="write full report as JSON") + args = parser.parse_args() + + provider = resolve_provider(args.provider) + names = list(fixtures.BLOCK_CASES) + if args.cases: + wanted = set(args.cases.split(",")) + unknown = wanted - set(names) + if unknown: + parser.error(f"unknown cases: {sorted(unknown)}") + names = [n for n in names if n in wanted] + + report: dict[str, Any] = { + "provider": provider.name, + "device": args.device, + "provenance": provider.provenance(), + "capabilities": provider.capabilities(), + "cases": {}, + } + failed = False + for name in names: + try: + rows = _run_block(provider, name, args.device) + rows += _run_batch_invariance(provider, name, args.device) + except NotImplementedError as exc: + rows = [ + {"boundary": "", "ok": False, "golden": "", "got": f"NotImplemented: {exc}"} + ] + report["cases"][name] = rows + case_ok = all(r["ok"] for r in rows) + failed = failed or not case_ok + print(f"[{'PASS' if case_ok else 'FAIL'}] {name}") + for r in rows: + print( + f"{' ok ' if r['ok'] else ' XX '} {r['boundary']:<28} " + f"golden={r['golden']} got={r['got']}" + ) + + print(f"\nprovider={provider.name} profile={provider.numeric_profile} device={args.device}") + if args.json_path: + with open(args.json_path, "w") as fh: + json.dump(report, fh, indent=2, sort_keys=True) + print(f"report written to {args.json_path}") + print( + "RESULT:", + "FAIL (strict boundaries diverged)" if failed else "PASS (all boundaries byte-equal)", + ) + return 1 if failed else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/fixtures/p1/golden_hashes.json b/tests/fixtures/p1/golden_hashes.json new file mode 100644 index 00000000..931ad49e --- /dev/null +++ b/tests/fixtures/p1/golden_hashes.json @@ -0,0 +1,184 @@ +{ + "cases": { + "fused_pre_norm": { + "bwd.d_c": "a0d671287e9d89450271cb9c6d1b75dc8ac1b1bef24a6fddfba6175b71336b92", + "bwd.d_gamma": "67d216275d0b662b70b423086ba11d1e7193685843ddd9cd6c6e106e06ddd4db", + "bwd.d_hidden": "bf734dd46102b092fba3d05b6e6a536ce2896d87a56a2afef0383cf3507235fb", + "bwd.d_post": "aa77b606f48bbc6d1a3b3e5bbc85b01fd8c531aa8986b018a1905c705fb50608", + "bwd.d_r_old": "a1b887409a74702a622ef408b7129c3ca30f0f81f3cc1b96070d5e565506c8aa", + "bwd.d_stream0": "4ce06c382fdc5297fedc05fec66fc208c5198cc5ded4c1b89d25229208bffe0c", + "bwd.d_stream1": "d3c2db80747749d03567fd59998fa886fb346adcfb43df83f45ecb84b6d72fdf", + "bwd.d_stream2": "ab46b2bd14c703c9f04408a9e6cfbf4111c4a63d2901b43fc1161b92ffd3de64", + "bwd.d_stream3": "e63d9a64ac1527bf8d8fe2fa74ad7485e653336c76984790d616ee9c43f74fc3", + "bwd.dy_sublayer": "444f7ddd3091f6c6419b69020040744a8ac18c9567ceaa30dc8280d89ed69282", + "controller.h": "f1c44c60446abef06f1e712f7319f8338258e9b0e3fefa2fc36deda16b34cc3f", + "controller.p": "73a3898bd58ab559879d3258686805bcd506817e29cff96b61cada200a55efbb", + "controller.r": "952f64296d063062e522f9d629ca6622f09aea37a498e35ea2e8b65d70a765e0", + "grad.d_alpha_post": "8a67d7c57d1ac30e8fe014e377c9e2c06f86e116c69966697e49843a25ae058d", + "grad.d_alpha_pre": "02627e3a60ffa5feebd4da31629123bfbf98b4bd3361b1c31bcd4a431a7f1977", + "grad.d_alpha_res": "9f137f5cbddc9bece69111964bdd9cd94c56a691cc2c610777b87309472e314c", + "grad.d_bias": "54786b589f363d34ee69e6c6242b465651ac77f6b6bb4b0ff5ef37a4ab0fe92b", + "grad.d_c": "a0d671287e9d89450271cb9c6d1b75dc8ac1b1bef24a6fddfba6175b71336b92", + "grad.d_controller_weight": "a2b206bf627cea6308dedcb75e227795e4ed9d3a8ca46a66e8c2911ea74d428f", + "grad.d_gamma": "67d216275d0b662b70b423086ba11d1e7193685843ddd9cd6c6e106e06ddd4db", + "grad.d_hidden": "bf734dd46102b092fba3d05b6e6a536ce2896d87a56a2afef0383cf3507235fb", + "grad.d_post": "aa77b606f48bbc6d1a3b3e5bbc85b01fd8c531aa8986b018a1905c705fb50608", + "grad.d_r_old": "a1b887409a74702a622ef408b7129c3ca30f0f81f3cc1b96070d5e565506c8aa", + "grad.dy_sublayer": "444f7ddd3091f6c6419b69020040744a8ac18c9567ceaa30dc8280d89ed69282", + "norm.normalized": "1a54037dee5233740a7c9df57d210811203a3ab8208cbafb8a000fdaeaee3cf4", + "norm.residual": "2978118490b3e19d0a7d2392072406a3c78ee43fbeec54174964d7b7df8264f2", + "post.r_new": "45135d2a142d346a740e42069c4f7021c2eed703f27e0db69960911bc6d6ce14", + "pre.hidden": "2978118490b3e19d0a7d2392072406a3c78ee43fbeec54174964d7b7df8264f2", + "r_new": "45135d2a142d346a740e42069c4f7021c2eed703f27e0db69960911bc6d6ce14", + "split.c": "eef8c4b78a4a97dcf73e9ee3e015b2979fcef50422e90e6f38c4a734e3806b52", + "split.post": "0784fa6420d21c5497084e2bb59a4e3e974c4bbb022a55abd0a638675c6c61af", + "split.pre": "53ea4a5bea502576e78149b9ac224a6de0253f88c720c5721eaa633312f4fa8c" + }, + "mixer_frozen": { + "bwd.d_c": "644b562d8446861f9a703af33ca08a56d6935f82a24435cfdaf269c505a487cf", + "bwd.d_gamma": "e9c922f0aa7292cc6275a506df1f027561fa78480bdff680fcafc04da496ea8c", + "bwd.d_hidden": "27eec72de166e1df1f627b6ef46dd15dba9fc54ce709723805782da2e5d8c70a", + "bwd.d_post": "00822250814f728cfae3e061462cbc24a66dc50070e67c4e1bb3ae21df714563", + "bwd.d_r_old": "be1207bf5b7572a108e1e3a83ba2fae9eafc1c330b8160c3725a552e6be609bf", + "bwd.d_stream0": "a5c42a8f20e2a91713aff052cc09d1cb4c90e6705bd45c00f39e31faaf154616", + "bwd.d_stream1": "c364d5b3572fd519465f9cea007f88b9d75e138e2846195b9ea6d16b0214c3a4", + "bwd.d_stream2": "b17f38a18573bc4cfb3f12935965e30e71422717bfb628103922a0d5827eff54", + "bwd.d_stream3": "08dbb342b118ffb0ee50a89a059f7016b66fee6fa78f1fdb81e3502651ee934d", + "bwd.dy_sublayer": "0921cf324295ee1e1f61d95be200e0caadd88ccaf0a8a2d828c91da42ea3c73f", + "controller.h": "8fd95ecc1112d33f4d056a6ca03dfe5fc518cc7cd00602e1d251362bf68a2753", + "controller.p": "5e63ce025161fb5b9197c0f9655e5964d37291576f46e50ff75711a11e422cf6", + "controller.r": "e3184f50bbeb6887ce5f3785cd8d036b31d6619c411a91ba09abe46da630ad17", + "grad.d_c": "644b562d8446861f9a703af33ca08a56d6935f82a24435cfdaf269c505a487cf", + "grad.d_gamma": "e9c922f0aa7292cc6275a506df1f027561fa78480bdff680fcafc04da496ea8c", + "grad.d_hidden": "27eec72de166e1df1f627b6ef46dd15dba9fc54ce709723805782da2e5d8c70a", + "grad.d_post": "00822250814f728cfae3e061462cbc24a66dc50070e67c4e1bb3ae21df714563", + "grad.d_r_old": "be1207bf5b7572a108e1e3a83ba2fae9eafc1c330b8160c3725a552e6be609bf", + "grad.dy_sublayer": "0921cf324295ee1e1f61d95be200e0caadd88ccaf0a8a2d828c91da42ea3c73f", + "norm.normalized": "b344899ce2180c86e8141022a2dd450fd57bad36203bdf89e818af165f3db26a", + "norm.residual": "29a14e4bd27107a76715ce15d3251441d9788c2c0a662bb811f4e647d1579000", + "post.r_new": "f7bef1505cfb18a0bf5dff091e2b847cfa9ec3f7f27236d9e3cad9096e5b6066", + "pre.hidden": "29a14e4bd27107a76715ce15d3251441d9788c2c0a662bb811f4e647d1579000", + "r_new": "f7bef1505cfb18a0bf5dff091e2b847cfa9ec3f7f27236d9e3cad9096e5b6066", + "split.c": "f3a5f798d811fddd55eb63fb15601eacbd02d774d7cf16ff72572b7d7b5a41e4", + "split.post": "d99a2eaa6a1ece877589528ae2c9f17076d1b93f46ad97987a248b9edd8637af", + "split.pre": "941cf739e0e6eb48a5de80f63e3be1c6e667dbe75a81ce84b7ffe599b38ec016" + }, + "one_row": { + "bwd.d_c": "2bf1db541dedf0efb695b6a4d2c4381740f4f0eb42a647edd7ceba8a4e817ebd", + "bwd.d_gamma": "51b20c678e10e5ee453fe448c2e7f8490dd0472ddda6f29f1e821293e4865200", + "bwd.d_hidden": "28ca22b277231a74721a26e8cd46bf3ae5c1905d43fdc76087a36e4234bc49cc", + "bwd.d_post": "1330c48fdaa948c9e111f503fce0e7bb058b1d79ff5811da6f9496c43c2afe8f", + "bwd.d_r_old": "c8fee6ea81ba5b3dae9d6ebb0ddfca0f9599fe504b85d24627e1f24546a184e5", + "bwd.d_stream0": "91ddf36fbd0b9fd52bba34a6ce57a41f5e657c00f2c5ef2988eee9911a36304f", + "bwd.d_stream1": "67b3619b050457cd55dd15ce39bad1c11d4b94b3e39b6ea61647831e51cab1de", + "bwd.d_stream2": "5beb40d831369a03afc4035a1495f5ccb1ccb063102b40cce7661c6fd8217164", + "bwd.d_stream3": "2d548bdcc5328d80b907aee943f6805335a4eebfc977bb6678d1a70aa9b52850", + "bwd.dy_sublayer": "9114c281233bea57d349f7fac5f74c9621e5aba91d9975081b0b0f3259ac6bd6", + "controller.h": "9eccb103cc6cc44679ff751427c488ecfb51e8e2f65ad77cb8ac73cdab87e1b3", + "controller.p": "6d2da237edf96ff10377dcf764806dac926b63b868e7bf2764e47ffc168953ec", + "controller.r": "0ed8fee556e17362f225ff27d1c82a31a62b81a78b0a71b171ed109b298e2ba2", + "grad.d_alpha_post": "f07b93cc8421529a00410c176a854a9e1f7fca0ff9101ba7e496ba4b8218aa2c", + "grad.d_alpha_pre": "40252d74346979825daa34d73a145744d3f4dde40923a1bf655260e4394d53d0", + "grad.d_alpha_res": "1d462ff3cb420cfa7742d055f4c12bc33d3cace8807a1f543fa64add65d8d592", + "grad.d_bias": "5bffc0eb57cd794fcbc1643174d8796fe5b46a9a07237b98f9641d70eabb5295", + "grad.d_c": "2bf1db541dedf0efb695b6a4d2c4381740f4f0eb42a647edd7ceba8a4e817ebd", + "grad.d_controller_weight": "ba2e502197129e625dad9ef66686aa95a133a2743d589c434f3b9c4c6f8311d1", + "grad.d_gamma": "51b20c678e10e5ee453fe448c2e7f8490dd0472ddda6f29f1e821293e4865200", + "grad.d_hidden": "28ca22b277231a74721a26e8cd46bf3ae5c1905d43fdc76087a36e4234bc49cc", + "grad.d_post": "1330c48fdaa948c9e111f503fce0e7bb058b1d79ff5811da6f9496c43c2afe8f", + "grad.d_r_old": "c8fee6ea81ba5b3dae9d6ebb0ddfca0f9599fe504b85d24627e1f24546a184e5", + "grad.dy_sublayer": "9114c281233bea57d349f7fac5f74c9621e5aba91d9975081b0b0f3259ac6bd6", + "norm.normalized": "81fd14fe825242d6c85dc88a69823b806927b1745b1cd05b0626a06594f4a671", + "norm.residual": "ea07cdd0ffec7f466120390cf091c3aaa93d8f9d3f593254271eccedfc84550a", + "post.r_new": "9e4fcda5fe27e82c3f838500e870f859af7124df969026da044968be95fbc10c", + "pre.hidden": "ea07cdd0ffec7f466120390cf091c3aaa93d8f9d3f593254271eccedfc84550a", + "r_new": "9e4fcda5fe27e82c3f838500e870f859af7124df969026da044968be95fbc10c", + "split.c": "c25cb9cd1919a8a4f571ade0a02e9e0e4e6c09a65f9819ef4e2e2ff878d696cf", + "split.post": "3cab2cba40ba5ce25a39cfd946c8ea1ac4aee13fa2484e278f50b6d7cb27161b", + "split.pre": "07e191600f963ce7ba12d0bbcc89015d68f66655a54d5710791e67ebf73de2da" + }, + "packed_t16": { + "bwd.d_c": "3f573a2a506bd2669ab318c84adc88c06efaa0274bcdd8117013555283c23ec9", + "bwd.d_gamma": "8d30d84cfb87751048d599c882b4c873c3e5f5e55d85aaf9596fc2eadb09edee", + "bwd.d_hidden": "6567631097691d5fa4c9e382a59ba68a262ee4a6290b3c0c86730cafa8d5cf0d", + "bwd.d_post": "edb376c01993aa5fff0e98713f1ab16a6aed8e94b941d1e7547da40eb2b556c4", + "bwd.d_r_old": "b87ea37c94216723b3e7278e5e19b82429a58e446472f29a81d65611e443612c", + "bwd.d_stream0": "2cbe56c127c1ac6302f8d1d20774b7b14a47439d4d5a2af1ea2ee430898beb2e", + "bwd.d_stream1": "39dfd0fb34a9c584b3ba99bf6daa50be15aadad3ba818fdcd52ccad98d7c7ec6", + "bwd.d_stream2": "b7e80f100634af45c5eab616d5f2c2c85193be1cc2bcb64467187204ea5c519a", + "bwd.d_stream3": "be08bf2bcd8a2710ccc9b4c8dfb62843fd5e86e09bae04c26aa9868911fa23e1", + "bwd.dy_sublayer": "81e0f374e53876f1fd4ef416a8836b7cb69f2e693a40ae76d8b94bc1b48de4a1", + "controller.h": "7b8157f4b4cce2623e6e5c5f53ae4f6192258ed927376809ac9c20ba6c11f3db", + "controller.p": "e423cb667df95f122c705682d22d33238935166156552c50c8ce24d48caaf43c", + "controller.r": "ebc5d43e3fda0a178bfe9131405dfd2ea16b6143daedb0fcb44df72a5b6beb9f", + "grad.d_alpha_post": "dc7a6974c8da4ad9716b297b0f1f5bc843542b521001e6857316274abeea814c", + "grad.d_alpha_pre": "6a003585d682fce6abf8d03c140bcd531cf03c3dd799a4be6a46c69efce19ec2", + "grad.d_alpha_res": "41f95d3857b1d7f5979ba8223f684e882fd131fb5d1826c2bfe437d1024ba87d", + "grad.d_bias": "49d3f3329e8b324771555ddf74382679c9924a753f080d3335add19fe35705b0", + "grad.d_c": "3f573a2a506bd2669ab318c84adc88c06efaa0274bcdd8117013555283c23ec9", + "grad.d_controller_weight": "9f60ac4ce393f66768d7977115960a1454c717a688129dafbba07f0515c36dc7", + "grad.d_gamma": "8d30d84cfb87751048d599c882b4c873c3e5f5e55d85aaf9596fc2eadb09edee", + "grad.d_hidden": "6567631097691d5fa4c9e382a59ba68a262ee4a6290b3c0c86730cafa8d5cf0d", + "grad.d_post": "edb376c01993aa5fff0e98713f1ab16a6aed8e94b941d1e7547da40eb2b556c4", + "grad.d_r_old": "b87ea37c94216723b3e7278e5e19b82429a58e446472f29a81d65611e443612c", + "grad.dy_sublayer": "81e0f374e53876f1fd4ef416a8836b7cb69f2e693a40ae76d8b94bc1b48de4a1", + "norm.normalized": "9fc48539e3a340b22ff0bfb1153b85bc1f67bbc0360f3264d65bf2bc6d9f9cd3", + "norm.residual": "46a4e631c4e08bb70f18f1ee0f004fc8e6e1711c37338a995d5113b4d6fc73dc", + "post.r_new": "49b63364af3a072ad95418796088e3e07ed6b63886d62b9d64377cf2c07d412f", + "pre.hidden": "46a4e631c4e08bb70f18f1ee0f004fc8e6e1711c37338a995d5113b4d6fc73dc", + "r_new": "49b63364af3a072ad95418796088e3e07ed6b63886d62b9d64377cf2c07d412f", + "split.c": "e80340fd56f6e52f0542fc3bac4ec806e10761d080791efd7a3bfdc02333d5ba", + "split.post": "b6f15c50289811652a74a0ff09bad85f0570327d4e74d629d3e41a2e0869a9f4", + "split.pre": "ddd70698cb86e35a6a6a0bbc267e7695469254245e593e763a2611e8d7fa060f" + }, + "packed_t7_odd": { + "bwd.d_c": "3d7cc5332f9349dac8e47efc9f0360d56e528b41919133beaae86fc3421380ac", + "bwd.d_gamma": "5055b0a729868cb304b393ef4ea54536b9acc26150e7411bdb8f33915be4e89d", + "bwd.d_hidden": "a77f31a38d1fc394e449a6d8ee73afb4c94583f1cb83bccdf329d197babaf29d", + "bwd.d_post": "049df2a605804eee375c008e2807d47be65c160927d49eed8d4b9c5963adf055", + "bwd.d_r_old": "2dd600c0599e4a49fc14eee7f28384aae5876fbd0555facbc7d02bad6640b7ec", + "bwd.d_stream0": "dc64eb548c0e990b9d71070403512a50cacdf092088ccce88e95bb9b8dd025ef", + "bwd.d_stream1": "0e7b45bd2d5920ded6ca9845913f0709351996dc14d4c91f28626ff68442c2f9", + "bwd.d_stream2": "4fef807b573c1c3a2f9f5b4d48c404d6bd0f0896d4039ad66b7fd2009f11ce14", + "bwd.d_stream3": "4202168e9ec69edf0af8746928bfbbac021e23f67897c3a41612911e4a1b960b", + "bwd.dy_sublayer": "cdfe144499bdea8666be824211f590640483548025cc8ba01948043744a2d35a", + "controller.h": "6cd7ecab852d5394b4dd5ebdba5990419103196cd35a8545970eee4eb495a0f3", + "controller.p": "3c568c36bd3a848bca0499afebb4e88350bb9095a63d2159092e6e8fee6cf777", + "controller.r": "1e0aebf1c1829bcc5b4fcfd08d7d7ccda679d2542e8d69fbaadd2417072156e6", + "grad.d_alpha_post": "094b1f01caad8a33d340884ffad37ab1597e1b6d9dd3096389359969f0970bf2", + "grad.d_alpha_pre": "f34ff84cc6aae4e88c37f5ff5f9335aa14fe3096b06fe042544a65feab3fd0e5", + "grad.d_alpha_res": "ff77def690d9b24c3fbf1d4ce4586a5622ab4e7df8b1ad5de45255f1df33d4ea", + "grad.d_bias": "fb7a00c36c2a695859a3bc7f2e26e5dbd5a9be6e4d8134f10050673a504f644d", + "grad.d_c": "3d7cc5332f9349dac8e47efc9f0360d56e528b41919133beaae86fc3421380ac", + "grad.d_controller_weight": "86bb3903b8aec54115c7ba1c2c7a2d6f0ab61fcdcc79554498a3b47c806deddb", + "grad.d_gamma": "5055b0a729868cb304b393ef4ea54536b9acc26150e7411bdb8f33915be4e89d", + "grad.d_hidden": "a77f31a38d1fc394e449a6d8ee73afb4c94583f1cb83bccdf329d197babaf29d", + "grad.d_post": "049df2a605804eee375c008e2807d47be65c160927d49eed8d4b9c5963adf055", + "grad.d_r_old": "2dd600c0599e4a49fc14eee7f28384aae5876fbd0555facbc7d02bad6640b7ec", + "grad.dy_sublayer": "cdfe144499bdea8666be824211f590640483548025cc8ba01948043744a2d35a", + "norm.normalized": "8d8566101c118b65c78416bdb416a8665e0b2c2bc9869ddc2b485be50bd3145e", + "norm.residual": "ebb08aa5e7e53462a797c017d87da787792dc35d56adb244c1a1683122abc996", + "post.r_new": "32b2f4369a4639b3ad8546a53d7398eb54490497b5423e293ccf7633aed5b65e", + "pre.hidden": "ebb08aa5e7e53462a797c017d87da787792dc35d56adb244c1a1683122abc996", + "r_new": "32b2f4369a4639b3ad8546a53d7398eb54490497b5423e293ccf7633aed5b65e", + "split.c": "ecb2a76f3438f88c2b7c3f45c6d5c5344032e9cda0f4efc410102ed539c02078", + "split.post": "dc42d166bfb37b491745a702d2570c08f02a15dd9766cf7166fbc011823fab0c", + "split.pre": "3f73eafd8f8e970983e8aef545900d9fb8c58d22f5d1d193bd26f9c9664cf416" + }, + "rms_edges": { + "grad.dgamma": "3555e915bc9d4e7a438e217d897b5365033518b01ef351dc4b2ffbc7d3c3ab43", + "grad.dx": "88c55e55bea359619895f27bd269d1e0470e6ddf950e437c12718fd6df6b411c", + "residual": "295b8bde70b086d451eb17ceb1213ce94cebffef15cffc6464aed39de530c763", + "y": "531e66529bb3b9d5947a9b6442e1bd6ce1acca2eed9908e0974aa3309e78649f" + }, + "sinkhorn_edges": { + "c": "23eecc3be38a3fb841ae4fd66c6549848eaf61b7c1dfb338f5b09fef5038500f", + "grad.dh": "f073cb7f935783621186e9478edfdf698d65fef59d98bf3c4b19c20f728a53b7", + "post": "73ed0fc6cd099bc2a1a0c4b99b00a3a07be480f40eee7035751db9c73549e45f", + "pre": "b017dde96bbfe73cccf192c8af9ba9eff171354f7ddc7f6d597aab527e098ce3" + } + }, + "fixture_hidden": 128, + "numeric_profile": "oracle-fp32-mhc-v1", + "schema_version": "p1-mhc-layer-v1" +} diff --git a/tests/test_p1_contract.py b/tests/test_p1_contract.py new file mode 100644 index 00000000..581e215c --- /dev/null +++ b/tests/test_p1_contract.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P1 schema, fingerprint, and trace tests (issue #2 contracts).""" + +from __future__ import annotations + +import dataclasses + +import pytest +import torch + +from rl_engine.mhc import fixtures +from rl_engine.mhc.contract import ( + PROD_CONTROLLER_N, + PROD_HIDDEN, + SCHEMA_VERSION, + SINKHORN_ITERS, + LayerContract, + tensor_sha256, +) +from rl_engine.mhc.trace import MHCTrace, first_divergence + + +def test_fixture_batches_validate() -> None: + for name in fixtures.BLOCK_CASES: + batch = fixtures.make_batch(name) + assert batch.contract.schema_version == SCHEMA_VERSION + batch.validate() + fixtures.make_grads(name, batch).validate(batch) + + +def test_production_contract_constants_are_frozen() -> None: + prod = LayerContract() + prod.validate() + prod.assert_production() + assert (prod.hidden, prod.controller_n, prod.flat_k) == (PROD_HIDDEN, PROD_CONTROLLER_N, 16384) + assert prod.hc_mult == 4 and prod.sinkhorn_iters == SINKHORN_ITERS == 20 + assert prod.mhc_eps == prod.rmsnorm_eps == 1e-6 + with pytest.raises(ValueError): + LayerContract(hidden=128).assert_production() + + +def test_frozen_constants_cannot_be_overridden() -> None: + for kwargs in ( + {"hc_mult": 2}, + {"sinkhorn_iters": 10}, + {"mhc_eps": 1e-5}, + {"rmsnorm_eps": 1e-5}, + {"controller_n": 20}, + ): + with pytest.raises(ValueError): + LayerContract(**kwargs).validate() + + +def test_default_contract_is_the_canonical_unfused_mode() -> None: + """Train/infer byte-equality is the goal, so unfused is what you get unless + an engine explicitly declares it cannot expose the intermediate.""" + from rl_engine.mhc.contract import CANONICAL_FUSION_MODE + + assert LayerContract().fusion_mode == CANONICAL_FUSION_MODE == "unfused" + assert LayerContract().trainability == "full" + + +def test_unknown_modes_fail_closed() -> None: + for kwargs in ( + {"fusion_mode": "fused-everything"}, + {"trainability": "everything-trainable"}, + {"placement": "expert-parallel"}, + ): + with pytest.raises(ValueError): + LayerContract(**kwargs).validate() + + +def test_weight_fingerprint_detects_tampering() -> None: + batch = fixtures.make_batch("packed_t16") + batch.validate() + batch.controller.weight[0, 0] += 1.0 # tamper one checkpoint value + with pytest.raises(ValueError, match="fingerprint"): + batch.validate() + + +def test_gamma_tampering_is_also_caught() -> None: + batch = fixtures.make_batch("packed_t16") + batch.norm.gamma[0] = torch.tensor(9.0, dtype=torch.bfloat16) + with pytest.raises(ValueError, match="fingerprint"): + batch.validate() + + +def test_bad_dtypes_and_shapes_fail_closed() -> None: + batch = fixtures.make_batch("packed_t16") + with pytest.raises(TypeError): + dataclasses.replace(batch, r_old=batch.r_old.float()).validate() + with pytest.raises(TypeError): + dataclasses.replace(batch, token_id=batch.token_id.to(torch.int32)).validate() + with pytest.raises(ValueError): + dataclasses.replace(batch, y_sublayer=batch.y_sublayer[:, :-1]).validate() + with pytest.raises(ValueError, match="one-row"): + dataclasses.replace(batch, row_geometry="one-row").validate() + + +def test_controller_params_must_be_fp32() -> None: + batch = fixtures.make_batch("packed_t16") + bad = dataclasses.replace(batch.controller, weight=batch.controller.weight.to(torch.bfloat16)) + with pytest.raises(TypeError, match="FP32"): + bad.validate(batch.contract) + + +def test_grad_boundary_shape_checks() -> None: + batch = fixtures.make_batch("packed_t16") + grads = fixtures.make_grads("packed_t16", batch) + with pytest.raises(ValueError): + dataclasses.replace(grads, d_normalized=grads.d_normalized[:-1]).validate(batch) + with pytest.raises(TypeError): + dataclasses.replace(grads, d_residual=grads.d_residual.float()).validate(batch) + + +def test_batch_serialization_roundtrip(tmp_path) -> None: + batch = fixtures.make_batch("packed_t16") + path = tmp_path / "batch.pt" + torch.save(batch, path) + loaded = torch.load(path, weights_only=False) + loaded.validate() + assert tensor_sha256(loaded.r_old) == tensor_sha256(batch.r_old) + assert loaded.weight_fingerprint == batch.weight_fingerprint + + +def test_contract_fingerprint_moves_with_every_frozen_field() -> None: + base = LayerContract(hidden=128) + for kwargs in ( + {"hidden": 256}, + {"fusion_mode": "fused-pre-norm"}, + {"trainability": "mixer-frozen"}, + ): + assert dataclasses.replace(base, **kwargs).fingerprint() != base.fingerprint() + # layer_index is identity, not arithmetic: it must NOT move the fingerprint. + assert dataclasses.replace(base, layer_index=17).fingerprint() == base.fingerprint() + + +def test_trace_first_divergence() -> None: + a = MHCTrace(numeric_profile="p") + b = MHCTrace(numeric_profile="p") + t1 = torch.arange(4, dtype=torch.float32) + t2 = t1 + 1 + for trace, second in ((a, t1), (b, t2)): + trace.record("s1", t1) + trace.record("s2", second) + assert first_divergence(a, b) == "s2" + b.records[1] = a.records[1] + assert first_divergence(a, b) is None + + +def test_trace_records_arithmetic_provenance() -> None: + from rl_engine.mhc import oracle + + batch = fixtures.make_batch("packed_t16") + trace = MHCTrace(numeric_profile="p") + oracle.mhc_block_forward(batch, trace) + for key in ("reduction_tree", "fma", "rsqrt", "downcast_points", "weight_fingerprint"): + assert key in trace.notes, f"trace is missing required provenance note {key!r}" + assert "rsqrt(mean+eps)" in trace.notes["rsqrt"] + assert trace.to_dict()["records"][0]["name"] == "controller.p" diff --git a/tests/test_p1_oracle.py b/tests/test_p1_oracle.py new file mode 100644 index 00000000..da6c7f3f --- /dev/null +++ b/tests/test_p1_oracle.py @@ -0,0 +1,406 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Self-consistency tests for the P1 FP32 oracle (P1-D1..P1-D6; issue #2). + +Every hand-written backward is cross-checked against a plain-torch autograd +graph built from the same frozen formulas. The autograd graph uses library +reductions on purpose: it validates the *math*, while the byte-level golden +manifest validates the *arithmetic order*. +""" + +from __future__ import annotations + +import dataclasses +import math + +import pytest +import torch + +from rl_engine.mhc import fixtures, oracle +from rl_engine.mhc.contract import LayerContract, ResidualBatch +from rl_engine.mhc.reduction import HC_MULT + +EPS = 1e-6 + + +def _close(got: torch.Tensor, want: torch.Tensor, rel: float = 2e-3) -> bool: + got, want = got.float(), want.float() + return bool((got - want).abs().max() <= rel * want.abs().max() + 1e-5) + + +# A BF16 output carries ~1 ulp = 2^-8 of relative error against an FP32 +# reference, so boundaries that round to BF16 get the looser bound. +BF16_ULP = 1e-2 + + +# --- P1-D1: hc_split_sinkhorn --------------------------------------------- + + +def _sinkhorn_autograd(h: torch.Tensor, iters: int = 20): + """The same frozen schedule, written with library ops, for autograd.""" + pre = torch.sigmoid(h[:, 0:4]) + EPS + post = 2.0 * torch.sigmoid(h[:, 4:8]) + logits = h[:, 8:24].reshape(-1, 4, 4) + m = torch.softmax(logits, dim=2) + EPS + m = m / (m.sum(dim=1, keepdim=True) + EPS) + for _ in range(iters - 1): + m = m / (m.sum(dim=2, keepdim=True) + EPS) + m = m / (m.sum(dim=1, keepdim=True) + EPS) + return pre, post, m + + +def test_sinkhorn_forward_matches_the_literal_schedule() -> None: + g = torch.Generator().manual_seed(11) + h = torch.randn(5, 24, generator=g) + pre, post, c, _ = oracle.hc_split_sinkhorn_fwd(h, LayerContract(hidden=128)) + pre_ref, post_ref, c_ref = _sinkhorn_autograd(h) + assert _close(pre, pre_ref) and _close(post, post_ref) and _close(c, c_ref) + + +def test_sinkhorn_backward_matches_autograd() -> None: + g = torch.Generator().manual_seed(12) + h = torch.randn(5, 24, generator=g).requires_grad_(True) + dpre = torch.randn(5, 4, generator=g) + dpost = torch.randn(5, 4, generator=g) + dc = torch.randn(5, 4, 4, generator=g) + + pre_r, post_r, c_r = _sinkhorn_autograd(h) + (pre_r * dpre).sum().add((post_r * dpost).sum()).add((c_r * dc).sum()).backward() + + _, _, _, saved = oracle.hc_split_sinkhorn_fwd(h.detach(), LayerContract(hidden=128)) + dh = oracle.hc_split_sinkhorn_bwd(dpre, dpost, dc, saved) + assert _close(dh, h.grad) + + +def test_sinkhorn_runs_exactly_20_column_normalizations() -> None: + _, _, _, saved = oracle.hc_split_sinkhorn_fwd(torch.zeros(1, 24), LayerContract(hidden=128)) + kinds = [k for k, _, _ in saved["steps"]] + assert kinds.count("col") == 20 and kinds.count("row") == 19 + assert kinds[0] == "col", "the schedule starts with a column normalize" + + +def test_sinkhorn_row_and_column_order_changes_bytes() -> None: + """Row-first and column-first both converge, so they agree to ~1e-7 -- but + they are not byte-equal. Under a strict-bytes contract that is a failure, + which is exactly why issue #2 forbids swapping the order.""" + g = torch.Generator().manual_seed(13) + h = torch.randn(3, 24, generator=g) + _, _, c, _ = oracle.hc_split_sinkhorn_fwd(h, LayerContract(hidden=128)) + m = torch.softmax(h[:, 8:24].reshape(-1, 4, 4), dim=2) + EPS + m = m / (m.sum(dim=2, keepdim=True) + EPS) # row first: the wrong order + for _ in range(19): + m = m / (m.sum(dim=1, keepdim=True) + EPS) + m = m / (m.sum(dim=2, keepdim=True) + EPS) + assert _close(c, m, rel=1e-3), "both orders converge, so a tolerance test would pass" + assert not torch.equal(c, m), "but the bytes differ -- the order is load-bearing" + + +def test_sinkhorn_eps_guard_is_not_a_clamp() -> None: + """``sum + eps`` and ``clamp(sum, min=eps)`` differ; issue #2 pins ``sum + eps``.""" + h = torch.zeros(1, 24) + _, _, c, _ = oracle.hc_split_sinkhorn_fwd(h, LayerContract(hidden=128)) + m = torch.softmax(h[:, 8:24].reshape(-1, 4, 4), dim=2) + EPS + m = m / torch.clamp(m.sum(dim=1, keepdim=True), min=EPS) + for _ in range(19): + m = m / torch.clamp(m.sum(dim=2, keepdim=True), min=EPS) + m = m / torch.clamp(m.sum(dim=1, keepdim=True), min=EPS) + assert not torch.equal(c, m) + + +def test_pre_and_post_use_their_own_activations() -> None: + h = torch.zeros(1, 24) + pre, post, _, _ = oracle.hc_split_sinkhorn_fwd(h, LayerContract(hidden=128)) + assert _close(pre, torch.full((1, 4), 0.5 + EPS)) + assert _close(post, torch.full((1, 4), 1.0)) + + +# --- P1-D2: fp32_gemm_rms ------------------------------------------------- + + +def test_gemm_rms_forward_and_backward_match_autograd() -> None: + g = torch.Generator().manual_seed(21) + k, n = 64, 24 + x = torch.randn(6, k, generator=g).requires_grad_(True) + w = torch.randn(n, k, generator=g).requires_grad_(True) + dp = torch.randn(6, n, generator=g) + dr = torch.randn(6, generator=g) + + p_ref = x @ w.t() + r_ref = 1.0 / (torch.sqrt((x * x).sum(dim=1)) / math.sqrt(k) + EPS) + (p_ref * dp).sum().add((r_ref * dr).sum()).backward() + + p, r, saved = oracle.fp32_gemm_rms_fwd(x.detach(), w.detach(), EPS) + dx, dw = oracle.fp32_gemm_rms_bwd(dp, dr, x.detach(), w.detach(), saved) + assert _close(p, p_ref) and _close(r, r_ref) + assert _close(dx, x.grad) and _close(dw, w.grad) + + +def test_controller_rms_is_not_the_rmsnorm_formula() -> None: + """``1/(sqrt(mean)+eps)`` vs ``rsqrt(mean+eps)`` -- issue #2 forbids mixing them.""" + x = torch.full((1, 64), 1e-3) + _, r, _ = oracle.fp32_gemm_rms_fwd(x, torch.zeros(24, 64), EPS) + wrong = torch.rsqrt((x * x).mean(dim=1) + EPS) + assert not torch.equal(r, wrong) + + +# --- P1-D3: mhc_post ------------------------------------------------------ + + +def test_mhc_post_forward_and_backward_match_autograd() -> None: + g = torch.Generator().manual_seed(31) + t, d = 5, 32 + r_old = torch.randn(t, 4, d, generator=g).to(torch.bfloat16) + y = torch.randn(t, d, generator=g).to(torch.bfloat16) + c = torch.randn(t, 4, 4, generator=g) + post = torch.rand(t, 4, generator=g) * 2.0 + dr_new = torch.randn(t, 4, d, generator=g) + + r32 = r_old.float().requires_grad_(True) + y32 = y.float().requires_grad_(True) + c_a = c.clone().requires_grad_(True) + post_a = post.clone().requires_grad_(True) + ref = torch.einsum("tij,tid->tjd", c_a, r32) + post_a.unsqueeze(2) * y32.unsqueeze(1) + (ref * dr_new).sum().backward() + + out = oracle.mhc_post_fwd(r_old, y, c, post) + assert _close(out, ref, rel=BF16_ULP) + d_r_old, dy, dc, dpost = oracle.mhc_post_bwd(dr_new, r_old, y, c, post) + assert _close(d_r_old, r32.grad) + assert _close(dy, y32.grad) + assert _close(dc, c_a.grad) + assert _close(dpost, post_a.grad) + + +def test_mhc_post_downcasts_exactly_once() -> None: + out = oracle.mhc_post_fwd( + torch.ones(1, 4, 8, dtype=torch.bfloat16), + torch.ones(1, 8, dtype=torch.bfloat16), + torch.full((1, 4, 4), 0.1), + torch.full((1, 4), 0.3), + ) + assert out.dtype == torch.bfloat16 + # 4*0.1 + 0.3 computed in FP32 then rounded once != rounding each partial. + assert out[0, 0, 0].float() == torch.tensor(0.7).to(torch.bfloat16).float() + + +# --- P1-D4: h_aggregate --------------------------------------------------- + + +def test_h_aggregate_forward_and_backward_match_autograd() -> None: + g = torch.Generator().manual_seed(41) + pre = torch.rand(4, 4, generator=g) + 0.1 + r_old = torch.randn(4, 4, 16, generator=g).to(torch.bfloat16) + dh = torch.randn(4, 16, generator=g) + + pre_a = pre.clone().requires_grad_(True) + r_a = r_old.float().requires_grad_(True) + ref = (pre_a.unsqueeze(2) * r_a).sum(dim=1) + (ref * dh).sum().backward() + + assert _close(oracle.h_aggregate_fwd(pre, r_old), ref, rel=BF16_ULP) + dr, dpre = oracle.h_aggregate_bwd(dh, pre, r_old) + assert _close(dr, r_a.grad) and _close(dpre, pre_a.grad) + + +# --- P1-D5: rmsnorm_residual ---------------------------------------------- + + +def test_rmsnorm_residual_forward_and_backward_match_autograd() -> None: + g = torch.Generator().manual_seed(51) + t, d = 6, 64 + x = torch.randn(t, d, generator=g).to(torch.bfloat16) + gamma = (1.0 + torch.randn(d, generator=g) * 0.1).to(torch.bfloat16) + dy = torch.randn(t, d, generator=g).to(torch.bfloat16) + d_res = torch.randn(t, d, generator=g).to(torch.bfloat16) + + x_a = x.float().requires_grad_(True) + gamma_a = gamma.float().requires_grad_(True) + r_ref = torch.rsqrt((x_a * x_a).mean(dim=1) + EPS) + y_ref = x_a * r_ref.unsqueeze(1) * gamma_a + ((y_ref * dy.float()).sum() + (x_a * d_res.float()).sum()).backward() + + y, residual, saved = oracle.rmsnorm_residual_fwd(x, gamma, EPS) + assert _close(y, y_ref, rel=BF16_ULP) + assert torch.equal(residual, x), "the residual fork keeps the original BF16 bytes" + dx, dgamma = oracle.rmsnorm_residual_bwd(dy, d_res, x, gamma, saved) + assert _close(dx, x_a.grad) + assert _close(dgamma, gamma_a.grad) + + +def test_rmsnorm_is_not_an_add_then_norm() -> None: + """The fork is of the *unnormalized* input; it is not ``x += residual``.""" + x = torch.full((1, 32), 2.0, dtype=torch.bfloat16) + gamma = torch.ones(32, dtype=torch.bfloat16) + _, residual, _ = oracle.rmsnorm_residual_fwd(x, gamma, EPS) + assert torch.equal(residual, x) + + +def test_rmsnorm_zero_row_is_finite_via_eps() -> None: + x = torch.zeros(1, 32, dtype=torch.bfloat16) + gamma = torch.ones(32, dtype=torch.bfloat16) + y, _, saved = oracle.rmsnorm_residual_fwd(x, gamma, EPS) + assert torch.isfinite(y).all() and torch.isfinite(saved["r"]).all() + + +def test_rmsnorm_uses_rsqrt_not_one_over_sqrt() -> None: + """``rsqrt(m+eps)`` and ``1/sqrt(m+eps)`` differ in the last bit on real rows. + + The oracle must take the ``rsqrt`` branch exactly, and the two forms must + genuinely disagree -- otherwise this test would be vacuous. + """ + from rl_engine.mhc.reduction import fixed_sumsq + + g = torch.Generator().manual_seed(52) + x = (torch.randn(256, 128, generator=g) * 3.0).to(torch.bfloat16) + _, _, saved = oracle.rmsnorm_residual_fwd(x, torch.ones(128, dtype=torch.bfloat16), EPS) + m = fixed_sumsq(x.float(), dim=1) / 128.0 + assert torch.equal(saved["r"], torch.rsqrt(m + EPS)) + assert not torch.equal( + torch.rsqrt(m + EPS), 1.0 / torch.sqrt(m + EPS) + ), "the two forms agreed on every row; pick a sharper fixture" + + +# --- P1-D6: fixed-K GEMM reference ---------------------------------------- + + +def test_fixed_k_gemm_backward_matches_autograd() -> None: + g = torch.Generator().manual_seed(61) + x = torch.randn(5, 40, generator=g).requires_grad_(True) + w = torch.randn(7, 40, generator=g).requires_grad_(True) + dy = torch.randn(5, 7, generator=g) + (x @ w.t() * dy).sum().backward() + dx, dw = oracle.fixed_k_gemm_bwd(dy, x.detach(), w.detach()) + assert _close(dx, x.grad) and _close(dw, w.grad) + + +# --- block composition ---------------------------------------------------- + + +def _block_autograd(batch: ResidualBatch, grads): + """Autograd model of the whole P1 block, from the same frozen formulas.""" + d = batch.hidden + r = batch.r_old.float().requires_grad_(True) + y = batch.y_sublayer.float().requires_grad_(True) + w = batch.controller.weight.clone().requires_grad_(True) + a_pre = batch.controller.alpha_pre.clone().requires_grad_(True) + a_post = batch.controller.alpha_post.clone().requires_grad_(True) + a_res = batch.controller.alpha_res.clone().requires_grad_(True) + bias = batch.controller.bias.clone().requires_grad_(True) + gamma = batch.norm.gamma.float().requires_grad_(True) + + x_flat = r.reshape(batch.tokens, batch.contract.flat_k) + p = x_flat @ w.t() + k = batch.contract.flat_k + scale = 1.0 / (torch.sqrt((x_flat * x_flat).sum(dim=1)) / math.sqrt(k) + EPS) + alpha = torch.cat([a_pre.expand(4), a_post.expand(4), a_res.expand(16)], dim=-1) + h = (scale.unsqueeze(1) * p) * alpha + bias + pre, post, c = _sinkhorn_autograd(h) + + hidden = (pre.unsqueeze(2) * r).sum(dim=1).to(torch.bfloat16).float() + rstd = torch.rsqrt((hidden * hidden).mean(dim=1) + EPS) + normalized = hidden * rstd.unsqueeze(1) * gamma + r_new = torch.einsum("tij,tid->tjd", c, r) + post.unsqueeze(2) * y.unsqueeze(1) + + loss = ( + (r_new * grads.d_r_new.float()).sum() + + (normalized * grads.d_normalized.float()).sum() + + (hidden * grads.d_residual.float()).sum() + ) + loss.backward() + del d + return { + "d_r_old": r.grad, + "dy_sublayer": y.grad, + "d_controller_weight": w.grad, + "d_alpha_pre": a_pre.grad, + "d_alpha_post": a_post.grad, + "d_alpha_res": a_res.grad, + "d_bias": bias.grad, + "d_gamma": gamma.grad, + } + + +@pytest.mark.parametrize("case", ["one_row", "packed_t16", "packed_t7_odd"]) +def test_block_backward_matches_autograd(case: str) -> None: + batch = fixtures.make_batch(case) + grads = fixtures.make_grads(case, batch) + _, saved = oracle.mhc_block_forward(batch) + got = oracle.mhc_block_backward(batch, saved, grads) + want = _block_autograd(batch, grads) + for key, ref in want.items(): + assert _close(got[key], ref, rel=3e-2), f"{key} diverges from autograd" + + +def test_fused_equals_unfused_bytes() -> None: + """The fused boundary is *defined* as the unfused composition, so this test + proves the oracle is self-consistent -- it does not prove that any real + fused kernel matches. That obligation lands on the kernel: a fused + implementation is byte-equal only if it keeps the same reduction layout and + the same downcast points. P1-D5 owns proving it for a TE-backed provider. + """ + unfused = fixtures.make_batch("packed_t16") + fused = dataclasses.replace( + unfused, contract=dataclasses.replace(unfused.contract, fusion_mode="fused-pre-norm") + ).sealed() + r_a, saved_a = oracle.mhc_block_forward(unfused) + r_b, saved_b = oracle.mhc_block_forward(fused) + assert torch.equal(r_a, r_b) + assert torch.equal(saved_a["normalized"], saved_b["normalized"]) + assert torch.equal(saved_a["residual"], saved_b["residual"]) + + +def test_mixer_frozen_leaks_no_controller_gradient() -> None: + batch = fixtures.make_batch("mixer_frozen") + grads = fixtures.make_grads("mixer_frozen", batch) + _, saved = oracle.mhc_block_forward(batch) + out = oracle.mhc_block_backward(batch, saved, grads) + for key in ("d_controller_weight", "d_alpha_pre", "d_alpha_post", "d_alpha_res", "d_bias"): + assert out[key] is None, f"stop-grad mixer leaked {key}" + assert out["d_r_old"] is not None and torch.isfinite(out["d_r_old"]).all() + + +def test_unsupported_modes_fail_closed() -> None: + batch = fixtures.make_batch("packed_t16") + bad = dataclasses.replace( + batch, contract=dataclasses.replace(batch.contract, fusion_mode="fused-everything") + ) + with pytest.raises(ValueError, match="fusion_mode"): + bad.validate() + with pytest.raises(NotImplementedError, match="fail-closed"): + oracle._check_modes(bad.contract) + + +# --- invariance (issue #2 acceptance: same row, different batch/pad/stride) - + + +def test_same_row_same_bytes_across_batch_and_padding() -> None: + batch = fixtures.make_batch("packed_t16") + full, _ = oracle.mhc_block_forward(batch) + for start, stop in ((5, 6), (0, 3), (9, 16), (4, 12)): + part, _ = oracle.mhc_block_forward(fixtures.slice_batch(batch, start, stop)) + assert torch.equal(part, full[start:stop]), f"rows {start}:{stop} moved with the batch" + + +def test_non_contiguous_stride_does_not_change_bytes() -> None: + batch = fixtures.make_batch("packed_t16") + want, _ = oracle.mhc_block_forward(batch) + padded_r = torch.zeros(batch.tokens, HC_MULT, batch.hidden * 2, dtype=torch.bfloat16) + padded_r[:, :, ::2] = batch.r_old + padded_y = torch.zeros(batch.tokens, batch.hidden * 2, dtype=torch.bfloat16) + padded_y[:, ::2] = batch.y_sublayer + strided = dataclasses.replace(batch, r_old=padded_r[:, :, ::2], y_sublayer=padded_y[:, ::2]) + assert not strided.r_old.is_contiguous() + got, _ = oracle.mhc_block_forward(strided) + assert torch.equal(got, want) + + +def test_layer_index_and_token_id_do_not_change_arithmetic() -> None: + batch = fixtures.make_batch("packed_t16") + want, _ = oracle.mhc_block_forward(batch) + moved = dataclasses.replace( + batch, + token_id=batch.token_id + 99999, + contract=dataclasses.replace(batch.contract, layer_index=41), + ).sealed() + got, _ = oracle.mhc_block_forward(moved) + assert torch.equal(got, want) diff --git a/tests/test_p1_provider.py b/tests/test_p1_provider.py new file mode 100644 index 00000000..35aadad1 --- /dev/null +++ b/tests/test_p1_provider.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Provider protocol, fail-closed stub, and golden-manifest anchor tests.""" + +from __future__ import annotations + +import dataclasses + +import pytest +import torch + +from rl_engine.mhc import fixtures, oracle +from rl_engine.mhc.contract import LayerContract +from rl_engine.mhc.provider import ( + ReferenceProvider, + StubProvider, + check_capability, + resolve_provider, +) +from rl_engine.mhc.trace import MHCTrace, first_divergence + + +def test_reference_provider_matches_oracle_bytes() -> None: + batch = fixtures.make_batch("packed_t16") + grads = fixtures.make_grads("packed_t16", batch) + gold, cand = MHCTrace("a"), MHCTrace("b") + _, saved_g = oracle.mhc_block_forward(batch, gold) + _, saved_c = oracle.mhc_block_forward(batch, cand, ops=ReferenceProvider()) + assert first_divergence(gold, cand) is None + out_g = oracle.mhc_block_backward(batch, saved_g, grads, gold) + out_c = oracle.mhc_block_backward(batch, saved_c, grads, cand, ops=ReferenceProvider()) + assert first_divergence(gold, cand) is None + for key, grad in out_g.items(): + other = out_c[key] + assert (grad is None) == (other is None) + if grad is not None: + assert torch.equal(grad, other), f"{key} diverged" + + +def test_stub_provider_fails_closed() -> None: + stub = StubProvider() + with pytest.raises(NotImplementedError, match=r"P1-1 \(#14\)"): + stub.hc_split_sinkhorn_fwd(torch.zeros(1, 24)) + # #15 absorbs the fixed-K GEMM reference, so it points at P1-2, not its own task + with pytest.raises(NotImplementedError, match=r"P1-2 \(#15\)"): + stub.fixed_k_gemm_fwd(torch.zeros(1, 4), torch.zeros(2, 4)) + with pytest.raises(NotImplementedError, match=r"P1-5 \(#18\)"): + stub.rmsnorm_residual_fwd(torch.zeros(1, 8), torch.zeros(8), 1e-6) + batch = fixtures.make_batch("one_row") + with pytest.raises(NotImplementedError): + oracle.mhc_block_forward(batch, ops=stub) + + +def test_a_partial_provider_keeps_the_rest_on_the_oracle() -> None: + """The pattern each D1..D6 PR uses: override one op, run full acceptance.""" + + calls: list[str] = [] + + class OnlyPost(ReferenceProvider): + name = "only-post" + numeric_profile = "test-only-post" + + @staticmethod + def mhc_post_fwd(r_old, y, c, post): + calls.append("fwd") + return oracle.mhc_post_fwd(r_old, y, c, post) + + batch = fixtures.make_batch("packed_t16") + gold, cand = MHCTrace("a"), MHCTrace("b") + oracle.mhc_block_forward(batch, gold) + oracle.mhc_block_forward(batch, cand, ops=OnlyPost()) + assert calls == ["fwd"] + assert first_divergence(gold, cand) is None + + +def test_resolve_provider() -> None: + assert resolve_provider("reference").name == "reference" + assert resolve_provider("rl_engine.mhc.provider:StubProvider").name == "stub" + with pytest.raises(ValueError): + resolve_provider("not-a-spec") + prov = resolve_provider("reference").provenance() + assert prov["requested_backend"] == prov["actual_backend"] == "reference" + + +def test_check_capability_fails_closed_on_unsupported_placement() -> None: + provider = ReferenceProvider() + check_capability(provider, LayerContract(hidden=128)) + sharded = dataclasses.replace(LayerContract(hidden=128), placement="tp-sharded") + with pytest.raises(NotImplementedError, match="fail-closed"): + check_capability(provider, sharded) + + +def test_reference_provider_declares_every_supported_mode() -> None: + caps = ReferenceProvider().capabilities() + assert set(caps["fusion_modes"]) == set(oracle.SUPPORTED_FUSION) + assert set(caps["trainability"]) == set(oracle.SUPPORTED_TRAINABILITY) + + +def test_golden_manifest_anchor() -> None: + """CI anchor: regenerated golden hashes must match the committed manifest. + + A failure here means the oracle's bytes drifted (a torch RNG/libm change, + or an intentional contract change) -- regenerate with + ``python -m rl_engine.mhc.fixtures --write-manifest`` and review the diff. + """ + assert fixtures.load_manifest() == fixtures.golden_manifest() diff --git a/tests/test_p1_reduction.py b/tests/test_p1_reduction.py new file mode 100644 index 00000000..9e6d03db --- /dev/null +++ b/tests/test_p1_reduction.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""The pinned reduction trees (issue #2: every P1 reduction goes through these).""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.mhc.reduction import ( + fixed_dot, + fixed_sum, + fixed_sumsq, + stream4_max, + stream4_sum, + stream4_sum_dim, +) + + +def test_fixed_sum_is_the_ascending_left_fold() -> None: + x = torch.tensor([[1e8, 1.0, -1e8, 1.0]], dtype=torch.float32) + acc = torch.zeros(1, dtype=torch.float32) + for i in range(4): + acc = acc + x[:, i] + assert torch.equal(fixed_sum(x, dim=1), acc) + + +def test_stream4_sum_is_the_balanced_tree_not_the_left_fold() -> None: + """``(a0+a1)+(a2+a3)`` and the left fold disagree in FP32 -- that is the point.""" + a = [torch.tensor([v], dtype=torch.float32) for v in (1e-10, 1.0, -1.0, 1e-10)] + balanced = (a[0] + a[1]) + (a[2] + a[3]) + left_fold = fixed_sum(torch.stack(a, dim=1), dim=1) + assert torch.equal(stream4_sum(a), balanced) + assert not torch.equal(balanced, left_fold) + + +def test_stream4_sum_dim_matches_stream4_sum() -> None: + x = torch.randn(3, 4, 5, generator=torch.Generator().manual_seed(1)) + got = stream4_sum_dim(x, dim=1) + want = stream4_sum([x[:, 0], x[:, 1], x[:, 2], x[:, 3]]) + assert torch.equal(got, want) + + +def test_fixed_dot_matches_a_manual_ascending_k_loop() -> None: + g = torch.Generator().manual_seed(2) + a = torch.randn(3, 9, generator=g) + b = torch.randn(5, 9, generator=g) + acc = torch.zeros(3, 5) + for k in range(9): + acc = acc + a[:, k].unsqueeze(1) * b[:, k].unsqueeze(0) + assert torch.equal(fixed_dot(a, b), acc) + + +def test_fixed_dot_is_split_k_free() -> None: + """A two-half Split-K merge changes bytes, so the reference must not do it.""" + g = torch.Generator().manual_seed(5) + a = torch.randn(2, 64, generator=g) * 1e4 + b = torch.randn(2, 64, generator=g) * 1e-4 + split = fixed_dot(a[:, :32], b[:, :32]) + fixed_dot(a[:, 32:], b[:, 32:]) + assert not torch.equal(fixed_dot(a, b), split) + + +def test_fixed_sumsq_rounds_the_square_before_accumulating() -> None: + x = torch.tensor([[1.0000001, 2.0, 3.0]], dtype=torch.float32) + acc = torch.zeros(1, dtype=torch.float32) + for i in range(3): + acc = acc + x[:, i] * x[:, i] + assert torch.equal(fixed_sumsq(x, dim=1), acc) + + +def test_stream4_max_is_the_balanced_tree() -> None: + x = torch.tensor([[[1.0, 7.0, 3.0, 5.0]]]) + assert torch.equal(stream4_max(x, dim=2), torch.tensor([[7.0]])) + + +def test_wrong_arity_fails_closed() -> None: + with pytest.raises(ValueError): + stream4_sum([torch.zeros(1)] * 3) + with pytest.raises(ValueError): + stream4_sum_dim(torch.zeros(2, 5), dim=1) + with pytest.raises(ValueError): + fixed_dot(torch.zeros(2, 4), torch.zeros(3, 5))