From 1604db47f8a0f1341eb2def1769f0c232ef1a808 Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Tue, 1 Sep 2026 11:49:49 +0800 Subject: [PATCH 1/2] p5 starter --- docs/design/dsv4_p5_expert_start_kit.md | 97 ++++++ rl_engine/moe/__init__.py | 40 +++ rl_engine/moe/contract.py | 224 ++++++++++++++ rl_engine/moe/fixtures.py | 224 ++++++++++++++ rl_engine/moe/mx_format.py | 196 ++++++++++++ rl_engine/moe/oracle.py | 386 ++++++++++++++++++++++++ rl_engine/moe/provider.py | 194 ++++++++++++ rl_engine/moe/trace.py | 76 +++++ scripts/check_p5.py | 143 +++++++++ tests/fixtures/p5/golden_hashes.json | 131 ++++++++ tests/test_p5_contract.py | 65 ++++ tests/test_p5_mx_format.py | 89 ++++++ tests/test_p5_oracle.py | 147 +++++++++ tests/test_p5_provider.py | 58 ++++ 14 files changed, 2070 insertions(+) create mode 100644 docs/design/dsv4_p5_expert_start_kit.md create mode 100644 rl_engine/moe/__init__.py create mode 100644 rl_engine/moe/contract.py create mode 100644 rl_engine/moe/fixtures.py create mode 100644 rl_engine/moe/mx_format.py create mode 100644 rl_engine/moe/oracle.py create mode 100644 rl_engine/moe/provider.py create mode 100644 rl_engine/moe/trace.py create mode 100755 scripts/check_p5.py create mode 100644 tests/fixtures/p5/golden_hashes.json create mode 100644 tests/test_p5_contract.py create mode 100644 tests/test_p5_mx_format.py create mode 100644 tests/test_p5_oracle.py create mode 100644 tests/test_p5_provider.py diff --git a/docs/design/dsv4_p5_expert_start_kit.md b/docs/design/dsv4_p5_expert_start_kit.md new file mode 100644 index 00000000..f42d0a14 --- /dev/null +++ b/docs/design/dsv4_p5_expert_start_kit.md @@ -0,0 +1,97 @@ +# P5 Expert Start Kit (`P5-S0`) + +The start kit unblocks every P5 sub-issue (P5-1…P5-9): it freezes the data +contract, provides a bit-exact FP32 oracle for the five WS1 operators, +generates seeded golden fixtures, and ships one acceptance command that any +backend PR can run independently. + +## Sub-issue naming (development order posted on #8) + +`P5-N` is the development-order label from the sequencing comment on #8; +GitHub issue numbers stay authoritative for links. + +| Label | Scope | +| --- | --- | +| P5-S0 | This start kit (contract, oracle, fixtures, acceptance command) | +| P5-1 | `mxfp8_act_quant` (fwd + STE bwd) | +| P5-2 | `clamp_swiglu_weighted` (fwd + dgate/dup/dp_s) | +| P5-3 | `shared_grouped_lora_delta` (fwd + dX/dA/dB) | +| P5-4 | `mxfp8_mxfp4_grouped_gemm` (fwd + dX only) | +| P5-5 | `shared_expert_mlp` (fwd + dX only) | +| P5-6 | `moe_provider_adapter` (Megatron + vLLM injection) | +| P5-7 | WS2: EP placement, `expert_tensor_parallel_size = 1` gate | +| P5-8 | WS2: shared expert TP/SP + shared-once gate | +| P5-9 | WS2: adapter fail-closed under EP>1 / placements | + +## What is in the kit + +| Module | Contents | +| --- | --- | +| `rl_engine/moe/mx_format.py` | OCP MX codecs: E8M0 / E4M3 / E2M1, block-32 quantize/dequantize, nibble packing. Defines the golden bytes for P5-1/P5-4. | +| `rl_engine/moe/contract.py` | `ExpertBatch`, `SharedBatch`, `LoRAParams`, clamp constants, tensor fingerprints. P5-local subset of the Foundation `ExpertBatch` ABI (`p5-expertbatch-v1`). | +| `rl_engine/moe/oracle.py` | FP32 reference for the five operators plus the full routed/shared forward–backward compositions. | +| `rl_engine/moe/provider.py` | `ExpertProvider` protocol, `ReferenceProvider` (oracle-backed), `StubProvider` (fail-closed). | +| `rl_engine/moe/fixtures.py` | Seeded fixture cases and the golden-hash manifest (`tests/fixtures/p5/golden_hashes.json`, the CI anchor). | +| `rl_engine/moe/trace.py` | Boundary hashes + `first_divergence` (P5-local stand-in for `TraceEnvelope`). | +| `scripts/check_p5.py` | The acceptance command. | + +## Frozen numeric contract (recap of #8 + decisions made here) + +From the issues: + +1. LoRA-only fine-tuning; base weights frozen — **no `dW` anywhere**. +2. Routed base is MXFP8 activation × MXFP4 frozen weight; block = 32, scale = + E8M0, elements = E4M3 / E2M1 (OCP Microscaling v1.0). +3. Backward is BF16 (no MXFP8 re-quant); every reduction uses FP32 accumulators. +4. Route weight `p_s` is applied in `clamp_swiglu_weighted` + (`h = SiLU(min(gate,10)) · clamp(up,−10,10) · p_s`), exactly once globally. +5. `mxfp8_act_quant` amax is a row-local 32-element reduction; backward is STE. +6. One-round SwiGLU: FP32 math, a single BF16 round on the output. + +Decisions this kit had to freeze (flagged for review on #8; changing any of +them requires regenerating the manifest and bumping the schema/profile id): + +| # | Decision | Rationale | +| --- | --- | --- | +| D1 | **E4M3 encode = clamp to ±448 in FP32, then RNE cast** (torch `float8_e4m3fn`). Bare torch cast maps overflow to NaN; clamp+cast equals PTX `cvt.satfinite`. | Matches hardware satfinite; pinned by golden tests. | +| D2 | **E8M0 scale recipe**: `shared_exp = floor(log2(amax)) − emax_elem` (8 for E4M3, 2 for E2M1); all-zero block → code 127 (scale 1). `floor(log2)` computed exactly via `frexp`. | OCP-recommended recipe; exact integer arithmetic. | +| D3 | **Oracle numeric profile `oracle-fp32-serial-v1`**: serial ascending-index accumulation, mul-then-add rounding (**no FMA fusion**). A strict CUDA kernel must use `__fmul_rn`/`__fadd_rn` to match, or register its own profile. | Reduction order must be pinned for byte-equality; serial ascending is auditable. | +| D4 | **LoRA inter-GEMM rounding**: `U = X·Aᵀ` rounds to BF16 before `Y = U·Bᵀ·α`; in backward, `dY·α` and `dU` also round to BF16 between GEMMs. | Matches a two-GEMM BF16 pipeline; must hold on both engines. | +| D5 | **Clamp subgradients are zero exactly at the bounds** (strict inequalities pass gradient). | Tie-break must be deterministic; pinned by tests. | +| D6 | **Shared expert applies no clamp** (`h = SiLU(gate)·up`), reusing the one-round SwiGLU with `p_s = None`, per the fixed math in P5-5 (#64). | P5-5 (#64) prose says "reuse clamp_swiglu_weighted (without p_s)" but its math shows no clamp — **open question raised on the issue**. | +| D7 | Gradients returned by backward are FP32 (the accumulator dtype); rounding at the next operator edge is BF16. | Consistent with "BF16 backward, FP32 reductions". | + +## Byte-equality scope + +Strict byte-equality is required **between train and infer on the same +numeric profile and device**. The committed manifest anchors the CPU x86 +oracle; `scripts/check_p5.py` recomputes the oracle on the provider's device, +so transcendentals (sigmoid) never cross devices inside a strict comparison. +Hardware without equivalent capability (no FP8 MMA, fnuz formats, native MX +instructions) must register its own profile with an explicit tolerance — +never silently relax (P5-4/P5-6 contract). + +## 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`. +2. Run `python scripts/check_p5.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: `base_only_one_row`, `base_only_packed`, `lora_only`, +`base_plus_lora`, `uneven_experts` (zero-row experts), `shared_t1`, +`shared_t16`, plus operator edge cases `act_quant_edges` (powers of two, +RNE ties, zero rows) and `swiglu_boundary` (values at/inside/beyond clamps). + +Regenerate the manifest after an intentional contract change: + +```bash +python -m rl_engine.moe.fixtures --write-manifest +``` + +## Non-goals of the kit + +No CUDA/Triton kernels, no Megatron/vLLM injection (P5-6), no EP transport or +combine (P4/P6), no multi-rank gates (P5-7…P5-9). `output_slot` is carried +through untouched for P6. diff --git a/rl_engine/moe/__init__.py b/rl_engine/moe/__init__.py new file mode 100644 index 00000000..cb73e1ec --- /dev/null +++ b/rl_engine/moe/__init__.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P5 start kit: MXFP4 Routed Expert + LoRA + Shared Expert contracts (issue #8).""" + +from rl_engine.moe.contract import ( + GATE_CLAMP_MAX, + ORACLE_PROFILE, + SCHEMA_VERSION, + UP_CLAMP_MAX, + UP_CLAMP_MIN, + ExpertBatch, + LoRAParams, + SharedBatch, + tensor_sha256, +) +from rl_engine.moe.mx_format import MX_BLOCK, MXTensor, mx_dequantize, mx_quantize +from rl_engine.moe.provider import ExpertProvider, ReferenceProvider, StubProvider, resolve_provider +from rl_engine.moe.trace import ExpertTrace, first_divergence + +__all__ = [ + "GATE_CLAMP_MAX", + "ORACLE_PROFILE", + "SCHEMA_VERSION", + "UP_CLAMP_MAX", + "UP_CLAMP_MIN", + "ExpertBatch", + "ExpertProvider", + "ExpertTrace", + "LoRAParams", + "MXTensor", + "MX_BLOCK", + "ReferenceProvider", + "SharedBatch", + "StubProvider", + "first_divergence", + "mx_dequantize", + "mx_quantize", + "resolve_provider", + "tensor_sha256", +] diff --git a/rl_engine/moe/contract.py b/rl_engine/moe/contract.py new file mode 100644 index 00000000..fc46c7fc --- /dev/null +++ b/rl_engine/moe/contract.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Data contracts: ExpertBatch, SharedBatch, LoRA params.""" + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass, field +from typing import Any + +import torch + +from rl_engine.moe.mx_format import MX_BLOCK, MXTensor + +SCHEMA_VERSION = "p5-expertbatch-v1" + +GATE_CLAMP_MAX = 10.0 +UP_CLAMP_MIN = -10.0 +UP_CLAMP_MAX = 10.0 + +# The oracle's numeric profile: FP32 math, serial ascending-k reduction, +# mul-then-add rounding (no FMA fusion). Kernel backends declare their own. +ORACLE_PROFILE = "oracle-fp32-serial-v1" + +ROW_GEOMETRIES = ("one-row", "packed") + + +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() + + +def mx_fingerprint(t: MXTensor) -> str: + h = hashlib.sha256() + h.update(t.elem_format.encode()) + h.update(tensor_bytes(t.codes)) + h.update(tensor_bytes(t.scales)) + return h.hexdigest() + + +@dataclass(frozen=True) +class LoRAParams: + """BF16 LoRA adapters shared across local experts (P5-3, P5-7). + + ``a1``/``b1`` insert after the packed gate/up projection (fc1) and + ``a2``/``b2`` after the down projection (fc2). Base weights stay packed; + the LoRA path never unpacks them. + """ + + a1: torch.Tensor # BF16 [r, hidden] + b1: torch.Tensor # BF16 [2*ffn, r] + a2: torch.Tensor # BF16 [r, ffn] + b2: torch.Tensor # BF16 [hidden, r] + alpha: float + + def validate(self, hidden: int, ffn: int) -> None: + for name, t in (("a1", self.a1), ("b1", self.b1), ("a2", self.a2), ("b2", self.b2)): + if t.dtype != torch.bfloat16: + raise TypeError(f"LoRA {name} must be BF16, got {t.dtype}") + rank = self.a1.shape[0] + expect = { + "a1": (rank, hidden), + "b1": (2 * ffn, rank), + "a2": (rank, ffn), + "b2": (hidden, rank), + } + for name, shape in expect.items(): + got = tuple(getattr(self, name).shape) + if got != shape: + raise ValueError(f"LoRA {name} shape {got} != expected {shape}") + + def fingerprint(self) -> str: + h = hashlib.sha256() + for t in (self.a1, self.b1, self.a2, self.b2): + h.update(tensor_bytes(t)) + h.update(repr(float(self.alpha)).encode()) + return h.hexdigest() + + +@dataclass(frozen=True) +class ExpertBatch: + """Offline routed-expert input following the P5 start-kit contract. + + Rows are already EP-dispatched and sorted by local expert: + rows ``expert_offsets[e] : expert_offsets[e + 1]`` belong to local expert + ``e``. ``p_s`` is the route weight travelling with each row and + ``output_slot`` is carried through untouched for the P6 combine. + """ + + x: torch.Tensor # BF16 [M, hidden] + expert_offsets: torch.Tensor # int32 [n_local_experts + 1] + p_s: torch.Tensor # FP32 [M] + w1: MXTensor # e2m1 [E, 2*ffn, hidden] frozen base (gate rows then up rows) + w2: MXTensor # e2m1 [E, hidden, ffn] frozen base + lora: LoRAParams | None + output_slot: torch.Tensor # int32 [M] + row_geometry: str = "packed" + schema_version: str = SCHEMA_VERSION + numeric_profile: str = ORACLE_PROFILE + weight_fingerprint: str = "" + + @property + def hidden(self) -> int: + return int(self.x.shape[1]) + + @property + def ffn(self) -> int: + return int(self.w2.shape[2]) + + @property + def rows(self) -> int: + return int(self.x.shape[0]) + + def validate(self) -> None: + if self.schema_version != SCHEMA_VERSION: + raise ValueError(f"schema {self.schema_version!r} != {SCHEMA_VERSION!r}") + if self.row_geometry not in ROW_GEOMETRIES: + raise ValueError(f"row_geometry {self.row_geometry!r} not in {ROW_GEOMETRIES}") + if self.x.dtype != torch.bfloat16: + raise TypeError(f"x must be BF16, got {self.x.dtype}") + if self.p_s.dtype != torch.float32: + raise TypeError(f"p_s must be FP32, got {self.p_s.dtype}") + if self.expert_offsets.dtype != torch.int32 or self.output_slot.dtype != torch.int32: + raise TypeError("expert_offsets/output_slot must be int32") + m, hidden = self.x.shape + if self.p_s.shape != (m,) or self.output_slot.shape != (m,): + raise ValueError("p_s/output_slot must have shape [M]") + if hidden % MX_BLOCK != 0: + raise ValueError(f"hidden {hidden} not divisible by {MX_BLOCK}") + offsets = self.expert_offsets + if int(offsets[0]) != 0 or int(offsets[-1]) != m: + raise ValueError("expert_offsets must start at 0 and end at M") + if bool((offsets[1:] < offsets[:-1]).any()): + raise ValueError("expert_offsets must be non-decreasing") + n_experts = offsets.numel() - 1 + ffn = self.ffn + if tuple(self.w1.shape) != (n_experts, 2 * ffn, hidden): + raise ValueError(f"w1 shape {self.w1.shape} != {(n_experts, 2 * ffn, hidden)}") + if tuple(self.w2.shape) != (n_experts, hidden, ffn): + raise ValueError(f"w2 shape {self.w2.shape} != {(n_experts, hidden, ffn)}") + if self.w1.elem_format != "e2m1" or self.w2.elem_format != "e2m1": + raise ValueError("base weights must be MXFP4 (e2m1)") + if self.lora is not None: + self.lora.validate(hidden, ffn) + expected = self.compute_weight_fingerprint() + if self.weight_fingerprint and self.weight_fingerprint != expected: + raise ValueError("weight_fingerprint mismatch: packed base bytes were modified") + + def compute_weight_fingerprint(self) -> str: + h = hashlib.sha256() + h.update(mx_fingerprint(self.w1).encode()) + h.update(mx_fingerprint(self.w2).encode()) + if self.lora is not None: + h.update(self.lora.fingerprint().encode()) + return h.hexdigest() + + def to(self, device: torch.device | str) -> "ExpertBatch": + lora = self.lora + if lora is not None: + lora = LoRAParams( + lora.a1.to(device), + lora.b1.to(device), + lora.a2.to(device), + lora.b2.to(device), + lora.alpha, + ) + return ExpertBatch( + x=self.x.to(device), + expert_offsets=self.expert_offsets.to(device), + p_s=self.p_s.to(device), + w1=self.w1.to(device), + w2=self.w2.to(device), + lora=lora, + output_slot=self.output_slot.to(device), + row_geometry=self.row_geometry, + schema_version=self.schema_version, + numeric_profile=self.numeric_profile, + weight_fingerprint=self.weight_fingerprint, + ) + + +@dataclass(frozen=True) +class SharedBatch: + """Shared-expert input: every valid token, no routing, no LoRA (P5-5, issue #64).""" + + x: torch.Tensor # BF16 [T, hidden] + w_fc1: torch.Tensor # BF16 [2*ffn, hidden] frozen (gate rows then up rows) + w_fc2: torch.Tensor # BF16 [hidden, ffn] frozen + placement: str = "replicated" + schema_version: str = SCHEMA_VERSION + numeric_profile: str = ORACLE_PROFILE + metadata: dict[str, Any] = field(default_factory=dict) + + def validate(self) -> None: + if self.x.dtype != torch.bfloat16: + raise TypeError(f"x must be BF16, got {self.x.dtype}") + if self.w_fc1.dtype != torch.bfloat16 or self.w_fc2.dtype != torch.bfloat16: + raise TypeError("shared weights must be BF16 in the v1 contract") + t, hidden = self.x.shape + two_ffn = self.w_fc1.shape[0] + if two_ffn % 2 != 0 or self.w_fc1.shape[1] != hidden: + raise ValueError(f"w_fc1 shape {tuple(self.w_fc1.shape)} inconsistent with x") + if tuple(self.w_fc2.shape) != (hidden, two_ffn // 2): + raise ValueError(f"w_fc2 shape {tuple(self.w_fc2.shape)} != {(hidden, two_ffn // 2)}") + if self.placement not in ("replicated", "tp-sharded"): + raise ValueError(f"unknown placement {self.placement!r}") + + def to(self, device: torch.device | str) -> "SharedBatch": + return SharedBatch( + x=self.x.to(device), + w_fc1=self.w_fc1.to(device), + w_fc2=self.w_fc2.to(device), + placement=self.placement, + schema_version=self.schema_version, + numeric_profile=self.numeric_profile, + metadata=dict(self.metadata), + ) diff --git a/rl_engine/moe/fixtures.py b/rl_engine/moe/fixtures.py new file mode 100644 index 00000000..6a6ab2d0 --- /dev/null +++ b/rl_engine/moe/fixtures.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Seeded P5 fixtures and the golden-hash manifest (start-kit acceptance data). + +Fixtures are regenerated deterministically from seeds; the committed manifest +``tests/fixtures/p5/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. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +from typing import Any + +import torch + +from rl_engine.moe import oracle +from rl_engine.moe.contract import ( + ORACLE_PROFILE, + SCHEMA_VERSION, + ExpertBatch, + LoRAParams, + SharedBatch, + tensor_sha256, +) +from rl_engine.moe.mx_format import MXTensor, mx_quantize +from rl_engine.moe.trace import ExpertTrace + +FIXTURE_HIDDEN = 128 +FIXTURE_FFN = 64 +FIXTURE_RANK = 8 +BASE_SEED = 2026 + +DEFAULT_MANIFEST_PATH = Path("tests/fixtures/p5/golden_hashes.json") + +E2E_CASES: dict[str, dict[str, Any]] = { + "base_only_one_row": {"rows": 1, "offsets": [0, 1, 1], "lora": False, "geometry": "one-row"}, + "base_only_packed": {"rows": 24, "offsets": [0, 6, 12, 18, 24], "lora": False}, + "lora_only": {"rows": 24, "offsets": [0, 6, 12, 18, 24], "lora": True, "base_zero": True}, + "base_plus_lora": {"rows": 24, "offsets": [0, 6, 12, 18, 24], "lora": True}, + "uneven_experts": {"rows": 24, "offsets": [0, 0, 17, 17, 24], "lora": True}, +} + +SHARED_CASES: dict[str, dict[str, Any]] = { + "shared_t1": {"tokens": 1}, + "shared_t16": {"tokens": 16}, +} + + +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_base_weights( + g: torch.Generator, n_experts: int, zero: bool = False +) -> tuple[MXTensor, MXTensor]: + h, f = FIXTURE_HIDDEN, FIXTURE_FFN + scale = 1.0 / float(h) ** 0.5 + w1 = _randn(g, n_experts, 2 * f, h, scale=scale) + w2 = _randn(g, n_experts, h, f, scale=1.0 / float(f) ** 0.5) + if zero: + w1 = torch.zeros_like(w1) + w2 = torch.zeros_like(w2) + return mx_quantize(w1, "e2m1"), mx_quantize(w2, "e2m1") + + +def _make_lora(g: torch.Generator) -> LoRAParams: + h, f, r = FIXTURE_HIDDEN, FIXTURE_FFN, FIXTURE_RANK + return LoRAParams( + a1=_randn(g, r, h, scale=0.1).to(torch.bfloat16), + b1=_randn(g, 2 * f, r, scale=0.1).to(torch.bfloat16), + a2=_randn(g, r, f, scale=0.1).to(torch.bfloat16), + b2=_randn(g, h, r, scale=0.1).to(torch.bfloat16), + alpha=0.5, + ) + + +def make_expert_batch(name: str) -> ExpertBatch: + spec = E2E_CASES[name] + g = _gen(name) + rows = spec["rows"] + offsets = torch.tensor(spec["offsets"], dtype=torch.int32) + n_experts = offsets.numel() - 1 + w1, w2 = _make_base_weights(g, n_experts, zero=spec.get("base_zero", False)) + lora = _make_lora(g) if spec.get("lora") else None + batch = ExpertBatch( + x=_randn(g, rows, FIXTURE_HIDDEN).to(torch.bfloat16), + expert_offsets=offsets, + p_s=torch.rand(rows, generator=g, dtype=torch.float32), + w1=w1, + w2=w2, + lora=lora, + output_slot=torch.arange(rows, dtype=torch.int32), + row_geometry=spec.get("geometry", "packed"), + ) + batch = ExpertBatch( + **{**batch.__dict__, "weight_fingerprint": batch.compute_weight_fingerprint()} + ) + batch.validate() + return batch + + +def make_shared_batch(name: str) -> SharedBatch: + spec = SHARED_CASES[name] + g = _gen(name) + h, f = FIXTURE_HIDDEN, FIXTURE_FFN + batch = SharedBatch( + x=_randn(g, spec["tokens"], h).to(torch.bfloat16), + w_fc1=_randn(g, 2 * f, h, scale=1.0 / float(h) ** 0.5).to(torch.bfloat16), + w_fc2=_randn(g, h, f, scale=1.0 / float(f) ** 0.5).to(torch.bfloat16), + ) + batch.validate() + return batch + + +def make_grad_output(name: str, shape: tuple[int, ...]) -> torch.Tensor: + g = _gen(name + ".grad") + return _randn(g, *shape).to(torch.bfloat16) + + +def make_act_quant_edge_inputs() -> torch.Tensor: + """Edge inputs for P5-1 (#60): powers of two, ties, zero rows, subnormal scales.""" + rows = [] + rows.append(torch.tensor([2.0**k for k in range(-16, 16)], dtype=torch.float32)) + rows.append(torch.tensor([17.0, 18.0, 19.0, 20.0] * 8, dtype=torch.float32)) + rows.append(torch.zeros(32, dtype=torch.float32)) + rows.append(torch.linspace(-6.0, 6.0, 32, dtype=torch.float32)) + x = torch.stack(rows) + return x.to(torch.bfloat16) + + +def make_swiglu_boundary_inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Edge inputs for P5-2 (#63): gate/up exactly at, inside, and beyond the clamps.""" + gate_vals = [-12.0, -10.0, -1.0, 0.0, 1.0, 9.5, 10.0, 10.5] + up_vals = [-10.5, -10.0, -9.5, 0.0, 0.5, 9.5, 10.0, 10.5] + gate = torch.tensor([gate_vals * 4] * 3, dtype=torch.float32) + up = torch.tensor([up_vals * 4] * 3, dtype=torch.float32) + p_s = torch.tensor([0.0, 0.5, 1.0], dtype=torch.float32) + return gate, up, p_s + + +def _mx_hashes(prefix: str, t: MXTensor) -> dict[str, str]: + return {f"{prefix}.codes": tensor_sha256(t.codes), f"{prefix}.scales": tensor_sha256(t.scales)} + + +def golden_manifest() -> dict[str, Any]: + """Recompute every golden hash from seeds with the FP32 oracle.""" + cases: dict[str, dict[str, str]] = {} + for name in E2E_CASES: + batch = make_expert_batch(name) + trace = ExpertTrace(numeric_profile=ORACLE_PROFILE) + y, saved = oracle.routed_expert_forward(batch, trace) + dy = make_grad_output(name, tuple(y.shape)) + grads = oracle.routed_expert_backward(batch, saved, dy, trace) + hashes = trace.hashes() + for key, grad in grads.items(): + if grad is not None: + hashes[f"grad.{key}"] = tensor_sha256(grad) + cases[name] = hashes + for name in SHARED_CASES: + shared = make_shared_batch(name) + y, saved = oracle.shared_expert_mlp_fwd(shared) + dy = make_grad_output(name, tuple(y.shape)) + dx = oracle.shared_expert_mlp_bwd(dy, shared, saved) + cases[name] = {"shared_out": tensor_sha256(y), "grad.dx": tensor_sha256(dx)} + q_edge = oracle.mxfp8_act_quant_fwd(make_act_quant_edge_inputs()) + cases["act_quant_edges"] = _mx_hashes("act_quant", q_edge) + gate, up, p_s = make_swiglu_boundary_inputs() + h, sw_saved = oracle.clamp_swiglu_weighted_fwd(gate, up, p_s) + dh = make_grad_output("swiglu_boundary", tuple(h.shape)) + dgate, dup, dp_s = oracle.clamp_swiglu_weighted_bwd(dh, sw_saved) + assert dp_s is not None + cases["swiglu_boundary"] = { + "h": tensor_sha256(h), + "grad.dgate": tensor_sha256(dgate), + "grad.dup": tensor_sha256(dup), + "grad.dp_s": tensor_sha256(dp_s), + } + return { + "schema_version": SCHEMA_VERSION, + "numeric_profile": ORACLE_PROFILE, + "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 main() -> None: + parser = argparse.ArgumentParser(description="P5 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: + out = write_manifest(args.path) + print(f"wrote {out}") + else: + print(json.dumps(golden_manifest(), indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/rl_engine/moe/mx_format.py b/rl_engine/moe/mx_format.py new file mode 100644 index 00000000..8b59c63b --- /dev/null +++ b/rl_engine/moe/mx_format.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Bit-exact CPU/GPU reference codecs for OCP Microscaling (MX) formats. + +Implements the P5 quantization contract (#8; sub-issues P5-1 (#60), P5-4 (#61)): + +- MX block size is fixed at 32 elements, blocked along the last dimension. +- Shared scales are E8M0 (8-bit power-of-two exponent, bias 127). +- MXFP8 elements are OCP E4M3 (torch ``float8_e4m3fn``); encode is + clamp-to-[-448, 448] followed by round-to-nearest-even ("satfinite"). +- MXFP4 elements are E2M1 with values {0, 0.5, 1, 1.5, 2, 3, 4, 6} per sign; + encode is clamp-to-[-6, 6] followed by round-to-nearest-even. +- Scale derivation: ``shared_exp = floor(log2(amax)) - emax_elem`` where + ``emax_elem`` is 8 for E4M3 and 2 for E2M1; an all-zero block gets code 127 + (scale 1.0). Non-finite inputs are rejected (fail-closed). +- FP4 codes are packed two per byte, low nibble first ("nibble-lo-first"). + +These functions define the golden bytes for the P5 fixtures; kernel backends +must reproduce them exactly or register an explicit numeric profile. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +MX_BLOCK = 32 +E8M0_BIAS = 127 +E4M3_MAX = 448.0 +E2M1_MAX = 6.0 +EMAX_ELEM = {"e4m3": 8, "e2m1": 2} +NIBBLE_PACKING = "nibble-lo-first" + +_E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +_E2M1_BOUNDARIES = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) +_E2M1_TIES_UP = (0.75, 1.75, 3.5) + + +@dataclass(frozen=True) +class MXTensor: + """A block-scaled MX tensor (codes + E8M0 scales). + + ``shape`` is the logical element shape. For ``e4m3`` the codes tensor has + exactly that shape (one byte per element); for ``e2m1`` the last dimension + of ``codes`` is halved (two nibbles per byte, low nibble first). + ``scales`` has the logical shape with the last dimension divided by 32. + """ + + codes: torch.Tensor + scales: torch.Tensor + elem_format: str + shape: tuple[int, ...] + packing: str = NIBBLE_PACKING + + def __post_init__(self) -> None: + if self.elem_format not in EMAX_ELEM: + raise ValueError(f"unsupported elem_format {self.elem_format!r}") + if self.codes.dtype != torch.uint8 or self.scales.dtype != torch.uint8: + raise TypeError("MXTensor codes/scales must be uint8") + if self.shape[-1] % MX_BLOCK != 0: + raise ValueError(f"last dim {self.shape[-1]} not divisible by MX block {MX_BLOCK}") + + def to(self, device: torch.device | str) -> "MXTensor": + return MXTensor( + self.codes.to(device), + self.scales.to(device), + self.elem_format, + self.shape, + self.packing, + ) + + +def _check_finite(x: torch.Tensor, what: str) -> None: + if not torch.isfinite(x).all(): + raise ValueError(f"non-finite values in {what}; P5 quantization is fail-closed") + + +def floor_log2(x: torch.Tensor) -> torch.Tensor: + """Exact floor(log2(x)) for positive x via frexp (no libm log2 rounding).""" + _, exp = torch.frexp(x) + return exp.to(torch.int32) - 1 + + +def e8m0_decode(code: torch.Tensor) -> torch.Tensor: + """E8M0 code -> FP32 scale = 2**(code - 127). Code 255 (NaN) is rejected.""" + if bool((code == 255).any()): + raise ValueError("E8M0 NaN code 255 is not allowed in the P5 contract") + return torch.ldexp( + torch.ones(code.shape, dtype=torch.float32, device=code.device), + code.to(torch.int32) - E8M0_BIAS, + ) + + +def e8m0_scale_from_amax(amax: torch.Tensor, elem_format: str) -> torch.Tensor: + """Derive the shared-scale code: floor(log2(amax)) - emax_elem, bias 127. + + All-zero blocks (amax == 0) get code 127 (scale 1.0). + """ + _check_finite(amax, "amax") + if bool((amax < 0).any()): + raise ValueError("amax must be non-negative") + emax = EMAX_ELEM[elem_format] + exp = floor_log2(torch.clamp(amax, min=torch.finfo(torch.float32).tiny)) - emax + exp = torch.clamp(exp, min=-E8M0_BIAS, max=E8M0_BIAS) + code = (exp + E8M0_BIAS).to(torch.uint8) + return torch.where(amax == 0, torch.full_like(code, E8M0_BIAS), code) + + +def e4m3_encode(x: torch.Tensor) -> torch.Tensor: + """FP32 -> OCP E4M3 byte codes: clamp to +/-448 then RNE cast (satfinite). + + The clamp-then-cast pair is the frozen contract; torch's bare cast maps + overflow to NaN, so the clamp must never be removed. + """ + _check_finite(x, "e4m3 input") + clamped = torch.clamp(x.to(torch.float32), min=-E4M3_MAX, max=E4M3_MAX) + return clamped.to(torch.float8_e4m3fn).view(torch.uint8) + + +def e4m3_decode(codes: torch.Tensor) -> torch.Tensor: + return codes.view(torch.float8_e4m3fn).to(torch.float32) + + +def e2m1_encode(x: torch.Tensor) -> torch.Tensor: + """FP32 -> E2M1 nibble codes (0..15, sign in bit 3), RNE with saturation.""" + _check_finite(x, "e2m1 input") + x32 = x.to(torch.float32) + sign = torch.signbit(x32) + a = torch.clamp(x32.abs(), max=E2M1_MAX) + boundaries = torch.tensor(_E2M1_BOUNDARIES, dtype=torch.float32, device=x32.device) + # side='left': exact midpoints land on the lower code ... + idx = torch.searchsorted(boundaries, a.reshape(-1), right=False).reshape(a.shape) + # ... then bump the three midpoints whose round-to-even target is the upper code. + for tie in _E2M1_TIES_UP: + idx = torch.where(a == tie, idx + 1, idx) + return (idx.to(torch.uint8)) | (sign.to(torch.uint8) << 3) + + +def e2m1_decode(codes: torch.Tensor) -> torch.Tensor: + table = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=codes.device) + mag = table[(codes & 0x7).to(torch.long)] + sign = torch.where((codes & 0x8) != 0, -1.0, 1.0).to(torch.float32) + return mag * sign + + +def pack_nibbles(codes: torch.Tensor) -> torch.Tensor: + """Pack 4-bit codes two per byte along the last dim, low nibble first.""" + if codes.shape[-1] % 2 != 0: + raise ValueError("last dim must be even to pack nibbles") + lo = codes[..., 0::2] + hi = codes[..., 1::2] + return lo | (hi << 4) + + +def unpack_nibbles(packed: torch.Tensor) -> torch.Tensor: + lo = packed & 0xF + hi = packed >> 4 + out = torch.stack((lo, hi), dim=-1) + return out.reshape(*packed.shape[:-1], packed.shape[-1] * 2) + + +def mx_quantize(x: torch.Tensor, elem_format: str) -> MXTensor: + """BF16/FP32 -> MX tensor with block-32 E8M0 scales along the last dim. + + The amax reduction is strictly within one 32-element block of one row; + it never crosses rows (and therefore never crosses ranks). + """ + if elem_format not in EMAX_ELEM: + raise ValueError(f"unsupported elem_format {elem_format!r}") + x32 = x.to(torch.float32) + _check_finite(x32, "mx_quantize input") + shape = tuple(x32.shape) + if shape[-1] % MX_BLOCK != 0: + raise ValueError(f"last dim {shape[-1]} not divisible by MX block {MX_BLOCK}") + blocked = x32.reshape(*shape[:-1], shape[-1] // MX_BLOCK, MX_BLOCK) + amax = blocked.abs().amax(dim=-1) + scale_codes = e8m0_scale_from_amax(amax, elem_format) + scale = e8m0_decode(scale_codes) + scaled = (blocked / scale.unsqueeze(-1)).reshape(shape) + if elem_format == "e4m3": + codes = e4m3_encode(scaled) + else: + codes = pack_nibbles(e2m1_encode(scaled)) + return MXTensor(codes=codes, scales=scale_codes, elem_format=elem_format, shape=shape) + + +def mx_dequantize(t: MXTensor) -> torch.Tensor: + """MX tensor -> FP32 (exact: element decode and power-of-two scale).""" + if t.elem_format == "e4m3": + elems = e4m3_decode(t.codes) + else: + elems = e2m1_decode(unpack_nibbles(t.codes)) + scale = e8m0_decode(t.scales) + blocked = elems.reshape(*t.shape[:-1], t.shape[-1] // MX_BLOCK, MX_BLOCK) + return (blocked * scale.unsqueeze(-1)).reshape(t.shape) diff --git a/rl_engine/moe/oracle.py b/rl_engine/moe/oracle.py new file mode 100644 index 00000000..37172861 --- /dev/null +++ b/rl_engine/moe/oracle.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""FP32 CPU oracle for the five P5 operators (P5-1..P5-5; issues #60-#64). + +Numeric profile ``oracle-fp32-serial-v1``: + +- All accumulations are FP32, serial, in ascending index order. +- Every multiply and add rounds separately (mul-then-add; no fused FMA). + A strict CUDA kernel must either reproduce this (``__fmul_rn``/``__fadd_rn``) + or register its own numeric profile. +- Backward is BF16 at operator boundaries (gradients round to BF16 when they + cross an operator edge) with FP32 accumulators inside, per issue #8. +- Base weights are frozen: no ``dW`` is ever computed (issue #1 s2.5 item 1). +- ``mxfp8_act_quant`` backward is a straight-through estimator (dX = dY). + +The oracle favors auditability over speed; use start-kit fixture sizes. +""" + +from __future__ import annotations + +import sys +from typing import Any + +import torch + +from rl_engine.moe.contract import ( + GATE_CLAMP_MAX, + UP_CLAMP_MAX, + UP_CLAMP_MIN, + ExpertBatch, + SharedBatch, +) +from rl_engine.moe.mx_format import ( + MX_BLOCK, + MXTensor, + e2m1_decode, + e4m3_decode, + e8m0_decode, + mx_quantize, + unpack_nibbles, +) +from rl_engine.moe.trace import ExpertTrace + + +def _serial_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]. + """ + a32 = a.to(torch.float32) + b32 = b.to(torch.float32) + m, k = a32.shape + n, kb = b32.shape + if kb != k: + raise ValueError(f"serial_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 _block_scaled_dot( + a_elems: torch.Tensor, # FP32 [M, K] decoded elements + a_scales: torch.Tensor, # FP32 [M, K/32] + w_elems: torch.Tensor, # FP32 [N, K] decoded elements + w_scales: torch.Tensor, # FP32 [N, K/32] +) -> torch.Tensor: + """P5-4 (#61) fixed math: per 32-wide chunk j, ``acc += partial_j * sa_j * sw_j``. + + ``partial_j`` is the serial ascending-k FP32 dot of the decoded elements; + the scale application order is ``(partial * scale_a) * scale_w``. + """ + m, k = a_elems.shape + n = w_elems.shape[0] + n_blocks = k // MX_BLOCK + acc = torch.zeros(m, n, dtype=torch.float32, device=a_elems.device) + for j in range(n_blocks): + partial = torch.zeros(m, n, dtype=torch.float32, device=a_elems.device) + for kk in range(j * MX_BLOCK, (j + 1) * MX_BLOCK): + partial = partial + a_elems[:, kk].unsqueeze(1) * w_elems[:, kk].unsqueeze(0) + scaled = (partial * a_scales[:, j].unsqueeze(1)) * w_scales[:, j].unsqueeze(0) + acc = acc + scaled + return acc + + +# 1. mxfp8_act_quant — P5-1 (#60) + + +def mxfp8_act_quant_fwd(x: torch.Tensor) -> MXTensor: + """BF16 [M, K] -> MXFP8 (block-32 E8M0 scales, row-local amax).""" + return mx_quantize(x, "e4m3") + + +def mxfp8_act_quant_bwd(dy: torch.Tensor) -> torch.Tensor: + """Straight-through estimator: dX = dY (not the true derivative).""" + return dy.clone() + + +# 2. mxfp8_mxfp4_grouped_gemm — P5-4 (#61) + + +def mxfp8_mxfp4_grouped_gemm_fwd( + a: MXTensor, w: MXTensor, expert_offsets: torch.Tensor +) -> torch.Tensor: + """Frozen-base grouped GEMM: MXFP8 activation x MXFP4 weight -> FP32 [M, N]. + + ``w`` holds one [N, K] weight per local expert ([E, N, K]); rows + ``expert_offsets[e] : expert_offsets[e+1]`` of ``a`` use expert ``e``. + """ + if a.elem_format != "e4m3" or w.elem_format != "e2m1": + raise ValueError("grouped GEMM expects e4m3 activation and e2m1 weight") + m, k = a.shape + n_experts, n, wk = w.shape + if wk != k: + raise ValueError(f"K mismatch: activation {k} vs weight {wk}") + a_elems = e4m3_decode(a.codes) + a_scales = e8m0_decode(a.scales) + w_elems = e2m1_decode(unpack_nibbles(w.codes)) + w_scales = e8m0_decode(w.scales) + out = torch.zeros(m, n, dtype=torch.float32, device=a_elems.device) + for e in range(n_experts): + lo, hi = int(expert_offsets[e]), int(expert_offsets[e + 1]) + if lo == hi: + continue + out[lo:hi] = _block_scaled_dot(a_elems[lo:hi], a_scales[lo:hi], w_elems[e], w_scales[e]) + return out + + +def mxfp8_mxfp4_grouped_gemm_bwd( + dy: torch.Tensor, w: MXTensor, expert_offsets: torch.Tensor +) -> torch.Tensor: + """dX = dY @ W, BF16 operands with FP32 accumulator. No dW (frozen base). + + The MXFP4 weight is dequantized to BF16 (exact: <= 2 mantissa bits times a + power-of-two scale) and the reduction runs serially over ascending n. + """ + m = dy.shape[0] + n_experts, n, k = w.shape + dy_bf16 = dy.to(torch.bfloat16) + w_elems = e2m1_decode(unpack_nibbles(w.codes)) + w_scales = e8m0_decode(w.scales) + blocked = w_elems.reshape(n_experts, n, k // MX_BLOCK, MX_BLOCK) + w_full = (blocked * w_scales.unsqueeze(-1)).reshape(n_experts, n, k) + w_bf16 = w_full.to(torch.bfloat16) + dx = torch.zeros(m, k, dtype=torch.float32, device=dy.device) + for e in range(n_experts): + lo, hi = int(expert_offsets[e]), int(expert_offsets[e + 1]) + if lo == hi: + continue + dx[lo:hi] = _serial_dot(dy_bf16[lo:hi], w_bf16[e].t()) + return dx + + +# 3. shared_grouped_lora_delta — P5-3 (#62) + + +def shared_grouped_lora_delta_fwd( + x: torch.Tensor, a: torch.Tensor, b: torch.Tensor, alpha: float +) -> tuple[torch.Tensor, torch.Tensor]: + """LoRA delta ``Y = (X @ A.T) @ B.T * alpha`` on the BF16 path. + + Returns ``(y_fp32, u_bf16)``; ``u_bf16`` is the saved inter-GEMM + activation (the intermediate rounds to BF16 between the two GEMMs). + """ + u = _serial_dot(x, a) # [M, r] FP32 + u_bf16 = u.to(torch.bfloat16) + y = _serial_dot(u_bf16, b) * float(alpha) + return y, u_bf16 + + +def shared_grouped_lora_delta_bwd( + dy: torch.Tensor, + x: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + alpha: float, + u_bf16: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Backward of the same graph: returns ``(dX, dA, dB)`` as FP32. + + ``dY' = dY * alpha`` rounds to BF16, then each GEMM runs BF16-in / + FP32-accumulate, serial ascending order; ``dU`` rounds to BF16 before + reuse. Association order is frozen as written. + """ + dys = (dy.to(torch.float32) * float(alpha)).to(torch.bfloat16) + du = _serial_dot(dys, b.t()) # [M, r]: dY' [M, N] x B [N, r] + du_bf16 = du.to(torch.bfloat16) + db = _serial_dot(dys.t(), u_bf16.t()) # [N, r] = dY'.T [N, M] x U.T [r, M] -> a @ b.T + da = _serial_dot(du_bf16.t(), x.t()) # [r, K] + dx = _serial_dot(du_bf16, a.t()) # [M, K] + return dx, da, db + + +# 4. clamp_swiglu_weighted — P5-2 (#63) + + +def clamp_swiglu_weighted_fwd( + gate: torch.Tensor, up: torch.Tensor, p_s: torch.Tensor | None +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """``h = SiLU(min(gate, 10)) * clamp(up, -10, 10) * p_s`` — one-round. + + All math in FP32; the only BF16 round is on the output. The association + order ``(SiLU(g) * u) * p_s`` is frozen. ``p_s=None`` means the unweighted + shared-expert variant (no clamp is applied in that variant, P5-5 (#64)). + """ + gate32 = gate.to(torch.float32) + up32 = up.to(torch.float32) + if p_s is None: + g = gate32 + u = up32 + else: + g = torch.clamp(gate32, max=GATE_CLAMP_MAX) + u = torch.clamp(up32, min=UP_CLAMP_MIN, max=UP_CLAMP_MAX) + sig = torch.sigmoid(g) + silu = g * sig + prod = silu * u + h32 = prod if p_s is None else prod * p_s.unsqueeze(1) + saved = {"gate32": gate32, "up32": up32, "g": g, "u": u, "sig": sig, "silu": silu} + if p_s is not None: + saved["p_s"] = p_s + return h32.to(torch.bfloat16), saved + + +def clamp_swiglu_weighted_bwd( + dh: torch.Tensor, saved: dict[str, torch.Tensor] +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Returns ``(dgate, dup, dp_s)``; ``dp_s`` is FP32 [rows] (or None). + + Clamp subgradients are zero exactly at the bounds (strict inequalities + pass gradient). ``dp_s`` is a row-local serial sum over ascending n. + """ + dh32 = dh.to(torch.float32) + g, u, sig, silu = saved["g"], saved["u"], saved["sig"], saved["silu"] + gate32, up32 = saved["gate32"], saved["up32"] + p_s = saved.get("p_s") + weighted = dh32 if p_s is None else dh32 * p_s.unsqueeze(1) + dsilu = sig * (1.0 + g * (1.0 - sig)) + if p_s is None: + gate_mask = torch.ones_like(g) + up_mask = torch.ones_like(u) + else: + gate_mask = (gate32 < GATE_CLAMP_MAX).to(torch.float32) + up_mask = ((up32 > UP_CLAMP_MIN) & (up32 < UP_CLAMP_MAX)).to(torch.float32) + dgate = ((weighted * u) * dsilu) * gate_mask + dup = (weighted * silu) * up_mask + dp_s: torch.Tensor | None = None + if p_s is not None: + rows = dh32.shape[0] + acc = torch.zeros(rows, dtype=torch.float32, device=dh32.device) + for n in range(dh32.shape[1]): + acc = acc + (dh32[:, n] * silu[:, n]) * u[:, n] + dp_s = acc + return dgate, dup, dp_s + + +# 5. shared_expert_mlp — P5-5 (#64) + + +def shared_expert_mlp_fwd( + batch: SharedBatch, +) -> tuple[torch.Tensor, dict[str, Any]]: + """Shared expert fc1 -> SwiGLU -> fc2 on every valid token. Returns (y, saved).""" + z = _serial_dot(batch.x, batch.w_fc1) # [T, 2F] FP32 + ffn = z.shape[1] // 2 + gate, up = z[:, :ffn], z[:, ffn:] + h_bf16, sw_saved = clamp_swiglu_weighted_fwd(gate, up, p_s=None) + y32 = _serial_dot(h_bf16, batch.w_fc2) # [T, H] FP32 + y = y32.to(torch.bfloat16) + saved: dict[str, Any] = {"swiglu": sw_saved, "h_bf16": h_bf16} + return y, saved + + +def shared_expert_mlp_bwd( + dy: torch.Tensor, batch: SharedBatch, saved: dict[str, Any] +) -> torch.Tensor: + """Returns dX (FP32). Shared base weights are frozen: no dW.""" + dy_bf16 = dy.to(torch.bfloat16) + dh = _serial_dot(dy_bf16, batch.w_fc2.t()).to(torch.bfloat16) # [T, F] + dgate, dup, _ = clamp_swiglu_weighted_bwd(dh, saved["swiglu"]) + dz = torch.cat([dgate, dup], dim=1).to(torch.bfloat16) # [T, 2F] + dx = _serial_dot(dz, batch.w_fc1.t()) # [T, H] FP32 + return dx + + +# Routed-expert composition (the full P5 forward/backward chain) + + +def routed_expert_forward( + batch: ExpertBatch, trace: ExpertTrace | None = None, ops: Any = None +) -> tuple[torch.Tensor, dict[str, Any]]: + """Full routed pipeline: quant -> base GEMM + LoRA -> clamp-SwiGLU(p_s) + -> quant -> base GEMM + LoRA -> BF16 routed output. Returns (y, saved).""" + ops = ops if ops is not None else sys.modules[__name__] + batch.validate() + ffn = batch.ffn + q1 = ops.mxfp8_act_quant_fwd(batch.x) + z_base = ops.mxfp8_mxfp4_grouped_gemm_fwd(q1, batch.w1, batch.expert_offsets) + if batch.lora is not None: + z_lora, u1_bf16 = ops.shared_grouped_lora_delta_fwd( + batch.x, batch.lora.a1, batch.lora.b1, batch.lora.alpha + ) + else: + z_lora = torch.zeros_like(z_base) + u1_bf16 = torch.zeros(batch.rows, 0, dtype=torch.bfloat16, device=batch.x.device) + z = z_base + z_lora + gate, up = z[:, :ffn], z[:, ffn:] + h_bf16, sw_saved = ops.clamp_swiglu_weighted_fwd(gate, up, batch.p_s) + q2 = ops.mxfp8_act_quant_fwd(h_bf16) + y_base = ops.mxfp8_mxfp4_grouped_gemm_fwd(q2, batch.w2, batch.expert_offsets) + if batch.lora is not None: + y_lora, u2_bf16 = ops.shared_grouped_lora_delta_fwd( + h_bf16, batch.lora.a2, batch.lora.b2, batch.lora.alpha + ) + else: + y_lora = torch.zeros_like(y_base) + u2_bf16 = torch.zeros(batch.rows, 0, dtype=torch.bfloat16, device=batch.x.device) + y = (y_base + y_lora).to(torch.bfloat16) + if trace is not None: + trace.note("act_quant_bwd", "ste") + trace.record("act_quant1.codes", q1.codes) + trace.record("act_quant1.scales", q1.scales) + trace.record("fc1_base", z_base) + trace.record("fc1_lora", z_lora) + trace.record("fc1_out", z) + trace.record("swiglu_h", h_bf16) + trace.record("act_quant2.codes", q2.codes) + trace.record("act_quant2.scales", q2.scales) + trace.record("fc2_base", y_base) + trace.record("fc2_lora", y_lora) + trace.record("routed_out", y) + saved: dict[str, Any] = { + "h_bf16": h_bf16, + "swiglu": sw_saved, + "u1_bf16": u1_bf16, + "u2_bf16": u2_bf16, + } + return y, saved + + +def routed_expert_backward( + batch: ExpertBatch, + saved: dict[str, Any], + dy: torch.Tensor, + trace: ExpertTrace | None = None, + ops: Any = None, +) -> dict[str, torch.Tensor | None]: + """Backward chain. Returns dx, dp_s and dA1/dB1/dA2/dB2 (None w/o LoRA). + + No base-weight gradient exists anywhere in this function (frozen base). + """ + ops = ops if ops is not None else sys.modules[__name__] + dy_bf16 = dy.to(torch.bfloat16) + dh_base = ops.mxfp8_mxfp4_grouped_gemm_bwd(dy_bf16, batch.w2, batch.expert_offsets) + if batch.lora is not None: + dh_lora, da2, db2 = ops.shared_grouped_lora_delta_bwd( + dy_bf16, + saved["h_bf16"], + batch.lora.a2, + batch.lora.b2, + batch.lora.alpha, + saved["u2_bf16"], + ) + else: + dh_lora, da2, db2 = torch.zeros_like(dh_base), None, None + dh = ops.mxfp8_act_quant_bwd((dh_base + dh_lora).to(torch.bfloat16)) # STE + dgate, dup, dp_s = ops.clamp_swiglu_weighted_bwd(dh, saved["swiglu"]) + dz = torch.cat([dgate, dup], dim=1).to(torch.bfloat16) + dx_base = ops.mxfp8_mxfp4_grouped_gemm_bwd(dz, batch.w1, batch.expert_offsets) + if batch.lora is not None: + dx_lora, da1, db1 = ops.shared_grouped_lora_delta_bwd( + dz, + batch.x, + batch.lora.a1, + batch.lora.b1, + batch.lora.alpha, + saved["u1_bf16"], + ) + else: + dx_lora, da1, db1 = torch.zeros_like(dx_base), None, None + dx = ops.mxfp8_act_quant_bwd(dx_base + dx_lora) # STE; FP32 accumulator output + if trace is not None: + trace.record("bwd.dh", dh) + trace.record("bwd.dp_s", dp_s if dp_s is not None else torch.zeros(0)) + trace.record("bwd.dz", dz) + trace.record("bwd.dx", dx) + return {"dx": dx, "dp_s": dp_s, "da1": da1, "db1": db1, "da2": da2, "db2": db2} diff --git a/rl_engine/moe/provider.py b/rl_engine/moe/provider.py new file mode 100644 index 00000000..22cffb92 --- /dev/null +++ b/rl_engine/moe/provider.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P5 provider interface (P5-6, issue #65) plus reference and stub implementations. + +A provider implements the five 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. + +Fail-closed contract: a provider must raise on unsupported input instead of +silently falling back to another implementation, and ``provenance()`` must +report the backend that actually ran. +""" + +from __future__ import annotations + +import importlib +from typing import Any, Protocol, runtime_checkable + +import torch + +from rl_engine.moe import oracle +from rl_engine.moe.contract import ORACLE_PROFILE, SharedBatch +from rl_engine.moe.mx_format import MXTensor + + +@runtime_checkable +class ExpertProvider(Protocol): + """The five P5 WS1 operators. See ``oracle`` for the frozen semantics.""" + + name: str + numeric_profile: str + + def capabilities(self) -> dict[str, Any]: ... + + def provenance(self) -> dict[str, Any]: ... + + def mxfp8_act_quant_fwd(self, x: torch.Tensor) -> MXTensor: ... + + def mxfp8_act_quant_bwd(self, dy: torch.Tensor) -> torch.Tensor: ... + + def mxfp8_mxfp4_grouped_gemm_fwd( + self, a: MXTensor, w: MXTensor, expert_offsets: torch.Tensor + ) -> torch.Tensor: ... + + def mxfp8_mxfp4_grouped_gemm_bwd( + self, dy: torch.Tensor, w: MXTensor, expert_offsets: torch.Tensor + ) -> torch.Tensor: ... + + def shared_grouped_lora_delta_fwd( + self, x: torch.Tensor, a: torch.Tensor, b: torch.Tensor, alpha: float + ) -> tuple[torch.Tensor, torch.Tensor]: ... + + def shared_grouped_lora_delta_bwd( + self, + dy: torch.Tensor, + x: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + alpha: float, + u_bf16: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ... + + def clamp_swiglu_weighted_fwd( + self, gate: torch.Tensor, up: torch.Tensor, p_s: torch.Tensor | None + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: ... + + def clamp_swiglu_weighted_bwd( + self, dh: torch.Tensor, saved: dict[str, torch.Tensor] + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: ... + + def shared_expert_mlp_fwd(self, batch: SharedBatch) -> tuple[torch.Tensor, dict[str, Any]]: ... + + def shared_expert_mlp_bwd( + self, dy: torch.Tensor, batch: SharedBatch, saved: dict[str, Any] + ) -> torch.Tensor: ... + + +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"], + } + + def provenance(self) -> dict[str, Any]: + return { + "requested_backend": self.name, + "actual_backend": self.name, + "numeric_profile": self.numeric_profile, + "torch_version": torch.__version__, + } + + mxfp8_act_quant_fwd = staticmethod(oracle.mxfp8_act_quant_fwd) + mxfp8_act_quant_bwd = staticmethod(oracle.mxfp8_act_quant_bwd) + mxfp8_mxfp4_grouped_gemm_fwd = staticmethod(oracle.mxfp8_mxfp4_grouped_gemm_fwd) + mxfp8_mxfp4_grouped_gemm_bwd = staticmethod(oracle.mxfp8_mxfp4_grouped_gemm_bwd) + shared_grouped_lora_delta_fwd = staticmethod(oracle.shared_grouped_lora_delta_fwd) + shared_grouped_lora_delta_bwd = staticmethod(oracle.shared_grouped_lora_delta_bwd) + clamp_swiglu_weighted_fwd = staticmethod(oracle.clamp_swiglu_weighted_fwd) + clamp_swiglu_weighted_bwd = staticmethod(oracle.clamp_swiglu_weighted_bwd) + shared_expert_mlp_fwd = staticmethod(oracle.shared_expert_mlp_fwd) + shared_expert_mlp_bwd = staticmethod(oracle.shared_expert_mlp_bwd) + + +class StubProvider(ReferenceProvider): + """Fail-closed placeholder: every operator raises until a backend claims it. + + This is deliberately NOT a fallback to the oracle — P5-6 (#65) forbids + silent fallback, so an unimplemented operator must be loud. + """ + + name = "stub" + numeric_profile = "unimplemented" + + @staticmethod + def _todo(issue: str) -> NotImplementedError: + return NotImplementedError( + f"P5 operator not implemented; claim it on issue {issue} " + "(fail-closed: no silent fallback to the oracle)" + ) + + def mxfp8_act_quant_fwd(self, x: torch.Tensor) -> MXTensor: + raise self._todo("P5-1 (#60)") + + def mxfp8_act_quant_bwd(self, dy: torch.Tensor) -> torch.Tensor: + raise self._todo("P5-1 (#60)") + + def mxfp8_mxfp4_grouped_gemm_fwd( + self, a: MXTensor, w: MXTensor, expert_offsets: torch.Tensor + ) -> torch.Tensor: + raise self._todo("P5-4 (#61)") + + def mxfp8_mxfp4_grouped_gemm_bwd( + self, dy: torch.Tensor, w: MXTensor, expert_offsets: torch.Tensor + ) -> torch.Tensor: + raise self._todo("P5-4 (#61)") + + def shared_grouped_lora_delta_fwd( + self, x: torch.Tensor, a: torch.Tensor, b: torch.Tensor, alpha: float + ) -> tuple[torch.Tensor, torch.Tensor]: + raise self._todo("P5-3 (#62)") + + def shared_grouped_lora_delta_bwd( + self, + dy: torch.Tensor, + x: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + alpha: float, + u_bf16: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + raise self._todo("P5-3 (#62)") + + def clamp_swiglu_weighted_fwd( + self, gate: torch.Tensor, up: torch.Tensor, p_s: torch.Tensor | None + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + raise self._todo("P5-2 (#63)") + + def clamp_swiglu_weighted_bwd( + self, dh: torch.Tensor, saved: dict[str, torch.Tensor] + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + raise self._todo("P5-2 (#63)") + + def shared_expert_mlp_fwd(self, batch: SharedBatch) -> tuple[torch.Tensor, dict[str, Any]]: + raise self._todo("P5-5 (#64)") + + def shared_expert_mlp_bwd( + self, dy: torch.Tensor, batch: SharedBatch, saved: dict[str, Any] + ) -> torch.Tensor: + raise self._todo("P5-5 (#64)") + + +def resolve_provider(spec: str) -> ExpertProvider: + """Instantiate a provider from ``"module.path:ClassName"`` (or an alias).""" + aliases = { + "reference": "rl_engine.moe.provider:ReferenceProvider", + "stub": "rl_engine.moe.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, ExpertProvider): + raise TypeError(f"{spec} does not implement the ExpertProvider protocol") + return instance diff --git a/rl_engine/moe/trace.py b/rl_engine/moe/trace.py new file mode 100644 index 00000000..05b709dc --- /dev/null +++ b/rl_engine/moe/trace.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Minimal boundary trace for P5 (first-divergence localization). + +Every operator boundary records (name, dtype, shape, sha256 of raw bytes). +This is the P5-local stand-in for the Foundation ``TraceEnvelope``. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + +from rl_engine.moe.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 ExpertTrace: + """Ordered boundary hashes plus provenance notes for one P5 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: ExpertTrace, b: ExpertTrace) -> str | None: + """Name of the first boundary whose hash differs, or None if identical.""" + for ra, rb in zip(a.records, b.records): + 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 diff --git a/scripts/check_p5.py b/scripts/check_p5.py new file mode 100755 index 00000000..9b7604fd --- /dev/null +++ b/scripts/check_p5.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P5 start-kit acceptance command (issue #8, ``P5-S0``). + +Runs a provider's operators through the frozen routed/shared pipelines 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_p5.py + python scripts/check_p5.py --provider mypkg.p5:CudaP5Provider --device cuda + python scripts/check_p5.py --cases base_plus_lora,uneven_experts --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.moe import fixtures, oracle # noqa: E402 +from rl_engine.moe.contract import tensor_sha256 # noqa: E402 +from rl_engine.moe.provider import ExpertProvider, resolve_provider # noqa: E402 +from rl_engine.moe.trace import ExpertTrace # 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 _run_e2e(provider: ExpertProvider, name: str, device: str) -> list[dict[str, Any]]: + batch = fixtures.make_expert_batch(name).to(device) + gold_trace = ExpertTrace(numeric_profile="oracle") + y_gold, saved_gold = oracle.routed_expert_forward(batch, gold_trace) + dy = fixtures.make_grad_output(name, tuple(y_gold.shape)).to(device) + grads_gold = oracle.routed_expert_backward(batch, saved_gold, dy, gold_trace) + + cand_trace = ExpertTrace(numeric_profile=provider.numeric_profile) + y_cand, saved_cand = oracle.routed_expert_forward(batch, cand_trace, ops=provider) + grads_cand = oracle.routed_expert_backward(batch, saved_cand, dy, cand_trace, ops=provider) + + golden = gold_trace.hashes() + candidate = cand_trace.hashes() + for key, grad in grads_gold.items(): + if grad is not None: + golden[f"grad.{key}"] = tensor_sha256(grad) + for key, grad in grads_cand.items(): + if grad is not None: + candidate[f"grad.{key}"] = tensor_sha256(grad) + return _compare(golden, candidate) + + +def _run_shared(provider: ExpertProvider, name: str, device: str) -> list[dict[str, Any]]: + batch = fixtures.make_shared_batch(name).to(device) + y_gold, saved_gold = oracle.shared_expert_mlp_fwd(batch) + dy = fixtures.make_grad_output(name, tuple(y_gold.shape)).to(device) + dx_gold = oracle.shared_expert_mlp_bwd(dy, batch, saved_gold) + y_cand, saved_cand = provider.shared_expert_mlp_fwd(batch) + dx_cand = provider.shared_expert_mlp_bwd(dy, batch, saved_cand) + golden = {"shared_out": tensor_sha256(y_gold), "grad.dx": tensor_sha256(dx_gold)} + candidate = {"shared_out": tensor_sha256(y_cand), "grad.dx": tensor_sha256(dx_cand)} + return _compare(golden, candidate) + + +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) + e2e_names = list(fixtures.E2E_CASES) + shared_names = list(fixtures.SHARED_CASES) + if args.cases: + wanted = set(args.cases.split(",")) + unknown = wanted - set(e2e_names) - set(shared_names) + if unknown: + parser.error(f"unknown cases: {sorted(unknown)}") + e2e_names = [n for n in e2e_names if n in wanted] + shared_names = [n for n in shared_names if n in wanted] + + report: dict[str, Any] = { + "provider": provider.name, + "device": args.device, + "provenance": provider.provenance(), + "cases": {}, + } + failed = False + for name in e2e_names + shared_names: + runner = _run_e2e if name in fixtures.E2E_CASES else _run_shared + try: + rows = runner(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 + status = "PASS" if case_ok else "FAIL" + print(f"[{status}] {name}") + for r in rows: + mark = " ok " if r["ok"] else " XX " + print(f"{mark} {r['boundary']:<24} 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/p5/golden_hashes.json b/tests/fixtures/p5/golden_hashes.json new file mode 100644 index 00000000..a6858930 --- /dev/null +++ b/tests/fixtures/p5/golden_hashes.json @@ -0,0 +1,131 @@ +{ + "cases": { + "act_quant_edges": { + "act_quant.codes": "9d7c82e684a4eac7a56315aa3da7c64422423961eead25ed5472f1b41f418e05", + "act_quant.scales": "a808551fd6e70818712b9c5311a5e403f89db35567f76f9f0735b28efc65b9f3" + }, + "base_only_one_row": { + "act_quant1.codes": "5d789158f179ebfa3eb067cf32b4344c295d762bd9edb24991fa1a8ea3540c35", + "act_quant1.scales": "2481a63c85a62cf889d2b149f1a52e985a9341750173fe01eff50cc27b5941b5", + "act_quant2.codes": "6e300a96664f323bfd1558423b39392d522fd19b2fe3492cd80057a8d8f19ee9", + "act_quant2.scales": "b84ff8057ee3a7f87deac4ae29ac59292f02e6c28f987031648011018384d888", + "bwd.dh": "276201253996250be403b2f4a4a2e22b828f0fafc9c71fb2dc04b1bb9b266e98", + "bwd.dp_s": "eb7ea4121b5b1f46ebce56e381e8fa22cdaa53345070f942544666c670fd8720", + "bwd.dx": "dce364334b2f8c03fbcd823393b18c4f15198faa18604342fe8bfd0f24c6f4be", + "bwd.dz": "1005ad59ff6629c0523d8321ace49ac714d4965cf0348dde77ed64a8bb3ba9da", + "fc1_base": "7e1bc0f05f2d437a961ed1c989ad713e6237dcdc777192459ec6f5a876775855", + "fc1_lora": "076a27c79e5ace2a3d47f9dd2e83e4ff6ea8872b3c2218f66c92b89b55f36560", + "fc1_out": "7e1bc0f05f2d437a961ed1c989ad713e6237dcdc777192459ec6f5a876775855", + "fc2_base": "d254405e4994e2189142a0b2cdcee6f7de604d823fb69b0a06e808cdaa2a24e0", + "fc2_lora": "076a27c79e5ace2a3d47f9dd2e83e4ff6ea8872b3c2218f66c92b89b55f36560", + "grad.dp_s": "eb7ea4121b5b1f46ebce56e381e8fa22cdaa53345070f942544666c670fd8720", + "grad.dx": "dce364334b2f8c03fbcd823393b18c4f15198faa18604342fe8bfd0f24c6f4be", + "routed_out": "a8d265b10da68991379eb98f567d3d954bdf9d24df5c34d9b4b0f9d1fdd97f1e", + "swiglu_h": "83c0f79e20832592f4a3fb39e3e20a68f0e58f0d53da90f863a7ec95902d2e17" + }, + "base_only_packed": { + "act_quant1.codes": "9ffa0ff276b2c8473d6aed5f6751416e2520d24b4a6e0b5b9c98e929057daf60", + "act_quant1.scales": "3f8aa66fb0887020cd49fa86e419eedf54dc9f0016f04ab8d3ef1a34a32bde6e", + "act_quant2.codes": "49ea25a7a0711c7306bdaf2ed2928010753e1109ac1d3e43abf04c34b12f7d4d", + "act_quant2.scales": "ad7f43b595153943abc826462baf8a49f455044719ae03fba4fac4cd4a318cc0", + "bwd.dh": "51753222bedf9cd386a1beac42d7bc7559de4511ff8b6c7da06c56feee868b9b", + "bwd.dp_s": "ec946b9c85f1a10353dd49efc53527731c1aa2db9e5c02ee0ef888e39ec8d8d8", + "bwd.dx": "b778a61bbe218b59b74e9ba2ed494ad845551e869f58202cb9a754a2cbb1507c", + "bwd.dz": "7ce5feb6387be036bd5438a6768716cb8e4de0a517b22bd68bf302bc90b8a527", + "fc1_base": "d881f71a860bcef3b74c9cc008a3690451a804bcc35a71d32160e88570674cc3", + "fc1_lora": "f3cc103136423a57975750907ebc1d367e2985ac6338976d4d5a439f50323f4a", + "fc1_out": "d881f71a860bcef3b74c9cc008a3690451a804bcc35a71d32160e88570674cc3", + "fc2_base": "86fde473ad93a204e15e9a3734abe46c1acf02e018d2ba2dbaa4826217940149", + "fc2_lora": "f3cc103136423a57975750907ebc1d367e2985ac6338976d4d5a439f50323f4a", + "grad.dp_s": "ec946b9c85f1a10353dd49efc53527731c1aa2db9e5c02ee0ef888e39ec8d8d8", + "grad.dx": "b778a61bbe218b59b74e9ba2ed494ad845551e869f58202cb9a754a2cbb1507c", + "routed_out": "a309aef1f494bfe1df5e9638a0fe5148db528165ce8f3a33e959f44ab22e41d6", + "swiglu_h": "02e0d62cdd538c6b50c76731b55f85a56bf5ce5e0fea0137e2f79684c1b5a7fb" + }, + "base_plus_lora": { + "act_quant1.codes": "d9672b8b4d8d79070b9166990ecd1042a9c12602bdeccd3c1c453b2ffde7bbbb", + "act_quant1.scales": "176927dd58e36a96aa3f9b6162a66413bd7f058932f9e4b0905108480577747f", + "act_quant2.codes": "3788abe9c822af3efe37ee8c66d435619da7d4ad650d0df7041e59773cf33e0e", + "act_quant2.scales": "2ed848e3fdae653b14ce3915f4edba2df321a88df5992814e841fbbb399a11d9", + "bwd.dh": "abf279557604bd392b9664a4fe9e5a20cf082df98513cc31a985c85ad1052e3e", + "bwd.dp_s": "4b7a810586e5c3fb905d0cb6f6996b0ecf2af18bafe3e422b6526a5a7974ded9", + "bwd.dx": "ec3664978e8b4e447d743d1f08d1945e4c748c4f68c7ba003f802474dd2ae04d", + "bwd.dz": "993acb067ab8376093067ecb3d88e26748f22f9ea28ec49bf45eab83259e1b05", + "fc1_base": "817efa836dc92c6dea7c61184eb14660f4cd62024df023ca8fe84c413ff71e4f", + "fc1_lora": "368e8ce9d0116c0a6670e7e9c9a873a5e5d48a112429a22a2c89b1cd3a3d4cb6", + "fc1_out": "f05ab64f4f58aaee4dc89a350cbf00ddc03ada2d6438f46530d539cebe569487", + "fc2_base": "e8ca08aa5cbec7db19fcc79f14f2208b3c1c02772ad585422fe6d5b82e68bd3c", + "fc2_lora": "4ebeb8e7a8a1d738729a8c59d889f0f76ba45386dd0a2573571683fad8f68a0b", + "grad.da1": "64d9a3b07636755f27396cdfcfdcdb24e9940f88b6b688f99b1a3ffd481b46cf", + "grad.da2": "57c6f284351d89893629528bf38e856aee006563cfa1ac04983947ac412e7415", + "grad.db1": "877722b8b05d632f5ce90723e64473a4fa597086edb5474079c76fb4d778dfc3", + "grad.db2": "57c01c26ef999be5bac73805d1b5f5e4d5251df4d5929a3480b19425db4cfd8e", + "grad.dp_s": "4b7a810586e5c3fb905d0cb6f6996b0ecf2af18bafe3e422b6526a5a7974ded9", + "grad.dx": "ec3664978e8b4e447d743d1f08d1945e4c748c4f68c7ba003f802474dd2ae04d", + "routed_out": "80f3ec12e28d797743ee95e8a16ad2dc1892d88ca3eee71591f2c01452932eb0", + "swiglu_h": "5affe1b53585c3a2e7014a2afc06d155fdec8781805aa215d0805ac3e125ced1" + }, + "lora_only": { + "act_quant1.codes": "422b8f3e0e2b4f97b2b1be92df05db69830cf636a406c85680680685a79b3b5f", + "act_quant1.scales": "6b53937c40065a9f376d10a51d1db0caeb897751c19be6d89e4e58305e3fdb40", + "act_quant2.codes": "f211478f9783e4434e125b3be97383116963a1a0e18a9e174e70464c8770c886", + "act_quant2.scales": "f908ef8f11ccf59cc2ffa5db7ec2992cc0f3fcd46eac4ed19ed7a0095f1430fb", + "bwd.dh": "40a5697c5b81cc848f18f4ea67a609e6221ae75853167dcc1bcd0502c8591599", + "bwd.dp_s": "40b86e55671689fa5f3078de28ca43fcffc356dd875b702ed41644a4242a03c5", + "bwd.dx": "25101a96377255effe430945b22fe8a220fd6b110e142b10b856cc0fcc33f6be", + "bwd.dz": "6298d841a39ea8e282606651bc2e1da0639eef0492d6df1d63430c03f407cc9d", + "fc1_base": "f3cc103136423a57975750907ebc1d367e2985ac6338976d4d5a439f50323f4a", + "fc1_lora": "cc1b0485d94a8b6559b7f8ea1188f6f15db161538bde0eac9cca9d3df20afd4f", + "fc1_out": "cc1b0485d94a8b6559b7f8ea1188f6f15db161538bde0eac9cca9d3df20afd4f", + "fc2_base": "f3cc103136423a57975750907ebc1d367e2985ac6338976d4d5a439f50323f4a", + "fc2_lora": "9d7634d0b95f1086ba085776ad5a1a2a5fe2d15de4608260e5246283f875d688", + "grad.da1": "c6defa6c9228ef37d5cafc033afbdd05f6171f64133a48fe45dd46c5cde0b6fd", + "grad.da2": "b4639c570429264fd6de63f9eb4e29ca195c677bf3c07ebda5fff2d2aa251720", + "grad.db1": "074895e352a26a5842927598bc904d26fbeb79621d4548198ef1e8b2874ff6f3", + "grad.db2": "9fcd531a06bdc67fce83bda21a618aeec9c120d63e70cf86e404aab0af21b453", + "grad.dp_s": "40b86e55671689fa5f3078de28ca43fcffc356dd875b702ed41644a4242a03c5", + "grad.dx": "25101a96377255effe430945b22fe8a220fd6b110e142b10b856cc0fcc33f6be", + "routed_out": "0f56a26529925493f19c57ef34969a81dc98e09310b0066652c86dc3de2d1543", + "swiglu_h": "ea2b680537330fc6d719e832e5f9fb50104533fa533ea4b6e0ae82a132abff78" + }, + "shared_t1": { + "grad.dx": "b1b4a19f22afd7dadb89c0d1da01073ac574bc5f5e69610d8636575280f3ef79", + "shared_out": "bd80378faee822ab834009694ef642ad4d6e631851dfb5360303a9672a75339d" + }, + "shared_t16": { + "grad.dx": "9ab8822677e4788142f6bd434d761d3afc6cb3fc988c8cf49ac4b0c1b662e88f", + "shared_out": "7ebe849ef5c7f3a335a937b2dd60478b7e55e31e9956ac675018e7713d941ae1" + }, + "swiglu_boundary": { + "grad.dgate": "f1a24145119559f9f744fb60a57485d01c4101c369c599cf85201f73317eeec2", + "grad.dp_s": "d60ea4c2e20f5c76c313edab01998b4ac526680eab9fbdbfd8585c1d02862d3a", + "grad.dup": "cff872ede58f268a84af736ec3c0b322d872c8aba2b4cc7cdf1278d646dac1ab", + "h": "4cac169bc0184dad9a83592673d05d14e0c039b8e8c070b1c3098e293d465a0f" + }, + "uneven_experts": { + "act_quant1.codes": "9c556c57c18a95b01c717ab71dae174879fd835ef91d7229e5c85ab79ef36d1f", + "act_quant1.scales": "c1a7f639395b5c174035a82dfc3df4876d003248005a5d23af4074999a51ca24", + "act_quant2.codes": "7e3ee0a5c87b49899c65a60829d8a69cc087e2f3c86b3062771c3e1acf22c005", + "act_quant2.scales": "76f1cb4075c0bf9d7aa816a60044be2b2b4bb0897334c4d9c0f03e5f16e377e9", + "bwd.dh": "b2c713e15085f4b50b72f7bfc8c7a75e588e32430a547eb8604fcf4609590386", + "bwd.dp_s": "53460e952d399b890e3e0ddee74e20283fabf2e5437b6780617adff47bb8cf0e", + "bwd.dx": "cbfe75c9f59a210f38a18863420734d4d9b099333b55cd0141ac8ea5adaa58ce", + "bwd.dz": "85e98c269dcff60b9bbe4c0fb93b5eba1930bc4d0e75b77a02766919bf0f34d5", + "fc1_base": "d9cabad3dc59e60f784aa152f68dc71c2e98d84f6747b890dbcd03956b71c515", + "fc1_lora": "9c55d939fd70abaae270a7cd023eab18bf8c2368914687a3be602e835febf82a", + "fc1_out": "c5c02f9024cf2059de4898f0e56f6b7914d1e12971ff9ef54877368daa185b27", + "fc2_base": "6d636bbeb973deedd7a0ac2dfabd953a6f18d4d8aa905669247d2e802fd6eefe", + "fc2_lora": "418b3ac74b8cbf7c5f9f1a36e6dbf8acc658beba7f23c4429d44c67f920a4b06", + "grad.da1": "3a33d6b915f81fc2ba1bfbae0eded5daae16bb582cedce8bc39b9c5c80111942", + "grad.da2": "605c63c781acc8167aa902e5d7ab7f37eb5e50ab8fa49750fa22c150ed02ccdf", + "grad.db1": "8eaafec3175f025d125a4d3d2ceffc4fbcf18317e126a95ad5c7e421093d9b3b", + "grad.db2": "10a893ba3c0a5626decdbd7e5a0eaac63850c917afcf080605bf404a8484edb4", + "grad.dp_s": "53460e952d399b890e3e0ddee74e20283fabf2e5437b6780617adff47bb8cf0e", + "grad.dx": "cbfe75c9f59a210f38a18863420734d4d9b099333b55cd0141ac8ea5adaa58ce", + "routed_out": "d659476793e37fb88fd5a34db732d41df9f5b0d8495650cd0cb3226b5dae137d", + "swiglu_h": "b50899008e43cb95903c686bc64e18cd88e79c46bc0d8496f7dfefb7f216f658" + } + }, + "numeric_profile": "oracle-fp32-serial-v1", + "schema_version": "p5-expertbatch-v1" +} diff --git a/tests/test_p5_contract.py b/tests/test_p5_contract.py new file mode 100644 index 00000000..993a60e6 --- /dev/null +++ b/tests/test_p5_contract.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P5 schema, fingerprint, and trace tests (issue #8 contracts).""" + +from __future__ import annotations + +import dataclasses + +import pytest +import torch + +from rl_engine.moe import fixtures +from rl_engine.moe.contract import SCHEMA_VERSION, tensor_sha256 +from rl_engine.moe.trace import ExpertTrace, first_divergence + + +def test_fixture_batches_validate() -> None: + for name in fixtures.E2E_CASES: + batch = fixtures.make_expert_batch(name) + assert batch.schema_version == SCHEMA_VERSION + batch.validate() + for name in fixtures.SHARED_CASES: + fixtures.make_shared_batch(name).validate() + + +def test_weight_fingerprint_detects_tampering() -> None: + batch = fixtures.make_expert_batch("base_plus_lora") + batch.validate() + batch.w1.codes[0, 0, 0] ^= 0xFF # tamper one packed byte + with pytest.raises(ValueError, match="fingerprint"): + batch.validate() + + +def test_bad_offsets_and_dtypes_fail_closed() -> None: + batch = fixtures.make_expert_batch("base_only_packed") + bad = dataclasses.replace(batch, expert_offsets=torch.tensor([0, 30, 24], dtype=torch.int32)) + with pytest.raises(ValueError): + bad.validate() + bad2 = dataclasses.replace(batch, p_s=batch.p_s.to(torch.bfloat16)) + with pytest.raises(TypeError): + bad2.validate() + + +def test_batch_serialization_roundtrip(tmp_path) -> None: + batch = fixtures.make_expert_batch("base_plus_lora") + path = tmp_path / "batch.pt" + torch.save(batch, path) + loaded = torch.load(path, weights_only=False) + loaded.validate() + assert tensor_sha256(loaded.x) == tensor_sha256(batch.x) + assert loaded.weight_fingerprint == batch.weight_fingerprint + + +def test_trace_first_divergence() -> None: + a = ExpertTrace(numeric_profile="p") + b = ExpertTrace(numeric_profile="p") + t1 = torch.arange(4, dtype=torch.float32) + t2 = torch.arange(4, dtype=torch.float32) + 1 + a.record("s1", t1) + a.record("s2", t1) + b.record("s1", t1) + b.record("s2", t2) + assert first_divergence(a, b) == "s2" + b.records[1] = a.records[1] + assert first_divergence(a, b) is None diff --git a/tests/test_p5_mx_format.py b/tests/test_p5_mx_format.py new file mode 100644 index 00000000..90c5c18e --- /dev/null +++ b/tests/test_p5_mx_format.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Golden-value tests for the P5 MX codecs (P5-1 (#60) contract).""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.moe import mx_format as mx + + +def test_e4m3_golden_codes() -> None: + vals = [448.0, 464.0, 500.0, 17.0, 18.0, 19.0, -464.0, 2**-9, 2**-10, 1.5 * 2**-9, 0.0] + want = [0x7E, 0x7E, 0x7E, 0x58, 0x59, 0x5A, 0xFE, 0x01, 0x00, 0x02, 0x00] + codes = mx.e4m3_encode(torch.tensor(vals, dtype=torch.float32)) + assert codes.tolist() == want + + +def test_e4m3_rejects_non_finite() -> None: + with pytest.raises(ValueError): + mx.e4m3_encode(torch.tensor([float("nan")])) + with pytest.raises(ValueError): + mx.e4m3_encode(torch.tensor([float("inf")])) + + +def test_e4m3_roundtrip_all_finite_codes() -> None: + codes = torch.arange(256, dtype=torch.uint8) + finite = (codes & 0x7F) != 0x7F # exclude NaN codes + decoded = mx.e4m3_decode(codes[finite]) + re_encoded = mx.e4m3_encode(decoded) + assert torch.equal(re_encoded, codes[finite]) + + +def test_e8m0_scale_recipe() -> None: + amax = torch.tensor([1.0, 448.0, 0.0, 2.0**-10]) + codes = mx.e8m0_scale_from_amax(amax, "e4m3") + # floor(log2(amax)) - 8, bias 127; amax==0 -> 127 + assert codes.tolist() == [127 - 8, 127, 127, 127 - 10 - 8] + codes4 = mx.e8m0_scale_from_amax(torch.tensor([1.0]), "e2m1") + assert codes4.tolist() == [127 - 2] + with pytest.raises(ValueError): + mx.e8m0_decode(torch.tensor([255], dtype=torch.uint8)) + + +def test_e2m1_tie_to_even_and_roundtrip() -> None: + ties = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0]) + assert mx.e2m1_encode(ties).tolist() == [0, 2, 2, 4, 4, 6, 6] + values = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + for sign in (1.0, -1.0): + codes = mx.e2m1_encode(values * sign) + assert torch.equal(mx.e2m1_decode(codes), values * sign) + # saturation + assert mx.e2m1_encode(torch.tensor([100.0, -100.0])).tolist() == [7, 15] + + +def test_nibble_pack_roundtrip() -> None: + g = torch.Generator().manual_seed(0) + codes = torch.randint(0, 16, (4, 32), generator=g, dtype=torch.uint8) + assert torch.equal(mx.unpack_nibbles(mx.pack_nibbles(codes)), codes) + # low nibble first + packed = mx.pack_nibbles(torch.tensor([[0x1, 0x2]], dtype=torch.uint8)) + assert packed.tolist() == [[0x21]] + + +def test_mx_quantize_row_invariant() -> None: + g = torch.Generator().manual_seed(1) + x = (torch.randn(8, 64, generator=g)).to(torch.bfloat16) + for fmt in ("e4m3", "e2m1"): + full = mx.mx_quantize(x, fmt) + one = mx.mx_quantize(x[3:4], fmt) + assert torch.equal(full.codes[3:4], one.codes) + assert torch.equal(full.scales[3:4], one.scales) + + +def test_mx_quantize_error_bounds_and_validation() -> None: + g = torch.Generator().manual_seed(2) + x = torch.randn(4, 64, generator=g).to(torch.bfloat16) + d8 = mx.mx_dequantize(mx.mx_quantize(x, "e4m3")) + d4 = mx.mx_dequantize(mx.mx_quantize(x, "e2m1")) + scale = x.float().abs().max() + # The OCP floor(log2) recipe saturates the top (448,512)*scale band, so the + # worst error at block amax is 12.5% for e4m3 (25% for e2m1) plus rounding. + assert (d8 - x.float()).abs().max() / scale < 0.13 + assert (d4 - x.float()).abs().max() / scale < 0.30 + with pytest.raises(ValueError): + mx.mx_quantize(torch.randn(4, 33), "e4m3") + with pytest.raises(ValueError): + mx.mx_quantize(torch.full((1, 32), float("inf")), "e4m3") diff --git a/tests/test_p5_oracle.py b/tests/test_p5_oracle.py new file mode 100644 index 00000000..feb6f7b9 --- /dev/null +++ b/tests/test_p5_oracle.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Self-consistency tests for the P5 FP32 oracle (P5-1..P5-5; issues #60-#64).""" + +from __future__ import annotations + +import torch + +from rl_engine.moe import fixtures, oracle +from rl_engine.moe.contract import tensor_sha256 + + +def test_act_quant_bwd_is_ste() -> None: + dy = torch.randn(4, 32) + dx = oracle.mxfp8_act_quant_bwd(dy) + assert torch.equal(dx, dy) and dx is not dy + + +def test_swiglu_bwd_matches_autograd_away_from_clamps() -> None: + g = torch.Generator().manual_seed(3) + gate = (torch.randn(4, 32, generator=g) * 2.0).requires_grad_(True) + up = (torch.randn(4, 32, generator=g) * 2.0).requires_grad_(True) + p_s = torch.rand(4, generator=g).requires_grad_(True) + ref = ( + torch.nn.functional.silu(torch.clamp(gate, max=10.0)) + * torch.clamp(up, -10.0, 10.0) + * p_s.unsqueeze(1) + ) + dh = torch.randn(4, 32, generator=g).to(torch.bfloat16) + ref.backward(dh.to(torch.float32)) + h, saved = oracle.clamp_swiglu_weighted_fwd(gate.detach(), up.detach(), p_s.detach()) + dgate, dup, dp_s = oracle.clamp_swiglu_weighted_bwd(dh, saved) + assert torch.allclose(dgate, gate.grad, atol=1e-5) + assert torch.allclose(dup, up.grad, atol=1e-5) + assert torch.allclose(dp_s, p_s.grad, atol=1e-4) + + +def test_swiglu_clamp_subgradient_zero_at_bounds() -> None: + gate = torch.tensor([[10.0, 10.5, 9.5]]) + up = torch.tensor([[-10.0, 10.0, 5.0]]) + p_s = torch.ones(1) + _, saved = oracle.clamp_swiglu_weighted_fwd(gate, up, p_s) + dh = torch.ones(1, 3, dtype=torch.bfloat16) + dgate, dup, _ = oracle.clamp_swiglu_weighted_bwd(dh, saved) + assert dgate[0, 0] == 0.0 and dgate[0, 1] == 0.0 and dgate[0, 2] != 0.0 + assert dup[0, 0] == 0.0 and dup[0, 1] == 0.0 and dup[0, 2] != 0.0 + + +def test_route_weight_applied_exactly_once() -> None: + gate = torch.full((2, 32), 1.5) + up = torch.full((2, 32), 2.0) + h1, _ = oracle.clamp_swiglu_weighted_fwd(gate, up, torch.tensor([1.0, 1.0])) + h2, _ = oracle.clamp_swiglu_weighted_fwd(gate, up, torch.tensor([2.0, 2.0])) + assert torch.allclose(h2.float(), h1.float() * 2.0, rtol=1e-2) + + +def test_lora_bwd_matches_autograd() -> None: + g = torch.Generator().manual_seed(4) + x = torch.randn(6, 32, generator=g).to(torch.bfloat16) + a = (torch.randn(4, 32, generator=g) * 0.2).to(torch.bfloat16).requires_grad_(True) + b = (torch.randn(16, 4, generator=g) * 0.2).to(torch.bfloat16).requires_grad_(True) + xg = x.detach().clone().requires_grad_(True) + y_ref = (xg.float() @ a.float().t() @ b.float().t()) * 0.5 + dy = torch.randn(6, 16, generator=g).to(torch.bfloat16) + y_ref.backward(dy.float()) + y, u = oracle.shared_grouped_lora_delta_fwd(x, a.detach(), b.detach(), 0.5) + dx, da, db = oracle.shared_grouped_lora_delta_bwd(dy, x, a.detach(), b.detach(), 0.5, u) + + def _close(got: torch.Tensor, want: torch.Tensor) -> bool: + # BF16 inter-GEMM rounding => compare normalized to the tensor scale. + return bool((got - want).abs().max() <= 2e-2 * want.abs().max() + 1e-6) + + assert _close(y, y_ref) + assert _close(dx, xg.grad.float()) + assert _close(da, a.grad.float()) + assert _close(db, b.grad.float()) + assert bool(da.abs().sum() > 0) and bool(db.abs().sum() > 0) + + +def test_geometry_gate_one_row_equals_packed() -> None: + """P5-4 (#61) acceptance: row-count=1 and packed multi-row give equal bytes.""" + import dataclasses + + batch = fixtures.make_expert_batch("base_plus_lora") + y_packed, _ = oracle.routed_expert_forward(batch) + offsets = batch.expert_offsets.tolist() + for row in range(batch.rows): + expert = sum(1 for o in offsets[1:-1] if o <= row) + single = dataclasses.replace( + batch, + x=batch.x[row : row + 1], + p_s=batch.p_s[row : row + 1], + output_slot=batch.output_slot[row : row + 1], + expert_offsets=torch.tensor( + [0] * (expert + 1) + [1] * (len(offsets) - expert - 1), dtype=torch.int32 + ), + row_geometry="one-row", + ) + y_one, _ = oracle.routed_expert_forward(single) + assert torch.equal(y_one[0], y_packed[row]), f"row {row} diverges from one-row" + + +def test_backward_has_no_dw_and_leaves_base_untouched() -> None: + batch = fixtures.make_expert_batch("base_plus_lora") + before = tensor_sha256(batch.w1.codes) + y, saved = oracle.routed_expert_forward(batch) + dy = fixtures.make_grad_output("t", tuple(y.shape)) + grads = oracle.routed_expert_backward(batch, saved, dy) + assert set(grads) == {"dx", "dp_s", "da1", "db1", "da2", "db2"} + assert tensor_sha256(batch.w1.codes) == before + assert grads["dp_s"] is not None and grads["dp_s"].shape == (batch.rows,) + assert grads["dp_s"].dtype == torch.float32 + for key in ("da1", "db1", "da2", "db2"): + grad = grads[key] + assert grad is not None and torch.isfinite(grad).all() and bool(grad.abs().sum() > 0) + + +def test_base_only_has_no_lora_grads() -> None: + batch = fixtures.make_expert_batch("base_only_packed") + y, saved = oracle.routed_expert_forward(batch) + grads = oracle.routed_expert_backward( + batch, saved, fixtures.make_grad_output("t2", tuple(y.shape)) + ) + assert grads["da1"] is None and grads["db2"] is None + + +def test_shared_expert_batch_invariant() -> None: + import dataclasses + + batch = fixtures.make_shared_batch("shared_t16") + y_full, _ = oracle.shared_expert_mlp_fwd(batch) + one = dataclasses.replace(batch, x=batch.x[5:6]) + y_one, _ = oracle.shared_expert_mlp_fwd(one) + assert torch.equal(y_one[0], y_full[5]) + + +def test_shared_bwd_matches_autograd() -> None: + batch = fixtures.make_shared_batch("shared_t16") + x = batch.x.float().requires_grad_(True) + z = x @ batch.w_fc1.float().t() + ffn = z.shape[1] // 2 + y_ref = (torch.nn.functional.silu(z[:, :ffn]) * z[:, ffn:]) @ batch.w_fc2.float().t() + y, saved = oracle.shared_expert_mlp_fwd(batch) + dy = fixtures.make_grad_output("sg", tuple(y.shape)) + y_ref.backward(dy.float()) + dx = oracle.shared_expert_mlp_bwd(dy, batch, saved) + assert (dx - x.grad).abs().max() <= 2e-2 * x.grad.abs().max() diff --git a/tests/test_p5_provider.py b/tests/test_p5_provider.py new file mode 100644 index 00000000..eb3512ce --- /dev/null +++ b/tests/test_p5_provider.py @@ -0,0 +1,58 @@ +# 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 pytest +import torch + +from rl_engine.moe import fixtures, oracle +from rl_engine.moe.provider import ReferenceProvider, StubProvider, resolve_provider +from rl_engine.moe.trace import ExpertTrace, first_divergence + + +def test_reference_provider_matches_oracle_bytes() -> None: + batch = fixtures.make_expert_batch("base_plus_lora") + gold, cand = ExpertTrace("a"), ExpertTrace("b") + _, saved_g = oracle.routed_expert_forward(batch, gold) + _, saved_c = oracle.routed_expert_forward(batch, cand, ops=ReferenceProvider()) + assert first_divergence(gold, cand) is None + dy = fixtures.make_grad_output("p", (batch.rows, batch.hidden)) + grads_g = oracle.routed_expert_backward(batch, saved_g, dy) + grads_c = oracle.routed_expert_backward(batch, saved_c, dy, ops=ReferenceProvider()) + for key, grad in grads_g.items(): + other = grads_c[key] + assert (grad is None) == (other is None) + if grad is not None: + assert torch.equal(grad, other) + + +def test_stub_provider_fails_closed() -> None: + stub = StubProvider() + with pytest.raises(NotImplementedError, match="#60"): + stub.mxfp8_act_quant_fwd(torch.zeros(1, 32, dtype=torch.bfloat16)) + batch = fixtures.make_expert_batch("base_only_one_row") + with pytest.raises(NotImplementedError): + oracle.routed_expert_forward(batch, ops=stub) + + +def test_resolve_provider() -> None: + assert resolve_provider("reference").name == "reference" + assert resolve_provider("rl_engine.moe.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_golden_manifest_anchor() -> None: + """CI anchor: regenerated golden hashes must match the committed manifest. + + A failure here means the oracle's bytes drifted (torch RNG/libm change or + an intentional contract change) — regenerate with + ``python -m rl_engine.moe.fixtures --write-manifest`` and review the diff. + """ + committed = fixtures.load_manifest() + regenerated = fixtures.golden_manifest() + assert committed == regenerated From b95ba80bba1296202d23afc1a60c6ee2f49bf8ee Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Tue, 1 Sep 2026 13:27:16 +0800 Subject: [PATCH 2/2] ci: fix black line-length drift in vllm_runtime/flash_attn; pin black/isort config in pyproject Signed-off-by: KJLdefeated --- pyproject.toml | 109 ++++++++++-------- rl_engine/integrations/vllm_runtime.py | 6 +- .../kernels/ops/cuda/attention/flash_attn.py | 10 +- 3 files changed, 61 insertions(+), 64 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ca3b0c5d..216eec05 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,55 +1,62 @@ -[build-system] -requires = ["setuptools>=64", "wheel"] -build-backend = "setuptools.build_meta" - -[project] -name = "RL-Kernel" -version = "0.1.0" -description = "High-performance RL training engine focused on kernel fusion and memory efficiency." -readme = "README.md" -requires-python = ">=3.10" -license = {text = "Apache-2.0"} -authors = [ - {name = "RL-Kernel Contributors"} -] -dependencies = [ - "torch>=2.4.1", - "tabulate", - "numpy", - "accelerate", - "transformers==5.13.1", -] - -[project.entry-points."vllm.general_plugins"] -rl_kernel = "rl_engine.integrations.vllm_runtime:register_vllm_plugin" - -[project.optional-dependencies] -cuda = ["flashinfer-python>=0.1.6", "nvidia-ml-py"] +[build-system] +requires = ["setuptools>=64", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "RL-Kernel" +version = "0.1.0" +description = "High-performance RL training engine focused on kernel fusion and memory efficiency." +readme = "README.md" +requires-python = ">=3.10" +license = {text = "Apache-2.0"} +authors = [ + {name = "RL-Kernel Contributors"} +] +dependencies = [ + "torch>=2.4.1", + "tabulate", + "numpy", + "accelerate", + "transformers==5.13.1", +] + +[project.entry-points."vllm.general_plugins"] +rl_kernel = "rl_engine.integrations.vllm_runtime:register_vllm_plugin" + +[project.optional-dependencies] +cuda = ["flashinfer-python>=0.1.6", "nvidia-ml-py"] rocm = ["aiter"] vllm = ["vllm>=0.6.0"] drift-viewer = ["Pillow>=10", "PySide6>=6.6"] dev = ["pytest", "black", "isort", "ruff", "mypy", "pre-commit"] - -[tool.setuptools.packages.find] -where = ["."] -include = ["rl_engine*"] - -[tool.ruff] -line-length = 100 - -[tool.ruff.lint] -select = ["E", "F", "B"] -ignore = [] - -[tool.ruff.lint.per-file-ignores] -"__init__.py" = ["F401"] - -[tool.mypy] -ignore_missing_imports = true -follow_imports = "silent" - -[tool.pytest.ini_options] -markers = [ - "smoke_operator: temporary smoke-only operator plumbing tests", - "unit: CPU-safe unit tests", -] + +[tool.setuptools.packages.find] +where = ["."] +include = ["rl_engine*"] + +[tool.black] +line-length = 100 + +[tool.isort] +profile = "black" +line_length = 100 + +[tool.ruff] +line-length = 100 + +[tool.ruff.lint] +select = ["E", "F", "B"] +ignore = [] + +[tool.ruff.lint.per-file-ignores] +"__init__.py" = ["F401"] + +[tool.mypy] +ignore_missing_imports = true +follow_imports = "silent" + +[tool.pytest.ini_options] +markers = [ + "smoke_operator: temporary smoke-only operator plumbing tests", + "unit: CPU-safe unit tests", +] diff --git a/rl_engine/integrations/vllm_runtime.py b/rl_engine/integrations/vllm_runtime.py index ade13cab..6351f0ab 100644 --- a/rl_engine/integrations/vllm_runtime.py +++ b/rl_engine/integrations/vllm_runtime.py @@ -475,11 +475,7 @@ def strict_row_parallel_forward(instance: Any, input_: torch.Tensor) -> Any: )[instance.tp_rank].contiguous() assert instance.quant_method is not None - bias_ = ( - None - if (instance.tp_rank > 0 or instance.skip_bias_add) - else instance.bias - ) + bias_ = None if (instance.tp_rank > 0 or instance.skip_bias_add) else instance.bias output_parallel = instance.quant_method.apply(instance, input_parallel, bias_) if instance.reduce_results and instance.tp_size > 1: diff --git a/rl_engine/kernels/ops/cuda/attention/flash_attn.py b/rl_engine/kernels/ops/cuda/attention/flash_attn.py index 9ad510b3..e57cdeb5 100644 --- a/rl_engine/kernels/ops/cuda/attention/flash_attn.py +++ b/rl_engine/kernels/ops/cuda/attention/flash_attn.py @@ -130,9 +130,7 @@ def precompile_training( tensors, RNG state, or distributed collectives. """ if torch.version.hip is not None: - raise StrictFlashAttentionUnavailable( - "FA4 CUDA precompile is unavailable on ROCm" - ) + raise StrictFlashAttentionUnavailable("FA4 CUDA precompile is unavailable on ROCm") if not torch.cuda.is_available(): raise StrictFlashAttentionUnavailable( "FA4 CUDA precompile requires an available CUDA device" @@ -144,11 +142,7 @@ def precompile_training( if head_dim <= 0 or sequence_length <= 0: raise ValueError("head_dim and sequence_length must be positive") - target = ( - torch.device("cuda", torch.cuda.current_device()) - if device is None - else device - ) + target = torch.device("cuda", torch.cuda.current_device()) if device is None else device if target.type != "cuda": raise ValueError("strict FA4 training precompile requires a CUDA device")