From 1604db47f8a0f1341eb2def1769f0c232ef1a808 Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Tue, 1 Sep 2026 11:49:49 +0800 Subject: [PATCH 01/13] 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 02/13] 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") From 0a9c88920b6f5b3b6cdeaf1a78fdc986cfe83eb4 Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Fri, 4 Sep 2026 03:26:48 -0700 Subject: [PATCH 03/13] [DSv4][P5-5] Shared Expert MLP: strict CUDA + Triton kernels (#64) Implements shared_expert_mlp_fwd/bwd per the P5-S0 contract: every valid token runs fc1 -> one-round SwiGLU -> fc2 on BF16 frozen weights, backward returns dX only (FP32 accumulator), and the shared output stays independent of the routed path. Both backends reproduce the FP32 oracle's numeric profile oracle-fp32-serial-v1 byte-for-byte on the same device: one lane owns one output element and reduces serially in ascending k, multiply and add rounded separately (__fmul_rn/__fadd_rn on CUDA, uncontracted IEEE fp32 in Triton), sigmoid computed as 1/(1+expf(-x)) to match torch.sigmoid on FP32 CUDA tensors. No cross-lane floating-point reduction exists anywhere, so results are batch/padding invariant by construction (fwd(x)[t] == fwd(x[t:t+1]) byte-equal). The one-round SwiGLU core runs in shared mode (p_s=None, no clamp, per S0 decision D6) and is the reuse point for P5-2 (#63). Providers subclass ReferenceProvider and override only the two shared-expert methods, so the full acceptance command runs unchanged; unsupported input (non-CUDA device, missing extension or triton, foreign numeric profile) raises instead of falling back (fail-closed). Provenance records split_k=1 / serial-ascending-k / no-FMA per the P5-5 provenance requirement. Acceptance: python scripts/check_p5.py --provider rl_engine.moe.backends.shared_expert:CudaSharedExpertProvider --device cuda python scripts/check_p5.py --provider rl_engine.moe.backends.shared_expert:TritonSharedExpertProvider --device cuda pytest tests/test_shared_expert_mlp.py python benchmarks/benchmark_shared_expert_mlp.py Signed-off-by: Yizheng Jiao --- benchmarks/benchmark_shared_expert_mlp.py | 133 ++++ csrc/cuda/moe/shared_expert_mlp.cu | 213 ++++++ csrc/ops.cpp | 13 + docs/operators/shared-expert-mlp.md | 44 ++ rl_engine/_C.pyi | 3 + rl_engine/kernels/ops/triton/moe/__init__.py | 2 + .../kernels/ops/triton/moe/shared_expert.py | 185 ++++++ rl_engine/moe/backends/__init__.py | 3 + rl_engine/moe/backends/shared_expert.py | 152 +++++ setup.py | 625 +++++++++--------- tests/test_shared_expert_mlp.py | 137 ++++ 11 files changed, 1198 insertions(+), 312 deletions(-) create mode 100644 benchmarks/benchmark_shared_expert_mlp.py create mode 100644 csrc/cuda/moe/shared_expert_mlp.cu create mode 100644 docs/operators/shared-expert-mlp.md create mode 100644 rl_engine/kernels/ops/triton/moe/__init__.py create mode 100644 rl_engine/kernels/ops/triton/moe/shared_expert.py create mode 100644 rl_engine/moe/backends/__init__.py create mode 100644 rl_engine/moe/backends/shared_expert.py create mode 100644 tests/test_shared_expert_mlp.py diff --git a/benchmarks/benchmark_shared_expert_mlp.py b/benchmarks/benchmark_shared_expert_mlp.py new file mode 100644 index 00000000..8ef6fe5a --- /dev/null +++ b/benchmarks/benchmark_shared_expert_mlp.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P5-5 (#64) shared_expert_mlp benchmark: torch-native vs Triton vs CUDA. + +torch-native is the non-deterministic cuBLAS/eager reference (speed ceiling); +the Triton and CUDA rows are the strict ``oracle-fp32-serial-v1`` kernels this +PR delivers. Alignment between the strict backends is asserted on every shape. + + python benchmarks/benchmark_shared_expert_mlp.py [--tokens 16,256,2048] +""" + +from __future__ import annotations + +import argparse +import pathlib +import sys + +import torch + +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.contract import SharedBatch, tensor_sha256 # noqa: E402 + + +def torch_native(batch: SharedBatch, dy: torch.Tensor): + """Eager BF16 reference (cuBLAS + fused silu): fast but not bit-stable.""" + x = batch.x.detach().requires_grad_(True) + z = x @ batch.w_fc1.t() + ffn = z.shape[1] // 2 + gate, up = z[:, :ffn], z[:, ffn:] + h = torch.nn.functional.silu(gate) * up + y = h @ batch.w_fc2.t() + y.backward(dy) + return y, x.grad + + +def make_runner(provider): + def run(batch: SharedBatch, dy: torch.Tensor): + y, saved = provider.shared_expert_mlp_fwd(batch) + dx = provider.shared_expert_mlp_bwd(dy, batch, saved) + return y, dx + + return run + + +def time_ms(fn, *args, warmup: int = 3, iters: int = 10) -> float: + for _ in range(warmup): + fn(*args) + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn(*args) + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--hidden", type=int, default=4096) + parser.add_argument("--ffn", type=int, default=2048) + parser.add_argument("--tokens", default="16,256,2048") + parser.add_argument("--iters", type=int, default=10) + args = parser.parse_args() + + if not torch.cuda.is_available(): + print("CUDA device required") + return 1 + + from rl_engine.moe.provider import resolve_provider + + runners: dict[str, object] = {"torch-native": torch_native} + strict: dict[str, object] = {} + for label, spec in ( + ("triton", "rl_engine.moe.backends.shared_expert:TritonSharedExpertProvider"), + ("cuda", "rl_engine.moe.backends.shared_expert:CudaSharedExpertProvider"), + ): + try: + runner = make_runner(resolve_provider(spec)) + runners[label] = runner + strict[label] = runner + except NotImplementedError as exc: + print(f"[skip] {label}: {exc}") + + device = torch.device("cuda") + gen = torch.Generator(device="cpu").manual_seed(2026) + header = f"{'T':>6} {'backend':>14} {'fwd+bwd ms':>12} {'vs native':>10}" + print(f"H={args.hidden} F={args.ffn} ({torch.cuda.get_device_name(0)})") + print(header) + for t in [int(v) for v in args.tokens.split(",")]: + x = torch.randn(t, args.hidden, generator=gen).to(torch.bfloat16).to(device) + w1 = ( + (torch.randn(2 * args.ffn, args.hidden, generator=gen) / args.hidden**0.5) + .to(torch.bfloat16) + .to(device) + ) + w2 = ( + (torch.randn(args.hidden, args.ffn, generator=gen) / args.ffn**0.5) + .to(torch.bfloat16) + .to(device) + ) + batch = SharedBatch(x=x, w_fc1=w1, w_fc2=w2) + dy = torch.randn(t, args.hidden, generator=gen).to(torch.bfloat16).to(device) + + outputs = {} + base_ms = None + for label, fn in runners.items(): + ms = time_ms(fn, batch, dy, iters=args.iters) + outputs[label] = fn(batch, dy) + if label == "torch-native": + base_ms = ms + rel = f"{ms / base_ms:8.2f}x" if base_ms else " -" + print(f"{t:>6} {label:>14} {ms:12.3f} {rel:>10}") + + strict_hashes = { + label: (tensor_sha256(outputs[label][0]), tensor_sha256(outputs[label][1])) + for label in strict + } + if len(strict_hashes) == 2 and len(set(strict_hashes.values())) != 1: + print(f" !! strict backends diverged at T={t}: {strict_hashes}") + return 1 + if strict_hashes: + print(f" strict backends byte-equal: {len(strict_hashes)}/{len(strict_hashes)}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/csrc/cuda/moe/shared_expert_mlp.cu b/csrc/cuda/moe/shared_expert_mlp.cu new file mode 100644 index 00000000..a1cd7b1b --- /dev/null +++ b/csrc/cuda/moe/shared_expert_mlp.cu @@ -0,0 +1,213 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// P5-5 (#64) Shared Expert MLP strict kernels: fc1 -> one-round SwiGLU -> fc2. +// +// Numeric profile ``oracle-fp32-serial-v1`` (see rl_engine/moe/oracle.py): +// every accumulation is FP32, serial, ascending-k, with multiply and add +// rounded separately (__fmul_rn / __fadd_rn; never contracted into FMA). +// One thread owns one output element, so batch size and padding cannot +// change a row's bytes (Axis-A bitwise invariance) and there is no +// cross-thread floating-point reduction anywhere. +// +// The one-round SwiGLU core (FP32 math, single BF16 round on the output) is +// shared with P5-2 (#63): the p_s / clamp variant extends the same device +// functions in this translation unit rather than forking the math. + +#include +#include +#include +#include +#include + +namespace { + +// out[m, n] = sum_{k ascending} fadd_rn(acc, fmul_rn(a[m, k], b(n, k))) +// A is BF16 [M, K]; B is BF16 [N, K] (TRANS_B = false) or [K, N] (true). +// Output stays FP32; the caller rounds to BF16 where the contract says so. +template +__global__ void p5_strict_gemm_kernel( + const __nv_bfloat16* __restrict__ a, + const __nv_bfloat16* __restrict__ b, + float* __restrict__ out, + const int64_t m_rows, + const int64_t n_cols, + const int64_t k_dim) { + const int64_t total = m_rows * n_cols; + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + for (int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; idx < total; + idx += stride) { + const int64_t m = idx / n_cols; + const int64_t n = idx - m * n_cols; + const __nv_bfloat16* a_row = a + m * k_dim; + float acc = 0.0f; + for (int64_t k = 0; k < k_dim; ++k) { + const float av = __bfloat162float(a_row[k]); + const float bv = + __bfloat162float(TRANS_B ? b[k * n_cols + n] : b[n * k_dim + k]); + acc = __fadd_rn(acc, __fmul_rn(av, bv)); + } + out[idx] = acc; + } +} + +// Matches torch.sigmoid on FP32 CUDA tensors: 1 / (1 + exp(-x)) with +// IEEE div.rn and the accurate expf (no fast-math in this build). +__device__ __forceinline__ float sigmoid_rn(float x) { + return 1.0f / (1.0f + expf(-x)); +} + +// One-round SwiGLU core, shared-expert mode (p_s = None: no clamp, no route +// weight). z is the packed FP32 fc1 output [T, 2F] (gate columns then up); +// h is the single BF16 round of SiLU(gate) * up. +__global__ void p5_swiglu_shared_forward_kernel( + const float* __restrict__ z, + __nv_bfloat16* __restrict__ h, + const int64_t n, + const int64_t width) { + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + for (int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; idx < n; + idx += stride) { + const int64_t row = idx / width; + const int64_t col = idx - row * width; + const int64_t gate_index = row * (2 * width) + col; + const float g = z[gate_index]; + const float u = z[gate_index + width]; + const float sig = sigmoid_rn(g); + const float silu = __fmul_rn(g, sig); + h[idx] = __float2bfloat16(__fmul_rn(silu, u)); + } +} + +// Backward of the same graph (p_s = None): recomputes sig/silu from the saved +// FP32 z with the identical instruction sequence, so the bits match the +// forward. dz packs (dgate | dup), each rounded to BF16 exactly once at the +// operator edge (mirrors the oracle's cat(...).to(bfloat16)). +__global__ void p5_swiglu_shared_backward_kernel( + const __nv_bfloat16* __restrict__ dh, + const float* __restrict__ z, + __nv_bfloat16* __restrict__ dz, + const int64_t n, + const int64_t width) { + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + for (int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; idx < n; + idx += stride) { + const int64_t row = idx / width; + const int64_t col = idx - row * width; + const int64_t gate_index = row * (2 * width) + col; + const float g = z[gate_index]; + const float u = z[gate_index + width]; + const float dh32 = __bfloat162float(dh[idx]); + const float sig = sigmoid_rn(g); + const float silu = __fmul_rn(g, sig); + // dsilu = sig * (1 + g * (1 - sig)), each op rounded separately. + float t = __fsub_rn(1.0f, sig); + t = __fmul_rn(g, t); + t = __fadd_rn(1.0f, t); + const float dsilu = __fmul_rn(sig, t); + const float dgate = __fmul_rn(__fmul_rn(dh32, u), dsilu); + const float dup = __fmul_rn(dh32, silu); + dz[gate_index] = __float2bfloat16(dgate); + dz[gate_index + width] = __float2bfloat16(dup); + } +} + +void launch_1d(int64_t n, int& threads, int64_t& blocks) { + threads = 256; + blocks = (n + threads - 1) / threads; + if (blocks == 0) { + blocks = 1; + } + if (blocks > 65535) { + blocks = 65535; // grid-stride loops cover the rest + } +} + +void check_cuda_2d(const torch::Tensor& t, at::ScalarType dtype, const char* name) { + TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(t.dim() == 2, name, " must be 2-D"); + TORCH_CHECK(t.scalar_type() == dtype, name, " must be ", dtype, ", got ", t.scalar_type()); +} + +} // namespace + +torch::Tensor p5_strict_gemm(torch::Tensor a, torch::Tensor b, bool trans_b) { + check_cuda_2d(a, at::kBFloat16, "a"); + check_cuda_2d(b, at::kBFloat16, "b"); + TORCH_CHECK(a.device() == b.device(), "a and b must be on the same CUDA device"); + const int64_t m_rows = a.size(0); + const int64_t k_dim = a.size(1); + const int64_t n_cols = trans_b ? b.size(1) : b.size(0); + const int64_t bk = trans_b ? b.size(0) : b.size(1); + TORCH_CHECK(bk == k_dim, "K mismatch: a has K=", k_dim, ", b has K=", bk); + const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); + auto out = torch::empty({m_rows, n_cols}, a.options().dtype(at::kFloat)); + const int64_t n = out.numel(); + if (n == 0 || k_dim == 0) { + return n == 0 ? out : out.zero_(); + } + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + auto stream = at::cuda::getCurrentCUDAStream(); + const auto* a_ptr = reinterpret_cast(a.data_ptr()); + const auto* b_ptr = reinterpret_cast(b.data_ptr()); + if (trans_b) { + p5_strict_gemm_kernel<<>>( + a_ptr, b_ptr, out.data_ptr(), m_rows, n_cols, k_dim); + } else { + p5_strict_gemm_kernel<<>>( + a_ptr, b_ptr, out.data_ptr(), m_rows, n_cols, k_dim); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +torch::Tensor p5_swiglu_shared_forward(torch::Tensor z) { + check_cuda_2d(z, at::kFloat, "z"); + TORCH_CHECK(z.size(1) % 2 == 0, "z width must be even (packed gate|up)"); + const at::cuda::OptionalCUDAGuard device_guard(device_of(z)); + const int64_t width = z.size(1) / 2; + auto h = torch::empty({z.size(0), width}, z.options().dtype(at::kBFloat16)); + const int64_t n = h.numel(); + if (n == 0) { + return h; + } + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + auto stream = at::cuda::getCurrentCUDAStream(); + p5_swiglu_shared_forward_kernel<<>>( + z.data_ptr(), reinterpret_cast<__nv_bfloat16*>(h.data_ptr()), n, width); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return h; +} + +torch::Tensor p5_swiglu_shared_backward(torch::Tensor dh, torch::Tensor z) { + check_cuda_2d(dh, at::kBFloat16, "dh"); + check_cuda_2d(z, at::kFloat, "z"); + TORCH_CHECK(dh.device() == z.device(), "dh and z must be on the same CUDA device"); + TORCH_CHECK(z.size(1) % 2 == 0, "z width must be even (packed gate|up)"); + TORCH_CHECK( + dh.size(0) == z.size(0) && dh.size(1) * 2 == z.size(1), + "dh shape must match the packed gate/up halves of z"); + const at::cuda::OptionalCUDAGuard device_guard(device_of(z)); + auto dz = torch::empty_like(z, z.options().dtype(at::kBFloat16)); + const int64_t n = dh.numel(); + if (n == 0) { + return dz; + } + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + auto stream = at::cuda::getCurrentCUDAStream(); + p5_swiglu_shared_backward_kernel<<>>( + reinterpret_cast(dh.data_ptr()), + z.data_ptr(), + reinterpret_cast<__nv_bfloat16*>(dz.data_ptr()), + n, + dh.size(1)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return dz; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index aecc30ed..33bb88bd 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -131,6 +131,11 @@ std::vector swiglu_packed_backward_cuda( torch::Tensor dy, torch::Tensor gate_up); +// P5-5 (#64) Shared Expert MLP strict kernels (oracle-fp32-serial-v1) +torch::Tensor p5_strict_gemm(torch::Tensor a, torch::Tensor b, bool trans_b); +torch::Tensor p5_swiglu_shared_forward(torch::Tensor z); +torch::Tensor p5_swiglu_shared_backward(torch::Tensor dh, torch::Tensor z); + // RMSNorm Declarations & Wrappers void rmsnorm_forward_cuda( @@ -503,6 +508,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("swiglu_packed_backward", &swiglu_packed_backward, "Batch-invariant SwiGLU backward for [rows, 2 * intermediate]"); + // P5-5 (#64) Shared Expert MLP strict kernels (oracle-fp32-serial-v1) + m.def("p5_strict_gemm", &p5_strict_gemm, + "Strict BF16-in/FP32-out GEMM, serial ascending-k, mul-then-add"); + m.def("p5_swiglu_shared_forward", &p5_swiglu_shared_forward, + "One-round SwiGLU forward, shared-expert mode (p_s = None)"); + m.def("p5_swiglu_shared_backward", &p5_swiglu_shared_backward, + "One-round SwiGLU backward, shared-expert mode (p_s = None)"); + // Deterministic standard-softmax attention (issue #147) m.def( "deterministic_attention_forward", diff --git a/docs/operators/shared-expert-mlp.md b/docs/operators/shared-expert-mlp.md new file mode 100644 index 00000000..c4ac1809 --- /dev/null +++ b/docs/operators/shared-expert-mlp.md @@ -0,0 +1,44 @@ +# Shared Expert MLP (P5-5, issue #64) + +Shared expert for the DSv4 MoE block: every valid token runs +`fc1 -> one-round SwiGLU -> fc2` once. BF16 frozen weights, backward returns +only `dX` (FP32 accumulator dtype); the shared output stays independent of the +routed path (the combine belongs to P6). + +## Fixed math (`oracle-fp32-serial-v1`) + +``` +z = x @ w_fc1.T # BF16 operands, FP32 serial ascending-k, mul-then-add +h = BF16(SiLU(gate) * up) # FP32 math, single round; no clamp, no p_s +y = BF16(h @ w_fc2.T) # FP32 accumulate, one round +dX = FP32(dz @ w_fc1) # dh, dz round BF16 at operator edges +``` + +Strict kernels reproduce the oracle byte-for-byte on the same device: one lane +owns one output element and reduces serially in ascending k with +`__fmul_rn`/`__fadd_rn` (CUDA) or uncontracted IEEE fp32 arith (Triton), so +there is no cross-lane floating-point reduction and results are +batch/padding invariant. + +## Backends + +| backend | entry point | +| --- | --- | +| CUDA | `rl_engine.moe.backends.shared_expert:CudaSharedExpertProvider` (`csrc/cuda/moe/shared_expert_mlp.cu`) | +| Triton | `rl_engine.moe.backends.shared_expert:TritonSharedExpertProvider` (`rl_engine/kernels/ops/triton/moe/shared_expert.py`) | + +The one-round SwiGLU core runs in shared mode (`p_s = None`, no clamp) and is +the reuse point for P5-2 (#63), which extends the same core with clamp, +route weight, and `dp_s`. + +## Acceptance + +```bash +python scripts/check_p5.py --provider rl_engine.moe.backends.shared_expert:CudaSharedExpertProvider --device cuda +python scripts/check_p5.py --provider rl_engine.moe.backends.shared_expert:TritonSharedExpertProvider --device cuda +pytest tests/test_shared_expert_mlp.py +python benchmarks/benchmark_shared_expert_mlp.py +``` + +Fail-closed: non-CUDA input, a missing extension/triton install, or a foreign +numeric profile raises instead of falling back to the oracle. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 8e0e865a..8f7879fb 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -205,6 +205,9 @@ def swiglu_backward( gate: torch.Tensor, up: torch.Tensor, ) -> list[torch.Tensor]: ... +def p5_strict_gemm(a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: ... +def p5_swiglu_shared_forward(z: torch.Tensor) -> torch.Tensor: ... +def p5_swiglu_shared_backward(dh: torch.Tensor, z: torch.Tensor) -> torch.Tensor: ... def rmsnorm_forward( x: torch.Tensor, weight: torch.Tensor, diff --git a/rl_engine/kernels/ops/triton/moe/__init__.py b/rl_engine/kernels/ops/triton/moe/__init__.py new file mode 100644 index 00000000..86cf4c9d --- /dev/null +++ b/rl_engine/kernels/ops/triton/moe/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors diff --git a/rl_engine/kernels/ops/triton/moe/shared_expert.py b/rl_engine/kernels/ops/triton/moe/shared_expert.py new file mode 100644 index 00000000..125d8a83 --- /dev/null +++ b/rl_engine/kernels/ops/triton/moe/shared_expert.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P5-5 (#64) Shared Expert MLP strict Triton kernels. + +Numeric profile ``oracle-fp32-serial-v1`` (see ``rl_engine/moe/oracle.py``): +FP32 accumulation, serial ascending-k, multiply and add rounded separately. +Each output element is owned by one lane and reduced serially, so there is no +cross-lane floating-point reduction and results are batch/padding invariant. + +The one-round SwiGLU core (FP32 math, single BF16 round on the output) is the +shared-mode (``p_s = None``, no clamp) variant shared with P5-2 (#63). +""" + +from __future__ import annotations + +import torch + +try: + import triton + import triton.language as tl + + TRITON_AVAILABLE = True +except ImportError: # pragma: no cover - exercised on non-GPU installs + TRITON_AVAILABLE = False + + +if TRITON_AVAILABLE: + + @triton.jit + def _strict_gemm_kernel( + A, + B, + C, + K, + N, + stride_bn, + stride_bk, + BLOCK_N: tl.constexpr, + ): + # C[m, n] = sum_{k ascending} A[m, k] * B(n, k); FP32 accumulator, + # one lane per output element, mul and add rounded separately + # (fp32 arith in Triton is IEEE by default: no FMA contraction). + m = tl.program_id(0) + pn = tl.program_id(1) + offs_n = pn * BLOCK_N + tl.arange(0, BLOCK_N) + mask_n = offs_n < N + acc = tl.zeros([BLOCK_N], dtype=tl.float32) + a_row = A + m * K + b_cols = B + offs_n * stride_bn + for k in range(0, K): + a = tl.load(a_row + k).to(tl.float32) + b = tl.load(b_cols + k * stride_bk, mask=mask_n, other=0.0).to(tl.float32) + prod = a * b + acc = acc + prod + tl.store(C + m * N + offs_n, acc, mask=mask_n) + + @triton.jit + def _swiglu_shared_fwd_kernel( + Z, + H, + n_elem, + width, + BLOCK: tl.constexpr, + ): + # One-round SwiGLU, shared mode: h = BF16(SiLU(gate) * up), FP32 math, + # gate = z[:, :F], up = z[:, F:] packed in one [T, 2F] FP32 tensor. + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elem + row = offs // width + col = offs - row * width + gate_index = row * (2 * width) + col + g = tl.load(Z + gate_index, mask=mask, other=0.0) + u = tl.load(Z + gate_index + width, mask=mask, other=0.0) + sig = 1.0 / (1.0 + tl.exp(-g)) + silu = g * sig + h = (silu * u).to(tl.bfloat16) + tl.store(H + offs, h, mask=mask) + + @triton.jit + def _swiglu_shared_bwd_kernel( + DH, + Z, + DZ, + n_elem, + width, + BLOCK: tl.constexpr, + ): + # dgate = ((dh * u) * dsilu); dup = dh * silu; both round to BF16 once + # at the operator edge (mirrors the oracle's cat(...).to(bfloat16)). + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elem + row = offs // width + col = offs - row * width + gate_index = row * (2 * width) + col + g = tl.load(Z + gate_index, mask=mask, other=0.0) + u = tl.load(Z + gate_index + width, mask=mask, other=0.0) + dh = tl.load(DH + offs, mask=mask, other=0.0).to(tl.float32) + sig = 1.0 / (1.0 + tl.exp(-g)) + silu = g * sig + # dsilu = sig * (1 + g * (1 - sig)), each op rounded separately. + t = 1.0 - sig + t = g * t + t = 1.0 + t + dsilu = sig * t + dgate = (dh * u) * dsilu + dup = dh * silu + tl.store(DZ + gate_index, dgate.to(tl.bfloat16), mask=mask) + tl.store(DZ + gate_index + width, dup.to(tl.bfloat16), mask=mask) + + +def _check_cuda_2d(t: torch.Tensor, dtype: torch.dtype, name: str) -> None: + if not t.is_cuda: + raise NotImplementedError(f"{name} must be a CUDA tensor for the Triton backend") + if not t.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + if t.dim() != 2: + raise ValueError(f"{name} must be 2-D") + if t.dtype != dtype: + raise TypeError(f"{name} must be {dtype}, got {t.dtype}") + + +def strict_gemm(a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: + """``a @ b.T`` (or ``a @ b`` when ``trans_b``): BF16 in, FP32 out, strict.""" + if not TRITON_AVAILABLE: + raise NotImplementedError("triton is not installed (fail-closed, no fallback)") + _check_cuda_2d(a, torch.bfloat16, "a") + _check_cuda_2d(b, torch.bfloat16, "b") + m, k = a.shape + if trans_b: + bk, n = b.shape + stride_bn, stride_bk = 1, n + else: + n, bk = b.shape + stride_bn, stride_bk = k, 1 + if bk != k: + raise ValueError(f"K mismatch: a has K={k}, b has K={bk}") + out = torch.empty(m, n, dtype=torch.float32, device=a.device) + if out.numel() == 0: + return out + if k == 0: + return out.zero_() + block_n = min(triton.next_power_of_2(n), 256) + grid = (m, triton.cdiv(n, block_n)) + _strict_gemm_kernel[grid](a, b, out, k, n, stride_bn, stride_bk, BLOCK_N=block_n) + return out + + +def swiglu_shared_fwd(z: torch.Tensor) -> torch.Tensor: + """One-round SwiGLU forward, shared mode: FP32 [T, 2F] -> BF16 [T, F].""" + if not TRITON_AVAILABLE: + raise NotImplementedError("triton is not installed (fail-closed, no fallback)") + _check_cuda_2d(z, torch.float32, "z") + if z.shape[1] % 2 != 0: + raise ValueError("z width must be even (packed gate|up)") + width = z.shape[1] // 2 + h = torch.empty(z.shape[0], width, dtype=torch.bfloat16, device=z.device) + n_elem = h.numel() + if n_elem == 0: + return h + block = 1024 + grid = (triton.cdiv(n_elem, block),) + _swiglu_shared_fwd_kernel[grid](z, h, n_elem, width, BLOCK=block) + return h + + +def swiglu_shared_bwd(dh: torch.Tensor, z: torch.Tensor) -> torch.Tensor: + """One-round SwiGLU backward, shared mode: returns packed BF16 dz [T, 2F].""" + if not TRITON_AVAILABLE: + raise NotImplementedError("triton is not installed (fail-closed, no fallback)") + _check_cuda_2d(dh, torch.bfloat16, "dh") + _check_cuda_2d(z, torch.float32, "z") + if z.shape[1] % 2 != 0: + raise ValueError("z width must be even (packed gate|up)") + if dh.shape[0] != z.shape[0] or dh.shape[1] * 2 != z.shape[1]: + raise ValueError("dh shape must match the packed gate/up halves of z") + dz = torch.empty_like(z, dtype=torch.bfloat16) + n_elem = dh.numel() + if n_elem == 0: + return dz + block = 1024 + grid = (triton.cdiv(n_elem, block),) + _swiglu_shared_bwd_kernel[grid](dh, z, dz, n_elem, dh.shape[1], BLOCK=block) + return dz diff --git a/rl_engine/moe/backends/__init__.py b/rl_engine/moe/backends/__init__.py new file mode 100644 index 00000000..6eec9430 --- /dev/null +++ b/rl_engine/moe/backends/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P5 kernel backends. Each sub-issue registers its providers here.""" diff --git a/rl_engine/moe/backends/shared_expert.py b/rl_engine/moe/backends/shared_expert.py new file mode 100644 index 00000000..ef11df76 --- /dev/null +++ b/rl_engine/moe/backends/shared_expert.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P5-5 (#64) Shared Expert MLP providers (CUDA and Triton strict backends). + +Both backends implement the frozen math ``fc1 -> one-round SwiGLU -> fc2`` +under the ``oracle-fp32-serial-v1`` numeric profile and are byte-equal to the +FP32 oracle running on the same device. Only the two shared-expert methods are +overridden; every other operator stays on the oracle per the S0 start kit, so +the full acceptance command runs unchanged: + + python scripts/check_p5.py \ + --provider rl_engine.moe.backends.shared_expert:CudaSharedExpertProvider \ + --device cuda + +Fail-closed: unsupported input (non-CUDA device, missing extension/triton, +schema violations) raises instead of falling back to another implementation. +The shared output is produced from ``SharedBatch`` alone -- no route weight, +no routed combine (that boundary belongs to P6). +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from rl_engine.moe.contract import ORACLE_PROFILE, SharedBatch +from rl_engine.moe.provider import ReferenceProvider + + +class _StrictSharedExpertProvider(ReferenceProvider): + """Common composite: strict GEMMs + one-round SwiGLU, dX only (frozen base).""" + + name = "shared-expert-strict" + numeric_profile = ORACLE_PROFILE + + # Backend hooks ------------------------------------------------------- + def _gemm(self, a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: + raise NotImplementedError + + def _swiglu_fwd(self, z: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + def _swiglu_bwd(self, dh: torch.Tensor, z: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + # Provider surface ---------------------------------------------------- + def capabilities(self) -> dict[str, Any]: + return { + "backend": self.name, + "operators": ["shared_expert_mlp_fwd", "shared_expert_mlp_bwd"], + "geometry": ["one-row", "packed"], + "devices": ["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__, + # Changing any of these changes the addition order (P5-5 s4). + "split_k": 1, + "reduction": "serial-ascending-k", + "rounding": "mul-then-add, no FMA", + "workspace": "none", + } + + def _check_batch(self, batch: SharedBatch) -> None: + batch.validate() + if batch.numeric_profile != ORACLE_PROFILE: + raise NotImplementedError( + f"{self.name} only implements {ORACLE_PROFILE!r}, " + f"got {batch.numeric_profile!r} (fail-closed, no fallback)" + ) + if not batch.x.is_cuda: + raise NotImplementedError( + f"{self.name} requires CUDA tensors, got device {batch.x.device} " + "(fail-closed, no fallback)" + ) + + def shared_expert_mlp_fwd(self, batch: SharedBatch) -> tuple[torch.Tensor, dict[str, Any]]: + self._check_batch(batch) + x = batch.x.contiguous() + w_fc1 = batch.w_fc1.contiguous() + w_fc2 = batch.w_fc2.contiguous() + z = self._gemm(x, w_fc1, False) # [T, 2F] FP32, kept for backward + h_bf16 = self._swiglu_fwd(z) # [T, F] BF16, the one round + y = self._gemm(h_bf16, w_fc2, False).to(torch.bfloat16) # [T, H] + saved: dict[str, Any] = {"z32": z, "h_bf16": h_bf16} + return y, saved + + def shared_expert_mlp_bwd( + self, dy: torch.Tensor, batch: SharedBatch, saved: dict[str, Any] + ) -> torch.Tensor: + self._check_batch(batch) + z = saved["z32"] + dy_bf16 = dy.to(torch.bfloat16).contiguous() + # dh = BF16(dY @ W2), dz = swiglu_bwd, dX = dz @ W1 (FP32 accumulator). + dh = self._gemm(dy_bf16, batch.w_fc2.contiguous(), True).to(torch.bfloat16) + dz = self._swiglu_bwd(dh, z) + dx = self._gemm(dz, batch.w_fc1.contiguous(), True) + return dx + + +class CudaSharedExpertProvider(_StrictSharedExpertProvider): + """CUDA backend: csrc/cuda/moe/shared_expert_mlp.cu via rl_engine._C.""" + + name = "shared-expert-cuda" + + def __init__(self) -> None: + try: + from rl_engine import _C + except ImportError as exc: # fail-closed: no oracle fallback + raise NotImplementedError( + "rl_engine._C is not built; install with RL_KERNEL_REQUIRE_EXT=1" + ) from exc + for symbol in ("p5_strict_gemm", "p5_swiglu_shared_forward", "p5_swiglu_shared_backward"): + if not hasattr(_C, symbol): + raise NotImplementedError(f"rl_engine._C lacks {symbol}; rebuild the extension") + self._ext = _C + + def _gemm(self, a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: + return self._ext.p5_strict_gemm(a, b, trans_b) + + def _swiglu_fwd(self, z: torch.Tensor) -> torch.Tensor: + return self._ext.p5_swiglu_shared_forward(z) + + def _swiglu_bwd(self, dh: torch.Tensor, z: torch.Tensor) -> torch.Tensor: + return self._ext.p5_swiglu_shared_backward(dh, z) + + +class TritonSharedExpertProvider(_StrictSharedExpertProvider): + """Triton backend: rl_engine/kernels/ops/triton/moe/shared_expert.py.""" + + name = "shared-expert-triton" + + def __init__(self) -> None: + from rl_engine.kernels.ops.triton.moe import shared_expert as tk + + if not tk.TRITON_AVAILABLE: + raise NotImplementedError("triton is not installed (fail-closed, no fallback)") + self._tk = tk + + def _gemm(self, a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: + return self._tk.strict_gemm(a, b, trans_b) + + def _swiglu_fwd(self, z: torch.Tensor) -> torch.Tensor: + return self._tk.swiglu_shared_fwd(z) + + def _swiglu_bwd(self, dh: torch.Tensor, z: torch.Tensor) -> torch.Tensor: + return self._tk.swiglu_shared_bwd(dh, z) diff --git a/setup.py b/setup.py index 79f882d9..6d8bbb4b 100644 --- a/setup.py +++ b/setup.py @@ -1,313 +1,314 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -import importlib.util -import os -import warnings -from pathlib import Path - -from setuptools import find_packages, setup - - -def _load_envs_module(): - envs_path = Path(__file__).with_name("envs.py") - spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) - if spec is None or spec.loader is None: - raise RuntimeError(f"failed to load environment helpers from {envs_path}") - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -envs = _load_envs_module() - - -def _load_torch_extension_tools(): - try: - import torch - except ModuleNotFoundError as exc: - if exc.name != "torch": - raise - return None, None, None - - from torch.utils.cpp_extension import BuildExtension, CUDAExtension - - # CUDAExtension is also the supported extension entry point for ROCm - # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when - # torch.version.hip is set. - return torch, BuildExtension, CUDAExtension - - -def _native_extension_required() -> bool: - """Whether the caller explicitly requested a native extension build.""" - return ( - envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) - or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) - or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) - or envs.env_flag("FORCE_CUDA") - ) - - -def _cuda_define_from_env(name: str, macro: str) -> list[str]: - value = os.environ.get(name) - if value is None: - return [] - parsed = int(value) - if parsed <= 0: - raise ValueError(f"{name} must be positive, got {value!r}") - return [f"-D{macro}={parsed}"] - - -_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( - "-Xfatbin", - "-compress-all", - "-gencode", - "--generate-code", - "--expt-", - "-lineinfo", - "-allow-unsupported-compiler", - "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", -) -_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { - "-Xfatbin", - "-gencode", - "--generate-code", -} - - -def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: - """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" - filtered_flags = [] - skip_next = False - for flag in flags: - if skip_next: - skip_next = False - continue - if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: - skip_next = True - continue - if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): - continue - filtered_flags.append(flag) - return filtered_flags - - -def get_extensions(): - torch, _, CUDAExtension = _load_torch_extension_tools() - if torch is None: - message = ( - "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " - "CUDA/ROCm PyTorch build first, then run " - "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." - ) - if _native_extension_required(): - raise RuntimeError(message) - warnings.warn( - f"{message} Continuing with the pure-Python fallback because no native extension " - "was explicitly requested.", - RuntimeWarning, - stacklevel=2, - ) - return [] - - extensions = [] - torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") - torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] - if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": - torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") - is_rocm = getattr(torch.version, "hip", None) is not None - - # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, - # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also - # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add - # --offload-arch. Do not require a visible GPU when a ROCm target was - # explicitly selected. - no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() - if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: - raise RuntimeError( - "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " - "Set one or more ';'-separated targets, for example " - "PYTORCH_ROCM_ARCH='gfx942;gfx950'." - ) - - if is_rocm or torch.cuda.is_available(): - cuda_sources = [ - "csrc/ops.cpp", - "csrc/fused_logp_kernel.cu", - "csrc/deterministic_logp_kernel.cu", - "csrc/cuda/gemm/det_gemm_kernel.cu", - "csrc/cuda/rmsnorm.cu", - "csrc/cuda/activation.cu", - "csrc/cuda/attention/deterministic_attention.cu", - "csrc/cuda/distributed/deterministic_collective.cu", - ] - if not is_rocm: - # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). - # The ROCm dispatcher falls back to PyTorch SDPA for this operator. - cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") - - nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] - if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): - nvcc_flags.append("--use_fast_math") - if not is_rocm: - cc_major, cc_minor = torch.cuda.get_device_capability() - enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" - if not enable_sm90: - # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. - nvcc_flags.append( - f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" - ) - nvcc_flags.append("--expt-relaxed-constexpr") - nvcc_flags.append("--expt-extended-lambda") - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - ) - ) - if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): - nvcc_flags.append("-lineinfo") - if ( - not is_rocm - and os.name == "nt" - and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) - ): - nvcc_flags.append("-allow-unsupported-compiler") - nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") - - cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] - extra_link_args = list(torch_rpath) - if os.name != "nt": - # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). - extra_link_args.append("-lcuda") - - if not is_rocm: - sm90_srcs = [ - "csrc/cuda/fused_logp_sm90.cu", - "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob - "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp - "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build - # Single-card batch-invariant embedding/lm-head. - "csrc/cuda/embedding_lm_head_sm90.cu", - ] - enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) - present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] - if enable_sm90 and present_sm90: - tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant - cuda_sources.extend(present_sm90) +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import importlib.util +import os +import warnings +from pathlib import Path + +from setuptools import find_packages, setup + + +def _load_envs_module(): + envs_path = Path(__file__).with_name("envs.py") + spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"failed to load environment helpers from {envs_path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +envs = _load_envs_module() + + +def _load_torch_extension_tools(): + try: + import torch + except ModuleNotFoundError as exc: + if exc.name != "torch": + raise + return None, None, None + + from torch.utils.cpp_extension import BuildExtension, CUDAExtension + + # CUDAExtension is also the supported extension entry point for ROCm + # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when + # torch.version.hip is set. + return torch, BuildExtension, CUDAExtension + + +def _native_extension_required() -> bool: + """Whether the caller explicitly requested a native extension build.""" + return ( + envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) + or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) + or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) + or envs.env_flag("FORCE_CUDA") + ) + + +def _cuda_define_from_env(name: str, macro: str) -> list[str]: + value = os.environ.get(name) + if value is None: + return [] + parsed = int(value) + if parsed <= 0: + raise ValueError(f"{name} must be positive, got {value!r}") + return [f"-D{macro}={parsed}"] + + +_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( + "-Xfatbin", + "-compress-all", + "-gencode", + "--generate-code", + "--expt-", + "-lineinfo", + "-allow-unsupported-compiler", + "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", +) +_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { + "-Xfatbin", + "-gencode", + "--generate-code", +} + + +def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: + """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" + filtered_flags = [] + skip_next = False + for flag in flags: + if skip_next: + skip_next = False + continue + if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: + skip_next = True + continue + if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): + continue + filtered_flags.append(flag) + return filtered_flags + + +def get_extensions(): + torch, _, CUDAExtension = _load_torch_extension_tools() + if torch is None: + message = ( + "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " + "CUDA/ROCm PyTorch build first, then run " + "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." + ) + if _native_extension_required(): + raise RuntimeError(message) + warnings.warn( + f"{message} Continuing with the pure-Python fallback because no native extension " + "was explicitly requested.", + RuntimeWarning, + stacklevel=2, + ) + return [] + + extensions = [] + torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") + torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] + if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": + torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") + is_rocm = getattr(torch.version, "hip", None) is not None + + # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, + # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also + # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add + # --offload-arch. Do not require a visible GPU when a ROCm target was + # explicitly selected. + no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() + if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: + raise RuntimeError( + "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " + "Set one or more ';'-separated targets, for example " + "PYTORCH_ROCM_ARCH='gfx942;gfx950'." + ) + + if is_rocm or torch.cuda.is_available(): + cuda_sources = [ + "csrc/ops.cpp", + "csrc/fused_logp_kernel.cu", + "csrc/deterministic_logp_kernel.cu", + "csrc/cuda/gemm/det_gemm_kernel.cu", + "csrc/cuda/rmsnorm.cu", + "csrc/cuda/activation.cu", + "csrc/cuda/moe/shared_expert_mlp.cu", + "csrc/cuda/attention/deterministic_attention.cu", + "csrc/cuda/distributed/deterministic_collective.cu", + ] + if not is_rocm: + # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). + # The ROCm dispatcher falls back to PyTorch SDPA for this operator. + cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") + + nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] + if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): + nvcc_flags.append("--use_fast_math") + if not is_rocm: + cc_major, cc_minor = torch.cuda.get_device_capability() + enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" + if not enable_sm90: + # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. + nvcc_flags.append( + f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" + ) + nvcc_flags.append("--expt-relaxed-constexpr") + nvcc_flags.append("--expt-extended-lambda") + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + ) + ) + if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): + nvcc_flags.append("-lineinfo") + if ( + not is_rocm + and os.name == "nt" + and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) + ): + nvcc_flags.append("-allow-unsupported-compiler") + nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") + + cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] + extra_link_args = list(torch_rpath) + if os.name != "nt": + # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). + extra_link_args.append("-lcuda") + + if not is_rocm: + sm90_srcs = [ + "csrc/cuda/fused_logp_sm90.cu", + "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob + "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp + "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build + # Single-card batch-invariant embedding/lm-head. + "csrc/cuda/embedding_lm_head_sm90.cu", + ] + enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) + present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] + if enable_sm90 and present_sm90: + tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant + cuda_sources.extend(present_sm90) nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") - cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - - # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp - # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in - # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. - enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" - if enable_det_gemm_sm90: - tma_arch = f"{cc_major}{cc_minor}a" - arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" - if arch_flag not in nvcc_flags: - nvcc_flags.append(arch_flag) - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") - cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") - - if is_rocm: - nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) - - extensions.append( - CUDAExtension( - name="rl_engine._C", - sources=cuda_sources, - include_dirs=[], - extra_compile_args={ - "cxx": cxx_flags, - "nvcc": nvcc_flags, - }, - extra_link_args=extra_link_args, - ) - ) - - if _native_extension_required() and not extensions: - raise RuntimeError( - "rl_engine._C was requested but no CUDA/ROCm build environment is available. " - "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " - "PYTORCH_ROCM_ARCH to the target architecture." - ) - - return extensions - - -def get_cmdclass(): - _, BuildExtension, _ = _load_torch_extension_tools() - if BuildExtension is None: - return {} - return {"build_ext": BuildExtension} - - -setup( - name="rl-engine", - version="0.1.0", - packages=find_packages(include=["rl_engine", "rl_engine.*"]), - install_requires=[ - "torch>=2.4.1", - "tabulate", - "numpy", - "accelerate", - "transformers==5.13.1", - ], - ext_modules=get_extensions(), - cmdclass=get_cmdclass(), - extras_require={ - "cuda": ["flashinfer"], - "rocm": ["aiter"], - "vllm": ["vllm>=0.6.0"], - "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], - }, - entry_points={ - "console_scripts": [ - "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", - ], - }, - python_requires=">=3.10", - include_package_data=True, - zip_safe=False, -) + cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + + # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp + # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in + # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. + enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" + if enable_det_gemm_sm90: + tma_arch = f"{cc_major}{cc_minor}a" + arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" + if arch_flag not in nvcc_flags: + nvcc_flags.append(arch_flag) + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") + cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") + + if is_rocm: + nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) + + extensions.append( + CUDAExtension( + name="rl_engine._C", + sources=cuda_sources, + include_dirs=[], + extra_compile_args={ + "cxx": cxx_flags, + "nvcc": nvcc_flags, + }, + extra_link_args=extra_link_args, + ) + ) + + if _native_extension_required() and not extensions: + raise RuntimeError( + "rl_engine._C was requested but no CUDA/ROCm build environment is available. " + "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " + "PYTORCH_ROCM_ARCH to the target architecture." + ) + + return extensions + + +def get_cmdclass(): + _, BuildExtension, _ = _load_torch_extension_tools() + if BuildExtension is None: + return {} + return {"build_ext": BuildExtension} + + +setup( + name="rl-engine", + version="0.1.0", + packages=find_packages(include=["rl_engine", "rl_engine.*"]), + install_requires=[ + "torch>=2.4.1", + "tabulate", + "numpy", + "accelerate", + "transformers==5.13.1", + ], + ext_modules=get_extensions(), + cmdclass=get_cmdclass(), + extras_require={ + "cuda": ["flashinfer"], + "rocm": ["aiter"], + "vllm": ["vllm>=0.6.0"], + "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], + }, + entry_points={ + "console_scripts": [ + "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", + ], + }, + python_requires=">=3.10", + include_package_data=True, + zip_safe=False, +) diff --git a/tests/test_shared_expert_mlp.py b/tests/test_shared_expert_mlp.py new file mode 100644 index 00000000..6590ceea --- /dev/null +++ b/tests/test_shared_expert_mlp.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P5-5 (#64) shared_expert_mlp: bit-wise alignment against the FP32 oracle. + +Every comparison is byte-equality (sha256 over raw little-endian bytes) with +the oracle executed on the same device, per the P5 start-kit acceptance rules. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.moe import fixtures, oracle +from rl_engine.moe.contract import SharedBatch, tensor_sha256 + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + +PROVIDER_SPECS = { + "cuda": "rl_engine.moe.backends.shared_expert:CudaSharedExpertProvider", + "triton": "rl_engine.moe.backends.shared_expert:TritonSharedExpertProvider", +} + + +@pytest.fixture(params=sorted(PROVIDER_SPECS)) +def provider(request): + from rl_engine.moe.provider import resolve_provider + + try: + return resolve_provider(PROVIDER_SPECS[request.param]) + except NotImplementedError as exc: + pytest.skip(f"{request.param} backend unavailable: {exc}") + + +def _run_oracle(batch: SharedBatch, dy: torch.Tensor): + y, saved = oracle.shared_expert_mlp_fwd(batch) + dx = oracle.shared_expert_mlp_bwd(dy, batch, saved) + return y, dx + + +def _run_provider(provider, batch: SharedBatch, dy: torch.Tensor): + y, saved = provider.shared_expert_mlp_fwd(batch) + dx = provider.shared_expert_mlp_bwd(dy, batch, saved) + return y, dx + + +@requires_cuda +@pytest.mark.parametrize("case", sorted(fixtures.SHARED_CASES)) +def test_shared_cases_byte_equal(provider, case): + batch = fixtures.make_shared_batch(case).to("cuda") + y_gold, saved_gold = oracle.shared_expert_mlp_fwd(batch) + dy = fixtures.make_grad_output(case, tuple(y_gold.shape)).to("cuda") + dx_gold = oracle.shared_expert_mlp_bwd(dy, batch, saved_gold) + y, dx = _run_provider(provider, batch, dy) + assert y.dtype == torch.bfloat16 and dx.dtype == torch.float32 + assert tensor_sha256(y) == tensor_sha256(y_gold) + assert tensor_sha256(dx) == tensor_sha256(dx_gold) + + +@requires_cuda +def test_batch_padding_invariance(provider): + """fwd(x)[t] must equal fwd(x[t:t+1]) byte-for-byte (Axis-A invariance).""" + batch = fixtures.make_shared_batch("shared_t16").to("cuda") + y_full, _ = provider.shared_expert_mlp_fwd(batch) + for t in range(batch.x.shape[0]): + row_batch = SharedBatch( + x=batch.x[t : t + 1].contiguous(), + w_fc1=batch.w_fc1, + w_fc2=batch.w_fc2, + ) + y_row, _ = provider.shared_expert_mlp_fwd(row_batch) + assert tensor_sha256(y_row) == tensor_sha256(y_full[t : t + 1]), f"row {t} diverged" + + +@requires_cuda +def test_frozen_weights_no_dw(provider): + """Backward returns only dX; the shared base weights stay frozen.""" + batch = fixtures.make_shared_batch("shared_t16").to("cuda") + w1_before = tensor_sha256(batch.w_fc1) + w2_before = tensor_sha256(batch.w_fc2) + y, saved = provider.shared_expert_mlp_fwd(batch) + dy = fixtures.make_grad_output("shared_t16", tuple(y.shape)).to("cuda") + dx = provider.shared_expert_mlp_bwd(dy, batch, saved) + assert dx.shape == batch.x.shape and dx.dtype == torch.float32 + assert batch.w_fc1.grad is None and batch.w_fc2.grad is None + assert not batch.w_fc1.requires_grad and not batch.w_fc2.requires_grad + assert tensor_sha256(batch.w_fc1) == w1_before + assert tensor_sha256(batch.w_fc2) == w2_before + + +@requires_cuda +def test_shared_output_independent_of_routed(provider): + """Shared output is not premixed with the routed path (boundary fixture). + + Running the full routed pipeline (any p_s, any expert batch) between two + shared calls must not change a single byte of the shared output, and the + shared output must equal the standalone oracle result (no p_s applied). + """ + shared = fixtures.make_shared_batch("shared_t16").to("cuda") + y_gold, _ = oracle.shared_expert_mlp_fwd(shared) + y_before, _ = provider.shared_expert_mlp_fwd(shared) + + routed = fixtures.make_expert_batch("base_plus_lora").to("cuda") + y_routed, saved_routed = oracle.routed_expert_forward(routed, ops=provider) + dy_routed = fixtures.make_grad_output("base_plus_lora", tuple(y_routed.shape)).to("cuda") + oracle.routed_expert_backward(routed, saved_routed, dy_routed, ops=provider) + + y_after, _ = provider.shared_expert_mlp_fwd(shared) + assert tensor_sha256(y_before) == tensor_sha256(y_gold) + assert tensor_sha256(y_after) == tensor_sha256(y_gold) + assert y_after.data_ptr() != shared.x.data_ptr() + + +@requires_cuda +def test_cuda_triton_byte_equal(): + """The two backends agree with each other bit-for-bit.""" + from rl_engine.moe.provider import resolve_provider + + providers = [] + for spec in PROVIDER_SPECS.values(): + try: + providers.append(resolve_provider(spec)) + except NotImplementedError as exc: + pytest.skip(f"backend unavailable: {exc}") + batch = fixtures.make_shared_batch("shared_t16").to("cuda") + dy = fixtures.make_grad_output("shared_t16", (batch.x.shape[0], batch.x.shape[1])).to("cuda") + results = [_run_provider(p, batch, dy) for p in providers] + (y_a, dx_a), (y_b, dx_b) = results + assert tensor_sha256(y_a) == tensor_sha256(y_b) + assert tensor_sha256(dx_a) == tensor_sha256(dx_b) + + +@requires_cuda +def test_fail_closed_on_cpu_input(provider): + batch = fixtures.make_shared_batch("shared_t1") # stays on CPU + with pytest.raises(NotImplementedError): + provider.shared_expert_mlp_fwd(batch) From 42aa92771379a0efd2d744ebf9c2468debdb9e8e Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Fri, 4 Sep 2026 04:02:02 -0700 Subject: [PATCH 04/13] fix(p5-5): use libdevice exp in the Triton SwiGLU core tl.exp is the fast exp2-based path and does not bit-match torch.sigmoid; libdevice __nv_expf does (0/4M mismatches on the device probe). Signed-off-by: Yizheng Jiao --- rl_engine/kernels/ops/triton/moe/shared_expert.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/rl_engine/kernels/ops/triton/moe/shared_expert.py b/rl_engine/kernels/ops/triton/moe/shared_expert.py index 125d8a83..95b415bc 100644 --- a/rl_engine/kernels/ops/triton/moe/shared_expert.py +++ b/rl_engine/kernels/ops/triton/moe/shared_expert.py @@ -18,6 +18,7 @@ try: import triton import triton.language as tl + import triton.language.extra.libdevice as tld TRITON_AVAILABLE = True except ImportError: # pragma: no cover - exercised on non-GPU installs @@ -72,7 +73,8 @@ def _swiglu_shared_fwd_kernel( gate_index = row * (2 * width) + col g = tl.load(Z + gate_index, mask=mask, other=0.0) u = tl.load(Z + gate_index + width, mask=mask, other=0.0) - sig = 1.0 / (1.0 + tl.exp(-g)) + # libdevice exp (__nv_expf) bit-matches torch.sigmoid; tl.exp does not. + sig = 1.0 / (1.0 + tld.exp(-g)) silu = g * sig h = (silu * u).to(tl.bfloat16) tl.store(H + offs, h, mask=mask) @@ -97,7 +99,7 @@ def _swiglu_shared_bwd_kernel( g = tl.load(Z + gate_index, mask=mask, other=0.0) u = tl.load(Z + gate_index + width, mask=mask, other=0.0) dh = tl.load(DH + offs, mask=mask, other=0.0).to(tl.float32) - sig = 1.0 / (1.0 + tl.exp(-g)) + sig = 1.0 / (1.0 + tld.exp(-g)) silu = g * sig # dsilu = sig * (1 + g * (1 - sig)), each op rounded separately. t = 1.0 - sig From d1ace117b3394b40ec25405d787094f790a8119e Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Fri, 4 Sep 2026 09:28:09 -0700 Subject: [PATCH 05/13] fix(p5-5): source Triton sigmoid from torch.sigmoid Neither tl.exp (exp2-based) nor libdevice __nv_expf bit-matches the nvcc expf inside torch.sigmoid (~45% / ~10% of fp32 values differ by 1 ulp); the tiny fixtures passed only because the BF16 round absorbed the difference, and the T=256 benchmark cross-check caught the divergence. The Triton path now takes torch.sigmoid(gate) as a kernel input and fuses the remaining SwiGLU math; a (256, 1024, 512) cross-backend byte-equality test locks the regression in. Signed-off-by: Yizheng Jiao --- docs/operators/shared-expert-mlp.md | 6 +++++ .../kernels/ops/triton/moe/shared_expert.py | 22 ++++++++++++++----- tests/test_shared_expert_mlp.py | 19 ++++++++++++---- 3 files changed, 37 insertions(+), 10 deletions(-) diff --git a/docs/operators/shared-expert-mlp.md b/docs/operators/shared-expert-mlp.md index c4ac1809..e06086e5 100644 --- a/docs/operators/shared-expert-mlp.md +++ b/docs/operators/shared-expert-mlp.md @@ -20,6 +20,12 @@ owns one output element and reduces serially in ascending k with there is no cross-lane floating-point reduction and results are batch/padding invariant. +Sigmoid is transcendental, so its bits follow the libm implementation: nvcc +`expf` (used by both `torch.sigmoid` and the CUDA kernel) bit-matches, while +`tl.exp` and libdevice `__nv_expf` do not (~10-45% of values differ by 1 ulp). +The Triton path therefore sources the sigmoid tensor from `torch.sigmoid` and +fuses the remaining SwiGLU math in the kernel. + ## Backends | backend | entry point | diff --git a/rl_engine/kernels/ops/triton/moe/shared_expert.py b/rl_engine/kernels/ops/triton/moe/shared_expert.py index 95b415bc..f510f239 100644 --- a/rl_engine/kernels/ops/triton/moe/shared_expert.py +++ b/rl_engine/kernels/ops/triton/moe/shared_expert.py @@ -9,6 +9,12 @@ The one-round SwiGLU core (FP32 math, single BF16 round on the output) is the shared-mode (``p_s = None``, no clamp) variant shared with P5-2 (#63). + +Sigmoid is transcendental and its bits depend on the libm implementation: +neither ``tl.exp`` (exp2-based) nor libdevice ``__nv_expf`` bit-matches the +nvcc ``expf`` inside ``torch.sigmoid``. The Triton path therefore sources the +sigmoid tensor from ``torch.sigmoid`` (same-device oracle parity, per the P5 +transcendental rule) and fuses all remaining SwiGLU math in the kernel. """ from __future__ import annotations @@ -18,7 +24,6 @@ try: import triton import triton.language as tl - import triton.language.extra.libdevice as tld TRITON_AVAILABLE = True except ImportError: # pragma: no cover - exercised on non-GPU installs @@ -58,6 +63,7 @@ def _strict_gemm_kernel( @triton.jit def _swiglu_shared_fwd_kernel( Z, + SIG, H, n_elem, width, @@ -65,6 +71,7 @@ def _swiglu_shared_fwd_kernel( ): # One-round SwiGLU, shared mode: h = BF16(SiLU(gate) * up), FP32 math, # gate = z[:, :F], up = z[:, F:] packed in one [T, 2F] FP32 tensor. + # SIG is torch.sigmoid(gate) (see module docstring). pid = tl.program_id(0) offs = pid * BLOCK + tl.arange(0, BLOCK) mask = offs < n_elem @@ -73,8 +80,7 @@ def _swiglu_shared_fwd_kernel( gate_index = row * (2 * width) + col g = tl.load(Z + gate_index, mask=mask, other=0.0) u = tl.load(Z + gate_index + width, mask=mask, other=0.0) - # libdevice exp (__nv_expf) bit-matches torch.sigmoid; tl.exp does not. - sig = 1.0 / (1.0 + tld.exp(-g)) + sig = tl.load(SIG + offs, mask=mask, other=0.0) silu = g * sig h = (silu * u).to(tl.bfloat16) tl.store(H + offs, h, mask=mask) @@ -83,6 +89,7 @@ def _swiglu_shared_fwd_kernel( def _swiglu_shared_bwd_kernel( DH, Z, + SIG, DZ, n_elem, width, @@ -99,7 +106,7 @@ def _swiglu_shared_bwd_kernel( g = tl.load(Z + gate_index, mask=mask, other=0.0) u = tl.load(Z + gate_index + width, mask=mask, other=0.0) dh = tl.load(DH + offs, mask=mask, other=0.0).to(tl.float32) - sig = 1.0 / (1.0 + tld.exp(-g)) + sig = tl.load(SIG + offs, mask=mask, other=0.0) silu = g * sig # dsilu = sig * (1 + g * (1 - sig)), each op rounded separately. t = 1.0 - sig @@ -161,9 +168,10 @@ def swiglu_shared_fwd(z: torch.Tensor) -> torch.Tensor: n_elem = h.numel() if n_elem == 0: return h + sig = torch.sigmoid(z[:, :width]).contiguous() block = 1024 grid = (triton.cdiv(n_elem, block),) - _swiglu_shared_fwd_kernel[grid](z, h, n_elem, width, BLOCK=block) + _swiglu_shared_fwd_kernel[grid](z, sig, h, n_elem, width, BLOCK=block) return h @@ -181,7 +189,9 @@ def swiglu_shared_bwd(dh: torch.Tensor, z: torch.Tensor) -> torch.Tensor: n_elem = dh.numel() if n_elem == 0: return dz + width = dh.shape[1] + sig = torch.sigmoid(z[:, :width]).contiguous() block = 1024 grid = (triton.cdiv(n_elem, block),) - _swiglu_shared_bwd_kernel[grid](dh, z, dz, n_elem, dh.shape[1], BLOCK=block) + _swiglu_shared_bwd_kernel[grid](dh, z, sig, dz, n_elem, width, BLOCK=block) return dz diff --git a/tests/test_shared_expert_mlp.py b/tests/test_shared_expert_mlp.py index 6590ceea..a114d711 100644 --- a/tests/test_shared_expert_mlp.py +++ b/tests/test_shared_expert_mlp.py @@ -112,8 +112,13 @@ def test_shared_output_independent_of_routed(provider): @requires_cuda -def test_cuda_triton_byte_equal(): - """The two backends agree with each other bit-for-bit.""" +@pytest.mark.parametrize("shape", [(16, 128, 64), (256, 1024, 512)]) +def test_cuda_triton_byte_equal(shape): + """The two backends agree with each other bit-for-bit. + + The larger shape samples enough values to expose rare transcendental + 1-ulp divergences that survive the BF16 round (caught once at T=256). + """ from rl_engine.moe.provider import resolve_provider providers = [] @@ -122,8 +127,14 @@ def test_cuda_triton_byte_equal(): providers.append(resolve_provider(spec)) except NotImplementedError as exc: pytest.skip(f"backend unavailable: {exc}") - batch = fixtures.make_shared_batch("shared_t16").to("cuda") - dy = fixtures.make_grad_output("shared_t16", (batch.x.shape[0], batch.x.shape[1])).to("cuda") + t, hidden, ffn = shape + gen = torch.Generator(device="cpu").manual_seed(hash(shape) % (2**31)) + batch = SharedBatch( + x=torch.randn(t, hidden, generator=gen).to(torch.bfloat16).cuda(), + w_fc1=(torch.randn(2 * ffn, hidden, generator=gen) / hidden**0.5).to(torch.bfloat16).cuda(), + w_fc2=(torch.randn(hidden, ffn, generator=gen) / ffn**0.5).to(torch.bfloat16).cuda(), + ) + dy = torch.randn(t, hidden, generator=gen).to(torch.bfloat16).cuda() results = [_run_provider(p, batch, dy) for p in providers] (y_a, dx_a), (y_b, dx_b) = results assert tensor_sha256(y_a) == tensor_sha256(y_b) From bef20ee886f4644795b7c26889a9d9dd602007e1 Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Fri, 4 Sep 2026 09:58:41 -0700 Subject: [PATCH 06/13] fix(p5-5): forbid FMA contraction in the Triton kernels via libdevice _rn ops The compiler may contract a * b + c into an FMA; at T=256 that rounded dsilu = sig * (1 + g * (1 - sig)) differently on 2/262144 dgate elements (1 ulp after the BF16 round). All mul/add/sub in the Triton strict GEMM and SwiGLU kernels now go through libdevice add_rn/mul_rn/sub_rn, the exact Triton spelling of the CUDA kernel's __fadd_rn/__fmul_rn/__fsub_rn. Signed-off-by: Yizheng Jiao --- .../kernels/ops/triton/moe/shared_expert.py | 27 +++++++++++-------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/rl_engine/kernels/ops/triton/moe/shared_expert.py b/rl_engine/kernels/ops/triton/moe/shared_expert.py index f510f239..d65a10f4 100644 --- a/rl_engine/kernels/ops/triton/moe/shared_expert.py +++ b/rl_engine/kernels/ops/triton/moe/shared_expert.py @@ -15,6 +15,11 @@ nvcc ``expf`` inside ``torch.sigmoid``. The Triton path therefore sources the sigmoid tensor from ``torch.sigmoid`` (same-device oracle parity, per the P5 transcendental rule) and fuses all remaining SwiGLU math in the kernel. + +All mul/add/sub go through libdevice ``*_rn`` (the Triton spelling of CUDA's +``__fmul_rn``/``__fadd_rn``/``__fsub_rn``): the compiler is allowed to +contract a plain ``a * b + c`` into an FMA, which changes the rounding (seen +as 1-ulp dgate drift at T=256), and the ``_rn`` intrinsics forbid that. """ from __future__ import annotations @@ -24,6 +29,7 @@ try: import triton import triton.language as tl + import triton.language.extra.libdevice as tld TRITON_AVAILABLE = True except ImportError: # pragma: no cover - exercised on non-GPU installs @@ -56,8 +62,7 @@ def _strict_gemm_kernel( for k in range(0, K): a = tl.load(a_row + k).to(tl.float32) b = tl.load(b_cols + k * stride_bk, mask=mask_n, other=0.0).to(tl.float32) - prod = a * b - acc = acc + prod + acc = tld.add_rn(acc, tld.mul_rn(a, b)) tl.store(C + m * N + offs_n, acc, mask=mask_n) @triton.jit @@ -81,8 +86,8 @@ def _swiglu_shared_fwd_kernel( g = tl.load(Z + gate_index, mask=mask, other=0.0) u = tl.load(Z + gate_index + width, mask=mask, other=0.0) sig = tl.load(SIG + offs, mask=mask, other=0.0) - silu = g * sig - h = (silu * u).to(tl.bfloat16) + silu = tld.mul_rn(g, sig) + h = tld.mul_rn(silu, u).to(tl.bfloat16) tl.store(H + offs, h, mask=mask) @triton.jit @@ -107,14 +112,14 @@ def _swiglu_shared_bwd_kernel( u = tl.load(Z + gate_index + width, mask=mask, other=0.0) dh = tl.load(DH + offs, mask=mask, other=0.0).to(tl.float32) sig = tl.load(SIG + offs, mask=mask, other=0.0) - silu = g * sig + silu = tld.mul_rn(g, sig) # dsilu = sig * (1 + g * (1 - sig)), each op rounded separately. - t = 1.0 - sig - t = g * t - t = 1.0 + t - dsilu = sig * t - dgate = (dh * u) * dsilu - dup = dh * silu + t = tld.sub_rn(1.0, sig) + t = tld.mul_rn(g, t) + t = tld.add_rn(1.0, t) + dsilu = tld.mul_rn(sig, t) + dgate = tld.mul_rn(tld.mul_rn(dh, u), dsilu) + dup = tld.mul_rn(dh, silu) tl.store(DZ + gate_index, dgate.to(tl.bfloat16), mask=mask) tl.store(DZ + gate_index + width, dup.to(tl.bfloat16), mask=mask) From e4f8c142487ff4d4da227f0a84fe4e05c943e8d1 Mon Sep 17 00:00:00 2001 From: Hsiu-I Liao Date: Sun, 6 Sep 2026 10:38:01 +0800 Subject: [PATCH 07/13] feat: add deterministic clamp swiglu weighted Signed-off-by: Hsiu-I Liao --- csrc/cuda/activation.cu | 384 ++++++++++++++++++++ csrc/ops.cpp | 51 +++ rl_engine/moe/cuda_provider.py | 196 ++++++++++ tests/test_p5_clamp_swiglu_weighted_cuda.py | 288 +++++++++++++++ 4 files changed, 919 insertions(+) create mode 100644 rl_engine/moe/cuda_provider.py create mode 100644 tests/test_p5_clamp_swiglu_weighted_cuda.py diff --git a/csrc/cuda/activation.cu b/csrc/cuda/activation.cu index 505654f6..07e5c89c 100644 --- a/csrc/cuda/activation.cu +++ b/csrc/cuda/activation.cu @@ -93,6 +93,164 @@ __global__ void swiglu_backward_kernel( d_gate[idx] = static_cast(dyv * uv * silu_grad_f32(gv)); } +__device__ __forceinline__ float sigmoid_f32_strict(float x) { // new + const float denominator = __fadd_rn(1.0f, expf(-x)); + return 1.0f / denominator; +} + +__global__ void clamp_swiglu_weighted_forward_kernel( + const float* __restrict__ gate, + const float* __restrict__ up, + const float* __restrict__ p_s, + at::BFloat16* __restrict__ h, + float* __restrict__ g_saved, + float* __restrict__ u_saved, + float* __restrict__ sig_saved, + float* __restrict__ silu_saved, + const int64_t n, + const int64_t width, + const bool weighted) { + const int64_t idx = + blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + + if (idx >= n) { + return; + } + + const float gate_value = gate[idx]; + const float up_value = up[idx]; + + float g = gate_value; + float u = up_value; + + if (weighted) { + if (g > 10.0f) { + g = 10.0f; + } + + if (u < -10.0f) { + u = -10.0f; + } else if (u > 10.0f) { + u = 10.0f; + } + } + + const float sig = sigmoid_f32_strict(g); + const float silu = __fmul_rn(g, sig); + const float product = __fmul_rn(silu, u); + + const float h32 = + weighted + ? __fmul_rn(product, p_s[idx / width]) + : product; + + h[idx] = static_cast(h32); + + g_saved[idx] = g; + u_saved[idx] = u; + sig_saved[idx] = sig; + silu_saved[idx] = silu; +} + +template +__global__ void clamp_swiglu_weighted_backward_kernel( + const scalar_t* __restrict__ dh, + const float* __restrict__ gate, + const float* __restrict__ up, + const float* __restrict__ g, + const float* __restrict__ u, + const float* __restrict__ sig, + const float* __restrict__ silu, + const float* __restrict__ p_s, + float* __restrict__ d_gate, + float* __restrict__ d_up, + const int64_t n, + const int64_t width, + const bool weighted) { + const int64_t idx = + blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + + if (idx >= n) { + return; + } + + const float dh32 = static_cast(dh[idx]); + + const float weighted_dh = + weighted + ? __fmul_rn(dh32, p_s[idx / width]) + : dh32; + + const float one_minus_sig = + __fadd_rn(1.0f, -sig[idx]); + + const float derivative_inner = + __fadd_rn( + 1.0f, + __fmul_rn(g[idx], one_minus_sig)); + + const float d_silu = + __fmul_rn(sig[idx], derivative_inner); + + const float gate_mask = + (!weighted || gate[idx] < 10.0f) + ? 1.0f + : 0.0f; + + const float up_mask = + (!weighted || + (up[idx] > -10.0f && up[idx] < 10.0f)) + ? 1.0f + : 0.0f; + + d_gate[idx] = + __fmul_rn( + __fmul_rn( + __fmul_rn(weighted_dh, u[idx]), + d_silu), + gate_mask); + + d_up[idx] = + __fmul_rn( + __fmul_rn(weighted_dh, silu[idx]), + up_mask); +} + +template +__global__ void clamp_swiglu_weighted_dp_s_kernel( + const scalar_t* __restrict__ dh, + const float* __restrict__ silu, + const float* __restrict__ u, + float* __restrict__ dp_s, + const int64_t rows, + const int64_t width) { + const int64_t row = + blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + + if (row >= rows) { + return; + } + + float acc = 0.0f; + const int64_t row_offset = row * width; + +#pragma unroll 1 + for (int64_t column = 0; column < width; ++column) { + const int64_t idx = row_offset + column; + + const float term = + __fmul_rn( + __fmul_rn( + static_cast(dh[idx]), + silu[idx]), + u[idx]); + + acc = __fadd_rn(acc, term); + } + + dp_s[row] = acc; +} + template __global__ void swiglu_packed_forward_kernel( const scalar_t* __restrict__ gate_up, @@ -170,6 +328,67 @@ static void check_same_device( rhs.device()); } +static void check_cuda_contig_fp32( // new + const torch::Tensor& t, + const char* name) { + TORCH_CHECK( + t.is_cuda(), + name, + " must be a CUDA tensor"); + + TORCH_CHECK( + t.is_contiguous(), + name, + " must be contiguous"); + + TORCH_CHECK( + t.scalar_type() == at::kFloat, + name, + " must be float32"); +} + +static void check_same_shape_2d( + const torch::Tensor& lhs, + const torch::Tensor& rhs, + const char* lhs_name, + const char* rhs_name) { + TORCH_CHECK( + lhs.dim() == 2, + lhs_name, + " must be 2D [rows, width]"); + + TORCH_CHECK( + rhs.dim() == 2, + rhs_name, + " must be 2D [rows, width]"); + + TORCH_CHECK( + lhs.sizes() == rhs.sizes(), + lhs_name, + " and ", + rhs_name, + " must share shape"); +} + +static void check_route_weights( + const torch::optional& p_s, + const torch::Tensor& reference) { + if (!p_s.has_value()) { + return; + } + + check_cuda_contig_fp32(*p_s, "p_s"); + check_same_device(*p_s, reference, "p_s", "gate"); + + TORCH_CHECK( + p_s->dim() == 1, + "p_s must be 1D [rows]"); + + TORCH_CHECK( + p_s->size(0) == reference.size(0), + "p_s must have shape [rows]"); +} + } // namespace torch::Tensor silu_forward_cuda(torch::Tensor x) { @@ -364,3 +583,168 @@ std::vector swiglu_packed_backward_cuda( C10_CUDA_KERNEL_LAUNCH_CHECK(); return {d_gate, d_up}; } + +std::vector clamp_swiglu_weighted_forward_cuda( // new + torch::Tensor gate, + torch::Tensor up, + torch::optional p_s) { + check_cuda_contig_fp32(gate, "gate"); + check_cuda_contig_fp32(up, "up"); + check_same_device(gate, up, "gate", "up"); + check_same_shape_2d(gate, up, "gate", "up"); + check_route_weights(p_s, gate); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(gate)); + + auto h = + torch::empty( + gate.sizes(), + gate.options().dtype(torch::kBFloat16)); + + auto g = torch::empty_like(gate); + auto u = torch::empty_like(up); + auto sig = torch::empty_like(gate); + auto silu = torch::empty_like(gate); + + const int64_t n = gate.numel(); + + if (n == 0) { + return {h, g, u, sig, silu}; + } + + int threads = 0; + int64_t blocks = 0; + + launch_1d(n, threads, blocks); + + auto stream = at::cuda::getCurrentCUDAStream(); + + clamp_swiglu_weighted_forward_kernel + <<>>( + gate.data_ptr(), + up.data_ptr(), + p_s.has_value() + ? p_s->data_ptr() + : nullptr, + h.data_ptr(), + g.data_ptr(), + u.data_ptr(), + sig.data_ptr(), + silu.data_ptr(), + n, + gate.size(1), + p_s.has_value()); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + return {h, g, u, sig, silu}; +} + +std::vector clamp_swiglu_weighted_backward_cuda( + torch::Tensor dh, + torch::Tensor gate, + torch::Tensor up, + torch::Tensor g, + torch::Tensor u, + torch::Tensor sig, + torch::Tensor silu, + torch::optional p_s) { + check_cuda_contig(dh, "dh"); + check_cuda_contig_fp32(gate, "gate"); + check_cuda_contig_fp32(up, "up"); + check_cuda_contig_fp32(g, "g"); + check_cuda_contig_fp32(u, "u"); + check_cuda_contig_fp32(sig, "sig"); + check_cuda_contig_fp32(silu, "silu"); + + check_same_device(dh, gate, "dh", "gate"); + check_same_device(gate, up, "gate", "up"); + check_same_device(gate, g, "gate", "g"); + check_same_device(gate, u, "gate", "u"); + check_same_device(gate, sig, "gate", "sig"); + check_same_device(gate, silu, "gate", "silu"); + + check_same_shape_2d(gate, up, "gate", "up"); + check_same_shape_2d(gate, dh, "gate", "dh"); + check_same_shape_2d(gate, g, "gate", "g"); + check_same_shape_2d(gate, u, "gate", "u"); + check_same_shape_2d(gate, sig, "gate", "sig"); + check_same_shape_2d(gate, silu, "gate", "silu"); + + check_route_weights(p_s, gate); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(gate)); + + auto d_gate = torch::empty_like(gate); + auto d_up = torch::empty_like(up); + + auto dp_s = + p_s.has_value() + ? torch::zeros( + {gate.size(0)}, + gate.options()) + : torch::empty( + {0}, + gate.options()); + + const int64_t n = gate.numel(); + + if (n == 0) { + return {d_gate, d_up, dp_s}; + } + + int threads = 0; + int64_t blocks = 0; + + launch_1d(n, threads, blocks); + + auto stream = at::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + dh.scalar_type(), + "clamp_swiglu_weighted_backward_cuda", + [&] { + clamp_swiglu_weighted_backward_kernel + <<>>( + dh.data_ptr(), + gate.data_ptr(), + up.data_ptr(), + g.data_ptr(), + u.data_ptr(), + sig.data_ptr(), + silu.data_ptr(), + p_s.has_value() + ? p_s->data_ptr() + : nullptr, + d_gate.data_ptr(), + d_up.data_ptr(), + n, + gate.size(1), + p_s.has_value()); + + if (p_s.has_value()) { + int dp_threads = 0; + int64_t dp_blocks = 0; + + launch_1d( + gate.size(0), + dp_threads, + dp_blocks); + + clamp_swiglu_weighted_dp_s_kernel + <<>>( + dh.data_ptr(), + silu.data_ptr(), + u.data_ptr(), + dp_s.data_ptr(), + gate.size(0), + gate.size(1)); + } + }); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + return {d_gate, d_up, dp_s}; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index aecc30ed..336b4498 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -130,7 +130,21 @@ torch::Tensor swiglu_packed_forward_cuda(torch::Tensor gate_up); std::vector swiglu_packed_backward_cuda( torch::Tensor dy, torch::Tensor gate_up); +// P5-2 ClampSwiGLU weighted declarations +std::vector clamp_swiglu_weighted_forward_cuda( + torch::Tensor gate, + torch::Tensor up, + torch::optional p_s); +std::vector clamp_swiglu_weighted_backward_cuda( + torch::Tensor dh, + torch::Tensor gate, + torch::Tensor up, + torch::Tensor g, + torch::Tensor u, + torch::Tensor sig, + torch::Tensor silu, + torch::optional p_s); // RMSNorm Declarations & Wrappers void rmsnorm_forward_cuda( @@ -293,6 +307,35 @@ std::vector swiglu_packed_backward( return swiglu_packed_backward_cuda(dy, gate_up); } +std::vector clamp_swiglu_weighted_forward( + torch::Tensor gate, + torch::Tensor up, + torch::optional p_s) { + return clamp_swiglu_weighted_forward_cuda( + gate, + up, + p_s); +} + +std::vector clamp_swiglu_weighted_backward( + torch::Tensor dh, + torch::Tensor gate, + torch::Tensor up, + torch::Tensor g, + torch::Tensor u, + torch::Tensor sig, + torch::Tensor silu, + torch::optional p_s) { + return clamp_swiglu_weighted_backward_cuda( + dh, + gate, + up, + g, + u, + sig, + silu, + p_s); +} // Deterministic standard-softmax attention (issue #147) std::vector deterministic_attention_forward( torch::Tensor q, @@ -502,7 +545,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Batch-invariant SwiGLU forward for [rows, 2 * intermediate]"); m.def("swiglu_packed_backward", &swiglu_packed_backward, "Batch-invariant SwiGLU backward for [rows, 2 * intermediate]"); + m.def( + "clamp_swiglu_weighted_forward", + &clamp_swiglu_weighted_forward, + "P5 clamp_swiglu_weighted forward CUDA"); + m.def( + "clamp_swiglu_weighted_backward", + &clamp_swiglu_weighted_backward, + "P5 clamp_swiglu_weighted backward CUDA"); // Deterministic standard-softmax attention (issue #147) m.def( "deterministic_attention_forward", diff --git a/rl_engine/moe/cuda_provider.py b/rl_engine/moe/cuda_provider.py new file mode 100644 index 00000000..7ab8592b --- /dev/null +++ b/rl_engine/moe/cuda_provider.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""CUDA provider for P5-2 clamp_swiglu_weighted.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.moe.contract import ORACLE_PROFILE +from rl_engine.moe.provider import ReferenceProvider + +_FORWARD_SYMBOL = "clamp_swiglu_weighted_forward" +_BACKWARD_SYMBOL = "clamp_swiglu_weighted_backward" + +_SUPPORTED_DTYPES = ( + torch.float16, + torch.bfloat16, + torch.float32, +) + +_SAVED_KEYS = ( + "gate32", + "up32", + "g", + "u", + "sig", + "silu", +) + + +def _require_cuda_extension() -> None: + if not _EXT_AVAILABLE or _C is None: + raise RuntimeError( + "ClampSwiGLUWeightedCudaProvider requires the compiled " "rl_engine._C extension." + ) + + missing = [name for name in (_FORWARD_SYMBOL, _BACKWARD_SYMBOL) if not hasattr(_C, name)] + + if missing: + raise RuntimeError( + "P5-2 CUDA symbols are unavailable: " + f"{', '.join(missing)}. " + "Rebuild rl_engine._C from this branch." + ) + + +def _validate_matrix( + x: torch.Tensor, + name: str, +) -> None: + if x.device.type != "cuda": + raise RuntimeError(f"{name} must be a CUDA tensor, got {x.device}.") + + if x.dim() != 2: + raise ValueError(f"{name} must be 2D [rows, width], " f"got shape {tuple(x.shape)}.") + + if x.dtype not in _SUPPORTED_DTYPES: + raise TypeError(f"{name} must have dtype fp16, bf16, or fp32, " f"got {x.dtype}.") + + +class ClampSwiGLUWeightedCudaProvider(ReferenceProvider): + """P5-2 CUDA provider with deterministic FP32 computation.""" + + name = "cuda-clamp-swiglu-weighted" + numeric_profile = ORACLE_PROFILE + + def __init__(self) -> None: + _require_cuda_extension() + + def capabilities(self) -> dict[str, Any]: + return { + "backend": "cuda", + "delivered_ops": ["clamp_swiglu_weighted"], + "geometry": ["one-row", "packed"], + "devices": ["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__, + } + + def clamp_swiglu_weighted_fwd( + self, + gate: torch.Tensor, + up: torch.Tensor, + p_s: torch.Tensor | None, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + _validate_matrix(gate, "gate") + _validate_matrix(up, "up") + + if gate.device != up.device: + raise RuntimeError( + "gate and up must share a CUDA device, " f"got {gate.device} and {up.device}." + ) + + if gate.shape != up.shape: + raise ValueError("gate and up must share shape, " f"got {gate.shape} and {up.shape}.") + + if p_s is not None: + if p_s.device != gate.device: + raise RuntimeError( + "p_s and gate must share a CUDA device, " f"got {p_s.device} and {gate.device}." + ) + + if p_s.dtype != torch.float32: + raise TypeError(f"p_s must have dtype fp32, got {p_s.dtype}.") + + if p_s.shape != (gate.shape[0],): + raise ValueError( + f"p_s must have shape ({gate.shape[0]},), " f"got {tuple(p_s.shape)}." + ) + + gate32 = gate.float().contiguous() + up32 = up.float().contiguous() + p_s32 = None if p_s is None else p_s.contiguous() + + h, g, u, sig, silu = getattr( + _C, + _FORWARD_SYMBOL, + )( + gate32, + up32, + p_s32, + ) + + saved = { + "gate32": gate32, + "up32": up32, + "g": g, + "u": u, + "sig": sig, + "silu": silu, + } + + if p_s32 is not None: + saved["p_s"] = p_s32 + + return h, saved + + def clamp_swiglu_weighted_bwd( + self, + dh: torch.Tensor, + saved: dict[str, torch.Tensor], + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + ]: + missing = [key for key in _SAVED_KEYS if key not in saved] + + if missing: + raise KeyError(f"saved state is missing: {', '.join(missing)}") + + _validate_matrix(dh, "dh") + + gate32 = saved["gate32"] + + if dh.device != gate32.device: + raise RuntimeError( + "dh and saved tensors must share a CUDA device, " + f"got {dh.device} and {gate32.device}." + ) + + if dh.shape != gate32.shape: + raise ValueError( + "dh and saved tensors must share shape, " f"got {dh.shape} and {gate32.shape}." + ) + + p_s = saved.get("p_s") + + dgate, dup, dp_s = getattr( + _C, + _BACKWARD_SYMBOL, + )( + dh.contiguous(), + gate32, + saved["up32"], + saved["g"], + saved["u"], + saved["sig"], + saved["silu"], + p_s, + ) + + return ( + dgate, + dup, + dp_s if p_s is not None else None, + ) diff --git a/tests/test_p5_clamp_swiglu_weighted_cuda.py b/tests/test_p5_clamp_swiglu_weighted_cuda.py new file mode 100644 index 00000000..22a60b85 --- /dev/null +++ b/tests/test_p5_clamp_swiglu_weighted_cuda.py @@ -0,0 +1,288 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""CUDA tests for P5-2 clamp_swiglu_weighted.""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.moe import fixtures, oracle +from rl_engine.moe.cuda_provider import ( + ClampSwiGLUWeightedCudaProvider, +) + +_HAS_P5_CUDA = bool( + torch.cuda.is_available() + and _EXT_AVAILABLE + and _C is not None + and hasattr(_C, "clamp_swiglu_weighted_forward") + and hasattr(_C, "clamp_swiglu_weighted_backward") +) + +requires_p5_cuda = pytest.mark.skipif( + not _HAS_P5_CUDA, + reason="P5-2 CUDA extension unavailable", +) + + +def _assert_exact( + got: torch.Tensor, + want: torch.Tensor, +) -> None: + assert got.dtype == want.dtype + assert got.shape == want.shape + assert torch.equal(got, want) + + +@requires_p5_cuda +def test_boundary_fixture_is_byte_exact() -> None: + provider = ClampSwiGLUWeightedCudaProvider() + + gate, up, p_s = (tensor.cuda() for tensor in fixtures.make_swiglu_boundary_inputs()) + + dh = fixtures.make_grad_output( + "swiglu_boundary", + tuple(gate.shape), + ).cuda() + + h_ref, saved_ref = oracle.clamp_swiglu_weighted_fwd( + gate, + up, + p_s, + ) + + grads_ref = oracle.clamp_swiglu_weighted_bwd( + dh, + saved_ref, + ) + + h_got, saved_got = provider.clamp_swiglu_weighted_fwd( + gate, + up, + p_s, + ) + + grads_got = provider.clamp_swiglu_weighted_bwd( + dh, + saved_got, + ) + + _assert_exact(h_got, h_ref) + + for key in ("g", "u", "sig", "silu"): + _assert_exact( + saved_got[key], + saved_ref[key], + ) + + for got, want in zip( + grads_got, + grads_ref, + strict=True, + ): + assert got is not None + assert want is not None + _assert_exact(got, want) + + +@requires_p5_cuda +def test_shared_variant_has_no_weight_no_clamp_and_no_dp_s() -> None: + provider = ClampSwiGLUWeightedCudaProvider() + + gate = torch.tensor( + [[12.0, -12.0, 10.0, -10.0]], + device="cuda", + ) + + up = torch.tensor( + [[12.0, -12.0, 10.0, -10.0]], + device="cuda", + ) + + dh = torch.tensor( + [[1.0, -0.5, 0.25, -2.0]], + device="cuda", + dtype=torch.bfloat16, + ) + + h_ref, saved_ref = oracle.clamp_swiglu_weighted_fwd( + gate, + up, + None, + ) + + dgate_ref, dup_ref, dp_ref = oracle.clamp_swiglu_weighted_bwd( + dh, + saved_ref, + ) + + h_got, saved_got = provider.clamp_swiglu_weighted_fwd( + gate, + up, + None, + ) + + dgate_got, dup_got, dp_got = provider.clamp_swiglu_weighted_bwd( + dh, + saved_got, + ) + + _assert_exact(h_got, h_ref) + _assert_exact(saved_got["g"], gate) + _assert_exact(saved_got["u"], up) + _assert_exact(dgate_got, dgate_ref) + _assert_exact(dup_got, dup_ref) + + assert dp_got is None + assert dp_ref is None + + +@requires_p5_cuda +def test_repeated_runs_are_byte_exact() -> None: + provider = ClampSwiGLUWeightedCudaProvider() + + generator = torch.Generator().manual_seed(63) + + gate = ( + torch.randn( + 7, + 257, + generator=generator, + ) + * 12.0 + ).cuda() + + up = ( + torch.randn( + 7, + 257, + generator=generator, + ) + * 12.0 + ).cuda() + + p_s = torch.rand( + 7, + generator=generator, + ).cuda() + + dh = ( + torch.randn( + 7, + 257, + generator=generator, + ) + .to(torch.bfloat16) + .cuda() + ) + + h_first, saved_first = provider.clamp_swiglu_weighted_fwd( + gate, + up, + p_s, + ) + + grads_first = provider.clamp_swiglu_weighted_bwd( + dh, + saved_first, + ) + + for _ in range(4): + h_repeat, saved_repeat = provider.clamp_swiglu_weighted_fwd( + gate, + up, + p_s, + ) + + grads_repeat = provider.clamp_swiglu_weighted_bwd( + dh, + saved_repeat, + ) + + _assert_exact(h_repeat, h_first) + + for got, want in zip( + grads_repeat, + grads_first, + strict=True, + ): + assert got is not None + assert want is not None + _assert_exact(got, want) + + +@requires_p5_cuda +def test_empty_batch_and_fail_closed_validation() -> None: + provider = ClampSwiGLUWeightedCudaProvider() + + empty = torch.empty( + 0, + 32, + device="cuda", + ) + + p_s = torch.empty( + 0, + device="cuda", + ) + + dh = torch.empty( + 0, + 32, + device="cuda", + dtype=torch.bfloat16, + ) + + h, saved = provider.clamp_swiglu_weighted_fwd( + empty, + empty, + p_s, + ) + + dgate, dup, dp_s = provider.clamp_swiglu_weighted_bwd( + dh, + saved, + ) + + assert h.shape == (0, 32) + assert h.dtype == torch.bfloat16 + assert dgate.shape == (0, 32) + assert dup.shape == (0, 32) + assert dp_s is not None + assert dp_s.shape == (0,) + + with pytest.raises( + RuntimeError, + match="CUDA tensor", + ): + provider.clamp_swiglu_weighted_fwd( + empty.cpu(), + empty.cpu(), + p_s.cpu(), + ) + + with pytest.raises( + ValueError, + match="share shape", + ): + provider.clamp_swiglu_weighted_fwd( + empty, + torch.empty( + 0, + 16, + device="cuda", + ), + p_s, + ) + + with pytest.raises( + TypeError, + match="p_s must have dtype fp32", + ): + provider.clamp_swiglu_weighted_fwd( + empty, + empty, + p_s.to(torch.bfloat16), + ) From 4c7f1b337f02e03b30a5590af475a9be8fae633e Mon Sep 17 00:00:00 2001 From: Hsiu-I Liao Date: Sun, 6 Sep 2026 15:58:22 +0800 Subject: [PATCH 08/13] fix: expose optional p_s in P5 bindings Signed-off-by: Hsiu-I Liao --- csrc/ops.cpp | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 336b4498..b712863a 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -546,9 +546,17 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("swiglu_packed_backward", &swiglu_packed_backward, "Batch-invariant SwiGLU backward for [rows, 2 * intermediate]"); m.def( - "clamp_swiglu_weighted_forward", - &clamp_swiglu_weighted_forward, - "P5 clamp_swiglu_weighted forward CUDA"); + "clamp_swiglu_weighted_backward", + &clamp_swiglu_weighted_backward, + py::arg("dh"), + py::arg("gate"), + py::arg("up"), + py::arg("g"), + py::arg("u"), + py::arg("sig"), + py::arg("silu"), + py::arg("p_s") = py::none(), + "P5 clamp_swiglu_weighted backward CUDA"); m.def( "clamp_swiglu_weighted_backward", From b0784e60715c698f2b8cb12c6ecd148061da4e37 Mon Sep 17 00:00:00 2001 From: Hsiu-I Liao Date: Sun, 6 Sep 2026 22:54:14 +0800 Subject: [PATCH 09/13] fix: register P5 forward binding Signed-off-by: Hsiu-I Liao --- csrc/ops.cpp | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/csrc/ops.cpp b/csrc/ops.cpp index b712863a..15233f10 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -546,21 +546,24 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("swiglu_packed_backward", &swiglu_packed_backward, "Batch-invariant SwiGLU backward for [rows, 2 * intermediate]"); m.def( - "clamp_swiglu_weighted_backward", - &clamp_swiglu_weighted_backward, - py::arg("dh"), - py::arg("gate"), - py::arg("up"), - py::arg("g"), - py::arg("u"), - py::arg("sig"), - py::arg("silu"), - py::arg("p_s") = py::none(), - "P5 clamp_swiglu_weighted backward CUDA"); + "clamp_swiglu_weighted_forward", + &clamp_swiglu_weighted_forward, + py::arg("gate"), + py::arg("up"), + py::arg("p_s") = py::none(), + "P5 clamp_swiglu_weighted forward CUDA"); m.def( "clamp_swiglu_weighted_backward", &clamp_swiglu_weighted_backward, + py::arg("dh"), + py::arg("gate"), + py::arg("up"), + py::arg("g"), + py::arg("u"), + py::arg("sig"), + py::arg("silu"), + py::arg("p_s") = py::none(), "P5 clamp_swiglu_weighted backward CUDA"); // Deterministic standard-softmax attention (issue #147) m.def( From 0e88933bc3fc3591e9d4f52a44b1fcdf944262f1 Mon Sep 17 00:00:00 2001 From: Hsiu-I Liao Date: Tue, 8 Sep 2026 01:08:01 +0800 Subject: [PATCH 10/13] perf: recompute P5 SwiGLU intermediates in backward Signed-off-by: Hsiu-I Liao --- csrc/cuda/activation.cu | 117 ++++++++++---------- csrc/ops.cpp | 30 ++--- rl_engine/moe/cuda_provider.py | 14 +-- tests/test_p5_clamp_swiglu_weighted_cuda.py | 75 +++++++++++-- 4 files changed, 132 insertions(+), 104 deletions(-) diff --git a/csrc/cuda/activation.cu b/csrc/cuda/activation.cu index 07e5c89c..6d5e2f4c 100644 --- a/csrc/cuda/activation.cu +++ b/csrc/cuda/activation.cu @@ -93,7 +93,7 @@ __global__ void swiglu_backward_kernel( d_gate[idx] = static_cast(dyv * uv * silu_grad_f32(gv)); } -__device__ __forceinline__ float sigmoid_f32_strict(float x) { // new +__device__ __forceinline__ float sigmoid_f32_strict(float x) { const float denominator = __fadd_rn(1.0f, expf(-x)); return 1.0f / denominator; } @@ -103,10 +103,6 @@ __global__ void clamp_swiglu_weighted_forward_kernel( const float* __restrict__ up, const float* __restrict__ p_s, at::BFloat16* __restrict__ h, - float* __restrict__ g_saved, - float* __restrict__ u_saved, - float* __restrict__ sig_saved, - float* __restrict__ silu_saved, const int64_t n, const int64_t width, const bool weighted) { @@ -145,11 +141,6 @@ __global__ void clamp_swiglu_weighted_forward_kernel( : product; h[idx] = static_cast(h32); - - g_saved[idx] = g; - u_saved[idx] = u; - sig_saved[idx] = sig; - silu_saved[idx] = silu; } template @@ -157,10 +148,6 @@ __global__ void clamp_swiglu_weighted_backward_kernel( const scalar_t* __restrict__ dh, const float* __restrict__ gate, const float* __restrict__ up, - const float* __restrict__ g, - const float* __restrict__ u, - const float* __restrict__ sig, - const float* __restrict__ silu, const float* __restrict__ p_s, float* __restrict__ d_gate, float* __restrict__ d_up, @@ -174,6 +161,27 @@ __global__ void clamp_swiglu_weighted_backward_kernel( return; } + const float gate_value = gate[idx]; + const float up_value = up[idx]; + + float g = gate_value; + float u = up_value; + + if (weighted) { + if (g > 10.0f) { + g = 10.0f; + } + + if (u < -10.0f) { + u = -10.0f; + } else if (u > 10.0f) { + u = 10.0f; + } + } + + const float sig = sigmoid_f32_strict(g); + const float silu = __fmul_rn(g, sig); + const float dh32 = static_cast(dh[idx]); const float weighted_dh = @@ -182,45 +190,45 @@ __global__ void clamp_swiglu_weighted_backward_kernel( : dh32; const float one_minus_sig = - __fadd_rn(1.0f, -sig[idx]); + __fadd_rn(1.0f, -sig); const float derivative_inner = __fadd_rn( 1.0f, - __fmul_rn(g[idx], one_minus_sig)); + __fmul_rn(g, one_minus_sig)); const float d_silu = - __fmul_rn(sig[idx], derivative_inner); + __fmul_rn(sig, derivative_inner); const float gate_mask = - (!weighted || gate[idx] < 10.0f) + (!weighted || gate_value < 10.0f) ? 1.0f : 0.0f; const float up_mask = (!weighted || - (up[idx] > -10.0f && up[idx] < 10.0f)) + (up_value > -10.0f && up_value < 10.0f)) ? 1.0f : 0.0f; d_gate[idx] = __fmul_rn( __fmul_rn( - __fmul_rn(weighted_dh, u[idx]), + __fmul_rn(weighted_dh, u), d_silu), gate_mask); d_up[idx] = __fmul_rn( - __fmul_rn(weighted_dh, silu[idx]), + __fmul_rn(weighted_dh, silu), up_mask); } template __global__ void clamp_swiglu_weighted_dp_s_kernel( const scalar_t* __restrict__ dh, - const float* __restrict__ silu, - const float* __restrict__ u, + const float* __restrict__ gate, + const float* __restrict__ up, float* __restrict__ dp_s, const int64_t rows, const int64_t width) { @@ -238,12 +246,28 @@ __global__ void clamp_swiglu_weighted_dp_s_kernel( for (int64_t column = 0; column < width; ++column) { const int64_t idx = row_offset + column; + float g = gate[idx]; + float u = up[idx]; + + if (g > 10.0f) { + g = 10.0f; + } + + if (u < -10.0f) { + u = -10.0f; + } else if (u > 10.0f) { + u = 10.0f; + } + + const float sig = sigmoid_f32_strict(g); + const float silu = __fmul_rn(g, sig); + const float term = __fmul_rn( __fmul_rn( static_cast(dh[idx]), - silu[idx]), - u[idx]); + silu), + u); acc = __fadd_rn(acc, term); } @@ -328,7 +352,7 @@ static void check_same_device( rhs.device()); } -static void check_cuda_contig_fp32( // new +static void check_cuda_contig_fp32( const torch::Tensor& t, const char* name) { TORCH_CHECK( @@ -584,7 +608,7 @@ std::vector swiglu_packed_backward_cuda( return {d_gate, d_up}; } -std::vector clamp_swiglu_weighted_forward_cuda( // new +std::vector clamp_swiglu_weighted_forward_cuda( torch::Tensor gate, torch::Tensor up, torch::optional p_s) { @@ -601,15 +625,10 @@ std::vector clamp_swiglu_weighted_forward_cuda( // new gate.sizes(), gate.options().dtype(torch::kBFloat16)); - auto g = torch::empty_like(gate); - auto u = torch::empty_like(up); - auto sig = torch::empty_like(gate); - auto silu = torch::empty_like(gate); - const int64_t n = gate.numel(); if (n == 0) { - return {h, g, u, sig, silu}; + return {h}; } int threads = 0; @@ -627,49 +646,29 @@ std::vector clamp_swiglu_weighted_forward_cuda( // new ? p_s->data_ptr() : nullptr, h.data_ptr(), - g.data_ptr(), - u.data_ptr(), - sig.data_ptr(), - silu.data_ptr(), n, gate.size(1), p_s.has_value()); C10_CUDA_KERNEL_LAUNCH_CHECK(); - return {h, g, u, sig, silu}; + return {h}; } std::vector clamp_swiglu_weighted_backward_cuda( torch::Tensor dh, torch::Tensor gate, torch::Tensor up, - torch::Tensor g, - torch::Tensor u, - torch::Tensor sig, - torch::Tensor silu, torch::optional p_s) { check_cuda_contig(dh, "dh"); check_cuda_contig_fp32(gate, "gate"); check_cuda_contig_fp32(up, "up"); - check_cuda_contig_fp32(g, "g"); - check_cuda_contig_fp32(u, "u"); - check_cuda_contig_fp32(sig, "sig"); - check_cuda_contig_fp32(silu, "silu"); check_same_device(dh, gate, "dh", "gate"); check_same_device(gate, up, "gate", "up"); - check_same_device(gate, g, "gate", "g"); - check_same_device(gate, u, "gate", "u"); - check_same_device(gate, sig, "gate", "sig"); - check_same_device(gate, silu, "gate", "silu"); check_same_shape_2d(gate, up, "gate", "up"); check_same_shape_2d(gate, dh, "gate", "dh"); - check_same_shape_2d(gate, g, "gate", "g"); - check_same_shape_2d(gate, u, "gate", "u"); - check_same_shape_2d(gate, sig, "gate", "sig"); - check_same_shape_2d(gate, silu, "gate", "silu"); check_route_weights(p_s, gate); @@ -711,10 +710,6 @@ std::vector clamp_swiglu_weighted_backward_cuda( dh.data_ptr(), gate.data_ptr(), up.data_ptr(), - g.data_ptr(), - u.data_ptr(), - sig.data_ptr(), - silu.data_ptr(), p_s.has_value() ? p_s->data_ptr() : nullptr, @@ -736,8 +731,8 @@ std::vector clamp_swiglu_weighted_backward_cuda( clamp_swiglu_weighted_dp_s_kernel <<>>( dh.data_ptr(), - silu.data_ptr(), - u.data_ptr(), + gate.data_ptr(), + up.data_ptr(), dp_s.data_ptr(), gate.size(0), gate.size(1)); @@ -747,4 +742,4 @@ std::vector clamp_swiglu_weighted_backward_cuda( C10_CUDA_KERNEL_LAUNCH_CHECK(); return {d_gate, d_up, dp_s}; -} +} \ No newline at end of file diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 15233f10..20fef37f 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -140,10 +140,6 @@ std::vector clamp_swiglu_weighted_backward_cuda( torch::Tensor dh, torch::Tensor gate, torch::Tensor up, - torch::Tensor g, - torch::Tensor u, - torch::Tensor sig, - torch::Tensor silu, torch::optional p_s); // RMSNorm Declarations & Wrappers @@ -321,19 +317,11 @@ std::vector clamp_swiglu_weighted_backward( torch::Tensor dh, torch::Tensor gate, torch::Tensor up, - torch::Tensor g, - torch::Tensor u, - torch::Tensor sig, - torch::Tensor silu, torch::optional p_s) { return clamp_swiglu_weighted_backward_cuda( dh, gate, up, - g, - u, - sig, - silu, p_s); } // Deterministic standard-softmax attention (issue #147) @@ -554,17 +542,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "P5 clamp_swiglu_weighted forward CUDA"); m.def( - "clamp_swiglu_weighted_backward", - &clamp_swiglu_weighted_backward, - py::arg("dh"), - py::arg("gate"), - py::arg("up"), - py::arg("g"), - py::arg("u"), - py::arg("sig"), - py::arg("silu"), - py::arg("p_s") = py::none(), - "P5 clamp_swiglu_weighted backward CUDA"); + "clamp_swiglu_weighted_backward", + &clamp_swiglu_weighted_backward, + py::arg("dh"), + py::arg("gate"), + py::arg("up"), + py::arg("p_s") = py::none(), + "P5 clamp_swiglu_weighted backward CUDA"); // Deterministic standard-softmax attention (issue #147) m.def( "deterministic_attention_forward", diff --git a/rl_engine/moe/cuda_provider.py b/rl_engine/moe/cuda_provider.py index 7ab8592b..f47164d3 100644 --- a/rl_engine/moe/cuda_provider.py +++ b/rl_engine/moe/cuda_provider.py @@ -24,10 +24,6 @@ _SAVED_KEYS = ( "gate32", "up32", - "g", - "u", - "sig", - "silu", ) @@ -121,7 +117,7 @@ def clamp_swiglu_weighted_fwd( up32 = up.float().contiguous() p_s32 = None if p_s is None else p_s.contiguous() - h, g, u, sig, silu = getattr( + (h,) = getattr( _C, _FORWARD_SYMBOL, )( @@ -133,10 +129,6 @@ def clamp_swiglu_weighted_fwd( saved = { "gate32": gate32, "up32": up32, - "g": g, - "u": u, - "sig": sig, - "silu": silu, } if p_s32 is not None: @@ -182,10 +174,6 @@ def clamp_swiglu_weighted_bwd( dh.contiguous(), gate32, saved["up32"], - saved["g"], - saved["u"], - saved["sig"], - saved["silu"], p_s, ) diff --git a/tests/test_p5_clamp_swiglu_weighted_cuda.py b/tests/test_p5_clamp_swiglu_weighted_cuda.py index 22a60b85..07777311 100644 --- a/tests/test_p5_clamp_swiglu_weighted_cuda.py +++ b/tests/test_p5_clamp_swiglu_weighted_cuda.py @@ -36,6 +36,39 @@ def _assert_exact( assert torch.equal(got, want) +def _assert_saved_state( + saved: dict[str, torch.Tensor], + gate: torch.Tensor, + up: torch.Tensor, + p_s: torch.Tensor | None, +) -> None: + expected_keys = { + "gate32", + "up32", + } + + if p_s is not None: + expected_keys.add("p_s") + + assert set(saved) == expected_keys + + _assert_exact( + saved["gate32"], + gate.float().contiguous(), + ) + + _assert_exact( + saved["up32"], + up.float().contiguous(), + ) + + if p_s is not None: + _assert_exact( + saved["p_s"], + p_s.contiguous(), + ) + + @requires_p5_cuda def test_boundary_fixture_is_byte_exact() -> None: provider = ClampSwiGLUWeightedCudaProvider() @@ -71,11 +104,12 @@ def test_boundary_fixture_is_byte_exact() -> None: _assert_exact(h_got, h_ref) - for key in ("g", "u", "sig", "silu"): - _assert_exact( - saved_got[key], - saved_ref[key], - ) + _assert_saved_state( + saved_got, + gate, + up, + p_s, + ) for got, want in zip( grads_got, @@ -130,8 +164,14 @@ def test_shared_variant_has_no_weight_no_clamp_and_no_dp_s() -> None: ) _assert_exact(h_got, h_ref) - _assert_exact(saved_got["g"], gate) - _assert_exact(saved_got["u"], up) + + _assert_saved_state( + saved_got, + gate, + up, + None, + ) + _assert_exact(dgate_got, dgate_ref) _assert_exact(dup_got, dup_ref) @@ -189,6 +229,13 @@ def test_repeated_runs_are_byte_exact() -> None: saved_first, ) + _assert_saved_state( + saved_first, + gate, + up, + p_s, + ) + for _ in range(4): h_repeat, saved_repeat = provider.clamp_swiglu_weighted_fwd( gate, @@ -203,6 +250,13 @@ def test_repeated_runs_are_byte_exact() -> None: _assert_exact(h_repeat, h_first) + _assert_saved_state( + saved_repeat, + gate, + up, + p_s, + ) + for got, want in zip( grads_repeat, grads_first, @@ -246,6 +300,13 @@ def test_empty_batch_and_fail_closed_validation() -> None: saved, ) + _assert_saved_state( + saved, + empty, + empty, + p_s, + ) + assert h.shape == (0, 32) assert h.dtype == torch.bfloat16 assert dgate.shape == (0, 32) From 85a32114f2183861915df2c0758b64cfced59f23 Mon Sep 17 00:00:00 2001 From: Hsiu-I Liao Date: Tue, 8 Sep 2026 01:41:09 +0800 Subject: [PATCH 11/13] perf: optimize P5 SwiGLU state and packed layout Signed-off-by: Hsiu-I Liao --- csrc/cuda/activation.cu | 202 ++++++++++++++++++-- csrc/ops.cpp | 39 ++++ rl_engine/moe/cuda_provider.py | 112 ++++++++++- tests/test_p5_clamp_swiglu_weighted_cuda.py | 83 ++++++++ 4 files changed, 423 insertions(+), 13 deletions(-) diff --git a/csrc/cuda/activation.cu b/csrc/cuda/activation.cu index 6d5e2f4c..59616fd7 100644 --- a/csrc/cuda/activation.cu +++ b/csrc/cuda/activation.cu @@ -93,6 +93,10 @@ __global__ void swiglu_backward_kernel( d_gate[idx] = static_cast(dyv * uv * silu_grad_f32(gv)); } +#if defined(__USE_FAST_MATH__) +#error "P5 clamp_swiglu_weighted requires precise FP32 math; disable --use_fast_math" +#endif + __device__ __forceinline__ float sigmoid_f32_strict(float x) { const float denominator = __fadd_rn(1.0f, expf(-x)); return 1.0f / denominator; @@ -105,6 +109,7 @@ __global__ void clamp_swiglu_weighted_forward_kernel( at::BFloat16* __restrict__ h, const int64_t n, const int64_t width, + const int64_t input_stride, const bool weighted) { const int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; @@ -113,8 +118,12 @@ __global__ void clamp_swiglu_weighted_forward_kernel( return; } - const float gate_value = gate[idx]; - const float up_value = up[idx]; + const int64_t row = idx / width; + const int64_t column = idx - row * width; + const int64_t input_offset = row * input_stride + column; + + const float gate_value = gate[input_offset]; + const float up_value = up[input_offset]; float g = gate_value; float u = up_value; @@ -137,7 +146,7 @@ __global__ void clamp_swiglu_weighted_forward_kernel( const float h32 = weighted - ? __fmul_rn(product, p_s[idx / width]) + ? __fmul_rn(product, p_s[row]) : product; h[idx] = static_cast(h32); @@ -153,6 +162,7 @@ __global__ void clamp_swiglu_weighted_backward_kernel( float* __restrict__ d_up, const int64_t n, const int64_t width, + const int64_t input_stride, const bool weighted) { const int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; @@ -161,8 +171,12 @@ __global__ void clamp_swiglu_weighted_backward_kernel( return; } - const float gate_value = gate[idx]; - const float up_value = up[idx]; + const int64_t row = idx / width; + const int64_t column = idx - row * width; + const int64_t input_offset = row * input_stride + column; + + const float gate_value = gate[input_offset]; + const float up_value = up[input_offset]; float g = gate_value; float u = up_value; @@ -186,7 +200,7 @@ __global__ void clamp_swiglu_weighted_backward_kernel( const float weighted_dh = weighted - ? __fmul_rn(dh32, p_s[idx / width]) + ? __fmul_rn(dh32, p_s[row]) : dh32; const float one_minus_sig = @@ -231,7 +245,8 @@ __global__ void clamp_swiglu_weighted_dp_s_kernel( const float* __restrict__ up, float* __restrict__ dp_s, const int64_t rows, - const int64_t width) { + const int64_t width, + const int64_t input_stride) { const int64_t row = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; @@ -240,14 +255,17 @@ __global__ void clamp_swiglu_weighted_dp_s_kernel( } float acc = 0.0f; - const int64_t row_offset = row * width; + + const int64_t input_row_offset = row * input_stride; + const int64_t output_row_offset = row * width; #pragma unroll 1 for (int64_t column = 0; column < width; ++column) { - const int64_t idx = row_offset + column; + const int64_t input_idx = input_row_offset + column; + const int64_t output_idx = output_row_offset + column; - float g = gate[idx]; - float u = up[idx]; + float g = gate[input_idx]; + float u = up[input_idx]; if (g > 10.0f) { g = 10.0f; @@ -265,7 +283,7 @@ __global__ void clamp_swiglu_weighted_dp_s_kernel( const float term = __fmul_rn( __fmul_rn( - static_cast(dh[idx]), + static_cast(dh[output_idx]), silu), u); @@ -648,6 +666,7 @@ std::vector clamp_swiglu_weighted_forward_cuda( h.data_ptr(), n, gate.size(1), + gate.size(1), p_s.has_value()); C10_CUDA_KERNEL_LAUNCH_CHECK(); @@ -717,6 +736,7 @@ std::vector clamp_swiglu_weighted_backward_cuda( d_up.data_ptr(), n, gate.size(1), + gate.size(1), p_s.has_value()); if (p_s.has_value()) { @@ -735,11 +755,169 @@ std::vector clamp_swiglu_weighted_backward_cuda( up.data_ptr(), dp_s.data_ptr(), gate.size(0), + gate.size(1), gate.size(1)); } }); C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {d_gate, d_up, dp_s}; +} + +std::vector clamp_swiglu_weighted_packed_forward_cuda( + torch::Tensor gate_up, + torch::optional p_s) { + check_cuda_contig_fp32(gate_up, "gate_up"); + + TORCH_CHECK( + gate_up.dim() == 2, + "gate_up must be 2D [rows, 2 * width]"); + + TORCH_CHECK( + gate_up.size(1) % 2 == 0, + "gate_up width must be even"); + + check_route_weights(p_s, gate_up); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(gate_up)); + + const int64_t rows = gate_up.size(0); + const int64_t width = gate_up.size(1) / 2; + + auto h = torch::empty( + {rows, width}, + gate_up.options().dtype(torch::kBFloat16)); + + const int64_t n = h.numel(); + + if (n == 0) { + return {h}; + } + + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + + auto stream = at::cuda::getCurrentCUDAStream(); + + const float* gate_ptr = gate_up.data_ptr(); + const float* up_ptr = gate_ptr + width; + + clamp_swiglu_weighted_forward_kernel + <<>>( + gate_ptr, + up_ptr, + p_s.has_value() + ? p_s->data_ptr() + : nullptr, + h.data_ptr(), + n, + width, + 2 * width, + p_s.has_value()); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + return {h}; +} + +std::vector clamp_swiglu_weighted_packed_backward_cuda( + torch::Tensor dh, + torch::Tensor gate_up, + torch::optional p_s) { + check_cuda_contig(dh, "dh"); + check_cuda_contig_fp32(gate_up, "gate_up"); + check_same_device(dh, gate_up, "dh", "gate_up"); + + TORCH_CHECK( + gate_up.dim() == 2, + "gate_up must be 2D [rows, 2 * width]"); + + TORCH_CHECK( + gate_up.size(1) % 2 == 0, + "gate_up width must be even"); + + const int64_t rows = gate_up.size(0); + const int64_t width = gate_up.size(1) / 2; + + TORCH_CHECK( + dh.dim() == 2 && + dh.size(0) == rows && + dh.size(1) == width, + "dh must have shape [rows, width]"); + + check_route_weights(p_s, gate_up); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(gate_up)); + + auto d_gate = torch::empty( + {rows, width}, + gate_up.options()); + + auto d_up = torch::empty( + {rows, width}, + gate_up.options()); + + auto dp_s = + p_s.has_value() + ? torch::zeros({rows}, gate_up.options()) + : torch::empty({0}, gate_up.options()); + + const int64_t n = dh.numel(); + + if (n == 0) { + return {d_gate, d_up, dp_s}; + } + + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + + auto stream = at::cuda::getCurrentCUDAStream(); + + const float* gate_ptr = gate_up.data_ptr(); + const float* up_ptr = gate_ptr + width; + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + dh.scalar_type(), + "clamp_swiglu_weighted_packed_backward_cuda", + [&] { + clamp_swiglu_weighted_backward_kernel + <<>>( + dh.data_ptr(), + gate_ptr, + up_ptr, + p_s.has_value() + ? p_s->data_ptr() + : nullptr, + d_gate.data_ptr(), + d_up.data_ptr(), + n, + width, + 2 * width, + p_s.has_value()); + + if (p_s.has_value()) { + int dp_threads = 0; + int64_t dp_blocks = 0; + launch_1d(rows, dp_threads, dp_blocks); + + clamp_swiglu_weighted_dp_s_kernel + <<>>( + dh.data_ptr(), + gate_ptr, + up_ptr, + dp_s.data_ptr(), + rows, + width, + 2 * width); + } + }); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {d_gate, d_up, dp_s}; } \ No newline at end of file diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 20fef37f..ded684a7 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -141,6 +141,14 @@ std::vector clamp_swiglu_weighted_backward_cuda( torch::Tensor gate, torch::Tensor up, torch::optional p_s); + std::vector clamp_swiglu_weighted_packed_forward_cuda( + torch::Tensor gate_up, + torch::optional p_s); + +std::vector clamp_swiglu_weighted_packed_backward_cuda( + torch::Tensor dh, + torch::Tensor gate_up, + torch::optional p_s); // RMSNorm Declarations & Wrappers void rmsnorm_forward_cuda( @@ -324,6 +332,23 @@ std::vector clamp_swiglu_weighted_backward( up, p_s); } +std::vector clamp_swiglu_weighted_packed_forward( + torch::Tensor gate_up, + torch::optional p_s) { + return clamp_swiglu_weighted_packed_forward_cuda( + gate_up, + p_s); +} + +std::vector clamp_swiglu_weighted_packed_backward( + torch::Tensor dh, + torch::Tensor gate_up, + torch::optional p_s) { + return clamp_swiglu_weighted_packed_backward_cuda( + dh, + gate_up, + p_s); +} // Deterministic standard-softmax attention (issue #147) std::vector deterministic_attention_forward( torch::Tensor q, @@ -549,6 +574,20 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("up"), py::arg("p_s") = py::none(), "P5 clamp_swiglu_weighted backward CUDA"); + m.def( + "clamp_swiglu_weighted_packed_forward", + &clamp_swiglu_weighted_packed_forward, + py::arg("gate_up"), + py::arg("p_s") = py::none(), + "P5 packed clamp_swiglu_weighted forward CUDA"); + + m.def( + "clamp_swiglu_weighted_packed_backward", + &clamp_swiglu_weighted_packed_backward, + py::arg("dh"), + py::arg("gate_up"), + py::arg("p_s") = py::none(), + "P5 packed clamp_swiglu_weighted backward CUDA"); // Deterministic standard-softmax attention (issue #147) m.def( "deterministic_attention_forward", diff --git a/rl_engine/moe/cuda_provider.py b/rl_engine/moe/cuda_provider.py index f47164d3..598b4ae6 100644 --- a/rl_engine/moe/cuda_provider.py +++ b/rl_engine/moe/cuda_provider.py @@ -14,6 +14,8 @@ _FORWARD_SYMBOL = "clamp_swiglu_weighted_forward" _BACKWARD_SYMBOL = "clamp_swiglu_weighted_backward" +_PACKED_FORWARD_SYMBOL = "clamp_swiglu_weighted_packed_forward" +_PACKED_BACKWARD_SYMBOL = "clamp_swiglu_weighted_packed_backward" _SUPPORTED_DTYPES = ( torch.float16, @@ -33,7 +35,14 @@ def _require_cuda_extension() -> None: "ClampSwiGLUWeightedCudaProvider requires the compiled " "rl_engine._C extension." ) - missing = [name for name in (_FORWARD_SYMBOL, _BACKWARD_SYMBOL) if not hasattr(_C, name)] + required_symbols = ( + _FORWARD_SYMBOL, + _BACKWARD_SYMBOL, + _PACKED_FORWARD_SYMBOL, + _PACKED_BACKWARD_SYMBOL, + ) + + missing = [name for name in required_symbols if not hasattr(_C, name)] if missing: raise RuntimeError( @@ -57,6 +66,20 @@ def _validate_matrix( raise TypeError(f"{name} must have dtype fp16, bf16, or fp32, " f"got {x.dtype}.") +def _validate_packed_matrix( + gate_up: torch.Tensor, +) -> None: + _validate_matrix( + gate_up, + "gate_up", + ) + + if gate_up.shape[1] % 2 != 0: + raise ValueError( + "gate_up must have shape [rows, 2 * width], " f"got shape {tuple(gate_up.shape)}." + ) + + class ClampSwiGLUWeightedCudaProvider(ReferenceProvider): """P5-2 CUDA provider with deterministic FP32 computation.""" @@ -71,6 +94,7 @@ def capabilities(self) -> dict[str, Any]: "backend": "cuda", "delivered_ops": ["clamp_swiglu_weighted"], "geometry": ["one-row", "packed"], + "layouts": ["separate", "gate-up-packed"], "devices": ["cuda"], } @@ -182,3 +206,89 @@ def clamp_swiglu_weighted_bwd( dup, dp_s if p_s is not None else None, ) + + def clamp_swiglu_weighted_packed_fwd( + self, + gate_up: torch.Tensor, + p_s: torch.Tensor | None, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + _validate_packed_matrix(gate_up) + + if p_s is not None: + if p_s.device != gate_up.device: + raise RuntimeError( + "p_s and gate_up must share a CUDA device, " + f"got {p_s.device} and {gate_up.device}." + ) + + if p_s.dtype != torch.float32: + raise TypeError(f"p_s must have dtype fp32, got {p_s.dtype}.") + + if p_s.shape != (gate_up.shape[0],): + raise ValueError( + f"p_s must have shape ({gate_up.shape[0]},), " f"got {tuple(p_s.shape)}." + ) + + gate_up32 = gate_up.float().contiguous() + p_s32 = None if p_s is None else p_s.contiguous() + + (h,) = getattr( + _C, + _PACKED_FORWARD_SYMBOL, + )( + gate_up32, + p_s32, + ) + + saved = { + "gate_up32": gate_up32, + } + + if p_s32 is not None: + saved["p_s"] = p_s32 + + return h, saved + + def clamp_swiglu_weighted_packed_bwd( + self, + dh: torch.Tensor, + saved: dict[str, torch.Tensor], + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + ]: + if "gate_up32" not in saved: + raise KeyError("saved state is missing: gate_up32") + + _validate_matrix(dh, "dh") + + gate_up32 = saved["gate_up32"] + rows = gate_up32.shape[0] + width = gate_up32.shape[1] // 2 + + if dh.device != gate_up32.device: + raise RuntimeError( + "dh and gate_up32 must share a CUDA device, " + f"got {dh.device} and {gate_up32.device}." + ) + + if dh.shape != (rows, width): + raise ValueError(f"dh must have shape ({rows}, {width}), " f"got {tuple(dh.shape)}.") + + p_s = saved.get("p_s") + + dgate, dup, dp_s = getattr( + _C, + _PACKED_BACKWARD_SYMBOL, + )( + dh.contiguous(), + gate_up32, + p_s, + ) + + return ( + dgate, + dup, + dp_s if p_s is not None else None, + ) diff --git a/tests/test_p5_clamp_swiglu_weighted_cuda.py b/tests/test_p5_clamp_swiglu_weighted_cuda.py index 07777311..f12e7b09 100644 --- a/tests/test_p5_clamp_swiglu_weighted_cuda.py +++ b/tests/test_p5_clamp_swiglu_weighted_cuda.py @@ -179,6 +179,89 @@ def test_shared_variant_has_no_weight_no_clamp_and_no_dp_s() -> None: assert dp_ref is None +@pytest.mark.parametrize( + "weighted", + [ + False, + True, + ], +) +@requires_p5_cuda +def test_packed_layout_is_byte_exact_and_zero_copy( + weighted: bool, +) -> None: + provider = ClampSwiGLUWeightedCudaProvider() + + gate, up, p_s = (tensor.cuda() for tensor in fixtures.make_swiglu_boundary_inputs()) + + gate_up = torch.cat( + ( + gate, + up, + ), + dim=1, + ).contiguous() + + route_weights = p_s if weighted else None + + dh = fixtures.make_grad_output( + "swiglu_boundary", + tuple(gate.shape), + ).cuda() + + h_ref, saved_ref = oracle.clamp_swiglu_weighted_fwd( + gate, + up, + route_weights, + ) + + grads_ref = oracle.clamp_swiglu_weighted_bwd( + dh, + saved_ref, + ) + + h_got, saved_got = provider.clamp_swiglu_weighted_packed_fwd( + gate_up, + route_weights, + ) + + grads_got = provider.clamp_swiglu_weighted_packed_bwd( + dh, + saved_got, + ) + + _assert_exact( + h_got, + h_ref, + ) + + expected_saved_keys = { + "gate_up32", + } + + if route_weights is not None: + expected_saved_keys.add("p_s") + + assert set(saved_got) == expected_saved_keys + + assert gate_up.dtype == torch.float32 + assert saved_got["gate_up32"].data_ptr() == gate_up.data_ptr() + + for got, want in zip( + grads_got, + grads_ref, + strict=True, + ): + if want is None: + assert got is None + else: + assert got is not None + _assert_exact( + got, + want, + ) + + @requires_p5_cuda def test_repeated_runs_are_byte_exact() -> None: provider = ClampSwiGLUWeightedCudaProvider() From da6b81bc13f8c76a24d146f568e4cc525f50a879 Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Tue, 8 Sep 2026 21:07:59 -0700 Subject: [PATCH 12/13] perf(p5-5): det_gemm / tl.dot performance profiles alongside the strict path Adds two explicitly-selected performance providers that keep the P5-5 round positions (fc1 output and dX stay FP32; y/dh round once to BF16) and swap only the GEMM reduction order, per review feedback on #387: - shared-expert-cuda-det (p5-det-gemm-v1): reuses det_gemm_kernel.cu (fixed K-order, no split-K; TMA+mma.sync on SM90+, scalar K-tree fallback elsewhere) through two new FP32-output wrappers that expose the existing gemm_dispatch output_fp32 path. - shared-expert-triton-det (p5-triton-dot-v1): tl.dot with fixed 64x64x32 tiles, ascending-k, no autotune and no split-K. Both are deterministic and batch-invariant but not byte-equal to oracle-fp32-serial-v1; selecting them by provider name is the explicit opt-in, and the strict providers remain the oracle-parity gate. New tests pin repeat-run byte-equality, per-row batch invariance, and closeness to the oracle; the benchmark now reports all four backends. Signed-off-by: Yizheng Jiao --- benchmarks/benchmark_shared_expert_mlp.py | 11 ++- csrc/cuda/gemm/det_gemm_kernel.cu | 20 +++++ csrc/ops.cpp | 10 +++ rl_engine/_C.pyi | 5 ++ .../kernels/ops/triton/moe/shared_expert.py | 87 +++++++++++++++++++ rl_engine/moe/backends/shared_expert.py | 65 +++++++++++++- tests/test_shared_expert_mlp.py | 72 ++++++++++++++- 7 files changed, 261 insertions(+), 9 deletions(-) diff --git a/benchmarks/benchmark_shared_expert_mlp.py b/benchmarks/benchmark_shared_expert_mlp.py index 8ef6fe5a..6ed904de 100644 --- a/benchmarks/benchmark_shared_expert_mlp.py +++ b/benchmarks/benchmark_shared_expert_mlp.py @@ -76,14 +76,17 @@ def main() -> int: runners: dict[str, object] = {"torch-native": torch_native} strict: dict[str, object] = {} - for label, spec in ( - ("triton", "rl_engine.moe.backends.shared_expert:TritonSharedExpertProvider"), - ("cuda", "rl_engine.moe.backends.shared_expert:CudaSharedExpertProvider"), + for label, spec, is_strict in ( + ("triton", "rl_engine.moe.backends.shared_expert:TritonSharedExpertProvider", True), + ("cuda", "rl_engine.moe.backends.shared_expert:CudaSharedExpertProvider", True), + ("triton-det", "rl_engine.moe.backends.shared_expert:TritonDetSharedExpertProvider", False), + ("cuda-det", "rl_engine.moe.backends.shared_expert:CudaDetSharedExpertProvider", False), ): try: runner = make_runner(resolve_provider(spec)) runners[label] = runner - strict[label] = runner + if is_strict: + strict[label] = runner except NotImplementedError as exc: print(f"[skip] {label}: {exc}") diff --git a/csrc/cuda/gemm/det_gemm_kernel.cu b/csrc/cuda/gemm/det_gemm_kernel.cu index 1b535d2a..a907534e 100644 --- a/csrc/cuda/gemm/det_gemm_kernel.cu +++ b/csrc/cuda/gemm/det_gemm_kernel.cu @@ -546,6 +546,26 @@ torch::Tensor det_gemm_fwd_fp32(torch::Tensor a, torch::Tensor b) { return gemm_dispatch(a, b); } +// FP32-output variants: same kernels and reduction order, but the FP32 +// accumulator is stored without the final BF16 round. Used by operators whose +// contract keeps an intermediate in FP32 (e.g. P5-5 fc1 output, dX). +torch::Tensor det_gemm_fwd_out_fp32(torch::Tensor a, torch::Tensor b) { + check_in(a, "A"); check_in(b, "B"); + a = a.contiguous(); b = b.contiguous(); + TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "det_gemm_fwd_out_fp32: expect 2D [M,K]@[K,N]"); + TORCH_CHECK(b.size(0) == a.size(1), "det_gemm_fwd_out_fp32: K mismatch"); + return gemm_dispatch(a, b, RhsLayout::kKN, OutputLayout::kMN, /*output_fp32=*/true); +} + +torch::Tensor det_gemm_fwd_rhs_transposed_out_fp32(torch::Tensor a, torch::Tensor bt) { + check_in(a, "A"); check_in(bt, "Bt"); + a = a.contiguous(); bt = bt.contiguous(); + TORCH_CHECK(a.dim() == 2 && bt.dim() == 2, + "det_gemm_fwd_rhs_transposed_out_fp32: expect A[M,K] and Bt[N,K]"); + TORCH_CHECK(bt.size(1) == a.size(1), "det_gemm_fwd_rhs_transposed_out_fp32: K mismatch"); + return gemm_dispatch(a, bt, RhsLayout::kNK, OutputLayout::kMN, /*output_fp32=*/true); +} + torch::Tensor det_gemm_da(torch::Tensor dc, torch::Tensor b) { check_in(dc, "dC"); check_in(b, "B"); dc = dc.contiguous(); b = b.contiguous(); diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 33bb88bd..fb8f8872 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -115,6 +115,8 @@ void deterministic_collective_all_gather_fused( bool det_gemm_sm90_compiled(); torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b); torch::Tensor det_gemm_fwd_rhs_transposed(torch::Tensor a, torch::Tensor bt); +torch::Tensor det_gemm_fwd_out_fp32(torch::Tensor a, torch::Tensor b); +torch::Tensor det_gemm_fwd_rhs_transposed_out_fp32(torch::Tensor a, torch::Tensor bt); torch::Tensor det_gemm_da(torch::Tensor dc, torch::Tensor b); torch::Tensor det_gemm_db(torch::Tensor a, torch::Tensor dc); torch::Tensor det_gemm_db_transposed(torch::Tensor a, torch::Tensor dc); @@ -481,6 +483,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "det_gemm_fwd_rhs_transposed", &det_gemm_fwd_rhs_transposed, "Batch-invariant deterministic GEMM with physical Bt[N,K] (C=A@Bt^T)"); + m.def( + "det_gemm_fwd_out_fp32", + &det_gemm_fwd_out_fp32, + "det_gemm_fwd storing the FP32 accumulator (no final BF16 round)"); + m.def( + "det_gemm_fwd_rhs_transposed_out_fp32", + &det_gemm_fwd_rhs_transposed_out_fp32, + "det_gemm_fwd_rhs_transposed storing the FP32 accumulator"); m.def("det_gemm_da", &det_gemm_da, "Batch-invariant deterministic GEMM backward dA (dC@B^T)"); m.def("det_gemm_db", &det_gemm_db, "Batch-invariant deterministic GEMM backward dB (A^T@dC)"); m.def( diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 8f7879fb..fff3f125 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -194,6 +194,11 @@ def det_gemm_fwd_rhs_transposed( a: torch.Tensor, bt: torch.Tensor, ) -> torch.Tensor: ... +def det_gemm_fwd_out_fp32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: ... +def det_gemm_fwd_rhs_transposed_out_fp32( + a: torch.Tensor, + bt: torch.Tensor, +) -> torch.Tensor: ... def det_gemm_da(dc: torch.Tensor, b: torch.Tensor) -> torch.Tensor: ... def det_gemm_db(a: torch.Tensor, dc: torch.Tensor) -> torch.Tensor: ... def det_gemm_db_transposed(a: torch.Tensor, dc: torch.Tensor) -> torch.Tensor: ... diff --git a/rl_engine/kernels/ops/triton/moe/shared_expert.py b/rl_engine/kernels/ops/triton/moe/shared_expert.py index d65a10f4..211f2022 100644 --- a/rl_engine/kernels/ops/triton/moe/shared_expert.py +++ b/rl_engine/kernels/ops/triton/moe/shared_expert.py @@ -65,6 +65,47 @@ def _strict_gemm_kernel( acc = tld.add_rn(acc, tld.mul_rn(a, b)) tl.store(C + m * N + offs_n, acc, mask=mask_n) + @triton.jit + def _det_dot_gemm_kernel( + A, + B, + C, + M, + N, + K, + stride_bn, + stride_bk, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + # Performance GEMM (profile p5-triton-dot-v1): tensor-core tl.dot with + # FP32 accumulators, fixed constexpr tiles, ascending-k tile order and + # NO split-K -> deterministic and batch-invariant (a row's product uses + # only its own A row plus B; masked padding rows are zero). The + # reduction order inside a tile differs from the strict serial path, + # so this is NOT byte-equal to oracle-fp32-serial-v1. + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + for k0 in range(0, K, BLOCK_K): + offs_k = k0 + tl.arange(0, BLOCK_K) + a = tl.load( + A + offs_m[:, None] * K + offs_k[None, :], + mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), + other=0.0, + ) + b = tl.load( + B + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn, + mask=(offs_k[:, None] < K) & (offs_n[None, :] < N), + other=0.0, + ) + acc = tl.dot(a, b, acc) + c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + tl.store(C + offs_m[:, None] * N + offs_n[None, :], acc, mask=c_mask) + @triton.jit def _swiglu_shared_fwd_kernel( Z, @@ -161,6 +202,52 @@ def strict_gemm(a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor return out +def det_dot_gemm(a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: + """Performance GEMM (p5-triton-dot-v1): tl.dot tiles, FP32 out, no split-K. + + Same signature and round positions as :func:`strict_gemm`; only the + reduction order inside a tile differs (tensor-core MMA vs serial). + """ + if not TRITON_AVAILABLE: + raise NotImplementedError("triton is not installed (fail-closed, no fallback)") + _check_cuda_2d(a, torch.bfloat16, "a") + _check_cuda_2d(b, torch.bfloat16, "b") + m, k = a.shape + if trans_b: + bk, n = b.shape + stride_bn, stride_bk = 1, n + else: + n, bk = b.shape + stride_bn, stride_bk = k, 1 + if bk != k: + raise ValueError(f"K mismatch: a has K={k}, b has K={bk}") + out = torch.empty(m, n, dtype=torch.float32, device=a.device) + if out.numel() == 0: + return out + if k == 0: + return out.zero_() + # Fixed tiles (no autotune): tuning by shape would change the reduction + # tree with batch size and break invariance. + block_m, block_n, block_k = 64, 64, 32 + grid = (triton.cdiv(m, block_m), triton.cdiv(n, block_n)) + _det_dot_gemm_kernel[grid]( + a, + b, + out, + m, + n, + k, + stride_bn, + stride_bk, + BLOCK_M=block_m, + BLOCK_N=block_n, + BLOCK_K=block_k, + num_warps=4, + num_stages=3, + ) + return out + + def swiglu_shared_fwd(z: torch.Tensor) -> torch.Tensor: """One-round SwiGLU forward, shared mode: FP32 [T, 2F] -> BF16 [T, F].""" if not TRITON_AVAILABLE: diff --git a/rl_engine/moe/backends/shared_expert.py b/rl_engine/moe/backends/shared_expert.py index ef11df76..5b65af64 100644 --- a/rl_engine/moe/backends/shared_expert.py +++ b/rl_engine/moe/backends/shared_expert.py @@ -33,6 +33,11 @@ class _StrictSharedExpertProvider(ReferenceProvider): name = "shared-expert-strict" numeric_profile = ORACLE_PROFILE + # Performance profiles keep the contract's round positions but change the + # in-GEMM reduction order, so they are not byte-equal to the oracle. + # Selecting such a provider by name is the explicit opt-in; it is never a + # silent substitute for the strict path (fail-closed rule, P5-6). + strict_profile = True # Backend hooks ------------------------------------------------------- def _gemm(self, a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: @@ -68,7 +73,7 @@ def provenance(self) -> dict[str, Any]: def _check_batch(self, batch: SharedBatch) -> None: batch.validate() - if batch.numeric_profile != ORACLE_PROFILE: + if self.strict_profile and batch.numeric_profile != ORACLE_PROFILE: raise NotImplementedError( f"{self.name} only implements {ORACLE_PROFILE!r}, " f"got {batch.numeric_profile!r} (fail-closed, no fallback)" @@ -150,3 +155,61 @@ def _swiglu_fwd(self, z: torch.Tensor) -> torch.Tensor: def _swiglu_bwd(self, dh: torch.Tensor, z: torch.Tensor) -> torch.Tensor: return self._tk.swiglu_shared_bwd(dh, z) + + +class CudaDetSharedExpertProvider(CudaSharedExpertProvider): + """Performance CUDA backend: det_gemm (csrc/cuda/gemm/det_gemm_kernel.cu). + + Deterministic and batch-invariant (fixed K order, no split-K; TMA+mma.sync + on SM90+, scalar K-tree fallback elsewhere). Round positions match the + P5-5 contract (fc1 out and dX stay FP32; y/dh round once to BF16), but the + in-GEMM reduction order differs from the oracle, so outputs are close, not + byte-equal. SwiGLU stays on the strict CUDA core. + """ + + name = "shared-expert-cuda-det" + numeric_profile = "p5-det-gemm-v1" + strict_profile = False + + def _gemm(self, a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: + if trans_b: # b is the logical [K, N] operand + return self._ext.det_gemm_fwd_out_fp32(a, b) + return self._ext.det_gemm_fwd_rhs_transposed_out_fp32(a, b) + + def provenance(self) -> dict[str, Any]: + info = super().provenance() + info.update( + { + "split_k": 1, + "reduction": "fixed-k-tile-tree (det_gemm)", + "rounding": "FP32 accumulate, FMA/MMA inside tiles", + "sm90_tensor_core": bool(getattr(self._ext, "det_gemm_sm90_compiled")()), + } + ) + return info + + +class TritonDetSharedExpertProvider(TritonSharedExpertProvider): + """Performance Triton backend: tl.dot tiles with fixed geometry. + + Same guarantees and caveats as :class:`CudaDetSharedExpertProvider`, with + profile ``p5-triton-dot-v1`` (tile reduction order differs per backend). + """ + + name = "shared-expert-triton-det" + numeric_profile = "p5-triton-dot-v1" + strict_profile = False + + def _gemm(self, a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: + return self._tk.det_dot_gemm(a, b, trans_b) + + def provenance(self) -> dict[str, Any]: + info = super().provenance() + info.update( + { + "split_k": 1, + "reduction": "tl.dot 64x64x32 tiles, ascending-k", + "rounding": "FP32 accumulate, MMA inside tiles", + } + ) + return info diff --git a/tests/test_shared_expert_mlp.py b/tests/test_shared_expert_mlp.py index a114d711..e2927f1a 100644 --- a/tests/test_shared_expert_mlp.py +++ b/tests/test_shared_expert_mlp.py @@ -21,15 +21,42 @@ "triton": "rl_engine.moe.backends.shared_expert:TritonSharedExpertProvider", } +# Performance profiles: deterministic and batch-invariant, close to (but not +# byte-equal with) the oracle -- the in-GEMM reduction order differs. +PERF_PROVIDER_SPECS = { + "cuda-det": "rl_engine.moe.backends.shared_expert:CudaDetSharedExpertProvider", + "triton-det": "rl_engine.moe.backends.shared_expert:TritonDetSharedExpertProvider", +} -@pytest.fixture(params=sorted(PROVIDER_SPECS)) -def provider(request): + +def _resolve_or_skip(label: str, spec: str): from rl_engine.moe.provider import resolve_provider try: - return resolve_provider(PROVIDER_SPECS[request.param]) + return resolve_provider(spec) except NotImplementedError as exc: - pytest.skip(f"{request.param} backend unavailable: {exc}") + pytest.skip(f"{label} backend unavailable: {exc}") + + +@pytest.fixture(params=sorted(PROVIDER_SPECS)) +def provider(request): + return _resolve_or_skip(request.param, PROVIDER_SPECS[request.param]) + + +@pytest.fixture(params=sorted(PERF_PROVIDER_SPECS)) +def perf_provider(request): + return _resolve_or_skip(request.param, PERF_PROVIDER_SPECS[request.param]) + + +def _random_batch(t: int, hidden: int, ffn: int, seed: int = 2026): + gen = torch.Generator(device="cpu").manual_seed(seed) + batch = SharedBatch( + x=torch.randn(t, hidden, generator=gen).to(torch.bfloat16).cuda(), + w_fc1=(torch.randn(2 * ffn, hidden, generator=gen) / hidden**0.5).to(torch.bfloat16).cuda(), + w_fc2=(torch.randn(hidden, ffn, generator=gen) / ffn**0.5).to(torch.bfloat16).cuda(), + ) + dy = torch.randn(t, hidden, generator=gen).to(torch.bfloat16).cuda() + return batch, dy def _run_oracle(batch: SharedBatch, dy: torch.Tensor): @@ -146,3 +173,40 @@ def test_fail_closed_on_cpu_input(provider): batch = fixtures.make_shared_batch("shared_t1") # stays on CPU with pytest.raises(NotImplementedError): provider.shared_expert_mlp_fwd(batch) + + +@requires_cuda +def test_perf_backend_deterministic(perf_provider): + """Two runs of the performance profile are byte-identical.""" + batch, dy = _random_batch(256, 1024, 512) + runs = [_run_provider(perf_provider, batch, dy) for _ in range(2)] + (y_a, dx_a), (y_b, dx_b) = runs + assert tensor_sha256(y_a) == tensor_sha256(y_b) + assert tensor_sha256(dx_a) == tensor_sha256(dx_b) + + +@requires_cuda +def test_perf_backend_batch_invariance(perf_provider): + """fwd(x)[t] == fwd(x[t:t+1]) byte-for-byte also on the performance path.""" + batch, _ = _random_batch(16, 256, 128) + y_full, _ = perf_provider.shared_expert_mlp_fwd(batch) + for t in range(batch.x.shape[0]): + row_batch = SharedBatch( + x=batch.x[t : t + 1].contiguous(), + w_fc1=batch.w_fc1, + w_fc2=batch.w_fc2, + ) + y_row, _ = perf_provider.shared_expert_mlp_fwd(row_batch) + assert tensor_sha256(y_row) == tensor_sha256(y_full[t : t + 1]), f"row {t} diverged" + + +@requires_cuda +def test_perf_backend_close_to_oracle(perf_provider): + """Same round positions, different reduction order: close, not byte-equal.""" + batch, dy = _random_batch(64, 512, 256) + y_gold, saved_gold = oracle.shared_expert_mlp_fwd(batch) + dx_gold = oracle.shared_expert_mlp_bwd(dy, batch, saved_gold) + y, dx = _run_provider(perf_provider, batch, dy) + assert y.dtype == torch.bfloat16 and dx.dtype == torch.float32 + torch.testing.assert_close(y.float(), y_gold.float(), rtol=2e-2, atol=2e-2) + torch.testing.assert_close(dx, dx_gold, rtol=2e-2, atol=2e-2) From 7abd3134f21cd7842a6cf5cdb9d6972945c4f5b5 Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Tue, 8 Sep 2026 22:56:00 -0700 Subject: [PATCH 13/13] test(p5-5): per-profile oracle tolerances for the performance backends det_gemm merges K with a BF16 mid-split tree (its TP-equivalence design), so its deviation from the FP32-serial oracle is BF16-tree-sized; the tl.dot path keeps FP32 accumulators and only carries reduction-order noise. Each provider now declares its own closeness tolerance instead of sharing one number. Signed-off-by: Yizheng Jiao --- rl_engine/moe/backends/shared_expert.py | 7 +++++++ tests/test_shared_expert_mlp.py | 12 +++++++++--- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/rl_engine/moe/backends/shared_expert.py b/rl_engine/moe/backends/shared_expert.py index 5b65af64..ba156c28 100644 --- a/rl_engine/moe/backends/shared_expert.py +++ b/rl_engine/moe/backends/shared_expert.py @@ -170,6 +170,10 @@ class CudaDetSharedExpertProvider(CudaSharedExpertProvider): name = "shared-expert-cuda-det" numeric_profile = "p5-det-gemm-v1" strict_profile = False + # det_gemm rounds each BK=32 partial to BF16 and merges the K dimension + # with a BF16 mid-split tree (its TP-equivalence design), so its deviation + # from the FP32-serial oracle is BF16-tree-sized, not FP32-sized. + oracle_tolerance = {"rtol": 1e-1, "atol": 6e-2} def _gemm(self, a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: if trans_b: # b is the logical [K, N] operand @@ -199,6 +203,9 @@ class TritonDetSharedExpertProvider(TritonSharedExpertProvider): name = "shared-expert-triton-det" numeric_profile = "p5-triton-dot-v1" strict_profile = False + # Full-FP32 accumulators (only the contract's BF16 rounds), so deviation + # from the oracle is reduction-order noise only. + oracle_tolerance = {"rtol": 2e-2, "atol": 2e-2} def _gemm(self, a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: return self._tk.det_dot_gemm(a, b, trans_b) diff --git a/tests/test_shared_expert_mlp.py b/tests/test_shared_expert_mlp.py index e2927f1a..549f5ac9 100644 --- a/tests/test_shared_expert_mlp.py +++ b/tests/test_shared_expert_mlp.py @@ -202,11 +202,17 @@ def test_perf_backend_batch_invariance(perf_provider): @requires_cuda def test_perf_backend_close_to_oracle(perf_provider): - """Same round positions, different reduction order: close, not byte-equal.""" + """Same round positions, different reduction order: close, not byte-equal. + + Each provider declares its own tolerance: det_gemm merges the K dimension + with a BF16 mid-split tree (BF16-tree-sized deviation), while the tl.dot + path keeps FP32 accumulators (reduction-order noise only). + """ + tol = perf_provider.oracle_tolerance batch, dy = _random_batch(64, 512, 256) y_gold, saved_gold = oracle.shared_expert_mlp_fwd(batch) dx_gold = oracle.shared_expert_mlp_bwd(dy, batch, saved_gold) y, dx = _run_provider(perf_provider, batch, dy) assert y.dtype == torch.bfloat16 and dx.dtype == torch.float32 - torch.testing.assert_close(y.float(), y_gold.float(), rtol=2e-2, atol=2e-2) - torch.testing.assert_close(dx, dx_gold, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(y.float(), y_gold.float(), **tol) + torch.testing.assert_close(dx, dx_gold, **tol)