From e6373c23e2b7125b4c750319fb5a25bd84516d5e Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Thu, 3 Sep 2026 10:03:23 +0000 Subject: [PATCH 1/4] [DSv4][P1-S0] Start kit for the P1 work package (mHC + RMSNorm) Freezes the contract, the reference answers and the checker for the P1 work package so all six sub-tasks (P1-D1..P1-D6) can start in parallel against the same frozen math and the same golden bytes. No GPU kernels. - rl_engine/mhc/reduction.py: the two pinned reduction trees (serial ascending left fold for long reductions, (a0+a1)+(a2+a3) for the four streams). Every P1 accumulation goes through here. - rl_engine/mhc/contract.py: LayerContract, ResidualBatch, controller and norm params, GradBoundary, checkpoint fingerprints. - rl_engine/mhc/oracle.py: FP32 reference for hc_split_sinkhorn, fp32_gemm_rms, mhc_post, mhc_pre/h_aggregate, rmsnorm_residual and the fixed-K GEMM, forward and backward, plus the block composition. - rl_engine/mhc/provider.py: MHCProvider protocol, oracle-backed reference, fail-closed stub, capability gate. - rl_engine/mhc/fixtures.py + tests/fixtures/p1/golden_hashes.json: seeded cases and the CI anchor for the golden bytes. - scripts/check_p1.py: the acceptance command. - 53 CPU tests; every hand-written backward is cross-checked against an autograd graph built from the same frozen formulas. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Kcu8vXve5YMZ5wqCCKRfvJ --- docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md | 162 +++++ rl_engine/mhc/__init__.py | 74 ++ rl_engine/mhc/contract.py | 335 ++++++++++ rl_engine/mhc/fixtures.py | 238 +++++++ rl_engine/mhc/oracle.py | 667 +++++++++++++++++++ rl_engine/mhc/provider.py | 262 ++++++++ rl_engine/mhc/reduction.py | 135 ++++ rl_engine/mhc/trace.py | 82 +++ scripts/check_p1.py | 160 +++++ tests/fixtures/p1/golden_hashes.json | 176 +++++ tests/test_p1_contract.py | 152 +++++ tests/test_p1_oracle.py | 397 +++++++++++ tests/test_p1_provider.py | 103 +++ tests/test_p1_reduction.py | 82 +++ 14 files changed, 3025 insertions(+) create mode 100644 docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md create mode 100644 rl_engine/mhc/__init__.py create mode 100644 rl_engine/mhc/contract.py create mode 100644 rl_engine/mhc/fixtures.py create mode 100644 rl_engine/mhc/oracle.py create mode 100644 rl_engine/mhc/provider.py create mode 100644 rl_engine/mhc/reduction.py create mode 100644 rl_engine/mhc/trace.py create mode 100755 scripts/check_p1.py create mode 100644 tests/fixtures/p1/golden_hashes.json create mode 100644 tests/test_p1_contract.py create mode 100644 tests/test_p1_oracle.py create mode 100644 tests/test_p1_provider.py create mode 100644 tests/test_p1_reduction.py 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..0deb5dc6 --- /dev/null +++ b/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md @@ -0,0 +1,162 @@ +# 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). + +## Sub-issue naming + +| Label | Scope | +| --- | --- | +| `P1-S0` | This start kit (contract, oracle, fixtures, acceptance command) | +| `P1-D1` | `hc_split_sinkhorn` (fwd + bwd through all 20 Sinkhorn rounds) | +| `P1-D2` | `fp32_gemm_rms` (controller projection + controller RMS, fwd + bwd) | +| `P1-D3` | `mhc_post` (fwd + `dR_old`/`dy`/`dC`/`dPOST`) | +| `P1-D4` | `mhc_pre` / `h_aggregate` + the TE/Megatron/Miles provider adapter | +| `P1-D5` | `rmsnorm_residual` (fwd + `dX`/`dGamma`) | +| `P1-D6` | Fixed-K / batch-invariant GEMM reference + equivalence harness | +| `P1-R0` | Review, package the P1 provider artifact, GPU CI | + +## 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. + +## 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)`.** The issue's forward text says `norm/sqrt(16384)` while its backward text says `q = sqrt(s/N)`; those differ in FP32. The forward wording wins, and backward reuses the saved `q`. | The two forms are not bit-equal; the forward is the authoritative statement. **Please confirm in review.** | +| 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` and `bias` are FP32 `[24]` vectors**, applied per controller output. | The general case; a scalar `alpha` is the broadcast special case and stays representable. | +| 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 | **The fused pre+norm boundary is defined as the unfused composition**, byte for byte. | Miles/XoRL fuse pre-mix and normalize into one launch. Defining the fused boundary this way turns "fused equals unfused" into a test rather than an assumption; a kernel whose fused residual store changes the reduction layout must register a different profile instead of presenting itself as the same kernel. | +| 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 + +- **D1** — `sqrt(s)/sqrt(K)` vs `sqrt(s/K)` in the controller RMS. The issue + text says both; the kit takes the forward wording. One line of Miles/Megatron + source settles it. +- **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. The kit implements the schedule exactly as written in + the issue; if Miles differs, D2/D4 change and the manifest is regenerated. +- **D8** — whether Miles can expose the pre-fusion intermediate at all. If it + cannot, `fused-pre-norm` becomes the only boundary for the inference side + and `unfused` stays the training-side definition, with the equivalence case + as the bridge. + +## 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..cc3bacba --- /dev/null +++ b/rl_engine/mhc/contract.py @@ -0,0 +1,335 @@ +# 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). +FUSION_MODES = ("unfused", "fused-pre-norm") +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: torch.Tensor # FP32 [controller_n] + bias: torch.Tensor # FP32 [controller_n] + + def validate(self, contract: LayerContract) -> None: + for name, t in (("weight", self.weight), ("alpha", self.alpha), ("bias", self.bias)): + 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.alpha.shape) != (n,) or tuple(self.bias.shape) != (n,): + raise ValueError("controller alpha/bias must have shape [controller_n]") + + def fingerprint(self) -> str: + h = hashlib.sha256() + for t in (self.weight, self.alpha, 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.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__ = [ + "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..9af29cb7 --- /dev/null +++ b/rl_engine/mhc/fixtures.py @@ -0,0 +1,238 @@ +# 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=_randn(g, n, 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..d8be84e7 --- /dev/null +++ b/rl_engine/mhc/oracle.py @@ -0,0 +1,667 @@ +# 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``.""" + 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) + h, m = _controller_affine(p, r, batch.controller.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.alpha + m, p, r = saved["m"], saved["p"], saved["r"] + dbias = fixed_sum(dh, dim=0) + dalpha = fixed_sum(dh * m, dim=0) + 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": None if frozen else dalpha, + "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. + """ + 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_grads["d_alpha"], + "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..8fd9eb30 --- /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-D1 (hc_split_sinkhorn)") + + def hc_split_sinkhorn_bwd(self, dpre, dpost, dc, saved): + raise self._todo("P1-D1 (hc_split_sinkhorn)") + + def fp32_gemm_rms_fwd(self, x_flat, weight, eps): + raise self._todo("P1-D2 (fp32_gemm_rms)") + + def fp32_gemm_rms_bwd(self, dp, dr, x_flat, weight, saved): + raise self._todo("P1-D2 (fp32_gemm_rms)") + + def mhc_post_fwd(self, r_old, y, c, post): + raise self._todo("P1-D3 (mhc_post)") + + def mhc_post_bwd(self, dr_new, r_old, y, c, post): + raise self._todo("P1-D3 (mhc_post)") + + def h_aggregate_fwd(self, pre, r_old): + raise self._todo("P1-D4 (mhc_pre / h_aggregate)") + + def h_aggregate_bwd(self, dh, pre, r_old): + raise self._todo("P1-D4 (mhc_pre / h_aggregate)") + + def rmsnorm_residual_fwd(self, x, gamma, eps): + raise self._todo("P1-D5 (rmsnorm_residual)") + + def rmsnorm_residual_bwd(self, dy, d_residual, x, gamma, saved): + raise self._todo("P1-D5 (rmsnorm_residual)") + + def fixed_k_gemm_fwd(self, x, w): + raise self._todo("P1-D6 (fixed-K GEMM reference)") + + def fixed_k_gemm_bwd(self, dy, x, w): + raise self._todo("P1-D6 (fixed-K GEMM reference)") + + +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..c424ca93 --- /dev/null +++ b/tests/fixtures/p1/golden_hashes.json @@ -0,0 +1,176 @@ +{ + "cases": { + "fused_pre_norm": { + "bwd.d_c": "a0d671287e9d89450271cb9c6d1b75dc8ac1b1bef24a6fddfba6175b71336b92", + "bwd.d_gamma": "2cbb9bf1d2a76c10da22e70c29df5c8ad68122a9edb9b3d7dbccd64b8fd13e5e", + "bwd.d_hidden": "7b9f76e104842bf09fdcd44f11800a510bbafb000b4373d329e10da783ce9a49", + "bwd.d_post": "aa77b606f48bbc6d1a3b3e5bbc85b01fd8c531aa8986b018a1905c705fb50608", + "bwd.d_r_old": "d9c23ad0c28d6f552f85bac27a8bb562b6eafe146badd710f2b63caa9f4f056e", + "bwd.d_stream0": "5b04030721cc487f80a5030c228db590e039474866c2328683555209e0944691", + "bwd.d_stream1": "01b83afb0d71ea384a24ccfa22c5e3fbb7f4060f21c71072a601f994ca67d64a", + "bwd.d_stream2": "455dceed73dd8efef3a69d49ad5469df893bb97f5a3e51c0bcd2302616b7ae42", + "bwd.d_stream3": "e699c4e85a1e0611b2fb1b35c2e44215cd1d3b26c4af38c7cc8dded4e1043888", + "bwd.dy_sublayer": "e89fca5228e6653ca20cb99bb51ee99fcfda4e76f128a51e78a5980f3a66f49a", + "controller.h": "bbe5fe0aea57a63f641bf93aecd96941abcc066b9fa3ed3c347cd3b56b98dd6c", + "controller.p": "73a3898bd58ab559879d3258686805bcd506817e29cff96b61cada200a55efbb", + "controller.r": "952f64296d063062e522f9d629ca6622f09aea37a498e35ea2e8b65d70a765e0", + "grad.d_alpha": "d86ffd43eeba3f7a9ceaac17a72b186682da9125e4d920a62133f7d73d1a82ee", + "grad.d_bias": "340d1d8382680bf98296c2a762a0a68af87c69a54507208e4a6b6c2027654a38", + "grad.d_c": "a0d671287e9d89450271cb9c6d1b75dc8ac1b1bef24a6fddfba6175b71336b92", + "grad.d_controller_weight": "0559a0e8022a70b2aa92e57b823d4a9a962814b540e253766de47d4c08e502f7", + "grad.d_gamma": "2cbb9bf1d2a76c10da22e70c29df5c8ad68122a9edb9b3d7dbccd64b8fd13e5e", + "grad.d_hidden": "7b9f76e104842bf09fdcd44f11800a510bbafb000b4373d329e10da783ce9a49", + "grad.d_post": "aa77b606f48bbc6d1a3b3e5bbc85b01fd8c531aa8986b018a1905c705fb50608", + "grad.d_r_old": "d9c23ad0c28d6f552f85bac27a8bb562b6eafe146badd710f2b63caa9f4f056e", + "grad.dy_sublayer": "e89fca5228e6653ca20cb99bb51ee99fcfda4e76f128a51e78a5980f3a66f49a", + "norm.normalized": "959d73eec2445fedf7721bf6a0a4e82c9021178b4d7c3e6913c33c9c6b3df09d", + "norm.residual": "28f5695c110eba5413488a4fb913a23d2c995aa38c7247e984a20ab8893b79b3", + "post.r_new": "c55b182a847ba286c421a6c00fd19038876ffc90044c66746cd4c094f2bee227", + "pre.hidden": "28f5695c110eba5413488a4fb913a23d2c995aa38c7247e984a20ab8893b79b3", + "r_new": "c55b182a847ba286c421a6c00fd19038876ffc90044c66746cd4c094f2bee227", + "split.c": "b5d09915a9948a4bb02a447bafa3bc21f5c3189b368c8b8e8e7f2519371d6757", + "split.post": "56a90fb2cb231d05e98c411028178234465b01b05a304e192b1464ba8efa99bd", + "split.pre": "3cf24b388aa6f03d6a48802ddcf3a3908bc63731f6b633742c0f7820279f0b2c" + }, + "mixer_frozen": { + "bwd.d_c": "644b562d8446861f9a703af33ca08a56d6935f82a24435cfdaf269c505a487cf", + "bwd.d_gamma": "4c79126d8ce2bbdec2da09959f87bd37d3c57bb6cea07002bfbe3f930b974b56", + "bwd.d_hidden": "24925df10fcedc74b03f7b9a03b8441cfd9bf81f728af397bb6a1689f392029d", + "bwd.d_post": "00822250814f728cfae3e061462cbc24a66dc50070e67c4e1bb3ae21df714563", + "bwd.d_r_old": "5d3b5882f78956dc9c4f2499508f147a5ba2278f49a438c7dd963559f7ef3ed4", + "bwd.d_stream0": "22616b202c8328551460c3610515b5e66317d98513485f7fc54eb6ff70a5dae2", + "bwd.d_stream1": "f950507bffd515220ffb58f68d497011c48da55768c2ace0557ca8f258538056", + "bwd.d_stream2": "082bf3086e17a938d24ea63140fdd9caac39362eee0524160c0a8c59b77cfa64", + "bwd.d_stream3": "f318a6ee434a4ff250b493d0d65f440b89e8e514c765b7003128d2d7e53f01c8", + "bwd.dy_sublayer": "71f264ef3ef1f4a1a3b47e1c695a5816b86e4eead38ae4dce000436f914437bf", + "controller.h": "6779debc9da9633da1a9a752adf9047e9b25d3e0696b31e3065856d3d1edcfc5", + "controller.p": "5e63ce025161fb5b9197c0f9655e5964d37291576f46e50ff75711a11e422cf6", + "controller.r": "e3184f50bbeb6887ce5f3785cd8d036b31d6619c411a91ba09abe46da630ad17", + "grad.d_c": "644b562d8446861f9a703af33ca08a56d6935f82a24435cfdaf269c505a487cf", + "grad.d_gamma": "4c79126d8ce2bbdec2da09959f87bd37d3c57bb6cea07002bfbe3f930b974b56", + "grad.d_hidden": "24925df10fcedc74b03f7b9a03b8441cfd9bf81f728af397bb6a1689f392029d", + "grad.d_post": "00822250814f728cfae3e061462cbc24a66dc50070e67c4e1bb3ae21df714563", + "grad.d_r_old": "5d3b5882f78956dc9c4f2499508f147a5ba2278f49a438c7dd963559f7ef3ed4", + "grad.dy_sublayer": "71f264ef3ef1f4a1a3b47e1c695a5816b86e4eead38ae4dce000436f914437bf", + "norm.normalized": "84bee319934eef7ec72a4269e7a66228b7b598bb7415e49774dc43a2e1ddace4", + "norm.residual": "e6ea03227ea97d857952435311c5eb8a5aa601d556e098535a2c15856efc732d", + "post.r_new": "63531fa1253c4c9a95b16587c47bb46bde499f22b9981f6de44a669dd1c2bbec", + "pre.hidden": "e6ea03227ea97d857952435311c5eb8a5aa601d556e098535a2c15856efc732d", + "r_new": "63531fa1253c4c9a95b16587c47bb46bde499f22b9981f6de44a669dd1c2bbec", + "split.c": "e3bc8f726b626e2e9fd704724d2f715b683f8e9f573cc0052876cacce31bd4b7", + "split.post": "8dc421dd34664d5906ec0770b11aa070e88f8cf8c258d55a5d25dbcb927f5d51", + "split.pre": "6c70243dea814984cb9b9fa91b78c55f785a677639223009bb7a466142932cbb" + }, + "one_row": { + "bwd.d_c": "2bf1db541dedf0efb695b6a4d2c4381740f4f0eb42a647edd7ceba8a4e817ebd", + "bwd.d_gamma": "6446e0a79a7f572645d2035edbdc8ba2ca80c0c8e1c2a73ca1a9f60ad8c56e4f", + "bwd.d_hidden": "5bb71358cf77a95e40bdd20581f0dd947fc26c5f207243cd32848169b25160fa", + "bwd.d_post": "1330c48fdaa948c9e111f503fce0e7bb058b1d79ff5811da6f9496c43c2afe8f", + "bwd.d_r_old": "c417f939f9b45fe95beacef07ec4f7f16de5871172866b750524b911bd012f2e", + "bwd.d_stream0": "1f3aaec02c8efb2dbbe2dc7b5209d23c8fab2ac7de59d8cbfb750ba1762f0e4d", + "bwd.d_stream1": "65f47151b2b6b35b3da54cc863d8fe07ffde4bf715b2c945739f2d1295bf6e86", + "bwd.d_stream2": "ae86ccefff2fc2a41b5a39ef87d1dc2f5c808b8161b2a972251c920878910c0c", + "bwd.d_stream3": "d05a20c9800f67804627c1fccfadf35c280145e5334e84369af6ea38400422ca", + "bwd.dy_sublayer": "d712cddba241d525fe30b12110ae83844d27a8f7f27af58a5c9d379de033c401", + "controller.h": "62b5eb67c84659e31b4f8ad98963ad85370353506dc13c2da64778a6ec7bfcdd", + "controller.p": "6d2da237edf96ff10377dcf764806dac926b63b868e7bf2764e47ffc168953ec", + "controller.r": "0ed8fee556e17362f225ff27d1c82a31a62b81a78b0a71b171ed109b298e2ba2", + "grad.d_alpha": "08af3310b0057e8d6719fed3310325c4069f9bc13faa18b7f653acd3d66a977e", + "grad.d_bias": "0a4891f9d439e8bc6f8e0eecd4a79ed3eb56dbf26b30e16013b6d9a5fe6bc6af", + "grad.d_c": "2bf1db541dedf0efb695b6a4d2c4381740f4f0eb42a647edd7ceba8a4e817ebd", + "grad.d_controller_weight": "f98d0f49f04752dcbffac7a355775668ee7d404ede7aae90b6936d93c4fff0c3", + "grad.d_gamma": "6446e0a79a7f572645d2035edbdc8ba2ca80c0c8e1c2a73ca1a9f60ad8c56e4f", + "grad.d_hidden": "5bb71358cf77a95e40bdd20581f0dd947fc26c5f207243cd32848169b25160fa", + "grad.d_post": "1330c48fdaa948c9e111f503fce0e7bb058b1d79ff5811da6f9496c43c2afe8f", + "grad.d_r_old": "c417f939f9b45fe95beacef07ec4f7f16de5871172866b750524b911bd012f2e", + "grad.dy_sublayer": "d712cddba241d525fe30b12110ae83844d27a8f7f27af58a5c9d379de033c401", + "norm.normalized": "daf839f63a03952403fa271d5908aa0c8113c98d061b03fbd4085dc8d38d226b", + "norm.residual": "e241696650c290471579a96c4fad24c49041696f2ddce026493f231dece1b43d", + "post.r_new": "123c2f784e6052882217288aceaa547231f3fce750c0859122e328b6a46ce6b8", + "pre.hidden": "e241696650c290471579a96c4fad24c49041696f2ddce026493f231dece1b43d", + "r_new": "123c2f784e6052882217288aceaa547231f3fce750c0859122e328b6a46ce6b8", + "split.c": "06fa72314507fe0e9413fbff2937c4f76f28c26968df30fce312236468601ac0", + "split.post": "fec5506651ed1406f5e526c081611b0ce23216beda926b23e5e9d74846169159", + "split.pre": "2511d4f9a82f34724c2822a326a20ea8e28600568f13047a822fff191f4a077d" + }, + "packed_t16": { + "bwd.d_c": "3f573a2a506bd2669ab318c84adc88c06efaa0274bcdd8117013555283c23ec9", + "bwd.d_gamma": "87ce69b941c0eed1475bcfaf608f23ac24e7c6582367d01b78e417fac036c103", + "bwd.d_hidden": "cc910a66cadf396885f5082c8969f7cc207910daa058d855dbbf1a48fd0d403b", + "bwd.d_post": "edb376c01993aa5fff0e98713f1ab16a6aed8e94b941d1e7547da40eb2b556c4", + "bwd.d_r_old": "645bc398d44ab31d92e04b173b1d68b91c5b7e82b77832f6941767118f93b1c3", + "bwd.d_stream0": "3b96ad5281ead49c0e596e4c26a1bf033266e8bbb556c6d7cc37d3699d37b7b7", + "bwd.d_stream1": "1fd810234d6587e2d7021ae030047b8c9d58f6f0d9ec9b76394b9bb391b80cf0", + "bwd.d_stream2": "6aa0159f5253249184ba9de35f3bd0a47a5b164180858e74822ca702c9c91970", + "bwd.d_stream3": "3433b559730345dc3d38d78837373cdba4cc3f63eefcc1b141cb1f8f6fdd6ad3", + "bwd.dy_sublayer": "e8028379f260744aee3e21889f1c955f38c363c063d267663c8103e1950dd502", + "controller.h": "2da829cd236c4c268a36165f2a796f650eeb86cfab3d2a70a9ca445ef861e0ad", + "controller.p": "e423cb667df95f122c705682d22d33238935166156552c50c8ce24d48caaf43c", + "controller.r": "ebc5d43e3fda0a178bfe9131405dfd2ea16b6143daedb0fcb44df72a5b6beb9f", + "grad.d_alpha": "74df4fe6c987cabe231706935ba0785f5001f6b78a9831664e93bae6821024f4", + "grad.d_bias": "6f4afe52afe88385ba02456cdaf5c2f71f6eae53dda737514f2e1abe31f40b3c", + "grad.d_c": "3f573a2a506bd2669ab318c84adc88c06efaa0274bcdd8117013555283c23ec9", + "grad.d_controller_weight": "3b9bcce8e798e1287f2a260409ceee3f3c610515743acf705d4dfda145484272", + "grad.d_gamma": "87ce69b941c0eed1475bcfaf608f23ac24e7c6582367d01b78e417fac036c103", + "grad.d_hidden": "cc910a66cadf396885f5082c8969f7cc207910daa058d855dbbf1a48fd0d403b", + "grad.d_post": "edb376c01993aa5fff0e98713f1ab16a6aed8e94b941d1e7547da40eb2b556c4", + "grad.d_r_old": "645bc398d44ab31d92e04b173b1d68b91c5b7e82b77832f6941767118f93b1c3", + "grad.dy_sublayer": "e8028379f260744aee3e21889f1c955f38c363c063d267663c8103e1950dd502", + "norm.normalized": "cd4ca1d9249a3987c285f0e4ec75a163d2c4bf1d16c158d01617c40a203374ba", + "norm.residual": "1de7fd9c90975c7a31c8b04921737d1494ef33a2e79f27fdedcf93f022414a99", + "post.r_new": "619463014e0852ab8bfa105ce8734b2f764162d2ab839a3668bffeabb434ca5a", + "pre.hidden": "1de7fd9c90975c7a31c8b04921737d1494ef33a2e79f27fdedcf93f022414a99", + "r_new": "619463014e0852ab8bfa105ce8734b2f764162d2ab839a3668bffeabb434ca5a", + "split.c": "12b25bd8db4ee0dba6fca2a012b36b58b01b822f554e5ffaa23e651cd127e3f6", + "split.post": "075e73f84f62aa372cb030365b046ca35f11b52a19d9691acc633c48e6bd0915", + "split.pre": "5c9f5b1f8fcc8f3a8f2019dfdb8394498769b14692689ff237ef77ae4ec60086" + }, + "packed_t7_odd": { + "bwd.d_c": "3d7cc5332f9349dac8e47efc9f0360d56e528b41919133beaae86fc3421380ac", + "bwd.d_gamma": "1d6a9bd5ea4f28e15486f05c23b91edf6fce9a5bb2be4f2b89acf8c3d91d8bdf", + "bwd.d_hidden": "76d1b779a3fe8c0a89f2b8506360d82cd0b574c8f3c3e8e09ee28cacf353550b", + "bwd.d_post": "049df2a605804eee375c008e2807d47be65c160927d49eed8d4b9c5963adf055", + "bwd.d_r_old": "5cc94c4ef9a303c0f603d3db0987d8c3ef8e23ae3be265305688bff9d63a0aad", + "bwd.d_stream0": "a113bf7982eb3df270365ad43e9b76988503038de9e3fc6e3fd6c1f64571968a", + "bwd.d_stream1": "87c978dd759a4dae8a627c725395f878bc5c269be3785d9cf6e49652f655d9f4", + "bwd.d_stream2": "96af2de9fe231c2ca75fb2cef9c5d8d650e36766bbffb70da8013df5c4285459", + "bwd.d_stream3": "58680ebc9b0c9c043bd73262a5d81432cade3073b24ae45b6b1cb2057acb9c66", + "bwd.dy_sublayer": "449aa7deb18ff1c70c2e4d89628a1759814693c9eb3e66e0974f7c7d75a59bb2", + "controller.h": "1a82827e4fc25f93405a1b862425cea2a9c36e66f6839dcb7b3b4843119a845c", + "controller.p": "3c568c36bd3a848bca0499afebb4e88350bb9095a63d2159092e6e8fee6cf777", + "controller.r": "1e0aebf1c1829bcc5b4fcfd08d7d7ccda679d2542e8d69fbaadd2417072156e6", + "grad.d_alpha": "0149d892249bae065b3cffe9d13c82ddf39b55927ff14e85496a2227f73b1a9e", + "grad.d_bias": "e4783ae2ccdaaedf845d4549ca0d1492e7363c2b94e27e76c80c59091217a798", + "grad.d_c": "3d7cc5332f9349dac8e47efc9f0360d56e528b41919133beaae86fc3421380ac", + "grad.d_controller_weight": "6844879775e886c24ddde3382898bd11ae4b043f37447d33eb4869e1c0a38fec", + "grad.d_gamma": "1d6a9bd5ea4f28e15486f05c23b91edf6fce9a5bb2be4f2b89acf8c3d91d8bdf", + "grad.d_hidden": "76d1b779a3fe8c0a89f2b8506360d82cd0b574c8f3c3e8e09ee28cacf353550b", + "grad.d_post": "049df2a605804eee375c008e2807d47be65c160927d49eed8d4b9c5963adf055", + "grad.d_r_old": "5cc94c4ef9a303c0f603d3db0987d8c3ef8e23ae3be265305688bff9d63a0aad", + "grad.dy_sublayer": "449aa7deb18ff1c70c2e4d89628a1759814693c9eb3e66e0974f7c7d75a59bb2", + "norm.normalized": "e8bd67df4e08da06948e6bb01bc487c2fccd27d9ff95a4e1e16bd85636f701ad", + "norm.residual": "eec3f0342ddf7353b47af652ed258595df38177741308c27211efb2b54a0a0df", + "post.r_new": "aca7cdeab1fc9670a944ea737986a439945f888ee8b4e396e0d370f750682840", + "pre.hidden": "eec3f0342ddf7353b47af652ed258595df38177741308c27211efb2b54a0a0df", + "r_new": "aca7cdeab1fc9670a944ea737986a439945f888ee8b4e396e0d370f750682840", + "split.c": "5b0d44b39615bc28c7c9137e2d8694747decf28c514e80bfa52882ddb2278b84", + "split.post": "d7dafd8501437d3106cd98b94c7f57b0184f266413d966f8bdfc027c802e1fcd", + "split.pre": "fb06566020b4ce8a069e189f373efd5280505ac48fd0d72c2dfeaadb97d07b28" + }, + "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..025f676b --- /dev/null +++ b/tests/test_p1_contract.py @@ -0,0 +1,152 @@ +# 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_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..2414e3bb --- /dev/null +++ b/tests/test_p1_oracle.py @@ -0,0 +1,397 @@ +# 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) + alpha = batch.controller.alpha.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) + 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": alpha.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: + """Miles fuses pre-mix and normalize into one launch; the kit's fused + boundary must be the same arithmetic, not merely the same math.""" + 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", "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..d646f55b --- /dev/null +++ b/tests/test_p1_provider.py @@ -0,0 +1,103 @@ +# 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="P1-D1"): + stub.hc_split_sinkhorn_fwd(torch.zeros(1, 24)) + with pytest.raises(NotImplementedError, match="P1-D6"): + stub.fixed_k_gemm_fwd(torch.zeros(1, 4), torch.zeros(2, 4)) + 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)) From 3e836d467bfec70326e6ba784f75d408d622dd2b Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Thu, 3 Sep 2026 10:55:22 +0000 Subject: [PATCH 2/4] [DSv4][P1-S0] Pin unfused as canonical; fix alpha to Megatron's 3 scalars Follow-up after reading the Megatron mHC source (megatron/core/transformer/hyper_connection.py and megatron/core/fusions/fused_mhc_kernels.py, dev). Megatron's native and fused paths disagree with each other: the fused kernel uses TF32 MMA, split_k=16 at K>=16384, and runtime autotuning, and computes sqrt(s/K) where native computes sqrt(s)/sqrt(K). Three of issue #2's explicit bans (Split-K off, TF32 off, fixed reduction tree) are violated by that path at exactly the production shape. So Megatron supplies the formulas and constants for this contract, never the reduction order. Changes: - fusion: 'unfused' is now CANONICAL and documented as such. Train/infer byte-equality needs every boundary hashable on its own, which only the unfused decomposition gives. v1 does not reuse TE's TEFusedResidualRMSNorm -- it refuses to expose the intermediate (raises on any forward hook), collapsing two boundaries into one. Fusion stays legal; changing the reduction layout or a downcast point does not. - alpha: Megatron holds three learnable scalars (alpha_pre/post/res) broadcast over the PRE/POST/COMB segments, not 24 independent gains. Forward was already equivalent; backward was not, so dAlpha is now three scalars with the segment fold pinned like every other reduction. - design doc: new section on what Megatron actually does and why it is not the byte reference; D1/D6/D8 rewritten with the source evidence; open questions updated (Miles must also run unfused -- needs their sign-off). - the fused/unfused test now says what it actually proves (oracle self-consistency, not that a real fused kernel matches). 54 tests pass; check_p1.py PASSes on cpu and cuda, stub still fails closed; ruff/flake8/black/isort/mypy clean. Golden manifest regenerated. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Kcu8vXve5YMZ5wqCCKRfvJ --- docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md | 76 +++++- rl_engine/mhc/contract.py | 65 ++++- rl_engine/mhc/fixtures.py | 4 +- rl_engine/mhc/oracle.py | 35 ++- tests/fixtures/p1/golden_hashes.json | 242 ++++++++++--------- tests/test_p1_contract.py | 9 + tests/test_p1_oracle.py | 19 +- 7 files changed, 302 insertions(+), 148 deletions(-) diff --git a/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md b/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md index 0deb5dc6..5816e0c4 100644 --- a/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md +++ b/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md @@ -16,7 +16,7 @@ Issue: [DSV4][P1/7] mHC 与 RMSNorm 确定性前向/反向 (#2). | `P1-D2` | `fp32_gemm_rms` (controller projection + controller RMS, fwd + bwd) | | `P1-D3` | `mhc_post` (fwd + `dR_old`/`dy`/`dC`/`dPOST`) | | `P1-D4` | `mhc_pre` / `h_aggregate` + the TE/Megatron/Miles provider adapter | -| `P1-D5` | `rmsnorm_residual` (fwd + `dX`/`dGamma`) | +| `P1-D5` | `rmsnorm_residual` (fwd + `dX`/`dGamma`), unfused; TE fast path only once proven byte-equal | | `P1-D6` | Fixed-K / batch-invariant GEMM reference + equivalence harness | | `P1-R0` | Review, package the P1 provider artifact, GPU CI | @@ -53,6 +53,50 @@ 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: @@ -76,14 +120,14 @@ them means regenerating the manifest and bumping the schema/profile id): | # | Decision | Rationale | | --- | --- | --- | -| D1 | **`q = sqrt(s) / sqrt(K)`, not `sqrt(s / K)`.** The issue's forward text says `norm/sqrt(16384)` while its backward text says `q = sqrt(s/N)`; those differ in FP32. The forward wording wins, and backward reuses the saved `q`. | The two forms are not bit-equal; the forward is the authoritative statement. **Please confirm in review.** | +| 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` and `bias` are FP32 `[24]` vectors**, applied per controller output. | The general case; a scalar `alpha` is the broadcast special case and stays representable. | +| 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 | **The fused pre+norm boundary is defined as the unfused composition**, byte for byte. | Miles/XoRL fuse pre-mix and normalize into one launch. Defining the fused boundary this way turns "fused equals unfused" into a test rather than an assumption; a kernel whose fused residual store changes the reduction layout must register a different profile instead of presenting itself as the same kernel. | +| D8 | **`unfused` is canonical and the default**, and v1 does not reuse TE's `TEFusedResidualRMSNorm`. `fused-pre-norm` stays supported as the exception 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 is the goal, and only the unfused decomposition lets every boundary be hashed separately, so a divergence localizes to one operator instead of one megakernel. TE's fused module *refuses* to expose the intermediate — it raises on any forward hook — which would collapse two boundaries into one. The extra pass is worth it; 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". | @@ -141,17 +185,23 @@ python -m rl_engine.mhc.fixtures --write-manifest ## Open questions for review -- **D1** — `sqrt(s)/sqrt(K)` vs `sqrt(s/K)` in the controller RMS. The issue - text says both; the kit takes the forward wording. One line of Miles/Megatron - source settles it. +- **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. The kit implements the schedule exactly as written in - the issue; if Miles differs, D2/D4 change and the manifest is regenerated. -- **D8** — whether Miles can expose the pre-fusion intermediate at all. If it - cannot, `fused-pre-norm` becomes the only boundary for the inference side - and `unfused` stays the training-side definition, with the equivalence case - as the bridge. + 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 diff --git a/rl_engine/mhc/contract.py b/rl_engine/mhc/contract.py index cc3bacba..1ceab82e 100644 --- a/rl_engine/mhc/contract.py +++ b/rl_engine/mhc/contract.py @@ -35,7 +35,18 @@ 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") @@ -143,27 +154,68 @@ class ControllerParams: """ weight: torch.Tensor # FP32 [controller_n, K] - alpha: torch.Tensor # FP32 [controller_n] + 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: - for name, t in (("weight", self.weight), ("alpha", self.alpha), ("bias", self.bias)): + 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.alpha.shape) != (n,) or tuple(self.bias.shape) != (n,): - raise ValueError("controller alpha/bias must have shape [controller_n]") + 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, self.bias): + 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.to(device), self.bias.to(device)) + 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) @@ -313,6 +365,7 @@ def to(self, device: torch.device | str) -> "GradBoundary": __all__ = [ + "CANONICAL_FUSION_MODE", "COMB_SLICE", "FUSION_MODES", "MHC_EPS", diff --git a/rl_engine/mhc/fixtures.py b/rl_engine/mhc/fixtures.py index 9af29cb7..d017698f 100644 --- a/rl_engine/mhc/fixtures.py +++ b/rl_engine/mhc/fixtures.py @@ -86,7 +86,9 @@ def make_batch(name: str) -> ResidualBatch: y_sublayer=_randn(g, t, d).to(torch.bfloat16), controller=ControllerParams( weight=_randn(g, n, k, scale=1.0 / float(k) ** 0.5), - alpha=_randn(g, n, scale=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)), diff --git a/rl_engine/mhc/oracle.py b/rl_engine/mhc/oracle.py index d8be84e7..6a087bc2 100644 --- a/rl_engine/mhc/oracle.py +++ b/rl_engine/mhc/oracle.py @@ -291,7 +291,12 @@ def h_aggregate_bwd( 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``.""" + """``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 @@ -310,7 +315,8 @@ def mhc_pre_fwd( 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) - h, m = _controller_affine(p, r, batch.controller.alpha, batch.controller.bias) + 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 = { @@ -349,10 +355,16 @@ def mhc_pre_bwd( 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.alpha + alpha = batch.controller.expanded_alpha(contract) m, p, r = saved["m"], saved["p"], saved["r"] dbias = fixed_sum(dh, dim=0) - dalpha = fixed_sum(dh * m, 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 @@ -367,7 +379,9 @@ def mhc_pre_bwd( return { "d_r_old": d_r_old, "d_controller_weight": None if frozen else dweight, - "d_alpha": None if frozen else dalpha, + "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, } @@ -453,6 +467,13 @@ def rmsnorm_residual_fwd( ``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. + + v1 does **not** reuse TE's ``TEFusedResidualRMSNorm``. That module fuses the + fork and the norm and refuses to expose the intermediate (it raises on any + forward hook), which would collapse two hashable boundaries into one and + leave a divergence unlocalizable. Keeping the two boundaries separate is + worth the extra pass; a TE fast path can be swapped back in later through + the provider hook once it is proven byte-equal. """ x32 = _f32(x) d = x32.shape[1] @@ -629,7 +650,9 @@ def mhc_block_backward( "d_c": dc, "d_post": dpost, "d_controller_weight": pre_grads["d_controller_weight"], - "d_alpha": pre_grads["d_alpha"], + "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: diff --git a/tests/fixtures/p1/golden_hashes.json b/tests/fixtures/p1/golden_hashes.json index c424ca93..931ad49e 100644 --- a/tests/fixtures/p1/golden_hashes.json +++ b/tests/fixtures/p1/golden_hashes.json @@ -2,160 +2,168 @@ "cases": { "fused_pre_norm": { "bwd.d_c": "a0d671287e9d89450271cb9c6d1b75dc8ac1b1bef24a6fddfba6175b71336b92", - "bwd.d_gamma": "2cbb9bf1d2a76c10da22e70c29df5c8ad68122a9edb9b3d7dbccd64b8fd13e5e", - "bwd.d_hidden": "7b9f76e104842bf09fdcd44f11800a510bbafb000b4373d329e10da783ce9a49", + "bwd.d_gamma": "67d216275d0b662b70b423086ba11d1e7193685843ddd9cd6c6e106e06ddd4db", + "bwd.d_hidden": "bf734dd46102b092fba3d05b6e6a536ce2896d87a56a2afef0383cf3507235fb", "bwd.d_post": "aa77b606f48bbc6d1a3b3e5bbc85b01fd8c531aa8986b018a1905c705fb50608", - "bwd.d_r_old": "d9c23ad0c28d6f552f85bac27a8bb562b6eafe146badd710f2b63caa9f4f056e", - "bwd.d_stream0": "5b04030721cc487f80a5030c228db590e039474866c2328683555209e0944691", - "bwd.d_stream1": "01b83afb0d71ea384a24ccfa22c5e3fbb7f4060f21c71072a601f994ca67d64a", - "bwd.d_stream2": "455dceed73dd8efef3a69d49ad5469df893bb97f5a3e51c0bcd2302616b7ae42", - "bwd.d_stream3": "e699c4e85a1e0611b2fb1b35c2e44215cd1d3b26c4af38c7cc8dded4e1043888", - "bwd.dy_sublayer": "e89fca5228e6653ca20cb99bb51ee99fcfda4e76f128a51e78a5980f3a66f49a", - "controller.h": "bbe5fe0aea57a63f641bf93aecd96941abcc066b9fa3ed3c347cd3b56b98dd6c", + "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": "d86ffd43eeba3f7a9ceaac17a72b186682da9125e4d920a62133f7d73d1a82ee", - "grad.d_bias": "340d1d8382680bf98296c2a762a0a68af87c69a54507208e4a6b6c2027654a38", + "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": "0559a0e8022a70b2aa92e57b823d4a9a962814b540e253766de47d4c08e502f7", - "grad.d_gamma": "2cbb9bf1d2a76c10da22e70c29df5c8ad68122a9edb9b3d7dbccd64b8fd13e5e", - "grad.d_hidden": "7b9f76e104842bf09fdcd44f11800a510bbafb000b4373d329e10da783ce9a49", + "grad.d_controller_weight": "a2b206bf627cea6308dedcb75e227795e4ed9d3a8ca46a66e8c2911ea74d428f", + "grad.d_gamma": "67d216275d0b662b70b423086ba11d1e7193685843ddd9cd6c6e106e06ddd4db", + "grad.d_hidden": "bf734dd46102b092fba3d05b6e6a536ce2896d87a56a2afef0383cf3507235fb", "grad.d_post": "aa77b606f48bbc6d1a3b3e5bbc85b01fd8c531aa8986b018a1905c705fb50608", - "grad.d_r_old": "d9c23ad0c28d6f552f85bac27a8bb562b6eafe146badd710f2b63caa9f4f056e", - "grad.dy_sublayer": "e89fca5228e6653ca20cb99bb51ee99fcfda4e76f128a51e78a5980f3a66f49a", - "norm.normalized": "959d73eec2445fedf7721bf6a0a4e82c9021178b4d7c3e6913c33c9c6b3df09d", - "norm.residual": "28f5695c110eba5413488a4fb913a23d2c995aa38c7247e984a20ab8893b79b3", - "post.r_new": "c55b182a847ba286c421a6c00fd19038876ffc90044c66746cd4c094f2bee227", - "pre.hidden": "28f5695c110eba5413488a4fb913a23d2c995aa38c7247e984a20ab8893b79b3", - "r_new": "c55b182a847ba286c421a6c00fd19038876ffc90044c66746cd4c094f2bee227", - "split.c": "b5d09915a9948a4bb02a447bafa3bc21f5c3189b368c8b8e8e7f2519371d6757", - "split.post": "56a90fb2cb231d05e98c411028178234465b01b05a304e192b1464ba8efa99bd", - "split.pre": "3cf24b388aa6f03d6a48802ddcf3a3908bc63731f6b633742c0f7820279f0b2c" + "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": "4c79126d8ce2bbdec2da09959f87bd37d3c57bb6cea07002bfbe3f930b974b56", - "bwd.d_hidden": "24925df10fcedc74b03f7b9a03b8441cfd9bf81f728af397bb6a1689f392029d", + "bwd.d_gamma": "e9c922f0aa7292cc6275a506df1f027561fa78480bdff680fcafc04da496ea8c", + "bwd.d_hidden": "27eec72de166e1df1f627b6ef46dd15dba9fc54ce709723805782da2e5d8c70a", "bwd.d_post": "00822250814f728cfae3e061462cbc24a66dc50070e67c4e1bb3ae21df714563", - "bwd.d_r_old": "5d3b5882f78956dc9c4f2499508f147a5ba2278f49a438c7dd963559f7ef3ed4", - "bwd.d_stream0": "22616b202c8328551460c3610515b5e66317d98513485f7fc54eb6ff70a5dae2", - "bwd.d_stream1": "f950507bffd515220ffb58f68d497011c48da55768c2ace0557ca8f258538056", - "bwd.d_stream2": "082bf3086e17a938d24ea63140fdd9caac39362eee0524160c0a8c59b77cfa64", - "bwd.d_stream3": "f318a6ee434a4ff250b493d0d65f440b89e8e514c765b7003128d2d7e53f01c8", - "bwd.dy_sublayer": "71f264ef3ef1f4a1a3b47e1c695a5816b86e4eead38ae4dce000436f914437bf", - "controller.h": "6779debc9da9633da1a9a752adf9047e9b25d3e0696b31e3065856d3d1edcfc5", + "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": "4c79126d8ce2bbdec2da09959f87bd37d3c57bb6cea07002bfbe3f930b974b56", - "grad.d_hidden": "24925df10fcedc74b03f7b9a03b8441cfd9bf81f728af397bb6a1689f392029d", + "grad.d_gamma": "e9c922f0aa7292cc6275a506df1f027561fa78480bdff680fcafc04da496ea8c", + "grad.d_hidden": "27eec72de166e1df1f627b6ef46dd15dba9fc54ce709723805782da2e5d8c70a", "grad.d_post": "00822250814f728cfae3e061462cbc24a66dc50070e67c4e1bb3ae21df714563", - "grad.d_r_old": "5d3b5882f78956dc9c4f2499508f147a5ba2278f49a438c7dd963559f7ef3ed4", - "grad.dy_sublayer": "71f264ef3ef1f4a1a3b47e1c695a5816b86e4eead38ae4dce000436f914437bf", - "norm.normalized": "84bee319934eef7ec72a4269e7a66228b7b598bb7415e49774dc43a2e1ddace4", - "norm.residual": "e6ea03227ea97d857952435311c5eb8a5aa601d556e098535a2c15856efc732d", - "post.r_new": "63531fa1253c4c9a95b16587c47bb46bde499f22b9981f6de44a669dd1c2bbec", - "pre.hidden": "e6ea03227ea97d857952435311c5eb8a5aa601d556e098535a2c15856efc732d", - "r_new": "63531fa1253c4c9a95b16587c47bb46bde499f22b9981f6de44a669dd1c2bbec", - "split.c": "e3bc8f726b626e2e9fd704724d2f715b683f8e9f573cc0052876cacce31bd4b7", - "split.post": "8dc421dd34664d5906ec0770b11aa070e88f8cf8c258d55a5d25dbcb927f5d51", - "split.pre": "6c70243dea814984cb9b9fa91b78c55f785a677639223009bb7a466142932cbb" + "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": "6446e0a79a7f572645d2035edbdc8ba2ca80c0c8e1c2a73ca1a9f60ad8c56e4f", - "bwd.d_hidden": "5bb71358cf77a95e40bdd20581f0dd947fc26c5f207243cd32848169b25160fa", + "bwd.d_gamma": "51b20c678e10e5ee453fe448c2e7f8490dd0472ddda6f29f1e821293e4865200", + "bwd.d_hidden": "28ca22b277231a74721a26e8cd46bf3ae5c1905d43fdc76087a36e4234bc49cc", "bwd.d_post": "1330c48fdaa948c9e111f503fce0e7bb058b1d79ff5811da6f9496c43c2afe8f", - "bwd.d_r_old": "c417f939f9b45fe95beacef07ec4f7f16de5871172866b750524b911bd012f2e", - "bwd.d_stream0": "1f3aaec02c8efb2dbbe2dc7b5209d23c8fab2ac7de59d8cbfb750ba1762f0e4d", - "bwd.d_stream1": "65f47151b2b6b35b3da54cc863d8fe07ffde4bf715b2c945739f2d1295bf6e86", - "bwd.d_stream2": "ae86ccefff2fc2a41b5a39ef87d1dc2f5c808b8161b2a972251c920878910c0c", - "bwd.d_stream3": "d05a20c9800f67804627c1fccfadf35c280145e5334e84369af6ea38400422ca", - "bwd.dy_sublayer": "d712cddba241d525fe30b12110ae83844d27a8f7f27af58a5c9d379de033c401", - "controller.h": "62b5eb67c84659e31b4f8ad98963ad85370353506dc13c2da64778a6ec7bfcdd", + "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": "08af3310b0057e8d6719fed3310325c4069f9bc13faa18b7f653acd3d66a977e", - "grad.d_bias": "0a4891f9d439e8bc6f8e0eecd4a79ed3eb56dbf26b30e16013b6d9a5fe6bc6af", + "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": "f98d0f49f04752dcbffac7a355775668ee7d404ede7aae90b6936d93c4fff0c3", - "grad.d_gamma": "6446e0a79a7f572645d2035edbdc8ba2ca80c0c8e1c2a73ca1a9f60ad8c56e4f", - "grad.d_hidden": "5bb71358cf77a95e40bdd20581f0dd947fc26c5f207243cd32848169b25160fa", + "grad.d_controller_weight": "ba2e502197129e625dad9ef66686aa95a133a2743d589c434f3b9c4c6f8311d1", + "grad.d_gamma": "51b20c678e10e5ee453fe448c2e7f8490dd0472ddda6f29f1e821293e4865200", + "grad.d_hidden": "28ca22b277231a74721a26e8cd46bf3ae5c1905d43fdc76087a36e4234bc49cc", "grad.d_post": "1330c48fdaa948c9e111f503fce0e7bb058b1d79ff5811da6f9496c43c2afe8f", - "grad.d_r_old": "c417f939f9b45fe95beacef07ec4f7f16de5871172866b750524b911bd012f2e", - "grad.dy_sublayer": "d712cddba241d525fe30b12110ae83844d27a8f7f27af58a5c9d379de033c401", - "norm.normalized": "daf839f63a03952403fa271d5908aa0c8113c98d061b03fbd4085dc8d38d226b", - "norm.residual": "e241696650c290471579a96c4fad24c49041696f2ddce026493f231dece1b43d", - "post.r_new": "123c2f784e6052882217288aceaa547231f3fce750c0859122e328b6a46ce6b8", - "pre.hidden": "e241696650c290471579a96c4fad24c49041696f2ddce026493f231dece1b43d", - "r_new": "123c2f784e6052882217288aceaa547231f3fce750c0859122e328b6a46ce6b8", - "split.c": "06fa72314507fe0e9413fbff2937c4f76f28c26968df30fce312236468601ac0", - "split.post": "fec5506651ed1406f5e526c081611b0ce23216beda926b23e5e9d74846169159", - "split.pre": "2511d4f9a82f34724c2822a326a20ea8e28600568f13047a822fff191f4a077d" + "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": "87ce69b941c0eed1475bcfaf608f23ac24e7c6582367d01b78e417fac036c103", - "bwd.d_hidden": "cc910a66cadf396885f5082c8969f7cc207910daa058d855dbbf1a48fd0d403b", + "bwd.d_gamma": "8d30d84cfb87751048d599c882b4c873c3e5f5e55d85aaf9596fc2eadb09edee", + "bwd.d_hidden": "6567631097691d5fa4c9e382a59ba68a262ee4a6290b3c0c86730cafa8d5cf0d", "bwd.d_post": "edb376c01993aa5fff0e98713f1ab16a6aed8e94b941d1e7547da40eb2b556c4", - "bwd.d_r_old": "645bc398d44ab31d92e04b173b1d68b91c5b7e82b77832f6941767118f93b1c3", - "bwd.d_stream0": "3b96ad5281ead49c0e596e4c26a1bf033266e8bbb556c6d7cc37d3699d37b7b7", - "bwd.d_stream1": "1fd810234d6587e2d7021ae030047b8c9d58f6f0d9ec9b76394b9bb391b80cf0", - "bwd.d_stream2": "6aa0159f5253249184ba9de35f3bd0a47a5b164180858e74822ca702c9c91970", - "bwd.d_stream3": "3433b559730345dc3d38d78837373cdba4cc3f63eefcc1b141cb1f8f6fdd6ad3", - "bwd.dy_sublayer": "e8028379f260744aee3e21889f1c955f38c363c063d267663c8103e1950dd502", - "controller.h": "2da829cd236c4c268a36165f2a796f650eeb86cfab3d2a70a9ca445ef861e0ad", + "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": "74df4fe6c987cabe231706935ba0785f5001f6b78a9831664e93bae6821024f4", - "grad.d_bias": "6f4afe52afe88385ba02456cdaf5c2f71f6eae53dda737514f2e1abe31f40b3c", + "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": "3b9bcce8e798e1287f2a260409ceee3f3c610515743acf705d4dfda145484272", - "grad.d_gamma": "87ce69b941c0eed1475bcfaf608f23ac24e7c6582367d01b78e417fac036c103", - "grad.d_hidden": "cc910a66cadf396885f5082c8969f7cc207910daa058d855dbbf1a48fd0d403b", + "grad.d_controller_weight": "9f60ac4ce393f66768d7977115960a1454c717a688129dafbba07f0515c36dc7", + "grad.d_gamma": "8d30d84cfb87751048d599c882b4c873c3e5f5e55d85aaf9596fc2eadb09edee", + "grad.d_hidden": "6567631097691d5fa4c9e382a59ba68a262ee4a6290b3c0c86730cafa8d5cf0d", "grad.d_post": "edb376c01993aa5fff0e98713f1ab16a6aed8e94b941d1e7547da40eb2b556c4", - "grad.d_r_old": "645bc398d44ab31d92e04b173b1d68b91c5b7e82b77832f6941767118f93b1c3", - "grad.dy_sublayer": "e8028379f260744aee3e21889f1c955f38c363c063d267663c8103e1950dd502", - "norm.normalized": "cd4ca1d9249a3987c285f0e4ec75a163d2c4bf1d16c158d01617c40a203374ba", - "norm.residual": "1de7fd9c90975c7a31c8b04921737d1494ef33a2e79f27fdedcf93f022414a99", - "post.r_new": "619463014e0852ab8bfa105ce8734b2f764162d2ab839a3668bffeabb434ca5a", - "pre.hidden": "1de7fd9c90975c7a31c8b04921737d1494ef33a2e79f27fdedcf93f022414a99", - "r_new": "619463014e0852ab8bfa105ce8734b2f764162d2ab839a3668bffeabb434ca5a", - "split.c": "12b25bd8db4ee0dba6fca2a012b36b58b01b822f554e5ffaa23e651cd127e3f6", - "split.post": "075e73f84f62aa372cb030365b046ca35f11b52a19d9691acc633c48e6bd0915", - "split.pre": "5c9f5b1f8fcc8f3a8f2019dfdb8394498769b14692689ff237ef77ae4ec60086" + "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": "1d6a9bd5ea4f28e15486f05c23b91edf6fce9a5bb2be4f2b89acf8c3d91d8bdf", - "bwd.d_hidden": "76d1b779a3fe8c0a89f2b8506360d82cd0b574c8f3c3e8e09ee28cacf353550b", + "bwd.d_gamma": "5055b0a729868cb304b393ef4ea54536b9acc26150e7411bdb8f33915be4e89d", + "bwd.d_hidden": "a77f31a38d1fc394e449a6d8ee73afb4c94583f1cb83bccdf329d197babaf29d", "bwd.d_post": "049df2a605804eee375c008e2807d47be65c160927d49eed8d4b9c5963adf055", - "bwd.d_r_old": "5cc94c4ef9a303c0f603d3db0987d8c3ef8e23ae3be265305688bff9d63a0aad", - "bwd.d_stream0": "a113bf7982eb3df270365ad43e9b76988503038de9e3fc6e3fd6c1f64571968a", - "bwd.d_stream1": "87c978dd759a4dae8a627c725395f878bc5c269be3785d9cf6e49652f655d9f4", - "bwd.d_stream2": "96af2de9fe231c2ca75fb2cef9c5d8d650e36766bbffb70da8013df5c4285459", - "bwd.d_stream3": "58680ebc9b0c9c043bd73262a5d81432cade3073b24ae45b6b1cb2057acb9c66", - "bwd.dy_sublayer": "449aa7deb18ff1c70c2e4d89628a1759814693c9eb3e66e0974f7c7d75a59bb2", - "controller.h": "1a82827e4fc25f93405a1b862425cea2a9c36e66f6839dcb7b3b4843119a845c", + "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": "0149d892249bae065b3cffe9d13c82ddf39b55927ff14e85496a2227f73b1a9e", - "grad.d_bias": "e4783ae2ccdaaedf845d4549ca0d1492e7363c2b94e27e76c80c59091217a798", + "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": "6844879775e886c24ddde3382898bd11ae4b043f37447d33eb4869e1c0a38fec", - "grad.d_gamma": "1d6a9bd5ea4f28e15486f05c23b91edf6fce9a5bb2be4f2b89acf8c3d91d8bdf", - "grad.d_hidden": "76d1b779a3fe8c0a89f2b8506360d82cd0b574c8f3c3e8e09ee28cacf353550b", + "grad.d_controller_weight": "86bb3903b8aec54115c7ba1c2c7a2d6f0ab61fcdcc79554498a3b47c806deddb", + "grad.d_gamma": "5055b0a729868cb304b393ef4ea54536b9acc26150e7411bdb8f33915be4e89d", + "grad.d_hidden": "a77f31a38d1fc394e449a6d8ee73afb4c94583f1cb83bccdf329d197babaf29d", "grad.d_post": "049df2a605804eee375c008e2807d47be65c160927d49eed8d4b9c5963adf055", - "grad.d_r_old": "5cc94c4ef9a303c0f603d3db0987d8c3ef8e23ae3be265305688bff9d63a0aad", - "grad.dy_sublayer": "449aa7deb18ff1c70c2e4d89628a1759814693c9eb3e66e0974f7c7d75a59bb2", - "norm.normalized": "e8bd67df4e08da06948e6bb01bc487c2fccd27d9ff95a4e1e16bd85636f701ad", - "norm.residual": "eec3f0342ddf7353b47af652ed258595df38177741308c27211efb2b54a0a0df", - "post.r_new": "aca7cdeab1fc9670a944ea737986a439945f888ee8b4e396e0d370f750682840", - "pre.hidden": "eec3f0342ddf7353b47af652ed258595df38177741308c27211efb2b54a0a0df", - "r_new": "aca7cdeab1fc9670a944ea737986a439945f888ee8b4e396e0d370f750682840", - "split.c": "5b0d44b39615bc28c7c9137e2d8694747decf28c514e80bfa52882ddb2278b84", - "split.post": "d7dafd8501437d3106cd98b94c7f57b0184f266413d966f8bdfc027c802e1fcd", - "split.pre": "fb06566020b4ce8a069e189f373efd5280505ac48fd0d72c2dfeaadb97d07b28" + "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", diff --git a/tests/test_p1_contract.py b/tests/test_p1_contract.py index 025f676b..581e215c 100644 --- a/tests/test_p1_contract.py +++ b/tests/test_p1_contract.py @@ -52,6 +52,15 @@ def test_frozen_constants_cannot_be_overridden() -> None: 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"}, diff --git a/tests/test_p1_oracle.py b/tests/test_p1_oracle.py index 2414e3bb..da6c7f3f 100644 --- a/tests/test_p1_oracle.py +++ b/tests/test_p1_oracle.py @@ -282,7 +282,9 @@ def _block_autograd(batch: ResidualBatch, grads): r = batch.r_old.float().requires_grad_(True) y = batch.y_sublayer.float().requires_grad_(True) w = batch.controller.weight.clone().requires_grad_(True) - alpha = batch.controller.alpha.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) @@ -290,6 +292,7 @@ def _block_autograd(batch: ResidualBatch, grads): 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) @@ -309,7 +312,9 @@ def _block_autograd(batch: ResidualBatch, grads): "d_r_old": r.grad, "dy_sublayer": y.grad, "d_controller_weight": w.grad, - "d_alpha": alpha.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, } @@ -327,8 +332,12 @@ def test_block_backward_matches_autograd(case: str) -> None: def test_fused_equals_unfused_bytes() -> None: - """Miles fuses pre-mix and normalize into one launch; the kit's fused - boundary must be the same arithmetic, not merely the same math.""" + """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") @@ -345,7 +354,7 @@ def test_mixer_frozen_leaks_no_controller_gradient() -> None: 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", "d_bias"): + 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() From 48bf9eb6191d3efe07eedd2bb95cd3b31822673e Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Thu, 3 Sep 2026 11:21:58 +0000 Subject: [PATCH 3/4] [DSv4][P1-S0] Correct the P1 task split to match issue #2 Issue #2 splits P1 by function, not by operator: P1-D1 owns the mHC forwards, P1-D2 the RMSNorm forwards, P1-D3 every custom backward across both, P1-D4 the provider adapter, P1-D5 trace/tests/benchmarks and P1-D6 the fixed-K GEMM reference -- eight tasks with P1-S0 and P1-R0. The kit had invented a one-task-per-operator split instead. That was wrong in the design doc and, worse, in StubProvider's messages, which told a developer to "claim it on P1-D1 (hc_split_sinkhorn)" and would have pointed them at the wrong task. - StubProvider now names the real owning task per (operator, direction); every *_bwd points at P1-D3. - design doc: the eight tasks with owners, the dependency line, and a note that operators and tasks are different axes. - provider docstring and tests updated to match. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Kcu8vXve5YMZ5wqCCKRfvJ --- docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md | 40 ++++++++++++++------ rl_engine/mhc/provider.py | 20 +++++----- tests/test_p1_provider.py | 5 +++ 3 files changed, 44 insertions(+), 21 deletions(-) diff --git a/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md b/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md index 5816e0c4..0e033d70 100644 --- a/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md +++ b/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md @@ -7,18 +7,36 @@ backend PR can run independently. Issue: [DSV4][P1/7] mHC 与 RMSNorm 确定性前向/反向 (#2). -## Sub-issue naming +## The eight P1 tasks -| Label | Scope | -| --- | --- | -| `P1-S0` | This start kit (contract, oracle, fixtures, acceptance command) | -| `P1-D1` | `hc_split_sinkhorn` (fwd + bwd through all 20 Sinkhorn rounds) | -| `P1-D2` | `fp32_gemm_rms` (controller projection + controller RMS, fwd + bwd) | -| `P1-D3` | `mhc_post` (fwd + `dR_old`/`dy`/`dC`/`dPOST`) | -| `P1-D4` | `mhc_pre` / `h_aggregate` + the TE/Megatron/Miles provider adapter | -| `P1-D5` | `rmsnorm_residual` (fwd + `dX`/`dGamma`), unfused; TE fast path only once proven byte-equal | -| `P1-D6` | Fixed-K / batch-invariant GEMM reference + equivalence harness | -| `P1-R0` | Review, package the P1 provider artifact, GPU CI | +Issue #2 splits the work by **function**, not by operator. Read that carefully +before claiming one: `P1-D3` owns *every* custom backward across both mHC and +RMSNorm, and `P1-D1`/`P1-D2` own only forwards. + +| Label | Owner | Scope | +| --- | --- | --- | +| `P1-S0` | maintainer | This start kit: `LayerContract`, CPU oracle, forward/saved-tensor golden fixtures, provider stub, standalone acceptance command | +| `P1-D1` | dev | `hc_pre` / `hc_post` **forward** — covers `hc_split_sinkhorn_fwd`, `fp32_gemm_rms_fwd`, `h_aggregate_fwd`, `mhc_pre_fwd`, `mhc_post_fwd` | +| `P1-D2` | dev | RMSNorm / residual-RMSNorm **forward**: extend the existing WS1 RMSNorm, add residual-RMSNorm, fixed reduction tree, `rsqrt` | +| `P1-D3` | dev | mHC/RMSNorm custom **backward**, built on the recorded saved-tensor fixtures — **does not wait on D1/D2 kernels** | +| `P1-D4` | dev | TE/Megatron/Miles provider adapter, capability and fallback checks — codes against the stub API first | +| `P1-D5` | dev | Trace, batch/padding/stride/recompute tests, correctness/performance benchmarks — runs against the oracle/stub first | +| `P1-D6` | dev | Fixed-K / batch-invariant GEMM reference + bit-equivalence harness (consumed by P2/P3/P5/P7) | +| `P1-R0` | maintainer | Review each `P1-D*` independently, package the P1 provider artifact, strict bytes / fallback / performance / local GPU CI | + +`Foundation v1 -> P1-S0 -> {P1-D1 || P1-D2 || P1-D3 || P1-D4 || P1-D5 || P1-D6} -> P1-R0` + +Every `P1-D*` depends only on `P1-S0`, never on each other. That is what the +kit buys: D3 differentiates against recorded saved tensors, D4 and D5 run +against the stub and the oracle, so none of them block on a forward kernel. + +### The five WS1 operators (a different axis) + +The operators to be written are `hc_split_sinkhorn`, `fp32_gemm_rms`, +`mhc_post`, `mhc_pre` (with `h_aggregate`) and `rmsnorm_residual`, plus the +`P1-D6` fixed-K GEMM. Operators are the *what*; the tasks above are the *who*. +`ReferenceProvider` is keyed by operator, so an owner overrides whichever +`_fwd`/`_bwd` entry points fall inside their task. ## What is in the kit diff --git a/rl_engine/mhc/provider.py b/rl_engine/mhc/provider.py index 8fd9eb30..27c84bc1 100644 --- a/rl_engine/mhc/provider.py +++ b/rl_engine/mhc/provider.py @@ -179,34 +179,34 @@ def _todo(task: str) -> NotImplementedError: ) def hc_split_sinkhorn_fwd(self, h, contract=None): - raise self._todo("P1-D1 (hc_split_sinkhorn)") + raise self._todo("P1-D1 (hc_pre/hc_post forward)") def hc_split_sinkhorn_bwd(self, dpre, dpost, dc, saved): - raise self._todo("P1-D1 (hc_split_sinkhorn)") + raise self._todo("P1-D3 (mHC/RMSNorm custom backward)") def fp32_gemm_rms_fwd(self, x_flat, weight, eps): - raise self._todo("P1-D2 (fp32_gemm_rms)") + raise self._todo("P1-D1 (hc_pre/hc_post forward)") def fp32_gemm_rms_bwd(self, dp, dr, x_flat, weight, saved): - raise self._todo("P1-D2 (fp32_gemm_rms)") + raise self._todo("P1-D3 (mHC/RMSNorm custom backward)") def mhc_post_fwd(self, r_old, y, c, post): - raise self._todo("P1-D3 (mhc_post)") + raise self._todo("P1-D1 (hc_pre/hc_post forward)") def mhc_post_bwd(self, dr_new, r_old, y, c, post): - raise self._todo("P1-D3 (mhc_post)") + raise self._todo("P1-D3 (mHC/RMSNorm custom backward)") def h_aggregate_fwd(self, pre, r_old): - raise self._todo("P1-D4 (mhc_pre / h_aggregate)") + raise self._todo("P1-D1 (hc_pre/hc_post forward)") def h_aggregate_bwd(self, dh, pre, r_old): - raise self._todo("P1-D4 (mhc_pre / h_aggregate)") + raise self._todo("P1-D3 (mHC/RMSNorm custom backward)") def rmsnorm_residual_fwd(self, x, gamma, eps): - raise self._todo("P1-D5 (rmsnorm_residual)") + raise self._todo("P1-D2 (RMSNorm/residual-RMSNorm forward)") def rmsnorm_residual_bwd(self, dy, d_residual, x, gamma, saved): - raise self._todo("P1-D5 (rmsnorm_residual)") + raise self._todo("P1-D3 (mHC/RMSNorm custom backward)") def fixed_k_gemm_fwd(self, x, w): raise self._todo("P1-D6 (fixed-K GEMM reference)") diff --git a/tests/test_p1_provider.py b/tests/test_p1_provider.py index d646f55b..02812f06 100644 --- a/tests/test_p1_provider.py +++ b/tests/test_p1_provider.py @@ -43,6 +43,11 @@ def test_stub_provider_fails_closed() -> None: stub.hc_split_sinkhorn_fwd(torch.zeros(1, 24)) with pytest.raises(NotImplementedError, match="P1-D6"): stub.fixed_k_gemm_fwd(torch.zeros(1, 4), torch.zeros(2, 4)) + # backward belongs to P1-D3 regardless of which operator it is + with pytest.raises(NotImplementedError, match="P1-D3"): + stub.rmsnorm_residual_bwd(None, None, None, None, {}) + with pytest.raises(NotImplementedError, match="P1-D2"): + 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) From 39c4aea31e359ffb0fd28fbad9cdf9067ba06690 Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Thu, 3 Sep 2026 11:25:25 +0000 Subject: [PATCH 4/4] [DSv4][P1-S0] Use the real sub-issues of #2 (#14-#21) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #2 has eight GitHub sub-issues, and they are split one per operator -- not by function as the older 任务分工 section in the issue body reads. The previous commit followed that superseded section and mislabelled everything; this restores the operator split and uses the authoritative issue numbers, matching how P5 labels its sub-tasks. - P1-1 (#14) hc_split_sinkhorn, P1-2 (#15) fp32_gemm_rms, P1-3 (#16) mhc_post, P1-4 (#17) mhc_pre/h_aggregate, P1-5 (#18) rmsnorm_residual -- each fwd + bwd together; P1-6..P1-8 (#19-#21) are the WS2 TP/SP/CP/ DP/PP tasks. - #15 explicitly absorbs the fixed-K / batch-invariant GEMM reference and its equivalence harness, so fixed_k_gemm points there rather than at a task of its own. - StubProvider names the owning sub-issue per operator. - D8 reframed against #18's actual wording: it says prefer TE's TEFusedResidualRMSNorm and self-write only if TE fails the deterministic contract. TE fails it -- it refuses to expose the pre-normalization intermediate -- so unfused v1 is the escape hatch #18 provides, not a departure from it. - design doc: the eight sub-issues with stage, issue number and the WS2 attachment points (placement + check_capability). 54 tests pass; check_p1.py reference PASS, stub fails closed; lint clean. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Kcu8vXve5YMZ5wqCCKRfvJ --- docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md | 57 ++++++++++---------- rl_engine/mhc/oracle.py | 13 ++--- rl_engine/mhc/provider.py | 24 ++++----- tests/test_p1_provider.py | 10 ++-- 4 files changed, 52 insertions(+), 52 deletions(-) diff --git a/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md b/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md index 0e033d70..eda802fd 100644 --- a/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md +++ b/docs/design/dsv4_p1_mhc_rmsnorm_start_kit.md @@ -7,36 +7,37 @@ backend PR can run independently. Issue: [DSV4][P1/7] mHC 与 RMSNorm 确定性前向/反向 (#2). -## The eight P1 tasks +## 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 | -Issue #2 splits the work by **function**, not by operator. Read that carefully -before claiming one: `P1-D3` owns *every* custom backward across both mHC and -RMSNorm, and `P1-D1`/`P1-D2` own only forwards. - -| Label | Owner | Scope | -| --- | --- | --- | -| `P1-S0` | maintainer | This start kit: `LayerContract`, CPU oracle, forward/saved-tensor golden fixtures, provider stub, standalone acceptance command | -| `P1-D1` | dev | `hc_pre` / `hc_post` **forward** — covers `hc_split_sinkhorn_fwd`, `fp32_gemm_rms_fwd`, `h_aggregate_fwd`, `mhc_pre_fwd`, `mhc_post_fwd` | -| `P1-D2` | dev | RMSNorm / residual-RMSNorm **forward**: extend the existing WS1 RMSNorm, add residual-RMSNorm, fixed reduction tree, `rsqrt` | -| `P1-D3` | dev | mHC/RMSNorm custom **backward**, built on the recorded saved-tensor fixtures — **does not wait on D1/D2 kernels** | -| `P1-D4` | dev | TE/Megatron/Miles provider adapter, capability and fallback checks — codes against the stub API first | -| `P1-D5` | dev | Trace, batch/padding/stride/recompute tests, correctness/performance benchmarks — runs against the oracle/stub first | -| `P1-D6` | dev | Fixed-K / batch-invariant GEMM reference + bit-equivalence harness (consumed by P2/P3/P5/P7) | -| `P1-R0` | maintainer | Review each `P1-D*` independently, package the P1 provider artifact, strict bytes / fallback / performance / local GPU CI | - -`Foundation v1 -> P1-S0 -> {P1-D1 || P1-D2 || P1-D3 || P1-D4 || P1-D5 || P1-D6} -> P1-R0` - -Every `P1-D*` depends only on `P1-S0`, never on each other. That is what the -kit buys: D3 differentiates against recorded saved tensors, D4 and D5 run -against the stub and the oracle, so none of them block on a forward kernel. +``` +Foundation v1 -> P1-S0 -> {P1-1 || P1-2 || P1-3 || P1-4 || P1-5} -> P1-R0 + | + +-> WS2: P1-6, P1-7, P1-8 +``` -### The five WS1 operators (a different axis) +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`. -The operators to be written are `hc_split_sinkhorn`, `fp32_gemm_rms`, -`mhc_post`, `mhc_pre` (with `h_aggregate`) and `rmsnorm_residual`, plus the -`P1-D6` fixed-K GEMM. Operators are the *what*; the tasks above are the *who*. -`ReferenceProvider` is keyed by operator, so an owner overrides whichever -`_fwd`/`_bwd` entry points fall inside their task. +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 @@ -145,7 +146,7 @@ them means regenerating the manifest and bumping the schema/profile id): | 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**, and v1 does not reuse TE's `TEFusedResidualRMSNorm`. `fused-pre-norm` stays supported as the exception 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 is the goal, and only the unfused decomposition lets every boundary be hashed separately, so a divergence localizes to one operator instead of one megakernel. TE's fused module *refuses* to expose the intermediate — it raises on any forward hook — which would collapse two boundaries into one. The extra pass is worth it; a fused fast path can be swapped back through the provider hook once proven byte-equal. | +| 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". | diff --git a/rl_engine/mhc/oracle.py b/rl_engine/mhc/oracle.py index 6a087bc2..1437830c 100644 --- a/rl_engine/mhc/oracle.py +++ b/rl_engine/mhc/oracle.py @@ -468,12 +468,13 @@ def rmsnorm_residual_fwd( 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. - v1 does **not** reuse TE's ``TEFusedResidualRMSNorm``. That module fuses the - fork and the norm and refuses to expose the intermediate (it raises on any - forward hook), which would collapse two hashable boundaries into one and - leave a divergence unlocalizable. Keeping the two boundaries separate is - worth the extra pass; a TE fast path can be swapped back in later through - the provider hook once it is proven byte-equal. + 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] diff --git a/rl_engine/mhc/provider.py b/rl_engine/mhc/provider.py index 27c84bc1..bf189cdc 100644 --- a/rl_engine/mhc/provider.py +++ b/rl_engine/mhc/provider.py @@ -179,40 +179,40 @@ def _todo(task: str) -> NotImplementedError: ) def hc_split_sinkhorn_fwd(self, h, contract=None): - raise self._todo("P1-D1 (hc_pre/hc_post forward)") + raise self._todo("P1-1 (#14)") def hc_split_sinkhorn_bwd(self, dpre, dpost, dc, saved): - raise self._todo("P1-D3 (mHC/RMSNorm custom backward)") + raise self._todo("P1-1 (#14)") def fp32_gemm_rms_fwd(self, x_flat, weight, eps): - raise self._todo("P1-D1 (hc_pre/hc_post forward)") + raise self._todo("P1-2 (#15)") def fp32_gemm_rms_bwd(self, dp, dr, x_flat, weight, saved): - raise self._todo("P1-D3 (mHC/RMSNorm custom backward)") + raise self._todo("P1-2 (#15)") def mhc_post_fwd(self, r_old, y, c, post): - raise self._todo("P1-D1 (hc_pre/hc_post forward)") + raise self._todo("P1-3 (#16)") def mhc_post_bwd(self, dr_new, r_old, y, c, post): - raise self._todo("P1-D3 (mHC/RMSNorm custom backward)") + raise self._todo("P1-3 (#16)") def h_aggregate_fwd(self, pre, r_old): - raise self._todo("P1-D1 (hc_pre/hc_post forward)") + raise self._todo("P1-4 (#17)") def h_aggregate_bwd(self, dh, pre, r_old): - raise self._todo("P1-D3 (mHC/RMSNorm custom backward)") + raise self._todo("P1-4 (#17)") def rmsnorm_residual_fwd(self, x, gamma, eps): - raise self._todo("P1-D2 (RMSNorm/residual-RMSNorm forward)") + raise self._todo("P1-5 (#18)") def rmsnorm_residual_bwd(self, dy, d_residual, x, gamma, saved): - raise self._todo("P1-D3 (mHC/RMSNorm custom backward)") + raise self._todo("P1-5 (#18)") def fixed_k_gemm_fwd(self, x, w): - raise self._todo("P1-D6 (fixed-K GEMM reference)") + raise self._todo("P1-2 (#15)") def fixed_k_gemm_bwd(self, dy, x, w): - raise self._todo("P1-D6 (fixed-K GEMM reference)") + raise self._todo("P1-2 (#15)") def resolve_provider(spec: str) -> MHCProvider: diff --git a/tests/test_p1_provider.py b/tests/test_p1_provider.py index 02812f06..35aadad1 100644 --- a/tests/test_p1_provider.py +++ b/tests/test_p1_provider.py @@ -39,14 +39,12 @@ def test_reference_provider_matches_oracle_bytes() -> None: def test_stub_provider_fails_closed() -> None: stub = StubProvider() - with pytest.raises(NotImplementedError, match="P1-D1"): + with pytest.raises(NotImplementedError, match=r"P1-1 \(#14\)"): stub.hc_split_sinkhorn_fwd(torch.zeros(1, 24)) - with pytest.raises(NotImplementedError, match="P1-D6"): + # #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)) - # backward belongs to P1-D3 regardless of which operator it is - with pytest.raises(NotImplementedError, match="P1-D3"): - stub.rmsnorm_residual_bwd(None, None, None, None, {}) - with pytest.raises(NotImplementedError, match="P1-D2"): + 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):